Prevent stale workflow workspaces

src/forge/orchestrator/worker.py

@@ -30,6 +30,7 @@
 from forge.skills.utils import extract_project_key
 from forge.utils.redaction import redact_secrets
 from forge.workflow.nodes.error_handler import notify_error
+from forge.workflow.nodes.workspace_setup import teardown_workspace
 from forge.workflow.pr_state import (
     activate_pull_request_for_event,
     all_pull_requests_merged,
@@ -83,6 +84,19 @@ async def _report_new_workflow_error(result: dict, error_before_invoke: str | No
     )
 
 
+async def _cleanup_terminal_workspace(result: dict[str, Any]) -> dict[str, Any]:
+    """Remove a workspace recreated after the normal post-PR teardown."""
+    if result.get("current_node") != "complete" or not result.get("workspace_path"):
+        return result
+
+    cleaned = await teardown_workspace(result)
+    return {
+        **cleaned,
+        "current_node": "complete",
+        "is_paused": False,
+    }
+
+
 _PRD_GATE_NODES = ("prd_approval_gate", "generate_prd", "regenerate_prd")
 _SPEC_GATE_NODES = ("spec_approval_gate", "generate_spec", "regenerate_spec")
 _REVIEW_GATES = ("human_review_gate", "review_response_gate")
@@ -486,6 +500,11 @@ async def _process_workflow(self, message: QueueMessage) -> None:
                 # Run the workflow from the beginning
                 result = await compiled_workflow.ainvoke(state, config=config)
 
+            cleaned_result = await _cleanup_terminal_workspace(result)
+            if cleaned_result != result:
+                await compiled_workflow.aupdate_state(config, cleaned_result)
+                result = cleaned_result
+
             # Nodes continue to use scalar PR fields as a compatibility view.
             # Persist that view back into the selected per-repository record
             # after every invocation so subsequent webhooks restore fresh CI,

src/forge/workflow/nodes/workspace_setup.py

@@ -45,7 +45,7 @@ def _remove_workspace_backup(path: Path) -> None:
     """Remove a replaced workspace, tolerating short-lived filesystem races."""
     for attempt in range(1, _BACKUP_CLEANUP_ATTEMPTS + 1):
         try:
-            shutil.rmtree(path)
+            WorkspaceManager.remove_path(path)
             return
         except FileNotFoundError:
             return

src/forge/workspace/manager.py

@@ -113,35 +113,40 @@ def destroy_workspace(self, workspace: Workspace) -> None:
         workspace_id = f"{workspace.ticket_key}:{workspace.repo_name}"
 
         if workspace.path.exists():
-            podman = shutil.which("podman")
-            if podman:
-                try:
-                    result = subprocess.run(
-                        [podman, "unshare", "rm", "-rf", str(workspace.path)],
-                        check=False,
-                    )
-                except OSError as exc:
-                    logger.warning(
-                        "podman unshare failed for %s, falling back to shutil: %s",
-                        workspace.path,
-                        exc,
-                    )
-                    shutil.rmtree(workspace.path)
-                else:
-                    if result.returncode != 0:
-                        logger.warning(
-                            "podman unshare failed for %s, falling back to shutil",
-                            workspace.path,
-                        )
-                        shutil.rmtree(workspace.path)
-            else:
-                shutil.rmtree(workspace.path)
+            self.remove_path(workspace.path)
             logger.info(f"Destroyed workspace: {workspace}")
 
         workspace.is_active = False
-        if workspace_id in self._workspaces:
+        if self._workspaces.get(workspace_id) is workspace:
             del self._workspaces[workspace_id]
 
+    @staticmethod
+    def remove_path(path: Path) -> None:
+        """Remove a workspace tree, including files owned by sandbox containers."""
+        podman = shutil.which("podman")
+        if podman:
+            try:
+                result = subprocess.run(
+                    [podman, "unshare", "rm", "-rf", str(path)],
+                    check=False,
+                    timeout=60,
+                )
+            except (OSError, subprocess.TimeoutExpired) as exc:
+                logger.warning(
+                    "podman unshare failed for %s, falling back to shutil: %s",
+                    path,
+                    exc,
+                )
+            else:
+                if result.returncode == 0 and not path.exists():
+                    return
+                logger.warning(
+                    "podman unshare failed to remove %s, falling back to shutil",
+                    path,
+                )
+
+        shutil.rmtree(path)
+
     def destroy_all_for_ticket(self, ticket_key: str) -> int:
         """Destroy all workspaces for a ticket.
 

tests/unit/orchestrator/test_worker.py

@@ -8,6 +8,7 @@
 from forge.models.events import EventSource
 from forge.orchestrator.worker import (
     OrchestratorWorker,
+    _cleanup_terminal_workspace,
     _has_new_reportable_error,
     _report_new_workflow_error,
 )
@@ -27,6 +28,44 @@ def test_has_new_reportable_error(result: dict, error_before_invoke: str | None,
     assert _has_new_reportable_error(result, error_before_invoke) is expected
 
 
+@pytest.mark.asyncio
+async def test_terminal_workflow_cleans_recreated_workspace():
+    result = {
+        "ticket_key": "TEST-1",
+        "current_node": "complete",
+        "workspace_path": "/tmp/forge-TEST-1-repo",
+        "is_paused": False,
+    }
+    torn_down = {
+        **result,
+        "workspace_path": None,
+        "current_node": "workspace_complete",
+    }
+
+    with patch(
+        "forge.orchestrator.worker.teardown_workspace",
+        AsyncMock(return_value=torn_down),
+    ) as teardown:
+        cleaned = await _cleanup_terminal_workspace(result)
+
+    teardown.assert_awaited_once_with(result)
+    assert cleaned["workspace_path"] is None
+    assert cleaned["current_node"] == "complete"
+
+
+@pytest.mark.asyncio
+async def test_nonterminal_workflow_keeps_workspace():
+    result = {
+        "current_node": "human_review_gate",
+        "workspace_path": "/tmp/forge-TEST-1-repo",
+    }
+
+    with patch("forge.orchestrator.worker.teardown_workspace", AsyncMock()) as teardown:
+        assert await _cleanup_terminal_workspace(result) is result
+
+    teardown.assert_not_awaited()
+
+
 @pytest.mark.asyncio
 async def test_report_new_workflow_error_posts_once():
     result = {
@@ -957,7 +996,9 @@ async def test_ensure_skills_skipped_gracefully_when_forge_skills_not_set(
         """
         ensure_skills_called = False
 
-        async def fake_ensure_skills_no_property(project_key, jira_client, _skills_dir, **_kw) -> None:
+        async def fake_ensure_skills_no_property(
+            project_key, jira_client, _skills_dir, **_kw
+        ) -> None:
             """Simulate ensure_skills when forge.skills property is absent (returns None)."""
             nonlocal ensure_skills_called
             ensure_skills_called = True

tests/unit/workflow/nodes/test_workspace_setup.py

@@ -12,6 +12,7 @@
 from forge.models.workflow import ForgeLabel
 from forge.workflow.feature.state import create_initial_feature_state
 from forge.workflow.nodes.workspace_setup import prepare_workspace, setup_workspace
+from forge.workspace.manager import WorkspaceManager
 
 
 def create_mock_jira_client():
@@ -507,6 +508,38 @@ def test_sync_failure_recreates_workspace_from_fork(self, tmp_path):
         new_git.checkout_branch.assert_called_once_with("forge/test-123", remote="fork")
         assert new_git.workspace_recreated is True
 
+    def test_recovery_uses_container_aware_cleanup_for_old_workspace(self, tmp_path):
+        workspace_path = tmp_path / "forge-TEST-125-org-repo"
+        workspace_path.mkdir()
+
+        state = create_initial_feature_state(
+            ticket_key="TEST-125",
+            current_repo="org/repo",
+            workspace_path=str(workspace_path),
+            fork_owner="forge-bot",
+            fork_repo="repo",
+            context={"branch_name": "forge/test-125"},
+        )
+
+        old_git = MagicMock()
+        old_git.pull_rebase.side_effect = RuntimeError("sync failed")
+        new_git = MagicMock()
+        settings = MagicMock(workspace_base_dir=str(tmp_path))
+
+        with (
+            patch("forge.workflow.nodes.workspace_setup.get_settings", return_value=settings),
+            patch(
+                "forge.workflow.nodes.workspace_setup.GitOperations",
+                side_effect=[old_git, new_git],
+            ),
+            patch.object(WorkspaceManager, "remove_path") as remove_path,
+        ):
+            result_path, _ = prepare_workspace(state)
+
+        assert result_path == str(workspace_path)
+        remove_path.assert_called_once()
+        assert "-old-" in remove_path.call_args.args[0].name
+
     def test_failed_replacement_preserves_existing_workspace(self, tmp_path):
         """A failed recovery clone must not delete the only local commit."""
         workspace_path = tmp_path / "forge-TEST-124-org-repo"
@@ -561,24 +594,21 @@ def test_backup_cleanup_retries_directory_not_empty(self, tmp_path):
         real_rmtree = shutil.rmtree
         cleanup_calls = 0
 
-        def transient_rmtree(path, *args, **kwargs):
+        def transient_remove(path):
             nonlocal cleanup_calls
             if Path(path).name.startswith(f".{workspace_path.name}-old-"):
                 cleanup_calls += 1
                 if cleanup_calls == 1:
                     raise OSError(errno.ENOTEMPTY, "Directory not empty", path)
-            return real_rmtree(path, *args, **kwargs)
+            return real_rmtree(path)
 
         with (
             patch("forge.workflow.nodes.workspace_setup.get_settings", return_value=settings),
             patch(
                 "forge.workflow.nodes.workspace_setup.GitOperations",
                 side_effect=[old_git, new_git],
             ),
-            patch(
-                "forge.workflow.nodes.workspace_setup.shutil.rmtree",
-                side_effect=transient_rmtree,
-            ),
+            patch.object(WorkspaceManager, "remove_path", side_effect=transient_remove),
             patch("forge.workflow.nodes.workspace_setup.time.sleep") as sleep,
         ):
             result_path, result_git = prepare_workspace(state)
@@ -607,21 +637,18 @@ def test_backup_cleanup_failure_does_not_fail_recovery(self, tmp_path, caplog):
         settings = MagicMock(workspace_base_dir=str(tmp_path))
         real_rmtree = shutil.rmtree
 
-        def persistent_rmtree(path, *args, **kwargs):
+        def persistent_remove(path):
             if Path(path).name.startswith(f".{workspace_path.name}-old-"):
                 raise OSError(errno.ENOTEMPTY, "Directory not empty", path)
-            return real_rmtree(path, *args, **kwargs)
+            return real_rmtree(path)
 
         with (
             patch("forge.workflow.nodes.workspace_setup.get_settings", return_value=settings),
             patch(
                 "forge.workflow.nodes.workspace_setup.GitOperations",
                 side_effect=[old_git, new_git],
             ),
-            patch(
-                "forge.workflow.nodes.workspace_setup.shutil.rmtree",
-                side_effect=persistent_rmtree,
-            ),
+            patch.object(WorkspaceManager, "remove_path", side_effect=persistent_remove),
             patch("forge.workflow.nodes.workspace_setup.time.sleep"),
             caplog.at_level("WARNING"),
         ):

tests/unit/workspace/test_manager_cleanup.py

@@ -29,6 +29,7 @@ def test_uses_podman_unshare_when_available(self, tmp_path):
         mock_run.assert_called_once_with(
             ["/usr/bin/podman", "unshare", "rm", "-rf", str(ws.path)],
             check=False,
+            timeout=60,
         )
         assert not ws.is_active
 
@@ -90,3 +91,22 @@ def test_workspace_marked_inactive_and_removed_from_registry(self, tmp_path):
 
         assert not ws.is_active
         assert "T-1:org/repo" not in manager._workspaces
+
+    def test_destroying_superseded_workspace_keeps_current_registry_entry(self, tmp_path):
+        old = _workspace(tmp_path)
+        current = Workspace(
+            path=tmp_path / "current",
+            repo_name=old.repo_name,
+            branch_name=old.branch_name,
+            ticket_key=old.ticket_key,
+        )
+        manager = WorkspaceManager()
+        manager._workspaces["T-1:org/repo"] = current
+
+        with (
+            patch("forge.workspace.manager.shutil.which", return_value=None),
+            patch("forge.workspace.manager.shutil.rmtree"),
+        ):
+            manager.destroy_workspace(old)
+
+        assert manager.get_workspace("T-1", "org/repo") is current