diff --git a/src/forge/workflow/nodes/workspace_setup.py b/src/forge/workflow/nodes/workspace_setup.py index 190a5b07..788b78ef 100644 --- a/src/forge/workflow/nodes/workspace_setup.py +++ b/src/forge/workflow/nodes/workspace_setup.py @@ -232,27 +232,39 @@ async def setup_workspace(state: WorkflowState) -> WorkflowState: git.clone() logger.info(f"Clone completed successfully for {current_repo}") - # Detect the upstream default branch (main, master, etc.) + # Detect the upstream default branch and ensure fork exists. default_branch = "main" + fork_owner = state.get("fork_owner", "") + fork_repo_name = state.get("fork_repo", "") + if current_repo and "/" in current_repo: owner, repo_name = current_repo.split("/", 1) github = GitHubClient() try: - repo_data = await github.get_repository(owner, repo_name) - default_branch = repo_data.get("default_branch", "main") - logger.info(f"Detected default branch for {current_repo}: {default_branch}") - except Exception as exc: - logger.warning(f"Could not detect default branch for {current_repo}: {exc}") + try: + repo_data = await github.get_repository(owner, repo_name) + default_branch = repo_data.get("default_branch", "main") + logger.info(f"Detected default branch for {current_repo}: {default_branch}") + except Exception as exc: + logger.warning(f"Could not detect default branch for {current_repo}: {exc}") + + if not fork_owner or not fork_repo_name: + try: + fork_data = await github.get_or_create_fork(owner, repo_name) + fork_owner = fork_data["owner"]["login"] + fork_repo_name = fork_data["name"] + await github.sync_fork_with_upstream(fork_owner, fork_repo_name) + logger.info( + f"Ensured fork {fork_owner}/{fork_repo_name} for workspace push" + ) + except Exception as exc: + logger.warning(f"Could not ensure fork for {current_repo}: {exc}") finally: await github.close() # Set up feature branch. - # If the workflow already created a PR (fork_owner/fork_repo in state), - # the branch lives on the fork. Add the fork remote, check whether the - # branch exists there, and check it out so we don't lose history. - fork_owner = state.get("fork_owner", "") - fork_repo_name = state.get("fork_repo", "") - + # If fork is available, add the fork remote so implementation pushes + # go to the fork instead of origin. if fork_owner and fork_repo_name: git.add_fork_remote(fork_owner, fork_repo_name) branch_exists_on_fork = git.remote_branch_exists(workspace.branch_name, remote="fork") @@ -298,16 +310,19 @@ async def setup_workspace(state: WorkflowState) -> WorkflowState: logger.info(f"Workspace ready: {workspace}") - return update_state_timestamp( - { - **state, - "workspace_path": str(workspace.path), - "current_repo": current_repo, - "context": context, - "current_node": "implementation", - "last_error": None, - } - ) + result_state = { + **state, + "workspace_path": str(workspace.path), + "current_repo": current_repo, + "context": context, + "current_node": "implementation", + "last_error": None, + } + if fork_owner and fork_repo_name: + result_state["fork_owner"] = fork_owner + result_state["fork_repo"] = fork_repo_name + + return update_state_timestamp(result_state) except Exception as e: logger.error(f"Workspace setup failed for {ticket_key}: {e}") diff --git a/tests/unit/workflow/nodes/test_workspace_setup.py b/tests/unit/workflow/nodes/test_workspace_setup.py index 3e42cd68..41e10c71 100644 --- a/tests/unit/workflow/nodes/test_workspace_setup.py +++ b/tests/unit/workflow/nodes/test_workspace_setup.py @@ -64,6 +64,18 @@ def create_mock_guardrails_loader(): return mock +def create_mock_github_client(): + """Create a mock GitHubClient for fork/repo operations.""" + mock = MagicMock() + mock.get_repository = AsyncMock(return_value={"default_branch": "main"}) + mock.get_or_create_fork = AsyncMock( + return_value={"owner": {"login": "fork-owner"}, "name": "test-repo"} + ) + mock.sync_fork_with_upstream = AsyncMock() + mock.close = AsyncMock() + return mock + + class TestWorkspaceSetupStatusComment: """Test cases for workspace setup posting status comments.""" @@ -74,6 +86,7 @@ async def test_workspace_setup_posts_status_comment(self): mock_manager, mock_workspace = create_mock_workspace_manager() mock_git = create_mock_git_operations() mock_guardrails_loader = create_mock_guardrails_loader() + mock_github = create_mock_github_client() state = create_initial_feature_state( ticket_key="TEST-123", @@ -89,6 +102,7 @@ async def test_workspace_setup_posts_status_comment(self): ), patch("forge.workflow.nodes.workspace_setup.GitOperations", return_value=mock_git), patch("forge.workflow.nodes.workspace_setup.GuardrailsLoader", mock_guardrails_loader), + patch("forge.workflow.nodes.workspace_setup.GitHubClient", return_value=mock_github), ): result = await setup_workspace(state) @@ -120,6 +134,7 @@ async def test_workspace_setup_uses_local_git_exclude_for_forge_dir(self, tmp_pa mock_jira = create_mock_jira_client() mock_git = create_mock_git_operations() mock_guardrails_loader = create_mock_guardrails_loader() + mock_github = create_mock_github_client() state = create_initial_feature_state( ticket_key="TEST-123", @@ -135,6 +150,7 @@ async def test_workspace_setup_uses_local_git_exclude_for_forge_dir(self, tmp_pa ), patch("forge.workflow.nodes.workspace_setup.GitOperations", return_value=mock_git), patch("forge.workflow.nodes.workspace_setup.GuardrailsLoader", mock_guardrails_loader), + patch("forge.workflow.nodes.workspace_setup.GitHubClient", return_value=mock_github), ): await setup_workspace(state) @@ -149,6 +165,7 @@ async def test_workspace_setup_handles_missing_repo_name(self): mock_manager, mock_workspace = create_mock_workspace_manager() mock_git = create_mock_git_operations() mock_guardrails_loader = create_mock_guardrails_loader() + mock_github = create_mock_github_client() state = create_initial_feature_state( ticket_key="TEST-456", @@ -164,6 +181,7 @@ async def test_workspace_setup_handles_missing_repo_name(self): ), patch("forge.workflow.nodes.workspace_setup.GitOperations", return_value=mock_git), patch("forge.workflow.nodes.workspace_setup.GuardrailsLoader", mock_guardrails_loader), + patch("forge.workflow.nodes.workspace_setup.GitHubClient", return_value=mock_github), ): await setup_workspace(state) @@ -184,6 +202,7 @@ async def test_workspace_setup_sets_implementing_label(self): mock_manager, mock_workspace = create_mock_workspace_manager() mock_git = create_mock_git_operations() mock_guardrails_loader = create_mock_guardrails_loader() + mock_github = create_mock_github_client() state = create_initial_feature_state( ticket_key="TEST-789", @@ -198,6 +217,7 @@ async def test_workspace_setup_sets_implementing_label(self): ), patch("forge.workflow.nodes.workspace_setup.GitOperations", return_value=mock_git), patch("forge.workflow.nodes.workspace_setup.GuardrailsLoader", mock_guardrails_loader), + patch("forge.workflow.nodes.workspace_setup.GitHubClient", return_value=mock_github), ): await setup_workspace(state) @@ -213,6 +233,7 @@ async def test_workspace_setup_transitions_tasks(self): mock_manager, mock_workspace = create_mock_workspace_manager() mock_git = create_mock_git_operations() mock_guardrails_loader = create_mock_guardrails_loader() + mock_github = create_mock_github_client() state = create_initial_feature_state( ticket_key="TEST-101", @@ -228,6 +249,7 @@ async def test_workspace_setup_transitions_tasks(self): ), patch("forge.workflow.nodes.workspace_setup.GitOperations", return_value=mock_git), patch("forge.workflow.nodes.workspace_setup.GuardrailsLoader", mock_guardrails_loader), + patch("forge.workflow.nodes.workspace_setup.GitHubClient", return_value=mock_github), ): await setup_workspace(state) @@ -249,6 +271,7 @@ async def test_workspace_setup_continues_on_jira_failure(self, caplog): mock_manager, mock_workspace = create_mock_workspace_manager() mock_git = create_mock_git_operations() mock_guardrails_loader = create_mock_guardrails_loader() + mock_github = create_mock_github_client() state = create_initial_feature_state( ticket_key="TEST-999", @@ -263,6 +286,7 @@ async def test_workspace_setup_continues_on_jira_failure(self, caplog): ), patch("forge.workflow.nodes.workspace_setup.GitOperations", return_value=mock_git), patch("forge.workflow.nodes.workspace_setup.GuardrailsLoader", mock_guardrails_loader), + patch("forge.workflow.nodes.workspace_setup.GitHubClient", return_value=mock_github), ): # Should not raise an exception result = await setup_workspace(state) @@ -279,6 +303,49 @@ async def test_workspace_setup_continues_on_jira_failure(self, caplog): mock_jira.close.assert_called_once() +class TestWorkspaceSetupForkCreation: + """Test that setup_workspace ensures fork when fork_owner/fork_repo not in state.""" + + @pytest.mark.asyncio + async def test_workspace_setup_creates_fork_when_not_in_state(self): + """Should call get_or_create_fork and persist fork info in state.""" + mock_jira = create_mock_jira_client() + mock_manager, mock_workspace = create_mock_workspace_manager() + mock_git = create_mock_git_operations() + mock_guardrails_loader = create_mock_guardrails_loader() + mock_github = create_mock_github_client() + + state = create_initial_feature_state( + ticket_key="TEST-FORK", + current_repo="upstream-org/my-repo", + ) + assert not state.get("fork_owner") + assert not state.get("fork_repo") + + with ( + patch("forge.workflow.nodes.workspace_setup.JiraClient", return_value=mock_jira), + patch( + "forge.workflow.nodes.workspace_setup.get_workspace_manager", + return_value=mock_manager, + ), + patch("forge.workflow.nodes.workspace_setup.GitOperations", return_value=mock_git), + patch("forge.workflow.nodes.workspace_setup.GuardrailsLoader", mock_guardrails_loader), + patch("forge.workflow.nodes.workspace_setup.GitHubClient", return_value=mock_github), + ): + result = await setup_workspace(state) + + mock_github.get_or_create_fork.assert_called_once_with( + "upstream-org", "my-repo" + ) + mock_github.sync_fork_with_upstream.assert_called_once_with( + "fork-owner", "test-repo" + ) + assert result["fork_owner"] == "fork-owner" + assert result["fork_repo"] == "test-repo" + mock_git.add_fork_remote.assert_called_once_with("fork-owner", "test-repo") + mock_github.close.assert_called_once() + + class TestPrepareWorkspaceRecovery: """Tests for prepare_workspace workspace sync/recreation behavior."""