From c255b5762082b5e551f9ab625fee0500ddcb6a00 Mon Sep 17 00:00:00 2001 From: eshulman2 Date: Sun, 26 Jul 2026 10:16:47 +0300 Subject: [PATCH] fix: preserve pending work during workspace recovery --- src/forge/workflow/nodes/implementation.py | 6 +- src/forge/workflow/nodes/local_reviewer.py | 6 +- .../workflow/nodes/task_takeover_execution.py | 6 +- src/forge/workflow/nodes/workspace_setup.py | 70 ++++++++++++++----- src/forge/workspace/git_ops.py | 5 ++ .../workflow/nodes/test_implementation.py | 27 +++++++ .../workflow/nodes/test_workspace_setup.py | 35 ++++++++++ 7 files changed, 136 insertions(+), 19 deletions(-) diff --git a/src/forge/workflow/nodes/implementation.py b/src/forge/workflow/nodes/implementation.py index 5514cc26..88d9846b 100644 --- a/src/forge/workflow/nodes/implementation.py +++ b/src/forge/workflow/nodes/implementation.py @@ -68,7 +68,11 @@ async def implement_task(state: WorkflowState) -> WorkflowState: "current_node": implementation_node, } - same_workspace_survived = local_workspace_survived and workspace_path == recorded_workspace + same_workspace_survived = ( + local_workspace_survived + and workspace_path == recorded_workspace + and git.workspace_recreated is not True + ) if state.get("implementation_push_pending") and same_workspace_survived: try: await push_to_fork_with_retry(git) diff --git a/src/forge/workflow/nodes/local_reviewer.py b/src/forge/workflow/nodes/local_reviewer.py index 26c1d0d7..eade16fc 100644 --- a/src/forge/workflow/nodes/local_reviewer.py +++ b/src/forge/workflow/nodes/local_reviewer.py @@ -114,7 +114,11 @@ async def local_review_changes(state: WorkflowState) -> WorkflowState: {**state, "current_node": "create_pr", "last_error": str(exc)} ) - same_workspace_survived = local_workspace_survived and workspace_path == recorded_workspace + same_workspace_survived = ( + local_workspace_survived + and workspace_path == recorded_workspace + and git.workspace_recreated is not True + ) if state.get("review_push_pending") and same_workspace_survived: try: await push_to_fork_with_retry(git) diff --git a/src/forge/workflow/nodes/task_takeover_execution.py b/src/forge/workflow/nodes/task_takeover_execution.py index 0a5b0747..b29d06b5 100644 --- a/src/forge/workflow/nodes/task_takeover_execution.py +++ b/src/forge/workflow/nodes/task_takeover_execution.py @@ -44,7 +44,11 @@ async def execute_task_changes(state: TaskTakeoverState) -> TaskTakeoverState: workspace_path, git = prepare_workspace(state) state = {**state, "workspace_path": workspace_path} - same_workspace_survived = local_workspace_survived and workspace_path == recorded_workspace + same_workspace_survived = ( + local_workspace_survived + and workspace_path == recorded_workspace + and git.workspace_recreated is not True + ) if state.get("implementation_push_pending") and same_workspace_survived: try: await push_to_fork_with_retry(git) diff --git a/src/forge/workflow/nodes/workspace_setup.py b/src/forge/workflow/nodes/workspace_setup.py index 190a5b07..476cf7eb 100644 --- a/src/forge/workflow/nodes/workspace_setup.py +++ b/src/forge/workflow/nodes/workspace_setup.py @@ -2,6 +2,7 @@ import logging import shutil +import tempfile from pathlib import Path from typing import Any @@ -38,24 +39,61 @@ def _recreate_workspace_from_fork( "missing branch_name, current_repo, fork_owner, or fork_repo in state" ) - if stale_workspace_path: - stale_path = Path(stale_workspace_path) - if stale_path.exists(): - logger.warning( - "Removing existing workspace for %s before recreating from fork: %s", - ticket_key, - stale_path, - ) - shutil.rmtree(stale_path) - manager = WorkspaceManager(base_dir=get_settings().workspace_base_dir) workspace_obj = manager.create_workspace(repo_name=current_repo, ticket_key=ticket_key) - git = GitOperations(workspace_obj) - git.clone() - git.add_fork_remote(fork_owner, fork_repo) - git.checkout_branch(branch_name, remote="fork") - logger.info(f"Workspace recreated at {workspace_obj.path} for {ticket_key}") - return str(workspace_obj.path), git + target_path = workspace_obj.path + stale_path = Path(stale_workspace_path) if stale_workspace_path else None + + # Build and validate the replacement beside the target. The existing + # workspace may contain the only copy of an unpushed commit, so it must not + # be removed merely because fetch/rebase failed. + target_path.parent.mkdir(parents=True, exist_ok=True) + replacement_path = Path( + tempfile.mkdtemp(prefix=f".{target_path.name}-replacement-", dir=target_path.parent) + ) + replacement_workspace = Workspace( + path=replacement_path, + repo_name=current_repo, + branch_name=branch_name, + ticket_key=ticket_key, + ) + git = GitOperations(replacement_workspace) + try: + git.clone() + git.add_fork_remote(fork_owner, fork_repo) + git.checkout_branch(branch_name, remote="fork") + except Exception: + shutil.rmtree(replacement_path, ignore_errors=True) + raise + + old_path = stale_path if stale_path and stale_path.exists() else None + backup_path: Path | None = None + try: + if old_path: + backup_path = Path( + tempfile.mkdtemp(prefix=f".{target_path.name}-old-", dir=target_path.parent) + ) + backup_path.rmdir() + old_path.rename(backup_path) + elif target_path.exists() and target_path != replacement_path: + # create_workspace creates the deterministic target directory when + # recovery starts without a stale workspace. + target_path.rmdir() + + replacement_path.rename(target_path) + except Exception: + if backup_path and backup_path.exists() and not target_path.exists(): + backup_path.rename(target_path) + shutil.rmtree(replacement_path, ignore_errors=True) + raise + + if backup_path and backup_path.exists(): + shutil.rmtree(backup_path) + + git.workspace.path = target_path + git.workspace_recreated = True + logger.info(f"Workspace recreated at {target_path} for {ticket_key}") + return str(target_path), git def prepare_workspace( diff --git a/src/forge/workspace/git_ops.py b/src/forge/workspace/git_ops.py index 6c8f8a7a..74f802a1 100644 --- a/src/forge/workspace/git_ops.py +++ b/src/forge/workspace/git_ops.py @@ -22,6 +22,11 @@ def __init__(self, workspace: Workspace): """ self.workspace = workspace self.settings = get_settings() + # Set by workspace recovery when this instance represents a replacement + # clone rather than the workspace recorded in workflow state. The path + # alone cannot identify that case because managed workspaces reuse a + # deterministic path for each ticket and repository. + self.workspace_recreated = False @property def repo_path(self) -> Path: diff --git a/tests/unit/workflow/nodes/test_implementation.py b/tests/unit/workflow/nodes/test_implementation.py index 90de8ba8..b2eed711 100644 --- a/tests/unit/workflow/nodes/test_implementation.py +++ b/tests/unit/workflow/nodes/test_implementation.py @@ -342,6 +342,33 @@ async def test_recreated_workspace_does_not_mark_pending_task_complete(self, tmp runner.run.assert_awaited_once() assert result["implementation_push_pending"] is False + @pytest.mark.asyncio + async def test_same_path_recreation_does_not_mark_pending_task_complete(self, tmp_path) -> None: + """A deterministic path does not prove that the original clone survived.""" + from forge.workflow.nodes.implementation import implement_task + + state = _make_state(workspace_path=str(tmp_path)) + state["implementation_push_pending"] = True + state["implementation_push_pending_task"] = "TASK-456" + mock_git = MagicMock() + mock_git.workspace_recreated = True + mock_jira = _make_mock_jira() + runner = _make_successful_runner() + + with ( + patch( + "forge.workflow.nodes.implementation.prepare_workspace", + return_value=(str(tmp_path), mock_git), + ), + patch("forge.workflow.nodes.implementation.JiraClient", return_value=mock_jira), + patch("forge.workflow.nodes.implementation.ContainerRunner", return_value=runner), + patch("forge.workflow.nodes.implementation.get_settings"), + ): + result = await implement_task(state) + + runner.run.assert_awaited_once() + assert result["implementation_push_pending"] is False + @pytest.mark.asyncio async def test_bug_container_failure_keeps_bug_implementation_node(self): """Bug container failures keep the bug graph retry node.""" diff --git a/tests/unit/workflow/nodes/test_workspace_setup.py b/tests/unit/workflow/nodes/test_workspace_setup.py index 3e42cd68..310920ef 100644 --- a/tests/unit/workflow/nodes/test_workspace_setup.py +++ b/tests/unit/workflow/nodes/test_workspace_setup.py @@ -319,3 +319,38 @@ def test_sync_failure_recreates_workspace_from_fork(self, tmp_path): new_git.clone.assert_called_once() new_git.add_fork_remote.assert_called_once_with("forge-bot", "repo") new_git.checkout_branch.assert_called_once_with("forge/test-123", remote="fork") + assert new_git.workspace_recreated is True + + 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" + workspace_path.mkdir() + local_commit = workspace_path / "local-commit.txt" + local_commit.write_text("not pushed yet") + + state = create_initial_feature_state( + ticket_key="TEST-124", + current_repo="org/repo", + workspace_path=str(workspace_path), + fork_owner="forge-bot", + fork_repo="repo", + context={"branch_name": "forge/test-124"}, + ) + + old_git = MagicMock() + old_git.pull_rebase.side_effect = RuntimeError("sync failed") + new_git = MagicMock() + new_git.clone.side_effect = RuntimeError("clone failed") + 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], + ), + pytest.raises(RuntimeError, match="clone failed"), + ): + prepare_workspace(state) + + assert local_commit.read_text() == "not pushed yet"