diff --git a/products/tasks/backend/temporal/execute_sandbox/workflow.py b/products/tasks/backend/temporal/execute_sandbox/workflow.py index 3f45d6beceeb..fbc2b62daedb 100644 --- a/products/tasks/backend/temporal/execute_sandbox/workflow.py +++ b/products/tasks/backend/temporal/execute_sandbox/workflow.py @@ -462,7 +462,8 @@ async def run(self, input: ExecuteSandboxInput) -> ExecuteSandboxOutput: self._posthog_mcp_scopes = input.posthog_mcp_scopes await self._update_task_run_status("in_progress") - await self._emit_progress("sandbox", "in_progress", "Setting up sandbox", "setup") + sandbox_label = "Restoring sandbox" if self.context.is_snapshot_resume else "Setting up sandbox" + await self._emit_progress("sandbox", "in_progress", sandbox_label, "setup") await self._track_workflow_event( "task_run_started", { diff --git a/products/tasks/backend/temporal/process_task/activities/get_task_processing_context.py b/products/tasks/backend/temporal/process_task/activities/get_task_processing_context.py index d05aeca6e84c..54722de3aace 100644 --- a/products/tasks/backend/temporal/process_task/activities/get_task_processing_context.py +++ b/products/tasks/backend/temporal/process_task/activities/get_task_processing_context.py @@ -133,6 +133,12 @@ def github_read_access(self) -> bool: def sandbox_environment_id(self) -> str | None: return (self.state or {}).get("sandbox_environment_id") + @property + def is_snapshot_resume(self) -> bool: + state = self.state or {} + has_resume_source = isinstance(state.get("resume_from_run_id"), str) or state.get("handoff_resumed") is True + return has_resume_source and isinstance(state.get("snapshot_external_id"), str) + @property def loop_id(self) -> str | None: """Set when this run was spawned by a loop firing (see products/tasks/backend/facade/loops.py).""" diff --git a/products/tasks/backend/temporal/process_task/activities/tests/test_get_task_processing_context.py b/products/tasks/backend/temporal/process_task/activities/tests/test_get_task_processing_context.py index bcd6f576f6e0..9f496e803a20 100644 --- a/products/tasks/backend/temporal/process_task/activities/tests/test_get_task_processing_context.py +++ b/products/tasks/backend/temporal/process_task/activities/tests/test_get_task_processing_context.py @@ -38,6 +38,45 @@ VM_FLAG_PAYLOAD_TARGET = "products.tasks.backend.constants.posthoganalytics.get_feature_flag_payload" +@pytest.mark.parametrize( + "state,expected", + [ + ({}, False), + ({"resume_from_run_id": "previous-run"}, False), + ({"handoff_resumed": True}, False), + ({"snapshot_external_id": "snapshot-id"}, False), + ( + { + "resume_from_run_id": "previous-run", + "snapshot_external_id": "snapshot-id", + }, + True, + ), + ( + { + "handoff_resumed": True, + "snapshot_external_id": "snapshot-id", + }, + True, + ), + ], +) +def test_snapshot_resume_requires_a_resume_marker_and_snapshot(state: dict[str, str | bool], expected: bool): + context = TaskProcessingContext( + task_id="task-id", + run_id="run-id", + team_id=1, + team_uuid="team-uuid", + organization_id="organization-id", + github_integration_id=None, + repository=None, + distinct_id="distinct-id", + state=state, + ) + + assert context.is_snapshot_resume is expected + + @pytest.mark.requires_secrets class TestIsAgentOtelTelemetryEnabled: @pytest.mark.parametrize( diff --git a/products/tasks/backend/temporal/process_task/workflow.py b/products/tasks/backend/temporal/process_task/workflow.py index 9147fd5cf9cb..731b0dc15f05 100644 --- a/products/tasks/backend/temporal/process_task/workflow.py +++ b/products/tasks/backend/temporal/process_task/workflow.py @@ -1024,7 +1024,8 @@ async def _provision_and_start_agent(self, input: ProcessTaskInput, run_id: str) # Announce the first progress step immediately so the desktop card # shows up before any provisioning log lines arrive. - await self._emit_progress("sandbox", "in_progress", "Setting up sandbox", "setup") + sandbox_label = "Restoring sandbox" if self.context.is_snapshot_resume else "Setting up sandbox" + await self._emit_progress("sandbox", "in_progress", sandbox_label, "setup") await self._track_workflow_event( "task_run_started",