diff --git a/.semgrep/rules/security/idor-team-scoped-models.yaml b/.semgrep/rules/security/idor-team-scoped-models.yaml index b3768d8ca814..330f88c223de 100644 --- a/.semgrep/rules/security/idor-team-scoped-models.yaml +++ b/.semgrep/rules/security/idor-team-scoped-models.yaml @@ -289,6 +289,7 @@ rules: |TaskAutomation |TaskPresence |TaskRun + |TaskSession |TaskThreadMessage |TaskThreadMessageMention |Channel @@ -606,6 +607,7 @@ rules: |TaskAutomation |TaskPresence |TaskRun + |TaskSession |TaskThreadMessage |TaskThreadMessageMention |Channel diff --git a/posthog/test/setup_receivers_baseline.txt b/posthog/test/setup_receivers_baseline.txt index 313b5d1aa263..c4e0e915e52a 100644 --- a/posthog/test/setup_receivers_baseline.txt +++ b/posthog/test/setup_receivers_baseline.txt @@ -91,6 +91,7 @@ post_delete:products.feature_flags.backend.local_evaluation.feature_flag_changed post_delete:products.feature_flags.backend.models.feature_flag.refresh_flag_cache_on_updates post_delete:products.slack_app.backend.signals.invalidate_repo_list_on_user_github_change post_delete:products.surveys.backend.models.survey_changed +post_delete:products.tasks.backend.models.delete_task_session_object post_delete:products.workflows.backend.models.hog_flow.hog_flow.hog_flow_deleted post_init:posthog.storage.gateway_credential_signal_handlers._snapshot_oauth post_init:posthog.storage.gateway_credential_signal_handlers._snapshot_secret_key diff --git a/products/tasks/backend/constants.py b/products/tasks/backend/constants.py index bcd9d0737352..56c0f9895b08 100644 --- a/products/tasks/backend/constants.py +++ b/products/tasks/backend/constants.py @@ -8,6 +8,7 @@ MODAL_VM_SANDBOX_FEATURE_FLAG = "tasks-modal-vm-sandbox" MODAL_NETWORK_ALLOWLIST_FEATURE_FLAG = "tasks-modal-network-allowlist" AGENT_RUN_OTEL_TELEMETRY_FEATURE_FLAG = "tasks-agent-run-otel-telemetry" +PI_CLOUD_RUNTIME_FEATURE_FLAG = "pi-harness" # Run-state key the telemetry flag decision is stamped under at dispatch (temporal/client.py). # Consumers read the stamp, so the decision stays stable for the run's whole lifetime. AGENT_OTEL_TELEMETRY_STATE_KEY = "agent_otel_telemetry_enabled" @@ -66,6 +67,8 @@ def vm_sandbox_allowed_origins(*, distinct_id: str, organization_id: str) -> set MAX_CUSTOM_IMAGES_PER_TEAM = 20 MAX_CUSTOM_IMAGES_PER_USER = 10 +TASK_SESSION_MAX_SIZE_BYTES = 10 * 1024 * 1024 +TASK_SESSION_UPLOAD_FORM_OVERHEAD_BYTES = 64 * 1024 MODAL_DIRECTORY_RESUME_SNAPSHOTS_FEATURE_FLAG = "tasks-modal-directory-resume-snapshots" STREAM_VIA_PROXY_FEATURE_FLAG = "tasks-stream-via-proxy" diff --git a/products/tasks/backend/facade/api.py b/products/tasks/backend/facade/api.py index a7199e03d551..e4b95fd2a24f 100644 --- a/products/tasks/backend/facade/api.py +++ b/products/tasks/backend/facade/api.py @@ -16,6 +16,7 @@ """ import re +import hashlib import logging from collections.abc import Iterable, Sequence from concurrent.futures import ThreadPoolExecutor @@ -40,7 +41,9 @@ AGENT_OTEL_TELEMETRY_STATE_KEY, MAX_CUSTOM_IMAGES_PER_TEAM, MAX_CUSTOM_IMAGES_PER_USER, + PI_CLOUD_RUNTIME_FEATURE_FLAG, RESERVED_SANDBOX_ENVIRONMENT_VARIABLE_KEYS, + TASK_SESSION_MAX_SIZE_BYTES, is_blocked_sandbox_env_key, ) from products.tasks.backend.error_telemetry import truncate_error_message @@ -57,11 +60,13 @@ CodeInviteRedemption, SandboxCustomImage, SandboxEnvironment, + SandboxSession, SandboxSnapshot, Task, TaskActivity, TaskAutomation, TaskRun, + TaskSession, TaskThreadMessage, TaskThreadMessageMention, ) @@ -105,6 +110,7 @@ "TaskRunEnvironment", "TaskRunStatus", "append_task_run_log", + "ensure_task_run_session", "beacon_task_presence", "bootstrap_task_run", "can_mint_readonly_github_token", @@ -150,6 +156,8 @@ "get_task_detail", "get_task_id_for_run", "get_task_run", + "get_task_run_session", + "sync_task_run_session", "get_task_run_detail", "get_task_run_sandbox_connection", "get_task_run_living_artifact", @@ -171,6 +179,7 @@ "list_task_repositories", "list_task_runs", "list_tasks", + "pi_cloud_runtime_enabled", "prepare_task_run_artifact_uploads", "prepare_task_staged_artifacts", "presign_task_run_artifact", @@ -208,6 +217,7 @@ "update_task_run_state", "upsert_internal_sandbox_env", "validate_set_output", + "validate_task_run_sandbox_token", "validate_task_run_artifact_ids", "warm_task_sandbox", ] @@ -2204,6 +2214,148 @@ def append_task_run_log( return _task_run_detail_to_dto(run) +def ensure_task_run_session(run_id: str | UUID) -> UUID: + with transaction.atomic(): + run = TaskRun.objects.select_for_update(of=("self",)).select_related("task__team").get(id=run_id) + if run.active_task_session_id is not None: + return run.active_task_session_id + + task_session = TaskSession.create_for_task(run.task) + run.active_task_session = task_session + run.save(update_fields=["active_task_session", "updated_at"]) + return task_session.id + + +def get_task_run_session( + run_id: str | UUID, task_id: str | UUID, team_id: int +) -> tuple[UUID, str | None, str | None] | None: + from posthog.storage import object_storage # noqa: PLC0415 + + run = _get_visible_run(run_id, task_id, team_id) + if run is None or run.active_task_session_id is None: + return None + task_session = TaskSession.objects.unscoped().get(id=run.active_task_session_id) + if task_session.object_storage_key is None: + return task_session.id, None, None + download_url = object_storage.get_presigned_url(task_session.object_storage_key, expiration=3600) + if not download_url: + raise RuntimeError("Unable to prepare task session download") + return task_session.id, download_url, task_session.content_sha256 + + +def _validate_task_session_content(content: bytes) -> None: + if not content or len(content) > TASK_SESSION_MAX_SIZE_BYTES: + raise ValueError("The task session content size is invalid") + + +def _get_open_sandbox_session(run_id: UUID, sandbox_id: str) -> SandboxSession | None: + return ( + SandboxSession.objects.unscoped() + .filter( + task_run_id=run_id, + sandbox_id=sandbox_id, + ended_at__isnull=True, + ) + .first() + ) + + +def _delete_task_session_object(task_session_id: UUID, object_storage_key: str) -> None: + from posthog.storage import object_storage # noqa: PLC0415 + + try: + object_storage.delete(object_storage_key) + except Exception as error: + logger.warning( + "task_session.failed_to_delete_object", + extra={ + "task_session_id": str(task_session_id), + "object_storage_key": object_storage_key, + "error": str(error), + }, + ) + + +def validate_task_run_sandbox_token( + token: str, + run_id: str | UUID, + task_id: str | UUID, + team_id: int, + sandbox_id: str, +) -> bool: + from jwt import InvalidTokenError # noqa: PLC0415 + + from products.tasks.backend.logic.services.connection_token import ( # noqa: PLC0415 + validate_sandbox_event_ingest_token, + ) + + try: + claims = validate_sandbox_event_ingest_token(token) + except (InvalidTokenError, ValueError): + return False + return ( + claims.run_id == str(run_id) + and claims.task_id == str(task_id) + and claims.team_id == team_id + and claims.sandbox_id == sandbox_id + ) + + +def sync_task_run_session( + run_id: str | UUID, + task_id: str | UUID, + team_id: int, + *, + sandbox_id: str, + expected_content_sha256: str | None, + content: bytes, +) -> tuple[UUID, str] | None: + from posthog.storage import object_storage # noqa: PLC0415 + + visible_run = _get_visible_run(run_id, task_id, team_id) + if visible_run is None or visible_run.active_task_session_id is None: + return None + _validate_task_session_content(content) + if _get_open_sandbox_session(visible_run.id, sandbox_id) is None: + raise ValueError("The task session writer is not the active sandbox") + + content_sha256 = hashlib.sha256(content).hexdigest() + object_storage_key = ( + f"task-sessions/{visible_run.task.team.organization_id}/{visible_run.task_id}/" + f"{visible_run.active_task_session_id}/{uuid4()}.jsonl" + ) + object_storage.write(object_storage_key, content) + + previous_object_storage_key: str | None = None + try: + with transaction.atomic(): + locked_run = TaskRun.objects.select_for_update(of=("self",)).get(id=visible_run.id) + if locked_run.active_task_session_id != visible_run.active_task_session_id: + raise ValueError("The task session sync is stale") + if _get_open_sandbox_session(locked_run.id, sandbox_id) is None: + raise ValueError("The task session writer is not the active sandbox") + + task_session = TaskSession.objects.unscoped().select_for_update().get(id=locked_run.active_task_session_id) + if task_session.content_sha256 == content_sha256: + transaction.on_commit(lambda: _delete_task_session_object(task_session.id, object_storage_key)) + return task_session.id, content_sha256 + if task_session.content_sha256 != expected_content_sha256: + raise ValueError("The task session content is stale") + + previous_object_storage_key = task_session.object_storage_key + task_session.object_storage_key = object_storage_key + task_session.content_sha256 = content_sha256 + task_session.size = len(content) + task_session.save(update_fields=["object_storage_key", "content_sha256", "size", "updated_at"]) + if previous_object_storage_key is not None: + transaction.on_commit(lambda: _delete_task_session_object(task_session.id, previous_object_storage_key)) + task_session.tag_object() + return task_session.id, content_sha256 + except Exception: + _delete_task_session_object(visible_run.active_task_session_id, object_storage_key) + raise + + def task_run_has_slack_mapping(run_id: str | UUID, task_id: str | UUID, team_id: int) -> bool | None: """Whether a run is mapped to a Slack thread. ``None`` if the run isn't found.""" from products.slack_app.backend.models import ( # noqa: PLC0415 — cross-product import kept off the api import path @@ -2691,7 +2843,11 @@ def create_task_run_stream_read_token(run_id: str | UUID, task_id: str | UUID, t return _create(task_run=run) -def resolve_stream_base_url(*, distinct_id: str, organization_id: str | UUID) -> str | None: +def task_uses_pi_runtime(task_id: str | UUID, team_id: int) -> bool: + return Task.objects.filter(id=task_id, team_id=team_id, runtime=Task.Runtime.PI).exists() + + +def resolve_stream_base_url(*, distinct_id: str, organization_id: str | UUID, force_proxy: bool = False) -> str | None: """Agent-proxy base URL for the read leg, or ``None`` to read from Django directly. Returns the configured agent-proxy URL only when it is set for this environment AND the @@ -2705,7 +2861,7 @@ def resolve_stream_base_url(*, distinct_id: str, organization_id: str | UUID) -> return None # Local dev disables the analytics SDK, so the rollout flag never evaluates; the URL setting # is the opt-in there. Prod (DEBUG off) still gates on the flag below. - if settings.DEBUG: + if settings.DEBUG or force_proxy: return proxy_url try: enabled = bool( @@ -3209,7 +3365,13 @@ def bootstrap_task_run( def _trigger_task_processing_workflow( - task: Task, run: TaskRun, user_id: int | None, *, raise_on_error: bool = False + task: Task, + run: TaskRun, + user_id: int | None, + *, + initial_message: str | None = None, + initial_artifact_ids: list[str] | None = None, + raise_on_error: bool = False, ) -> None: from products.tasks.backend.temporal.client import ( # noqa: PLC0415 — keep temporalio off the api import path execute_task_processing_workflow, @@ -3218,6 +3380,7 @@ def _trigger_task_processing_workflow( RunSource, parse_run_state, ) + from products.tasks.backend.temporal.process_task.workflow import PendingFollowup # noqa: PLC0415 # SIGNAL_REPORT: implementation runs log their work on the report (notes, code references) # via the task:write artefact tools. @@ -3226,13 +3389,28 @@ def _trigger_task_processing_workflow( posthog_mcp_scopes: Literal["read_only", "full"] = "full" if run_source in full_mcp_run_sources else "read_only" try: logger.info("Attempting to trigger task processing workflow for task %s, run %s", task.id, run.id) - execute_task_processing_workflow( - task_id=str(task.id), - run_id=str(run.id), - team_id=task.team.id, - user_id=user_id, - posthog_mcp_scopes=posthog_mcp_scopes, - ) + if initial_message or initial_artifact_ids: + execute_task_processing_workflow( + task_id=str(task.id), + run_id=str(run.id), + team_id=task.team.id, + user_id=user_id, + posthog_mcp_scopes=posthog_mcp_scopes, + initial_message=PendingFollowup( + message=initial_message, + artifact_ids=initial_artifact_ids or [], + actor_user_id=user_id, + message_id=str(uuid4()), + ), + ) + else: + execute_task_processing_workflow( + task_id=str(task.id), + run_id=str(run.id), + team_id=task.team.id, + user_id=user_id, + posthog_mcp_scopes=posthog_mcp_scopes, + ) logger.info("Workflow trigger completed for task %s, run %s", task.id, run.id) except Exception as e: logger.exception("Failed to trigger task processing workflow for task %s, run %s: %s", task.id, run.id, e) @@ -3253,8 +3431,6 @@ def check_task_run_startable(run_id: str | UUID, task_id: str | UUID, team_id: i run = _get_visible_run(run_id, task_id, team_id) if run is None: return "not_found" - if run.task.runtime == Task.Runtime.PI: - return "unsupported_runtime" if run.environment != TaskRun.Environment.CLOUD: return "not_cloud" if run.status not in _STARTABLE_TASK_RUN_STATUSES: @@ -3290,10 +3466,11 @@ def start_task_run( return "missing_artifacts:" + ",".join(missing_artifact_ids), None state_updates: dict = {} - if pending_user_message is not None: - state_updates["pending_user_message"] = pending_user_message - if pending_user_artifact_ids: - state_updates["pending_user_artifact_ids"] = pending_user_artifact_ids + if task.runtime != Task.Runtime.PI: + if pending_user_message is not None: + state_updates["pending_user_message"] = pending_user_message + if pending_user_artifact_ids: + state_updates["pending_user_artifact_ids"] = pending_user_artifact_ids previous_state = dict(run.state or {}) try: @@ -3301,7 +3478,16 @@ def start_task_run( TaskRun.update_state_atomic(run.id, updates=state_updates) run.refresh_from_db() logger.info("Triggering workflow for task %s, existing run %s", task.id, run.id) - _trigger_task_processing_workflow(task, run, user_id, raise_on_error=True) + _trigger_task_processing_workflow( + task, + run, + user_id, + initial_message=(pending_user_message or task.description or None) + if task.runtime == Task.Runtime.PI + else None, + initial_artifact_ids=pending_user_artifact_ids if task.runtime == Task.Runtime.PI else None, + raise_on_error=True, + ) except Exception: if state_updates: rollback_updates = { @@ -3323,6 +3509,8 @@ def resume_task_run_in_cloud( ``"already_active"`` (400), ``"auth_error:"`` (400, github auth), ``"workflow_failed"`` (502), or ``"resumed"`` (run_dto set). Mirrors ``TaskRunViewSet.resume_in_cloud``. """ + from products.tasks.backend.facade.streams import reset_task_run_stream # noqa: PLC0415 + from products.tasks.backend.redis import run_uses_dedicated_stream # noqa: PLC0415 from products.tasks.backend.temporal.client import ( # noqa: PLC0415 — keep temporalio off the api import path resume_task_in_cloud_workflow, ) @@ -3334,9 +3522,6 @@ def resume_task_run_in_cloud( run = _get_visible_run(run_id, task_id, team_id) if run is None: return "not_found", None, None - if run.task.runtime == Task.Runtime.PI: - return "unsupported_runtime", None, None - logger.info( "resume_in_cloud_called", extra={ @@ -3391,6 +3576,11 @@ def resume_task_run_in_cloud( logger.info("Resuming task run in cloud", extra={"task_run_id": str(run.id), "task_id": str(run.task_id)}) try: + if not reset_task_run_stream( + str(run.id), + use_dedicated=run_uses_dedicated_stream(run.state), + ): + raise RuntimeError("Failed to reset task run event stream") resume_task_in_cloud_workflow(str(run.id), run.workflow_id) except Exception as e: logger.exception("Failed to trigger handoff workflow", extra={"task_run_id": str(run.id), "error": str(e)}) @@ -3508,6 +3698,25 @@ def get_conversation_task_dtos(task_ids: Sequence[str | UUID], team_id: int) -> return {task.id: _task_detail_to_dto(task, include_latest_run=False) for task in tasks} +def pi_cloud_runtime_enabled(team: Team, user: User) -> bool: + distinct_id = user.distinct_id or f"user_{user.id}" + organization_id = str(team.organization_id) + try: + return bool( + posthoganalytics.feature_enabled( + PI_CLOUD_RUNTIME_FEATURE_FLAG, + distinct_id, + groups={"organization": organization_id}, + group_properties={"organization": {"id": organization_id}}, + only_evaluate_locally=False, + send_feature_flag_events=False, + ) + ) + except Exception: + logger.exception("pi-harness flag check failed; treating as disabled") + return False + + def task_runtime(task_id: str | UUID, team_id: int, user_id: int | None, *, for_control: bool = False) -> str | None: return ( _visible_task_qs(team_id, user_id, for_control=for_control) @@ -4382,18 +4591,12 @@ def run_task( task = _visible_task_qs(team_id, user_id, for_control=True).filter(id=task_id).first() if task is None: return None - if task.runtime == Task.Runtime.PI: - return contracts.TaskRunResult( - error=contracts.TaskValidationError( - kind="detail", detail="Pi tasks cannot be run through the ACP task workflow." - ) - ) - mode = validated_data.get("mode", "background") branch = validated_data.get("branch") resume_from_run_id = validated_data.get("resume_from_run_id") pending_user_message = validated_data.get("pending_user_message") pending_user_artifact_ids = validated_data.get("pending_user_artifact_ids") or [] + is_pi_task = task.runtime == Task.Runtime.PI if not resume_from_run_id: warm_run = _idling_warm_run_for_task(task) @@ -4478,9 +4681,9 @@ def run_task( } extra_state: dict | None = None - if pending_user_message is not None: + if pending_user_message is not None and not is_pi_task: extra_state = {"pending_user_message": pending_user_message} - if pending_user_artifact_ids: + if pending_user_artifact_ids and not is_pi_task: extra_state = extra_state or {} extra_state["pending_user_artifact_ids"] = pending_user_artifact_ids if initial_permission_mode is not None: @@ -4500,8 +4703,9 @@ def run_task( prev_state = parse_run_state(previous_run.state) extra_state = extra_state or {} - extra_state["resume_from_run_id"] = str(resume_from_run_id) - extra_state.update(prev_state.resume_snapshot_carry_state()) + if not is_pi_task: + extra_state["resume_from_run_id"] = str(resume_from_run_id) + extra_state.update(prev_state.resume_snapshot_carry_state()) # The resumed agent still pushes the head branch baked into the original prompt, so the # PR webhook must be able to match this run, not the terminal predecessor. @@ -4540,7 +4744,7 @@ def run_task( provider = get_provider_for_runtime_adapter(runtime_adapter) - for key, value in { + run_state_values = { "pr_base_branch": branch, "pr_authorship_mode": pr_authorship_mode, "auto_publish": auto_publish, @@ -4552,7 +4756,11 @@ def run_task( "reasoning_effort": reasoning_effort, "context_window": context_window, "fast_mode": fast_mode, - }.items(): + } + if is_pi_task: + for key in ("runtime_adapter", "provider", "model", "reasoning_effort"): + run_state_values.pop(key) + for key, value in run_state_values.items(): if value is not None: extra_state = extra_state or {} extra_state[key] = value.value if hasattr(value, "value") else value @@ -4650,6 +4858,9 @@ def run_task( logger.info("Creating task run for task %s with mode=%s, branch=%s", task.id, mode, branch) task_run = task.create_run(mode=mode, branch=branch, extra_state=extra_state) + if is_pi_task and resume_from_run_id: + task_run.active_task_session = previous_run.active_task_session + task_run.save(update_fields=["active_task_session", "updated_at"]) if imported_mcp_servers or relayed_mcp_servers: update_fields = ["updated_at"] @@ -4671,7 +4882,20 @@ def run_task( cache_github_user_token(str(task_run.id), github_user_token) logger.info("Triggering workflow for task %s, run %s", task.id, task_run.id) - _trigger_task_processing_workflow(task, task_run, user_id, raise_on_error=False) + if is_pi_task: + initial_message = ( + pending_user_message if resume_from_run_id else pending_user_message or task.description or None + ) + _trigger_task_processing_workflow( + task, + task_run, + user_id, + initial_message=initial_message, + initial_artifact_ids=pending_user_artifact_ids, + raise_on_error=False, + ) + else: + _trigger_task_processing_workflow(task, task_run, user_id, raise_on_error=False) return contracts.TaskRunResult(task=get_task_detail(task.id, team_id, user_id)) diff --git a/products/tasks/backend/facade/streams.py b/products/tasks/backend/facade/streams.py index aaccea943072..c4e9da1b0308 100644 --- a/products/tasks/backend/facade/streams.py +++ b/products/tasks/backend/facade/streams.py @@ -16,6 +16,7 @@ TaskRunRedisStream, TaskRunStreamError, get_task_run_stream_key, + reset_task_run_stream, ) from products.tasks.backend.redis import run_uses_dedicated_stream @@ -28,5 +29,6 @@ "TaskRunStreamError", "get_task_run_stream_key", "handle_task_run_event_ingest", + "reset_task_run_stream", "run_uses_dedicated_stream", ] diff --git a/products/tasks/backend/logic/services/connection_token.py b/products/tasks/backend/logic/services/connection_token.py index b2c3046492e6..a49da81d47c5 100644 --- a/products/tasks/backend/logic/services/connection_token.py +++ b/products/tasks/backend/logic/services/connection_token.py @@ -38,6 +38,7 @@ class SandboxEventIngestTokenPayload: run_id: str task_id: str team_id: int + sandbox_id: str | None @dataclass(frozen=True) @@ -216,7 +217,12 @@ def create_sandbox_connection_token(task_run: TaskRun, user_id: int, distinct_id return jwt.encode(payload, key.private_key_pem, algorithm="RS256", headers={"kid": key.kid}) -def _encode_run_scoped_token(task_run: TaskRun, audience: str, ttl: timedelta) -> str: +def _encode_run_scoped_token( + task_run: TaskRun, + audience: str, + ttl: timedelta, + extra_claims: dict[str, object] | None = None, +) -> str: """Encode a run-scoped JWT carrying no user identity, signed with the run's key. Shared by the event-ingest and stream-read tokens; they stay distinct capabilities @@ -232,18 +238,33 @@ def _encode_run_scoped_token(task_run: TaskRun, audience: str, ttl: timedelta) - "exp": now + ttl, "aud": audience, } + if extra_claims: + payload.update(extra_claims) key = _signing_key_for_run(task_run) return jwt.encode(payload, key.private_key_pem, algorithm="RS256", headers={"kid": key.kid}) -def create_sandbox_event_ingest_token(task_run: TaskRun, ttl: timedelta = SANDBOX_EVENT_INGEST_TOKEN_TTL) -> str: +def create_sandbox_event_ingest_token( + task_run: TaskRun, + ttl: timedelta = SANDBOX_EVENT_INGEST_TOKEN_TTL, + *, + sandbox_id: str | None = None, +) -> str: """ Create a run-scoped JWT token for sandbox-to-Django live event ingest. This token intentionally carries no user identity and grants one capability: appending ordered live events for this task run. """ - return _encode_run_scoped_token(task_run, SANDBOX_EVENT_INGEST_AUDIENCE, ttl) + active_sandbox_id = sandbox_id or (task_run.state or {}).get("sandbox_id") + if not isinstance(active_sandbox_id, str) or not active_sandbox_id: + raise ValueError("Task run has no active sandbox identity") + return _encode_run_scoped_token( + task_run, + SANDBOX_EVENT_INGEST_AUDIENCE, + ttl, + {"sandbox_id": active_sandbox_id}, + ) def validate_sandbox_event_ingest_token(token: str) -> SandboxEventIngestTokenPayload: @@ -252,11 +273,19 @@ def validate_sandbox_event_ingest_token(token: str) -> SandboxEventIngestTokenPa run_id = payload.get("run_id") task_id = payload.get("task_id") team_id = payload.get("team_id") + sandbox_id = payload.get("sandbox_id") if not isinstance(run_id, str) or not isinstance(task_id, str) or type(team_id) is not int: raise jwt.InvalidTokenError("Sandbox event ingest token has invalid claims") + if sandbox_id is not None and (not isinstance(sandbox_id, str) or not sandbox_id): + raise jwt.InvalidTokenError("Sandbox event ingest token has invalid claims") - return SandboxEventIngestTokenPayload(run_id=run_id, task_id=task_id, team_id=team_id) + return SandboxEventIngestTokenPayload( + run_id=run_id, + task_id=task_id, + team_id=team_id, + sandbox_id=sandbox_id, + ) def create_stream_read_token(task_run: TaskRun, ttl: timedelta = STREAM_READ_TOKEN_TTL) -> str: diff --git a/products/tasks/backend/logic/services/docker_sandbox.py b/products/tasks/backend/logic/services/docker_sandbox.py index d782e55ce551..0a91d4451daa 100644 --- a/products/tasks/backend/logic/services/docker_sandbox.py +++ b/products/tasks/backend/logic/services/docker_sandbox.py @@ -164,43 +164,35 @@ def _run(args: list[str], check: bool = False, timeout: int | None = None) -> su return result @staticmethod - def _get_local_posthog_code_packages() -> tuple[str, str, str, str] | None: - """ - Get paths to local PostHog Desktop packages for development builds. - - Configure via LOCAL_POSTHOG_CODE_MONOREPO_ROOT pointing to the PostHog Desktop monorepo root. - Returns tuple of (agent_path, shared_path, git_path, enricher_path) or None if not configured. - """ + def _get_local_posthog_code_root() -> str | None: monorepo_root = os.environ.get( "LOCAL_POSTHOG_CODE_MONOREPO_ROOT", os.environ.get("LOCAL_TWIG_MONOREPO_ROOT", "") ) - if not monorepo_root or not os.path.isdir(monorepo_root): + if not monorepo_root: return None monorepo_root = os.path.abspath(monorepo_root) - agent_path = os.path.join(monorepo_root, "packages", "agent") - shared_path = os.path.join(monorepo_root, "packages", "shared") - git_path = os.path.join(monorepo_root, "packages", "git") - enricher_path = os.path.join(monorepo_root, "packages", "enricher") - - missing = [] - if not os.path.isdir(agent_path): - missing.append(f"agent: {agent_path}") - if not os.path.isdir(shared_path): - missing.append(f"shared: {shared_path}") - if not os.path.isdir(git_path): - missing.append(f"git: {git_path}") - if not os.path.isdir(enricher_path): - missing.append(f"enricher: {enricher_path}") - + required_paths = [ + os.path.join(monorepo_root, ".npmrc"), + os.path.join(monorepo_root, "package.json"), + os.path.join(monorepo_root, "pnpm-workspace.yaml"), + os.path.join(monorepo_root, "pnpm-lock.yaml"), + os.path.join(monorepo_root, "patches"), + *[ + os.path.join(monorepo_root, "packages", package_name, "package.json") + for package_name in ("agent", "harness", "shared", "git", "enricher") + ], + ] + missing = [path for path in required_paths if not os.path.exists(path)] if missing: + missing_paths = ", ".join(missing) raise SandboxProvisionError( - f"LOCAL_POSTHOG_CODE_MONOREPO_ROOT is set but required packages not found: {', '.join(missing)}", + f"LOCAL_POSTHOG_CODE_MONOREPO_ROOT is invalid: {missing_paths}", {"monorepo_root": monorepo_root, "missing": missing}, - cause=RuntimeError(f"Missing packages: {', '.join(missing)}"), + cause=RuntimeError(f"Missing paths: {missing_paths}"), ) - return agent_path, shared_path, git_path, enricher_path + return monorepo_root @staticmethod def _build_image_if_needed( @@ -253,34 +245,27 @@ def _build_image_if_needed( DockerSandbox._run(argv, check=True) @staticmethod - def _build_local_image(agent_path: str, shared_path: str, git_path: str, enricher_path: str) -> None: - """Build the local sandbox image with local PostHog Desktop packages.""" + def _build_local_image(monorepo_root: str) -> None: logger.info("Building posthog-sandbox-base-local image with local PostHog Desktop packages...") dockerfile_path = os.path.join( settings.BASE_DIR, "products/tasks/backend/sandbox/images/Dockerfile.sandbox-local" ) with tempfile.TemporaryDirectory() as tmpdir: - shutil.copytree( - agent_path, - os.path.join(tmpdir, "local-agent"), - ignore=shutil.ignore_patterns("node_modules"), - ) - shutil.copytree( - shared_path, - os.path.join(tmpdir, "local-shared"), - ignore=shutil.ignore_patterns("node_modules"), - ) - shutil.copytree( - git_path, - os.path.join(tmpdir, "local-git"), - ignore=shutil.ignore_patterns("node_modules"), - ) - shutil.copytree( - enricher_path, - os.path.join(tmpdir, "local-enricher"), - ignore=shutil.ignore_patterns("node_modules"), - ) + workspace_path = os.path.join(tmpdir, "local-workspace") + packages_path = os.path.join(workspace_path, "packages") + os.makedirs(packages_path) + + for file_name in (".npmrc", "package.json", "pnpm-lock.yaml", "pnpm-workspace.yaml"): + shutil.copy2(os.path.join(monorepo_root, file_name), workspace_path) + shutil.copytree(os.path.join(monorepo_root, "patches"), os.path.join(workspace_path, "patches")) + + for package_name in ("agent", "harness", "shared", "git", "enricher"): + shutil.copytree( + os.path.join(monorepo_root, "packages", package_name), + os.path.join(packages_path, package_name), + ignore=shutil.ignore_patterns("node_modules", ".turbo"), + ) DockerSandbox._run( [ @@ -338,10 +323,9 @@ def _ensure_image_exists(template: SandboxTemplate) -> str: ) return PI_IMAGE_NAME - local_packages = DockerSandbox._get_local_posthog_code_packages() - if local_packages: - agent_path, shared_path, git_path, enricher_path = local_packages - DockerSandbox._build_local_image(agent_path, shared_path, git_path, enricher_path) + local_monorepo_root = DockerSandbox._get_local_posthog_code_root() + if local_monorepo_root: + DockerSandbox._build_local_image(local_monorepo_root) return "posthog-sandbox-base-local" return DEFAULT_IMAGE_NAME @@ -790,6 +774,7 @@ def _build_agent_server_command( auto_publish: bool = False, interaction_origin: str | None = None, branch: str | None = None, + agent_runtime: str | None = None, runtime_adapter: str | None = None, provider: str | None = None, model: str | None = None, @@ -801,6 +786,7 @@ def _build_agent_server_command( relay_mcp_servers_arg: str = "", allowed_domains: list[str] | None = None, event_ingest_token: str | None = None, + task_run_session_token: str | None = None, event_ingest_url: str | None = None, event_ingest_keep_stream_open: bool = False, repo_ready_file: str | None = None, @@ -813,6 +799,8 @@ def _build_agent_server_command( event_ingest_url = DockerSandbox._transform_url_for_docker(event_ingest_url) env_prefix = build_agent_runtime_env_prefix( interaction_origin=interaction_origin, + agent_runtime=agent_runtime, + sandbox_id=self.id, runtime_adapter=runtime_adapter, provider=provider, model=model, @@ -821,6 +809,7 @@ def _build_agent_server_command( fast_mode=fast_mode, initial_permission_mode=initial_permission_mode, event_ingest_token=event_ingest_token, + task_run_session_token=task_run_session_token, event_ingest_url=event_ingest_url, event_ingest_keep_stream_open=event_ingest_keep_stream_open, rtk_enabled=rtk_enabled, @@ -899,6 +888,7 @@ def start_agent_server( auto_publish: bool = False, interaction_origin: str | None = None, branch: str | None = None, + agent_runtime: str | None = None, runtime_adapter: str | None = None, provider: str | None = None, model: str | None = None, @@ -910,6 +900,7 @@ def start_agent_server( relayed_mcp_servers: list[str] | None = None, allowed_domains: list[str] | None = None, event_ingest_token: str | None = None, + task_run_session_token: str | None = None, event_ingest_url: str | None = None, event_ingest_keep_stream_open: bool = False, repo_ready_file: str | None = None, @@ -951,6 +942,9 @@ def start_agent_server( if relayed_mcp_servers: relay_mcp_servers_arg = f" --relayMcpServers {shlex.quote(json.dumps(relayed_mcp_servers))}" + if agent_runtime == "pi" and not self.agent_server_supports_pi_runtime(): + raise RuntimeError("Installed sandbox agent-server does not support the Pi runtime") + if auto_publish and not self.agent_server_supports_auto_publish(): logger.warning(f"Installed agent-server in sandbox {self.id} predates --autoPublish; starting review-first") auto_publish = False @@ -972,6 +966,7 @@ def start_agent_server( auto_publish, interaction_origin, branch, + agent_runtime, runtime_adapter, provider, model, @@ -983,6 +978,7 @@ def start_agent_server( relay_mcp_servers_arg=relay_mcp_servers_arg, allowed_domains=allowed_domains, event_ingest_token=event_ingest_token, + task_run_session_token=task_run_session_token, event_ingest_url=event_ingest_url, event_ingest_keep_stream_open=event_ingest_keep_stream_open, repo_ready_file=repo_ready_file, @@ -1025,6 +1021,7 @@ def start_agent_server( auto_publish, interaction_origin, branch=None, + agent_runtime=agent_runtime, runtime_adapter=runtime_adapter, provider=provider, model=model, @@ -1036,6 +1033,7 @@ def start_agent_server( relay_mcp_servers_arg=relay_mcp_servers_arg, allowed_domains=allowed_domains, event_ingest_token=event_ingest_token, + task_run_session_token=task_run_session_token, event_ingest_url=event_ingest_url, event_ingest_keep_stream_open=event_ingest_keep_stream_open, repo_ready_file=repo_ready_file, diff --git a/products/tasks/backend/logic/services/modal_sandbox.py b/products/tasks/backend/logic/services/modal_sandbox.py index 69ca7cfdaf7b..2d81f541c106 100644 --- a/products/tasks/backend/logic/services/modal_sandbox.py +++ b/products/tasks/backend/logic/services/modal_sandbox.py @@ -68,6 +68,7 @@ read_gh_guard_script, ) from products.tasks.backend.logic.services.local_packages import ( + LocalPackage, get_local_package_runtime_dependencies, get_local_posthog_code_packages, ) @@ -301,6 +302,31 @@ def _merge_runtime_dependency_specs(name: str, existing: str, candidate: str) -> return f"{existing} {candidate}" +def _local_package_bin_link_commands(packages: tuple[LocalPackage, ...]) -> list[str]: + commands: list[str] = [] + bin_root = "/scripts/node_modules/.bin" + + for package in packages: + manifest = json.loads((package.source_path / "package.json").read_text()) + package_name = manifest.get("name") + package_bin = manifest.get("bin", {}) + if not isinstance(package_name, str): + continue + if isinstance(package_bin, str): + package_bin = {package_name.rsplit("/", 1)[-1]: package_bin} + if not isinstance(package_bin, dict): + continue + + for executable, target in package_bin.items(): + if not isinstance(executable, str) or not isinstance(target, str): + continue + target_path = f"../{package_name}/{target.removeprefix('./')}" + executable_path = f"{bin_root}/{executable}" + commands.append(f"ln -sfn {shlex.quote(target_path)} {shlex.quote(executable_path)}") + + return commands + + def _attach_local_package_mounts(image: modal.Image, template: SandboxTemplate) -> modal.Image: """Overlay each local package's built `dist/` dir onto the installed package via add_local_dir(copy=False). No-op unless `template` bundles the agent-server @@ -345,6 +371,10 @@ def _attach_local_package_mounts(image: modal.Image, template: SandboxTemplate) ) image = image.run_commands(install_command) + bin_link_commands = _local_package_bin_link_commands(packages) + if bin_link_commands: + image = image.run_commands(*bin_link_commands) + for package in packages: image = image.add_local_dir( str(package.build_output_path), @@ -945,6 +975,7 @@ def _build_agent_server_command( auto_publish: bool = False, interaction_origin: str | None = None, branch: str | None = None, + agent_runtime: str | None = None, runtime_adapter: str | None = None, provider: str | None = None, model: str | None = None, @@ -956,6 +987,7 @@ def _build_agent_server_command( relay_mcp_servers_arg: str = "", allowed_domains: list[str] | None = None, event_ingest_token: str | None = None, + task_run_session_token: str | None = None, event_ingest_url: str | None = None, event_ingest_keep_stream_open: bool = False, repo_ready_file: str | None = None, @@ -964,6 +996,8 @@ def _build_agent_server_command( ) -> str: env_prefix = build_agent_runtime_env_prefix( interaction_origin=interaction_origin, + agent_runtime=agent_runtime, + sandbox_id=self.id, runtime_adapter=runtime_adapter, provider=provider, model=model, @@ -972,6 +1006,7 @@ def _build_agent_server_command( fast_mode=fast_mode, initial_permission_mode=initial_permission_mode, event_ingest_token=event_ingest_token, + task_run_session_token=task_run_session_token, event_ingest_url=event_ingest_url, event_ingest_keep_stream_open=event_ingest_keep_stream_open, rtk_enabled=rtk_enabled, @@ -1076,6 +1111,7 @@ def start_agent_server( auto_publish: bool = False, interaction_origin: str | None = None, branch: str | None = None, + agent_runtime: str | None = None, runtime_adapter: str | None = None, provider: str | None = None, model: str | None = None, @@ -1087,6 +1123,7 @@ def start_agent_server( relayed_mcp_servers: list[str] | None = None, allowed_domains: list[str] | None = None, event_ingest_token: str | None = None, + task_run_session_token: str | None = None, event_ingest_url: str | None = None, event_ingest_keep_stream_open: bool = False, repo_ready_file: str | None = None, @@ -1131,6 +1168,9 @@ def start_agent_server( if relayed_mcp_servers: relay_mcp_servers_arg = f" --relayMcpServers {shlex.quote(json.dumps(relayed_mcp_servers))}" + if agent_runtime == "pi" and not self.agent_server_supports_pi_runtime(): + raise RuntimeError("Installed sandbox agent-server does not support the Pi runtime") + if auto_publish and not self.agent_server_supports_auto_publish(): logger.warning(f"Installed agent-server in sandbox {self.id} predates --autoPublish; starting review-first") auto_publish = False @@ -1152,6 +1192,7 @@ def start_agent_server( auto_publish, interaction_origin, branch, + agent_runtime, runtime_adapter, provider, model, @@ -1163,6 +1204,7 @@ def start_agent_server( relay_mcp_servers_arg=relay_mcp_servers_arg, allowed_domains=allowed_domains, event_ingest_token=event_ingest_token, + task_run_session_token=task_run_session_token, event_ingest_url=event_ingest_url, event_ingest_keep_stream_open=event_ingest_keep_stream_open, repo_ready_file=repo_ready_file, diff --git a/products/tasks/backend/logic/services/sandbox.py b/products/tasks/backend/logic/services/sandbox.py index 2dbd31c63da6..db46d7464bcb 100644 --- a/products/tasks/backend/logic/services/sandbox.py +++ b/products/tasks/backend/logic/services/sandbox.py @@ -143,7 +143,9 @@ def is_vm(self) -> bool: """Repos the sandbox is allowed to clone unauthenticated, even when the team has no GitHub integration""" # TODO: Remove `posthog/.github` when we switch repo discovery to repo-less agent (now it works as a lightweight dummy) -SENSITIVE_AGENT_RUNTIME_ENV_NAMES: frozenset[str] = frozenset({"POSTHOG_TASK_RUN_EVENT_INGEST_TOKEN"}) +SENSITIVE_AGENT_RUNTIME_ENV_NAMES: frozenset[str] = frozenset( + {"POSTHOG_TASK_RUN_EVENT_INGEST_TOKEN", "POSTHOG_TASK_RUN_SESSION_TOKEN"} +) SENSITIVE_AGENT_RUNTIME_ENV_PATTERN = re.compile( r"(?P" + "|".join(re.escape(name) for name in SENSITIVE_AGENT_RUNTIME_ENV_NAMES) + r")=" r"(?P'(?:[^']|'\"'\"')*'|\"(?:\\.|[^\"])*\"|\S+)" @@ -167,6 +169,8 @@ def redact_sandbox_command(command: str) -> str: def build_agent_runtime_env_prefix( *, interaction_origin: str | None = None, + agent_runtime: str | None = None, + sandbox_id: str | None = None, runtime_adapter: str | None = None, provider: str | None = None, model: str | None = None, @@ -175,12 +179,15 @@ def build_agent_runtime_env_prefix( fast_mode: bool | None = None, initial_permission_mode: str | None = None, event_ingest_token: str | None = None, + task_run_session_token: str | None = None, event_ingest_url: str | None = None, event_ingest_keep_stream_open: bool = False, rtk_enabled: bool = True, ) -> str: env_vars = { "POSTHOG_CODE_INTERACTION_ORIGIN": interaction_origin, + "POSTHOG_AGENT_RUNTIME": agent_runtime, + "POSTHOG_SANDBOX_ID": sandbox_id, "POSTHOG_CODE_RUNTIME_ADAPTER": runtime_adapter, "POSTHOG_CODE_PROVIDER": provider, "POSTHOG_CODE_MODEL": model, @@ -190,6 +197,7 @@ def build_agent_runtime_env_prefix( "POSTHOG_CODE_FAST_MODE": None if fast_mode is None else ("true" if fast_mode else "false"), "POSTHOG_CODE_INITIAL_PERMISSION_MODE": initial_permission_mode, "POSTHOG_TASK_RUN_EVENT_INGEST_TOKEN": event_ingest_token, + "POSTHOG_TASK_RUN_SESSION_TOKEN": task_run_session_token, "POSTHOG_TASK_RUN_EVENT_INGEST_URL": event_ingest_url, "POSTHOG_TASK_RUN_EVENT_INGEST_KEEP_STREAM_OPEN": "true" if event_ingest_keep_stream_open else None, # Set explicitly in both states: "0" opts the run out, "1" pins auto-detection on @@ -264,6 +272,13 @@ def agent_server_supports_exec_permission_regex(self) -> bool: ) return result.exit_code == 0 + def agent_server_supports_pi_runtime(self) -> bool: + result = self.execute( + "grep -q POSTHOG_AGENT_RUNTIME /scripts/node_modules/.bin/agent-server", + timeout_seconds=10, + ) + return result.exit_code == 0 + def clone_repository( self, repository: str, @@ -334,6 +349,7 @@ def start_agent_server( auto_publish: bool = False, interaction_origin: str | None = None, branch: str | None = None, + agent_runtime: str | None = None, runtime_adapter: str | None = None, provider: str | None = None, model: str | None = None, @@ -345,6 +361,7 @@ def start_agent_server( relayed_mcp_servers: list[str] | None = None, allowed_domains: list[str] | None = None, event_ingest_token: str | None = None, + task_run_session_token: str | None = None, event_ingest_url: str | None = None, event_ingest_keep_stream_open: bool = False, repo_ready_file: str | None = None, diff --git a/products/tasks/backend/logic/services/sandbox_usage.py b/products/tasks/backend/logic/services/sandbox_usage.py index 277f0d832378..ffb8516e6aa0 100644 --- a/products/tasks/backend/logic/services/sandbox_usage.py +++ b/products/tasks/backend/logic/services/sandbox_usage.py @@ -18,6 +18,7 @@ from typing import ParamSpec, TypeVar from uuid import UUID +from django.db import transaction from django.db.models import Q from django.utils import timezone @@ -44,48 +45,52 @@ def wrapper(*args: P.args, **kwargs: P.kwargs) -> R | None: return wrapper -@_best_effort def open_sandbox_session( - *, run_id: str | UUID, sandbox_id: str, config: SandboxConfig, sandbox_created_at: datetime | None = None + *, + run_id: str | UUID, + sandbox_id: str, + config: SandboxConfig, + sandbox_created_at: datetime | None = None, + required: bool = False, ) -> None: - """Record a freshly provisioned sandbox against its run. - - ``sandbox_created_at`` is the ``Sandbox.create()`` boundary — the provider's TTL - clock starts there, minutes before repo setup finishes and this row is opened, so - the TTL deadline must anchor on it rather than on insert time. - - Reads the live ``TaskRun`` row rather than any workflow-start snapshot: a warm - run claimed while its sandbox was still provisioning has already lost the - ``await_user_message`` marker, so the session is created attributed. Upserts on - ``sandbox_id`` so activity retries stay idempotent, and never regresses - ``user_attributed_at`` on an existing row. - """ - run = TaskRun.objects.select_related("task").only("id", "team_id", "state", "task__origin_product").get(id=run_id) - state = run.state or {} - created_at = sandbox_created_at or timezone.now() - shape = { - "team_id": run.team_id, - "task_run_id": run.id, - "origin_product": run.task.origin_product, - "prewarmed": bool(state.get("prewarmed")), - "vm_runtime": config.is_vm, - "cpu_cores": config.cpu_cores, - "memory_gb": config.memory_gb, - "ttl_seconds": config.ttl_seconds, - "burstable": config.burstable_resources, - "cpu_request_cores": config.cpu_request_cores if config.burstable_resources else None, - "memory_request_mb": config.memory_request_mb if config.burstable_resources else None, - "created_at": created_at, - "ttl_expires_at": created_at + timedelta(seconds=config.ttl_seconds), - } - SandboxSession.objects.for_team(run.team_id).update_or_create( - sandbox_id=sandbox_id, - defaults=shape, - create_defaults={ - **shape, - "user_attributed_at": None if state.get("await_user_message") else timezone.now(), - }, - ) + """Record a freshly provisioned sandbox against its run.""" + try: + with transaction.atomic(): + run = ( + TaskRun.objects.select_for_update(of=("self",)) + .select_related("task") + .only("id", "team_id", "state", "task__origin_product") + .get(id=run_id) + ) + state = run.state or {} + created_at = sandbox_created_at or timezone.now() + shape = { + "team_id": run.team_id, + "task_run_id": run.id, + "origin_product": run.task.origin_product, + "prewarmed": bool(state.get("prewarmed")), + "vm_runtime": config.is_vm, + "cpu_cores": config.cpu_cores, + "memory_gb": config.memory_gb, + "ttl_seconds": config.ttl_seconds, + "burstable": config.burstable_resources, + "cpu_request_cores": config.cpu_request_cores if config.burstable_resources else None, + "memory_request_mb": config.memory_request_mb if config.burstable_resources else None, + "created_at": created_at, + "ttl_expires_at": created_at + timedelta(seconds=config.ttl_seconds), + } + SandboxSession.objects.for_team(run.team_id).update_or_create( + sandbox_id=sandbox_id, + defaults=shape, + create_defaults={ + **shape, + "user_attributed_at": None if state.get("await_user_message") else timezone.now(), + }, + ) + except Exception: + logger.exception("sandbox_usage.ledger_write_failed", helper="open_sandbox_session") + if required: + raise @_best_effort @@ -93,9 +98,15 @@ def close_sandbox_session(sandbox_id: str, *, reason: str) -> None: """Stamp the sandbox's end. Idempotent — the first stamp wins.""" # Unscoped: cleanup/reap activities only carry the globally-unique provider # sandbox id, not team context. - SandboxSession.objects.unscoped().filter(sandbox_id=sandbox_id, ended_at__isnull=True).update( - ended_at=timezone.now(), ended_reason=reason - ) + sandbox_session = SandboxSession.objects.unscoped().filter(sandbox_id=sandbox_id).first() + if sandbox_session is None: + return + with transaction.atomic(): + TaskRun.objects.select_for_update().get(id=sandbox_session.task_run_id) + SandboxSession.objects.unscoped().filter( + id=sandbox_session.id, + ended_at__isnull=True, + ).update(ended_at=timezone.now(), ended_reason=reason) @_best_effort diff --git a/products/tasks/backend/logic/services/test_docker_sandbox.py b/products/tasks/backend/logic/services/test_docker_sandbox.py index 6eac741cd9eb..d9899c925f62 100644 --- a/products/tasks/backend/logic/services/test_docker_sandbox.py +++ b/products/tasks/backend/logic/services/test_docker_sandbox.py @@ -140,6 +140,44 @@ def test_transform_url_for_docker(self, input_url, expected_url): result = DockerSandbox._transform_url_for_docker(input_url) assert result == expected_url + def test_get_local_posthog_code_root(self, tmp_path, monkeypatch): + for file_name in (".npmrc", "package.json", "pnpm-workspace.yaml", "pnpm-lock.yaml"): + (tmp_path / file_name).touch() + (tmp_path / "patches").mkdir() + for package_name in ("agent", "harness", "shared", "git", "enricher"): + package_path = tmp_path / "packages" / package_name + package_path.mkdir(parents=True) + (package_path / "package.json").touch() + monkeypatch.setenv("LOCAL_POSTHOG_CODE_MONOREPO_ROOT", str(tmp_path)) + + assert DockerSandbox._get_local_posthog_code_root() == str(tmp_path) + + def test_build_local_image_copies_minimal_workspace_into_docker_context(self, tmp_path): + monorepo_path = tmp_path / "code" + context_path = tmp_path / "context" + for file_name in (".npmrc", "package.json", "pnpm-workspace.yaml", "pnpm-lock.yaml"): + monorepo_path.mkdir(exist_ok=True) + (monorepo_path / file_name).touch() + (monorepo_path / "patches").mkdir() + for package_name in ("agent", "harness", "shared", "git", "enricher"): + package_path = monorepo_path / "packages" / package_name + package_path.mkdir(parents=True) + (package_path / "package.json").touch() + + with ( + patch.object(DockerSandbox, "_run") as run, + patch("products.tasks.backend.logic.services.docker_sandbox.tempfile.TemporaryDirectory") as temporary, + ): + temporary.return_value.__enter__.return_value = str(context_path) + DockerSandbox._build_local_image(str(monorepo_path)) + + workspace_path = context_path / "local-workspace" + assert (workspace_path / "pnpm-workspace.yaml").is_file() + assert (workspace_path / "packages" / "harness" / "package.json").is_file() + command = run.call_args.args[0] + assert command[0:2] == ["docker", "build"] + assert command[-1] == str(context_path) + @patch("products.tasks.backend.logic.services.docker_sandbox.subprocess.run") @patch("products.tasks.backend.logic.services.docker_sandbox.os.path.exists") def test_create_transforms_posthog_api_url(self, mock_exists, mock_run): @@ -371,9 +409,61 @@ def test_start_agent_server_command_escaping(self, repository, task_id, run_id, assert shlex.quote(repo_path) in command assert shlex.quote(task_id) in command assert shlex.quote(run_id) in command + assert "POSTHOG_SANDBOX_ID=abc123" in command + assert "--sandboxId" not in command assert shlex.quote(mode) in command assert "--createPr true" in command + def test_start_agent_server_preserves_pi_protocol_without_branch_retry(self) -> None: + sandbox = DockerSandbox.__new__(DockerSandbox) + sandbox._container_id = "abc123" + sandbox.id = "abc123" + sandbox.config = SandboxConfig(name="test") + sandbox._host_port = 12345 + + with ( + patch.object(sandbox, "is_running", return_value=True), + patch.object(sandbox, "write_file"), + patch.object(sandbox, "agent_server_supports_auto_publish", return_value=True), + patch.object(sandbox, "agent_server_supports_pi_runtime", return_value=True), + patch.object(sandbox, "execute") as mock_execute, + patch.object(sandbox, "_launch_and_check", side_effect=[False, True]), + patch.object( + sandbox, "_build_agent_server_command", wraps=sandbox._build_agent_server_command + ) as mock_build, + ): + mock_execute.return_value = ExecutionResult(stdout="", stderr="", exit_code=0, error=None) + sandbox.start_agent_server( + "posthog/posthog", + "task-123", + "run-456", + branch="main", + agent_runtime="pi", + ) + + assert mock_build.call_count == 2 + assert mock_build.call_args_list[1].kwargs["agent_runtime"] == "pi" + + def test_start_agent_server_rejects_an_image_without_pi_support(self) -> None: + sandbox = DockerSandbox.__new__(DockerSandbox) + sandbox._container_id = "abc123" + sandbox.id = "abc123" + sandbox.config = SandboxConfig(name="test") + sandbox._host_port = 12345 + + with ( + patch.object(sandbox, "is_running", return_value=True), + patch.object(sandbox, "write_file"), + patch.object(sandbox, "agent_server_supports_pi_runtime", return_value=False), + pytest.raises(RuntimeError, match="does not support the Pi runtime"), + ): + sandbox.start_agent_server( + "posthog/posthog", + "task-123", + "run-456", + agent_runtime="pi", + ) + def test_parse_repo_mount_map_empty(self): with patch.dict(os.environ, {}, clear=True): assert parse_sandbox_repo_mount_map() == {} @@ -586,6 +676,7 @@ def test_start_agent_server_includes_runtime_environment_variables(self): "task-123", "run-456", "background", + agent_runtime="pi", runtime_adapter="codex", provider="openai", model="gpt-5.3-codex", @@ -598,6 +689,8 @@ def test_start_agent_server_includes_runtime_environment_variables(self): ) command = _agent_server_launch_command(mock_execute) + assert "POSTHOG_AGENT_RUNTIME=pi" in command + assert "POSTHOG_SANDBOX_ID=abc123" in command assert "POSTHOG_CODE_RUNTIME_ADAPTER=codex" in command assert "POSTHOG_CODE_PROVIDER=openai" in command assert "POSTHOG_CODE_MODEL=gpt-5.3-codex" in command diff --git a/products/tasks/backend/logic/services/tests/test_modal_sandbox.py b/products/tasks/backend/logic/services/tests/test_modal_sandbox.py index 6dd6a863a370..a54baff039ed 100644 --- a/products/tasks/backend/logic/services/tests/test_modal_sandbox.py +++ b/products/tasks/backend/logic/services/tests/test_modal_sandbox.py @@ -300,12 +300,14 @@ def test_installs_runtime_dependencies_before_mounting_local_builds(self, tmp_pa (source_path / "package.json").write_text( json.dumps( { + "name": "@posthog/agent", + "bin": {"agent-server": "./dist/server/bin.js"}, "dependencies": { "@openai/codex": "0.140.0", "custom-runtime": "github:example/custom-runtime#v1.2.3", "@posthog/shared": "workspace:*", "zod": "^4.2.0", - } + }, } ) ) @@ -326,11 +328,13 @@ def test_installs_runtime_dependencies_before_mounting_local_builds(self, tmp_pa base_image = MagicMock() system_dependency_image = MagicMock() dependency_image = MagicMock() + linked_image = MagicMock() mounted_image = MagicMock() final_image = MagicMock() base_image.apt_install.return_value = system_dependency_image system_dependency_image.run_commands.return_value = dependency_image - dependency_image.add_local_dir.return_value = mounted_image + dependency_image.run_commands.return_value = linked_image + linked_image.add_local_dir.return_value = mounted_image mounted_image.add_local_dir.return_value = final_image with patch( @@ -353,7 +357,10 @@ def test_installs_runtime_dependencies_before_mounting_local_builds(self, tmp_pa assert "@openai/codex@0.140.0" not in command assert "custom-runtime@github:" not in command assert "@posthog/shared" not in command - dependency_image.add_local_dir.assert_called_once_with( + dependency_image.run_commands.assert_called_once_with( + "ln -sfn ../@posthog/agent/dist/server/bin.js /scripts/node_modules/.bin/agent-server" + ) + linked_image.add_local_dir.assert_called_once_with( str(build_output_path), "/scripts/node_modules/@posthog/agent/dist", copy=False, @@ -443,6 +450,8 @@ def test_start_agent_server_success_without_domains_skips_agentsh(self, mock_san assert f"--repositoryPath {shlex.quote('/tmp/workspace/repos/posthog/posthog')}" in command assert f"--taskId {shlex.quote('task-123')}" in command assert f"--runId {shlex.quote('run-456')}" in command + assert f"POSTHOG_SANDBOX_ID={shlex.quote(mock_sandbox.id)}" in command + assert "--sandboxId" not in command assert f"--mode {shlex.quote('background')}" in command assert "--createPr true" in command assert "agentsh exec" not in command @@ -545,6 +554,7 @@ def test_start_agent_server_includes_runtime_environment_variables(self, mock_sa task_id="task-123", run_id="run-456", mode="background", + agent_runtime="pi", runtime_adapter="codex", provider="openai", model="gpt-5.3-codex", @@ -557,6 +567,7 @@ def test_start_agent_server_includes_runtime_environment_variables(self, mock_sa ) command = _agent_server_launch_command(mock_sandbox.execute) + assert "POSTHOG_AGENT_RUNTIME=pi" in command assert "POSTHOG_CODE_RUNTIME_ADAPTER=codex" in command assert "POSTHOG_CODE_PROVIDER=openai" in command assert "POSTHOG_CODE_MODEL=gpt-5.3-codex" in command diff --git a/products/tasks/backend/logic/stream/redis_stream.py b/products/tasks/backend/logic/stream/redis_stream.py index ba916b6ca208..01fc469c14ee 100644 --- a/products/tasks/backend/logic/stream/redis_stream.py +++ b/products/tasks/backend/logic/stream/redis_stream.py @@ -554,6 +554,28 @@ async def delete_stream(self) -> bool: return False +def reset_task_run_stream(run_id: str, use_dedicated: bool = False) -> bool: + stream_key = get_task_run_stream_key(run_id) + sequence_key = get_task_run_stream_sequence_key(stream_key) + completed_key = get_task_run_stream_completed_key(stream_key) + agent_active_key = get_task_run_stream_agent_active_key(stream_key) + heartbeat_key = get_task_run_stream_heartbeat_key(stream_key) + client = get_tasks_stream_redis_sync(use_dedicated) + + try: + client.delete( + stream_key, + sequence_key, + completed_key, + agent_active_key, + heartbeat_key, + ) + return True + except Exception: + logger.exception("task_run_stream_reset_failed", run_id=run_id) + return False + + def publish_task_run_stream_event(run_id: str, event: dict, use_dedicated: bool = False) -> str | None: """Synchronously publish a task-run event to Redis. diff --git a/products/tasks/backend/migrations/0074_task_session.py b/products/tasks/backend/migrations/0074_task_session.py new file mode 100644 index 000000000000..589f7e7e5e76 --- /dev/null +++ b/products/tasks/backend/migrations/0074_task_session.py @@ -0,0 +1,97 @@ +# Generated by Django 5.2.14 on 2026-07-17 13:24 + +import django.utils.timezone +import django.db.models.deletion +from django.db import migrations, models + +import posthog.models.utils + + +class Migration(migrations.Migration): + dependencies = [ + ("posthog", "1256_userproductlist_default_reason"), + ("tasks", "0073_task_activity"), + ] + + operations = [ + migrations.CreateModel( + name="TaskSession", + fields=[ + ( + "id", + models.UUIDField( + default=posthog.models.utils.uuid7, + editable=False, + primary_key=True, + serialize=False, + ), + ), + ( + "object_storage_key", + models.CharField(blank=True, max_length=512, null=True, unique=True), + ), + ("content_sha256", models.CharField(blank=True, max_length=64, null=True)), + ("size", models.PositiveIntegerField(default=0)), + ("created_at", models.DateTimeField(default=django.utils.timezone.now)), + ("updated_at", models.DateTimeField(auto_now=True)), + ( + "organization", + models.ForeignKey( + db_constraint=False, + db_index=False, + on_delete=django.db.models.deletion.CASCADE, + related_name="+", + to="posthog.organization", + ), + ), + ( + "team", + models.ForeignKey( + db_constraint=False, + db_index=False, + on_delete=django.db.models.deletion.CASCADE, + related_name="+", + to="posthog.team", + ), + ), + ( + "task", + models.ForeignKey( + db_index=False, + on_delete=django.db.models.deletion.CASCADE, + related_name="task_sessions", + to="tasks.task", + ), + ), + ], + options={ + "db_table": "posthog_task_session", + }, + ), + migrations.AddField( + model_name="taskrun", + name="active_task_session", + field=models.ForeignKey( + blank=True, + null=True, + on_delete=django.db.models.deletion.SET_NULL, + related_name="active_runs", + to="tasks.tasksession", + ), + ), + migrations.AddIndex( + model_name="tasksession", + index=models.Index( + fields=["organization", "-updated_at"], + name="task_session_org_updated_idx", + ), + ), + migrations.AddIndex( + model_name="tasksession", + index=models.Index(fields=["team", "-updated_at"], name="task_session_team_updated_idx"), + ), + migrations.AddIndex( + model_name="tasksession", + index=models.Index(fields=["task", "-updated_at"], name="task_session_task_updated_idx"), + ), + ] diff --git a/products/tasks/backend/migrations/max_migration.txt b/products/tasks/backend/migrations/max_migration.txt index 2f9da34a1494..4fc010358639 100644 --- a/products/tasks/backend/migrations/max_migration.txt +++ b/products/tasks/backend/migrations/max_migration.txt @@ -1 +1 @@ -0073_task_activity +0074_task_session diff --git a/products/tasks/backend/models.py b/products/tasks/backend/models.py index 6bef26ea4d09..67772255cc88 100644 --- a/products/tasks/backend/models.py +++ b/products/tasks/backend/models.py @@ -1,6 +1,5 @@ import os import re -import json import uuid import string import secrets @@ -8,7 +7,7 @@ from datetime import datetime from typing import TYPE_CHECKING, Any, Literal, Optional -from django.db.models.signals import post_save +from django.db.models.signals import post_delete, post_save from django.dispatch import receiver from pydantic import BaseModel @@ -48,6 +47,7 @@ from products.tasks.backend.logic.stream.redis_stream import publish_task_run_stream_event from products.tasks.backend.metrics import observe_task_run_created, observe_task_run_dispatch_callback from products.tasks.backend.redis import evaluate_dedicated_stream_flag, run_uses_dedicated_stream +from products.tasks.backend.storage import append_jsonl_object logger = structlog.get_logger(__name__) @@ -387,7 +387,7 @@ def create_run( extra_state: dict | None = None, branch: str | None = None, ) -> "TaskRun": - state: dict = {"mode": mode} + state: dict = {} if self.runtime == Task.Runtime.PI else {"mode": mode} if extra_state: state.update({k: v for k, v in extra_state.items() if k != "mode"}) # Pin the stream-routing decision once so every reader/writer agrees for this run's life. @@ -870,6 +870,71 @@ def _dispatch() -> None: return task +class TaskSession(TeamScopedRootMixin, UUIDModel): + organization = models.ForeignKey( + "posthog.Organization", + on_delete=models.CASCADE, + related_name="+", + db_constraint=False, + db_index=False, + ) + team = models.ForeignKey( + "posthog.Team", + on_delete=models.CASCADE, + related_name="+", + db_constraint=False, + db_index=False, + ) + task = models.ForeignKey(Task, on_delete=models.CASCADE, related_name="task_sessions", db_index=False) + object_storage_key = models.CharField(max_length=512, null=True, blank=True, unique=True) + content_sha256 = models.CharField(max_length=64, null=True, blank=True) + size = models.PositiveIntegerField(default=0) + created_at = models.DateTimeField(default=django_timezone.now) + updated_at = models.DateTimeField(auto_now=True) + + class Meta: + db_table = "posthog_task_session" + indexes = [ + models.Index(fields=["organization", "-updated_at"], name="task_session_org_updated_idx"), + models.Index(fields=["team", "-updated_at"], name="task_session_team_updated_idx"), + models.Index(fields=["task", "-updated_at"], name="task_session_task_updated_idx"), + ] + + @classmethod + def create_for_task(cls, task: Task) -> "TaskSession": + return cls.objects.unscoped().create( + organization_id=task.team.organization_id, + team_id=task.team_id, + task=task, + ) + + def read_jsonl(self) -> str: + if self.object_storage_key is None: + return "" + return object_storage.read(self.object_storage_key, missing_ok=True) or "" + + def tag_object(self) -> None: + if self.object_storage_key is None: + return + try: + object_storage.tag( + self.object_storage_key, + { + "data_class": "task_session", + "organization_id": str(self.organization_id), + "team_id": str(self.team_id), + "task_id": str(self.task_id), + }, + ) + except Exception as error: + logger.warning( + "task_session.failed_to_tag_object", + task_session_id=str(self.id), + object_storage_key=self.object_storage_key, + error=str(error), + ) + + class TaskThreadMessage(TeamScopedRootMixin): """One message in a task's thread — the side conversation channel members have around a task. Human messages never reach the agent unless the task author @@ -1449,6 +1514,13 @@ class Environment(models.TextChoices): id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False) task = models.ForeignKey(Task, on_delete=models.CASCADE, related_name="runs") team = models.ForeignKey("posthog.Team", on_delete=models.CASCADE) + active_task_session = models.ForeignKey( + TaskSession, + on_delete=models.SET_NULL, + null=True, + blank=True, + related_name="active_runs", + ) branch = models.CharField(max_length=255, blank=True, null=True, help_text="Branch name for the run") @@ -1592,6 +1664,9 @@ def prepare_for_cloud_handoff(self) -> None: state.pop("pending_user_message", None) state.pop("pending_user_message_id", None) state.pop("pending_user_message_ts", None) + state.pop("sandbox_id", None) + state.pop("sandbox_url", None) + state.pop("sandbox_jwt_kid", None) self.state = state logger.info( @@ -1653,6 +1728,21 @@ def _mutator(state: dict[str, Any]) -> None: return cls.mutate_state_atomic(run_id, _mutator) + @classmethod + def clear_sandbox_connection_state_atomic( + cls, + run_id: str | uuid.UUID, + sandbox_id: str, + ) -> dict[str, Any]: + def _mutator(state: dict[str, Any]) -> None: + if state.get("sandbox_id") != sandbox_id: + return + + for key in ("sandbox_id", "sandbox_url", "sandbox_connect_token", "sandbox_jwt_kid"): + state.pop(key, None) + + return cls.mutate_state_atomic(run_id, _mutator) + @staticmethod def get_workflow_id( task_id: str | uuid.UUID, run_id: str | uuid.UUID, workflow_id_prefix: str | None = None @@ -1790,13 +1880,7 @@ def append_log(self, entries: list[dict], *, ttl_days: int | None = DEFAULT_LOG_ if not entries: return - existing_content = object_storage.read(self.log_url, missing_ok=True) or "" - is_new_file = not existing_content - - new_lines = "\n".join(json.dumps(entry) for entry in entries) - content = existing_content + ("\n" if existing_content else "") + new_lines - - object_storage.write(self.log_url, content) + is_new_file = append_jsonl_object(self.log_url, entries) self._mirror_logs_to_posthog_logs(entries) @@ -2689,6 +2773,27 @@ def __str__(self): return f"Presence: user {self.user_id} on task {self.task_id} via device {self.push_token_id}" +@receiver(post_delete, sender=TaskSession) +def delete_task_session_object(sender: type[TaskSession], instance: TaskSession, **kwargs: Any) -> None: + if instance.object_storage_key is None: + return + object_storage_key = instance.object_storage_key + task_session_id = str(instance.id) + + def delete_object() -> None: + try: + object_storage.delete(object_storage_key) + except Exception as error: + logger.warning( + "task_session.failed_to_delete_object", + task_session_id=task_session_id, + object_storage_key=object_storage_key, + error=str(error), + ) + + transaction.on_commit(delete_object) + + @receiver(post_save, sender=TaskRun) def track_task_run_completion(sender, instance: TaskRun, created: bool, **kwargs): try: diff --git a/products/tasks/backend/presentation/serializers.py b/products/tasks/backend/presentation/serializers.py index 2d8b56bfbc05..af48b3ad07af 100644 --- a/products/tasks/backend/presentation/serializers.py +++ b/products/tasks/backend/presentation/serializers.py @@ -730,6 +730,17 @@ def validate_entries(self, value): return value +class TaskSessionResponseSerializer(serializers.Serializer): + id = serializers.UUIDField(help_text="Task session identifier") + download_url = serializers.URLField(allow_null=True, help_text="Temporary URL for downloading the session") + content_sha256 = serializers.CharField(allow_null=True, help_text="SHA-256 digest of the current session content") + + +class TaskSessionSyncResponseSerializer(serializers.Serializer): + id = serializers.UUIDField(help_text="Task session identifier") + content_sha256 = serializers.CharField(help_text="SHA-256 digest of the uploaded session content") + + class TaskRunRelayMessageResponseSerializer(serializers.Serializer): status = serializers.CharField(help_text="Relay status: 'accepted' or 'skipped'") relay_id = serializers.CharField(required=False, help_text="Relay workflow ID when accepted") @@ -2424,6 +2435,9 @@ class TaskRunCommandRequestSerializer(serializers.Serializer): "permission_response", "set_config_option", "mcp_response", + "pi/rpc", + "queue_get", + "queue_clear", ] # Cap on the serialized mcp_response params (docs/cloud-mcp-relay.md): the relayed JSON-RPC @@ -2451,7 +2465,7 @@ class TaskRunCommandRequestSerializer(serializers.Serializer): ) def validate_id(self, value): - if value is not None and not isinstance(value, (str, int, float)): + if value is not None and not isinstance(value, str | int | float): raise serializers.ValidationError("id must be a string or number") return value @@ -2496,6 +2510,18 @@ def validate(self, attrs): raise serializers.ValidationError( {"params": "user_message requires a non-empty content string, artifact_ids, or both"} ) + elif method == "pi/rpc": + command = params.get("command") + if not isinstance(command, dict): + raise serializers.ValidationError({"params": "command must be an object"}) + command_type = command.get("type") + if not isinstance(command_type, str) or not command_type: + raise serializers.ValidationError({"params": "command.type must be a non-empty string"}) + command_id = command.get("id") + if not isinstance(command_id, str) or not command_id: + raise serializers.ValidationError({"params": "command.id must be a non-empty string"}) + if attrs.get("id") != command_id: + raise serializers.ValidationError({"id": "id must match params.command.id"}) elif method == "permission_response": self._require_nonempty_string(params, "requestId") self._require_nonempty_string(params, "optionId") @@ -2530,7 +2556,7 @@ class TaskRunCommandResponseSerializer(serializers.Serializer): jsonrpc = serializers.CharField(help_text="JSON-RPC version") id = serializers.JSONField(required=False, default=None, help_text="Request ID echoed back (string or number)") - result = serializers.DictField(required=False, help_text="Command result on success") + result = serializers.JSONField(required=False, help_text="Command result on success") error = serializers.DictField(required=False, help_text="Error details on failure") diff --git a/products/tasks/backend/presentation/views/api.py b/products/tasks/backend/presentation/views/api.py index 31939ceaae5c..a2c7be89ab49 100644 --- a/products/tasks/backend/presentation/views/api.py +++ b/products/tasks/backend/presentation/views/api.py @@ -19,8 +19,9 @@ from rest_framework import status, viewsets from rest_framework.authentication import SessionAuthentication from rest_framework.decorators import action -from rest_framework.exceptions import NotFound, PermissionDenied, ValidationError +from rest_framework.exceptions import NotFound, ParseError, PermissionDenied, ValidationError from rest_framework.pagination import LimitOffsetPagination +from rest_framework.parsers import BaseParser from rest_framework.permissions import IsAuthenticated from rest_framework.response import Response @@ -107,6 +108,8 @@ TaskRunStartRequestSerializer, TaskRunUpdateSerializer, TaskSerializer, + TaskSessionResponseSerializer, + TaskSessionSyncResponseSerializer, TaskStagedArtifactsFinalizeUploadRequestSerializer, TaskStagedArtifactsFinalizeUploadResponseSerializer, TaskStagedArtifactsPrepareUploadRequestSerializer, @@ -121,8 +124,38 @@ from ee.hogai.utils.aio import async_to_sync + +class OctetStreamParser(BaseParser): + media_type = "application/octet-stream" + + def parse(self, stream, media_type=None, parser_context=None): + request = (parser_context or {}).get("request") + raw_content_length = request.META.get("CONTENT_LENGTH") if request is not None else None + if not isinstance(raw_content_length, str): + raise ParseError("A valid Content-Length header is required") + try: + content_length = int(raw_content_length) + except ValueError as error: + raise ParseError("A valid Content-Length header is required") from error + if content_length < 0 or content_length > tasks_facade.TASK_SESSION_MAX_SIZE_BYTES: + raise ParseError("The task session content size is invalid") + + content = stream.read(tasks_facade.TASK_SESSION_MAX_SIZE_BYTES + 1) + if len(content) > tasks_facade.TASK_SESSION_MAX_SIZE_BYTES: + raise ParseError("The task session content size is invalid") + return content + + logger = logging.getLogger(__name__) + +def _pi_cloud_runtime_disabled_response() -> Response: + return Response( + TaskRunErrorResponseSerializer({"error": "Pi cloud runtime is disabled"}).data, + status=status.HTTP_403_FORBIDDEN, + ) + + TASKS_PREWARM_SANDBOX_FLAG = "tasks-prewarm-sandbox" TASK_RUN_STREAM_KEEPALIVE_INTERVAL_SECONDS = 20.0 @@ -572,6 +605,7 @@ def staged_artifacts_finalize_upload(self, request, pk=None, **kwargs): responses={ 200: OpenApiResponse(response=TaskSerializer, description="Task with updated latest run"), 400: OpenApiResponse(response=TaskRunErrorResponseSerializer, description="Invalid task run payload"), + 403: OpenApiResponse(response=TaskRunErrorResponseSerializer, description="Pi cloud runtime is disabled"), 404: OpenApiResponse(description="Task not found"), 429: OpenApiResponse( response=TaskRunErrorResponseSerializer, description="Team is over its posthog_code usage limit" @@ -586,6 +620,10 @@ def run(self, request, pk=None, **kwargs): # Original order: 404 if the task isn't visible, then gate (always cloud) before the run. if not tasks_facade.task_visible(pk, self.team_id, self._user_id(), for_control=True): raise NotFound() + if tasks_facade.task_runtime( + pk, self.team_id, self._user_id(), for_control=True + ) == tasks_facade.TaskRuntime.PI and not tasks_facade.pi_cloud_runtime_enabled(self.team, request.user): + return _pi_cloud_runtime_disabled_response() # Self-driving report tasks (Inbox "Create PR" / "Discuss") are entitled through the Inbox # (`product-autonomy`), not the PostHog Code (`tasks`) product, so they skip the Code @@ -852,6 +890,7 @@ def _user_id(self) -> int | None: "retrieve", "logs", "session_logs", + "task_session", "stream", "stream_token", "artifacts_presign", @@ -923,6 +962,7 @@ def _validation_error_response(self, error: tasks_contracts.TaskRunValidationErr responses={ 201: OpenApiResponse(response=TaskRunDetailSerializer, description="Created task run"), 400: OpenApiResponse(response=TaskRunErrorResponseSerializer, description="Invalid task run payload"), + 403: OpenApiResponse(response=TaskRunErrorResponseSerializer, description="Pi cloud runtime is disabled"), 429: OpenApiResponse( response=TaskRunErrorResponseSerializer, description="Team is over its posthog_code usage limit" ), @@ -937,6 +977,10 @@ def create(self, request, *args, **kwargs): # Gate cloud runs before the run row is created; local runs aren't limited. if environment == tasks_facade.TaskRunEnvironment.CLOUD: + if tasks_facade.task_runtime( + task_id, self.team_id, self._user_id(), for_control=True + ) == tasks_facade.TaskRuntime.PI and not tasks_facade.pi_cloud_runtime_enabled(self.team, request.user): + return _pi_cloud_runtime_disabled_response() if (limit_response := cloud_usage_limit_response(request.user, self.team_id)) is not None: return limit_response @@ -954,6 +998,7 @@ def create(self, request, *args, **kwargs): responses={ 200: OpenApiResponse(response=TaskSerializer, description="Task with updated latest run"), 400: OpenApiResponse(response=TaskRunErrorResponseSerializer, description="Invalid start payload"), + 403: OpenApiResponse(response=TaskRunErrorResponseSerializer, description="Pi cloud runtime is disabled"), 404: OpenApiResponse(description="Task run not found"), 429: OpenApiResponse( response=TaskRunErrorResponseSerializer, description="Team is over its posthog_code usage limit" @@ -965,15 +1010,15 @@ def create(self, request, *args, **kwargs): @action(detail=True, methods=["post"], url_path="start", required_scopes=["task:write"]) def start(self, request, pk=None, **kwargs): task_id = self._ensure_task_accessible() - startable = tasks_facade.check_task_run_startable(pk, task_id, self.team_id) if startable == "not_found": raise NotFound() - if startable == "unsupported_runtime": - return Response( - TaskRunErrorResponseSerializer({"error": "Pi tasks cannot be run through the ACP task workflow."}).data, - status=status.HTTP_400_BAD_REQUEST, - ) + + if tasks_facade.task_runtime( + task_id, self.team_id, self._user_id(), for_control=True + ) == tasks_facade.TaskRuntime.PI and not tasks_facade.pi_cloud_runtime_enabled(self.team, request.user): + return _pi_cloud_runtime_disabled_response() + if startable == "not_cloud": return Response( TaskRunErrorResponseSerializer({"error": "Only cloud runs can be started via this endpoint"}).data, @@ -1165,6 +1210,111 @@ def append_log(self, request, pk=None, **kwargs): response["Server-Timing"] = timer.to_header_string() return response + @extend_schema( + responses={ + 200: TaskSessionResponseSerializer, + 404: OpenApiResponse(description="Task session not found"), + }, + summary="Get active task session storage access", + ) + @action(detail=True, methods=["get"], url_path="task_session", required_scopes=["task:read"]) + def task_session(self, request, pk=None, **kwargs): + task_id = self._ensure_task_accessible() + result = tasks_facade.get_task_run_session(pk, task_id, self.team_id) + if result is None: + raise NotFound() + session_id, download_url, content_sha256 = result + return Response( + TaskSessionResponseSerializer( + { + "id": session_id, + "download_url": download_url, + "content_sha256": content_sha256, + } + ).data + ) + + @extend_schema( + request=OpenApiTypes.BINARY, + parameters=[ + OpenApiParameter( + name="X-Sandbox-ID", + type=OpenApiTypes.STR, + location=OpenApiParameter.HEADER, + required=True, + description="Active sandbox identifier", + ), + OpenApiParameter( + name="X-Task-Run-Token", + type=OpenApiTypes.STR, + location=OpenApiParameter.HEADER, + required=True, + description="Sandbox-scoped task run token", + ), + OpenApiParameter( + name="If-Match", + type=OpenApiTypes.STR, + location=OpenApiParameter.HEADER, + required=True, + description="Expected current content SHA-256 digest, or none for an empty session", + ), + ], + responses={ + 200: TaskSessionSyncResponseSerializer, + 400: OpenApiResponse(description="Missing required header"), + 403: OpenApiResponse(description="Invalid task run token"), + 404: OpenApiResponse(description="Task session not found"), + 409: OpenApiResponse(response=TaskRunErrorResponseSerializer), + }, + summary="Replace the active native task session", + ) + @action( + detail=True, + methods=["post"], + url_path="task_session_sync", + required_scopes=["task:write"], + parser_classes=[OctetStreamParser], + ) + def sync_task_session(self, request, pk=None, **kwargs): + task_id = self._ensure_task_accessible() + sandbox_id = request.headers.get("X-Sandbox-ID") + if not sandbox_id: + raise ValidationError({"X-Sandbox-ID": "This header is required."}) + task_run_token = request.headers.get("X-Task-Run-Token") + if not task_run_token or not tasks_facade.validate_task_run_sandbox_token( + task_run_token, + pk, + task_id, + self.team_id, + sandbox_id, + ): + raise PermissionDenied("The task run token is invalid") + if_match = request.headers.get("If-Match") + if if_match is None: + raise ValidationError({"If-Match": "This header is required."}) + expected_content_sha256 = if_match.strip().strip('"') + if expected_content_sha256 == "none": + expected_content_sha256 = None + + try: + result = tasks_facade.sync_task_run_session( + pk, + task_id, + self.team_id, + sandbox_id=sandbox_id, + expected_content_sha256=expected_content_sha256, + content=request.data, + ) + except ValueError as error: + return Response( + TaskRunErrorResponseSerializer({"error": str(error)}).data, + status=status.HTTP_409_CONFLICT, + ) + if result is None: + raise NotFound() + session_id, content_sha256 = result + return Response(TaskSessionSyncResponseSerializer({"id": session_id, "content_sha256": content_sha256}).data) + @validated_request( request_serializer=TaskRunRelayMessageRequestSerializer, responses={ @@ -1485,7 +1635,9 @@ def stream_token(self, request, pk=None, **kwargs): if token is None: raise NotFound() stream_base_url = tasks_facade.resolve_stream_base_url( - distinct_id=request.user.distinct_id, organization_id=self.team.organization_id + distinct_id=request.user.distinct_id, + organization_id=self.team.organization_id, + force_proxy=tasks_facade.task_uses_pi_runtime(task_id, self.team_id), ) return Response(StreamReadTokenResponseSerializer({"token": token, "stream_base_url": stream_base_url}).data) @@ -1506,7 +1658,7 @@ def stream_token(self, request, pk=None, **kwargs): summary="Send command to task run", description="Queue user_message JSON-RPC commands through the task workflow and forward sandbox control " "commands to the agent server. Supports user_message, cancel, close, permission_response, " - "set_config_option, and mcp_response commands.", + "set_config_option, mcp_response, native Pi RPC commands, and Pi queue operations.", strict_request_validation=True, ) @action( @@ -1517,16 +1669,26 @@ def stream_token(self, request, pk=None, **kwargs): ) def command(self, request, pk=None, **kwargs): task_id = self._ensure_task_accessible() + method = request.validated_data["method"] + task_runtime = tasks_facade.task_runtime(task_id, self.team_id, self._user_id(), for_control=True) if ( - tasks_facade.task_runtime(task_id, self.team_id, self._user_id(), for_control=True) - == tasks_facade.TaskRuntime.PI - ): + method.startswith("pi/") or method in {"queue_get", "queue_clear"} + ) and task_runtime != tasks_facade.TaskRuntime.PI: return Response( - TaskRunErrorResponseSerializer({"error": "Pi tasks do not support ACP task commands."}).data, + TaskRunErrorResponseSerializer({"error": "Pi commands require a Pi task."}).data, + status=status.HTTP_400_BAD_REQUEST, + ) + if task_runtime == tasks_facade.TaskRuntime.PI and method not in { + "user_message", + "cancel", + "pi/rpc", + "queue_get", + "queue_clear", + }: + return Response( + TaskRunErrorResponseSerializer({"error": f"{method} is not supported for Pi tasks."}).data, status=status.HTTP_400_BAD_REQUEST, ) - - method = request.validated_data["method"] request_id = request.validated_data.get("id") params = request.validated_data.get("params") @@ -1558,6 +1720,7 @@ def command(self, request, pk=None, **kwargs): content=command_params.get("content"), artifact_ids=artifact_ids, actor_user_id=request.user.id, + message_id=str(request_id) if request_id is not None else None, steer=command_params.get("steer", False), ) except Exception: @@ -1851,6 +2014,7 @@ def session_logs(self, request, pk=None, **kwargs): 400: OpenApiResponse( response=TaskRunErrorResponseSerializer, description="Run already active or workflow failed" ), + 403: OpenApiResponse(response=TaskRunErrorResponseSerializer, description="Pi cloud runtime is disabled"), 429: OpenApiResponse( response=TaskRunErrorResponseSerializer, description="Team is over its posthog_code usage limit" ), @@ -1868,6 +2032,10 @@ def resume_in_cloud(self, request, pk=None, **kwargs): task_id = self._ensure_task_accessible() if tasks_facade.get_task_run_detail(pk, task_id, self.team_id) is None: raise NotFound() + if tasks_facade.task_runtime( + task_id, self.team_id, self._user_id(), for_control=True + ) == tasks_facade.TaskRuntime.PI and not tasks_facade.pi_cloud_runtime_enabled(self.team, request.user): + return _pi_cloud_runtime_disabled_response() # Resume also runs in cloud: gate before handoff. if (limit_response := cloud_usage_limit_response(request.user, self.team_id)) is not None: @@ -1876,11 +2044,6 @@ def resume_in_cloud(self, request, pk=None, **kwargs): outcome, run, _ = tasks_facade.resume_task_run_in_cloud(pk, task_id, self.team_id, self._user_id()) if outcome == "not_found": raise NotFound() - if outcome == "unsupported_runtime": - return Response( - TaskRunErrorResponseSerializer({"error": "Pi tasks cannot be run through the ACP task workflow."}).data, - status=status.HTTP_400_BAD_REQUEST, - ) if outcome == "already_active": return Response( TaskRunErrorResponseSerializer({"error": "Run is already active in cloud"}).data, diff --git a/products/tasks/backend/sandbox/images/Dockerfile.sandbox-local b/products/tasks/backend/sandbox/images/Dockerfile.sandbox-local index 7c66546dd5dd..221bbee12e9d 100644 --- a/products/tasks/backend/sandbox/images/Dockerfile.sandbox-local +++ b/products/tasks/backend/sandbox/images/Dockerfile.sandbox-local @@ -4,26 +4,12 @@ FROM posthog-sandbox-base -COPY local-shared /local-shared -COPY local-git /local-git -COPY local-enricher /local-enricher -COPY local-agent /local-agent +COPY local-workspace /local-workspace -# Pack shared first (no workspace deps), then enricher and git (both depend on -# shared), then agent (depends on shared, git, and enricher). Every package with -# a "workspace:*" dependency must have it rewritten to the packed tgz before -# `pnpm pack`, otherwise pnpm fails with ERR_PNPM_CANNOT_RESOLVE_WORKSPACE_PROTOCOL. -RUN cd /local-shared && pnpm pack && \ - cd /local-enricher && \ - sed -i 's/"@posthog\/shared": "workspace:\*"/"@posthog\/shared": "file:\/local-shared\/posthog-shared-1.0.0.tgz"/' package.json && \ - pnpm pack && \ - cd /local-git && \ - sed -i 's/"@posthog\/shared": "workspace:\*"/"@posthog\/shared": "file:\/local-shared\/posthog-shared-1.0.0.tgz"/' package.json && \ - pnpm pack && \ - cd /local-agent && \ - sed -i 's/"@posthog\/shared": "workspace:\*"/"@posthog\/shared": "file:\/local-shared\/posthog-shared-1.0.0.tgz"/' package.json && \ - sed -i 's/"@posthog\/git": "workspace:\*"/"@posthog\/git": "file:\/local-git\/posthog-git-1.0.0.tgz"/' package.json && \ - sed -i 's/"@posthog\/enricher": "workspace:\*"/"@posthog\/enricher": "file:\/local-enricher\/posthog-enricher-1.0.0.tgz"/' package.json && \ - pnpm pack && \ - cd /scripts && pnpm install /local-agent/*.tgz --config.dangerouslyAllowAllBuilds=true && \ - rm -rf /local-agent /local-shared /local-git /local-enricher +RUN cd /local-workspace && \ + pnpm install --filter @posthog/agent... --ignore-scripts --frozen-lockfile && \ + cd /local-workspace/packages/agent && \ + pnpm pack --pack-destination /tmp && \ + cd /scripts && \ + pnpm install /tmp/posthog-agent-*.tgz --config.dangerouslyAllowAllBuilds=true && \ + rm -rf /tmp/posthog-agent-*.tgz /local-workspace diff --git a/products/tasks/backend/storage.py b/products/tasks/backend/storage.py new file mode 100644 index 000000000000..8a260721c7b8 --- /dev/null +++ b/products/tasks/backend/storage.py @@ -0,0 +1,15 @@ +import json +from typing import Any + +from posthog.storage import object_storage + + +def append_jsonl_object(object_storage_key: str, entries: list[dict[str, Any]]) -> bool: + existing_content = object_storage.read(object_storage_key, missing_ok=True) or "" + is_new_object = not existing_content + new_lines = "\n".join(json.dumps(entry) for entry in entries) + content = existing_content + ("\n" if existing_content else "") + new_lines + + object_storage.write(object_storage_key, content) + + return is_new_object diff --git a/products/tasks/backend/temporal/client.py b/products/tasks/backend/temporal/client.py index 624201901db5..2eefc31a9724 100644 --- a/products/tasks/backend/temporal/client.py +++ b/products/tasks/backend/temporal/client.py @@ -27,7 +27,7 @@ STEERING_PROTOCOL_QUERY_TIMEOUT, STEERING_PROTOCOL_VERSION, ) -from products.tasks.backend.temporal.process_task.workflow import ProcessTaskInput +from products.tasks.backend.temporal.process_task.workflow import PendingFollowup, ProcessTaskInput from products.tasks.backend.temporal.slack_relay.activities import RelaySlackMessageInput if TYPE_CHECKING: @@ -198,6 +198,7 @@ async def execute_task_processing_workflow_async( posthog_mcp_scopes: PosthogMcpScopes = "read_only", prewarmed: bool = False, workflow_id_prefix: Optional[str] = None, + initial_message: PendingFollowup | None = None, ) -> None: """ Start the task processing workflow asynchronously. Fire-and-forget. @@ -229,6 +230,7 @@ async def execute_task_processing_workflow_async( slack_thread_context=slack_context_dict, posthog_mcp_scopes=posthog_mcp_scopes, prewarmed=prewarmed, + initial_message=initial_message, ) logger.info( @@ -282,6 +284,7 @@ def execute_task_processing_workflow( posthog_mcp_scopes: PosthogMcpScopes = "read_only", prewarmed: bool = False, workflow_id_prefix: Optional[str] = None, + initial_message: PendingFollowup | None = None, ) -> None: """ Start the task processing workflow synchronously. Fire-and-forget. @@ -312,6 +315,7 @@ def execute_task_processing_workflow( slack_thread_context=slack_context_dict, posthog_mcp_scopes=posthog_mcp_scopes, prewarmed=prewarmed, + initial_message=initial_message, ) logger.info( diff --git a/products/tasks/backend/temporal/process_task/activities/get_sandbox_for_repository.py b/products/tasks/backend/temporal/process_task/activities/get_sandbox_for_repository.py index 3ca704025039..54c176bcdb8f 100644 --- a/products/tasks/backend/temporal/process_task/activities/get_sandbox_for_repository.py +++ b/products/tasks/backend/temporal/process_task/activities/get_sandbox_for_repository.py @@ -377,21 +377,28 @@ def get_sandbox_for_repository(input: GetSandboxForRepositoryInput) -> GetSandbo credentials = sandbox.get_connect_credentials() - sandbox_state = { - "sandbox_id": sandbox.id, - "sandbox_url": credentials.url, - SANDBOX_JWT_STATE_KID_KEY: get_primary_sandbox_jwt_kid(), - } - if credentials.token: - sandbox_state["sandbox_connect_token"] = credentials.token - TaskRun.update_state_atomic(ctx.run_id, updates=sandbox_state) - - # Best-effort usage-ledger row (swallows its own failures). Only after the - # sandbox is fully reachable, mirroring create_sandbox_for_repository: the - # failure paths above destroy the sandbox, and those must not enter the ledger. - open_sandbox_session( - run_id=ctx.run_id, sandbox_id=sandbox.id, config=sandbox.config, sandbox_created_at=sandbox_created_at - ) + try: + sandbox_state = { + "sandbox_id": sandbox.id, + "sandbox_url": credentials.url, + SANDBOX_JWT_STATE_KID_KEY: get_primary_sandbox_jwt_kid(), + } + if credentials.token: + sandbox_state["sandbox_connect_token"] = credentials.token + TaskRun.update_state_atomic(ctx.run_id, updates=sandbox_state) + open_sandbox_session( + run_id=ctx.run_id, + sandbox_id=sandbox.id, + config=sandbox.config, + sandbox_created_at=sandbox_created_at, + required=ctx.task_runtime == "pi", + ) + except Exception: + try: + sandbox.destroy() + finally: + TaskRun.clear_sandbox_connection_state_atomic(ctx.run_id, sandbox.id) + raise activity.logger.info(f"Created sandbox {sandbox.id} (used_snapshot={used_snapshot})") 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 6f4482f842dc..d05aeca6e84c 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 @@ -24,6 +24,7 @@ vm_sandbox_default_base_origin_products, ) from products.tasks.backend.exceptions import TaskInvalidStateError, TaskRunNotReadyError +from products.tasks.backend.facade.api import ensure_task_run_session from products.tasks.backend.feature_flags import is_agent_otel_telemetry_enabled from products.tasks.backend.logic.services.sandbox_config import ( MAX_SANDBOX_CPU_CORES, @@ -65,6 +66,7 @@ class TaskProcessingContext: repository: str | None distinct_id: str origin_product: str | None = None + task_runtime: str = Task.Runtime.ACP environment: str | None = None github_user_integration_id: str | None = None task_created_by_id: int | None = None @@ -107,7 +109,6 @@ class TaskProcessingContext: @property def mode(self) -> str: - """Get the execution mode from state. Defaults to 'background'.""" return (self.state or {}).get("mode", "background") @property @@ -679,6 +680,9 @@ def get_task_processing_context(input: GetTaskProcessingContextInput) -> TaskPro emit_agent_log(run_id, "debug", "Fetching task details") task: Task = task_run.task + if task.runtime == Task.Runtime.PI: + ensure_task_run_session(task_run.id) + team: Team = task.team organization_id = str(team.organization_id) if not task.created_by: @@ -797,12 +801,17 @@ def get_task_processing_context(input: GetTaskProcessingContextInput) -> TaskPro or False ) # Ensure we get a boolean value even if the flag is missing emit_agent_log(run_id, "debug", f"pr_loop_enabled: {pr_loop_enabled} for this task run") - sandbox_event_ingest_enabled = _is_sandbox_event_ingest_enabled( - distinct_id=distinct_id, - organization_id=organization_id, - run_id=run_id, - state=state, - ) + pi_persistent_streaming = task.runtime == Task.Runtime.PI and not is_slack_interaction_state(state) + sandbox_event_ingest_override = state.get("sandbox_event_ingest_enabled") + if pi_persistent_streaming and not isinstance(sandbox_event_ingest_override, bool): + sandbox_event_ingest_enabled = True + else: + sandbox_event_ingest_enabled = _is_sandbox_event_ingest_enabled( + distinct_id=distinct_id, + organization_id=organization_id, + run_id=run_id, + state=state, + ) emit_agent_log( run_id, "debug", @@ -927,6 +936,7 @@ def get_task_processing_context(input: GetTaskProcessingContextInput) -> TaskPro repository=task.repository, distinct_id=distinct_id, origin_product=task.origin_product, + task_runtime=task.runtime, environment=task_run.environment, task_created_by_id=task.created_by_id, create_pr=input.create_pr, 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 2c4ec913182f..82d21921ea43 100644 --- a/products/tasks/backend/temporal/process_task/activities/provision_sandbox.py +++ b/products/tasks/backend/temporal/process_task/activities/provision_sandbox.py @@ -623,17 +623,20 @@ def create_sandbox_for_repository(input: CreateSandboxForRepositoryInput) -> Cre if credentials.token: sandbox_state["sandbox_connect_token"] = credentials.token TaskRun.update_state_atomic(ctx.run_id, updates=sandbox_state) + open_sandbox_session( + run_id=ctx.run_id, + sandbox_id=sandbox.id, + config=sandbox.config, + sandbox_created_at=sandbox_created_at, + required=ctx.task_runtime == "pi", + ) except Exception: - sandbox.destroy() + try: + sandbox.destroy() + finally: + TaskRun.clear_sandbox_connection_state_atomic(ctx.run_id, sandbox.id) raise - # Best-effort usage-ledger row (swallows its own failures). After the state - # write on purpose: the except branch above destroys sandboxes that never - # became reachable, and those must not enter the ledger. - open_sandbox_session( - run_id=ctx.run_id, sandbox_id=sandbox.id, config=sandbox.config, sandbox_created_at=sandbox_created_at - ) - emit_agent_log(ctx.run_id, "debug", f"Sandbox provisioned: {sandbox.id}") activity.logger.info(f"Created sandbox {sandbox.id} (used_snapshot={actual_used_snapshot})") diff --git a/products/tasks/backend/temporal/process_task/activities/relay_sandbox_events.py b/products/tasks/backend/temporal/process_task/activities/relay_sandbox_events.py index e7f5c2ba2b23..62212f8fe4cc 100644 --- a/products/tasks/backend/temporal/process_task/activities/relay_sandbox_events.py +++ b/products/tasks/backend/temporal/process_task/activities/relay_sandbox_events.py @@ -592,8 +592,24 @@ def _is_session_update(event_data: dict) -> bool: ) +def _pi_conversation_event(event_data: dict) -> dict | None: + if event_data.get("type") != "pi_event": + return None + event = event_data.get("event") + return event if isinstance(event, dict) else None + + def _is_active_agent_update(event_data: dict) -> bool: - """True only for session/update events where the agent is actively generating.""" + """True only for events where the agent is actively generating.""" + pi_event = _pi_conversation_event(event_data) + if pi_event is not None: + return pi_event.get("type") in { + "assistant_message_chunk", + "assistant_thought_chunk", + "tool_call_started", + "tool_call_updated", + "user_message", + } if not _is_session_update(event_data): return False update = (event_data.get("notification", {}).get("params") or {}).get("update") or {} @@ -675,7 +691,15 @@ def _tool_args_preview(raw_input: Any) -> str | None: def _extract_agent_message_text(event_data: dict) -> str | None: - """Text delta from an ACP agent_message_chunk session/update, else None.""" + """Text delta from an agent message event, else None.""" + pi_event = _pi_conversation_event(event_data) + if pi_event is not None and pi_event.get("type") == "assistant_message_chunk": + content = pi_event.get("content") + if isinstance(content, dict) and content.get("type") == "text": + text = content.get("text") + return text if isinstance(text, str) else None + return None + notification = event_data.get("notification", {}) if notification.get("method") != "session/update": return None @@ -730,7 +754,11 @@ def _is_keepalive_event(event_data: dict) -> bool: return event_data.get("type") == "keepalive" -_is_end_of_turn = is_turn_complete +def _is_end_of_turn(event_data: dict) -> bool: + pi_event = _pi_conversation_event(event_data) + if pi_event is not None: + return pi_event.get("type") == "turn_completed" + return is_turn_complete(event_data) async def _emit_agentsh_events(sandbox_id: str, run_id: str, last_ts_ns: list[int]) -> None: diff --git a/products/tasks/backend/temporal/process_task/activities/send_followup_to_sandbox.py b/products/tasks/backend/temporal/process_task/activities/send_followup_to_sandbox.py index 45fbf91cbeb2..d68fd93fd924 100644 --- a/products/tasks/backend/temporal/process_task/activities/send_followup_to_sandbox.py +++ b/products/tasks/backend/temporal/process_task/activities/send_followup_to_sandbox.py @@ -74,8 +74,6 @@ class SendFollowupToSandboxInput: message: str | None = None posthog_mcp_scopes: PosthogMcpScopes = "read_only" artifact_ids: list[str] | None = None - # Idempotency key, stable across retries and redeliveries; the - # agent-server drops a duplicate it already accepted. message_id: str | None = None # Sender of this message; None (older senders, pre-rollout histories) # falls back to the run-state actor. @@ -83,6 +81,7 @@ class SendFollowupToSandboxInput: # Signal context, passed through from PendingFollowup. context: dict[str, Any] | None = None steer: bool = False + max_attempts: int = SEND_FOLLOWUP_MAX_ATTEMPTS @activity.defn @@ -292,7 +291,7 @@ def _deliver_followup(input: SendFollowupToSandboxInput) -> str | None: # releases the id when a delivered turn fails before completion. attempt = _current_attempt() failure_kind = "delivery unknown" if result.status_code == 504 else "retryable failure" - if attempt < SEND_FOLLOWUP_MAX_ATTEMPTS: + if attempt < input.max_attempts: logger.warning( "send_followup_retrying", run_id=input.run_id, diff --git a/products/tasks/backend/temporal/process_task/activities/start_agent_server.py b/products/tasks/backend/temporal/process_task/activities/start_agent_server.py index cbd2be5df441..459c0652d200 100644 --- a/products/tasks/backend/temporal/process_task/activities/start_agent_server.py +++ b/products/tasks/backend/temporal/process_task/activities/start_agent_server.py @@ -174,6 +174,7 @@ class _LaunchParams: agentsh_domains: list[str] | None protected_base_branch: str | None event_ingest_token: str | None + task_run_session_token: str | None event_ingest_url: str | None event_ingest_keep_stream_open: bool @@ -193,7 +194,7 @@ def _include_personal_mcp_for_task(task: Task) -> bool: return not task.internal -def _prepare_launch(ctx: TaskProcessingContext, scopes: PosthogMcpScopes) -> _LaunchParams: +def _prepare_launch(ctx: TaskProcessingContext, scopes: PosthogMcpScopes, sandbox_id: str) -> _LaunchParams: try: task = Task.objects.select_related("created_by", "team").get(id=ctx.task_id) actor_user = get_task_run_credential_user(task, ctx.state) @@ -221,15 +222,20 @@ def _prepare_launch(ctx: TaskProcessingContext, scopes: PosthogMcpScopes) -> _La {"task_id": ctx.task_id, "run_id": ctx.run_id}, cause=TaskRun.DoesNotExist(f"TaskRun {ctx.run_id} not found"), ) - if event_stream_ingest_enabled: + task_run_session_token: str | None = None + if event_stream_ingest_enabled or task.runtime == Task.Runtime.PI: try: - event_ingest_token = create_sandbox_event_ingest_token(task_run) + run_token = create_sandbox_event_ingest_token(task_run, sandbox_id=sandbox_id) except Exception as e: raise SandboxExecutionError( - "Failed to create sandbox event ingest token", + "Failed to create sandbox task run token", {"task_id": ctx.task_id, "run_id": ctx.run_id, "error": str(e)}, cause=e, ) + if event_stream_ingest_enabled: + event_ingest_token = run_token + if task.runtime == Task.Runtime.PI: + task_run_session_token = run_token mcp_configs = get_sandbox_ph_mcp_configs( token=access_token, @@ -307,6 +313,7 @@ def _prepare_launch(ctx: TaskProcessingContext, scopes: PosthogMcpScopes) -> _La agentsh_domains=agentsh_domains, protected_base_branch=protected_base_branch, event_ingest_token=event_ingest_token, + task_run_session_token=task_run_session_token, event_ingest_url=event_ingest_url, event_ingest_keep_stream_open=ctx.agent_proxy_keep_stream_open, ) @@ -330,6 +337,7 @@ def _invoke_start_agent_server( auto_publish=ctx.auto_publish, interaction_origin=ctx.interaction_origin, branch=params.protected_base_branch, + agent_runtime=ctx.task_runtime, runtime_adapter=ctx.runtime_adapter, provider=ctx.provider, model=ctx.model, @@ -341,6 +349,7 @@ def _invoke_start_agent_server( relayed_mcp_servers=params.relayed_mcp_servers or None, allowed_domains=params.agentsh_domains, event_ingest_token=params.event_ingest_token, + task_run_session_token=params.task_run_session_token, event_ingest_url=params.event_ingest_url, event_ingest_keep_stream_open=params.event_ingest_keep_stream_open, repo_ready_file=repo_ready_file, @@ -442,7 +451,7 @@ def start_agent_server(input: StartAgentServerInput) -> StartAgentServerOutput: # repo directory can never appear later. The deferred/overlap path clones in parallel # and gates the session on the repo-ready barrier instead. _ensure_repository_on_disk(ctx, sandbox) - params = _prepare_launch(ctx, input.posthog_mcp_scopes) + params = _prepare_launch(ctx, input.posthog_mcp_scopes, input.sandbox_id) with StepTimer("agent_server_ready", boot_path=input.boot_path) as ready_timer: _invoke_start_agent_server(sandbox, ctx, params, repo_ready_file=None, wait_for_health=True) @@ -480,7 +489,7 @@ def launch_agent_server(input: StartAgentServerInput) -> StartAgentServerOutput: emit_agent_log(ctx.run_id, "debug", "Launching agent server (deferred readiness)") sandbox = Sandbox.get_by_id(input.sandbox_id) - params = _prepare_launch(ctx, input.posthog_mcp_scopes) + params = _prepare_launch(ctx, input.posthog_mcp_scopes, input.sandbox_id) repo_ready_file = REPO_READY_FILE if input.defer_for_clone else None with StepTimer("agent_server_launch", boot_path=input.boot_path) as launch_timer: 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 cab7cff82544..bcd6f576f6e0 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 @@ -360,6 +360,41 @@ def feature_enabled(flag_key, **kwargs): sandbox_args, _sandbox_kwargs = feature_enabled_mock.call_args_list[1] assert sandbox_args[0] == SANDBOX_EVENT_INGEST_FEATURE_FLAG + @pytest.mark.django_db(transaction=True) + def test_pi_runtime_enables_event_ingest_without_bypassing_persistent_upload_rollout( + self, activity_environment, test_task + ): + test_task.runtime = Task.Runtime.PI + test_task.save(update_fields=["runtime"]) + task_run = test_task.create_run() + input_data = GetTaskProcessingContextInput(run_id=str(task_run.id)) + + with patch( + "products.tasks.backend.temporal.process_task.activities.get_task_processing_context.posthoganalytics.feature_enabled", + return_value=False, + ): + result = async_to_sync(activity_environment.run)(get_task_processing_context, input_data) + + assert result.sandbox_event_ingest_enabled is True + assert result.agent_proxy_keep_stream_open is False + + @pytest.mark.django_db(transaction=True) + def test_pi_runtime_respects_persistent_event_streaming_kill_switches(self, activity_environment, test_task): + test_task.runtime = Task.Runtime.PI + test_task.save(update_fields=["runtime"]) + task_run = test_task.create_run( + extra_state={ + "sandbox_event_ingest_enabled": False, + "agent_proxy_keep_stream_open": False, + } + ) + input_data = GetTaskProcessingContextInput(run_id=str(task_run.id)) + + result = async_to_sync(activity_environment.run)(get_task_processing_context, input_data) + + assert result.sandbox_event_ingest_enabled is False + assert result.agent_proxy_keep_stream_open is False + @pytest.mark.django_db(transaction=True) def test_pr_loop_enabled_for_signal_report_origin_ignores_flag(self, activity_environment, test_task): # Signals implementation PRs are bot-authored and always opt into the PR @@ -994,22 +1029,21 @@ def test_get_task_processing_context_exposes_ci_prompt(self, activity_environmen assert result.ci_prompt == custom_prompt @pytest.mark.django_db(transaction=True) - def test_get_task_processing_context_exposes_runtime_metadata(self, activity_environment, test_task): - task_run = test_task.create_run( - extra_state={ - "runtime_adapter": "codex", - "provider": "openai", - "model": "gpt-5.3-codex", - "reasoning_effort": "high", - "initial_permission_mode": "plan", - } - ) + def test_get_task_processing_context_creates_native_pi_session(self, activity_environment, test_task): + test_task.runtime = Task.Runtime.PI + test_task.save(update_fields=["runtime"]) + task_run = test_task.create_run() input_data = GetTaskProcessingContextInput(run_id=str(task_run.id)) result = async_to_sync(activity_environment.run)(get_task_processing_context, input_data) - assert result.runtime_adapter == "codex" - assert result.provider == "openai" - assert result.model == "gpt-5.3-codex" - assert result.reasoning_effort == "high" - assert result.initial_permission_mode == "plan" + task_run.refresh_from_db() + assert task_run.active_task_session is not None + assert task_run.active_task_session.object_storage_key is None + assert task_run.active_task_session.team_id == test_task.team_id + assert result.task_runtime == "pi" + assert result.runtime_adapter is None + assert result.provider is None + assert result.model is None + assert result.reasoning_effort is None + assert result.initial_permission_mode is None diff --git a/products/tasks/backend/temporal/process_task/activities/tests/test_relay_sandbox_events.py b/products/tasks/backend/temporal/process_task/activities/tests/test_relay_sandbox_events.py index 022c4170c7da..9c48c63dce23 100644 --- a/products/tasks/backend/temporal/process_task/activities/tests/test_relay_sandbox_events.py +++ b/products/tasks/backend/temporal/process_task/activities/tests/test_relay_sandbox_events.py @@ -65,6 +65,11 @@ class TestIsEndOfTurn: {"type": "event", "notification": {"result": {"stopReason": "end_turn"}}}, False, ), + ( + "pi_turn_complete", + {"type": "pi_event", "event": {"type": "turn_completed"}}, + True, + ), ] ) def test_is_end_of_turn(self, _name: str, event_data: dict, expected: bool): @@ -150,6 +155,9 @@ def _su(sub_type: str) -> dict: def test_session_update_sub_types(self, _name: str, sub_type: str, expected: bool) -> None: assert _is_active_agent_update(self._su(sub_type)) is expected + def test_pi_generation_event_is_active(self) -> None: + assert _is_active_agent_update({"type": "pi_event", "event": {"type": "assistant_message_chunk"}}) + @parameterized.expand( [ ("missing_session_update_key", {"update": {}}), diff --git a/products/tasks/backend/temporal/process_task/activities/tests/test_start_agent_server.py b/products/tasks/backend/temporal/process_task/activities/tests/test_start_agent_server.py index 29773c1cff61..907390a5b75b 100644 --- a/products/tasks/backend/temporal/process_task/activities/tests/test_start_agent_server.py +++ b/products/tasks/backend/temporal/process_task/activities/tests/test_start_agent_server.py @@ -177,6 +177,7 @@ async def test_start_agent_server_uses_captured_sandbox_event_ingest_flag(mocker assert result.sandbox_url == "https://sandbox.example" assert result.connect_token == "connect-token" create_event_ingest_token.assert_called_once() + assert create_event_ingest_token.call_args.kwargs == {"sandbox_id": "sandbox-id"} sandbox.start_agent_server.assert_called_once() assert sandbox.start_agent_server.call_args.kwargs["event_ingest_token"] == "event-ingest-token" diff --git a/products/tasks/backend/temporal/process_task/tests/test_send_followup_to_sandbox.py b/products/tasks/backend/temporal/process_task/tests/test_send_followup_to_sandbox.py index e5f0e86fee5f..09b2e8451144 100644 --- a/products/tasks/backend/temporal/process_task/tests/test_send_followup_to_sandbox.py +++ b/products/tasks/backend/temporal/process_task/tests/test_send_followup_to_sandbox.py @@ -709,6 +709,25 @@ def test_retryable_failure_retries_without_sentinel(self, _patches): _patches["error"].assert_not_called() _patches["turn_complete"].assert_not_called() + def test_pi_retryable_failure_writes_sentinel_on_its_only_attempt(self, _patches): + _patches["user_msg"].return_value = CommandResult( + success=False, status_code=502, error="Connection to sandbox failed", retryable=True + ) + + with pytest.raises(ApplicationError, match="send_followup failed") as exc_info: + send_followup_to_sandbox( + SendFollowupToSandboxInput( + run_id="run-1", + message="hi", + message_id="m-1", + max_attempts=1, + ) + ) + + assert exc_info.value.non_retryable is True + _patches["error"].assert_called_once() + _patches["turn_complete"].assert_not_called() + def test_retryable_stream_error_final_attempt_writes_actionable_sentinel(self, _patches): _patches["user_msg"].return_value = CommandResult( success=False, diff --git a/products/tasks/backend/temporal/process_task/tests/test_workflow.py b/products/tasks/backend/temporal/process_task/tests/test_workflow.py index 77715f4bdf05..6f552b702ed3 100644 --- a/products/tasks/backend/temporal/process_task/tests/test_workflow.py +++ b/products/tasks/backend/temporal/process_task/tests/test_workflow.py @@ -377,6 +377,41 @@ async def fake_send_followup( assert deliveries == [("from Slack", "message-123", False)] + async def test_duplicate_sender_message_id_is_queued_once(self, monkeypatch): + workflow = ProcessTaskWorkflow() + workflow._context = _build_context(github_integration_id=123) + monkeypatch.setattr(process_task_workflow_module.workflow, "deprecate_patch", Mock()) + monkeypatch.setattr(process_task_workflow_module.workflow, "logger", Mock()) + + await workflow.send_followup_message("first", [], "message-123") + await workflow.send_followup_message("duplicate", [], "message-123") + + assert [(item.message, item.message_id) for item in workflow._pending_followups] == [("first", "message-123")] + + async def test_sender_message_id_deduplication_is_bounded(self, monkeypatch): + workflow = ProcessTaskWorkflow() + workflow._context = _build_context(github_integration_id=123) + monkeypatch.setattr(process_task_workflow_module.workflow, "deprecate_patch", Mock()) + monkeypatch.setattr(process_task_workflow_module.workflow, "logger", Mock()) + + for index in range(501): + await workflow.send_followup_message(f"message {index}", [], f"message-{index}") + await workflow.send_followup_message("duplicate", [], "message-1") + + assert len(workflow._accepted_message_ids) == 500 + assert len(workflow._pending_followups) == 501 + + async def test_same_message_id_from_different_senders_is_not_deduplicated(self, monkeypatch): + workflow = ProcessTaskWorkflow() + workflow._context = _build_context(github_integration_id=123) + monkeypatch.setattr(process_task_workflow_module.workflow, "deprecate_patch", Mock()) + monkeypatch.setattr(process_task_workflow_module.workflow, "logger", Mock()) + + await workflow.send_followup_message("first", [], "message-1", actor_user_id=1) + await workflow.send_followup_message("second", [], "message-1", actor_user_id=2) + + assert [item.message for item in workflow._pending_followups] == ["first", "second"] + async def test_native_steer_preserves_sender_identity(self, monkeypatch): workflow = ProcessTaskWorkflow() workflow._context = _build_context(github_integration_id=123) diff --git a/products/tasks/backend/temporal/process_task/workflow.py b/products/tasks/backend/temporal/process_task/workflow.py index 843879d046f7..9147fd5cf9cb 100644 --- a/products/tasks/backend/temporal/process_task/workflow.py +++ b/products/tasks/backend/temporal/process_task/workflow.py @@ -100,6 +100,7 @@ from .slack_agent_design_relay import SlackAgentDesignRelayInput, SlackAgentDesignRelayWorkflow DEAD_SANDBOX_ERROR_TYPES = ("SandboxNotRunningError", "SandboxNotFoundError") +MAX_ACCEPTED_MESSAGE_IDS = 500 def _is_dead_sandbox_failure(error: BaseException) -> bool: @@ -107,6 +108,16 @@ def _is_dead_sandbox_failure(error: BaseException) -> bool: return isinstance(cause, temporalio.exceptions.ApplicationError) and cause.type in DEAD_SANDBOX_ERROR_TYPES +def _message_dedupe_key( + message_id: str, + actor_user_id: int | None, + message_context: dict[str, Any] | None, +) -> str: + slack_user_id = (message_context or {}).get("actor_slack_user_id") + actor_slack_user_id = slack_user_id if isinstance(slack_user_id, str) else "" + return f"{actor_user_id or ''}:{actor_slack_user_id}:{message_id}" + + @dataclass class ResumedSandboxState: """Loop state carried across continue_as_new to re-attach without re-provisioning.""" @@ -122,6 +133,7 @@ class ResumedSandboxState: last_active_time: Optional[str] # ISO8601, or None if never active # Defaulted so continue_as_new payloads from pre-rollout runs deserialize. pr_unresolved_threads: int = 0 + accepted_message_ids: list[str] = field(default_factory=list) @dataclass @@ -131,6 +143,7 @@ class ProcessTaskInput: slack_thread_context: Optional[dict[str, Any]] = None posthog_mcp_scopes: PosthogMcpScopes = "read_only" prewarmed: bool = False + initial_message: Optional["PendingFollowup"] = None # Set only on a continue_as_new continuation, to skip provisioning and re-attach. resumed_sandbox: Optional[ResumedSandboxState] = None @@ -140,8 +153,6 @@ class PendingFollowup: message: str | None artifact_ids: list[str] actor_user_id: int | None = None - # Sender-supplied idempotency key (stable across the sender's retries); - # None falls back to a workflow-generated id. message_id: str | None = None # Signal context carried verbatim (e.g. actor_slack_user_id for reply # tagging); consumers validate the keys they read. @@ -287,6 +298,8 @@ def __init__(self) -> None: self._pending_followup: PendingFollowup | None = None self._pending_followups: list[PendingFollowup] = [] self._next_followup_sequence: int = 0 + self._accepted_message_ids: list[str] = [] + self._accepted_message_id_set: set[str] = set() self._active_followup_task: asyncio.Task[None] | None = None self._shutting_down: bool = False self._pending_permission_responses: list[PendingPermissionResponse] = [] @@ -329,6 +342,11 @@ def parse_inputs(inputs: list[str]) -> ProcessTaskInput: slack_thread_context=loaded.get("slack_thread_context"), posthog_mcp_scopes=loaded.get("posthog_mcp_scopes", "read_only"), prewarmed=loaded.get("prewarmed", False), + initial_message=( + PendingFollowup(**loaded["initial_message"]) + if isinstance(loaded.get("initial_message"), dict) + else None + ), resumed_sandbox=ResumedSandboxState(**resumed) if resumed else None, ) @@ -725,7 +743,10 @@ async def run(self, input: ProcessTaskInput) -> ProcessTaskOutput: ) # A continuation already delivered the first user message in a prior execution. - if input.resumed_sandbox is None and self._should_forward_pending_user_message(): + if input.resumed_sandbox is None and input.initial_message is not None: + self._pending_followups.append(input.initial_message) + await self._dispatch_next_followup() + elif input.resumed_sandbox is None and self._should_forward_pending_user_message(): await self._forward_pending_user_message() # Wait for completion signal or inactivity timeout. @@ -1111,6 +1132,7 @@ def _build_resumed_input(self, input: ProcessTaskInput, sandbox_id: str) -> Proc first_user_message_received=self._first_user_message_received, is_agent_design_enabled=self._is_agent_design_enabled, last_active_time=self._last_active_time.isoformat() if self._last_active_time else None, + accepted_message_ids=self._accepted_message_ids, ), ) @@ -1121,15 +1143,18 @@ def _restore_resumed_state(self, resumed: ResumedSandboxState) -> None: self._pr_unresolved_threads = resumed.pr_unresolved_threads self._pr_progress_emitted = resumed.pr_progress_emitted self._first_user_message_received = resumed.first_user_message_received + self._accepted_message_ids = resumed.accepted_message_ids + self._accepted_message_id_set = set(resumed.accepted_message_ids) self._last_active_time = datetime.fromisoformat(resumed.last_active_time) if resumed.last_active_time else None async def _get_task_processing_context(self, input: ProcessTaskInput) -> TaskProcessingContext: - return await workflow.execute_activity( + context = await workflow.execute_activity( get_task_processing_context, GetTaskProcessingContextInput(run_id=input.run_id, create_pr=input.create_pr), start_to_close_timeout=timedelta(minutes=2), retry_policy=RetryPolicy(maximum_attempts=3), ) + return context async def _get_sandbox_for_repository(self) -> GetSandboxForRepositoryOutput: prepared = await workflow.execute_activity( @@ -2005,6 +2030,16 @@ def _queue_followup_message( extra={"run_id": context.run_id if context is not None else None}, ) return + if message_id: + dedupe_key = _message_dedupe_key(message_id, actor_user_id, message_context) + if dedupe_key in self._accepted_message_id_set: + return + if len(self._accepted_message_ids) >= MAX_ACCEPTED_MESSAGE_IDS: + oldest_key = self._accepted_message_ids.pop(0) + self._accepted_message_id_set.discard(oldest_key) + self._accepted_message_ids.append(dedupe_key) + self._accepted_message_id_set.add(dedupe_key) + pending_followup = PendingFollowup( message=message, artifact_ids=artifact_ids or [], @@ -2088,6 +2123,7 @@ async def _send_followup_to_sandbox( }, ) try: + max_attempts = 1 if self.context.task_runtime == "pi" else SEND_FOLLOWUP_MAX_ATTEMPTS return await workflow.execute_activity( send_followup_to_sandbox, SendFollowupToSandboxInput( @@ -2099,17 +2135,13 @@ async def _send_followup_to_sandbox( actor_user_id=actor_user_id, context=context, steer=steer, + max_attempts=max_attempts, ), start_to_close_timeout=timedelta(minutes=35), - # The activity heartbeats while blocked on the sync delivery - # call, so a worker restart is detected here instead of at - # start_to_close. Retries are safe: message_id lets the - # agent-server drop a redelivery it already accepted, and - # sentinel-writing failures raise non-retryable. heartbeat_timeout=timedelta(minutes=1), retry_policy=RetryPolicy( initial_interval=timedelta(seconds=5), - maximum_attempts=SEND_FOLLOWUP_MAX_ATTEMPTS, + maximum_attempts=max_attempts, ), ) except Exception as e: diff --git a/products/tasks/backend/tests/test_agent_proxy_callback.py b/products/tasks/backend/tests/test_agent_proxy_callback.py index fbf2be386173..9616d8684fb8 100644 --- a/products/tasks/backend/tests/test_agent_proxy_callback.py +++ b/products/tasks/backend/tests/test_agent_proxy_callback.py @@ -40,7 +40,10 @@ def _url(self, run_id: str | None = None) -> str: return f"/internal/tasks/runs/{run_id or self.task_run.id}/agent-proxy-callback/" def _token(self, run: TaskRun | None = None) -> str: - return create_sandbox_event_ingest_token(run or self.task_run) + token_run = run or self.task_run + token_run.state = {**(token_run.state or {}), "sandbox_id": f"sandbox-{token_run.id}"} + token_run.save(update_fields=["state", "updated_at"]) + return create_sandbox_event_ingest_token(token_run) def _body(self, **overrides: Any) -> dict[str, Any]: body: dict[str, Any] = { diff --git a/products/tasks/backend/tests/test_agentsh.py b/products/tasks/backend/tests/test_agentsh.py index a5bf7e0bbb1e..545ccc80beaf 100644 --- a/products/tasks/backend/tests/test_agentsh.py +++ b/products/tasks/backend/tests/test_agentsh.py @@ -398,6 +398,7 @@ def test_command_without_domains_skips_agentsh_exec(self): from products.tasks.backend.logic.services.modal_sandbox import ModalSandbox sandbox = ModalSandbox.__new__(ModalSandbox) + sandbox.id = "sb-test" cmd = sandbox._build_agent_server_command( repo_path="/tmp/workspace/repos/org/repo", task_id="test-task", @@ -427,6 +428,7 @@ def test_command_includes_auto_publish_flag_only_when_opted_in(self, provider, a sandbox = ModalSandbox.__new__(ModalSandbox) else: sandbox = DockerSandbox.__new__(DockerSandbox) + sandbox.id = "sb-test" cmd = sandbox._build_agent_server_command( repo_path="/tmp/workspace/repos/org/repo", task_id="test-task", @@ -500,6 +502,7 @@ def test_command_includes_allowed_domains(self): from products.tasks.backend.logic.services.modal_sandbox import ModalSandbox sandbox = ModalSandbox.__new__(ModalSandbox) + sandbox.id = "sb-test" cmd = sandbox._build_agent_server_command( repo_path="/tmp/workspace/repos/org/repo", task_id="test-task", @@ -518,6 +521,7 @@ def test_command_includes_runtime_environment_variables(self): from products.tasks.backend.logic.services.modal_sandbox import ModalSandbox sandbox = ModalSandbox.__new__(ModalSandbox) + sandbox.id = "sb-test" cmd = sandbox._build_agent_server_command( repo_path="/tmp/workspace/repos/org/repo", task_id="test-task", diff --git a/products/tasks/backend/tests/test_api.py b/products/tasks/backend/tests/test_api.py index fb5fa4ab19ec..3bb8b99da24a 100644 --- a/products/tasks/backend/tests/test_api.py +++ b/products/tasks/backend/tests/test_api.py @@ -3,6 +3,7 @@ import uuid import base64 import asyncio +import hashlib from collections.abc import AsyncGenerator, Iterator from datetime import timedelta from typing import Any, ClassVar, cast @@ -39,6 +40,7 @@ get_posthog_code_usage, ) from products.tasks.backend.logic.services.connection_token import ( + create_sandbox_event_ingest_token, get_sandbox_jwt_public_key, reset_sandbox_jwt_key_cache, ) @@ -54,14 +56,17 @@ get_task_run_stream_key, ) from products.tasks.backend.models import ( + Channel, CodeInvite, CodeInviteRedemption, SandboxCustomImage, SandboxEnvironment, + SandboxSession, Task, TaskArtifact, TaskAutomation, TaskRun, + TaskSession, ) from products.tasks.backend.presentation.serializers import ( TASK_RUN_ARTIFACT_MAX_SIZE_BYTES, @@ -164,7 +169,7 @@ def set_tasks_feature_flag(self, enabled=True): self.mock_feature_flag = self.feature_flag_patcher.start() def check_flag(flag_name, *_args, **_kwargs): - if flag_name == "tasks": + if flag_name in {"tasks", "pi-harness"}: return enabled return False @@ -1631,14 +1636,34 @@ def test_run_endpoint_triggers_workflow(self, mock_workflow): self.assertEqual(latest_run["environment"], "cloud") @patch("products.tasks.backend.temporal.client.execute_task_processing_workflow") - def test_run_endpoint_rejects_pi_task(self, mock_workflow): + def test_run_endpoint_starts_pi_task(self, mock_workflow): task = self.create_task(runtime=Task.Runtime.PI) response = self.client.post(f"/api/projects/@current/tasks/{task.id}/run/") - self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) - self.assertEqual(response.json()["detail"], "Pi tasks cannot be run through the ACP task workflow.") - self.assertFalse(task.runs.exists()) + self.assertEqual(response.status_code, status.HTTP_200_OK) + run = task.runs.get() + mock_workflow.assert_called_once() + workflow_input = mock_workflow.call_args.kwargs + self.assertEqual(workflow_input["task_id"], str(task.id)) + self.assertEqual(workflow_input["run_id"], str(run.id)) + self.assertEqual(workflow_input["team_id"], task.team.id) + self.assertEqual(workflow_input["user_id"], self.user.id) + self.assertEqual(workflow_input["posthog_mcp_scopes"], "full") + self.assertEqual(workflow_input["initial_message"].message, "Test Description") + self.assertEqual(workflow_input["initial_message"].artifact_ids, []) + self.assertNotIn("mode", run.state) + self.assertNotIn("pending_user_message", run.state) + + @patch("products.tasks.backend.temporal.client.execute_task_processing_workflow") + def test_run_endpoint_rejects_pi_task_when_disabled(self, mock_workflow): + task = self.create_task(runtime=Task.Runtime.PI) + self.set_tasks_feature_flag(False) + + response = self.client.post(f"/api/projects/@current/tasks/{task.id}/run/") + + self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN) + self.assertEqual(response.json()["error"], "Pi cloud runtime is disabled") mock_workflow.assert_not_called() @parameterized.expand( @@ -2250,14 +2275,29 @@ def test_start_run_endpoint_triggers_workflow_for_existing_cloud_run(self, mock_ ) @patch("products.tasks.backend.temporal.client.execute_task_processing_workflow") - def test_start_run_endpoint_rejects_pi_task(self, mock_workflow): + def test_start_run_endpoint_starts_pi_task(self, mock_workflow): task = self.create_task(runtime=Task.Runtime.PI) task_run = task.create_run(environment=TaskRun.Environment.CLOUD) response = self.client.post(f"/api/projects/@current/tasks/{task.id}/runs/{task_run.id}/start/") - self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) - self.assertEqual(response.json()["error"], "Pi tasks cannot be run through the ACP task workflow.") + self.assertEqual(response.status_code, status.HTTP_200_OK) + mock_workflow.assert_called_once() + workflow_input = mock_workflow.call_args.kwargs + self.assertEqual(workflow_input["task_id"], str(task.id)) + self.assertEqual(workflow_input["run_id"], str(task_run.id)) + self.assertEqual(workflow_input["initial_message"].message, "Test Description") + self.assertEqual(workflow_input["initial_message"].artifact_ids, []) + self.assertNotIn("pending_user_message", task_run.state) + + @patch("products.tasks.backend.temporal.client.execute_task_processing_workflow") + def test_start_run_endpoint_returns_not_found_before_pi_runtime_gate(self, mock_workflow): + task = self.create_task(runtime=Task.Runtime.PI) + self.set_tasks_feature_flag(False) + + response = self.client.post(f"/api/projects/@current/tasks/{task.id}/runs/{uuid.uuid4()}/start/") + + self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND) mock_workflow.assert_not_called() @patch("products.tasks.backend.temporal.client.execute_task_processing_workflow") @@ -4512,21 +4552,27 @@ def test_update_run_status_to_completed_signals_workflow(self, mock_signal): self.assertEqual(run.status, TaskRun.Status.COMPLETED) self.assertIsNotNone(run.completed_at) + @patch("products.tasks.backend.presentation.views.api.tasks_facade.pi_cloud_runtime_enabled", return_value=True) @patch("products.tasks.backend.temporal.client.resume_task_in_cloud_workflow") - def test_resume_in_cloud_rejects_pi_task(self, mock_resume): + @patch("products.tasks.backend.facade.streams.reset_task_run_stream", return_value=True) + def test_resume_in_cloud_starts_pi_task(self, mock_reset_stream, mock_resume, _mock_pi_enabled): task = self.create_task(runtime=Task.Runtime.PI) run = TaskRun.objects.create( task=task, team=self.team, environment=TaskRun.Environment.LOCAL, status=TaskRun.Status.COMPLETED, + state={"pr_authorship_mode": "bot"}, ) response = self.client.post(f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/resume_in_cloud/") - self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) - self.assertEqual(response.json()["error"], "Pi tasks cannot be run through the ACP task workflow.") - mock_resume.assert_not_called() + self.assertEqual(response.status_code, status.HTTP_200_OK) + run.refresh_from_db() + self.assertEqual(run.environment, TaskRun.Environment.CLOUD) + self.assertEqual(run.status, TaskRun.Status.QUEUED) + mock_reset_stream.assert_called_once_with(str(run.id), use_dedicated=False) + mock_resume.assert_called_once_with(str(run.id), run.workflow_id) @patch("products.tasks.backend.temporal.client.resume_task_in_cloud_workflow") def test_resume_in_cloud_rejects_user_authorship_without_github_identity_when_no_repo(self, mock_resume): @@ -6650,6 +6696,25 @@ def flag_enabled(key, *args, **kwargs): self.assertEqual(response.status_code, status.HTTP_200_OK) self.assertIsNone(response.json()["stream_base_url"]) + def test_stream_token_returns_proxy_url_for_pi_when_flag_disabled(self): + task = self.create_task(runtime=Task.Runtime.PI) + run = TaskRun.objects.create(task=task, team=self.team, status=TaskRun.Status.IN_PROGRESS) + + def flag_enabled(key, *args, **kwargs): + return key != "tasks-stream-via-proxy" + + with ( + self.settings(TASKS_AGENT_PROXY_PUBLIC_URL="https://agent-proxy.example.com", DEBUG=False), + patch( + "products.tasks.backend.facade.api.posthoganalytics.feature_enabled", + side_effect=flag_enabled, + ), + ): + response = self.client.get(f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/stream_token/") + + self.assertEqual(response.status_code, status.HTTP_200_OK) + self.assertEqual(response.json()["stream_base_url"], "https://agent-proxy.example.com") + def test_stream_token_returns_proxy_url_in_debug_without_flag(self): # Local dev (DEBUG) disables the analytics SDK, so the URL setting alone opts in. task = self.create_task() @@ -7968,6 +8033,7 @@ def test_repository_readiness_requires_repository(self): self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) +@override_settings(SANDBOX_JWT_PRIVATE_KEY=TEST_RSA_PRIVATE_KEY) class TestTaskRunCommandAPI(BaseTaskAPITest): def _command_url(self, task, run): return f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/command/" @@ -7994,6 +8060,21 @@ def _create_run_with_sandbox(self, task, sandbox_url="http://localhost:9999", co state=state, ) + def _open_sandbox_session(self, run, sandbox_id="sandbox-1"): + now = django_timezone.now() + run.state = {**(run.state or {}), "sandbox_id": sandbox_id} + run.save(update_fields=["state", "updated_at"]) + return SandboxSession.objects.unscoped().create( + team=self.team, + task_run=run, + sandbox_id=sandbox_id, + cpu_cores=1, + memory_gb=1, + ttl_seconds=3600, + created_at=now, + ttl_expires_at=now + timedelta(hours=1), + ) + def _mock_agent_response(self, mock_post, body, status_code=200): mock_resp = MagicMock() mock_resp.status_code = status_code @@ -8002,20 +8083,23 @@ def _mock_agent_response(self, mock_post, body, status_code=200): mock_resp.text = json.dumps(body) if isinstance(body, dict) else str(body) mock_post.return_value = mock_resp - @patch("products.tasks.backend.temporal.client.signal_task_followup_message") - def test_command_rejects_pi_task(self, mock_signal_followup): + def test_command_rejects_unsupported_acp_method_for_pi_task(self): task = self.create_task(runtime=Task.Runtime.PI) run = self._create_run_with_sandbox(task) response = self.client.post( self._command_url(task, run), - self._make_user_message(), + { + "jsonrpc": "2.0", + "method": "permission_response", + "params": {"requestId": "request-1", "optionId": "allow"}, + "id": "req-1", + }, format="json", ) self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) - self.assertEqual(response.json()["error"], "Pi tasks do not support ACP task commands.") - mock_signal_followup.assert_not_called() + self.assertEqual(response.json()["error"], "permission_response is not supported for Pi tasks.") @patch("products.tasks.backend.temporal.client.signal_task_followup_message") def test_command_signals_user_message(self, mock_signal_followup): @@ -8034,7 +8118,7 @@ def test_command_signals_user_message(self, mock_signal_followup): self.assertTrue(data["result"]["queued"]) mock_signal_followup.assert_called_once_with( - run.workflow_id, "Hello agent", [], None, self.user.id, None, steer=False + run.workflow_id, "Hello agent", [], "req-1", self.user.id, None, steer=False ) @patch("products.tasks.backend.temporal.client.signal_task_followup_message") @@ -8055,7 +8139,7 @@ def test_command_signals_steer_intent(self, mock_signal_followup): self.assertEqual(response.status_code, status.HTTP_200_OK) mock_signal_followup.assert_called_once_with( - run.workflow_id, "Change direction", [], None, self.user.id, None, steer=True + run.workflow_id, "Change direction", [], "req-steer", self.user.id, None, steer=True ) @patch("products.tasks.backend.temporal.client.signal_task_followup_message") @@ -8093,7 +8177,7 @@ def test_command_signals_user_message_without_active_sandbox(self, mock_signal_f self.assertEqual(response.status_code, status.HTTP_200_OK) self.assertTrue(response.json()["result"]["queued"]) mock_signal_followup.assert_called_once_with( - run.workflow_id, "Hello agent", [], None, self.user.id, None, steer=False + run.workflow_id, "Hello agent", [], "req-1", self.user.id, None, steer=False ) @patch("products.tasks.backend.temporal.client.signal_task_followup_message") @@ -8127,7 +8211,13 @@ def test_command_signals_user_message_artifact_ids(self, mock_signal_followup): self.assertEqual(response.status_code, status.HTTP_200_OK) self.assertTrue(response.json()["result"]["queued"]) mock_signal_followup.assert_called_once_with( - run.workflow_id, "See attached", ["artifact-123"], None, self.user.id, None, steer=False + run.workflow_id, + "See attached", + ["artifact-123"], + "req-attachments", + self.user.id, + None, + steer=False, ) @patch("products.tasks.backend.temporal.client.signal_task_followup_message") @@ -8193,6 +8283,265 @@ def test_command_proxies_cancel(self, mock_post): self.assertEqual(response.status_code, status.HTTP_200_OK) self.assertTrue(response.json()["result"]["cancelled"]) + def test_empty_task_session_returns_read_only_storage_access(self): + task = self.create_task(runtime=Task.Runtime.PI) + run = self._create_run_with_sandbox(task) + task_session = TaskSession.create_for_task(task) + run.active_task_session = task_session + run.save(update_fields=["active_task_session"]) + + response = self.client.get(f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/task_session/") + + self.assertEqual(response.status_code, status.HTTP_200_OK) + self.assertEqual( + response.json(), + { + "id": str(task_session.id), + "download_url": None, + "content_sha256": None, + }, + ) + + @patch("products.tasks.backend.models.object_storage.tag") + @patch("posthog.storage.object_storage.write") + def test_task_session_sync_atomically_replaces_opaque_content(self, mock_write, mock_tag): + task = self.create_task(runtime=Task.Runtime.PI) + run = self._create_run_with_sandbox(task) + task_session = TaskSession.create_for_task(task) + run.active_task_session = task_session + run.save(update_fields=["active_task_session"]) + self._open_sandbox_session(run) + content = b"opaque native session content" + + response = self.client.generic( + "POST", + f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/task_session_sync/", + cast(str, content), + content_type="application/octet-stream", + HTTP_IF_MATCH='"none"', + HTTP_X_SANDBOX_ID="sandbox-1", + HTTP_X_TASK_RUN_TOKEN=create_sandbox_event_ingest_token(run), + ) + + expected_sha256 = hashlib.sha256(content).hexdigest() + self.assertEqual(response.status_code, status.HTTP_200_OK) + self.assertEqual(response.json(), {"id": str(task_session.id), "content_sha256": expected_sha256}) + task_session.refresh_from_db() + self.assertEqual(task_session.content_sha256, expected_sha256) + self.assertEqual(task_session.size, len(content)) + self.assertIsNotNone(task_session.object_storage_key) + mock_write.assert_called_once_with(task_session.object_storage_key, content) + mock_tag.assert_called_once() + + @patch("posthog.storage.object_storage.delete") + @patch("posthog.storage.object_storage.write") + def test_task_session_sync_rejects_stale_content_hash(self, mock_write, mock_delete): + task = self.create_task(runtime=Task.Runtime.PI) + run = self._create_run_with_sandbox(task) + task_session = TaskSession.create_for_task(task) + task_session.object_storage_key = "task-sessions/current.jsonl" + task_session.content_sha256 = "current-hash" + task_session.size = 32 + task_session.save(update_fields=["object_storage_key", "content_sha256", "size"]) + run.active_task_session = task_session + run.save(update_fields=["active_task_session"]) + self._open_sandbox_session(run) + + response = self.client.generic( + "POST", + f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/task_session_sync/", + cast(str, b'{"type":"session","version":3}\n'), + content_type="application/octet-stream", + HTTP_IF_MATCH='"stale-hash"', + HTTP_X_SANDBOX_ID="sandbox-1", + HTTP_X_TASK_RUN_TOKEN=create_sandbox_event_ingest_token(run), + ) + + self.assertEqual(response.status_code, status.HTTP_409_CONFLICT) + self.assertEqual(response.json()["error"], "The task session content is stale") + task_session.refresh_from_db() + self.assertEqual(task_session.object_storage_key, "task-sessions/current.jsonl") + self.assertEqual(task_session.content_sha256, "current-hash") + mock_write.assert_called_once() + mock_delete.assert_called_once() + + @patch("posthog.storage.object_storage.write") + def test_task_session_sync_rejects_a_closed_or_different_sandbox(self, mock_write): + task = self.create_task(runtime=Task.Runtime.PI) + run = self._create_run_with_sandbox(task) + task_session = TaskSession.create_for_task(task) + run.active_task_session = task_session + run.save(update_fields=["active_task_session"]) + self._open_sandbox_session(run, "active-sandbox") + + response = self.client.generic( + "POST", + f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/task_session_sync/", + cast(str, b'{"type":"session"}\n'), + content_type="application/octet-stream", + HTTP_IF_MATCH='"none"', + HTTP_X_SANDBOX_ID="stale-sandbox", + HTTP_X_TASK_RUN_TOKEN=create_sandbox_event_ingest_token(run), + ) + + self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN) + self.assertEqual(response.json()["detail"], "The task run token is invalid") + mock_write.assert_not_called() + + @patch("posthog.storage.object_storage.write") + def test_task_session_sync_rejects_empty_content_before_upload(self, mock_write): + task = self.create_task(runtime=Task.Runtime.PI) + run = self._create_run_with_sandbox(task) + task_session = TaskSession.create_for_task(task) + run.active_task_session = task_session + run.save(update_fields=["active_task_session"]) + self._open_sandbox_session(run) + + response = self.client.generic( + "POST", + f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/task_session_sync/", + cast(str, b""), + content_type="application/octet-stream", + HTTP_IF_MATCH='"none"', + HTTP_X_SANDBOX_ID="sandbox-1", + HTTP_X_TASK_RUN_TOKEN=create_sandbox_event_ingest_token(run), + ) + + self.assertEqual(response.status_code, status.HTTP_409_CONFLICT) + self.assertEqual(response.json()["error"], "The task session content size is invalid") + mock_write.assert_not_called() + + def test_task_session_sync_rejects_missing_task_run_token(self): + task = self.create_task(runtime=Task.Runtime.PI) + run = self._create_run_with_sandbox(task) + task_session = TaskSession.create_for_task(task) + run.active_task_session = task_session + run.save(update_fields=["active_task_session"]) + self._open_sandbox_session(run) + + response = self.client.generic( + "POST", + f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/task_session_sync/", + cast(str, b'{"type":"session"}\n'), + content_type="application/octet-stream", + HTTP_IF_MATCH='"none"', + HTTP_X_SANDBOX_ID="sandbox-1", + ) + + self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN) + self.assertEqual(response.json()["detail"], "The task run token is invalid") + + def test_task_session_sync_rejects_oversized_content_length(self): + task = self.create_task(runtime=Task.Runtime.PI) + run = self._create_run_with_sandbox(task) + + response = self.client.generic( + "POST", + f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/task_session_sync/", + cast(str, b"x"), + content_type="application/octet-stream", + CONTENT_LENGTH=str(tasks_facade.TASK_SESSION_MAX_SIZE_BYTES + 1), + ) + + self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) + self.assertEqual(response.json()["detail"], "The task session content size is invalid") + + @patch("posthog.storage.object_storage.get_presigned_url") + def test_task_session_is_readable_for_a_public_channel_task(self, mock_download_url): + other_user = self.create_organization_user("task-owner") + channel = Channel.objects.unscoped().create(team=self.team, name="shared", created_by=other_user) + task = self.create_task(created_by=other_user, runtime=Task.Runtime.PI) + task.channel = channel + task.save(update_fields=["channel"]) + run = self._create_run_with_sandbox(task) + task_session = TaskSession.create_for_task(task) + run.active_task_session = task_session + run.save(update_fields=["active_task_session"]) + mock_download_url.return_value = "https://storage.example/session.jsonl" + + response = self.client.get(f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/task_session/") + + self.assertEqual(response.status_code, status.HTTP_200_OK) + self.assertEqual(response.json()["id"], str(task_session.id)) + + @patch("products.tasks.backend.temporal.client.signal_task_followup_message") + def test_command_signals_pi_user_message(self, mock_signal_followup): + task = self.create_task(runtime=Task.Runtime.PI) + run = self._create_run_with_sandbox(task) + mock_signal_followup.return_value = True + + response = self.client.post( + self._command_url(task, run), + self._make_user_message(), + format="json", + ) + + self.assertEqual(response.status_code, status.HTTP_200_OK) + self.assertEqual(mock_signal_followup.call_args.args[3], "req-1") + + @override_settings(SANDBOX_JWT_PRIVATE_KEY=TEST_RSA_PRIVATE_KEY) + @patch("products.tasks.backend.presentation.views.api.http_requests.post") + def test_command_proxies_pi_rpc(self, mock_post): + reset_sandbox_jwt_key_cache() + rpc_response = { + "type": "response", + "command": "future_native_command", + "success": True, + "data": {"accepted": True}, + } + self._mock_agent_response(mock_post, {"jsonrpc": "2.0", "id": "native", "result": rpc_response}) + task = self.create_task(runtime=Task.Runtime.PI) + run = self._create_run_with_sandbox(task) + + response = self.client.post( + self._command_url(task, run), + { + "jsonrpc": "2.0", + "method": "pi/rpc", + "params": {"command": {"id": "native", "type": "future_native_command"}}, + "id": "native", + }, + format="json", + ) + + self.assertEqual(response.status_code, status.HTTP_200_OK) + self.assertEqual(response.json()["result"], rpc_response) + + @parameterized.expand([("queue_get",), ("queue_clear",)]) + @override_settings(SANDBOX_JWT_PRIVATE_KEY=TEST_RSA_PRIVATE_KEY) + @patch("products.tasks.backend.presentation.views.api.http_requests.post") + def test_command_proxies_pi_queue_operations(self, method, mock_post): + reset_sandbox_jwt_key_cache() + queue = {"steering": ["fix this"], "followUp": ["then summarize"]} + self._mock_agent_response(mock_post, {"jsonrpc": "2.0", "id": "queue", "result": queue}) + task = self.create_task(runtime=Task.Runtime.PI) + run = self._create_run_with_sandbox(task) + + response = self.client.post( + self._command_url(task, run), + {"jsonrpc": "2.0", "method": method, "id": "queue"}, + format="json", + ) + + self.assertEqual(response.status_code, status.HTTP_200_OK) + self.assertEqual(response.json()["result"], queue) + + @parameterized.expand([("queue_get",), ("queue_clear",)]) + @patch("products.tasks.backend.presentation.views.api.http_requests.post") + def test_command_rejects_pi_queue_operations_for_acp(self, method, mock_post): + task = self.create_task() + run = self._create_run_with_sandbox(task) + + response = self.client.post( + self._command_url(task, run), + {"jsonrpc": "2.0", "method": method, "id": "queue"}, + format="json", + ) + + self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) + self.assertEqual(response.json()["error"], "Pi commands require a Pi task.") + mock_post.assert_not_called() + @override_settings(SANDBOX_JWT_PRIVATE_KEY=TEST_RSA_PRIVATE_KEY) @patch("products.tasks.backend.presentation.views.api.http_requests.post") def test_command_proxies_close(self, mock_post): diff --git a/products/tasks/backend/tests/test_connection_token.py b/products/tasks/backend/tests/test_connection_token.py index b6aae52e60b6..2e811fca0430 100644 --- a/products/tasks/backend/tests/test_connection_token.py +++ b/products/tasks/backend/tests/test_connection_token.py @@ -1,5 +1,5 @@ import uuid -from datetime import timedelta +from datetime import UTC, datetime, timedelta from types import SimpleNamespace from typing import cast @@ -14,6 +14,7 @@ from products.tasks.backend.logic.services.connection_token import ( SANDBOX_CONNECTION_AUDIENCE, + SANDBOX_EVENT_INGEST_AUDIENCE, SANDBOX_JWT_STATE_KID_KEY, _compute_kid, _derive_public_key_pem, @@ -54,7 +55,7 @@ def _fake_run(state: dict | None = None) -> TaskRun: task_id=uuid.uuid4(), team_id=1, mode="background", - state=state if state is not None else {}, + state={"sandbox_id": "sandbox-1", **(state or {})}, ), ) @@ -111,6 +112,43 @@ def test_ingest_token_signed_with_run_stored_kid(self) -> None: self.assertEqual(jwt.get_unverified_header(token)["kid"], KID_A) payload = validate_sandbox_event_ingest_token(token) self.assertEqual(payload.team_id, 1) + self.assertEqual(payload.sandbox_id, "sandbox-1") + + @override_settings(SANDBOX_JWT_PRIVATE_KEY=KEY_A, SANDBOX_JWT_PRIVATE_KEY_SECONDARY=None) + def test_ingest_token_uses_explicit_sandbox_identity(self) -> None: + reset_sandbox_jwt_key_cache() + + token = create_sandbox_event_ingest_token( + _fake_run({"sandbox_id": None}), + sandbox_id="sandbox-explicit", + ) + + payload = validate_sandbox_event_ingest_token(token) + self.assertEqual(payload.sandbox_id, "sandbox-explicit") + + @override_settings(SANDBOX_JWT_PRIVATE_KEY=KEY_A, SANDBOX_JWT_PRIVATE_KEY_SECONDARY=None) + def test_legacy_ingest_token_without_sandbox_id_remains_valid(self) -> None: + reset_sandbox_jwt_key_cache() + run = _fake_run() + now = datetime.now(tz=UTC) + token = jwt.encode( + { + "run_id": str(run.id), + "task_id": str(run.task_id), + "team_id": run.team_id, + "iat": now, + "exp": now + timedelta(minutes=5), + "aud": SANDBOX_EVENT_INGEST_AUDIENCE, + }, + KEY_A, + algorithm="RS256", + headers={"kid": KID_A}, + ) + + payload = validate_sandbox_event_ingest_token(token) + + self.assertEqual(payload.run_id, str(run.id)) + self.assertIsNone(payload.sandbox_id) def test_ingest_token_validates_after_primary_rotation(self) -> None: # Ingest is rotation-safe: a token signed under the old primary keeps validating after the @@ -180,7 +218,15 @@ def test_duplicate_primary_and_secondary_collapse_to_one_key(self) -> None: def _fake_task_run() -> TaskRun: - return cast(TaskRun, SimpleNamespace(id=_RUN_ID, task_id=_TASK_ID, team_id=7, state={})) + return cast( + TaskRun, + SimpleNamespace( + id=_RUN_ID, + task_id=_TASK_ID, + team_id=7, + state={"sandbox_id": "sandbox-1"}, + ), + ) @override_settings(SANDBOX_JWT_PRIVATE_KEY=TEST_RSA_PRIVATE_KEY, SANDBOX_JWT_PUBLIC_KEY=None) diff --git a/products/tasks/backend/tests/test_event_ingest.py b/products/tasks/backend/tests/test_event_ingest.py index 5e18cf77916a..1fd072f1abd9 100644 --- a/products/tasks/backend/tests/test_event_ingest.py +++ b/products/tasks/backend/tests/test_event_ingest.py @@ -81,7 +81,10 @@ def _ingest_url( return f"/api/projects/{project_id}/tasks/{task.id}/runs/{run.id}/event_stream/" def _create_token(self, run: TaskRun | None = None) -> str: - return create_sandbox_event_ingest_token(run or self.task_run) + token_run = run or self.task_run + token_run.state = {**(token_run.state or {}), "sandbox_id": f"sandbox-{token_run.id}"} + token_run.save(update_fields=["state", "updated_at"]) + return create_sandbox_event_ingest_token(token_run) def _call_ingest( self, diff --git a/products/tasks/backend/tests/test_models.py b/products/tasks/backend/tests/test_models.py index 75fb54b58335..971c31d98b05 100644 --- a/products/tasks/backend/tests/test_models.py +++ b/products/tasks/backend/tests/test_models.py @@ -652,6 +652,28 @@ def test_create_run_does_not_inject_permission_mode_by_default(self): self.assertNotIn("initial_permission_mode", run.state) + @patch("products.tasks.backend.models.TaskRun.publish_stream_state_event") + def test_prepare_for_cloud_handoff_clears_stale_sandbox_routing(self, _publish): + run = TaskRun.objects.create( + task=self.task, + team=self.team, + status=TaskRun.Status.COMPLETED, + state={ + "sandbox_id": "old-sandbox", + "sandbox_url": "https://old-sandbox.test", + "sandbox_jwt_kid": "old-key", + "snapshot_external_id": "snapshot-1", + }, + ) + + run.prepare_for_cloud_handoff() + + self.assertNotIn("sandbox_id", run.state) + self.assertNotIn("sandbox_url", run.state) + self.assertNotIn("sandbox_jwt_kid", run.state) + self.assertEqual(run.state["snapshot_external_id"], "snapshot-1") + self.assertTrue(run.state["handoff_resumed"]) + def test_s3_prefixes_keep_existing_logs_and_artifact_paths(self): run = TaskRun.objects.create( task=self.task, @@ -697,6 +719,42 @@ def test_update_state_atomic_merges_against_latest_state(self): self.assertNotIn("pending_user_message", run.state) self.assertNotIn("pending_user_artifact_ids", run.state) + def test_clear_sandbox_connection_state_atomic_removes_matching_sandbox(self): + run = TaskRun.objects.create( + task=self.task, + team=self.team, + state={ + "sandbox_id": "sandbox-123", + "sandbox_url": "https://sandbox.example.com", + "sandbox_connect_token": "token", + "sandbox_jwt_kid": "key", + "mode": "interactive", + }, + ) + + TaskRun.clear_sandbox_connection_state_atomic(run.id, "sandbox-123") + + run.refresh_from_db() + self.assertEqual(run.state, {"mode": "interactive"}) + + def test_clear_sandbox_connection_state_atomic_preserves_newer_sandbox(self): + run = TaskRun.objects.create( + task=self.task, + team=self.team, + state={ + "sandbox_id": "new-sandbox", + "sandbox_url": "https://new-sandbox.example.com", + "sandbox_connect_token": "new-token", + "sandbox_jwt_kid": "new-key", + }, + ) + + TaskRun.clear_sandbox_connection_state_atomic(run.id, "old-sandbox") + + run.refresh_from_db() + self.assertEqual(run.state["sandbox_id"], "new-sandbox") + self.assertEqual(run.state["sandbox_connect_token"], "new-token") + def test_mutate_state_atomic_can_derive_values_under_lock(self): run = TaskRun.objects.create( task=self.task, diff --git a/products/tasks/backend/tests/test_run_log_mirror.py b/products/tasks/backend/tests/test_run_log_mirror.py index 7efdde6dc8e4..4ec910624838 100644 --- a/products/tasks/backend/tests/test_run_log_mirror.py +++ b/products/tasks/backend/tests/test_run_log_mirror.py @@ -288,7 +288,7 @@ def _create_run(self, origin_product: str, telemetry_enabled: bool | None = True ] ) @patch("products.tasks.backend.logic.services.run_log_mirror.logger") - @patch("products.tasks.backend.models.object_storage") + @patch("products.tasks.backend.storage.object_storage") def test_mirrors_only_allowlisted_origin_products(self, origin_product, expect_mirrored, mock_storage, mock_logger): mock_storage.read.return_value = None run = self._create_run(origin_product) @@ -307,7 +307,7 @@ def test_mirrors_only_allowlisted_origin_products(self, origin_product, expect_m mock_logger.info.assert_not_called() @patch("products.tasks.backend.logic.services.run_log_mirror.logger") - @patch("products.tasks.backend.models.object_storage") + @patch("products.tasks.backend.storage.object_storage") def test_no_mirroring_without_telemetry_flag_stamp(self, mock_storage, mock_logger): mock_storage.read.return_value = None run = self._create_run(Task.OriginProduct.SIGNALS_SCOUT, telemetry_enabled=None) @@ -319,7 +319,7 @@ def test_no_mirroring_without_telemetry_flag_stamp(self, mock_storage, mock_logg @override_settings(TASK_RUN_LOGS_MIRROR_ORIGIN_PRODUCTS=[]) @patch("products.tasks.backend.logic.services.run_log_mirror.logger") - @patch("products.tasks.backend.models.object_storage") + @patch("products.tasks.backend.storage.object_storage") def test_no_mirroring_when_disabled(self, mock_storage, mock_logger): mock_storage.read.return_value = None run = self._create_run(Task.OriginProduct.SIGNALS_SCOUT) @@ -333,7 +333,7 @@ def test_no_mirroring_when_disabled(self, mock_storage, mock_logger): "products.tasks.backend.logic.services.run_log_mirror.mirror_entries", side_effect=RuntimeError("kaboom"), ) - @patch("products.tasks.backend.models.object_storage") + @patch("products.tasks.backend.storage.object_storage") def test_mirror_failure_does_not_break_log_write(self, mock_storage, mock_mirror): mock_storage.read.return_value = None run = self._create_run(Task.OriginProduct.SIGNALS_SCOUT) diff --git a/products/tasks/backend/tests/test_storage.py b/products/tasks/backend/tests/test_storage.py new file mode 100644 index 000000000000..e8b093fa28b1 --- /dev/null +++ b/products/tasks/backend/tests/test_storage.py @@ -0,0 +1,32 @@ +from unittest.mock import MagicMock, patch + +from django.test import SimpleTestCase + +from parameterized import parameterized + +from products.tasks.backend.storage import append_jsonl_object + + +class TestAppendJsonlObject(SimpleTestCase): + @parameterized.expand( + [ + ("", True, '{"type": "session"}'), + ('{"type": "session"}', False, '{"type": "session"}\n{"type": "message"}'), + ] + ) + @patch("products.tasks.backend.storage.object_storage.write") + @patch("products.tasks.backend.storage.object_storage.read") + def test_appends_complete_json_lines( + self, + existing_content: str, + expected_is_new: bool, + expected_content: str, + mock_read: MagicMock, + mock_write: MagicMock, + ) -> None: + mock_read.return_value = existing_content + + is_new = append_jsonl_object("sessions/example.jsonl", [{"type": "session" if expected_is_new else "message"}]) + + self.assertEqual(is_new, expected_is_new) + mock_write.assert_called_once_with("sessions/example.jsonl", expected_content) diff --git a/products/tasks/frontend/generated/api.schemas.ts b/products/tasks/frontend/generated/api.schemas.ts index f818c6f1dc2d..2ef2e7534103 100644 --- a/products/tasks/frontend/generated/api.schemas.ts +++ b/products/tasks/frontend/generated/api.schemas.ts @@ -2838,6 +2838,9 @@ export const JsonrpcEnumApi = { * * `permission_response` - permission_response * * `set_config_option` - set_config_option * * `mcp_response` - mcp_response + * * `pi/rpc` - pi/rpc + * * `queue_get` - queue_get + * * `queue_clear` - queue_clear */ export type MethodEnumApi = (typeof MethodEnumApi)[keyof typeof MethodEnumApi] @@ -2848,6 +2851,9 @@ export const MethodEnumApi = { PermissionResponse: 'permission_response', SetConfigOption: 'set_config_option', McpResponse: 'mcp_response', + PiRpc: 'pi/rpc', + QueueGet: 'queue_get', + QueueClear: 'queue_clear', } as const /** @@ -2865,7 +2871,10 @@ export interface TaskRunCommandRequestApi { * * `close` - close * * `permission_response` - permission_response * * `set_config_option` - set_config_option - * * `mcp_response` - mcp_response */ + * * `mcp_response` - mcp_response + * * `pi/rpc` - pi/rpc + * * `queue_get` - queue_get + * * `queue_clear` - queue_clear */ method: MethodEnumApi /** Parameters for the command */ params?: TaskRunCommandRequestApiParams @@ -2873,11 +2882,6 @@ export interface TaskRunCommandRequestApi { id?: unknown } -/** - * Command result on success - */ -export type TaskRunCommandResponseApiResult = { [key: string]: unknown } - /** * Error details on failure */ @@ -2892,7 +2896,7 @@ export interface TaskRunCommandResponseApi { /** Request ID echoed back (string or number) */ id?: unknown /** Command result on success */ - result?: TaskRunCommandResponseApiResult + result?: unknown /** Error details on failure */ error?: TaskRunCommandResponseApiError } @@ -2959,6 +2963,28 @@ export interface StreamReadTokenResponseApi { stream_base_url: string | null } +export interface TaskSessionResponseApi { + /** Task session identifier */ + id: string + /** + * Temporary URL for downloading the session + * @nullable + */ + download_url: string | null + /** + * SHA-256 digest of the current session content + * @nullable + */ + content_sha256: string | null +} + +export interface TaskSessionSyncResponseApi { + /** Task session identifier */ + id: string + /** SHA-256 digest of the uploaded session content */ + content_sha256: string +} + /** * * `slack_message` - slack_message * * `slack_canvas` - slack_canvas diff --git a/products/tasks/frontend/generated/api.ts b/products/tasks/frontend/generated/api.ts index e54734405cf8..73d2cc0ee28d 100644 --- a/products/tasks/frontend/generated/api.ts +++ b/products/tasks/frontend/generated/api.ts @@ -93,6 +93,8 @@ import type { TaskRunRelayMessageRequestApi, TaskRunRelayMessageResponseApi, TaskRunStartRequestApi, + TaskSessionResponseApi, + TaskSessionSyncResponseApi, TaskStagedArtifactsFinalizeUploadRequestApi, TaskStagedArtifactsFinalizeUploadResponseApi, TaskStagedArtifactsPrepareUploadRequestApi, @@ -1490,7 +1492,7 @@ export const getTasksRunsCommandCreateUrl = (projectId: string, taskId: string, } /** - * Queue user_message JSON-RPC commands through the task workflow and forward sandbox control commands to the agent server. Supports user_message, cancel, close, permission_response, set_config_option, and mcp_response commands. + * Queue user_message JSON-RPC commands through the task workflow and forward sandbox control commands to the agent server. Supports user_message, cancel, close, permission_response, set_config_option, mcp_response, native Pi RPC commands, and Pi queue operations. * @summary Send command to task run */ export const tasksRunsCommandCreate = async ( @@ -1718,6 +1720,49 @@ export const tasksRunsStreamTokenRetrieve = async ( }) } +export const getTasksRunsTaskSessionRetrieveUrl = (projectId: string, taskId: string, id: string) => { + return `/api/projects/${projectId}/tasks/${taskId}/runs/${id}/task_session/` +} + +/** + * API for managing task runs. Each run represents an execution of a task. + * @summary Get active task session storage access + */ +export const tasksRunsTaskSessionRetrieve = async ( + projectId: string, + taskId: string, + id: string, + options?: RequestInit +): Promise => { + return apiMutator(getTasksRunsTaskSessionRetrieveUrl(projectId, taskId, id), { + ...options, + method: 'GET', + }) +} + +export const getTasksRunsTaskSessionSyncCreateUrl = (projectId: string, taskId: string, id: string) => { + return `/api/projects/${projectId}/tasks/${taskId}/runs/${id}/task_session_sync/` +} + +/** + * API for managing task runs. Each run represents an execution of a task. + * @summary Replace the active native task session + */ +export const tasksRunsTaskSessionSyncCreate = async ( + projectId: string, + taskId: string, + id: string, + tasksRunsTaskSessionSyncCreateBody?: Blob, + options?: RequestInit +): Promise => { + return apiMutator(getTasksRunsTaskSessionSyncCreateUrl(projectId, taskId, id), { + ...options, + method: 'POST', + headers: { 'Content-Type': 'application/octet-stream', ...options?.headers }, + body: tasksRunsTaskSessionSyncCreateBody, + }) +} + export const getTasksRunsLivingArtifactsListUrl = (projectId: string, taskId: string, runId: string) => { return `/api/projects/${projectId}/tasks/${taskId}/runs/${runId}/living_artifacts/` } diff --git a/products/tasks/frontend/generated/api.zod.ts b/products/tasks/frontend/generated/api.zod.ts index eec19c31c0e6..21d2d5d863cc 100644 --- a/products/tasks/frontend/generated/api.zod.ts +++ b/products/tasks/frontend/generated/api.zod.ts @@ -2604,7 +2604,7 @@ export const TasksRunsCancelCreateBody = /* @__PURE__ */ zod.object({ }) /** - * Queue user_message JSON-RPC commands through the task workflow and forward sandbox control commands to the agent server. Supports user_message, cancel, close, permission_response, set_config_option, and mcp_response commands. + * Queue user_message JSON-RPC commands through the task workflow and forward sandbox control commands to the agent server. Supports user_message, cancel, close, permission_response, set_config_option, mcp_response, native Pi RPC commands, and Pi queue operations. * @summary Send command to task run */ export const TasksRunsCommandCreateBody = /* @__PURE__ */ zod @@ -2614,12 +2614,22 @@ export const TasksRunsCommandCreateBody = /* @__PURE__ */ zod .describe('\* `2.0` - 2.0') .describe("JSON-RPC version, must be '2.0'\n\n\* `2.0` - 2.0"), method: zod - .enum(['user_message', 'cancel', 'close', 'permission_response', 'set_config_option', 'mcp_response']) + .enum([ + 'user_message', + 'cancel', + 'close', + 'permission_response', + 'set_config_option', + 'mcp_response', + 'pi/rpc', + 'queue_get', + 'queue_clear', + ]) .describe( - '\* `user_message` - user_message\n\* `cancel` - cancel\n\* `close` - close\n\* `permission_response` - permission_response\n\* `set_config_option` - set_config_option\n\* `mcp_response` - mcp_response' + '\* `user_message` - user_message\n\* `cancel` - cancel\n\* `close` - close\n\* `permission_response` - permission_response\n\* `set_config_option` - set_config_option\n\* `mcp_response` - mcp_response\n\* `pi\/rpc` - pi\/rpc\n\* `queue_get` - queue_get\n\* `queue_clear` - queue_clear' ) .describe( - 'Command method to execute on the agent server\n\n\* `user_message` - user_message\n\* `cancel` - cancel\n\* `close` - close\n\* `permission_response` - permission_response\n\* `set_config_option` - set_config_option\n\* `mcp_response` - mcp_response' + 'Command method to execute on the agent server\n\n\* `user_message` - user_message\n\* `cancel` - cancel\n\* `close` - close\n\* `permission_response` - permission_response\n\* `set_config_option` - set_config_option\n\* `mcp_response` - mcp_response\n\* `pi\/rpc` - pi\/rpc\n\* `queue_get` - queue_get\n\* `queue_clear` - queue_clear' ), params: zod.record(zod.string(), zod.unknown()).optional().describe('Parameters for the command'), id: zod.unknown().optional().describe('Optional JSON-RPC request ID (string or number)'), diff --git a/products/tasks/mcp/tools.yaml b/products/tasks/mcp/tools.yaml index b548c9ee6351..d9d67d7c76cd 100644 --- a/products/tasks/mcp/tools.yaml +++ b/products/tasks/mcp/tools.yaml @@ -480,6 +480,12 @@ tools: tasks-runs-stream-token-retrieve: operation: tasks_runs_stream_token_retrieve enabled: false + tasks-runs-task-session-retrieve: + operation: tasks_runs_task_session_retrieve + enabled: false + tasks-runs-task-session-sync-create: + operation: tasks_runs_task_session_sync_create + enabled: false tasks-slack-thread-context-retrieve: operation: tasks_slack_thread_context_retrieve enabled: false diff --git a/services/mcp/src/api/generated.ts b/services/mcp/src/api/generated.ts index 3dcd99759174..a8ae06ae6bc6 100644 --- a/services/mcp/src/api/generated.ts +++ b/services/mcp/src/api/generated.ts @@ -39477,6 +39477,9 @@ export namespace Schemas { * * `permission_response` - permission_response * * `set_config_option` - set_config_option * * `mcp_response` - mcp_response + * * `pi/rpc` - pi/rpc + * * `queue_get` - queue_get + * * `queue_clear` - queue_clear */ export type MethodEnum = typeof MethodEnum[keyof typeof MethodEnum]; @@ -39488,6 +39491,9 @@ export namespace Schemas { PermissionResponse: 'permission_response', SetConfigOption: 'set_config_option', McpResponse: 'mcp_response', + PiRpc: 'pi/rpc', + QueueGet: 'queue_get', + QueueClear: 'queue_clear', } as const; /** @@ -67900,7 +67906,10 @@ export namespace Schemas { * * `close` - close * * `permission_response` - permission_response * * `set_config_option` - set_config_option - * * `mcp_response` - mcp_response */ + * * `mcp_response` - mcp_response + * * `pi/rpc` - pi/rpc + * * `queue_get` - queue_get + * * `queue_clear` - queue_clear */ method: MethodEnum; /** Parameters for the command */ params?: TaskRunCommandRequestParams; @@ -67908,11 +67917,6 @@ export namespace Schemas { id?: unknown; } - /** - * Command result on success - */ - export type TaskRunCommandResponseResult = { [key: string]: unknown }; - /** * Error details on failure */ @@ -67927,7 +67931,7 @@ export namespace Schemas { /** Request ID echoed back (string or number) */ id?: unknown; /** Command result on success */ - result?: TaskRunCommandResponseResult; + result?: unknown; /** Error details on failure */ error?: TaskRunCommandResponseError; } @@ -68232,6 +68236,28 @@ export namespace Schemas { pending_user_artifact_ids?: string[]; } + export interface TaskSessionResponse { + /** Task session identifier */ + id: string; + /** + * Temporary URL for downloading the session + * @nullable + */ + download_url: string | null; + /** + * SHA-256 digest of the current session content + * @nullable + */ + content_sha256: string | null; + } + + export interface TaskSessionSyncResponse { + /** Task session identifier */ + id: string; + /** SHA-256 digest of the uploaded session content */ + content_sha256: string; + } + export interface TaskStagedArtifactFinalizeUpload { /** Stable identifier returned by the staged prepare upload endpoint */ id: string;