diff --git a/src/forge/workflow/nodes/implementation.py b/src/forge/workflow/nodes/implementation.py index 9e3cf2b4..f6429f21 100644 --- a/src/forge/workflow/nodes/implementation.py +++ b/src/forge/workflow/nodes/implementation.py @@ -11,7 +11,9 @@ - Orchestrator (this node) handles git push after container exits """ +import asyncio import logging +import shutil from pathlib import Path from forge.config import get_settings @@ -31,6 +33,23 @@ logger = logging.getLogger(__name__) +_MAX_IMPLEMENTATION_RETRIES = 3 +_RETRY_BACKOFF_BASE_SECONDS = 1.0 + + +def _cleanup_failed_workspace(path: str | None) -> None: + """Best-effort removal of a workspace directory left by a failed prepare.""" + if not path: + return + workspace = Path(path) + if not workspace.exists(): + return + try: + shutil.rmtree(workspace, ignore_errors=True) + logger.info("Cleaned up failed workspace at %s", workspace) + except OSError as exc: + logger.warning("Failed to clean up workspace %s: %s", workspace, exc) + async def implement_task(state: WorkflowState) -> WorkflowState: """Implement a single Task using container sandbox. @@ -55,6 +74,33 @@ async def implement_task(state: WorkflowState) -> WorkflowState: implementation_node = _implementation_node_name(state) recorded_workspace = state.get("workspace_path") local_workspace_survived = bool(recorded_workspace and Path(recorded_workspace).exists()) + retry_count = state.get("retry_count", 0) + + # Hard stop before allocating another temp workspace when the graph-level + # retry budget is already exhausted. + if state.get("last_error") and retry_count >= _MAX_IMPLEMENTATION_RETRIES: + logger.error( + "Implementation retry limit (%s) already reached for %s; not preparing workspace", + _MAX_IMPLEMENTATION_RETRIES, + ticket_key, + ) + return update_state_timestamp( + { + **state, + "current_node": "escalate_blocked", + "last_error": state.get("last_error"), + } + ) + + if state.get("last_error") and retry_count > 0: + delay = min(_RETRY_BACKOFF_BASE_SECONDS * (2 ** (retry_count - 1)), 8.0) + logger.info( + "Backing off %.1fs before implementation retry %s for %s", + delay, + retry_count, + ticket_key, + ) + await asyncio.sleep(delay) try: git: GitOperations @@ -62,8 +108,13 @@ async def implement_task(state: WorkflowState) -> WorkflowState: state = {**state, "workspace_path": workspace_path} except Exception as exc: logger.error("Unable to prepare implementation workspace for %s: %s", ticket_key, exc) + # prepare_workspace cleans ephemeral dirs it creates; also drop a + # recorded path that no longer represents a usable workspace. + if recorded_workspace and not (Path(recorded_workspace).exists()): + _cleanup_failed_workspace(recorded_workspace) return { **state, + "workspace_path": None, "last_error": str(exc), "current_node": implementation_node, "retry_count": state.get("retry_count", 0) + 1, diff --git a/src/forge/workflow/nodes/workspace_setup.py b/src/forge/workflow/nodes/workspace_setup.py index 15f28b3d..d5424326 100644 --- a/src/forge/workflow/nodes/workspace_setup.py +++ b/src/forge/workflow/nodes/workspace_setup.py @@ -56,6 +56,9 @@ def _recreate_workspace_from_fork( workspace_obj = manager.create_workspace(repo_name=current_repo, ticket_key=ticket_key) target_path = workspace_obj.path stale_path = Path(stale_workspace_path) if stale_workspace_path else None + # create_workspace may mkdtemp a brand-new empty directory that must be + # removed if clone/checkout fails; otherwise retries fill /tmp. + created_empty_target = target_path.exists() and not any(target_path.iterdir()) # Build and validate the replacement beside the target. The existing # workspace may contain the only copy of an unpushed commit, so it must not @@ -77,6 +80,9 @@ def _recreate_workspace_from_fork( git.checkout_branch(branch_name, remote="fork") except Exception: shutil.rmtree(replacement_path, ignore_errors=True) + if created_empty_target and target_path.exists(): + shutil.rmtree(target_path, ignore_errors=True) + manager.destroy_workspace(workspace_obj) raise old_path = stale_path if stale_path and stale_path.exists() else None @@ -98,6 +104,8 @@ def _recreate_workspace_from_fork( if backup_path and backup_path.exists() and not target_path.exists(): backup_path.rename(target_path) shutil.rmtree(replacement_path, ignore_errors=True) + if created_empty_target and target_path.exists() and not any(target_path.iterdir()): + shutil.rmtree(target_path, ignore_errors=True) raise if backup_path and backup_path.exists(): diff --git a/tests/unit/workflow/nodes/test_workspace_recreate_cleanup.py b/tests/unit/workflow/nodes/test_workspace_recreate_cleanup.py new file mode 100644 index 00000000..e373865c --- /dev/null +++ b/tests/unit/workflow/nodes/test_workspace_recreate_cleanup.py @@ -0,0 +1,66 @@ +"""Regression tests for failed workspace recreation cleanup (#191).""" + +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + +from forge.workflow.nodes.workspace_setup import _recreate_workspace_from_fork + + +def test_recreate_cleans_empty_target_when_checkout_fails(tmp_path): + """Failed fork checkout must not leave orphan forge-* directories behind.""" + ticket = "AISOS-2119" + created_dirs: list[Path] = [] + + def fake_create_workspace(*, repo_name, ticket_key, branch_name=None): + path = Path(tmp_path) / f"forge-{ticket_key}-orphan" + path.mkdir(parents=True, exist_ok=True) + created_dirs.append(path) + ws = MagicMock() + ws.path = path + ws.repo_name = repo_name + ws.ticket_key = ticket_key + ws.branch_name = branch_name or f"forge/{ticket_key.lower()}" + return ws + + mock_git = MagicMock() + mock_git.clone = MagicMock() + mock_git.add_fork_remote = MagicMock() + mock_git.checkout_branch = MagicMock( + side_effect=RuntimeError( + "fatal: 'fork/forge/aisos-2119' is not a commit and a branch " + "'forge/aisos-2119' cannot be created from it" + ) + ) + + with ( + patch( + "forge.workflow.nodes.workspace_setup.get_settings", + return_value=MagicMock(workspace_base_dir=None), + ), + patch( + "forge.workflow.nodes.workspace_setup.WorkspaceManager" + ) as mock_manager_cls, + patch( + "forge.workflow.nodes.workspace_setup.GitOperations", + return_value=mock_git, + ), + ): + manager = MagicMock() + manager.create_workspace = MagicMock(side_effect=fake_create_workspace) + manager.destroy_workspace = MagicMock() + mock_manager_cls.return_value = manager + + with pytest.raises(RuntimeError, match="is not a commit"): + _recreate_workspace_from_fork( + ticket_key=ticket, + current_repo="org/repo", + branch_name="forge/aisos-2119", + fork_owner="bot", + fork_repo="repo", + ) + + assert created_dirs, "expected create_workspace to allocate a target dir" + assert not created_dirs[0].exists(), "orphan target directory must be removed" + manager.destroy_workspace.assert_called_once()