From b0f18f63a3cdc706414431357d92a633a6e7b028 Mon Sep 17 00:00:00 2001 From: eshulman2 Date: Sun, 9 Aug 2026 16:26:09 +0300 Subject: [PATCH] Prevent stale workflow workspaces --- src/forge/orchestrator/worker.py | 19 +++++++ src/forge/workflow/nodes/workspace_setup.py | 2 +- src/forge/workspace/manager.py | 53 ++++++++++--------- tests/unit/orchestrator/test_worker.py | 43 ++++++++++++++- .../workflow/nodes/test_workspace_setup.py | 51 +++++++++++++----- tests/unit/workspace/test_manager_cleanup.py | 20 +++++++ 6 files changed, 150 insertions(+), 38 deletions(-) diff --git a/src/forge/orchestrator/worker.py b/src/forge/orchestrator/worker.py index 74c72682b..0089b72c6 100644 --- a/src/forge/orchestrator/worker.py +++ b/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, diff --git a/src/forge/workflow/nodes/workspace_setup.py b/src/forge/workflow/nodes/workspace_setup.py index 78a143088..b88dcba66 100644 --- a/src/forge/workflow/nodes/workspace_setup.py +++ b/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 diff --git a/src/forge/workspace/manager.py b/src/forge/workspace/manager.py index 7bc39a935..c47bdead8 100644 --- a/src/forge/workspace/manager.py +++ b/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. diff --git a/tests/unit/orchestrator/test_worker.py b/tests/unit/orchestrator/test_worker.py index 824c6906d..c5e19caba 100644 --- a/tests/unit/orchestrator/test_worker.py +++ b/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 diff --git a/tests/unit/workflow/nodes/test_workspace_setup.py b/tests/unit/workflow/nodes/test_workspace_setup.py index 016134810..0b4567981 100644 --- a/tests/unit/workflow/nodes/test_workspace_setup.py +++ b/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,13 +594,13 @@ 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), @@ -575,10 +608,7 @@ def transient_rmtree(path, *args, **kwargs): "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,10 +637,10 @@ 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), @@ -618,10 +648,7 @@ def persistent_rmtree(path, *args, **kwargs): "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"), ): diff --git a/tests/unit/workspace/test_manager_cleanup.py b/tests/unit/workspace/test_manager_cleanup.py index e016e7162..5fa31b582 100644 --- a/tests/unit/workspace/test_manager_cleanup.py +++ b/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