From 3505588a609380d29bf6b60b7bc772b7d4fb8546 Mon Sep 17 00:00:00 2001 From: eshulman2 Date: Wed, 29 Jul 2026 20:25:44 +0300 Subject: [PATCH 1/3] fix: clean up unregistered workspaces safely --- src/forge/workflow/nodes/workspace_setup.py | 59 ++++++++++++++ tests/sandbox/test_task_execution.py | 86 +++++++++++++++++++++ 2 files changed, 145 insertions(+) diff --git a/src/forge/workflow/nodes/workspace_setup.py b/src/forge/workflow/nodes/workspace_setup.py index 58fb1564..db45eef0 100644 --- a/src/forge/workflow/nodes/workspace_setup.py +++ b/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,47 @@ 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 + ): + 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( { diff --git a/tests/sandbox/test_task_execution.py b/tests/sandbox/test_task_execution.py index bca99bbd..1b0a8f0c 100644 --- a/tests/sandbox/test_task_execution.py +++ b/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,88 @@ 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"}) + ) + 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"}) + ) + 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() From 8ea25f49348709e8e66acf2c72e009b95f7db576 Mon Sep 17 00:00:00 2001 From: eshulman2 Date: Thu, 30 Jul 2026 10:12:00 +0300 Subject: [PATCH 2/3] Address review feedback: symlink test, comment, sort_keys consistency - Add test for symlink-to-valid-workspace rejection - Add clarifying comment at the raise ValueError site - Match sort_keys=True in test identity file writes Co-Authored-By: Claude Opus 4.6 (1M context) --- src/forge/workflow/nodes/workspace_setup.py | 2 ++ tests/sandbox/test_task_execution.py | 35 +++++++++++++++++++-- 2 files changed, 35 insertions(+), 2 deletions(-) diff --git a/src/forge/workflow/nodes/workspace_setup.py b/src/forge/workflow/nodes/workspace_setup.py index db45eef0..48c5c017 100644 --- a/src/forge/workflow/nodes/workspace_setup.py +++ b/src/forge/workflow/nodes/workspace_setup.py @@ -446,6 +446,8 @@ async def teardown_workspace(state: WorkflowState) -> WorkflowState: 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}" ) diff --git a/tests/sandbox/test_task_execution.py b/tests/sandbox/test_task_execution.py index 1b0a8f0c..ca04e4f7 100644 --- a/tests/sandbox/test_task_execution.py +++ b/tests/sandbox/test_task_execution.py @@ -286,7 +286,7 @@ async def test_teardown_workspace_destroys_unregistered_state_path( forge_dir = workspace_path / ".forge" forge_dir.mkdir() (forge_dir / "workspace.json").write_text( - json.dumps({"ticket_key": "TASK-123", "repo_name": "acme/backend"}) + json.dumps({"ticket_key": "TASK-123", "repo_name": "acme/backend"}, sort_keys=True) ) state = _make_state(workspace_path=str(workspace_path)) mock_manager = MagicMock() @@ -342,7 +342,7 @@ async def test_teardown_workspace_rejects_other_repo_identity( forge_dir = workspace_path / ".forge" forge_dir.mkdir() (forge_dir / "workspace.json").write_text( - json.dumps({"ticket_key": "TASK-123", "repo_name": "acme/frontend"}) + 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() @@ -355,3 +355,34 @@ async def test_teardown_workspace_rejects_other_repo_identity( 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() From 07c03b2c40f6e5c2742027d73aff46fdce233370 Mon Sep 17 00:00:00 2001 From: eshulman2 Date: Thu, 30 Jul 2026 10:22:55 +0300 Subject: [PATCH 3/3] fix: persist identity for rebase workspaces --- src/forge/workflow/nodes/rebase.py | 10 ++++- src/forge/workflow/nodes/workspace_setup.py | 6 +-- tests/unit/workflow/nodes/test_rebase.py | 47 +++++++++++++++++++++ 3 files changed, 59 insertions(+), 4 deletions(-) create mode 100644 tests/unit/workflow/nodes/test_rebase.py diff --git a/src/forge/workflow/nodes/rebase.py b/src/forge/workflow/nodes/rebase.py index e6b1ef42..2c759c67 100644 --- a/src/forge/workflow/nodes/rebase.py +++ b/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"): diff --git a/src/forge/workflow/nodes/workspace_setup.py b/src/forge/workflow/nodes/workspace_setup.py index 48c5c017..15f28b3d 100644 --- a/src/forge/workflow/nodes/workspace_setup.py +++ b/src/forge/workflow/nodes/workspace_setup.py @@ -28,7 +28,7 @@ _WORKSPACE_IDENTITY_FILE = ".forge/workspace.json" -def _write_workspace_identity(path: Path, *, ticket_key: str, repo_name: str) -> None: +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) @@ -105,7 +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) + 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 @@ -335,7 +335,7 @@ 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( + write_workspace_identity( workspace.path, ticket_key=ticket_key, repo_name=current_repo, diff --git a/tests/unit/workflow/nodes/test_rebase.py b/tests/unit/workflow/nodes/test_rebase.py new file mode 100644 index 00000000..d7142869 --- /dev/null +++ b/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()