diff --git a/products/tasks/backend/temporal/process_task/activities/provision_sandbox.py b/products/tasks/backend/temporal/process_task/activities/provision_sandbox.py index 4c2860f5b1b8..685f7a2776c7 100644 --- a/products/tasks/backend/temporal/process_task/activities/provision_sandbox.py +++ b/products/tasks/backend/temporal/process_task/activities/provision_sandbox.py @@ -24,7 +24,7 @@ get_primary_sandbox_jwt_kid, get_sandbox_jwt_public_key, ) -from products.tasks.backend.logic.services.sandbox import Sandbox, SandboxConfig, SandboxTemplate +from products.tasks.backend.logic.services.sandbox import ExecutionResult, Sandbox, SandboxConfig, SandboxTemplate from products.tasks.backend.logic.services.sandbox_usage import open_sandbox_session from products.tasks.backend.models import SandboxSnapshot, Task, TaskRun from products.tasks.backend.temporal.metrics import ( @@ -689,12 +689,37 @@ def clone_repository_in_sandbox(input: CloneRepositoryInSandboxInput) -> CloneRe branch=ctx.branch if is_resume else None, ) + if is_resume and ctx.branch and _is_missing_remote_branch_clone_error(clone_result): + emit_agent_log( + ctx.run_id, + "debug", + f"Resume branch {ctx.branch} is unavailable; cloning the repository default branch so the agent can restore its git checkpoint", + ) + clone_result = sandbox.clone_repository( + input.repository, + github_token=input.github_token, + shallow=input.shallow_clone, + branch=None, + ) + if clone_result.exit_code != 0: raise RuntimeError(f"Failed to clone repository {input.repository}: {clone_result.stderr}") return CloneRepositoryInSandboxOutput(clone_ms=clone_timer.elapsed_ms) +def _is_missing_remote_branch_clone_error(result: ExecutionResult) -> bool: + if result.exit_code == 0: + return False + + output = f"{result.stdout}\n{result.stderr}".casefold() + return ( + "could not find remote branch" in output + or ("remote branch" in output and "not found in upstream origin" in output) + or "couldn't find remote ref" in output + ) + + @activity.defn @asyncify def checkout_branch_in_sandbox(input: CheckoutBranchInSandboxInput) -> CheckoutBranchInSandboxOutput: diff --git a/products/tasks/backend/temporal/process_task/activities/tests/test_provision_sandbox.py b/products/tasks/backend/temporal/process_task/activities/tests/test_provision_sandbox.py index 54de6b15a2d8..6af3e5e2c3ac 100644 --- a/products/tasks/backend/temporal/process_task/activities/tests/test_provision_sandbox.py +++ b/products/tasks/backend/temporal/process_task/activities/tests/test_provision_sandbox.py @@ -52,3 +52,57 @@ def test_clone_repository_uses_saved_branch_only_for_resumes(mocker, activity_en shallow=True, branch=expected_branch, ) + + +def test_resume_clone_falls_back_to_default_branch_when_saved_branch_is_missing(mocker, activity_environment): + context = TaskProcessingContext( + task_id="task-id", + run_id="run-id", + team_id=1, + team_uuid="team-uuid", + organization_id="organization-id", + github_integration_id=123, + repository="posthog/posthog", + distinct_id="distinct-id", + state={"resume_from_run_id": "previous-run-id"}, + _branch="branch-from-a-sibling-repository", + ) + sandbox = mocker.Mock() + sandbox.clone_repository.side_effect = [ + ExecutionResult( + stdout="", + stderr=( + "warning: Could not find remote branch branch-from-a-sibling-repository to clone.\n" + "fatal: Remote branch branch-from-a-sibling-repository not found in upstream origin" + ), + exit_code=128, + ), + ExecutionResult(stdout="", stderr="", exit_code=0), + ] + mocker.patch.object(Sandbox, "get_by_id", return_value=sandbox) + + async_to_sync(activity_environment.run)( + clone_repository_in_sandbox, + CloneRepositoryInSandboxInput( + context=context, + sandbox_id="sandbox-id", + repository="posthog/posthog", + github_token="github-token", + shallow_clone=True, + ), + ) + + assert sandbox.clone_repository.call_args_list == [ + mocker.call( + "posthog/posthog", + github_token="github-token", + shallow=True, + branch="branch-from-a-sibling-repository", + ), + mocker.call( + "posthog/posthog", + github_token="github-token", + shallow=True, + branch=None, + ), + ]