From ccbcf72c2d23253edd9fd3db5a37caefaa27853f Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Wed, 22 Jul 2026 17:21:06 +0200 Subject: [PATCH 01/22] feat(tasks): add durable Pi task sessions --- products/tasks/backend/constants.py | 2 + products/tasks/backend/facade/api.py | 262 ++++++++++++- .../backend/logic/services/docker_sandbox.py | 94 ++--- .../backend/logic/services/modal_sandbox.py | 5 + .../tasks/backend/logic/services/sandbox.py | 5 + .../logic/services/test_docker_sandbox.py | 72 ++++ .../services/tests/test_modal_sandbox.py | 4 + .../backend/migrations/0073_task_session.py | 78 ++++ .../migrations/0074_task_session_revisions.py | 23 ++ .../backend/migrations/max_migration.txt | 2 +- products/tasks/backend/models.py | 115 +++++- .../tasks/backend/presentation/serializers.py | 40 +- .../tasks/backend/presentation/views/api.py | 115 +++++- .../sandbox/images/Dockerfile.sandbox-local | 30 +- products/tasks/backend/storage.py | 15 + .../activities/get_task_processing_context.py | 6 + .../activities/relay_sandbox_events.py | 34 +- .../activities/start_agent_server.py | 1 + .../tests/test_get_task_processing_context.py | 8 + .../tests/test_relay_sandbox_events.py | 8 + products/tasks/backend/tests/test_api.py | 358 +++++++++++++++++- products/tasks/backend/tests/test_models.py | 22 ++ products/tasks/backend/tests/test_storage.py | 68 ++++ .../tasks/frontend/generated/api.schemas.ts | 46 ++- products/tasks/frontend/generated/api.ts | 74 ++++ products/tasks/frontend/generated/api.zod.ts | 37 +- services/mcp/src/api/generated.ts | 46 ++- 27 files changed, 1422 insertions(+), 148 deletions(-) create mode 100644 products/tasks/backend/migrations/0073_task_session.py create mode 100644 products/tasks/backend/migrations/0074_task_session_revisions.py create mode 100644 products/tasks/backend/storage.py create mode 100644 products/tasks/backend/tests/test_storage.py diff --git a/products/tasks/backend/constants.py b/products/tasks/backend/constants.py index bcd9d0737352..c0a1a3c379da 100644 --- a/products/tasks/backend/constants.py +++ b/products/tasks/backend/constants.py @@ -66,6 +66,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 cb5a8cd7b40f..b4b3e4af6bcc 100644 --- a/products/tasks/backend/facade/api.py +++ b/products/tasks/backend/facade/api.py @@ -16,6 +16,7 @@ """ import re +import json import logging from collections.abc import Iterable, Sequence from concurrent.futures import ThreadPoolExecutor @@ -41,6 +42,8 @@ MAX_CUSTOM_IMAGES_PER_TEAM, MAX_CUSTOM_IMAGES_PER_USER, RESERVED_SANDBOX_ENVIRONMENT_VARIABLE_KEYS, + TASK_SESSION_MAX_SIZE_BYTES, + TASK_SESSION_UPLOAD_FORM_OVERHEAD_BYTES, is_blocked_sandbox_env_key, ) from products.tasks.backend.error_telemetry import truncate_error_message @@ -61,6 +64,7 @@ Task, TaskAutomation, TaskRun, + TaskSession, TaskThreadMessage, TaskThreadMessageMention, ) @@ -104,6 +108,7 @@ "TaskRunEnvironment", "TaskRunStatus", "append_task_run_log", + "ensure_task_run_session", "beacon_task_presence", "bootstrap_task_run", "can_mint_readonly_github_token", @@ -149,6 +154,8 @@ "get_task_detail", "get_task_id_for_run", "get_task_run", + "get_task_run_session", + "prepare_task_run_session_sync", "get_task_run_detail", "get_task_run_sandbox_connection", "get_task_run_living_artifact", @@ -204,6 +211,7 @@ "update_task", "update_task_automation", "update_task_run", + "finalize_task_run_session_sync", "update_task_run_state", "upsert_internal_sandbox_env", "validate_set_output", @@ -2203,6 +2211,248 @@ 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, int] | 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.get(id=run.active_task_session_id) + 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.revision + + +def prepare_task_run_session_sync( + run_id: str | UUID, + task_id: str | UUID, + team_id: int, + *, + sandbox_id: str, + expected_revision: int, +) -> tuple[UUID, UUID, dict] | 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 + + with transaction.atomic(): + run = TaskRun.objects.select_for_update(of=("self",)).get(id=visible_run.id) + if (run.state or {}).get("sandbox_id") != sandbox_id: + raise ValueError("The task session writer is not the active sandbox") + task_session = TaskSession.objects.select_for_update().get(id=run.active_task_session_id) + if task_session.revision != expected_revision: + raise ValueError("The task session revision is stale") + + stale_pending_key = task_session.pending_object_storage_key + sync_id = uuid4() + object_storage_key = ( + f"task-sessions/{task_session.organization_id}/{task_session.task_id}/{task_session.id}/" + f"uploads/{task_session.revision + 1}-{sync_id}.jsonl" + ) + task_session.pending_sync_id = sync_id + task_session.pending_object_storage_key = object_storage_key + task_session.save(update_fields=["pending_sync_id", "pending_object_storage_key", "updated_at"]) + + if stale_pending_key: + + def delete_stale_pending_object() -> None: + try: + object_storage.delete(stale_pending_key) + except Exception as error: + logger.warning( + "task_session.failed_to_delete_stale_pending_object", + extra={ + "task_session_id": str(task_session.id), + "object_storage_key": stale_pending_key, + "error": str(error), + }, + ) + + transaction.on_commit(delete_stale_pending_object) + + upload = object_storage.get_presigned_post( + object_storage_key, + conditions=[["content-length-range", 1, TASK_SESSION_MAX_SIZE_BYTES + TASK_SESSION_UPLOAD_FORM_OVERHEAD_BYTES]], + ) + if not upload: + raise RuntimeError("Unable to prepare task session upload") + return task_session.id, sync_id, upload + + +def _delete_task_session_objects(task_session_id: UUID, object_storage_keys: Sequence[str]) -> None: + from posthog.storage import object_storage # noqa: PLC0415 + + for object_storage_key in object_storage_keys: + 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 _reject_task_session_sync( + task_session_id: UUID, + sync_id: UUID, + pending_key: str, + promoted_key: str, +) -> None: + with transaction.atomic(): + task_session = TaskSession.objects.select_for_update().get(id=task_session_id) + if task_session.pending_sync_id != sync_id or task_session.pending_object_storage_key != pending_key: + return + + task_session.pending_sync_id = None + task_session.pending_object_storage_key = None + task_session.save(update_fields=["pending_sync_id", "pending_object_storage_key", "updated_at"]) + transaction.on_commit(lambda: _delete_task_session_objects(task_session_id, (pending_key, promoted_key))) + + +class _TaskSessionUploadRejected(ValueError): + pass + + +def finalize_task_run_session_sync( + run_id: str | UUID, + task_id: str | UUID, + team_id: int, + *, + sandbox_id: str, + sync_id: UUID, + expected_revision: int, +) -> tuple[UUID, int] | 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 + task_session_id = visible_run.active_task_session_id + pending_key: str | None = None + promoted_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.state or {}).get("sandbox_id") != sandbox_id: + raise ValueError("The task session writer is not the active sandbox") + if locked_run.active_task_session_id != task_session_id: + raise ValueError("The task session sync is stale") + + locked_session = TaskSession.objects.select_for_update().get(id=task_session_id) + promoted_object_storage_key = ( + f"task-sessions/{locked_session.organization_id}/{locked_session.task_id}/{locked_session.id}/" + f"revisions/{expected_revision + 1}-{sync_id}.jsonl" + ) + promoted_key = promoted_object_storage_key + if ( + locked_session.revision == expected_revision + 1 + and locked_session.object_storage_key == promoted_object_storage_key + ): + return locked_session.id, locked_session.revision + pending_object_storage_key = locked_session.pending_object_storage_key + if ( + locked_session.revision != expected_revision + or locked_session.pending_sync_id != sync_id + or not pending_object_storage_key + ): + raise ValueError("The task session sync is stale") + + pending_key = pending_object_storage_key + pending_metadata = object_storage.head_object(pending_object_storage_key) + if pending_metadata is None: + raise ValueError("The uploaded task session is missing") + pending_size = pending_metadata.get("ContentLength") + if not isinstance(pending_size, int): + raise ValueError("Unable to determine the uploaded task session size") + if pending_size < 1: + raise _TaskSessionUploadRejected("The uploaded task session is not valid JSONL") + if pending_size > TASK_SESSION_MAX_SIZE_BYTES: + raise _TaskSessionUploadRejected("The uploaded task session exceeds the maximum size") + + object_storage.copy(pending_object_storage_key, promoted_object_storage_key) + promoted_metadata = object_storage.head_object(promoted_object_storage_key) + if promoted_metadata is None: + raise ValueError("The uploaded task session is missing") + promoted_size = promoted_metadata.get("ContentLength") + if not isinstance(promoted_size, int): + raise ValueError("Unable to determine the uploaded task session size") + if promoted_size < 1: + raise _TaskSessionUploadRejected("The uploaded task session is not valid JSONL") + if promoted_size > TASK_SESSION_MAX_SIZE_BYTES: + raise _TaskSessionUploadRejected("The uploaded task session exceeds the maximum size") + + try: + content = object_storage.read(promoted_object_storage_key, missing_ok=True) + except UnicodeDecodeError as error: + raise _TaskSessionUploadRejected("The uploaded task session is not valid JSONL") from error + if not content: + raise ValueError("The uploaded task session is missing") + if len(content.encode("utf-8")) > TASK_SESSION_MAX_SIZE_BYTES: + raise _TaskSessionUploadRejected("The uploaded task session exceeds the maximum size") + try: + entries = [json.loads(line) for line in content.splitlines() if line.strip()] + except json.JSONDecodeError as error: + raise _TaskSessionUploadRejected("The uploaded task session is not valid JSONL") from error + if not entries or any(not isinstance(entry, dict) for entry in entries): + raise _TaskSessionUploadRejected("The uploaded task session is not valid JSONL") + if entries[0].get("type") != "session": + raise _TaskSessionUploadRejected("The uploaded task session has no session header") + + previous_key = locked_session.object_storage_key + locked_session.object_storage_key = promoted_object_storage_key + locked_session.revision += 1 + locked_session.pending_sync_id = None + locked_session.pending_object_storage_key = None + locked_session.save( + update_fields=[ + "object_storage_key", + "revision", + "pending_sync_id", + "pending_object_storage_key", + "updated_at", + ] + ) + transaction.on_commit( + lambda: _delete_task_session_objects(locked_session.id, (previous_key, pending_object_storage_key)) + ) + locked_session.mark_synced() + return locked_session.id, locked_session.revision + except _TaskSessionUploadRejected: + if pending_key is None or promoted_key is None: + raise + _reject_task_session_sync(task_session_id, sync_id, pending_key, promoted_key) + raise + except Exception: + if promoted_key is not None: + finalized_session = TaskSession.objects.get(id=task_session_id) + if ( + finalized_session.revision == expected_revision + 1 + and finalized_session.object_storage_key == promoted_key + ): + return finalized_session.id, finalized_session.revision + 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 @@ -3248,8 +3498,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: @@ -3329,9 +3577,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={ @@ -4377,13 +4622,6 @@ 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") diff --git a/products/tasks/backend/logic/services/docker_sandbox.py b/products/tasks/backend/logic/services/docker_sandbox.py index 3fe76eac06e8..c2fb0140b65c 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_protocol: str | None = None, runtime_adapter: str | None = None, provider: str | None = None, model: str | None = None, @@ -811,6 +796,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_protocol=agent_protocol, + sandbox_id=self.id, runtime_adapter=runtime_adapter, provider=provider, model=model, @@ -895,6 +882,7 @@ def start_agent_server( auto_publish: bool = False, interaction_origin: str | None = None, branch: str | None = None, + agent_protocol: str | None = None, runtime_adapter: str | None = None, provider: str | None = None, model: str | None = None, @@ -966,6 +954,7 @@ def start_agent_server( auto_publish, interaction_origin, branch, + agent_protocol, runtime_adapter, provider, model, @@ -1017,6 +1006,7 @@ def start_agent_server( auto_publish, interaction_origin, branch=None, + agent_protocol=agent_protocol, runtime_adapter=runtime_adapter, provider=provider, model=model, diff --git a/products/tasks/backend/logic/services/modal_sandbox.py b/products/tasks/backend/logic/services/modal_sandbox.py index 9ef32c2b8674..ffb787a18698 100644 --- a/products/tasks/backend/logic/services/modal_sandbox.py +++ b/products/tasks/backend/logic/services/modal_sandbox.py @@ -945,6 +945,7 @@ def _build_agent_server_command( auto_publish: bool = False, interaction_origin: str | None = None, branch: str | None = None, + agent_protocol: str | None = None, runtime_adapter: str | None = None, provider: str | None = None, model: str | None = None, @@ -962,6 +963,8 @@ def _build_agent_server_command( ) -> str: env_prefix = build_agent_runtime_env_prefix( interaction_origin=interaction_origin, + agent_protocol=agent_protocol, + sandbox_id=self.id, runtime_adapter=runtime_adapter, provider=provider, model=model, @@ -1072,6 +1075,7 @@ def start_agent_server( auto_publish: bool = False, interaction_origin: str | None = None, branch: str | None = None, + agent_protocol: str | None = None, runtime_adapter: str | None = None, provider: str | None = None, model: str | None = None, @@ -1146,6 +1150,7 @@ def start_agent_server( auto_publish, interaction_origin, branch, + agent_protocol, runtime_adapter, provider, model, diff --git a/products/tasks/backend/logic/services/sandbox.py b/products/tasks/backend/logic/services/sandbox.py index 7503f439e97f..d68315ce8a1f 100644 --- a/products/tasks/backend/logic/services/sandbox.py +++ b/products/tasks/backend/logic/services/sandbox.py @@ -167,6 +167,8 @@ def redact_sandbox_command(command: str) -> str: def build_agent_runtime_env_prefix( *, interaction_origin: str | None = None, + agent_protocol: str | None = None, + sandbox_id: str | None = None, runtime_adapter: str | None = None, provider: str | None = None, model: str | None = None, @@ -179,6 +181,8 @@ def build_agent_runtime_env_prefix( ) -> str: env_vars = { "POSTHOG_CODE_INTERACTION_ORIGIN": interaction_origin, + "POSTHOG_AGENT_PROTOCOL": agent_protocol, + "POSTHOG_SANDBOX_ID": sandbox_id, "POSTHOG_CODE_RUNTIME_ADAPTER": runtime_adapter, "POSTHOG_CODE_PROVIDER": provider, "POSTHOG_CODE_MODEL": model, @@ -329,6 +333,7 @@ def start_agent_server( auto_publish: bool = False, interaction_origin: str | None = None, branch: str | None = None, + agent_protocol: str | None = None, runtime_adapter: str | None = None, provider: str | None = None, model: str | None = None, diff --git a/products/tasks/backend/logic/services/test_docker_sandbox.py b/products/tasks/backend/logic/services/test_docker_sandbox.py index a9c9eae4be08..f114a21a6f1a 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,40 @@ 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, "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_protocol="pi", + ) + + assert mock_build.call_count == 2 + assert mock_build.call_args_list[1].kwargs["agent_protocol"] == "pi" + def test_parse_repo_mount_map_empty(self): with patch.dict(os.environ, {}, clear=True): assert parse_sandbox_repo_mount_map() == {} @@ -586,6 +655,7 @@ def test_start_agent_server_includes_runtime_environment_variables(self): "task-123", "run-456", "background", + agent_protocol="pi", runtime_adapter="codex", provider="openai", model="gpt-5.3-codex", @@ -596,6 +666,8 @@ def test_start_agent_server_includes_runtime_environment_variables(self): ) command = _agent_server_launch_command(mock_execute) + assert "POSTHOG_AGENT_PROTOCOL=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 898ae7340135..2942afd49492 100644 --- a/products/tasks/backend/logic/services/tests/test_modal_sandbox.py +++ b/products/tasks/backend/logic/services/tests/test_modal_sandbox.py @@ -443,6 +443,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 +547,7 @@ def test_start_agent_server_includes_runtime_environment_variables(self, mock_sa task_id="task-123", run_id="run-456", mode="background", + agent_protocol="pi", runtime_adapter="codex", provider="openai", model="gpt-5.3-codex", @@ -555,6 +558,7 @@ def test_start_agent_server_includes_runtime_environment_variables(self, mock_sa ) command = _agent_server_launch_command(mock_sandbox.execute) + assert "POSTHOG_AGENT_PROTOCOL=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/migrations/0073_task_session.py b/products/tasks/backend/migrations/0073_task_session.py new file mode 100644 index 000000000000..81b376cf4cad --- /dev/null +++ b/products/tasks/backend/migrations/0073_task_session.py @@ -0,0 +1,78 @@ +# Generated by Django 5.2.14 on 2026-07-17 13:24 + +import uuid + +import django.utils.timezone +import django.db.models.deletion +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [ + ("posthog", "1256_userproductlist_default_reason"), + ("tasks", "0072_loop_skill_bundles"), + ] + + operations = [ + migrations.CreateModel( + name="TaskSession", + fields=[ + ( + "id", + models.UUIDField( + default=uuid.uuid4, + editable=False, + primary_key=True, + serialize=False, + ), + ), + ("object_storage_key", models.CharField(max_length=512, unique=True)), + ("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", + ), + ), + ( + "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=["task", "-updated_at"], name="task_session_task_updated_idx"), + ), + ] diff --git a/products/tasks/backend/migrations/0074_task_session_revisions.py b/products/tasks/backend/migrations/0074_task_session_revisions.py new file mode 100644 index 000000000000..235932867b6d --- /dev/null +++ b/products/tasks/backend/migrations/0074_task_session_revisions.py @@ -0,0 +1,23 @@ +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [("tasks", "0073_task_session")] + + operations = [ + migrations.AddField( + model_name="tasksession", + name="revision", + field=models.PositiveBigIntegerField(default=0), + ), + migrations.AddField( + model_name="tasksession", + name="pending_sync_id", + field=models.UUIDField(blank=True, null=True), + ), + migrations.AddField( + model_name="tasksession", + name="pending_object_storage_key", + field=models.CharField(blank=True, max_length=512, null=True, unique=True), + ), + ] diff --git a/products/tasks/backend/migrations/max_migration.txt b/products/tasks/backend/migrations/max_migration.txt index e12dffe2bd4e..26c8a0ed1326 100644 --- a/products/tasks/backend/migrations/max_migration.txt +++ b/products/tasks/backend/migrations/max_migration.txt @@ -1 +1 @@ -0072_loop_skill_bundles +0074_task_session_revisions diff --git a/products/tasks/backend/models.py b/products/tasks/backend/models.py index 6a64428598b4..06cce3952f1f 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 @@ -47,6 +46,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__) @@ -865,6 +865,79 @@ def _dispatch() -> None: return task +class TaskSession(models.Model): + id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False) + organization = models.ForeignKey( + "posthog.Organization", + 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, unique=True) + revision = models.PositiveBigIntegerField(default=0) + pending_sync_id = models.UUIDField(null=True, blank=True) + pending_object_storage_key = models.CharField(max_length=512, null=True, blank=True, unique=True) + 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=["task", "-updated_at"], name="task_session_task_updated_idx"), + ] + + @classmethod + def create_for_task(cls, task: Task) -> "TaskSession": + session_id = uuid.uuid4() + object_storage_key = f"task-sessions/{task.team.organization_id}/{task.id}/{session_id}.jsonl" + return cls.objects.create( + id=session_id, + organization_id=task.team.organization_id, + task=task, + object_storage_key=object_storage_key, + ) + + def read_jsonl(self) -> str: + return object_storage.read(self.object_storage_key, missing_ok=True) or "" + + def write_jsonl(self, content: str | bytes) -> None: + object_storage.write(self.object_storage_key, content) + self._tag_object() + + def mark_synced(self) -> None: + self._tag_object() + self.save(update_fields=["updated_at"]) + + def append_entries(self, entries: list[dict[str, Any]]) -> None: + if not entries: + return + + is_new_object = append_jsonl_object(self.object_storage_key, entries) + if is_new_object: + self._tag_object() + + def _tag_object(self) -> None: + try: + object_storage.tag( + self.object_storage_key, + { + "data_class": "task_session", + "organization_id": str(self.organization_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 @@ -1342,6 +1415,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") @@ -1485,6 +1565,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( @@ -1683,13 +1766,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) @@ -2582,6 +2659,26 @@ 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: + object_storage_keys = {key for key in (instance.object_storage_key, instance.pending_object_storage_key) if key} + task_session_id = str(instance.id) + + def delete_objects() -> None: + for object_storage_key in object_storage_keys: + 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_objects) + + @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 8b267f3edcfa..3f69d4368e76 100644 --- a/products/tasks/backend/presentation/serializers.py +++ b/products/tasks/backend/presentation/serializers.py @@ -727,6 +727,34 @@ def validate_entries(self, value): return value +class TaskSessionResponseSerializer(serializers.Serializer): + id = serializers.UUIDField() + download_url = serializers.URLField() + revision = serializers.IntegerField(min_value=0) + + +class TaskSessionSyncPrepareSerializer(serializers.Serializer): + sandbox_id = serializers.CharField() + expected_revision = serializers.IntegerField(min_value=0) + + +class TaskSessionSyncPrepareResponseSerializer(serializers.Serializer): + id = serializers.UUIDField() + sync_id = serializers.UUIDField() + upload = serializers.DictField() + + +class TaskSessionSyncSerializer(serializers.Serializer): + sandbox_id = serializers.CharField() + sync_id = serializers.UUIDField() + expected_revision = serializers.IntegerField(min_value=0) + + +class TaskSessionSyncResponseSerializer(serializers.Serializer): + id = serializers.UUIDField() + revision = serializers.IntegerField(min_value=1) + + 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") @@ -2280,6 +2308,7 @@ class TaskRunCommandRequestSerializer(serializers.Serializer): "permission_response", "set_config_option", "mcp_response", + "pi/rpc", ] # Cap on the serialized mcp_response params (docs/cloud-mcp-relay.md): the relayed JSON-RPC @@ -2307,7 +2336,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 @@ -2352,6 +2381,13 @@ 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"}) elif method == "permission_response": self._require_nonempty_string(params, "requestId") self._require_nonempty_string(params, "optionId") @@ -2386,7 +2422,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..d603b66aa26e 100644 --- a/products/tasks/backend/presentation/views/api.py +++ b/products/tasks/backend/presentation/views/api.py @@ -107,6 +107,11 @@ TaskRunStartRequestSerializer, TaskRunUpdateSerializer, TaskSerializer, + TaskSessionResponseSerializer, + TaskSessionSyncPrepareResponseSerializer, + TaskSessionSyncPrepareSerializer, + TaskSessionSyncResponseSerializer, + TaskSessionSyncSerializer, TaskStagedArtifactsFinalizeUploadRequestSerializer, TaskStagedArtifactsFinalizeUploadResponseSerializer, TaskStagedArtifactsPrepareUploadRequestSerializer, @@ -852,6 +857,7 @@ def _user_id(self) -> int | None: "retrieve", "logs", "session_logs", + "task_session", "stream", "stream_token", "artifacts_presign", @@ -969,11 +975,6 @@ def start(self, request, pk=None, **kwargs): 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 startable == "not_cloud": return Response( TaskRunErrorResponseSerializer({"error": "Only cloud runs can be started via this endpoint"}).data, @@ -1165,6 +1166,94 @@ 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="Prepare 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, revision = result + return Response( + TaskSessionResponseSerializer({"id": session_id, "download_url": download_url, "revision": revision}).data + ) + + @validated_request( + request_serializer=TaskSessionSyncPrepareSerializer, + responses={ + 200: OpenApiResponse(response=TaskSessionSyncPrepareResponseSerializer), + 404: OpenApiResponse(description="Task session not found"), + 409: OpenApiResponse(response=TaskRunErrorResponseSerializer), + }, + summary="Prepare an active task session sync", + strict_request_validation=True, + ) + @action( + detail=True, + methods=["post"], + url_path="task_session_sync_prepare", + required_scopes=["task:write"], + ) + def prepare_task_session_sync(self, request, pk=None, **kwargs): + task_id = self._ensure_task_accessible() + try: + result = tasks_facade.prepare_task_run_session_sync( + pk, + task_id, + self.team_id, + sandbox_id=request.validated_data["sandbox_id"], + expected_revision=request.validated_data["expected_revision"], + ) + except ValueError as error: + return Response( + TaskRunErrorResponseSerializer({"error": str(error)}).data, + status=status.HTTP_409_CONFLICT, + ) + if result is None: + raise NotFound() + session_id, sync_id, upload = result + return Response( + TaskSessionSyncPrepareResponseSerializer({"id": session_id, "sync_id": sync_id, "upload": upload}).data + ) + + @validated_request( + request_serializer=TaskSessionSyncSerializer, + responses={ + 200: OpenApiResponse(response=TaskSessionSyncResponseSerializer), + 404: OpenApiResponse(description="Task session not found"), + 409: OpenApiResponse(response=TaskRunErrorResponseSerializer), + }, + summary="Finalize an active task session sync", + strict_request_validation=True, + ) + @action(detail=True, methods=["post"], url_path="task_session_sync", required_scopes=["task:write"]) + def sync_task_session(self, request, pk=None, **kwargs): + task_id = self._ensure_task_accessible() + try: + result = tasks_facade.finalize_task_run_session_sync( + pk, + task_id, + self.team_id, + sandbox_id=request.validated_data["sandbox_id"], + sync_id=request.validated_data["sync_id"], + expected_revision=request.validated_data["expected_revision"], + ) + except (json.JSONDecodeError, ValueError) as error: + return Response( + TaskRunErrorResponseSerializer({"error": str(error)}).data, + status=status.HTTP_409_CONFLICT, + ) + if result is None: + raise NotFound() + session_id, revision = result + return Response(TaskSessionSyncResponseSerializer({"id": session_id, "revision": revision}).data) + @validated_request( request_serializer=TaskRunRelayMessageRequestSerializer, responses={ @@ -1517,16 +1606,13 @@ def stream_token(self, request, pk=None, **kwargs): ) def command(self, request, pk=None, **kwargs): task_id = self._ensure_task_accessible() - if ( - tasks_facade.task_runtime(task_id, self.team_id, self._user_id(), for_control=True) - == tasks_facade.TaskRuntime.PI - ): + method = request.validated_data["method"] + task_runtime = tasks_facade.task_runtime(task_id, self.team_id, self._user_id(), for_control=True) + if method.startswith("pi/") 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, ) - - method = request.validated_data["method"] request_id = request.validated_data.get("id") params = request.validated_data.get("params") @@ -1876,11 +1962,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/process_task/activities/get_task_processing_context.py b/products/tasks/backend/temporal/process_task/activities/get_task_processing_context.py index d5da4c9457ad..56683cc62580 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 @@ -663,6 +665,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: @@ -911,6 +916,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/relay_sandbox_events.py b/products/tasks/backend/temporal/process_task/activities/relay_sandbox_events.py index 7f83fb51a35d..c75484157fd1 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 @@ -562,8 +562,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 {} @@ -645,7 +661,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 @@ -700,7 +724,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/start_agent_server.py b/products/tasks/backend/temporal/process_task/activities/start_agent_server.py index e2f86ba16cdf..acfd39bf4ce6 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 @@ -330,6 +330,7 @@ def _invoke_start_agent_server( auto_publish=ctx.auto_publish, interaction_origin=ctx.interaction_origin, branch=params.protected_base_branch, + agent_protocol=ctx.task_runtime, runtime_adapter=ctx.runtime_adapter, provider=ctx.provider, model=ctx.model, 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..9e0bae765d9b 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 @@ -995,6 +995,8 @@ def test_get_task_processing_context_exposes_ci_prompt(self, activity_environmen @pytest.mark.django_db(transaction=True) def test_get_task_processing_context_exposes_runtime_metadata(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={ "runtime_adapter": "codex", @@ -1008,6 +1010,12 @@ def test_get_task_processing_context_exposes_runtime_metadata(self, activity_env input_data = GetTaskProcessingContextInput(run_id=str(task_run.id)) result = async_to_sync(activity_environment.run)(get_task_processing_context, input_data) + task_run.refresh_from_db() + assert task_run.active_task_session is not None + assert task_run.active_task_session.object_storage_key.startswith( + f"task-sessions/{test_task.team.organization_id}/{test_task.id}/" + ) + assert result.task_runtime == "pi" assert result.runtime_adapter == "codex" assert result.provider == "openai" assert result.model == "gpt-5.3-codex" 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 ee7a779743d6..a26de5e86d18 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 @@ -62,6 +62,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): @@ -147,6 +152,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/tests/test_api.py b/products/tasks/backend/tests/test_api.py index fbf566c3204d..d5cfce55aa51 100644 --- a/products/tasks/backend/tests/test_api.py +++ b/products/tasks/backend/tests/test_api.py @@ -30,6 +30,7 @@ from posthog.storage import object_storage from products.slack_app.backend.models import SlackThreadTaskMapping +from products.tasks.backend.constants import TASK_SESSION_MAX_SIZE_BYTES, TASK_SESSION_UPLOAD_FORM_OVERHEAD_BYTES from products.tasks.backend.facade import api as tasks_facade from products.tasks.backend.logic.services.code_usage_gate import ( CodeUsageStatus, @@ -54,6 +55,7 @@ get_task_run_stream_key, ) from products.tasks.backend.models import ( + Channel, CodeInvite, CodeInviteRedemption, SandboxCustomImage, @@ -62,6 +64,7 @@ TaskArtifact, TaskAutomation, TaskRun, + TaskSession, ) from products.tasks.backend.presentation.serializers import ( TASK_RUN_ARTIFACT_MAX_SIZE_BYTES, @@ -1631,15 +1634,20 @@ 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()) - mock_workflow.assert_not_called() + self.assertEqual(response.status_code, status.HTTP_200_OK) + run = task.runs.get() + mock_workflow.assert_called_once_with( + task_id=str(task.id), + run_id=str(run.id), + team_id=task.team.id, + user_id=self.user.id, + posthog_mcp_scopes="full", + ) @parameterized.expand( [ @@ -2250,15 +2258,20 @@ 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.") - mock_workflow.assert_not_called() + self.assertEqual(response.status_code, status.HTTP_200_OK) + mock_workflow.assert_called_once_with( + task_id=str(task.id), + run_id=str(task_run.id), + team_id=task.team.id, + user_id=self.user.id, + posthog_mcp_scopes="full", + ) @patch("products.tasks.backend.temporal.client.execute_task_processing_workflow") def test_start_run_endpoint_rejects_missing_run_artifacts(self, mock_workflow): @@ -4454,7 +4467,7 @@ def test_update_run_status_to_completed_signals_workflow(self, mock_signal): self.assertIsNotNone(run.completed_at) @patch("products.tasks.backend.temporal.client.resume_task_in_cloud_workflow") - def test_resume_in_cloud_rejects_pi_task(self, mock_resume): + def test_resume_in_cloud_starts_pi_task(self, mock_resume): task = self.create_task(runtime=Task.Runtime.PI) run = TaskRun.objects.create( task=task, @@ -4465,9 +4478,11 @@ def test_resume_in_cloud_rejects_pi_task(self, mock_resume): 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_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): @@ -8134,6 +8149,323 @@ def test_command_proxies_cancel(self, mock_post): self.assertEqual(response.status_code, status.HTTP_200_OK) self.assertTrue(response.json()["result"]["cancelled"]) + @patch("posthog.storage.object_storage.get_presigned_url") + def test_task_session_returns_read_only_storage_access(self, mock_download_url): + 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"]) + 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)) + self.assertEqual(response.json()["download_url"], "https://storage.example/session.jsonl") + self.assertEqual(response.json()["revision"], 0) + self.assertNotIn("upload", response.json()) + + @patch("posthog.storage.object_storage.get_presigned_post") + @patch("posthog.storage.object_storage.head_object") + @patch("posthog.storage.object_storage.copy") + @patch("posthog.storage.object_storage.read") + @patch("products.tasks.backend.models.object_storage.tag") + def test_task_session_sync_validates_and_promotes_an_immutable_revision( + self, mock_tag, mock_read, mock_copy, mock_head, mock_upload + ): + 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.state = {**run.state, "sandbox_id": "sandbox-1"} + run.save(update_fields=["active_task_session", "state"]) + mock_upload.return_value = { + "url": "https://storage.example/upload", + "fields": {"key": "pending-session.jsonl"}, + } + mock_head.return_value = {"ContentLength": 31} + mock_read.return_value = '{"type":"session","version":3}\n' + + prepare_response = self.client.post( + f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/task_session_sync_prepare/", + {"sandbox_id": "sandbox-1", "expected_revision": 0}, + format="json", + ) + self.assertEqual(prepare_response.status_code, status.HTTP_200_OK) + task_session.refresh_from_db() + pending_key = task_session.pending_object_storage_key + self.assertIsNotNone(pending_key) + mock_upload.assert_called_once_with( + pending_key, + conditions=[ + ["content-length-range", 1, TASK_SESSION_MAX_SIZE_BYTES + TASK_SESSION_UPLOAD_FORM_OVERHEAD_BYTES] + ], + ) + + response = self.client.post( + f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/task_session_sync/", + { + "sandbox_id": "sandbox-1", + "sync_id": prepare_response.json()["sync_id"], + "expected_revision": 0, + }, + format="json", + ) + + self.assertEqual(response.status_code, status.HTTP_200_OK) + self.assertEqual(response.json()["id"], str(task_session.id)) + self.assertEqual(response.json()["revision"], 1) + task_session.refresh_from_db() + self.assertEqual(task_session.revision, 1) + self.assertIsNone(task_session.pending_sync_id) + self.assertNotEqual(task_session.object_storage_key, pending_key) + mock_copy.assert_called_once_with(pending_key, task_session.object_storage_key) + mock_read.assert_called_once_with(task_session.object_storage_key, missing_ok=True) + mock_tag.assert_called_once() + + @patch("posthog.storage.object_storage.get_presigned_post") + @patch("posthog.storage.object_storage.head_object") + @patch("posthog.storage.object_storage.copy") + @patch("posthog.storage.object_storage.read") + @patch("posthog.storage.object_storage.delete") + def test_task_session_sync_rejects_an_oversized_upload( + self, mock_delete, mock_read, mock_copy, mock_head, mock_upload + ): + 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.state = {**run.state, "sandbox_id": "sandbox-1"} + run.save(update_fields=["active_task_session", "state"]) + mock_upload.return_value = { + "url": "https://storage.example/upload", + "fields": {"key": "pending-session.jsonl"}, + } + mock_head.return_value = {"ContentLength": TASK_SESSION_MAX_SIZE_BYTES + 1} + + prepare_response = self.client.post( + f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/task_session_sync_prepare/", + {"sandbox_id": "sandbox-1", "expected_revision": 0}, + format="json", + ) + task_session.refresh_from_db() + pending_key = task_session.pending_object_storage_key + sync_id = prepare_response.json()["sync_id"] + with self.captureOnCommitCallbacks(execute=True): + response = self.client.post( + f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/task_session_sync/", + { + "sandbox_id": "sandbox-1", + "sync_id": sync_id, + "expected_revision": 0, + }, + format="json", + ) + + self.assertEqual(response.status_code, status.HTTP_409_CONFLICT) + self.assertEqual(response.json()["error"], "The uploaded task session exceeds the maximum size") + task_session.refresh_from_db() + self.assertIsNone(task_session.pending_sync_id) + self.assertIsNone(task_session.pending_object_storage_key) + mock_copy.assert_not_called() + mock_read.assert_not_called() + self.assertIn(call(pending_key), mock_delete.call_args_list) + + @patch("posthog.storage.object_storage.get_presigned_post") + @patch("posthog.storage.object_storage.head_object", return_value={"ContentLength": 31}) + @patch("posthog.storage.object_storage.copy") + @patch("posthog.storage.object_storage.read", return_value="not-json\n") + @patch("posthog.storage.object_storage.delete") + def test_task_session_sync_rejects_invalid_jsonl_and_clears_the_pending_upload( + self, mock_delete, mock_read, mock_copy, mock_head, mock_upload + ): + 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.state = {**run.state, "sandbox_id": "sandbox-1"} + run.save(update_fields=["active_task_session", "state"]) + mock_upload.return_value = { + "url": "https://storage.example/upload", + "fields": {"key": "pending-session.jsonl"}, + } + + prepare_response = self.client.post( + f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/task_session_sync_prepare/", + {"sandbox_id": "sandbox-1", "expected_revision": 0}, + format="json", + ) + task_session.refresh_from_db() + pending_key = task_session.pending_object_storage_key + sync_id = prepare_response.json()["sync_id"] + promoted_key = ( + f"task-sessions/{task_session.organization_id}/{task_session.task_id}/{task_session.id}/" + f"revisions/1-{sync_id}.jsonl" + ) + + with self.captureOnCommitCallbacks(execute=True): + response = self.client.post( + f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/task_session_sync/", + {"sandbox_id": "sandbox-1", "sync_id": sync_id, "expected_revision": 0}, + format="json", + ) + + self.assertEqual(response.status_code, status.HTTP_409_CONFLICT) + task_session.refresh_from_db() + self.assertIsNone(task_session.pending_sync_id) + self.assertIsNone(task_session.pending_object_storage_key) + self.assertCountEqual(mock_delete.call_args_list, [call(pending_key), call(promoted_key)]) + + @patch("posthog.storage.object_storage.get_presigned_post") + @patch("posthog.storage.object_storage.head_object") + @patch("posthog.storage.object_storage.delete") + def test_task_session_sync_preserves_pending_upload_after_storage_error(self, mock_delete, mock_head, mock_upload): + 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.state = {**run.state, "sandbox_id": "sandbox-1"} + run.save(update_fields=["active_task_session", "state"]) + mock_upload.return_value = { + "url": "https://storage.example/upload", + "fields": {"key": "pending-session.jsonl"}, + } + + prepare_response = self.client.post( + f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/task_session_sync_prepare/", + {"sandbox_id": "sandbox-1", "expected_revision": 0}, + format="json", + ) + task_session.refresh_from_db() + pending_sync_id = task_session.pending_sync_id + pending_key = task_session.pending_object_storage_key + mock_head.return_value = None + + response = self.client.post( + f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/task_session_sync/", + { + "sandbox_id": "sandbox-1", + "sync_id": prepare_response.json()["sync_id"], + "expected_revision": 0, + }, + format="json", + ) + + self.assertEqual(response.status_code, status.HTTP_409_CONFLICT) + task_session.refresh_from_db() + self.assertEqual(task_session.pending_sync_id, pending_sync_id) + self.assertEqual(task_session.pending_object_storage_key, pending_key) + mock_delete.assert_not_called() + + @patch("posthog.storage.object_storage.get_presigned_post") + @patch("posthog.storage.object_storage.head_object", return_value={"ContentLength": 31}) + @patch("posthog.storage.object_storage.copy") + @patch("posthog.storage.object_storage.read", return_value='{"type":"session","version":3}\n') + @patch("posthog.storage.object_storage.delete") + def test_serialized_duplicate_finalize_recovers_without_rewriting_the_promoted_revision( + self, mock_delete, mock_read, mock_copy, mock_head, mock_upload + ): + 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.state = {**run.state, "sandbox_id": "sandbox-1"} + run.save(update_fields=["active_task_session", "state"]) + mock_upload.return_value = { + "url": "https://storage.example/upload", + "fields": {"key": "pending-session.jsonl"}, + } + + prepare_response = self.client.post( + f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/task_session_sync_prepare/", + {"sandbox_id": "sandbox-1", "expected_revision": 0}, + format="json", + ) + sync_id = prepare_response.json()["sync_id"] + finalize_url = f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/task_session_sync/" + finalize_payload = {"sandbox_id": "sandbox-1", "sync_id": sync_id, "expected_revision": 0} + + first_response = self.client.post(finalize_url, finalize_payload, format="json") + self.assertEqual(first_response.status_code, status.HTTP_200_OK) + + task_session.refresh_from_db() + promoted_key = task_session.object_storage_key + mock_head.reset_mock() + mock_copy.reset_mock() + mock_read.reset_mock() + + duplicate_response = self.client.post(finalize_url, finalize_payload, format="json") + + self.assertEqual(duplicate_response.status_code, status.HTTP_200_OK) + self.assertEqual(duplicate_response.json()["revision"], 1) + task_session.refresh_from_db() + self.assertEqual(task_session.object_storage_key, promoted_key) + mock_head.assert_not_called() + mock_copy.assert_not_called() + mock_read.assert_not_called() + mock_delete.assert_not_called() + + @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.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) + mock_signal_followup.assert_called_once() + + @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": {"type": "future_native_command"}}, + "id": "native", + }, + format="json", + ) + + self.assertEqual(response.status_code, status.HTTP_200_OK) + self.assertEqual(response.json()["result"], rpc_response) + @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_models.py b/products/tasks/backend/tests/test_models.py index 42e2ce8a3a2b..6d765c83e92b 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, diff --git a/products/tasks/backend/tests/test_storage.py b/products/tasks/backend/tests/test_storage.py new file mode 100644 index 000000000000..7090d078ce51 --- /dev/null +++ b/products/tasks/backend/tests/test_storage.py @@ -0,0 +1,68 @@ +import uuid + +from unittest.mock import MagicMock, patch + +from django.test import SimpleTestCase + +from parameterized import parameterized + +from products.tasks.backend.models import TaskSession +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) + + +class TestTaskSessionStorage(SimpleTestCase): + @patch("products.tasks.backend.models.object_storage.tag") + @patch("products.tasks.backend.storage.object_storage.write") + @patch("products.tasks.backend.storage.object_storage.read") + def test_new_session_object_is_attributed_without_expiry( + self, + mock_read: MagicMock, + mock_write: MagicMock, + mock_tag: MagicMock, + ) -> None: + organization_id = uuid.uuid4() + task_id = uuid.uuid4() + task_session = TaskSession( + id=uuid.uuid4(), + organization_id=organization_id, + task_id=task_id, + object_storage_key="sessions/example.jsonl", + ) + mock_read.return_value = None + + task_session.append_entries([{"type": "session"}]) + + mock_write.assert_called_once_with("sessions/example.jsonl", '{"type": "session"}') + mock_tag.assert_called_once_with( + "sessions/example.jsonl", + { + "data_class": "task_session", + "organization_id": str(organization_id), + "task_id": str(task_id), + }, + ) diff --git a/products/tasks/frontend/generated/api.schemas.ts b/products/tasks/frontend/generated/api.schemas.ts index 4e529073ac9f..d8ead9d350e0 100644 --- a/products/tasks/frontend/generated/api.schemas.ts +++ b/products/tasks/frontend/generated/api.schemas.ts @@ -2694,6 +2694,7 @@ export const JsonrpcEnumApi = { * * `permission_response` - permission_response * * `set_config_option` - set_config_option * * `mcp_response` - mcp_response + * * `pi/rpc` - pi/rpc */ export type MethodEnumApi = (typeof MethodEnumApi)[keyof typeof MethodEnumApi] @@ -2704,6 +2705,7 @@ export const MethodEnumApi = { PermissionResponse: 'permission_response', SetConfigOption: 'set_config_option', McpResponse: 'mcp_response', + PiRpc: 'pi/rpc', } as const /** @@ -2721,7 +2723,8 @@ 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 */ method: MethodEnumApi /** Parameters for the command */ params?: TaskRunCommandRequestApiParams @@ -2729,11 +2732,6 @@ export interface TaskRunCommandRequestApi { id?: unknown } -/** - * Command result on success - */ -export type TaskRunCommandResponseApiResult = { [key: string]: unknown } - /** * Error details on failure */ @@ -2748,7 +2746,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 } @@ -2815,6 +2813,40 @@ export interface StreamReadTokenResponseApi { stream_base_url: string | null } +export interface TaskSessionResponseApi { + id: string + download_url: string + /** @minimum 0 */ + revision: number +} + +export interface TaskSessionSyncApi { + sandbox_id: string + sync_id: string + /** @minimum 0 */ + expected_revision: number +} + +export interface TaskSessionSyncResponseApi { + id: string + /** @minimum 1 */ + revision: number +} + +export interface TaskSessionSyncPrepareApi { + sandbox_id: string + /** @minimum 0 */ + expected_revision: number +} + +export type TaskSessionSyncPrepareResponseApiUpload = { [key: string]: unknown } + +export interface TaskSessionSyncPrepareResponseApi { + id: string + sync_id: string + upload: TaskSessionSyncPrepareResponseApiUpload +} + /** * * `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 700f657f4f23..5b9b82dbeeb8 100644 --- a/products/tasks/frontend/generated/api.ts +++ b/products/tasks/frontend/generated/api.ts @@ -89,6 +89,11 @@ import type { TaskRunRelayMessageRequestApi, TaskRunRelayMessageResponseApi, TaskRunStartRequestApi, + TaskSessionResponseApi, + TaskSessionSyncApi, + TaskSessionSyncPrepareApi, + TaskSessionSyncPrepareResponseApi, + TaskSessionSyncResponseApi, TaskStagedArtifactsFinalizeUploadRequestApi, TaskStagedArtifactsFinalizeUploadResponseApi, TaskStagedArtifactsPrepareUploadRequestApi, @@ -1662,6 +1667,75 @@ 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 Prepare 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 Finalize an active task session sync + */ +export const tasksRunsTaskSessionSyncCreate = async ( + projectId: string, + taskId: string, + id: string, + taskSessionSyncApi: TaskSessionSyncApi, + options?: RequestInit +): Promise => { + return apiMutator(getTasksRunsTaskSessionSyncCreateUrl(projectId, taskId, id), { + ...options, + method: 'POST', + headers: { 'Content-Type': 'application/json', ...options?.headers }, + body: JSON.stringify(taskSessionSyncApi), + }) +} + +export const getTasksRunsTaskSessionSyncPrepareCreateUrl = (projectId: string, taskId: string, id: string) => { + return `/api/projects/${projectId}/tasks/${taskId}/runs/${id}/task_session_sync_prepare/` +} + +/** + * API for managing task runs. Each run represents an execution of a task. + * @summary Prepare an active task session sync + */ +export const tasksRunsTaskSessionSyncPrepareCreate = async ( + projectId: string, + taskId: string, + id: string, + taskSessionSyncPrepareApi: TaskSessionSyncPrepareApi, + options?: RequestInit +): Promise => { + return apiMutator( + getTasksRunsTaskSessionSyncPrepareCreateUrl(projectId, taskId, id), + { + ...options, + method: 'POST', + headers: { 'Content-Type': 'application/json', ...options?.headers }, + body: JSON.stringify(taskSessionSyncPrepareApi), + } + ) +} + 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 74fc911c81df..2a71c3941c1c 100644 --- a/products/tasks/frontend/generated/api.zod.ts +++ b/products/tasks/frontend/generated/api.zod.ts @@ -2564,12 +2564,20 @@ 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', + ]) .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' ) .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' ), 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)'), @@ -2632,6 +2640,29 @@ export const TasksRunsStartCreateBody = /* @__PURE__ */ zod.object({ ), }) +/** + * API for managing task runs. Each run represents an execution of a task. + * @summary Finalize an active task session sync + */ +export const tasksRunsTaskSessionSyncCreateBodyExpectedRevisionMin = 0 + +export const TasksRunsTaskSessionSyncCreateBody = /* @__PURE__ */ zod.object({ + sandbox_id: zod.string(), + sync_id: zod.uuid(), + expected_revision: zod.number().min(tasksRunsTaskSessionSyncCreateBodyExpectedRevisionMin), +}) + +/** + * API for managing task runs. Each run represents an execution of a task. + * @summary Prepare an active task session sync + */ +export const tasksRunsTaskSessionSyncPrepareCreateBodyExpectedRevisionMin = 0 + +export const TasksRunsTaskSessionSyncPrepareCreateBody = /* @__PURE__ */ zod.object({ + sandbox_id: zod.string(), + expected_revision: zod.number().min(tasksRunsTaskSessionSyncPrepareCreateBodyExpectedRevisionMin), +}) + /** * Create a stable, editable artifact handle from direct markdown/text content or an existing run artifact. Slack adapters deliver into the mapped Slack thread; document artifacts use external connector storage when available. * @summary Create a living artifact for a task run diff --git a/services/mcp/src/api/generated.ts b/services/mcp/src/api/generated.ts index 060c3c7d28f5..0d06a48556c1 100644 --- a/services/mcp/src/api/generated.ts +++ b/services/mcp/src/api/generated.ts @@ -39308,6 +39308,7 @@ export namespace Schemas { * * `permission_response` - permission_response * * `set_config_option` - set_config_option * * `mcp_response` - mcp_response + * * `pi/rpc` - pi/rpc */ export type MethodEnum = typeof MethodEnum[keyof typeof MethodEnum]; @@ -39319,6 +39320,7 @@ export namespace Schemas { PermissionResponse: 'permission_response', SetConfigOption: 'set_config_option', McpResponse: 'mcp_response', + PiRpc: 'pi/rpc', } as const; /** @@ -67543,7 +67545,8 @@ 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 */ method: MethodEnum; /** Parameters for the command */ params?: TaskRunCommandRequestParams; @@ -67551,11 +67554,6 @@ export namespace Schemas { id?: unknown; } - /** - * Command result on success - */ - export type TaskRunCommandResponseResult = { [key: string]: unknown }; - /** * Error details on failure */ @@ -67570,7 +67568,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; } @@ -67875,6 +67873,40 @@ export namespace Schemas { pending_user_artifact_ids?: string[]; } + export interface TaskSessionResponse { + id: string; + download_url: string; + /** @minimum 0 */ + revision: number; + } + + export interface TaskSessionSync { + sandbox_id: string; + sync_id: string; + /** @minimum 0 */ + expected_revision: number; + } + + export interface TaskSessionSyncPrepare { + sandbox_id: string; + /** @minimum 0 */ + expected_revision: number; + } + + export type TaskSessionSyncPrepareResponseUpload = { [key: string]: unknown }; + + export interface TaskSessionSyncPrepareResponse { + id: string; + sync_id: string; + upload: TaskSessionSyncPrepareResponseUpload; + } + + export interface TaskSessionSyncResponse { + id: string; + /** @minimum 1 */ + revision: number; + } + export interface TaskStagedArtifactFinalizeUpload { /** Stable identifier returned by the staged prepare upload endpoint */ id: string; From b0794761342d43cc17ff6d28033c7117b22ee3f9 Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Wed, 22 Jul 2026 17:47:57 +0200 Subject: [PATCH 02/22] fix(tasks): tighten Pi session cleanup and commands --- products/tasks/backend/facade/api.py | 23 ++++--- .../tasks/backend/presentation/views/api.py | 7 ++- products/tasks/backend/tests/test_api.py | 63 +++++++++++++++++-- 3 files changed, 80 insertions(+), 13 deletions(-) diff --git a/products/tasks/backend/facade/api.py b/products/tasks/backend/facade/api.py index b4b3e4af6bcc..c0aa10535f5a 100644 --- a/products/tasks/backend/facade/api.py +++ b/products/tasks/backend/facade/api.py @@ -2236,6 +2236,13 @@ def get_task_run_session(run_id: str | UUID, task_id: str | UUID, team_id: int) return task_session.id, download_url, task_session.revision +def _task_session_sync_key(task_session: TaskSession, directory: str, revision: int, sync_id: UUID) -> str: + return ( + f"task-sessions/{task_session.organization_id}/{task_session.task_id}/{task_session.id}/" + f"{directory}/{revision}-{sync_id}.jsonl" + ) + + def prepare_task_run_session_sync( run_id: str | UUID, task_id: str | UUID, @@ -2260,10 +2267,7 @@ def prepare_task_run_session_sync( stale_pending_key = task_session.pending_object_storage_key sync_id = uuid4() - object_storage_key = ( - f"task-sessions/{task_session.organization_id}/{task_session.task_id}/{task_session.id}/" - f"uploads/{task_session.revision + 1}-{sync_id}.jsonl" - ) + object_storage_key = _task_session_sync_key(task_session, "uploads", task_session.revision + 1, sync_id) task_session.pending_sync_id = sync_id task_session.pending_object_storage_key = object_storage_key task_session.save(update_fields=["pending_sync_id", "pending_object_storage_key", "updated_at"]) @@ -2349,6 +2353,7 @@ def finalize_task_run_session_sync( task_session_id = visible_run.active_task_session_id pending_key: str | None = None promoted_key: str | None = None + stale_upload_key: str | None = None try: with transaction.atomic(): @@ -2359,9 +2364,9 @@ def finalize_task_run_session_sync( raise ValueError("The task session sync is stale") locked_session = TaskSession.objects.select_for_update().get(id=task_session_id) - promoted_object_storage_key = ( - f"task-sessions/{locked_session.organization_id}/{locked_session.task_id}/{locked_session.id}/" - f"revisions/{expected_revision + 1}-{sync_id}.jsonl" + submitted_upload_key = _task_session_sync_key(locked_session, "uploads", expected_revision + 1, sync_id) + promoted_object_storage_key = _task_session_sync_key( + locked_session, "revisions", expected_revision + 1, sync_id ) promoted_key = promoted_object_storage_key if ( @@ -2375,6 +2380,8 @@ def finalize_task_run_session_sync( or locked_session.pending_sync_id != sync_id or not pending_object_storage_key ): + if locked_session.pending_sync_id != sync_id or pending_object_storage_key != submitted_upload_key: + stale_upload_key = submitted_upload_key raise ValueError("The task session sync is stale") pending_key = pending_object_storage_key @@ -2443,6 +2450,8 @@ def finalize_task_run_session_sync( _reject_task_session_sync(task_session_id, sync_id, pending_key, promoted_key) raise except Exception: + if stale_upload_key is not None: + _delete_task_session_objects(task_session_id, (stale_upload_key,)) if promoted_key is not None: finalized_session = TaskSession.objects.get(id=task_session_id) if ( diff --git a/products/tasks/backend/presentation/views/api.py b/products/tasks/backend/presentation/views/api.py index d603b66aa26e..82a8c12526b7 100644 --- a/products/tasks/backend/presentation/views/api.py +++ b/products/tasks/backend/presentation/views/api.py @@ -1595,7 +1595,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, and native Pi RPC commands.", strict_request_validation=True, ) @action( @@ -1613,6 +1613,11 @@ def command(self, request, pk=None, **kwargs): 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"}: + return Response( + TaskRunErrorResponseSerializer({"error": f"{method} is not supported for Pi tasks."}).data, + status=status.HTTP_400_BAD_REQUEST, + ) request_id = request.validated_data.get("id") params = request.validated_data.get("params") diff --git a/products/tasks/backend/tests/test_api.py b/products/tasks/backend/tests/test_api.py index d5cfce55aa51..d2d618b7a3a6 100644 --- a/products/tasks/backend/tests/test_api.py +++ b/products/tasks/backend/tests/test_api.py @@ -7958,20 +7958,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): @@ -8357,6 +8360,56 @@ def test_task_session_sync_preserves_pending_upload_after_storage_error(self, mo self.assertEqual(task_session.pending_object_storage_key, pending_key) mock_delete.assert_not_called() + @patch("posthog.storage.object_storage.get_presigned_post") + @patch("posthog.storage.object_storage.delete") + def test_stale_finalize_deletes_the_superseded_upload(self, mock_delete, mock_upload): + 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.state = {**run.state, "sandbox_id": "sandbox-1"} + run.save(update_fields=["active_task_session", "state"]) + mock_upload.return_value = { + "url": "https://storage.example/upload", + "fields": {"key": "pending-session.jsonl"}, + } + prepare_url = f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/task_session_sync_prepare/" + + first_prepare = self.client.post( + prepare_url, + {"sandbox_id": "sandbox-1", "expected_revision": 0}, + format="json", + ) + task_session.refresh_from_db() + first_pending_key = task_session.pending_object_storage_key + + with self.captureOnCommitCallbacks(execute=True): + second_prepare = self.client.post( + prepare_url, + {"sandbox_id": "sandbox-1", "expected_revision": 0}, + format="json", + ) + self.assertEqual(second_prepare.status_code, status.HTTP_200_OK) + task_session.refresh_from_db() + current_sync_id = task_session.pending_sync_id + mock_delete.reset_mock() + + response = self.client.post( + f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/task_session_sync/", + { + "sandbox_id": "sandbox-1", + "sync_id": first_prepare.json()["sync_id"], + "expected_revision": 0, + }, + format="json", + ) + + self.assertEqual(response.status_code, status.HTTP_409_CONFLICT) + self.assertEqual(response.json()["error"], "The task session sync is stale") + task_session.refresh_from_db() + self.assertEqual(task_session.pending_sync_id, current_sync_id) + mock_delete.assert_called_once_with(first_pending_key) + @patch("posthog.storage.object_storage.get_presigned_post") @patch("posthog.storage.object_storage.head_object", return_value={"ContentLength": 31}) @patch("posthog.storage.object_storage.copy") From ff419ea7b63165b4d4f78c6b79611f52cb5093bd Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Thu, 23 Jul 2026 13:44:49 +0200 Subject: [PATCH 03/22] fix(tasks): harden Pi session sync and local overlays --- products/tasks/backend/facade/api.py | 31 ++++--------- .../backend/logic/services/modal_sandbox.py | 30 +++++++++++++ .../services/tests/test_modal_sandbox.py | 13 ++++-- products/tasks/backend/tests/test_api.py | 43 +++++-------------- 4 files changed, 59 insertions(+), 58 deletions(-) diff --git a/products/tasks/backend/facade/api.py b/products/tasks/backend/facade/api.py index c0aa10535f5a..0454bf8aa1e0 100644 --- a/products/tasks/backend/facade/api.py +++ b/products/tasks/backend/facade/api.py @@ -2265,29 +2265,14 @@ def prepare_task_run_session_sync( if task_session.revision != expected_revision: raise ValueError("The task session revision is stale") - stale_pending_key = task_session.pending_object_storage_key - sync_id = uuid4() - object_storage_key = _task_session_sync_key(task_session, "uploads", task_session.revision + 1, sync_id) - task_session.pending_sync_id = sync_id - task_session.pending_object_storage_key = object_storage_key - task_session.save(update_fields=["pending_sync_id", "pending_object_storage_key", "updated_at"]) - - if stale_pending_key: - - def delete_stale_pending_object() -> None: - try: - object_storage.delete(stale_pending_key) - except Exception as error: - logger.warning( - "task_session.failed_to_delete_stale_pending_object", - extra={ - "task_session_id": str(task_session.id), - "object_storage_key": stale_pending_key, - "error": str(error), - }, - ) - - transaction.on_commit(delete_stale_pending_object) + sync_id = task_session.pending_sync_id + object_storage_key = task_session.pending_object_storage_key + if sync_id is None or object_storage_key is None: + sync_id = uuid4() + object_storage_key = _task_session_sync_key(task_session, "uploads", task_session.revision + 1, sync_id) + task_session.pending_sync_id = sync_id + task_session.pending_object_storage_key = object_storage_key + task_session.save(update_fields=["pending_sync_id", "pending_object_storage_key", "updated_at"]) upload = object_storage.get_presigned_post( object_storage_key, diff --git a/products/tasks/backend/logic/services/modal_sandbox.py b/products/tasks/backend/logic/services/modal_sandbox.py index ffb787a18698..8732e8697ae9 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), 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 2942afd49492..0eed9965c506 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, diff --git a/products/tasks/backend/tests/test_api.py b/products/tasks/backend/tests/test_api.py index d2d618b7a3a6..5fa904ee20a0 100644 --- a/products/tasks/backend/tests/test_api.py +++ b/products/tasks/backend/tests/test_api.py @@ -8361,8 +8361,7 @@ def test_task_session_sync_preserves_pending_upload_after_storage_error(self, mo mock_delete.assert_not_called() @patch("posthog.storage.object_storage.get_presigned_post") - @patch("posthog.storage.object_storage.delete") - def test_stale_finalize_deletes_the_superseded_upload(self, mock_delete, mock_upload): + def test_repeated_prepare_reuses_the_pending_upload(self, mock_upload): task = self.create_task(runtime=Task.Runtime.PI) run = self._create_run_with_sandbox(task) task_session = TaskSession.create_for_task(task) @@ -8374,41 +8373,21 @@ def test_stale_finalize_deletes_the_superseded_upload(self, mock_delete, mock_up "fields": {"key": "pending-session.jsonl"}, } prepare_url = f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/task_session_sync_prepare/" + payload = {"sandbox_id": "sandbox-1", "expected_revision": 0} - first_prepare = self.client.post( - prepare_url, - {"sandbox_id": "sandbox-1", "expected_revision": 0}, - format="json", - ) + first_prepare = self.client.post(prepare_url, payload, format="json") task_session.refresh_from_db() - first_pending_key = task_session.pending_object_storage_key - - with self.captureOnCommitCallbacks(execute=True): - second_prepare = self.client.post( - prepare_url, - {"sandbox_id": "sandbox-1", "expected_revision": 0}, - format="json", - ) - self.assertEqual(second_prepare.status_code, status.HTTP_200_OK) - task_session.refresh_from_db() - current_sync_id = task_session.pending_sync_id - mock_delete.reset_mock() + pending_key = task_session.pending_object_storage_key - response = self.client.post( - f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/task_session_sync/", - { - "sandbox_id": "sandbox-1", - "sync_id": first_prepare.json()["sync_id"], - "expected_revision": 0, - }, - format="json", - ) + second_prepare = self.client.post(prepare_url, payload, format="json") - self.assertEqual(response.status_code, status.HTTP_409_CONFLICT) - self.assertEqual(response.json()["error"], "The task session sync is stale") + self.assertEqual(first_prepare.status_code, status.HTTP_200_OK) + self.assertEqual(second_prepare.status_code, status.HTTP_200_OK) + self.assertEqual(second_prepare.json()["sync_id"], first_prepare.json()["sync_id"]) task_session.refresh_from_db() - self.assertEqual(task_session.pending_sync_id, current_sync_id) - mock_delete.assert_called_once_with(first_pending_key) + self.assertEqual(task_session.pending_object_storage_key, pending_key) + self.assertEqual(mock_upload.call_count, 2) + self.assertEqual(mock_upload.call_args_list[0], mock_upload.call_args_list[1]) @patch("posthog.storage.object_storage.get_presigned_post") @patch("posthog.storage.object_storage.head_object", return_value={"ContentLength": 31}) From 225ca7575bbbf4021a167d65fcf04f3f8c5df652 Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Thu, 23 Jul 2026 19:14:29 +0200 Subject: [PATCH 04/22] refactor(tasks): align agent runtime naming --- .../tasks/backend/logic/services/docker_sandbox.py | 10 +++++----- products/tasks/backend/logic/services/modal_sandbox.py | 8 ++++---- products/tasks/backend/logic/services/sandbox.py | 6 +++--- .../backend/logic/services/test_docker_sandbox.py | 8 ++++---- .../backend/logic/services/tests/test_modal_sandbox.py | 4 ++-- .../process_task/activities/start_agent_server.py | 2 +- 6 files changed, 19 insertions(+), 19 deletions(-) diff --git a/products/tasks/backend/logic/services/docker_sandbox.py b/products/tasks/backend/logic/services/docker_sandbox.py index c2fb0140b65c..fdb9ef6207d6 100644 --- a/products/tasks/backend/logic/services/docker_sandbox.py +++ b/products/tasks/backend/logic/services/docker_sandbox.py @@ -774,7 +774,7 @@ def _build_agent_server_command( auto_publish: bool = False, interaction_origin: str | None = None, branch: str | None = None, - agent_protocol: str | None = None, + agent_runtime: str | None = None, runtime_adapter: str | None = None, provider: str | None = None, model: str | None = None, @@ -796,7 +796,7 @@ 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_protocol=agent_protocol, + agent_runtime=agent_runtime, sandbox_id=self.id, runtime_adapter=runtime_adapter, provider=provider, @@ -882,7 +882,7 @@ def start_agent_server( auto_publish: bool = False, interaction_origin: str | None = None, branch: str | None = None, - agent_protocol: str | None = None, + agent_runtime: str | None = None, runtime_adapter: str | None = None, provider: str | None = None, model: str | None = None, @@ -954,7 +954,7 @@ def start_agent_server( auto_publish, interaction_origin, branch, - agent_protocol, + agent_runtime, runtime_adapter, provider, model, @@ -1006,7 +1006,7 @@ def start_agent_server( auto_publish, interaction_origin, branch=None, - agent_protocol=agent_protocol, + agent_runtime=agent_runtime, runtime_adapter=runtime_adapter, provider=provider, model=model, diff --git a/products/tasks/backend/logic/services/modal_sandbox.py b/products/tasks/backend/logic/services/modal_sandbox.py index 8732e8697ae9..2ca36cc3f661 100644 --- a/products/tasks/backend/logic/services/modal_sandbox.py +++ b/products/tasks/backend/logic/services/modal_sandbox.py @@ -975,7 +975,7 @@ def _build_agent_server_command( auto_publish: bool = False, interaction_origin: str | None = None, branch: str | None = None, - agent_protocol: str | None = None, + agent_runtime: str | None = None, runtime_adapter: str | None = None, provider: str | None = None, model: str | None = None, @@ -993,7 +993,7 @@ def _build_agent_server_command( ) -> str: env_prefix = build_agent_runtime_env_prefix( interaction_origin=interaction_origin, - agent_protocol=agent_protocol, + agent_runtime=agent_runtime, sandbox_id=self.id, runtime_adapter=runtime_adapter, provider=provider, @@ -1105,7 +1105,7 @@ def start_agent_server( auto_publish: bool = False, interaction_origin: str | None = None, branch: str | None = None, - agent_protocol: str | None = None, + agent_runtime: str | None = None, runtime_adapter: str | None = None, provider: str | None = None, model: str | None = None, @@ -1180,7 +1180,7 @@ def start_agent_server( auto_publish, interaction_origin, branch, - agent_protocol, + agent_runtime, runtime_adapter, provider, model, diff --git a/products/tasks/backend/logic/services/sandbox.py b/products/tasks/backend/logic/services/sandbox.py index d68315ce8a1f..91d5bb533c90 100644 --- a/products/tasks/backend/logic/services/sandbox.py +++ b/products/tasks/backend/logic/services/sandbox.py @@ -167,7 +167,7 @@ def redact_sandbox_command(command: str) -> str: def build_agent_runtime_env_prefix( *, interaction_origin: str | None = None, - agent_protocol: str | None = None, + agent_runtime: str | None = None, sandbox_id: str | None = None, runtime_adapter: str | None = None, provider: str | None = None, @@ -181,7 +181,7 @@ def build_agent_runtime_env_prefix( ) -> str: env_vars = { "POSTHOG_CODE_INTERACTION_ORIGIN": interaction_origin, - "POSTHOG_AGENT_PROTOCOL": agent_protocol, + "POSTHOG_AGENT_RUNTIME": agent_runtime, "POSTHOG_SANDBOX_ID": sandbox_id, "POSTHOG_CODE_RUNTIME_ADAPTER": runtime_adapter, "POSTHOG_CODE_PROVIDER": provider, @@ -333,7 +333,7 @@ def start_agent_server( auto_publish: bool = False, interaction_origin: str | None = None, branch: str | None = None, - agent_protocol: str | None = None, + agent_runtime: str | None = None, runtime_adapter: str | None = None, provider: str | None = None, model: str | None = None, diff --git a/products/tasks/backend/logic/services/test_docker_sandbox.py b/products/tasks/backend/logic/services/test_docker_sandbox.py index f114a21a6f1a..fef1143b2d7e 100644 --- a/products/tasks/backend/logic/services/test_docker_sandbox.py +++ b/products/tasks/backend/logic/services/test_docker_sandbox.py @@ -437,11 +437,11 @@ def test_start_agent_server_preserves_pi_protocol_without_branch_retry(self) -> "task-123", "run-456", branch="main", - agent_protocol="pi", + agent_runtime="pi", ) assert mock_build.call_count == 2 - assert mock_build.call_args_list[1].kwargs["agent_protocol"] == "pi" + assert mock_build.call_args_list[1].kwargs["agent_runtime"] == "pi" def test_parse_repo_mount_map_empty(self): with patch.dict(os.environ, {}, clear=True): @@ -655,7 +655,7 @@ def test_start_agent_server_includes_runtime_environment_variables(self): "task-123", "run-456", "background", - agent_protocol="pi", + agent_runtime="pi", runtime_adapter="codex", provider="openai", model="gpt-5.3-codex", @@ -666,7 +666,7 @@ def test_start_agent_server_includes_runtime_environment_variables(self): ) command = _agent_server_launch_command(mock_execute) - assert "POSTHOG_AGENT_PROTOCOL=pi" in command + 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 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 0eed9965c506..258737214c74 100644 --- a/products/tasks/backend/logic/services/tests/test_modal_sandbox.py +++ b/products/tasks/backend/logic/services/tests/test_modal_sandbox.py @@ -554,7 +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_protocol="pi", + agent_runtime="pi", runtime_adapter="codex", provider="openai", model="gpt-5.3-codex", @@ -565,7 +565,7 @@ def test_start_agent_server_includes_runtime_environment_variables(self, mock_sa ) command = _agent_server_launch_command(mock_sandbox.execute) - assert "POSTHOG_AGENT_PROTOCOL=pi" in command + 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/temporal/process_task/activities/start_agent_server.py b/products/tasks/backend/temporal/process_task/activities/start_agent_server.py index acfd39bf4ce6..0abac3538d11 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 @@ -330,7 +330,7 @@ def _invoke_start_agent_server( auto_publish=ctx.auto_publish, interaction_origin=ctx.interaction_origin, branch=params.protected_base_branch, - agent_protocol=ctx.task_runtime, + agent_runtime=ctx.task_runtime, runtime_adapter=ctx.runtime_adapter, provider=ctx.provider, model=ctx.model, From 468f7e5d412b507e52e28d19236c9ba52e23491e Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Fri, 24 Jul 2026 16:20:05 +0200 Subject: [PATCH 05/22] refactor(tasks): simplify pi session persistence --- products/tasks/backend/facade/api.py | 332 +++++++--------- .../backend/logic/services/sandbox_usage.py | 105 +++--- .../backend/migrations/0073_task_session.py | 20 +- .../0074_sandbox_session_open_constraint.py | 16 + .../migrations/0074_task_session_revisions.py | 23 -- .../backend/migrations/max_migration.txt | 2 +- products/tasks/backend/models.py | 76 ++-- .../tasks/backend/presentation/serializers.py | 23 +- .../tasks/backend/presentation/views/api.py | 83 ++-- products/tasks/backend/temporal/client.py | 6 +- .../activities/get_sandbox_for_repository.py | 34 +- .../activities/get_task_processing_context.py | 1 - .../activities/provision_sandbox.py | 14 +- .../activities/send_followup_to_sandbox.py | 2 - .../tests/test_get_task_processing_context.py | 27 +- .../backend/temporal/process_task/workflow.py | 23 +- products/tasks/backend/tests/test_api.py | 353 ++++++------------ products/tasks/backend/tests/test_storage.py | 36 -- 18 files changed, 458 insertions(+), 718 deletions(-) create mode 100644 products/tasks/backend/migrations/0074_sandbox_session_open_constraint.py delete mode 100644 products/tasks/backend/migrations/0074_task_session_revisions.py diff --git a/products/tasks/backend/facade/api.py b/products/tasks/backend/facade/api.py index 0454bf8aa1e0..8ab09d6dbfe6 100644 --- a/products/tasks/backend/facade/api.py +++ b/products/tasks/backend/facade/api.py @@ -16,7 +16,7 @@ """ import re -import json +import hashlib import logging from collections.abc import Iterable, Sequence from concurrent.futures import ThreadPoolExecutor @@ -43,7 +43,6 @@ MAX_CUSTOM_IMAGES_PER_USER, RESERVED_SANDBOX_ENVIRONMENT_VARIABLE_KEYS, TASK_SESSION_MAX_SIZE_BYTES, - TASK_SESSION_UPLOAD_FORM_OVERHEAD_BYTES, is_blocked_sandbox_env_key, ) from products.tasks.backend.error_telemetry import truncate_error_message @@ -60,6 +59,7 @@ CodeInviteRedemption, SandboxCustomImage, SandboxEnvironment, + SandboxSession, SandboxSnapshot, Task, TaskAutomation, @@ -155,7 +155,7 @@ "get_task_id_for_run", "get_task_run", "get_task_run_session", - "prepare_task_run_session_sync", + "sync_task_run_session", "get_task_run_detail", "get_task_run_sandbox_connection", "get_task_run_living_artifact", @@ -211,7 +211,6 @@ "update_task", "update_task_automation", "update_task_run", - "finalize_task_run_session_sync", "update_task_run_state", "upsert_internal_sandbox_env", "validate_set_output", @@ -2223,227 +2222,108 @@ def ensure_task_run_session(run_id: str | UUID) -> UUID: return task_session.id -def get_task_run_session(run_id: str | UUID, task_id: str | UUID, team_id: int) -> tuple[UUID, str, int] | None: +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.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.revision + return task_session.id, download_url, task_session.content_sha256 -def _task_session_sync_key(task_session: TaskSession, directory: str, revision: int, sync_id: UUID) -> str: - return ( - f"task-sessions/{task_session.organization_id}/{task_session.task_id}/{task_session.id}/" - f"{directory}/{revision}-{sync_id}.jsonl" - ) +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 prepare_task_run_session_sync( - run_id: str | UUID, - task_id: str | UUID, - team_id: int, - *, - sandbox_id: str, - expected_revision: int, -) -> tuple[UUID, UUID, dict] | None: - from posthog.storage import object_storage # noqa: PLC0415 +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() + ) - 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 - with transaction.atomic(): - run = TaskRun.objects.select_for_update(of=("self",)).get(id=visible_run.id) - if (run.state or {}).get("sandbox_id") != sandbox_id: - raise ValueError("The task session writer is not the active sandbox") - task_session = TaskSession.objects.select_for_update().get(id=run.active_task_session_id) - if task_session.revision != expected_revision: - raise ValueError("The task session revision is stale") - - sync_id = task_session.pending_sync_id - object_storage_key = task_session.pending_object_storage_key - if sync_id is None or object_storage_key is None: - sync_id = uuid4() - object_storage_key = _task_session_sync_key(task_session, "uploads", task_session.revision + 1, sync_id) - task_session.pending_sync_id = sync_id - task_session.pending_object_storage_key = object_storage_key - task_session.save(update_fields=["pending_sync_id", "pending_object_storage_key", "updated_at"]) - - upload = object_storage.get_presigned_post( - object_storage_key, - conditions=[["content-length-range", 1, TASK_SESSION_MAX_SIZE_BYTES + TASK_SESSION_UPLOAD_FORM_OVERHEAD_BYTES]], - ) - if not upload: - raise RuntimeError("Unable to prepare task session upload") - return task_session.id, sync_id, upload - - -def _delete_task_session_objects(task_session_id: UUID, object_storage_keys: Sequence[str]) -> None: +def _delete_task_session_object(task_session_id: UUID, object_storage_key: str) -> None: from posthog.storage import object_storage # noqa: PLC0415 - for object_storage_key in object_storage_keys: - 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 _reject_task_session_sync( - task_session_id: UUID, - sync_id: UUID, - pending_key: str, - promoted_key: str, -) -> None: - with transaction.atomic(): - task_session = TaskSession.objects.select_for_update().get(id=task_session_id) - if task_session.pending_sync_id != sync_id or task_session.pending_object_storage_key != pending_key: - return - - task_session.pending_sync_id = None - task_session.pending_object_storage_key = None - task_session.save(update_fields=["pending_sync_id", "pending_object_storage_key", "updated_at"]) - transaction.on_commit(lambda: _delete_task_session_objects(task_session_id, (pending_key, promoted_key))) - - -class _TaskSessionUploadRejected(ValueError): - pass + 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 finalize_task_run_session_sync( +def sync_task_run_session( run_id: str | UUID, task_id: str | UUID, team_id: int, *, sandbox_id: str, - sync_id: UUID, - expected_revision: int, -) -> tuple[UUID, int] | None: + 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 - task_session_id = visible_run.active_task_session_id - pending_key: str | None = None - promoted_key: str | None = None - stale_upload_key: str | None = 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.state or {}).get("sandbox_id") != sandbox_id: - raise ValueError("The task session writer is not the active sandbox") - if locked_run.active_task_session_id != task_session_id: - raise ValueError("The task session sync is stale") - - locked_session = TaskSession.objects.select_for_update().get(id=task_session_id) - submitted_upload_key = _task_session_sync_key(locked_session, "uploads", expected_revision + 1, sync_id) - promoted_object_storage_key = _task_session_sync_key( - locked_session, "revisions", expected_revision + 1, sync_id - ) - promoted_key = promoted_object_storage_key - if ( - locked_session.revision == expected_revision + 1 - and locked_session.object_storage_key == promoted_object_storage_key - ): - return locked_session.id, locked_session.revision - pending_object_storage_key = locked_session.pending_object_storage_key - if ( - locked_session.revision != expected_revision - or locked_session.pending_sync_id != sync_id - or not pending_object_storage_key - ): - if locked_session.pending_sync_id != sync_id or pending_object_storage_key != submitted_upload_key: - stale_upload_key = submitted_upload_key + 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") - pending_key = pending_object_storage_key - pending_metadata = object_storage.head_object(pending_object_storage_key) - if pending_metadata is None: - raise ValueError("The uploaded task session is missing") - pending_size = pending_metadata.get("ContentLength") - if not isinstance(pending_size, int): - raise ValueError("Unable to determine the uploaded task session size") - if pending_size < 1: - raise _TaskSessionUploadRejected("The uploaded task session is not valid JSONL") - if pending_size > TASK_SESSION_MAX_SIZE_BYTES: - raise _TaskSessionUploadRejected("The uploaded task session exceeds the maximum size") - - object_storage.copy(pending_object_storage_key, promoted_object_storage_key) - promoted_metadata = object_storage.head_object(promoted_object_storage_key) - if promoted_metadata is None: - raise ValueError("The uploaded task session is missing") - promoted_size = promoted_metadata.get("ContentLength") - if not isinstance(promoted_size, int): - raise ValueError("Unable to determine the uploaded task session size") - if promoted_size < 1: - raise _TaskSessionUploadRejected("The uploaded task session is not valid JSONL") - if promoted_size > TASK_SESSION_MAX_SIZE_BYTES: - raise _TaskSessionUploadRejected("The uploaded task session exceeds the maximum size") - - try: - content = object_storage.read(promoted_object_storage_key, missing_ok=True) - except UnicodeDecodeError as error: - raise _TaskSessionUploadRejected("The uploaded task session is not valid JSONL") from error - if not content: - raise ValueError("The uploaded task session is missing") - if len(content.encode("utf-8")) > TASK_SESSION_MAX_SIZE_BYTES: - raise _TaskSessionUploadRejected("The uploaded task session exceeds the maximum size") - try: - entries = [json.loads(line) for line in content.splitlines() if line.strip()] - except json.JSONDecodeError as error: - raise _TaskSessionUploadRejected("The uploaded task session is not valid JSONL") from error - if not entries or any(not isinstance(entry, dict) for entry in entries): - raise _TaskSessionUploadRejected("The uploaded task session is not valid JSONL") - if entries[0].get("type") != "session": - raise _TaskSessionUploadRejected("The uploaded task session has no session header") - - previous_key = locked_session.object_storage_key - locked_session.object_storage_key = promoted_object_storage_key - locked_session.revision += 1 - locked_session.pending_sync_id = None - locked_session.pending_object_storage_key = None - locked_session.save( - update_fields=[ - "object_storage_key", - "revision", - "pending_sync_id", - "pending_object_storage_key", - "updated_at", - ] - ) - transaction.on_commit( - lambda: _delete_task_session_objects(locked_session.id, (previous_key, pending_object_storage_key)) - ) - locked_session.mark_synced() - return locked_session.id, locked_session.revision - except _TaskSessionUploadRejected: - if pending_key is None or promoted_key is None: - raise - _reject_task_session_sync(task_session_id, sync_id, pending_key, promoted_key) - raise + task_session = TaskSession.objects.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: - if stale_upload_key is not None: - _delete_task_session_objects(task_session_id, (stale_upload_key,)) - if promoted_key is not None: - finalized_session = TaskSession.objects.get(id=task_session_id) - if ( - finalized_session.revision == expected_revision + 1 - and finalized_session.object_storage_key == promoted_key - ): - return finalized_session.id, finalized_session.revision + _delete_task_session_object(visible_run.active_task_session_id, object_storage_key) raise @@ -3448,7 +3328,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, @@ -3457,6 +3343,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. @@ -3471,6 +3358,16 @@ def _trigger_task_processing_workflow( 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()), + ) + if initial_message or initial_artifact_ids + else None + ), ) logger.info("Workflow trigger completed for task %s, run %s", task.id, run.id) except Exception as e: @@ -3527,10 +3424,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: @@ -3538,7 +3436,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 = { @@ -4621,6 +4528,7 @@ def run_task( 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) @@ -4697,9 +4605,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: @@ -4719,8 +4627,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. @@ -4757,7 +4666,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, @@ -4767,7 +4676,11 @@ def run_task( "provider": provider, "model": model, "reasoning_effort": reasoning_effort, - }.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 @@ -4865,6 +4778,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"] @@ -4886,7 +4802,19 @@ 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) + initial_message = None + 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 if is_pi_task else None, + raise_on_error=False, + ) return contracts.TaskRunResult(task=get_task_detail(task.id, team_id, user_id)) diff --git a/products/tasks/backend/logic/services/sandbox_usage.py b/products/tasks/backend/logic/services/sandbox_usage.py index 277f0d832378..3a4d0478e9a5 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,60 @@ 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() + .select_related("task") + .only("id", "team_id", "state", "task__origin_product") + .get(id=run_id) + ) + SandboxSession.objects.for_team(run.team_id).filter( + task_run_id=run.id, + ended_at__isnull=True, + ).exclude(sandbox_id=sandbox_id).update( + ended_at=timezone.now(), + ended_reason=SandboxSession.EndedReason.CLEANUP, + ) + + 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 +106,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/migrations/0073_task_session.py b/products/tasks/backend/migrations/0073_task_session.py index 81b376cf4cad..0addc4340261 100644 --- a/products/tasks/backend/migrations/0073_task_session.py +++ b/products/tasks/backend/migrations/0073_task_session.py @@ -26,7 +26,12 @@ class Migration(migrations.Migration): serialize=False, ), ), - ("object_storage_key", models.CharField(max_length=512, unique=True)), + ( + "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)), ( @@ -39,6 +44,15 @@ class Migration(migrations.Migration): to="posthog.organization", ), ), + ( + "team", + models.ForeignKey( + db_constraint=False, + on_delete=django.db.models.deletion.CASCADE, + related_name="+", + to="posthog.team", + ), + ), ( "task", models.ForeignKey( @@ -71,6 +85,10 @@ class Migration(migrations.Migration): 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/0074_sandbox_session_open_constraint.py b/products/tasks/backend/migrations/0074_sandbox_session_open_constraint.py new file mode 100644 index 000000000000..fefd0d04157f --- /dev/null +++ b/products/tasks/backend/migrations/0074_sandbox_session_open_constraint.py @@ -0,0 +1,16 @@ +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [("tasks", "0073_task_session")] + + operations = [ + migrations.AddConstraint( + model_name="sandboxsession", + constraint=models.UniqueConstraint( + condition=models.Q(ended_at__isnull=True), + fields=("task_run",), + name="sandbox_session_one_open_per_run", + ), + ), + ] diff --git a/products/tasks/backend/migrations/0074_task_session_revisions.py b/products/tasks/backend/migrations/0074_task_session_revisions.py deleted file mode 100644 index 235932867b6d..000000000000 --- a/products/tasks/backend/migrations/0074_task_session_revisions.py +++ /dev/null @@ -1,23 +0,0 @@ -from django.db import migrations, models - - -class Migration(migrations.Migration): - dependencies = [("tasks", "0073_task_session")] - - operations = [ - migrations.AddField( - model_name="tasksession", - name="revision", - field=models.PositiveBigIntegerField(default=0), - ), - migrations.AddField( - model_name="tasksession", - name="pending_sync_id", - field=models.UUIDField(blank=True, null=True), - ), - migrations.AddField( - model_name="tasksession", - name="pending_object_storage_key", - field=models.CharField(blank=True, max_length=512, null=True, unique=True), - ), - ] diff --git a/products/tasks/backend/migrations/max_migration.txt b/products/tasks/backend/migrations/max_migration.txt index 26c8a0ed1326..db2cbb9c071b 100644 --- a/products/tasks/backend/migrations/max_migration.txt +++ b/products/tasks/backend/migrations/max_migration.txt @@ -1 +1 @@ -0074_task_session_revisions +0074_sandbox_session_open_constraint diff --git a/products/tasks/backend/models.py b/products/tasks/backend/models.py index 06cce3952f1f..cb25f98f0695 100644 --- a/products/tasks/backend/models.py +++ b/products/tasks/backend/models.py @@ -386,7 +386,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. @@ -874,11 +874,16 @@ class TaskSession(models.Model): db_constraint=False, db_index=False, ) + team = models.ForeignKey( + "posthog.Team", + on_delete=models.CASCADE, + related_name="+", + db_constraint=False, + ) task = models.ForeignKey(Task, on_delete=models.CASCADE, related_name="task_sessions", db_index=False) - object_storage_key = models.CharField(max_length=512, unique=True) - revision = models.PositiveBigIntegerField(default=0) - pending_sync_id = models.UUIDField(null=True, blank=True) - pending_object_storage_key = models.CharField(max_length=512, null=True, blank=True, unique=True) + 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) @@ -886,46 +891,33 @@ 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": - session_id = uuid.uuid4() - object_storage_key = f"task-sessions/{task.team.organization_id}/{task.id}/{session_id}.jsonl" return cls.objects.create( - id=session_id, organization_id=task.team.organization_id, + team_id=task.team_id, task=task, - object_storage_key=object_storage_key, ) 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 write_jsonl(self, content: str | bytes) -> None: - object_storage.write(self.object_storage_key, content) - self._tag_object() - - def mark_synced(self) -> None: - self._tag_object() - self.save(update_fields=["updated_at"]) - - def append_entries(self, entries: list[dict[str, Any]]) -> None: - if not entries: + def tag_object(self) -> None: + if self.object_storage_key is None: return - - is_new_object = append_jsonl_object(self.object_storage_key, entries) - if is_new_object: - self._tag_object() - - def _tag_object(self) -> None: 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), }, ) @@ -2187,6 +2179,13 @@ class EndedReason(models.TextChoices): class Meta: db_table = "posthog_task_sandbox_session" + constraints = [ + models.UniqueConstraint( + fields=["task_run"], + condition=models.Q(ended_at__isnull=True), + name="sandbox_session_one_open_per_run", + ), + ] indexes = [ # The usage report scans sessions overlapping the period instance-wide: # closed recently (ended_at > begin) or open and not yet past their TTL @@ -2661,22 +2660,23 @@ def __str__(self): @receiver(post_delete, sender=TaskSession) def delete_task_session_object(sender: type[TaskSession], instance: TaskSession, **kwargs: Any) -> None: - object_storage_keys = {key for key in (instance.object_storage_key, instance.pending_object_storage_key) if key} + if instance.object_storage_key is None: + return + object_storage_key = instance.object_storage_key task_session_id = str(instance.id) - def delete_objects() -> None: - for object_storage_key in object_storage_keys: - 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), - ) + 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_objects) + transaction.on_commit(delete_object) @receiver(post_save, sender=TaskRun) diff --git a/products/tasks/backend/presentation/serializers.py b/products/tasks/backend/presentation/serializers.py index 3f69d4368e76..255a824fb2a4 100644 --- a/products/tasks/backend/presentation/serializers.py +++ b/products/tasks/backend/presentation/serializers.py @@ -729,30 +729,13 @@ def validate_entries(self, value): class TaskSessionResponseSerializer(serializers.Serializer): id = serializers.UUIDField() - download_url = serializers.URLField() - revision = serializers.IntegerField(min_value=0) - - -class TaskSessionSyncPrepareSerializer(serializers.Serializer): - sandbox_id = serializers.CharField() - expected_revision = serializers.IntegerField(min_value=0) - - -class TaskSessionSyncPrepareResponseSerializer(serializers.Serializer): - id = serializers.UUIDField() - sync_id = serializers.UUIDField() - upload = serializers.DictField() - - -class TaskSessionSyncSerializer(serializers.Serializer): - sandbox_id = serializers.CharField() - sync_id = serializers.UUIDField() - expected_revision = serializers.IntegerField(min_value=0) + download_url = serializers.URLField(allow_null=True) + content_sha256 = serializers.CharField(allow_null=True) class TaskSessionSyncResponseSerializer(serializers.Serializer): id = serializers.UUIDField() - revision = serializers.IntegerField(min_value=1) + content_sha256 = serializers.CharField() class TaskRunRelayMessageResponseSerializer(serializers.Serializer): diff --git a/products/tasks/backend/presentation/views/api.py b/products/tasks/backend/presentation/views/api.py index 82a8c12526b7..9c066acf7cef 100644 --- a/products/tasks/backend/presentation/views/api.py +++ b/products/tasks/backend/presentation/views/api.py @@ -108,10 +108,7 @@ TaskRunUpdateSerializer, TaskSerializer, TaskSessionResponseSerializer, - TaskSessionSyncPrepareResponseSerializer, - TaskSessionSyncPrepareSerializer, TaskSessionSyncResponseSerializer, - TaskSessionSyncSerializer, TaskStagedArtifactsFinalizeUploadRequestSerializer, TaskStagedArtifactsFinalizeUploadResponseSerializer, TaskStagedArtifactsPrepareUploadRequestSerializer, @@ -1171,7 +1168,7 @@ def append_log(self, request, pk=None, **kwargs): 200: TaskSessionResponseSerializer, 404: OpenApiResponse(description="Task session not found"), }, - summary="Prepare active task session storage access", + 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): @@ -1179,80 +1176,62 @@ def task_session(self, request, pk=None, **kwargs): result = tasks_facade.get_task_run_session(pk, task_id, self.team_id) if result is None: raise NotFound() - session_id, download_url, revision = result + session_id, download_url, content_sha256 = result return Response( - TaskSessionResponseSerializer({"id": session_id, "download_url": download_url, "revision": revision}).data + TaskSessionResponseSerializer( + { + "id": session_id, + "download_url": download_url, + "content_sha256": content_sha256, + } + ).data ) - @validated_request( - request_serializer=TaskSessionSyncPrepareSerializer, + @extend_schema( + request=OpenApiTypes.BINARY, responses={ - 200: OpenApiResponse(response=TaskSessionSyncPrepareResponseSerializer), + 200: TaskSessionSyncResponseSerializer, 404: OpenApiResponse(description="Task session not found"), 409: OpenApiResponse(response=TaskRunErrorResponseSerializer), }, - summary="Prepare an active task session sync", - strict_request_validation=True, + summary="Replace the active native task session", ) @action( detail=True, methods=["post"], - url_path="task_session_sync_prepare", + url_path="task_session_sync", required_scopes=["task:write"], ) - def prepare_task_session_sync(self, request, pk=None, **kwargs): - task_id = self._ensure_task_accessible() - try: - result = tasks_facade.prepare_task_run_session_sync( - pk, - task_id, - self.team_id, - sandbox_id=request.validated_data["sandbox_id"], - expected_revision=request.validated_data["expected_revision"], - ) - except ValueError as error: - return Response( - TaskRunErrorResponseSerializer({"error": str(error)}).data, - status=status.HTTP_409_CONFLICT, - ) - if result is None: - raise NotFound() - session_id, sync_id, upload = result - return Response( - TaskSessionSyncPrepareResponseSerializer({"id": session_id, "sync_id": sync_id, "upload": upload}).data - ) - - @validated_request( - request_serializer=TaskSessionSyncSerializer, - responses={ - 200: OpenApiResponse(response=TaskSessionSyncResponseSerializer), - 404: OpenApiResponse(description="Task session not found"), - 409: OpenApiResponse(response=TaskRunErrorResponseSerializer), - }, - summary="Finalize an active task session sync", - strict_request_validation=True, - ) - @action(detail=True, methods=["post"], url_path="task_session_sync", required_scopes=["task:write"]) 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."}) + 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.finalize_task_run_session_sync( + result = tasks_facade.sync_task_run_session( pk, task_id, self.team_id, - sandbox_id=request.validated_data["sandbox_id"], - sync_id=request.validated_data["sync_id"], - expected_revision=request.validated_data["expected_revision"], + sandbox_id=sandbox_id, + expected_content_sha256=expected_content_sha256, + content=request.body, ) - except (json.JSONDecodeError, ValueError) as error: + except ValueError as error: return Response( TaskRunErrorResponseSerializer({"error": str(error)}).data, status=status.HTTP_409_CONFLICT, ) if result is None: raise NotFound() - session_id, revision = result - return Response(TaskSessionSyncResponseSerializer({"id": session_id, "revision": revision}).data) + session_id, content_sha256 = result + return Response(TaskSessionSyncResponseSerializer({"id": session_id, "content_sha256": content_sha256}).data) @validated_request( request_serializer=TaskRunRelayMessageRequestSerializer, 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..a52bfb235204 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,25 @@ 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: + sandbox.destroy() + 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 56683cc62580..e670639e37c2 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 @@ -109,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 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..849b9af7697f 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,17 @@ 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() 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/send_followup_to_sandbox.py b/products/tasks/backend/temporal/process_task/activities/send_followup_to_sandbox.py index 45fbf91cbeb2..7f10e3ca1e42 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. 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 9e0bae765d9b..54dc585c7ad0 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 @@ -994,30 +994,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): + 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( - extra_state={ - "runtime_adapter": "codex", - "provider": "openai", - "model": "gpt-5.3-codex", - "reasoning_effort": "high", - "initial_permission_mode": "plan", - } - ) + 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) task_run.refresh_from_db() assert task_run.active_task_session is not None - assert task_run.active_task_session.object_storage_key.startswith( - f"task-sessions/{test_task.team.organization_id}/{test_task.id}/" - ) + 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 == "codex" - assert result.provider == "openai" - assert result.model == "gpt-5.3-codex" - assert result.reasoning_effort == "high" - assert result.initial_permission_mode == "plan" + 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/workflow.py b/products/tasks/backend/temporal/process_task/workflow.py index 99ac5ffbafc2..364fd2e42a23 100644 --- a/products/tasks/backend/temporal/process_task/workflow.py +++ b/products/tasks/backend/temporal/process_task/workflow.py @@ -131,6 +131,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 +141,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. @@ -329,6 +328,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 +729,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. @@ -1120,12 +1127,13 @@ def _restore_resumed_state(self, resumed: ResumedSandboxState) -> None: 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( @@ -2097,15 +2105,10 @@ async def _send_followup_to_sandbox( steer=steer, ), 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=(1 if self.context.task_runtime == "pi" else SEND_FOLLOWUP_MAX_ATTEMPTS), ), ) except Exception as e: diff --git a/products/tasks/backend/tests/test_api.py b/products/tasks/backend/tests/test_api.py index 5fa904ee20a0..8224005f4ae6 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 @@ -30,7 +31,6 @@ from posthog.storage import object_storage from products.slack_app.backend.models import SlackThreadTaskMapping -from products.tasks.backend.constants import TASK_SESSION_MAX_SIZE_BYTES, TASK_SESSION_UPLOAD_FORM_OVERHEAD_BYTES from products.tasks.backend.facade import api as tasks_facade from products.tasks.backend.logic.services.code_usage_gate import ( CodeUsageStatus, @@ -60,6 +60,7 @@ CodeInviteRedemption, SandboxCustomImage, SandboxEnvironment, + SandboxSession, Task, TaskArtifact, TaskAutomation, @@ -1641,13 +1642,17 @@ def test_run_endpoint_starts_pi_task(self, mock_workflow): self.assertEqual(response.status_code, status.HTTP_200_OK) run = task.runs.get() - mock_workflow.assert_called_once_with( - task_id=str(task.id), - run_id=str(run.id), - team_id=task.team.id, - user_id=self.user.id, - posthog_mcp_scopes="full", - ) + 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) @parameterized.expand( [ @@ -2265,13 +2270,13 @@ def test_start_run_endpoint_starts_pi_task(self, mock_workflow): response = self.client.post(f"/api/projects/@current/tasks/{task.id}/runs/{task_run.id}/start/") self.assertEqual(response.status_code, status.HTTP_200_OK) - mock_workflow.assert_called_once_with( - task_id=str(task.id), - run_id=str(task_run.id), - team_id=task.team.id, - user_id=self.user.id, - posthog_mcp_scopes="full", - ) + 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_rejects_missing_run_artifacts(self, mock_workflow): @@ -7950,6 +7955,19 @@ 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() + return SandboxSession.objects.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 @@ -8152,290 +8170,129 @@ def test_command_proxies_cancel(self, mock_post): self.assertEqual(response.status_code, status.HTTP_200_OK) self.assertTrue(response.json()["result"]["cancelled"]) - @patch("posthog.storage.object_storage.get_presigned_url") - def test_task_session_returns_read_only_storage_access(self, mock_download_url): + 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"]) - 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)) - self.assertEqual(response.json()["download_url"], "https://storage.example/session.jsonl") - self.assertEqual(response.json()["revision"], 0) - self.assertNotIn("upload", response.json()) + self.assertEqual( + response.json(), + { + "id": str(task_session.id), + "download_url": None, + "content_sha256": None, + }, + ) - @patch("posthog.storage.object_storage.get_presigned_post") - @patch("posthog.storage.object_storage.head_object") - @patch("posthog.storage.object_storage.copy") - @patch("posthog.storage.object_storage.read") @patch("products.tasks.backend.models.object_storage.tag") - def test_task_session_sync_validates_and_promotes_an_immutable_revision( - self, mock_tag, mock_read, mock_copy, mock_head, mock_upload - ): + @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.state = {**run.state, "sandbox_id": "sandbox-1"} - run.save(update_fields=["active_task_session", "state"]) - mock_upload.return_value = { - "url": "https://storage.example/upload", - "fields": {"key": "pending-session.jsonl"}, - } - mock_head.return_value = {"ContentLength": 31} - mock_read.return_value = '{"type":"session","version":3}\n' - - prepare_response = self.client.post( - f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/task_session_sync_prepare/", - {"sandbox_id": "sandbox-1", "expected_revision": 0}, - format="json", - ) - self.assertEqual(prepare_response.status_code, status.HTTP_200_OK) - task_session.refresh_from_db() - pending_key = task_session.pending_object_storage_key - self.assertIsNotNone(pending_key) - mock_upload.assert_called_once_with( - pending_key, - conditions=[ - ["content-length-range", 1, TASK_SESSION_MAX_SIZE_BYTES + TASK_SESSION_UPLOAD_FORM_OVERHEAD_BYTES] - ], - ) + run.save(update_fields=["active_task_session"]) + self._open_sandbox_session(run) + content = b"opaque native session content" - response = self.client.post( + response = self.client.generic( + "POST", f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/task_session_sync/", - { - "sandbox_id": "sandbox-1", - "sync_id": prepare_response.json()["sync_id"], - "expected_revision": 0, - }, - format="json", + content, + content_type="application/octet-stream", + HTTP_IF_MATCH='"none"', + HTTP_X_SANDBOX_ID="sandbox-1", ) + expected_sha256 = hashlib.sha256(content).hexdigest() self.assertEqual(response.status_code, status.HTTP_200_OK) - self.assertEqual(response.json()["id"], str(task_session.id)) - self.assertEqual(response.json()["revision"], 1) + self.assertEqual(response.json(), {"id": str(task_session.id), "content_sha256": expected_sha256}) task_session.refresh_from_db() - self.assertEqual(task_session.revision, 1) - self.assertIsNone(task_session.pending_sync_id) - self.assertNotEqual(task_session.object_storage_key, pending_key) - mock_copy.assert_called_once_with(pending_key, task_session.object_storage_key) - mock_read.assert_called_once_with(task_session.object_storage_key, missing_ok=True) + 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.get_presigned_post") - @patch("posthog.storage.object_storage.head_object") - @patch("posthog.storage.object_storage.copy") - @patch("posthog.storage.object_storage.read") @patch("posthog.storage.object_storage.delete") - def test_task_session_sync_rejects_an_oversized_upload( - self, mock_delete, mock_read, mock_copy, mock_head, mock_upload - ): + @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.state = {**run.state, "sandbox_id": "sandbox-1"} - run.save(update_fields=["active_task_session", "state"]) - mock_upload.return_value = { - "url": "https://storage.example/upload", - "fields": {"key": "pending-session.jsonl"}, - } - mock_head.return_value = {"ContentLength": TASK_SESSION_MAX_SIZE_BYTES + 1} + run.save(update_fields=["active_task_session"]) + self._open_sandbox_session(run) - prepare_response = self.client.post( - f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/task_session_sync_prepare/", - {"sandbox_id": "sandbox-1", "expected_revision": 0}, - format="json", + response = self.client.generic( + "POST", + f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/task_session_sync/", + b'{"type":"session","version":3}\n', + content_type="application/octet-stream", + HTTP_IF_MATCH='"stale-hash"', + HTTP_X_SANDBOX_ID="sandbox-1", ) - task_session.refresh_from_db() - pending_key = task_session.pending_object_storage_key - sync_id = prepare_response.json()["sync_id"] - with self.captureOnCommitCallbacks(execute=True): - response = self.client.post( - f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/task_session_sync/", - { - "sandbox_id": "sandbox-1", - "sync_id": sync_id, - "expected_revision": 0, - }, - format="json", - ) self.assertEqual(response.status_code, status.HTTP_409_CONFLICT) - self.assertEqual(response.json()["error"], "The uploaded task session exceeds the maximum size") + self.assertEqual(response.json()["error"], "The task session content is stale") task_session.refresh_from_db() - self.assertIsNone(task_session.pending_sync_id) - self.assertIsNone(task_session.pending_object_storage_key) - mock_copy.assert_not_called() - mock_read.assert_not_called() - self.assertIn(call(pending_key), mock_delete.call_args_list) + 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.get_presigned_post") - @patch("posthog.storage.object_storage.head_object", return_value={"ContentLength": 31}) - @patch("posthog.storage.object_storage.copy") - @patch("posthog.storage.object_storage.read", return_value="not-json\n") - @patch("posthog.storage.object_storage.delete") - def test_task_session_sync_rejects_invalid_jsonl_and_clears_the_pending_upload( - self, mock_delete, mock_read, mock_copy, mock_head, mock_upload - ): + @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.state = {**run.state, "sandbox_id": "sandbox-1"} - run.save(update_fields=["active_task_session", "state"]) - mock_upload.return_value = { - "url": "https://storage.example/upload", - "fields": {"key": "pending-session.jsonl"}, - } + run.save(update_fields=["active_task_session"]) + self._open_sandbox_session(run, "active-sandbox") - prepare_response = self.client.post( - f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/task_session_sync_prepare/", - {"sandbox_id": "sandbox-1", "expected_revision": 0}, - format="json", - ) - task_session.refresh_from_db() - pending_key = task_session.pending_object_storage_key - sync_id = prepare_response.json()["sync_id"] - promoted_key = ( - f"task-sessions/{task_session.organization_id}/{task_session.task_id}/{task_session.id}/" - f"revisions/1-{sync_id}.jsonl" + response = self.client.generic( + "POST", + f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/task_session_sync/", + b'{"type":"session"}\n', + content_type="application/octet-stream", + HTTP_IF_MATCH='"none"', + HTTP_X_SANDBOX_ID="stale-sandbox", ) - with self.captureOnCommitCallbacks(execute=True): - response = self.client.post( - f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/task_session_sync/", - {"sandbox_id": "sandbox-1", "sync_id": sync_id, "expected_revision": 0}, - format="json", - ) - self.assertEqual(response.status_code, status.HTTP_409_CONFLICT) - task_session.refresh_from_db() - self.assertIsNone(task_session.pending_sync_id) - self.assertIsNone(task_session.pending_object_storage_key) - self.assertCountEqual(mock_delete.call_args_list, [call(pending_key), call(promoted_key)]) + self.assertEqual(response.json()["error"], "The task session writer is not the active sandbox") + mock_write.assert_not_called() - @patch("posthog.storage.object_storage.get_presigned_post") - @patch("posthog.storage.object_storage.head_object") - @patch("posthog.storage.object_storage.delete") - def test_task_session_sync_preserves_pending_upload_after_storage_error(self, mock_delete, mock_head, mock_upload): + @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.state = {**run.state, "sandbox_id": "sandbox-1"} - run.save(update_fields=["active_task_session", "state"]) - mock_upload.return_value = { - "url": "https://storage.example/upload", - "fields": {"key": "pending-session.jsonl"}, - } - - prepare_response = self.client.post( - f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/task_session_sync_prepare/", - {"sandbox_id": "sandbox-1", "expected_revision": 0}, - format="json", - ) - task_session.refresh_from_db() - pending_sync_id = task_session.pending_sync_id - pending_key = task_session.pending_object_storage_key - mock_head.return_value = None + run.save(update_fields=["active_task_session"]) + self._open_sandbox_session(run) - response = self.client.post( + response = self.client.generic( + "POST", f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/task_session_sync/", - { - "sandbox_id": "sandbox-1", - "sync_id": prepare_response.json()["sync_id"], - "expected_revision": 0, - }, - format="json", + b"", + content_type="application/octet-stream", + HTTP_IF_MATCH='"none"', + HTTP_X_SANDBOX_ID="sandbox-1", ) self.assertEqual(response.status_code, status.HTTP_409_CONFLICT) - task_session.refresh_from_db() - self.assertEqual(task_session.pending_sync_id, pending_sync_id) - self.assertEqual(task_session.pending_object_storage_key, pending_key) - mock_delete.assert_not_called() - - @patch("posthog.storage.object_storage.get_presigned_post") - def test_repeated_prepare_reuses_the_pending_upload(self, mock_upload): - 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.state = {**run.state, "sandbox_id": "sandbox-1"} - run.save(update_fields=["active_task_session", "state"]) - mock_upload.return_value = { - "url": "https://storage.example/upload", - "fields": {"key": "pending-session.jsonl"}, - } - prepare_url = f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/task_session_sync_prepare/" - payload = {"sandbox_id": "sandbox-1", "expected_revision": 0} - - first_prepare = self.client.post(prepare_url, payload, format="json") - task_session.refresh_from_db() - pending_key = task_session.pending_object_storage_key - - second_prepare = self.client.post(prepare_url, payload, format="json") - - self.assertEqual(first_prepare.status_code, status.HTTP_200_OK) - self.assertEqual(second_prepare.status_code, status.HTTP_200_OK) - self.assertEqual(second_prepare.json()["sync_id"], first_prepare.json()["sync_id"]) - task_session.refresh_from_db() - self.assertEqual(task_session.pending_object_storage_key, pending_key) - self.assertEqual(mock_upload.call_count, 2) - self.assertEqual(mock_upload.call_args_list[0], mock_upload.call_args_list[1]) - - @patch("posthog.storage.object_storage.get_presigned_post") - @patch("posthog.storage.object_storage.head_object", return_value={"ContentLength": 31}) - @patch("posthog.storage.object_storage.copy") - @patch("posthog.storage.object_storage.read", return_value='{"type":"session","version":3}\n') - @patch("posthog.storage.object_storage.delete") - def test_serialized_duplicate_finalize_recovers_without_rewriting_the_promoted_revision( - self, mock_delete, mock_read, mock_copy, mock_head, mock_upload - ): - 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.state = {**run.state, "sandbox_id": "sandbox-1"} - run.save(update_fields=["active_task_session", "state"]) - mock_upload.return_value = { - "url": "https://storage.example/upload", - "fields": {"key": "pending-session.jsonl"}, - } - - prepare_response = self.client.post( - f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/task_session_sync_prepare/", - {"sandbox_id": "sandbox-1", "expected_revision": 0}, - format="json", - ) - sync_id = prepare_response.json()["sync_id"] - finalize_url = f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/task_session_sync/" - finalize_payload = {"sandbox_id": "sandbox-1", "sync_id": sync_id, "expected_revision": 0} - - first_response = self.client.post(finalize_url, finalize_payload, format="json") - self.assertEqual(first_response.status_code, status.HTTP_200_OK) - - task_session.refresh_from_db() - promoted_key = task_session.object_storage_key - mock_head.reset_mock() - mock_copy.reset_mock() - mock_read.reset_mock() - - duplicate_response = self.client.post(finalize_url, finalize_payload, format="json") - - self.assertEqual(duplicate_response.status_code, status.HTTP_200_OK) - self.assertEqual(duplicate_response.json()["revision"], 1) - task_session.refresh_from_db() - self.assertEqual(task_session.object_storage_key, promoted_key) - mock_head.assert_not_called() - mock_copy.assert_not_called() - mock_read.assert_not_called() - mock_delete.assert_not_called() + self.assertEqual(response.json()["error"], "The task session content size is invalid") + mock_write.assert_not_called() @patch("posthog.storage.object_storage.get_presigned_url") def test_task_session_is_readable_for_a_public_channel_task(self, mock_download_url): diff --git a/products/tasks/backend/tests/test_storage.py b/products/tasks/backend/tests/test_storage.py index 7090d078ce51..e8b093fa28b1 100644 --- a/products/tasks/backend/tests/test_storage.py +++ b/products/tasks/backend/tests/test_storage.py @@ -1,12 +1,9 @@ -import uuid - from unittest.mock import MagicMock, patch from django.test import SimpleTestCase from parameterized import parameterized -from products.tasks.backend.models import TaskSession from products.tasks.backend.storage import append_jsonl_object @@ -33,36 +30,3 @@ def test_appends_complete_json_lines( self.assertEqual(is_new, expected_is_new) mock_write.assert_called_once_with("sessions/example.jsonl", expected_content) - - -class TestTaskSessionStorage(SimpleTestCase): - @patch("products.tasks.backend.models.object_storage.tag") - @patch("products.tasks.backend.storage.object_storage.write") - @patch("products.tasks.backend.storage.object_storage.read") - def test_new_session_object_is_attributed_without_expiry( - self, - mock_read: MagicMock, - mock_write: MagicMock, - mock_tag: MagicMock, - ) -> None: - organization_id = uuid.uuid4() - task_id = uuid.uuid4() - task_session = TaskSession( - id=uuid.uuid4(), - organization_id=organization_id, - task_id=task_id, - object_storage_key="sessions/example.jsonl", - ) - mock_read.return_value = None - - task_session.append_entries([{"type": "session"}]) - - mock_write.assert_called_once_with("sessions/example.jsonl", '{"type": "session"}') - mock_tag.assert_called_once_with( - "sessions/example.jsonl", - { - "data_class": "task_session", - "organization_id": str(organization_id), - "task_id": str(task_id), - }, - ) From 765ebbd8e5538fdf67ba08c5d65a97bb8708bf00 Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Fri, 24 Jul 2026 19:16:01 +0200 Subject: [PATCH 06/22] fix(tasks): preserve Pi command identity --- .../tasks/backend/presentation/serializers.py | 5 ++++ .../tasks/backend/presentation/views/api.py | 1 + .../process_task/tests/test_workflow.py | 24 +++++++++++++++++++ .../backend/temporal/process_task/workflow.py | 12 ++++++++++ products/tasks/backend/tests/test_api.py | 4 ++-- 5 files changed, 44 insertions(+), 2 deletions(-) diff --git a/products/tasks/backend/presentation/serializers.py b/products/tasks/backend/presentation/serializers.py index 255a824fb2a4..2fa39d47b134 100644 --- a/products/tasks/backend/presentation/serializers.py +++ b/products/tasks/backend/presentation/serializers.py @@ -2371,6 +2371,11 @@ def validate(self, attrs): 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") diff --git a/products/tasks/backend/presentation/views/api.py b/products/tasks/backend/presentation/views/api.py index 9c066acf7cef..67ac47fbb105 100644 --- a/products/tasks/backend/presentation/views/api.py +++ b/products/tasks/backend/presentation/views/api.py @@ -1628,6 +1628,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: 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..4378f157fc29 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,30 @@ 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_ids_remain_deduplicated_for_the_run(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-0") + + assert len(workflow._pending_followups) == 501 + assert workflow._pending_followups[0].message == "message 0" + 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 364fd2e42a23..a0b613b28410 100644 --- a/products/tasks/backend/temporal/process_task/workflow.py +++ b/products/tasks/backend/temporal/process_task/workflow.py @@ -122,6 +122,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 @@ -286,6 +287,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] = [] @@ -1114,6 +1117,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, ), ) @@ -1124,6 +1128,8 @@ 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: @@ -2009,6 +2015,12 @@ def _queue_followup_message( extra={"run_id": context.run_id if context is not None else None}, ) return + if message_id and message_id in self._accepted_message_id_set: + return + if message_id: + self._accepted_message_ids.append(message_id) + self._accepted_message_id_set.add(message_id) + pending_followup = PendingFollowup( message=message, artifact_ids=artifact_ids or [], diff --git a/products/tasks/backend/tests/test_api.py b/products/tasks/backend/tests/test_api.py index 8224005f4ae6..562b90339578 100644 --- a/products/tasks/backend/tests/test_api.py +++ b/products/tasks/backend/tests/test_api.py @@ -8325,7 +8325,7 @@ def test_command_signals_pi_user_message(self, mock_signal_followup): ) self.assertEqual(response.status_code, status.HTTP_200_OK) - mock_signal_followup.assert_called_once() + 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") @@ -8346,7 +8346,7 @@ def test_command_proxies_pi_rpc(self, mock_post): { "jsonrpc": "2.0", "method": "pi/rpc", - "params": {"command": {"type": "future_native_command"}}, + "params": {"command": {"id": "native", "type": "future_native_command"}}, "id": "native", }, format="json", From d190d07adc3ec6e6a5fb0b21c833ee2281e04e32 Mon Sep 17 00:00:00 2001 From: Alessandro Pogliaghi Date: Mon, 27 Jul 2026 11:32:03 +0100 Subject: [PATCH 07/22] chore(tasks): refresh generated api types --- .../tasks/frontend/generated/api.schemas.ts | 31 +++----------- products/tasks/frontend/generated/api.ts | 40 +++---------------- products/tasks/frontend/generated/api.zod.ts | 23 ++--------- services/mcp/src/api/generated.ts | 31 +++----------- 4 files changed, 18 insertions(+), 107 deletions(-) diff --git a/products/tasks/frontend/generated/api.schemas.ts b/products/tasks/frontend/generated/api.schemas.ts index d8ead9d350e0..ffe97e58c6cc 100644 --- a/products/tasks/frontend/generated/api.schemas.ts +++ b/products/tasks/frontend/generated/api.schemas.ts @@ -2815,36 +2815,15 @@ export interface StreamReadTokenResponseApi { export interface TaskSessionResponseApi { id: string - download_url: string - /** @minimum 0 */ - revision: number -} - -export interface TaskSessionSyncApi { - sandbox_id: string - sync_id: string - /** @minimum 0 */ - expected_revision: number + /** @nullable */ + download_url: string | null + /** @nullable */ + content_sha256: string | null } export interface TaskSessionSyncResponseApi { id: string - /** @minimum 1 */ - revision: number -} - -export interface TaskSessionSyncPrepareApi { - sandbox_id: string - /** @minimum 0 */ - expected_revision: number -} - -export type TaskSessionSyncPrepareResponseApiUpload = { [key: string]: unknown } - -export interface TaskSessionSyncPrepareResponseApi { - id: string - sync_id: string - upload: TaskSessionSyncPrepareResponseApiUpload + content_sha256: string } /** diff --git a/products/tasks/frontend/generated/api.ts b/products/tasks/frontend/generated/api.ts index 5b9b82dbeeb8..4357cf1ab4a8 100644 --- a/products/tasks/frontend/generated/api.ts +++ b/products/tasks/frontend/generated/api.ts @@ -90,9 +90,6 @@ import type { TaskRunRelayMessageResponseApi, TaskRunStartRequestApi, TaskSessionResponseApi, - TaskSessionSyncApi, - TaskSessionSyncPrepareApi, - TaskSessionSyncPrepareResponseApi, TaskSessionSyncResponseApi, TaskStagedArtifactsFinalizeUploadRequestApi, TaskStagedArtifactsFinalizeUploadResponseApi, @@ -1439,7 +1436,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, and native Pi RPC commands. * @summary Send command to task run */ export const tasksRunsCommandCreate = async ( @@ -1673,7 +1670,7 @@ export const getTasksRunsTaskSessionRetrieveUrl = (projectId: string, taskId: st /** * API for managing task runs. Each run represents an execution of a task. - * @summary Prepare active task session storage access + * @summary Get active task session storage access */ export const tasksRunsTaskSessionRetrieve = async ( projectId: string, @@ -1693,49 +1690,22 @@ export const getTasksRunsTaskSessionSyncCreateUrl = (projectId: string, taskId: /** * API for managing task runs. Each run represents an execution of a task. - * @summary Finalize an active task session sync + * @summary Replace the active native task session */ export const tasksRunsTaskSessionSyncCreate = async ( projectId: string, taskId: string, id: string, - taskSessionSyncApi: TaskSessionSyncApi, + tasksRunsTaskSessionSyncCreateBody?: Blob | string, options?: RequestInit ): Promise => { return apiMutator(getTasksRunsTaskSessionSyncCreateUrl(projectId, taskId, id), { ...options, method: 'POST', - headers: { 'Content-Type': 'application/json', ...options?.headers }, - body: JSON.stringify(taskSessionSyncApi), + body: JSON.stringify(tasksRunsTaskSessionSyncCreateBody), }) } -export const getTasksRunsTaskSessionSyncPrepareCreateUrl = (projectId: string, taskId: string, id: string) => { - return `/api/projects/${projectId}/tasks/${taskId}/runs/${id}/task_session_sync_prepare/` -} - -/** - * API for managing task runs. Each run represents an execution of a task. - * @summary Prepare an active task session sync - */ -export const tasksRunsTaskSessionSyncPrepareCreate = async ( - projectId: string, - taskId: string, - id: string, - taskSessionSyncPrepareApi: TaskSessionSyncPrepareApi, - options?: RequestInit -): Promise => { - return apiMutator( - getTasksRunsTaskSessionSyncPrepareCreateUrl(projectId, taskId, id), - { - ...options, - method: 'POST', - headers: { 'Content-Type': 'application/json', ...options?.headers }, - body: JSON.stringify(taskSessionSyncPrepareApi), - } - ) -} - 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 2a71c3941c1c..e7fdcee2705a 100644 --- a/products/tasks/frontend/generated/api.zod.ts +++ b/products/tasks/frontend/generated/api.zod.ts @@ -2554,7 +2554,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, and native Pi RPC commands. * @summary Send command to task run */ export const TasksRunsCommandCreateBody = /* @__PURE__ */ zod @@ -2642,26 +2642,9 @@ export const TasksRunsStartCreateBody = /* @__PURE__ */ zod.object({ /** * API for managing task runs. Each run represents an execution of a task. - * @summary Finalize an active task session sync + * @summary Replace the active native task session */ -export const tasksRunsTaskSessionSyncCreateBodyExpectedRevisionMin = 0 - -export const TasksRunsTaskSessionSyncCreateBody = /* @__PURE__ */ zod.object({ - sandbox_id: zod.string(), - sync_id: zod.uuid(), - expected_revision: zod.number().min(tasksRunsTaskSessionSyncCreateBodyExpectedRevisionMin), -}) - -/** - * API for managing task runs. Each run represents an execution of a task. - * @summary Prepare an active task session sync - */ -export const tasksRunsTaskSessionSyncPrepareCreateBodyExpectedRevisionMin = 0 - -export const TasksRunsTaskSessionSyncPrepareCreateBody = /* @__PURE__ */ zod.object({ - sandbox_id: zod.string(), - expected_revision: zod.number().min(tasksRunsTaskSessionSyncPrepareCreateBodyExpectedRevisionMin), -}) +export const TasksRunsTaskSessionSyncCreateBody = /* @__PURE__ */ zod.instanceof(File) /** * Create a stable, editable artifact handle from direct markdown/text content or an existing run artifact. Slack adapters deliver into the mapped Slack thread; document artifacts use external connector storage when available. diff --git a/services/mcp/src/api/generated.ts b/services/mcp/src/api/generated.ts index 0d06a48556c1..e97d8bbd994f 100644 --- a/services/mcp/src/api/generated.ts +++ b/services/mcp/src/api/generated.ts @@ -67875,36 +67875,15 @@ export namespace Schemas { export interface TaskSessionResponse { id: string; - download_url: string; - /** @minimum 0 */ - revision: number; - } - - export interface TaskSessionSync { - sandbox_id: string; - sync_id: string; - /** @minimum 0 */ - expected_revision: number; - } - - export interface TaskSessionSyncPrepare { - sandbox_id: string; - /** @minimum 0 */ - expected_revision: number; - } - - export type TaskSessionSyncPrepareResponseUpload = { [key: string]: unknown }; - - export interface TaskSessionSyncPrepareResponse { - id: string; - sync_id: string; - upload: TaskSessionSyncPrepareResponseUpload; + /** @nullable */ + download_url: string | null; + /** @nullable */ + content_sha256: string | null; } export interface TaskSessionSyncResponse { id: string; - /** @minimum 1 */ - revision: number; + content_sha256: string; } export interface TaskStagedArtifactFinalizeUpload { From 061eada0f3fb308d830892a3846071638053415e Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Mon, 27 Jul 2026 13:57:31 +0200 Subject: [PATCH 08/22] fix(tasks): address Pi runtime review feedback --- .../security/idor-team-scoped-models.yaml | 2 + products/tasks/backend/constants.py | 1 + products/tasks/backend/facade/api.py | 4 +- .../backend/logic/services/docker_sandbox.py | 3 + .../backend/logic/services/modal_sandbox.py | 3 + .../tasks/backend/logic/services/sandbox.py | 7 +++ .../logic/services/test_docker_sandbox.py | 21 +++++++ .../backend/migrations/0073_task_session.py | 6 +- products/tasks/backend/models.py | 5 +- .../tasks/backend/presentation/views/api.py | 56 ++++++++++++++++++- .../process_task/tests/test_workflow.py | 17 +++++- .../backend/temporal/process_task/workflow.py | 23 ++++++-- products/tasks/backend/tests/test_api.py | 39 +++++++++---- 13 files changed, 160 insertions(+), 27 deletions(-) diff --git a/.semgrep/rules/security/idor-team-scoped-models.yaml b/.semgrep/rules/security/idor-team-scoped-models.yaml index 77a87128dd31..789488b3ac85 100644 --- a/.semgrep/rules/security/idor-team-scoped-models.yaml +++ b/.semgrep/rules/security/idor-team-scoped-models.yaml @@ -287,6 +287,7 @@ rules: |TaskAutomation |TaskPresence |TaskRun + |TaskSession |TaskThreadMessage |TaskThreadMessageMention |Channel @@ -602,6 +603,7 @@ rules: |TaskAutomation |TaskPresence |TaskRun + |TaskSession |TaskThreadMessage |TaskThreadMessageMention |Channel diff --git a/products/tasks/backend/constants.py b/products/tasks/backend/constants.py index c0a1a3c379da..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" diff --git a/products/tasks/backend/facade/api.py b/products/tasks/backend/facade/api.py index 8ab09d6dbfe6..8af79fdf8809 100644 --- a/products/tasks/backend/facade/api.py +++ b/products/tasks/backend/facade/api.py @@ -2230,7 +2230,7 @@ def get_task_run_session( 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.get(id=run.active_task_session_id) + 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) @@ -2306,7 +2306,7 @@ def sync_task_run_session( 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.select_for_update().get(id=locked_run.active_task_session_id) + 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 diff --git a/products/tasks/backend/logic/services/docker_sandbox.py b/products/tasks/backend/logic/services/docker_sandbox.py index fdb9ef6207d6..33d31522b2f7 100644 --- a/products/tasks/backend/logic/services/docker_sandbox.py +++ b/products/tasks/backend/logic/services/docker_sandbox.py @@ -933,6 +933,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 diff --git a/products/tasks/backend/logic/services/modal_sandbox.py b/products/tasks/backend/logic/services/modal_sandbox.py index 2ca36cc3f661..3f30890960a2 100644 --- a/products/tasks/backend/logic/services/modal_sandbox.py +++ b/products/tasks/backend/logic/services/modal_sandbox.py @@ -1159,6 +1159,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 diff --git a/products/tasks/backend/logic/services/sandbox.py b/products/tasks/backend/logic/services/sandbox.py index 91d5bb533c90..47bcf4bae7f7 100644 --- a/products/tasks/backend/logic/services/sandbox.py +++ b/products/tasks/backend/logic/services/sandbox.py @@ -263,6 +263,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, diff --git a/products/tasks/backend/logic/services/test_docker_sandbox.py b/products/tasks/backend/logic/services/test_docker_sandbox.py index fef1143b2d7e..88f7df2e357c 100644 --- a/products/tasks/backend/logic/services/test_docker_sandbox.py +++ b/products/tasks/backend/logic/services/test_docker_sandbox.py @@ -425,6 +425,7 @@ def test_start_agent_server_preserves_pi_protocol_without_branch_retry(self) -> 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( @@ -443,6 +444,26 @@ def test_start_agent_server_preserves_pi_protocol_without_branch_retry(self) -> 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() == {} diff --git a/products/tasks/backend/migrations/0073_task_session.py b/products/tasks/backend/migrations/0073_task_session.py index 0addc4340261..5d8ba44f4067 100644 --- a/products/tasks/backend/migrations/0073_task_session.py +++ b/products/tasks/backend/migrations/0073_task_session.py @@ -1,11 +1,11 @@ # Generated by Django 5.2.14 on 2026-07-17 13:24 -import uuid - import django.utils.timezone import django.db.models.deletion from django.db import migrations, models +import posthog.models.utils + class Migration(migrations.Migration): dependencies = [ @@ -20,7 +20,7 @@ class Migration(migrations.Migration): ( "id", models.UUIDField( - default=uuid.uuid4, + default=posthog.models.utils.uuid7, editable=False, primary_key=True, serialize=False, diff --git a/products/tasks/backend/models.py b/products/tasks/backend/models.py index cb25f98f0695..6c1a95cb476a 100644 --- a/products/tasks/backend/models.py +++ b/products/tasks/backend/models.py @@ -865,8 +865,7 @@ def _dispatch() -> None: return task -class TaskSession(models.Model): - id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False) +class TaskSession(TeamScopedRootMixin, UUIDModel): organization = models.ForeignKey( "posthog.Organization", on_delete=models.CASCADE, @@ -897,7 +896,7 @@ class Meta: @classmethod def create_for_task(cls, task: Task) -> "TaskSession": - return cls.objects.create( + return cls.objects.unscoped().create( organization_id=task.team.organization_id, team_id=task.team_id, task=task, diff --git a/products/tasks/backend/presentation/views/api.py b/products/tasks/backend/presentation/views/api.py index 67ac47fbb105..c460b98abc3f 100644 --- a/products/tasks/backend/presentation/views/api.py +++ b/products/tasks/backend/presentation/views/api.py @@ -21,6 +21,7 @@ from rest_framework.decorators import action from rest_framework.exceptions import NotFound, 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 @@ -33,6 +34,7 @@ from posthog.rate_limit import CodeInviteThrottle from posthog.renderers import ServerSentEventRenderer +from products.tasks.backend.constants import PI_CLOUD_RUNTIME_FEATURE_FLAG from products.tasks.backend.facade import ( access as tasks_access, api as tasks_facade, @@ -123,8 +125,43 @@ 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): + return stream.read() + + logger = logging.getLogger(__name__) + +def _pi_cloud_runtime_enabled(user: Any, team: Any) -> bool: + distinct_id = getattr(user, "distinct_id", None) or str(getattr(user, "uuid", "")) + organization_id = str(getattr(team, "organization_id", "") or "") + 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 _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 @@ -588,6 +625,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 _pi_cloud_runtime_enabled(request.user, self.team): + 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 @@ -940,6 +981,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 _pi_cloud_runtime_enabled(request.user, self.team): + return _pi_cloud_runtime_disabled_response() if (limit_response := cloud_usage_limit_response(request.user, self.team_id)) is not None: return limit_response @@ -968,6 +1013,10 @@ 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() + if tasks_facade.task_runtime( + task_id, self.team_id, self._user_id(), for_control=True + ) == tasks_facade.TaskRuntime.PI and not _pi_cloud_runtime_enabled(request.user, self.team): + return _pi_cloud_runtime_disabled_response() startable = tasks_facade.check_task_run_startable(pk, task_id, self.team_id) if startable == "not_found": @@ -1201,6 +1250,7 @@ def task_session(self, request, pk=None, **kwargs): 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() @@ -1221,7 +1271,7 @@ def sync_task_session(self, request, pk=None, **kwargs): self.team_id, sandbox_id=sandbox_id, expected_content_sha256=expected_content_sha256, - content=request.body, + content=request.data, ) except ValueError as error: return Response( @@ -1939,6 +1989,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 _pi_cloud_runtime_enabled(request.user, self.team): + 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: 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 4378f157fc29..6f552b702ed3 100644 --- a/products/tasks/backend/temporal/process_task/tests/test_workflow.py +++ b/products/tasks/backend/temporal/process_task/tests/test_workflow.py @@ -388,7 +388,7 @@ async def test_duplicate_sender_message_id_is_queued_once(self, monkeypatch): assert [(item.message, item.message_id) for item in workflow._pending_followups] == [("first", "message-123")] - async def test_sender_message_ids_remain_deduplicated_for_the_run(self, monkeypatch): + 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()) @@ -396,10 +396,21 @@ async def test_sender_message_ids_remain_deduplicated_for_the_run(self, monkeypa for index in range(501): await workflow.send_followup_message(f"message {index}", [], f"message-{index}") - await workflow.send_followup_message("duplicate", [], "message-0") + await workflow.send_followup_message("duplicate", [], "message-1") + assert len(workflow._accepted_message_ids) == 500 assert len(workflow._pending_followups) == 501 - assert workflow._pending_followups[0].message == "message 0" + + 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() diff --git a/products/tasks/backend/temporal/process_task/workflow.py b/products/tasks/backend/temporal/process_task/workflow.py index a0b613b28410..20b198197126 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.""" @@ -2015,11 +2026,15 @@ def _queue_followup_message( extra={"run_id": context.run_id if context is not None else None}, ) return - if message_id and message_id in self._accepted_message_id_set: - return if message_id: - self._accepted_message_ids.append(message_id) - self._accepted_message_id_set.add(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, diff --git a/products/tasks/backend/tests/test_api.py b/products/tasks/backend/tests/test_api.py index 562b90339578..dfa202118c9f 100644 --- a/products/tasks/backend/tests/test_api.py +++ b/products/tasks/backend/tests/test_api.py @@ -168,7 +168,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 @@ -1654,6 +1654,17 @@ def test_run_endpoint_starts_pi_task(self, mock_workflow): 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( [ ("run_source_omitted", None, "full"), @@ -7957,7 +7968,7 @@ def _create_run_with_sandbox(self, task, sandbox_url="http://localhost:9999", co def _open_sandbox_session(self, run, sandbox_id="sandbox-1"): now = django_timezone.now() - return SandboxSession.objects.create( + return SandboxSession.objects.unscoped().create( team=self.team, task_run=run, sandbox_id=sandbox_id, @@ -8011,7 +8022,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") @@ -8032,7 +8043,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") @@ -8070,7 +8081,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") @@ -8104,7 +8115,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") @@ -8203,7 +8220,7 @@ def test_task_session_sync_atomically_replaces_opaque_content(self, mock_write, response = self.client.generic( "POST", f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/task_session_sync/", - content, + cast(str, content), content_type="application/octet-stream", HTTP_IF_MATCH='"none"', HTTP_X_SANDBOX_ID="sandbox-1", @@ -8236,7 +8253,7 @@ def test_task_session_sync_rejects_stale_content_hash(self, mock_write, mock_del response = self.client.generic( "POST", f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/task_session_sync/", - b'{"type":"session","version":3}\n', + cast(str, b'{"type":"session","version":3}\n'), content_type="application/octet-stream", HTTP_IF_MATCH='"stale-hash"', HTTP_X_SANDBOX_ID="sandbox-1", @@ -8262,7 +8279,7 @@ def test_task_session_sync_rejects_a_closed_or_different_sandbox(self, mock_writ response = self.client.generic( "POST", f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/task_session_sync/", - b'{"type":"session"}\n', + cast(str, b'{"type":"session"}\n'), content_type="application/octet-stream", HTTP_IF_MATCH='"none"', HTTP_X_SANDBOX_ID="stale-sandbox", @@ -8284,7 +8301,7 @@ def test_task_session_sync_rejects_empty_content_before_upload(self, mock_write) response = self.client.generic( "POST", f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/task_session_sync/", - b"", + cast(str, b""), content_type="application/octet-stream", HTTP_IF_MATCH='"none"', HTTP_X_SANDBOX_ID="sandbox-1", @@ -8297,7 +8314,7 @@ def test_task_session_sync_rejects_empty_content_before_upload(self, mock_write) @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.create(team=self.team, name="shared", created_by=other_user) + 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"]) From 4d8bb7be9ce238cd49f88e34dcee9e3ca0f4cfb1 Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Mon, 27 Jul 2026 14:05:19 +0200 Subject: [PATCH 09/22] fix(tasks): keep Pi flag check behind facade --- products/tasks/backend/facade/api.py | 21 ++++++++++++++ .../tasks/backend/presentation/views/api.py | 28 +++---------------- 2 files changed, 25 insertions(+), 24 deletions(-) diff --git a/products/tasks/backend/facade/api.py b/products/tasks/backend/facade/api.py index 8af79fdf8809..cbee1a32cee7 100644 --- a/products/tasks/backend/facade/api.py +++ b/products/tasks/backend/facade/api.py @@ -41,6 +41,7 @@ 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, @@ -177,6 +178,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", @@ -3649,6 +3651,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 str(user.uuid) + 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) diff --git a/products/tasks/backend/presentation/views/api.py b/products/tasks/backend/presentation/views/api.py index c460b98abc3f..4417ada45032 100644 --- a/products/tasks/backend/presentation/views/api.py +++ b/products/tasks/backend/presentation/views/api.py @@ -34,7 +34,6 @@ from posthog.rate_limit import CodeInviteThrottle from posthog.renderers import ServerSentEventRenderer -from products.tasks.backend.constants import PI_CLOUD_RUNTIME_FEATURE_FLAG from products.tasks.backend.facade import ( access as tasks_access, api as tasks_facade, @@ -136,25 +135,6 @@ def parse(self, stream, media_type=None, parser_context=None): logger = logging.getLogger(__name__) -def _pi_cloud_runtime_enabled(user: Any, team: Any) -> bool: - distinct_id = getattr(user, "distinct_id", None) or str(getattr(user, "uuid", "")) - organization_id = str(getattr(team, "organization_id", "") or "") - 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 _pi_cloud_runtime_disabled_response() -> Response: return Response( TaskRunErrorResponseSerializer({"error": "Pi cloud runtime is disabled"}).data, @@ -627,7 +607,7 @@ def run(self, request, pk=None, **kwargs): raise NotFound() if tasks_facade.task_runtime( pk, self.team_id, self._user_id(), for_control=True - ) == tasks_facade.TaskRuntime.PI and not _pi_cloud_runtime_enabled(request.user, self.team): + ) == 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 @@ -983,7 +963,7 @@ def create(self, request, *args, **kwargs): 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 _pi_cloud_runtime_enabled(request.user, self.team): + ) == 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 @@ -1015,7 +995,7 @@ def start(self, request, pk=None, **kwargs): task_id = self._ensure_task_accessible() if tasks_facade.task_runtime( task_id, self.team_id, self._user_id(), for_control=True - ) == tasks_facade.TaskRuntime.PI and not _pi_cloud_runtime_enabled(request.user, self.team): + ) == tasks_facade.TaskRuntime.PI and not tasks_facade.pi_cloud_runtime_enabled(self.team, request.user): return _pi_cloud_runtime_disabled_response() startable = tasks_facade.check_task_run_startable(pk, task_id, self.team_id) @@ -1991,7 +1971,7 @@ def resume_in_cloud(self, request, pk=None, **kwargs): raise NotFound() if tasks_facade.task_runtime( task_id, self.team_id, self._user_id(), for_control=True - ) == tasks_facade.TaskRuntime.PI and not _pi_cloud_runtime_enabled(request.user, self.team): + ) == 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. From af7781b8c9fbd4d11c387d73657d41b76700baa3 Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Mon, 27 Jul 2026 14:07:47 +0200 Subject: [PATCH 10/22] fix(tasks): align Pi flag identity --- products/tasks/backend/facade/api.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/products/tasks/backend/facade/api.py b/products/tasks/backend/facade/api.py index cbee1a32cee7..2c6527458ff7 100644 --- a/products/tasks/backend/facade/api.py +++ b/products/tasks/backend/facade/api.py @@ -3652,7 +3652,7 @@ def get_conversation_task_dtos(task_ids: Sequence[str | UUID], team_id: int) -> def pi_cloud_runtime_enabled(team: Team, user: User) -> bool: - distinct_id = user.distinct_id or str(user.uuid) + distinct_id = user.distinct_id or user.email or str(user.uuid) organization_id = str(team.organization_id) try: return bool( From 8e87fe09ef913218d8fe784505ab7d562a226b5a Mon Sep 17 00:00:00 2001 From: "tests-posthog[bot]" <250237707+tests-posthog[bot]@users.noreply.github.com> Date: Mon, 27 Jul 2026 12:12:13 +0000 Subject: [PATCH 11/22] chore: update OpenAPI generated types --- products/tasks/frontend/generated/api.ts | 5 +++-- products/tasks/frontend/generated/api.zod.ts | 6 ------ products/tasks/mcp/tools.yaml | 6 ++++++ 3 files changed, 9 insertions(+), 8 deletions(-) diff --git a/products/tasks/frontend/generated/api.ts b/products/tasks/frontend/generated/api.ts index 4357cf1ab4a8..1c3cb908e092 100644 --- a/products/tasks/frontend/generated/api.ts +++ b/products/tasks/frontend/generated/api.ts @@ -1696,13 +1696,14 @@ export const tasksRunsTaskSessionSyncCreate = async ( projectId: string, taskId: string, id: string, - tasksRunsTaskSessionSyncCreateBody?: Blob | string, + tasksRunsTaskSessionSyncCreateBody?: Blob, options?: RequestInit ): Promise => { return apiMutator(getTasksRunsTaskSessionSyncCreateUrl(projectId, taskId, id), { ...options, method: 'POST', - body: JSON.stringify(tasksRunsTaskSessionSyncCreateBody), + headers: { 'Content-Type': 'application/octet-stream', ...options?.headers }, + body: tasksRunsTaskSessionSyncCreateBody, }) } diff --git a/products/tasks/frontend/generated/api.zod.ts b/products/tasks/frontend/generated/api.zod.ts index e7fdcee2705a..f86085134c9c 100644 --- a/products/tasks/frontend/generated/api.zod.ts +++ b/products/tasks/frontend/generated/api.zod.ts @@ -2640,12 +2640,6 @@ export const TasksRunsStartCreateBody = /* @__PURE__ */ zod.object({ ), }) -/** - * API for managing task runs. Each run represents an execution of a task. - * @summary Replace the active native task session - */ -export const TasksRunsTaskSessionSyncCreateBody = /* @__PURE__ */ zod.instanceof(File) - /** * Create a stable, editable artifact handle from direct markdown/text content or an existing run artifact. Slack adapters deliver into the mapped Slack thread; document artifacts use external connector storage when available. * @summary Create a living artifact for a task run diff --git a/products/tasks/mcp/tools.yaml b/products/tasks/mcp/tools.yaml index 31ac82fe78a8..ebe9980e9ae0 100644 --- a/products/tasks/mcp/tools.yaml +++ b/products/tasks/mcp/tools.yaml @@ -474,6 +474,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 From f526cdbd5b9668ee6b27878cb6aa5b3f411f80b2 Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Mon, 27 Jul 2026 14:33:44 +0200 Subject: [PATCH 12/22] fix(tasks): authenticate Pi session sync --- products/tasks/backend/facade/api.py | 26 ++++++++++ .../logic/services/connection_token.py | 38 +++++++++++++-- .../backend/logic/services/docker_sandbox.py | 5 ++ .../backend/logic/services/modal_sandbox.py | 4 ++ .../tasks/backend/logic/services/sandbox.py | 7 ++- .../tasks/backend/presentation/views/api.py | 25 +++++++++- .../activities/start_agent_server.py | 14 ++++-- .../tests/test_agent_proxy_callback.py | 5 +- products/tasks/backend/tests/test_api.py | 47 ++++++++++++++++++- .../backend/tests/test_connection_token.py | 13 ++++- .../tasks/backend/tests/test_event_ingest.py | 5 +- 11 files changed, 172 insertions(+), 17 deletions(-) diff --git a/products/tasks/backend/facade/api.py b/products/tasks/backend/facade/api.py index 2c6527458ff7..29229c1a64c1 100644 --- a/products/tasks/backend/facade/api.py +++ b/products/tasks/backend/facade/api.py @@ -216,6 +216,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", ] @@ -2274,6 +2275,31 @@ def _delete_task_session_object(task_session_id: UUID, object_storage_key: str) ) +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, diff --git a/products/tasks/backend/logic/services/connection_token.py b/products/tasks/backend/logic/services/connection_token.py index b2c3046492e6..6da30046e7b9 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 @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,6 +238,8 @@ 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}) @@ -243,7 +251,15 @@ def create_sandbox_event_ingest_token(task_run: TaskRun, ttl: timedelta = SANDBO 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) + sandbox_id = (task_run.state or {}).get("sandbox_id") + if not isinstance(sandbox_id, str) or not 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": sandbox_id}, + ) def validate_sandbox_event_ingest_token(token: str) -> SandboxEventIngestTokenPayload: @@ -252,11 +268,23 @@ 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") - - if not isinstance(run_id, str) or not isinstance(task_id, str) or type(team_id) is not int: + 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 + or 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 33d31522b2f7..e900a85e6b59 100644 --- a/products/tasks/backend/logic/services/docker_sandbox.py +++ b/products/tasks/backend/logic/services/docker_sandbox.py @@ -784,6 +784,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, @@ -804,6 +805,7 @@ def _build_agent_server_command( reasoning_effort=reasoning_effort, 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, @@ -892,6 +894,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, @@ -967,6 +970,7 @@ def start_agent_server( 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, @@ -1019,6 +1023,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 3f30890960a2..7dfe24603b1e 100644 --- a/products/tasks/backend/logic/services/modal_sandbox.py +++ b/products/tasks/backend/logic/services/modal_sandbox.py @@ -985,6 +985,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, @@ -1001,6 +1002,7 @@ def _build_agent_server_command( reasoning_effort=reasoning_effort, 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, @@ -1115,6 +1117,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, @@ -1193,6 +1196,7 @@ def start_agent_server( 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 47bcf4bae7f7..b2c7b4b2a25d 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+)" @@ -175,6 +177,7 @@ def build_agent_runtime_env_prefix( reasoning_effort: str | 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, @@ -189,6 +192,7 @@ def build_agent_runtime_env_prefix( "POSTHOG_CODE_REASONING_EFFORT": reasoning_effort, "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 @@ -350,6 +354,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/presentation/views/api.py b/products/tasks/backend/presentation/views/api.py index 4417ada45032..86e5f4c96763 100644 --- a/products/tasks/backend/presentation/views/api.py +++ b/products/tasks/backend/presentation/views/api.py @@ -19,7 +19,7 @@ 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 @@ -129,7 +129,19 @@ class OctetStreamParser(BaseParser): media_type = "application/octet-stream" def parse(self, stream, media_type=None, parser_context=None): - return stream.read() + request = (parser_context or {}).get("request") + raw_content_length = request.META.get("CONTENT_LENGTH") if request is not None else None + try: + content_length = int(raw_content_length) + except (TypeError, 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__) @@ -1237,6 +1249,15 @@ def sync_task_session(self, request, pk=None, **kwargs): 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."}) 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 0abac3538d11..1bf58f00f0e6 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 @@ -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) 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, ) @@ -340,6 +347,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, diff --git a/products/tasks/backend/tests/test_agent_proxy_callback.py b/products/tasks/backend/tests/test_agent_proxy_callback.py index 66be24830178..afab104440d5 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_api.py b/products/tasks/backend/tests/test_api.py index dfa202118c9f..d91fdcea0985 100644 --- a/products/tasks/backend/tests/test_api.py +++ b/products/tasks/backend/tests/test_api.py @@ -40,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, ) @@ -7940,6 +7941,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/" @@ -7968,6 +7970,8 @@ def _create_run_with_sandbox(self, task, sandbox_url="http://localhost:9999", co 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, @@ -8224,6 +8228,7 @@ def test_task_session_sync_atomically_replaces_opaque_content(self, mock_write, 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() @@ -8257,6 +8262,7 @@ def test_task_session_sync_rejects_stale_content_hash(self, mock_write, mock_del 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) @@ -8283,10 +8289,11 @@ def test_task_session_sync_rejects_a_closed_or_different_sandbox(self, mock_writ 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_409_CONFLICT) - self.assertEqual(response.json()["error"], "The task session writer is not the active sandbox") + 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") @@ -8305,12 +8312,48 @@ def test_task_session_sync_rejects_empty_content_before_upload(self, mock_write) 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") diff --git a/products/tasks/backend/tests/test_connection_token.py b/products/tasks/backend/tests/test_connection_token.py index b6aae52e60b6..257f293ef89e 100644 --- a/products/tasks/backend/tests/test_connection_token.py +++ b/products/tasks/backend/tests/test_connection_token.py @@ -54,7 +54,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 +111,7 @@ 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") 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 +181,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 15c92daa2263..1fd1f2a63691 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, From a76c2aa98e57412c1e5b88558739ccc6f539bdf6 Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Mon, 27 Jul 2026 16:01:56 +0200 Subject: [PATCH 13/22] feat(tasks): enable persistent Pi event streams --- products/tasks/backend/facade/api.py | 8 +++-- .../tasks/backend/presentation/views/api.py | 4 ++- .../activities/get_task_processing_context.py | 33 ++++++++++++------- .../tests/test_get_task_processing_context.py | 33 +++++++++++++++++++ products/tasks/backend/tests/test_api.py | 19 +++++++++++ 5 files changed, 82 insertions(+), 15 deletions(-) diff --git a/products/tasks/backend/facade/api.py b/products/tasks/backend/facade/api.py index 29229c1a64c1..a8ee7303b024 100644 --- a/products/tasks/backend/facade/api.py +++ b/products/tasks/backend/facade/api.py @@ -2842,7 +2842,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 @@ -2856,7 +2860,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( diff --git a/products/tasks/backend/presentation/views/api.py b/products/tasks/backend/presentation/views/api.py index 86e5f4c96763..847a781b0ee1 100644 --- a/products/tasks/backend/presentation/views/api.py +++ b/products/tasks/backend/presentation/views/api.py @@ -1604,7 +1604,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) 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 e670639e37c2..f64f1f30501a 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 @@ -785,12 +785,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", @@ -869,12 +874,16 @@ def get_task_processing_context(input: GetTaskProcessingContextInput) -> TaskPro "debug", f"use_modal_directory_resume_snapshots: {use_modal_directory_resume_snapshots} for this task run", ) - agent_proxy_keep_stream_open = _is_agent_proxy_keep_stream_open_enabled( - distinct_id=distinct_id, - organization_id=organization_id, - run_id=run_id, - state=state, - ) + keep_stream_open_override = state.get("agent_proxy_keep_stream_open") + if pi_persistent_streaming and not isinstance(keep_stream_open_override, bool): + agent_proxy_keep_stream_open = True + else: + agent_proxy_keep_stream_open = _is_agent_proxy_keep_stream_open_enabled( + distinct_id=distinct_id, + organization_id=organization_id, + run_id=run_id, + state=state, + ) emit_agent_log( run_id, "debug", 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 54dc585c7ad0..54554902e5cf 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,39 @@ 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_persistent_event_streaming_without_rollout_flags(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 True + + @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 diff --git a/products/tasks/backend/tests/test_api.py b/products/tasks/backend/tests/test_api.py index d91fdcea0985..cb0a3beb7299 100644 --- a/products/tasks/backend/tests/test_api.py +++ b/products/tasks/backend/tests/test_api.py @@ -6623,6 +6623,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() From c8afacdaf71493d12e241214d7de4563ea2f3d63 Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Mon, 27 Jul 2026 17:28:40 +0200 Subject: [PATCH 14/22] feat(tasks): proxy Pi queue operations --- .../tasks/backend/presentation/serializers.py | 2 ++ .../tasks/backend/presentation/views/api.py | 10 ++++++++-- .../activities/get_task_processing_context.py | 16 ++++++---------- .../tests/test_get_task_processing_context.py | 6 ++++-- products/tasks/backend/tests/test_api.py | 19 +++++++++++++++++++ 5 files changed, 39 insertions(+), 14 deletions(-) diff --git a/products/tasks/backend/presentation/serializers.py b/products/tasks/backend/presentation/serializers.py index 2fa39d47b134..9f12a06bfcb6 100644 --- a/products/tasks/backend/presentation/serializers.py +++ b/products/tasks/backend/presentation/serializers.py @@ -2292,6 +2292,8 @@ class TaskRunCommandRequestSerializer(serializers.Serializer): "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 diff --git a/products/tasks/backend/presentation/views/api.py b/products/tasks/backend/presentation/views/api.py index 847a781b0ee1..b6305f4c4cac 100644 --- a/products/tasks/backend/presentation/views/api.py +++ b/products/tasks/backend/presentation/views/api.py @@ -1627,7 +1627,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, mcp_response, and native Pi RPC commands.", + "set_config_option, mcp_response, native Pi RPC commands, and Pi queue operations.", strict_request_validation=True, ) @action( @@ -1645,7 +1645,13 @@ def command(self, request, pk=None, **kwargs): 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"}: + 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, 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 f64f1f30501a..c2b2ba5d2bb3 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 @@ -874,16 +874,12 @@ def get_task_processing_context(input: GetTaskProcessingContextInput) -> TaskPro "debug", f"use_modal_directory_resume_snapshots: {use_modal_directory_resume_snapshots} for this task run", ) - keep_stream_open_override = state.get("agent_proxy_keep_stream_open") - if pi_persistent_streaming and not isinstance(keep_stream_open_override, bool): - agent_proxy_keep_stream_open = True - else: - agent_proxy_keep_stream_open = _is_agent_proxy_keep_stream_open_enabled( - distinct_id=distinct_id, - organization_id=organization_id, - run_id=run_id, - state=state, - ) + agent_proxy_keep_stream_open = _is_agent_proxy_keep_stream_open_enabled( + distinct_id=distinct_id, + organization_id=organization_id, + run_id=run_id, + state=state, + ) emit_agent_log( run_id, "debug", 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 54554902e5cf..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 @@ -361,7 +361,9 @@ def feature_enabled(flag_key, **kwargs): assert sandbox_args[0] == SANDBOX_EVENT_INGEST_FEATURE_FLAG @pytest.mark.django_db(transaction=True) - def test_pi_runtime_enables_persistent_event_streaming_without_rollout_flags(self, activity_environment, test_task): + 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() @@ -374,7 +376,7 @@ def test_pi_runtime_enables_persistent_event_streaming_without_rollout_flags(sel 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 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): diff --git a/products/tasks/backend/tests/test_api.py b/products/tasks/backend/tests/test_api.py index cb0a3beb7299..76936b604f09 100644 --- a/products/tasks/backend/tests/test_api.py +++ b/products/tasks/backend/tests/test_api.py @@ -8434,6 +8434,25 @@ def test_command_proxies_pi_rpc(self, mock_post): 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) + @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): From 488c5168e3223f3dd4fd08d2471db101b0f88d81 Mon Sep 17 00:00:00 2001 From: "tests-posthog[bot]" <250237707+tests-posthog[bot]@users.noreply.github.com> Date: Tue, 28 Jul 2026 08:17:25 +0000 Subject: [PATCH 15/22] chore: update OpenAPI generated types --- products/tasks/frontend/generated/api.schemas.ts | 8 +++++++- products/tasks/frontend/generated/api.ts | 2 +- products/tasks/frontend/generated/api.zod.ts | 8 +++++--- services/mcp/src/api/generated.ts | 8 +++++++- 4 files changed, 20 insertions(+), 6 deletions(-) diff --git a/products/tasks/frontend/generated/api.schemas.ts b/products/tasks/frontend/generated/api.schemas.ts index 1b63b178d308..3795c0e23681 100644 --- a/products/tasks/frontend/generated/api.schemas.ts +++ b/products/tasks/frontend/generated/api.schemas.ts @@ -2839,6 +2839,8 @@ export const JsonrpcEnumApi = { * * `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] @@ -2850,6 +2852,8 @@ export const MethodEnumApi = { SetConfigOption: 'set_config_option', McpResponse: 'mcp_response', PiRpc: 'pi/rpc', + QueueGet: 'queue_get', + QueueClear: 'queue_clear', } as const /** @@ -2868,7 +2872,9 @@ export interface TaskRunCommandRequestApi { * * `permission_response` - permission_response * * `set_config_option` - set_config_option * * `mcp_response` - mcp_response - * * `pi/rpc` - pi/rpc */ + * * `pi/rpc` - pi/rpc + * * `queue_get` - queue_get + * * `queue_clear` - queue_clear */ method: MethodEnumApi /** Parameters for the command */ params?: TaskRunCommandRequestApiParams diff --git a/products/tasks/frontend/generated/api.ts b/products/tasks/frontend/generated/api.ts index ab0394fb3197..73d2cc0ee28d 100644 --- a/products/tasks/frontend/generated/api.ts +++ b/products/tasks/frontend/generated/api.ts @@ -1492,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, mcp_response, and native Pi RPC 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 ( diff --git a/products/tasks/frontend/generated/api.zod.ts b/products/tasks/frontend/generated/api.zod.ts index e4b0514f5886..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, mcp_response, and native Pi RPC 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 @@ -2622,12 +2622,14 @@ export const TasksRunsCommandCreateBody = /* @__PURE__ */ zod '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\n\* `pi\/rpc` - pi\/rpc' + '\* `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\n\* `pi\/rpc` - pi\/rpc' + '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/services/mcp/src/api/generated.ts b/services/mcp/src/api/generated.ts index ca971c66e5d6..1ba62a7821da 100644 --- a/services/mcp/src/api/generated.ts +++ b/services/mcp/src/api/generated.ts @@ -39478,6 +39478,8 @@ export namespace Schemas { * * `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]; @@ -39490,6 +39492,8 @@ export namespace Schemas { SetConfigOption: 'set_config_option', McpResponse: 'mcp_response', PiRpc: 'pi/rpc', + QueueGet: 'queue_get', + QueueClear: 'queue_clear', } as const; /** @@ -67903,7 +67907,9 @@ export namespace Schemas { * * `permission_response` - permission_response * * `set_config_option` - set_config_option * * `mcp_response` - mcp_response - * * `pi/rpc` - pi/rpc */ + * * `pi/rpc` - pi/rpc + * * `queue_get` - queue_get + * * `queue_clear` - queue_clear */ method: MethodEnum; /** Parameters for the command */ params?: TaskRunCommandRequestParams; From 08c952c8433e8cdfd76cc5edee350d8f1abab4b7 Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Tue, 28 Jul 2026 10:30:05 +0200 Subject: [PATCH 16/22] fix(tasks): restrict Pi RPC commands --- .../tasks/backend/presentation/serializers.py | 13 ++++++++++ products/tasks/backend/tests/test_api.py | 25 ++++++++++++++++--- 2 files changed, 35 insertions(+), 3 deletions(-) diff --git a/products/tasks/backend/presentation/serializers.py b/products/tasks/backend/presentation/serializers.py index 182081228b31..98973d8f3c14 100644 --- a/products/tasks/backend/presentation/serializers.py +++ b/products/tasks/backend/presentation/serializers.py @@ -2444,6 +2444,17 @@ class TaskRunCommandRequestSerializer(serializers.Serializer): # response payload plus envelope must fit in 300 KB. Params are forwarded to the sandbox # verbatim and never persisted or captured — they carry data from the user's private systems. MAX_MCP_RESPONSE_PARAMS_BYTES = 300_000 + ALLOWED_PI_RPC_COMMANDS = { + "abort", + "compact", + "get_available_models", + "get_available_thinking_levels", + "get_commands", + "get_session_stats", + "get_state", + "set_model", + "set_thinking_level", + } jsonrpc = serializers.ChoiceField( choices=["2.0"], @@ -2517,6 +2528,8 @@ def validate(self, attrs): 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"}) + if command_type not in self.ALLOWED_PI_RPC_COMMANDS: + raise serializers.ValidationError({"params": f"Pi RPC command is not allowed: {command_type}"}) 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"}) diff --git a/products/tasks/backend/tests/test_api.py b/products/tasks/backend/tests/test_api.py index 667b0b874b2e..2d8ab4dbc987 100644 --- a/products/tasks/backend/tests/test_api.py +++ b/products/tasks/backend/tests/test_api.py @@ -8471,9 +8471,9 @@ def test_command_proxies_pi_rpc(self, mock_post): reset_sandbox_jwt_key_cache() rpc_response = { "type": "response", - "command": "future_native_command", + "command": "get_state", "success": True, - "data": {"accepted": True}, + "data": {"isStreaming": False}, } self._mock_agent_response(mock_post, {"jsonrpc": "2.0", "id": "native", "result": rpc_response}) task = self.create_task(runtime=Task.Runtime.PI) @@ -8484,7 +8484,7 @@ def test_command_proxies_pi_rpc(self, mock_post): { "jsonrpc": "2.0", "method": "pi/rpc", - "params": {"command": {"id": "native", "type": "future_native_command"}}, + "params": {"command": {"id": "native", "type": "get_state"}}, "id": "native", }, format="json", @@ -8493,6 +8493,25 @@ def test_command_proxies_pi_rpc(self, mock_post): self.assertEqual(response.status_code, status.HTTP_200_OK) self.assertEqual(response.json()["result"], rpc_response) + @parameterized.expand([("bash",), ("future_native_command",)]) + def test_command_rejects_unsafe_or_unknown_pi_rpc(self, command_type): + 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": command_type}}, + "id": "native", + }, + format="json", + ) + + self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) + self.assertEqual(response.json()["type"], "validation_error") + @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") From 954a3df859855f1801558d3c1b9b02a406e9fbf3 Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Tue, 28 Jul 2026 12:18:04 +0200 Subject: [PATCH 17/22] Revert "fix(tasks): restrict Pi RPC commands" This reverts commit 8b5ee3f540502b746e491981952a21136de55db9. --- .../tasks/backend/presentation/serializers.py | 13 ---------- products/tasks/backend/tests/test_api.py | 25 +++---------------- 2 files changed, 3 insertions(+), 35 deletions(-) diff --git a/products/tasks/backend/presentation/serializers.py b/products/tasks/backend/presentation/serializers.py index 98973d8f3c14..182081228b31 100644 --- a/products/tasks/backend/presentation/serializers.py +++ b/products/tasks/backend/presentation/serializers.py @@ -2444,17 +2444,6 @@ class TaskRunCommandRequestSerializer(serializers.Serializer): # response payload plus envelope must fit in 300 KB. Params are forwarded to the sandbox # verbatim and never persisted or captured — they carry data from the user's private systems. MAX_MCP_RESPONSE_PARAMS_BYTES = 300_000 - ALLOWED_PI_RPC_COMMANDS = { - "abort", - "compact", - "get_available_models", - "get_available_thinking_levels", - "get_commands", - "get_session_stats", - "get_state", - "set_model", - "set_thinking_level", - } jsonrpc = serializers.ChoiceField( choices=["2.0"], @@ -2528,8 +2517,6 @@ def validate(self, attrs): 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"}) - if command_type not in self.ALLOWED_PI_RPC_COMMANDS: - raise serializers.ValidationError({"params": f"Pi RPC command is not allowed: {command_type}"}) 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"}) diff --git a/products/tasks/backend/tests/test_api.py b/products/tasks/backend/tests/test_api.py index 2d8ab4dbc987..667b0b874b2e 100644 --- a/products/tasks/backend/tests/test_api.py +++ b/products/tasks/backend/tests/test_api.py @@ -8471,9 +8471,9 @@ def test_command_proxies_pi_rpc(self, mock_post): reset_sandbox_jwt_key_cache() rpc_response = { "type": "response", - "command": "get_state", + "command": "future_native_command", "success": True, - "data": {"isStreaming": False}, + "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) @@ -8484,7 +8484,7 @@ def test_command_proxies_pi_rpc(self, mock_post): { "jsonrpc": "2.0", "method": "pi/rpc", - "params": {"command": {"id": "native", "type": "get_state"}}, + "params": {"command": {"id": "native", "type": "future_native_command"}}, "id": "native", }, format="json", @@ -8493,25 +8493,6 @@ def test_command_proxies_pi_rpc(self, mock_post): self.assertEqual(response.status_code, status.HTTP_200_OK) self.assertEqual(response.json()["result"], rpc_response) - @parameterized.expand([("bash",), ("future_native_command",)]) - def test_command_rejects_unsafe_or_unknown_pi_rpc(self, command_type): - 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": command_type}}, - "id": "native", - }, - format="json", - ) - - self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) - self.assertEqual(response.json()["type"], "validation_error") - @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") From 279a58285344452653ff59a683b5ceaf3b6045ba Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Tue, 28 Jul 2026 15:40:12 +0200 Subject: [PATCH 18/22] fix(tasks): align Pi sessions with current task runtime --- posthog/test/setup_receivers_baseline.txt | 1 + products/tasks/backend/facade/api.py | 19 ++++++++++--------- .../backend/logic/services/sandbox_usage.py | 8 -------- .../0075_sandbox_session_open_constraint.py | 16 ---------------- .../backend/migrations/max_migration.txt | 2 +- products/tasks/backend/models.py | 7 ------- .../tasks/backend/presentation/views/api.py | 4 +++- products/tasks/backend/tests/test_agentsh.py | 4 ++++ products/tasks/backend/tests/test_api.py | 3 ++- .../backend/tests/test_run_log_mirror.py | 8 ++++---- 10 files changed, 25 insertions(+), 47 deletions(-) delete mode 100644 products/tasks/backend/migrations/0075_sandbox_session_open_constraint.py 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/facade/api.py b/products/tasks/backend/facade/api.py index 5f518b44882a..82fc1d488d9d 100644 --- a/products/tasks/backend/facade/api.py +++ b/products/tasks/backend/facade/api.py @@ -4870,19 +4870,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) - initial_message = None 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 if is_pi_task else None, - raise_on_error=False, - ) + _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/logic/services/sandbox_usage.py b/products/tasks/backend/logic/services/sandbox_usage.py index 3a4d0478e9a5..164b9924eef0 100644 --- a/products/tasks/backend/logic/services/sandbox_usage.py +++ b/products/tasks/backend/logic/services/sandbox_usage.py @@ -62,14 +62,6 @@ def open_sandbox_session( .only("id", "team_id", "state", "task__origin_product") .get(id=run_id) ) - SandboxSession.objects.for_team(run.team_id).filter( - task_run_id=run.id, - ended_at__isnull=True, - ).exclude(sandbox_id=sandbox_id).update( - ended_at=timezone.now(), - ended_reason=SandboxSession.EndedReason.CLEANUP, - ) - state = run.state or {} created_at = sandbox_created_at or timezone.now() shape = { diff --git a/products/tasks/backend/migrations/0075_sandbox_session_open_constraint.py b/products/tasks/backend/migrations/0075_sandbox_session_open_constraint.py deleted file mode 100644 index 74708e7957d2..000000000000 --- a/products/tasks/backend/migrations/0075_sandbox_session_open_constraint.py +++ /dev/null @@ -1,16 +0,0 @@ -from django.db import migrations, models - - -class Migration(migrations.Migration): - dependencies = [("tasks", "0074_task_session")] - - operations = [ - migrations.AddConstraint( - model_name="sandboxsession", - constraint=models.UniqueConstraint( - condition=models.Q(ended_at__isnull=True), - fields=("task_run",), - name="sandbox_session_one_open_per_run", - ), - ), - ] diff --git a/products/tasks/backend/migrations/max_migration.txt b/products/tasks/backend/migrations/max_migration.txt index f829b35d3578..4fc010358639 100644 --- a/products/tasks/backend/migrations/max_migration.txt +++ b/products/tasks/backend/migrations/max_migration.txt @@ -1 +1 @@ -0075_sandbox_session_open_constraint +0074_task_session diff --git a/products/tasks/backend/models.py b/products/tasks/backend/models.py index 5bab124c51fd..1051ea285e25 100644 --- a/products/tasks/backend/models.py +++ b/products/tasks/backend/models.py @@ -2285,13 +2285,6 @@ class EndedReason(models.TextChoices): class Meta: db_table = "posthog_task_sandbox_session" - constraints = [ - models.UniqueConstraint( - fields=["task_run"], - condition=models.Q(ended_at__isnull=True), - name="sandbox_session_one_open_per_run", - ), - ] indexes = [ # The usage report scans sessions overlapping the period instance-wide: # closed recently (ended_at > begin) or open and not yet past their TTL diff --git a/products/tasks/backend/presentation/views/api.py b/products/tasks/backend/presentation/views/api.py index b6305f4c4cac..ee45bfd6fea3 100644 --- a/products/tasks/backend/presentation/views/api.py +++ b/products/tasks/backend/presentation/views/api.py @@ -131,9 +131,11 @@ class OctetStreamParser(BaseParser): 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 (TypeError, ValueError) as error: + 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") 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 667b0b874b2e..3727b21c7b1e 100644 --- a/products/tasks/backend/tests/test_api.py +++ b/products/tasks/backend/tests/test_api.py @@ -4542,8 +4542,9 @@ 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.facade.api.pi_cloud_runtime_enabled", return_value=True) @patch("products.tasks.backend.temporal.client.resume_task_in_cloud_workflow") - def test_resume_in_cloud_starts_pi_task(self, mock_resume): + def test_resume_in_cloud_starts_pi_task(self, mock_resume, _mock_pi_enabled): task = self.create_task(runtime=Task.Runtime.PI) run = TaskRun.objects.create( task=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) From bf81141c45147ea18b204f5efcb201fb7f5ef3e4 Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Tue, 28 Jul 2026 16:06:00 +0200 Subject: [PATCH 19/22] fix(tasks): preserve ACP workflow dispatch shape --- products/tasks/backend/facade/api.py | 31 ++++++++++++++---------- products/tasks/backend/tests/test_api.py | 2 +- 2 files changed, 19 insertions(+), 14 deletions(-) diff --git a/products/tasks/backend/facade/api.py b/products/tasks/backend/facade/api.py index 82fc1d488d9d..d3c3cb2803e3 100644 --- a/products/tasks/backend/facade/api.py +++ b/products/tasks/backend/facade/api.py @@ -3389,23 +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, - initial_message=( - PendingFollowup( + 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()), - ) - if initial_message or initial_artifact_ids - else None - ), - ) + ), + ) + 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) diff --git a/products/tasks/backend/tests/test_api.py b/products/tasks/backend/tests/test_api.py index 3727b21c7b1e..1a1a6b681464 100644 --- a/products/tasks/backend/tests/test_api.py +++ b/products/tasks/backend/tests/test_api.py @@ -4542,7 +4542,7 @@ 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.facade.api.pi_cloud_runtime_enabled", return_value=True) + @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_starts_pi_task(self, mock_resume, _mock_pi_enabled): task = self.create_task(runtime=Task.Runtime.PI) From 84d884aa8c7e736f8b9dde78b2780330b378b01b Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Tue, 28 Jul 2026 18:16:13 +0200 Subject: [PATCH 20/22] fix(tasks): address Pi cloud review findings --- products/tasks/backend/facade/api.py | 2 +- .../logic/services/connection_token.py | 25 ++++++------ .../backend/logic/services/sandbox_usage.py | 2 +- .../backend/migrations/0074_task_session.py | 1 + products/tasks/backend/models.py | 16 ++++++++ .../tasks/backend/presentation/serializers.py | 10 ++--- .../tasks/backend/presentation/views/api.py | 40 +++++++++++++++++-- .../activities/get_sandbox_for_repository.py | 5 ++- .../activities/provision_sandbox.py | 5 ++- .../activities/send_followup_to_sandbox.py | 3 +- .../activities/start_agent_server.py | 8 ++-- .../tests/test_start_agent_server.py | 1 + .../tests/test_send_followup_to_sandbox.py | 19 +++++++++ .../backend/temporal/process_task/workflow.py | 4 +- products/tasks/backend/tests/test_api.py | 26 ++++++++++++ .../backend/tests/test_connection_token.py | 39 +++++++++++++++++- products/tasks/backend/tests/test_models.py | 36 +++++++++++++++++ .../tasks/frontend/generated/api.schemas.ts | 13 +++++- services/mcp/src/api/generated.ts | 13 +++++- 19 files changed, 232 insertions(+), 36 deletions(-) diff --git a/products/tasks/backend/facade/api.py b/products/tasks/backend/facade/api.py index d3c3cb2803e3..7d5c9fbb2032 100644 --- a/products/tasks/backend/facade/api.py +++ b/products/tasks/backend/facade/api.py @@ -3692,7 +3692,7 @@ def get_conversation_task_dtos(task_ids: Sequence[str | UUID], team_id: int) -> def pi_cloud_runtime_enabled(team: Team, user: User) -> bool: - distinct_id = user.distinct_id or user.email or str(user.uuid) + distinct_id = user.distinct_id or f"user_{user.id}" organization_id = str(team.organization_id) try: return bool( diff --git a/products/tasks/backend/logic/services/connection_token.py b/products/tasks/backend/logic/services/connection_token.py index 6da30046e7b9..a49da81d47c5 100644 --- a/products/tasks/backend/logic/services/connection_token.py +++ b/products/tasks/backend/logic/services/connection_token.py @@ -38,7 +38,7 @@ class SandboxEventIngestTokenPayload: run_id: str task_id: str team_id: int - sandbox_id: str + sandbox_id: str | None @dataclass(frozen=True) @@ -244,21 +244,26 @@ def _encode_run_scoped_token( 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. """ - sandbox_id = (task_run.state or {}).get("sandbox_id") - if not isinstance(sandbox_id, str) or not sandbox_id: + 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": sandbox_id}, + {"sandbox_id": active_sandbox_id}, ) @@ -270,13 +275,9 @@ def validate_sandbox_event_ingest_token(token: str) -> SandboxEventIngestTokenPa 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 - or not isinstance(sandbox_id, str) - or not 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( diff --git a/products/tasks/backend/logic/services/sandbox_usage.py b/products/tasks/backend/logic/services/sandbox_usage.py index 164b9924eef0..ffb8516e6aa0 100644 --- a/products/tasks/backend/logic/services/sandbox_usage.py +++ b/products/tasks/backend/logic/services/sandbox_usage.py @@ -57,7 +57,7 @@ def open_sandbox_session( try: with transaction.atomic(): run = ( - TaskRun.objects.select_for_update() + TaskRun.objects.select_for_update(of=("self",)) .select_related("task") .only("id", "team_id", "state", "task__origin_product") .get(id=run_id) diff --git a/products/tasks/backend/migrations/0074_task_session.py b/products/tasks/backend/migrations/0074_task_session.py index 54dbfd2f3e3f..589f7e7e5e76 100644 --- a/products/tasks/backend/migrations/0074_task_session.py +++ b/products/tasks/backend/migrations/0074_task_session.py @@ -48,6 +48,7 @@ class Migration(migrations.Migration): "team", models.ForeignKey( db_constraint=False, + db_index=False, on_delete=django.db.models.deletion.CASCADE, related_name="+", to="posthog.team", diff --git a/products/tasks/backend/models.py b/products/tasks/backend/models.py index 1051ea285e25..67772255cc88 100644 --- a/products/tasks/backend/models.py +++ b/products/tasks/backend/models.py @@ -883,6 +883,7 @@ class TaskSession(TeamScopedRootMixin, UUIDModel): 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) @@ -1727,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 diff --git a/products/tasks/backend/presentation/serializers.py b/products/tasks/backend/presentation/serializers.py index 182081228b31..af48b3ad07af 100644 --- a/products/tasks/backend/presentation/serializers.py +++ b/products/tasks/backend/presentation/serializers.py @@ -731,14 +731,14 @@ def validate_entries(self, value): class TaskSessionResponseSerializer(serializers.Serializer): - id = serializers.UUIDField() - download_url = serializers.URLField(allow_null=True) - content_sha256 = serializers.CharField(allow_null=True) + 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() - content_sha256 = serializers.CharField() + 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): diff --git a/products/tasks/backend/presentation/views/api.py b/products/tasks/backend/presentation/views/api.py index ee45bfd6fea3..a2c7be89ab49 100644 --- a/products/tasks/backend/presentation/views/api.py +++ b/products/tasks/backend/presentation/views/api.py @@ -605,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" @@ -961,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" ), @@ -996,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" @@ -1007,14 +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 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() - startable = tasks_facade.check_task_run_startable(pk, task_id, self.team_id) - if startable == "not_found": - raise NotFound() if startable == "not_cloud": return Response( TaskRunErrorResponseSerializer({"error": "Only cloud runs can be started via this endpoint"}).data, @@ -1232,8 +1236,33 @@ def task_session(self, request, pk=None, **kwargs): @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), }, @@ -1642,7 +1671,9 @@ 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 method.startswith("pi/") and task_runtime != tasks_facade.TaskRuntime.PI: + if ( + method.startswith("pi/") or method in {"queue_get", "queue_clear"} + ) and task_runtime != tasks_facade.TaskRuntime.PI: return Response( TaskRunErrorResponseSerializer({"error": "Pi commands require a Pi task."}).data, status=status.HTTP_400_BAD_REQUEST, @@ -1983,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" ), 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 a52bfb235204..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 @@ -394,7 +394,10 @@ def get_sandbox_for_repository(input: GetSandboxForRepositoryInput) -> GetSandbo 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 activity.logger.info(f"Created sandbox {sandbox.id} (used_snapshot={used_snapshot})") 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 849b9af7697f..82d21921ea43 100644 --- a/products/tasks/backend/temporal/process_task/activities/provision_sandbox.py +++ b/products/tasks/backend/temporal/process_task/activities/provision_sandbox.py @@ -631,7 +631,10 @@ def create_sandbox_for_repository(input: CreateSandboxForRepositoryInput) -> Cre 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 emit_agent_log(ctx.run_id, "debug", f"Sandbox provisioned: {sandbox.id}") 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 7f10e3ca1e42..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 @@ -81,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 @@ -290,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 ad4fd3a18446..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 @@ -194,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) @@ -225,7 +225,7 @@ def _prepare_launch(ctx: TaskProcessingContext, scopes: PosthogMcpScopes) -> _La task_run_session_token: str | None = None if event_stream_ingest_enabled or task.runtime == Task.Runtime.PI: try: - run_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 task run token", @@ -451,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) @@ -489,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_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/workflow.py b/products/tasks/backend/temporal/process_task/workflow.py index e5ec11476b0c..9147fd5cf9cb 100644 --- a/products/tasks/backend/temporal/process_task/workflow.py +++ b/products/tasks/backend/temporal/process_task/workflow.py @@ -2123,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( @@ -2134,12 +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), heartbeat_timeout=timedelta(minutes=1), retry_policy=RetryPolicy( initial_interval=timedelta(seconds=5), - maximum_attempts=(1 if self.context.task_runtime == "pi" else SEND_FOLLOWUP_MAX_ATTEMPTS), + maximum_attempts=max_attempts, ), ) except Exception as e: diff --git a/products/tasks/backend/tests/test_api.py b/products/tasks/backend/tests/test_api.py index 1a1a6b681464..4bed40ab0eb4 100644 --- a/products/tasks/backend/tests/test_api.py +++ b/products/tasks/backend/tests/test_api.py @@ -2290,6 +2290,16 @@ def test_start_run_endpoint_starts_pi_task(self, mock_workflow): 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") def test_start_run_endpoint_rejects_missing_run_artifacts(self, mock_workflow): task = self.create_task() @@ -8513,6 +8523,22 @@ def test_command_proxies_pi_queue_operations(self, method, mock_post): 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 257f293ef89e..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, @@ -113,6 +114,42 @@ def test_ingest_token_signed_with_run_stored_kid(self) -> None: 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 # primary rotates, as long as the old key is retained as the secondary. diff --git a/products/tasks/backend/tests/test_models.py b/products/tasks/backend/tests/test_models.py index 33648a74cb5b..971c31d98b05 100644 --- a/products/tasks/backend/tests/test_models.py +++ b/products/tasks/backend/tests/test_models.py @@ -719,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/frontend/generated/api.schemas.ts b/products/tasks/frontend/generated/api.schemas.ts index 3795c0e23681..2ef2e7534103 100644 --- a/products/tasks/frontend/generated/api.schemas.ts +++ b/products/tasks/frontend/generated/api.schemas.ts @@ -2964,15 +2964,24 @@ export interface StreamReadTokenResponseApi { } export interface TaskSessionResponseApi { + /** Task session identifier */ id: string - /** @nullable */ + /** + * Temporary URL for downloading the session + * @nullable + */ download_url: string | null - /** @nullable */ + /** + * 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 } diff --git a/services/mcp/src/api/generated.ts b/services/mcp/src/api/generated.ts index 1ba62a7821da..a8ae06ae6bc6 100644 --- a/services/mcp/src/api/generated.ts +++ b/services/mcp/src/api/generated.ts @@ -68237,15 +68237,24 @@ export namespace Schemas { } export interface TaskSessionResponse { + /** Task session identifier */ id: string; - /** @nullable */ + /** + * Temporary URL for downloading the session + * @nullable + */ download_url: string | null; - /** @nullable */ + /** + * 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; } From 87b3c214756104ecb6335d01a255e8c58bfec6e7 Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Tue, 28 Jul 2026 19:52:19 +0200 Subject: [PATCH 21/22] fix(tasks): reset event stream when resuming cloud runs --- products/tasks/backend/facade/api.py | 7 ++++++ products/tasks/backend/facade/streams.py | 2 ++ .../backend/logic/stream/redis_stream.py | 22 +++++++++++++++++++ products/tasks/backend/tests/test_api.py | 4 +++- 4 files changed, 34 insertions(+), 1 deletion(-) diff --git a/products/tasks/backend/facade/api.py b/products/tasks/backend/facade/api.py index 7d5c9fbb2032..e4b95fd2a24f 100644 --- a/products/tasks/backend/facade/api.py +++ b/products/tasks/backend/facade/api.py @@ -3509,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, ) @@ -3574,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)}) 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/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/tests/test_api.py b/products/tasks/backend/tests/test_api.py index 4bed40ab0eb4..5272102f7b13 100644 --- a/products/tasks/backend/tests/test_api.py +++ b/products/tasks/backend/tests/test_api.py @@ -4554,7 +4554,8 @@ def test_update_run_status_to_completed_signals_workflow(self, mock_signal): @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_starts_pi_task(self, mock_resume, _mock_pi_enabled): + @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, @@ -4569,6 +4570,7 @@ def test_resume_in_cloud_starts_pi_task(self, mock_resume, _mock_pi_enabled): 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") From 68826c4a56620df62171aa170522e29c00a1bb85 Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Wed, 29 Jul 2026 12:50:52 +0200 Subject: [PATCH 22/22] test(tasks): use bot authorship for cloud resume --- products/tasks/backend/tests/test_api.py | 1 + 1 file changed, 1 insertion(+) diff --git a/products/tasks/backend/tests/test_api.py b/products/tasks/backend/tests/test_api.py index 5272102f7b13..3bb8b99da24a 100644 --- a/products/tasks/backend/tests/test_api.py +++ b/products/tasks/backend/tests/test_api.py @@ -4562,6 +4562,7 @@ def test_resume_in_cloud_starts_pi_task(self, mock_reset_stream, mock_resume, _m 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/")