Fix leaked workspaces after worker restarts

src/forge/workflow/nodes/rebase.py

@@ -15,7 +15,10 @@
 from forge.prompts import load_prompt
 from forge.sandbox import ContainerRunner
 from forge.workflow.feature.state import FeatureState as WorkflowState
-from forge.workflow.nodes.workspace_setup import get_workspace_manager
+from forge.workflow.nodes.workspace_setup import (
+    get_workspace_manager,
+    write_workspace_identity,
+)
 from forge.workflow.utils import merge_review_exhaustion, update_state_timestamp
 from forge.workflow.utils.jira_status import post_status_comment
 from forge.workspace.git_ops import GitOperations
@@ -61,6 +64,11 @@ async def rebase_pr(state: WorkflowState) -> WorkflowState:
         workspace = manager.create_workspace(repo_name=current_repo, ticket_key=ticket_key)
         git = GitOperations(workspace)
         git.clone()
+        write_workspace_identity(
+            workspace.path,
+            ticket_key=ticket_key,
+            repo_name=current_repo,
+        )
         git.add_fork_remote(fork_owner, fork_repo)
 
         if git.remote_branch_exists(workspace.branch_name, remote="fork"):

src/forge/workflow/nodes/workspace_setup.py

@@ -1,5 +1,6 @@
 """Workspace setup node for LangGraph workflow."""
 
+import json
 import logging
 import shutil
 import tempfile
@@ -24,6 +25,17 @@
 
 logger = logging.getLogger(__name__)
 
