Skip to content

fix: ensure fork remote in setup_workspace before implementation push #201

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
danchild wants to merge 2 commits into from
+104 −22
Closed

fix: ensure fork remote in setup_workspace before implementation push #201

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a7ee831
fix: ensure fork remote in setup_workspace before implementation push…
danchild Jul 24, 2026
759496b
fix: restore finally guard on GitHubClient.close and add fork-creatio…
danchild Jul 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
59 changes: 37 additions & 22 deletions src/forge/workflow/nodes/workspace_setup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down Expand Up @@ -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}")
Expand Down
67 changes: 67 additions & 0 deletions tests/unit/workflow/nodes/test_workspace_setup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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."""

Expand All @@ -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",
Expand All @@ -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)

Expand Down Expand Up @@ -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",
Expand All @@ -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)

Expand All @@ -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",
Expand All @@ -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)

Expand All @@ -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",
Expand All @@ -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)

Expand All @@ -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",
Expand All @@ -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)

Expand All @@ -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",
Expand All @@ -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)
Expand All @@ -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."""

Expand Down
Loading