+_WORKSPACE_IDENTITY_FILE = ".forge/workspace.json"
+
+
+def write_workspace_identity(path: Path, *, ticket_key: str, repo_name: str) -> None:
+    """Persist identity used to safely recover a workspace after process loss."""
+    identity_path = path / _WORKSPACE_IDENTITY_FILE
+    identity_path.parent.mkdir(parents=True, exist_ok=True)
+    identity_path.write_text(
+        json.dumps({"ticket_key": ticket_key, "repo_name": repo_name}, sort_keys=True) + "\n"
+    )
+
 
 def _recreate_workspace_from_fork(
     *,
@@ -93,6 +105,7 @@ def _recreate_workspace_from_fork(
 
     git.workspace.path = target_path
     git.workspace_recreated = True
+    write_workspace_identity(target_path, ticket_key=ticket_key, repo_name=current_repo)
     logger.info(f"Workspace recreated at {target_path} for {ticket_key}")
     return str(target_path), git
 
@@ -322,6 +335,11 @@ async def setup_workspace(state: WorkflowState) -> WorkflowState:
         forge_dir = workspace.path / ".forge"
         forge_dir.mkdir(exist_ok=True)
         (forge_dir / "history").mkdir(exist_ok=True)
+        write_workspace_identity(
+            workspace.path,
+            ticket_key=ticket_key,
+            repo_name=current_repo,
+        )
 
         # Keep Forge handoff files local to this clone without modifying the
         # target repository's tracked .gitignore.
@@ -399,6 +417,49 @@ async def teardown_workspace(state: WorkflowState) -> WorkflowState:
         if workspace:
             manager.destroy_workspace(workspace)
             logger.info(f"Workspace destroyed: {workspace}")
+        else:
+            # The manager registry is process-local. A restart or worker handoff
+            # can therefore leave a valid workspace on disk without a registry
+            # entry. Fall back to the path persisted in workflow state, but only
+            # after verifying that it is one of Forge's workspace directories.
+            path = Path(workspace_path)
+            settings = get_settings()
+            allowed_parent = (
+                Path(settings.workspace_base_dir)
+                if settings.workspace_base_dir
+                else Path(tempfile.gettempdir())
+            ).resolve()
+            resolved_path = path.resolve()
+            expected_prefix = f"forge-{ticket_key}-"
+            identity_path = resolved_path / _WORKSPACE_IDENTITY_FILE
+            try:
+                identity = json.loads(identity_path.read_text())
+            except (OSError, json.JSONDecodeError):
+                identity = None
+
+            if (
+                path.is_symlink()
+                or not path.is_dir()
+                or resolved_path.parent != allowed_parent
+                or not resolved_path.name.startswith(expected_prefix)
+                or not isinstance(identity, dict)
+                or identity.get("ticket_key") != ticket_key
+                or identity.get("repo_name") != current_repo
+            ):
+                # Caught by the outer except — records last_error and
+                # clears workspace_path without deleting the directory.
+                raise ValueError(
+                    f"Refusing to destroy unrecognized workspace path: {workspace_path}"
+                )
+
+            recovered_workspace = Workspace(
+                path=resolved_path,
+                repo_name=current_repo,
+                branch_name=state.get("context", {}).get("branch_name", ""),
+                ticket_key=ticket_key,
+            )
+            manager.destroy_workspace(recovered_workspace)
+            logger.info("Destroyed unregistered workspace: %s", resolved_path)
 
         return update_state_timestamp(
             {

tests/sandbox/test_task_execution.py

@@ -1,5 +1,6 @@
 """Integrated and sandbox tests for task execution in container environments."""
 
+import json
 import tempfile
 from collections.abc import Generator
 from pathlib import Path
@@ -269,3 +270,119 @@ async def test_teardown_workspace_secure_destruction(self, mock_get_manager: Mag
         assert teardown_state["current_node"] == "workspace_complete"
         mock_manager.get_workspace.assert_called_once_with("TASK-123", "acme/backend")
         mock_manager.destroy_workspace.assert_called_once_with(mock_workspace)
+
+    @pytest.mark.asyncio
+    @patch("forge.workflow.nodes.workspace_setup.get_settings")
+    @patch("forge.workflow.nodes.workspace_setup.get_workspace_manager")
+    async def test_teardown_workspace_destroys_unregistered_state_path(
+        self,
+        mock_get_manager: MagicMock,
+        mock_get_settings: MagicMock,
+        tmp_path: Path,
+    ) -> None:
+        """A worker restart must not leak a workspace missing from the registry."""
+        workspace_path = tmp_path / "forge-TASK-123-restarted"
+        workspace_path.mkdir()
+        forge_dir = workspace_path / ".forge"
+        forge_dir.mkdir()
+        (forge_dir / "workspace.json").write_text(
+            json.dumps({"ticket_key": "TASK-123", "repo_name": "acme/backend"}, sort_keys=True)
+        )
+        state = _make_state(workspace_path=str(workspace_path))
+        mock_manager = MagicMock()
+        mock_manager.get_workspace.return_value = None
+        mock_get_manager.return_value = mock_manager
+        mock_get_settings.return_value.workspace_base_dir = str(tmp_path)
+
+        teardown_state = await teardown_workspace(state)
+
+        assert teardown_state["workspace_path"] is None
+        recovered = mock_manager.destroy_workspace.call_args.args[0]
+        assert recovered.path == workspace_path
+        assert recovered.ticket_key == "TASK-123"
+        assert recovered.repo_name == "acme/backend"
+
+    @pytest.mark.asyncio
+    @patch("forge.workflow.nodes.workspace_setup.get_settings")
+    @patch("forge.workflow.nodes.workspace_setup.get_workspace_manager")
+    async def test_teardown_workspace_rejects_unrecognized_state_path(
+        self,
+        mock_get_manager: MagicMock,
+        mock_get_settings: MagicMock,
+        tmp_path: Path,
+    ) -> None:
+        """The registry fallback cannot recursively delete an arbitrary path."""
+        workspace_path = tmp_path / "not-a-forge-workspace"
+        workspace_path.mkdir()
+        state = _make_state(workspace_path=str(workspace_path))
+        mock_manager = MagicMock()
+        mock_manager.get_workspace.return_value = None
+        mock_get_manager.return_value = mock_manager
+        mock_get_settings.return_value.workspace_base_dir = str(tmp_path)
+
+        teardown_state = await teardown_workspace(state)
+
+        assert workspace_path.exists()
+        assert teardown_state["workspace_path"] is None
+        assert "Refusing to destroy unrecognized workspace path" in teardown_state["last_error"]
+        mock_manager.destroy_workspace.assert_not_called()
+
+    @pytest.mark.asyncio
+    @patch("forge.workflow.nodes.workspace_setup.get_settings")
+    @patch("forge.workflow.nodes.workspace_setup.get_workspace_manager")
+    async def test_teardown_workspace_rejects_other_repo_identity(
+        self,
+        mock_get_manager: MagicMock,
+        mock_get_settings: MagicMock,
+        tmp_path: Path,
+    ) -> None:
+        """Cross-repo feature state cannot delete a different repo's workspace."""
+        workspace_path = tmp_path / "forge-TASK-123-other-repo"
+        workspace_path.mkdir()
+        forge_dir = workspace_path / ".forge"
+        forge_dir.mkdir()
+        (forge_dir / "workspace.json").write_text(
+            json.dumps({"ticket_key": "TASK-123", "repo_name": "acme/frontend"}, sort_keys=True)
+        )
+        state = _make_state(workspace_path=str(workspace_path), current_repo="acme/backend")
+        mock_manager = MagicMock()
+        mock_manager.get_workspace.return_value = None
+        mock_get_manager.return_value = mock_manager
+        mock_get_settings.return_value.workspace_base_dir = str(tmp_path)
+
+        teardown_state = await teardown_workspace(state)
+
+        assert workspace_path.exists()
+        assert "Refusing to destroy unrecognized workspace path" in teardown_state["last_error"]
+        mock_manager.destroy_workspace.assert_not_called()
+
+    @pytest.mark.asyncio
+    @patch("forge.workflow.nodes.workspace_setup.get_settings")
+    @patch("forge.workflow.nodes.workspace_setup.get_workspace_manager")
+    async def test_teardown_workspace_rejects_symlink_to_valid_workspace(
+        self,
+        mock_get_manager: MagicMock,
+        mock_get_settings: MagicMock,
+        tmp_path: Path,
+    ) -> None:
+        """A symlink pointing at a valid workspace must be rejected."""
+        real_path = tmp_path / "forge-TASK-123-real"
+        real_path.mkdir()
+        forge_dir = real_path / ".forge"
+        forge_dir.mkdir()
+        (forge_dir / "workspace.json").write_text(
+            json.dumps({"ticket_key": "TASK-123", "repo_name": "acme/backend"}, sort_keys=True)
+        )
+        symlink_path = tmp_path / "forge-TASK-123-link"
+        symlink_path.symlink_to(real_path)
+        state = _make_state(workspace_path=str(symlink_path))
+        mock_manager = MagicMock()
+        mock_manager.get_workspace.return_value = None
+        mock_get_manager.return_value = mock_manager
+        mock_get_settings.return_value.workspace_base_dir = str(tmp_path)
+
+        teardown_state = await teardown_workspace(state)
+
+        assert real_path.exists()
+        assert "Refusing to destroy unrecognized workspace path" in teardown_state["last_error"]
+        mock_manager.destroy_workspace.assert_not_called()

tests/unit/workflow/nodes/test_rebase.py

@@ -0,0 +1,47 @@
+"""Tests for the pull-request rebase node."""
+
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+from forge.workflow.nodes.rebase import rebase_pr
+
+
+@pytest.mark.asyncio
+async def test_rebase_workspace_persists_identity_after_clone(tmp_path: Path) -> None:
+    """A restarted worker must be able to safely tear down a rebase workspace."""
+    workspace = SimpleNamespace(
+        path=tmp_path / "forge-TASK-123-acme-backend",
+        branch_name="forge/task-123",
+    )
+    workspace.path.mkdir()
+    manager = MagicMock()
+    manager.create_workspace.return_value = workspace
+    git = MagicMock()
+    git.remote_branch_exists.return_value = False
+    jira = MagicMock(close=AsyncMock())
+    github = MagicMock(close=AsyncMock())
+    state = {
+        "ticket_key": "TASK-123",
+        "current_repo": "acme/backend",
+        "fork_owner": "forge-bot",
+        "fork_repo": "backend",
+        "current_pr_number": 237,
+        "rebase_return_node": "ci_evaluator",
+    }
+
+    with (
+        patch("forge.workflow.nodes.rebase.get_settings"),
+        patch("forge.workflow.nodes.rebase.JiraClient", return_value=jira),
+        patch("forge.workflow.nodes.rebase.GitHubClient", return_value=github),
+        patch("forge.workflow.nodes.rebase.get_workspace_manager", return_value=manager),
+        patch("forge.workflow.nodes.rebase.GitOperations", return_value=git),
+    ):
+        await rebase_pr(state)
+
+    assert (workspace.path / ".forge" / "workspace.json").read_text() == (
+        '{"repo_name": "acme/backend", "ticket_key": "TASK-123"}\n'
+    )
+    git.clone.assert_called_once_with()