From 7e2d270e9acff440209ef917b022f2de7955bb8c Mon Sep 17 00:00:00 2001 From: Vojta Bartos Date: Mon, 13 Jul 2026 15:29:39 +0200 Subject: [PATCH 1/5] fix(tasks): rebind sandbox MCP session on actor transitions The follow-up MCP refresh gate skipped whenever any token had been issued for the run within the freshness window, so when a different user spoke next (multiplayer Slack threads), the live session kept the previous actor's OAuth token for up to the full window. Key the freshness mark per (sandbox, user) and remember which identity the session was last bound to: a transition bypasses the window and rebinds immediately, while repeat messages from the same actor still skip the redundant refresh. Scoping the marks to the sandbox id also means a replacement sandbox (retry, snapshot restore) starts unmarked and gets a fresh token instead of inheriting a stale run-keyed mark. --- .../activities/send_followup_to_sandbox.py | 47 +++++-- .../activities/start_agent_server.py | 16 ++- .../tests/test_send_followup_to_sandbox.py | 123 ++++++++++++++++-- .../backend/temporal/process_task/utils.py | 69 ++++++++-- 4 files changed, 220 insertions(+), 35 deletions(-) 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 b770a7d03207..5627e723b911 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 @@ -27,11 +27,14 @@ from products.tasks.backend.temporal.oauth import create_oauth_access_token_for_run from products.tasks.backend.temporal.process_task.utils import ( get_actor_distinct_id, + get_last_sandbox_identity, get_sandbox_ph_mcp_configs, get_task_run_credential_user, get_user_mcp_server_configs, is_slack_interaction_state, mark_mcp_token_issued, + mark_sandbox_identity, + sandbox_identity_scope, should_refresh_mcp_token, ) @@ -216,6 +219,13 @@ def _deliver_followup(input: SendFollowupToSandboxInput) -> None: raise ApplicationError(f"send_followup failed: {error_msg}", non_retryable=True) +def _mark_sandbox_session(scope: str, user_id: int) -> None: + """Record whose token the sandbox session now holds and start that user's + freshness window.""" + mark_mcp_token_issued(scope, user_id) + mark_sandbox_identity(scope, "mcp", user_id) + + def _refresh_sandbox_mcp( task_run: TaskRun, scopes: PosthogMcpScopes, @@ -226,17 +236,32 @@ def _refresh_sandbox_mcp( Best-effort: retries once on failure, then logs and returns. Never raises — a failed refresh should not block an otherwise-valid follow-up. - Skipped entirely if a token was issued for this run within the last - MCP_TOKEN_REFRESH_INTERVAL_SECONDS — the in-sandbox token is still fresh. + Skipped only when the sandbox's live session already holds a fresh token + for *this message's actor*: the freshness window is keyed per + (sandbox, user), and an actor transition (the run-state actor differs from + the identity last pushed to the sandbox) bypasses the window entirely, so + the session rebinds before the new speaker's turn is delivered. """ run_id = str(task_run.id) - if not should_refresh_mcp_token(run_id): - logger.info("refresh_mcp_skipped_within_interval", run_id=run_id) - return - task = task_run.task + scope = sandbox_identity_scope(run_id, task_run.state) try: actor_user = get_task_run_credential_user(task, task_run.state) + if actor_user is not None: + # Until a refresh records otherwise, the sandbox holds the + # boot-time token, which was minted for the task creator. + last_identity = get_last_sandbox_identity(scope, "mcp") or task.created_by_id + identity_changed = actor_user.id != last_identity + if not identity_changed and not should_refresh_mcp_token(scope, actor_user.id): + logger.info("refresh_mcp_skipped_within_interval", run_id=run_id, user_id=actor_user.id) + return + if identity_changed: + logger.info( + "refresh_mcp_identity_transition", + run_id=run_id, + previous_user_id=last_identity, + user_id=actor_user.id, + ) access_token = create_oauth_access_token_for_run(task, task_run.state, scopes=scopes) except Exception as e: logger.warning("refresh_mcp_token_mint_failed", run_id=run_id, error=str(e)) @@ -260,6 +285,10 @@ def _refresh_sandbox_mcp( mcp_configs = mcp_configs + user_mcp_configs if not mcp_configs: + # Nothing to push means there is no MCP session to rebind — mark the + # window anyway so we don't re-mint a token on every message. + if actor_user is not None: + _mark_sandbox_session(scope, actor_user.id) logger.info("refresh_mcp_skipped_no_configs", run_id=run_id) return @@ -272,7 +301,8 @@ def _refresh_sandbox_mcp( timeout=REFRESH_TIMEOUT_SECONDS, ) if result.success: - mark_mcp_token_issued(run_id) + if actor_user is not None: + _mark_sandbox_session(scope, actor_user.id) logger.info("refresh_mcp_delivered", run_id=run_id, attempts=1) return @@ -290,7 +320,8 @@ def _refresh_sandbox_mcp( timeout=REFRESH_TIMEOUT_SECONDS, ) if retry.success: - mark_mcp_token_issued(run_id) + if actor_user is not None: + _mark_sandbox_session(scope, actor_user.id) logger.info("refresh_mcp_delivered", run_id=run_id, attempts=2) return 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 19c6581b8c71..134122d42d1b 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 @@ -29,6 +29,7 @@ get_task_run_credential_user, get_user_mcp_server_configs, mark_mcp_token_issued, + mark_sandbox_identity, ) from .get_task_processing_context import TaskProcessingContext @@ -164,6 +165,9 @@ class StartAgentServerOutput: @dataclass class _LaunchParams: mcp_configs: list[McpServerConfig] + # The user the boot-time MCP OAuth token was minted for, recorded as the + # sandbox's initial session identity once the agent server starts. + mcp_actor_user_id: int | None agentsh_domains: list[str] | None protected_base_branch: str | None event_ingest_token: str | None @@ -275,6 +279,7 @@ def _prepare_launch(ctx: TaskProcessingContext, scopes: PosthogMcpScopes) -> _La return _LaunchParams( mcp_configs=mcp_configs, + mcp_actor_user_id=actor_user.id if actor_user else None, agentsh_domains=agentsh_domains, protected_base_branch=protected_base_branch, event_ingest_token=event_ingest_token, @@ -316,10 +321,13 @@ def _invoke_start_agent_server( rtk_enabled=ctx.rtk_enabled, ) - # Mark startup-time token issuance so follow-ups within the next - # 30m window skip the redundant refresh. - if params.mcp_configs: - mark_mcp_token_issued(ctx.run_id) + # Record the sandbox's boot-time session identity and start its + # freshness window, so follow-ups from the same actor within + # MCP_TOKEN_REFRESH_INTERVAL_SECONDS skip the redundant refresh. + # Keyed on the sandbox id: a replacement sandbox starts unmarked. + if params.mcp_configs and params.mcp_actor_user_id is not None: + mark_mcp_token_issued(sandbox.id, params.mcp_actor_user_id) + mark_sandbox_identity(sandbox.id, "mcp", params.mcp_actor_user_id) # Persist the effective rtk posture the agent launched with, so terminal # analytics can cohort runs by it (the state override alone misses the 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 5ff23e9eaf7c..a0f7a1b4a656 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 @@ -16,19 +16,29 @@ from products.tasks.backend.temporal.process_task.utils import ( McpServerConfig, _mcp_token_issued_cache_key, + _sandbox_identity_cache_key, mark_mcp_token_issued, + mark_sandbox_identity, ) pytestmark = pytest.mark.django_db +_GATE_CACHE_KEYS = [ + _mcp_token_issued_cache_key("run-1", 42), + _mcp_token_issued_cache_key("run-1", 99), + _mcp_token_issued_cache_key("sb-2", 42), + _sandbox_identity_cache_key("run-1", "mcp"), + _sandbox_identity_cache_key("sb-2", "mcp"), +] + @pytest.fixture(autouse=True) def _clear_mcp_token_cache(): - """Ensure each test starts with no recorded token issuances so the - refresh gate doesn't carry state between tests.""" - cache.delete(_mcp_token_issued_cache_key("run-1")) + """Ensure each test starts with no recorded token issuances or session + identities so the refresh gate doesn't carry state between tests.""" + cache.delete_many(_GATE_CACHE_KEYS) yield - cache.delete(_mcp_token_issued_cache_key("run-1")) + cache.delete_many(_GATE_CACHE_KEYS) def _make_mcp_config(name: str = "posthog", token: str = "tok") -> McpServerConfig: @@ -235,15 +245,15 @@ def test_scopes_propagate_to_oauth_and_configs( class TestRefreshIntervalGate: """Refreshes within MCP_TOKEN_REFRESH_INTERVAL_SECONDS of a previous - successful issuance must be skipped without minting a new token or - contacting the sandbox.""" + successful issuance for the same actor must be skipped without minting a + new token or contacting the sandbox.""" @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.send_refresh_session") @patch( "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.create_oauth_access_token_for_run" ) def test_skipped_when_token_recently_issued(self, mock_oauth, mock_send_refresh): - mark_mcp_token_issued("run-1") + mark_mcp_token_issued("run-1", 42) _refresh_sandbox_mcp(_make_task_run_mock(), "read_only", auth_token=None) @@ -268,8 +278,10 @@ def test_marks_after_successful_refresh(self, mock_oauth, mock_ph_configs, mock_ _refresh_sandbox_mcp(_make_task_run_mock(), "read_only", auth_token=None) - # Cache entry now exists → next refresh within the interval is gated. - assert cache.get(_mcp_token_issued_cache_key("run-1")) is True + # Cache entries now exist → next refresh for this actor within the + # interval is gated, and the session identity is recorded. + assert cache.get(_mcp_token_issued_cache_key("run-1", 42)) is True + assert cache.get(_sandbox_identity_cache_key("run-1", "mcp")) == 42 @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.time.sleep") @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.send_refresh_session") @@ -295,7 +307,7 @@ def test_marks_after_successful_retry( _refresh_sandbox_mcp(_make_task_run_mock(), "read_only", auth_token=None) - assert cache.get(_mcp_token_issued_cache_key("run-1")) is True + assert cache.get(_mcp_token_issued_cache_key("run-1", 42)) is True @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.time.sleep") @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.send_refresh_session") @@ -319,7 +331,96 @@ def test_does_not_mark_after_two_failures( _refresh_sandbox_mcp(_make_task_run_mock(), "read_only", auth_token=None) # Cache stays empty so the next follow-up retries the dispatch. - assert cache.get(_mcp_token_issued_cache_key("run-1")) is None + assert cache.get(_mcp_token_issued_cache_key("run-1", 42)) is None + assert cache.get(_sandbox_identity_cache_key("run-1", "mcp")) is None + + +def _patch_actor(user_id: int): + """Pin the resolved credential user so tests can drive actor transitions + without building real run state.""" + return patch( + "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_task_run_credential_user", + return_value=MagicMock(id=user_id), + ) + + +@patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.send_refresh_session") +@patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_user_mcp_server_configs") +@patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_sandbox_ph_mcp_configs") +@patch( + "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.create_oauth_access_token_for_run" +) +class TestIdentityTransitionGate: + """An actor transition (run-state actor differs from the identity last + pushed to the sandbox) must bypass the freshness window and rebind the + session; the marks are keyed per sandbox so a replacement sandbox starts + unmarked.""" + + def _arm_success(self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh): + mock_oauth.return_value = "fresh-token" + mock_ph_configs.return_value = [_make_mcp_config(token="fresh-token")] + mock_user_configs.return_value = [] + mock_send_refresh.return_value = CommandResult(success=True, status_code=200) + + def test_actor_change_bypasses_freshness_window( + self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh + ): + self._arm_success(mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh) + # The creator's token is fresh and the session is bound to them… + mark_mcp_token_issued("run-1", 42) + mark_sandbox_identity("run-1", "mcp", 42) + + # …but the next message comes from a different actor. + with _patch_actor(99): + _refresh_sandbox_mcp(_make_task_run_mock(), "read_only", auth_token=None) + + mock_send_refresh.assert_called_once() + assert cache.get(_sandbox_identity_cache_key("run-1", "mcp")) == 99 + assert cache.get(_mcp_token_issued_cache_key("run-1", 99)) is True + + def test_switch_back_to_creator_refreshes_despite_fresh_window( + self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh + ): + self._arm_success(mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh) + # The creator's window is still warm, but the session was last bound + # to another user — the creator speaking again is a transition. + mark_mcp_token_issued("run-1", 42) + mark_sandbox_identity("run-1", "mcp", 99) + + with _patch_actor(42): + _refresh_sandbox_mcp(_make_task_run_mock(), "read_only", auth_token=None) + + mock_send_refresh.assert_called_once() + assert cache.get(_sandbox_identity_cache_key("run-1", "mcp")) == 42 + + def test_same_actor_within_window_is_skipped( + self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh + ): + self._arm_success(mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh) + mark_mcp_token_issued("run-1", 99) + mark_sandbox_identity("run-1", "mcp", 99) + + with _patch_actor(99): + _refresh_sandbox_mcp(_make_task_run_mock(), "read_only", auth_token=None) + + mock_oauth.assert_not_called() + mock_send_refresh.assert_not_called() + + def test_replacement_sandbox_starts_unmarked( + self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh + ): + self._arm_success(mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh) + # Marks recorded against the run id (legacy scope) must not gate a run + # whose state now points at a concrete sandbox. + mark_mcp_token_issued("run-1", 42) + mark_sandbox_identity("run-1", "mcp", 42) + + with _patch_actor(42): + _refresh_sandbox_mcp(_make_task_run_mock(state={"sandbox_id": "sb-2"}), "read_only", auth_token=None) + + mock_send_refresh.assert_called_once() + assert cache.get(_mcp_token_issued_cache_key("sb-2", 42)) is True + assert cache.get(_sandbox_identity_cache_key("sb-2", "mcp")) == 42 class TestSendFollowupActivityRefreshOrdering: diff --git a/products/tasks/backend/temporal/process_task/utils.py b/products/tasks/backend/temporal/process_task/utils.py index c7f3632f69f8..9dfb89341417 100644 --- a/products/tasks/backend/temporal/process_task/utils.py +++ b/products/tasks/backend/temporal/process_task/utils.py @@ -3,7 +3,7 @@ import logging from dataclasses import dataclass, field from enum import StrEnum -from typing import TYPE_CHECKING, Any, Optional +from typing import TYPE_CHECKING, Any, Literal, Optional from urllib.parse import urlparse from django.conf import settings @@ -356,23 +356,68 @@ def get_sandbox_snapshot_metadata(snapshot: SandboxSnapshot) -> SnapshotMetadata MCP_TOKEN_REFRESH_INTERVAL_SECONDS = TOKEN_EXPIRATION_SECONDS / 2 # 3 hours -def _mcp_token_issued_cache_key(run_id: str) -> str: - return f"posthog_ai:task-run-mcp-token-issued:{run_id}" +def sandbox_identity_scope(run_id: str, state: dict[str, Any] | None) -> str: + """Cache scope for the marks describing what a run's sandbox holds. + + The freshness and identity marks below describe the state of a *sandbox*, + so they key on the sandbox id: a replacement sandbox (fresh provision, + snapshot restore, mid-run workflow retry) starts unmarked by construction + and therefore falls back to the boot-time default — nothing ever needs + clearing. Falls back to the run id for runs that haven't recorded a + sandbox id in their state yet. + """ + return (state or {}).get("sandbox_id") or run_id + + +def _mcp_token_issued_cache_key(scope: str, user_id: int) -> str: + return f"posthog_ai:sandbox-mcp-token-issued:{scope}:{user_id}" + + +def mark_mcp_token_issued(scope: str, user_id: int) -> None: + """Record that a fresh MCP token for ``user_id`` was issued to the sandbox. + + ``scope`` comes from ``sandbox_identity_scope``. The entry self-expires + after MCP_TOKEN_REFRESH_INTERVAL_SECONDS, so ``should_refresh_mcp_token`` + returns True again past that window. + """ + get_tasks_cache().set(_mcp_token_issued_cache_key(scope, user_id), True, timeout=MCP_TOKEN_REFRESH_INTERVAL_SECONDS) + + +def should_refresh_mcp_token(scope: str, user_id: int) -> bool: + """True when no MCP token for ``user_id`` was issued to the sandbox within + the last MCP_TOKEN_REFRESH_INTERVAL_SECONDS window.""" + return get_tasks_cache().get(_mcp_token_issued_cache_key(scope, user_id)) is None + + +# How long the sandbox's session identity is remembered — comfortably past any +# plausible sandbox lifetime. On eviction the identity is assumed to be the +# boot-time one (the task creator). +SANDBOX_IDENTITY_TTL_SECONDS = 7 * 24 * 60 * 60 + +SandboxIdentityKind = Literal["mcp"] + + +def _sandbox_identity_cache_key(scope: str, kind: SandboxIdentityKind) -> str: + return f"posthog_ai:sandbox-{kind}-identity:{scope}" -def mark_mcp_token_issued(run_id: str) -> None: - """Record that a fresh MCP token was issued to the sandbox for this run. +def mark_sandbox_identity(scope: str, kind: SandboxIdentityKind, value: int | str) -> None: + """Record which identity the sandbox's live session currently holds. - The cache entry self-expires after MCP_TOKEN_REFRESH_INTERVAL_SECONDS, so - `should_refresh_mcp_token` returns True again past that window. + ``mcp`` stores the user id the last-pushed OAuth token was minted for. + The run-state actor (``get_task_run_credential_user``) says who *should* + be acting; this mark says whose token the session actually has — the gap + between the two is an identity transition that must bypass the freshness + rate limit. ``scope`` comes from ``sandbox_identity_scope``, so a + replacement sandbox starts unmarked. """ - get_tasks_cache().set(_mcp_token_issued_cache_key(run_id), True, timeout=MCP_TOKEN_REFRESH_INTERVAL_SECONDS) + get_tasks_cache().set(_sandbox_identity_cache_key(scope, kind), value, timeout=SANDBOX_IDENTITY_TTL_SECONDS) -def should_refresh_mcp_token(run_id: str) -> bool: - """Return True if no MCP token has been issued for this run within the - last MCP_TOKEN_REFRESH_INTERVAL_SECONDS window.""" - return get_tasks_cache().get(_mcp_token_issued_cache_key(run_id)) is None +def get_last_sandbox_identity(scope: str, kind: SandboxIdentityKind) -> int | str | None: + """Return the identity the sandbox session was last bound to for a kind, + or None when unknown (never pushed, or the entry was evicted).""" + return get_tasks_cache().get(_sandbox_identity_cache_key(scope, kind)) @dataclass(frozen=True) From b3c85c20484cb906f3c15173e691b79463c2520c Mon Sep 17 00:00:00 2001 From: Vojta Bartos Date: Tue, 14 Jul 2026 13:57:14 +0200 Subject: [PATCH 2/5] refactor(tasks): rename mcp_actor_user_id to actor_user_id The field records whose credentials the sandbox booted with; the MCP prefix undersells it as more surfaces (GitHub) adopt the same mark. --- .../process_task/activities/start_agent_server.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) 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 134122d42d1b..76c56b8b79ba 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 @@ -165,9 +165,9 @@ class StartAgentServerOutput: @dataclass class _LaunchParams: mcp_configs: list[McpServerConfig] - # The user the boot-time MCP OAuth token was minted for, recorded as the + # The user the boot-time credentials were minted for, recorded as the # sandbox's initial session identity once the agent server starts. - mcp_actor_user_id: int | None + actor_user_id: int | None agentsh_domains: list[str] | None protected_base_branch: str | None event_ingest_token: str | None @@ -279,7 +279,7 @@ def _prepare_launch(ctx: TaskProcessingContext, scopes: PosthogMcpScopes) -> _La return _LaunchParams( mcp_configs=mcp_configs, - mcp_actor_user_id=actor_user.id if actor_user else None, + actor_user_id=actor_user.id if actor_user else None, agentsh_domains=agentsh_domains, protected_base_branch=protected_base_branch, event_ingest_token=event_ingest_token, @@ -325,9 +325,9 @@ def _invoke_start_agent_server( # freshness window, so follow-ups from the same actor within # MCP_TOKEN_REFRESH_INTERVAL_SECONDS skip the redundant refresh. # Keyed on the sandbox id: a replacement sandbox starts unmarked. - if params.mcp_configs and params.mcp_actor_user_id is not None: - mark_mcp_token_issued(sandbox.id, params.mcp_actor_user_id) - mark_sandbox_identity(sandbox.id, "mcp", params.mcp_actor_user_id) + if params.mcp_configs and params.actor_user_id is not None: + mark_mcp_token_issued(sandbox.id, params.actor_user_id) + mark_sandbox_identity(sandbox.id, "mcp", params.actor_user_id) # Persist the effective rtk posture the agent launched with, so terminal # analytics can cohort runs by it (the state override alone misses the From f01338ba408a7b91ea0b100555f03b86b2e17708 Mon Sep 17 00:00:00 2001 From: Vojta Bartos Date: Tue, 14 Jul 2026 14:10:00 +0200 Subject: [PATCH 3/5] refactor(tasks): use tasks: prefix for sandbox identity cache keys The marks belong to the tasks product, matching the existing tasks: key convention (e.g. the GitHub token rotation lock), not posthog_ai. --- products/tasks/backend/temporal/process_task/utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/products/tasks/backend/temporal/process_task/utils.py b/products/tasks/backend/temporal/process_task/utils.py index 9dfb89341417..0980d32aac41 100644 --- a/products/tasks/backend/temporal/process_task/utils.py +++ b/products/tasks/backend/temporal/process_task/utils.py @@ -370,7 +370,7 @@ def sandbox_identity_scope(run_id: str, state: dict[str, Any] | None) -> str: def _mcp_token_issued_cache_key(scope: str, user_id: int) -> str: - return f"posthog_ai:sandbox-mcp-token-issued:{scope}:{user_id}" + return f"tasks:sandbox-mcp-token-issued:{scope}:{user_id}" def mark_mcp_token_issued(scope: str, user_id: int) -> None: @@ -398,7 +398,7 @@ def should_refresh_mcp_token(scope: str, user_id: int) -> bool: def _sandbox_identity_cache_key(scope: str, kind: SandboxIdentityKind) -> str: - return f"posthog_ai:sandbox-{kind}-identity:{scope}" + return f"tasks:sandbox-{kind}-identity:{scope}" def mark_sandbox_identity(scope: str, kind: SandboxIdentityKind, value: int | str) -> None: From 1b397f92be75163e5695fa419f82f340c019819f Mon Sep 17 00:00:00 2001 From: Vojta Bartos Date: Tue, 14 Jul 2026 12:13:55 +0200 Subject: [PATCH 4/5] feat(tasks): rebind sandbox GitHub identity on actor transitions Unify the per-message identity reconciliation for MCP and GitHub in one gate (ensure_sandbox_identity), called right before a follow-up turn is delivered. The MCP surface keeps its transition-or-freshness semantics; the GitHub surface rebinds on actor transitions only, re-injecting the actor's GitHub token and git author into the live sandbox and notifying the agent-server with the effective PR authorship. Token TTL between messages stays owned by the credential-refresh loop. Also fix the credential-refresh loop reading the workflow-start snapshot of the run state: it now re-reads the state live, so cadence refreshes mint for the run's current actor instead of the boot-time one. --- .../temporal/execute_sandbox/workflow.py | 1 + .../activities/refresh_sandbox_credentials.py | 11 +- .../activities/send_followup_to_sandbox.py | 151 +----- .../activities/start_agent_server.py | 16 +- .../temporal/process_task/sandbox_identity.py | 285 ++++++++++ .../tests/test_sandbox_identity.py | 489 ++++++++++++++++++ .../tests/test_send_followup_to_sandbox.py | 405 +-------------- .../backend/temporal/process_task/utils.py | 5 +- .../backend/temporal/process_task/workflow.py | 1 + 9 files changed, 822 insertions(+), 542 deletions(-) create mode 100644 products/tasks/backend/temporal/process_task/sandbox_identity.py create mode 100644 products/tasks/backend/temporal/process_task/tests/test_sandbox_identity.py diff --git a/products/tasks/backend/temporal/execute_sandbox/workflow.py b/products/tasks/backend/temporal/execute_sandbox/workflow.py index 0e2ebaff7907..5855a2b5ab09 100644 --- a/products/tasks/backend/temporal/execute_sandbox/workflow.py +++ b/products/tasks/backend/temporal/execute_sandbox/workflow.py @@ -1203,6 +1203,7 @@ async def _send_followup_to_sandbox(self, message: str | None, artifact_ids: lis posthog_mcp_scopes=self._posthog_mcp_scopes, artifact_ids=artifact_ids, message_id=str(workflow.uuid4()), + context=self.context, ), start_to_close_timeout=timedelta(minutes=35), # See process_task: heartbeat detects worker restarts, message_id diff --git a/products/tasks/backend/temporal/process_task/activities/refresh_sandbox_credentials.py b/products/tasks/backend/temporal/process_task/activities/refresh_sandbox_credentials.py index 9adab94feea9..4279d6938923 100644 --- a/products/tasks/backend/temporal/process_task/activities/refresh_sandbox_credentials.py +++ b/products/tasks/backend/temporal/process_task/activities/refresh_sandbox_credentials.py @@ -1,4 +1,4 @@ -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from temporalio import activity @@ -98,6 +98,15 @@ def refresh_sandbox_credentials(input: RefreshSandboxCredentialsInput) -> Refres except Task.DoesNotExist as e: raise TaskNotFoundError(f"Task {ctx.task_id} not found", {"task_id": ctx.task_id}, cause=e) + # The workflow-start context carries a boot-time snapshot of the run + # state, but credential resolution must follow the *current* actor + # (multiplayer follow-ups update it mid-run) — re-read it live. + try: + live_state = TaskRun.objects.values_list("state", flat=True).get(id=ctx.run_id) + ctx = replace(ctx, state=live_state) + except TaskRun.DoesNotExist: + logger.warning("sandbox_credentials_refresh_run_missing", run_id=ctx.run_id) + refreshed_kinds: list[str] = [] orphaned_kinds: list[str] = [] next_refresh = DEFAULT_REFRESH_INTERVAL_SECONDS 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 5627e723b911..38051ff5e6ae 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 @@ -1,5 +1,4 @@ import json -import time import threading import contextvars from dataclasses import dataclass @@ -12,38 +11,24 @@ from posthog.temporal.common.utils import close_db_connections from posthog.temporal.oauth import PosthogMcpScopes -from products.tasks.backend.logic.services.agent_command import ( - FOLLOWUP_TIMEOUT_SECONDS, - REFRESH_TIMEOUT_SECONDS, - CommandResult, - send_refresh_session, - send_user_message, -) +from products.tasks.backend.logic.services.agent_command import FOLLOWUP_TIMEOUT_SECONDS, send_user_message from products.tasks.backend.logic.services.connection_token import create_sandbox_connection_token from products.tasks.backend.logic.services.staged_artifacts import get_task_run_artifacts_by_id from products.tasks.backend.logic.stream.redis_stream import get_task_run_stream_key from products.tasks.backend.models import TaskRun from products.tasks.backend.redis import get_tasks_stream_redis_sync, run_uses_dedicated_stream -from products.tasks.backend.temporal.oauth import create_oauth_access_token_for_run +from products.tasks.backend.temporal.process_task.activities.get_task_processing_context import TaskProcessingContext +from products.tasks.backend.temporal.process_task.sandbox_identity import ensure_sandbox_identity from products.tasks.backend.temporal.process_task.utils import ( get_actor_distinct_id, - get_last_sandbox_identity, - get_sandbox_ph_mcp_configs, get_task_run_credential_user, - get_user_mcp_server_configs, is_slack_interaction_state, - mark_mcp_token_issued, - mark_sandbox_identity, - sandbox_identity_scope, - should_refresh_mcp_token, ) from ee.hogai.sandbox import STOP_REASON_END_TURN, TURN_COMPLETE_METHOD logger = structlog.get_logger(__name__) -REFRESH_RETRY_DELAY_SECONDS = 0.5 - # Retries exist for attempt-level deaths (worker restart kills the in-flight # attempt, detected via heartbeat timeout) and for delivery-unknown failures. # Application failures that write an error sentinel raise non-retryable. @@ -60,6 +45,11 @@ class SendFollowupToSandboxInput: # Workflow-generated idempotency key. Stable across activity retries, so # the agent-server can drop a redelivery of a message it already accepted. message_id: str | None = None + # The run's processing context, needed to rebind GitHub credentials on an + # actor transition. Optional so activities scheduled by pre-rollout + # workflow histories (without it) still deserialize; they skip the GitHub + # surface until the run's next workflow. + context: TaskProcessingContext | None = None @activity.defn @@ -131,10 +121,15 @@ def _deliver_followup(input: SendFollowupToSandboxInput) -> None: task_run, user_id=actor_user.id, distinct_id=get_actor_distinct_id(actor_user) ) - # Push a fresh MCP config before the turn so the agent-server rebinds its - # ACP session to a non-stale OAuth token. Non-fatal: if refresh fails we + # Rebind the sandbox session (MCP OAuth token, GitHub credentials) to this + # message's actor before the turn runs. Non-fatal: if the rebind fails we # still deliver the follow-up with the existing (possibly stale) creds. - _refresh_sandbox_mcp(task_run, input.posthog_mcp_scopes, auth_token) + ensure_sandbox_identity( + task_run, + posthog_mcp_scopes=input.posthog_mcp_scopes, + auth_token=auth_token, + processing_context=input.context, + ) artifacts = None artifact_ids = input.artifact_ids or [] if artifact_ids: @@ -219,120 +214,6 @@ def _deliver_followup(input: SendFollowupToSandboxInput) -> None: raise ApplicationError(f"send_followup failed: {error_msg}", non_retryable=True) -def _mark_sandbox_session(scope: str, user_id: int) -> None: - """Record whose token the sandbox session now holds and start that user's - freshness window.""" - mark_mcp_token_issued(scope, user_id) - mark_sandbox_identity(scope, "mcp", user_id) - - -def _refresh_sandbox_mcp( - task_run: TaskRun, - scopes: PosthogMcpScopes, - auth_token: str | None, -) -> None: - """Mint a fresh OAuth token and push updated MCP configs to the sandbox. - - Best-effort: retries once on failure, then logs and returns. Never raises - — a failed refresh should not block an otherwise-valid follow-up. - - Skipped only when the sandbox's live session already holds a fresh token - for *this message's actor*: the freshness window is keyed per - (sandbox, user), and an actor transition (the run-state actor differs from - the identity last pushed to the sandbox) bypasses the window entirely, so - the session rebinds before the new speaker's turn is delivered. - """ - run_id = str(task_run.id) - task = task_run.task - scope = sandbox_identity_scope(run_id, task_run.state) - try: - actor_user = get_task_run_credential_user(task, task_run.state) - if actor_user is not None: - # Until a refresh records otherwise, the sandbox holds the - # boot-time token, which was minted for the task creator. - last_identity = get_last_sandbox_identity(scope, "mcp") or task.created_by_id - identity_changed = actor_user.id != last_identity - if not identity_changed and not should_refresh_mcp_token(scope, actor_user.id): - logger.info("refresh_mcp_skipped_within_interval", run_id=run_id, user_id=actor_user.id) - return - if identity_changed: - logger.info( - "refresh_mcp_identity_transition", - run_id=run_id, - previous_user_id=last_identity, - user_id=actor_user.id, - ) - access_token = create_oauth_access_token_for_run(task, task_run.state, scopes=scopes) - except Exception as e: - logger.warning("refresh_mcp_token_mint_failed", run_id=run_id, error=str(e)) - return - - mcp_configs = get_sandbox_ph_mcp_configs( - token=access_token, - project_id=task_run.team_id, - scopes=scopes, - interaction_origin=(task_run.state or {}).get("interaction_origin"), - task_id=str(task_run.task_id), - ) - if actor_user and actor_user.id: - user_mcp_configs = get_user_mcp_server_configs( - token=access_token, - team_id=task_run.team_id, - user_id=actor_user.id, - interaction_origin=(task_run.state or {}).get("interaction_origin"), - ) - if user_mcp_configs: - mcp_configs = mcp_configs + user_mcp_configs - - if not mcp_configs: - # Nothing to push means there is no MCP session to rebind — mark the - # window anyway so we don't re-mint a token on every message. - if actor_user is not None: - _mark_sandbox_session(scope, actor_user.id) - logger.info("refresh_mcp_skipped_no_configs", run_id=run_id) - return - - mcp_servers = [config.to_dict() for config in mcp_configs] - - result = send_refresh_session( - task_run, - mcp_servers, - auth_token=auth_token, - timeout=REFRESH_TIMEOUT_SECONDS, - ) - if result.success: - if actor_user is not None: - _mark_sandbox_session(scope, actor_user.id) - logger.info("refresh_mcp_delivered", run_id=run_id, attempts=1) - return - - logger.info( - "refresh_mcp_retrying", - run_id=run_id, - error=result.error, - status_code=result.status_code, - ) - time.sleep(REFRESH_RETRY_DELAY_SECONDS) - retry: CommandResult = send_refresh_session( - task_run, - mcp_servers, - auth_token=auth_token, - timeout=REFRESH_TIMEOUT_SECONDS, - ) - if retry.success: - if actor_user is not None: - _mark_sandbox_session(scope, actor_user.id) - logger.info("refresh_mcp_delivered", run_id=run_id, attempts=2) - return - - logger.warning( - "refresh_mcp_failed", - run_id=run_id, - error=retry.error, - status_code=retry.status_code, - ) - - def _get_stop_reason(result_data: dict[str, Any] | None) -> str: if not isinstance(result_data, dict): return STOP_REASON_END_TURN 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 76c56b8b79ba..39cad165be42 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 @@ -165,8 +165,9 @@ class StartAgentServerOutput: @dataclass class _LaunchParams: mcp_configs: list[McpServerConfig] - # The user the boot-time credentials were minted for, recorded as the - # sandbox's initial session identity once the agent server starts. + # The user the boot-time credentials (MCP OAuth token, GitHub token) were + # minted for, recorded as the sandbox's initial session identities once + # the agent server starts. actor_user_id: int | None agentsh_domains: list[str] | None protected_base_branch: str | None @@ -321,13 +322,16 @@ def _invoke_start_agent_server( rtk_enabled=ctx.rtk_enabled, ) - # Record the sandbox's boot-time session identity and start its + # Record the sandbox's boot-time session identities and start the MCP # freshness window, so follow-ups from the same actor within # MCP_TOKEN_REFRESH_INTERVAL_SECONDS skip the redundant refresh. # Keyed on the sandbox id: a replacement sandbox starts unmarked. - if params.mcp_configs and params.actor_user_id is not None: - mark_mcp_token_issued(sandbox.id, params.actor_user_id) - mark_sandbox_identity(sandbox.id, "mcp", params.actor_user_id) + if params.actor_user_id is not None: + if params.mcp_configs: + mark_mcp_token_issued(sandbox.id, params.actor_user_id) + mark_sandbox_identity(sandbox.id, "mcp", params.actor_user_id) + if ctx.has_github_credentials: + mark_sandbox_identity(sandbox.id, "github", params.actor_user_id) # Persist the effective rtk posture the agent launched with, so terminal # analytics can cohort runs by it (the state override alone misses the diff --git a/products/tasks/backend/temporal/process_task/sandbox_identity.py b/products/tasks/backend/temporal/process_task/sandbox_identity.py new file mode 100644 index 000000000000..2f7bffa977c5 --- /dev/null +++ b/products/tasks/backend/temporal/process_task/sandbox_identity.py @@ -0,0 +1,285 @@ +"""Reconciles the sandbox's live session identities to the run's current actor. + +A sandbox holds two per-user credential surfaces: the ACP session's MCP OAuth +token and the git/gh credentials (token, author, PR authorship) frozen into the +workspace. Each is minted for one user at a time, while the run's actor can +change between messages (multiplayer Slack threads). ``ensure_sandbox_identity`` +is the single choke point where, right before a follow-up turn is delivered, +every surface is compared against the identity last pushed to this sandbox and +rebound when the speaker changed. + +Freshness differs per surface: the MCP token has no other refresh path, so its +TTL window is enforced here too; GitHub tokens are kept alive between messages +by the credential-refresh loop, so the GitHub gate only reacts to identity +transitions. +""" + +import time +from dataclasses import replace +from typing import TYPE_CHECKING + +import structlog + +from posthog.temporal.oauth import PosthogMcpScopes + +from products.tasks.backend.logic.services.agent_command import ( + REFRESH_TIMEOUT_SECONDS, + CommandResult, + send_refresh_session, +) +from products.tasks.backend.logic.services.sandbox import Sandbox +from products.tasks.backend.models import TaskRun +from products.tasks.backend.temporal.oauth import create_oauth_access_token_for_run +from products.tasks.backend.temporal.process_task.sandbox_credentials import ( + GitHubSandboxCredential, + update_sandbox_env_file, +) +from products.tasks.backend.temporal.process_task.utils import ( + get_git_identity_env_vars, + get_last_sandbox_identity, + get_sandbox_ph_mcp_configs, + get_task_run_credential_user, + get_user_mcp_server_configs, + mark_mcp_token_issued, + mark_sandbox_identity, + sandbox_identity_scope, + should_refresh_mcp_token, +) + +if TYPE_CHECKING: + from posthog.models.user import User + + from products.tasks.backend.temporal.process_task.activities.get_task_processing_context import ( + TaskProcessingContext, + ) + +logger = structlog.get_logger(__name__) + +REFRESH_RETRY_DELAY_SECONDS = 0.5 + + +def ensure_sandbox_identity( + task_run: TaskRun, + *, + posthog_mcp_scopes: PosthogMcpScopes, + auth_token: str | None, + processing_context: "TaskProcessingContext | None" = None, +) -> None: + """Rebind the sandbox's live session to the run's current actor. + + Best-effort and never raises: a failed rebind must not block an + otherwise-valid follow-up, and an unmarked failure is retried on the next + message. The GitHub surface is only reconciled when the caller supplies + the run's processing context (older in-flight workflows don't). + """ + run_id = str(task_run.id) + scope = sandbox_identity_scope(run_id, task_run.state) + try: + actor_user = get_task_run_credential_user(task_run.task, task_run.state) + except Exception as e: + logger.warning("sandbox_identity_actor_resolution_failed", run_id=run_id, error=str(e)) + return + + try: + _ensure_mcp_identity(task_run, actor_user, scope, posthog_mcp_scopes, auth_token) + except Exception: + logger.warning("sandbox_identity_reconcile_failed", kind="mcp", run_id=run_id, exc_info=True) + if actor_user is not None and processing_context is not None: + try: + _ensure_github_identity(task_run, actor_user, scope, processing_context, auth_token) + except Exception: + logger.warning("sandbox_identity_reconcile_failed", kind="github", run_id=run_id, exc_info=True) + + +def _ensure_mcp_identity( + task_run: TaskRun, + actor_user: "User | None", + scope: str, + scopes: PosthogMcpScopes, + auth_token: str | None, +) -> None: + """Skip only when the session already holds a fresh token for this actor: + the freshness window is keyed per (sandbox, user), and an actor transition + bypasses the window entirely. A missing actor (Slack fail-closed) falls + through to the mint, which surfaces the standard warning.""" + run_id = str(task_run.id) + if actor_user is not None: + # Until a refresh records otherwise, the sandbox holds the boot-time + # token, which was minted for the task creator. + last_identity = get_last_sandbox_identity(scope, "mcp") or task_run.task.created_by_id + identity_changed = actor_user.id != last_identity + if not identity_changed and not should_refresh_mcp_token(scope, actor_user.id): + logger.info("refresh_mcp_skipped_within_interval", run_id=run_id, user_id=actor_user.id) + return + if identity_changed: + logger.info( + "refresh_mcp_identity_transition", + run_id=run_id, + previous_user_id=last_identity, + user_id=actor_user.id, + ) + _rebind_mcp(task_run, actor_user, scope, scopes, auth_token) + + +def _mark_mcp_session(scope: str, user_id: int) -> None: + """Record whose token the MCP session now holds and start that user's + freshness window.""" + mark_mcp_token_issued(scope, user_id) + mark_sandbox_identity(scope, "mcp", user_id) + + +def _rebind_mcp( + task_run: TaskRun, + actor_user: "User | None", + scope: str, + scopes: PosthogMcpScopes, + auth_token: str | None, +) -> None: + """Mint a fresh OAuth token for the actor and push updated MCP configs to + the sandbox. Retries once on failure, then logs and returns.""" + run_id = str(task_run.id) + task = task_run.task + try: + access_token = create_oauth_access_token_for_run(task, task_run.state, scopes=scopes) + except Exception as e: + logger.warning("refresh_mcp_token_mint_failed", run_id=run_id, error=str(e)) + return + + mcp_configs = get_sandbox_ph_mcp_configs( + token=access_token, + project_id=task_run.team_id, + scopes=scopes, + interaction_origin=(task_run.state or {}).get("interaction_origin"), + task_id=str(task_run.task_id), + ) + if actor_user and actor_user.id: + user_mcp_configs = get_user_mcp_server_configs( + token=access_token, + team_id=task_run.team_id, + user_id=actor_user.id, + interaction_origin=(task_run.state or {}).get("interaction_origin"), + ) + if user_mcp_configs: + mcp_configs = mcp_configs + user_mcp_configs + + if not mcp_configs: + # Nothing to push means there is no MCP session to rebind — mark the + # window anyway so we don't re-mint a token on every message. + if actor_user is not None: + _mark_mcp_session(scope, actor_user.id) + logger.info("refresh_mcp_skipped_no_configs", run_id=run_id) + return + + mcp_servers = [config.to_dict() for config in mcp_configs] + + result = send_refresh_session( + task_run, + mcp_servers, + auth_token=auth_token, + timeout=REFRESH_TIMEOUT_SECONDS, + ) + if result.success: + if actor_user is not None: + _mark_mcp_session(scope, actor_user.id) + logger.info("refresh_mcp_delivered", run_id=run_id, attempts=1) + return + + logger.info( + "refresh_mcp_retrying", + run_id=run_id, + error=result.error, + status_code=result.status_code, + ) + time.sleep(REFRESH_RETRY_DELAY_SECONDS) + retry: CommandResult = send_refresh_session( + task_run, + mcp_servers, + auth_token=auth_token, + timeout=REFRESH_TIMEOUT_SECONDS, + ) + if retry.success: + if actor_user is not None: + _mark_mcp_session(scope, actor_user.id) + logger.info("refresh_mcp_delivered", run_id=run_id, attempts=2) + return + + logger.warning( + "refresh_mcp_failed", + run_id=run_id, + error=retry.error, + status_code=retry.status_code, + ) + + +def _ensure_github_identity( + task_run: TaskRun, + actor_user: "User", + scope: str, + processing_context: "TaskProcessingContext", + auth_token: str | None, +) -> None: + """Transition-only gate: token TTL between messages is owned by the + credential-refresh loop, so this fires only when the speaker changed.""" + run_id = str(task_run.id) + if not processing_context.has_github_credentials: + return + last_identity = get_last_sandbox_identity(scope, "github") or task_run.task.created_by_id + if actor_user.id == last_identity: + return + logger.info( + "refresh_github_identity_transition", + run_id=run_id, + previous_user_id=last_identity, + user_id=actor_user.id, + ) + if _rebind_github(task_run, processing_context, auth_token): + mark_sandbox_identity(scope, "github", actor_user.id) + + +def _rebind_github( + task_run: TaskRun, + processing_context: "TaskProcessingContext", + auth_token: str | None, +) -> bool: + """Re-inject the actor's GitHub credentials and git author into the live + sandbox. Returns True when the sandbox now reflects the actor (or there is + nothing to rebind), False when the rebind should be retried next message.""" + run_id = str(task_run.id) + task = task_run.task + sandbox_id = (task_run.state or {}).get("sandbox_id") + if not sandbox_id: + return False + + # The workflow-start context carries a boot-time snapshot of the run + # state; credential resolution must see the current actor. + live_context = replace(processing_context, state=task_run.state) + try: + sandbox = Sandbox.get_by_id(sandbox_id) + outcome = GitHubSandboxCredential().refresh(sandbox, live_context, task) + if outcome.refreshed: + git_identity = get_git_identity_env_vars(task, task_run.state) + if git_identity: + update_sandbox_env_file(sandbox, git_identity) + except Exception: + logger.warning("refresh_github_identity_failed", run_id=run_id, exc_info=True) + return False + + if not outcome.refreshed: + # No refreshable GitHub credential in play (e.g. caller-token run) — + # nothing to diverge on; mark so we don't retry every message. + return True + + authorship = (task_run.state or {}).get("pr_authorship_mode") + notify = send_refresh_session( + task_run, + [], + auth_token=auth_token, + refreshed_credentials=["github"], + authorship=authorship, + ) + if not notify.success: + # Credentials already landed in the sandbox; the notification only + # feeds the agent-server's debug log. + logger.info("refresh_github_notify_failed", run_id=run_id, error=notify.error) + logger.info("refresh_github_delivered", run_id=run_id) + return True diff --git a/products/tasks/backend/temporal/process_task/tests/test_sandbox_identity.py b/products/tasks/backend/temporal/process_task/tests/test_sandbox_identity.py new file mode 100644 index 000000000000..e311bfcdfbf7 --- /dev/null +++ b/products/tasks/backend/temporal/process_task/tests/test_sandbox_identity.py @@ -0,0 +1,489 @@ +import pytest +from unittest.mock import MagicMock, patch + +from django.core.cache import cache + +from products.tasks.backend.logic.services.agent_command import CommandResult +from products.tasks.backend.temporal.process_task.activities.get_task_processing_context import TaskProcessingContext +from products.tasks.backend.temporal.process_task.sandbox_credentials import CredentialRefreshOutcome +from products.tasks.backend.temporal.process_task.sandbox_identity import ( + REFRESH_RETRY_DELAY_SECONDS, + ensure_sandbox_identity, +) +from products.tasks.backend.temporal.process_task.utils import ( + McpServerConfig, + _mcp_token_issued_cache_key, + _sandbox_identity_cache_key, + mark_mcp_token_issued, + mark_sandbox_identity, +) + +pytestmark = pytest.mark.django_db + +_GATE_CACHE_KEYS = [ + _mcp_token_issued_cache_key("run-1", 42), + _mcp_token_issued_cache_key("run-1", 99), + _mcp_token_issued_cache_key("sb-2", 42), + _mcp_token_issued_cache_key("sb-2", 99), + _sandbox_identity_cache_key("run-1", "mcp"), + _sandbox_identity_cache_key("sb-2", "mcp"), + _sandbox_identity_cache_key("run-1", "github"), + _sandbox_identity_cache_key("sb-2", "github"), +] + + +@pytest.fixture(autouse=True) +def _clear_identity_cache(): + """Ensure each test starts with no recorded token issuances or session + identities so the gates don't carry state between tests.""" + cache.delete_many(_GATE_CACHE_KEYS) + yield + cache.delete_many(_GATE_CACHE_KEYS) + + +def _make_mcp_config(name: str = "posthog", token: str = "tok") -> McpServerConfig: + return McpServerConfig( + type="http", + name=name, + url="https://mcp.posthog.com/mcp", + headers=[{"name": "Authorization", "value": f"Bearer {token}"}], + ) + + +def _make_task_run_mock(team_id: int = 7, created_by_id: int | None = 42, state: dict | None = None) -> MagicMock: + task = MagicMock() + task.created_by_id = created_by_id + if created_by_id is not None: + task.created_by = MagicMock(id=created_by_id, distinct_id=f"user-{created_by_id}") + else: + task.created_by = None + task_run = MagicMock() + task_run.id = "run-1" + task_run.team_id = team_id + task_run.task = task + task_run.task_id = "task-1" + # Default to None so `(task_run.state or {}).get(...)` returns None cleanly. + # MagicMock auto-attributes would otherwise return further MagicMock objects + # and leak into kwargs passed to `get_sandbox_ph_mcp_configs`. + task_run.state = state + return task_run + + +def _make_processing_context(**overrides: object) -> TaskProcessingContext: + defaults: dict = { + "task_id": "task-1", + "run_id": "run-1", + "team_id": 7, + "team_uuid": "team-uuid", + "organization_id": "org-id", + "github_integration_id": 11, + "repository": "posthog/example", + "distinct_id": "d-1", + } + defaults.update(overrides) + return TaskProcessingContext(**defaults) + + +def _ensure(task_run, scopes="read_only", auth_token=None, processing_context=None) -> None: + ensure_sandbox_identity( + task_run, + posthog_mcp_scopes=scopes, + auth_token=auth_token, + processing_context=processing_context, + ) + + +def _patch_actor(user_id: int): + """Pin the resolved credential user so tests can drive actor transitions + without building real run state.""" + return patch( + "products.tasks.backend.temporal.process_task.sandbox_identity.get_task_run_credential_user", + return_value=MagicMock(id=user_id), + ) + + +class TestRebindMcp: + @patch("products.tasks.backend.temporal.process_task.sandbox_identity.send_refresh_session") + @patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_user_mcp_server_configs") + @patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_sandbox_ph_mcp_configs") + @patch("products.tasks.backend.temporal.process_task.sandbox_identity.create_oauth_access_token_for_run") + def test_success_path_single_call(self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh): + mock_oauth.return_value = "fresh-token" + mock_ph_configs.return_value = [_make_mcp_config(token="fresh-token")] + mock_user_configs.return_value = [] + mock_send_refresh.return_value = CommandResult(success=True, status_code=200) + + task_run = _make_task_run_mock() + _ensure(task_run, auth_token="jwt") + + mock_oauth.assert_called_once_with(task_run.task, task_run.state, scopes="read_only") + mock_ph_configs.assert_called_once_with( + token="fresh-token", project_id=7, scopes="read_only", interaction_origin=None, task_id="task-1" + ) + mock_user_configs.assert_called_once_with(token="fresh-token", team_id=7, user_id=42, interaction_origin=None) + mock_send_refresh.assert_called_once() + _, kwargs = mock_send_refresh.call_args + assert kwargs["auth_token"] == "jwt" + assert mock_send_refresh.call_args.args[0] is task_run + # mcpServers payload is serialized McpServerConfig shape + mcp_servers = mock_send_refresh.call_args.args[1] + assert mcp_servers == [_make_mcp_config(token="fresh-token").to_dict()] + + @patch("products.tasks.backend.temporal.process_task.sandbox_identity.time.sleep") + @patch("products.tasks.backend.temporal.process_task.sandbox_identity.send_refresh_session") + @patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_user_mcp_server_configs") + @patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_sandbox_ph_mcp_configs") + @patch("products.tasks.backend.temporal.process_task.sandbox_identity.create_oauth_access_token_for_run") + def test_retries_once_on_first_failure( + self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh, mock_sleep + ): + mock_oauth.return_value = "fresh-token" + mock_ph_configs.return_value = [_make_mcp_config()] + mock_user_configs.return_value = [] + mock_send_refresh.side_effect = [ + CommandResult(success=False, status_code=502, error="transient", retryable=True), + CommandResult(success=True, status_code=200), + ] + + _ensure(_make_task_run_mock()) + + assert mock_send_refresh.call_count == 2 + mock_sleep.assert_called_once_with(REFRESH_RETRY_DELAY_SECONDS) + + @patch("products.tasks.backend.temporal.process_task.sandbox_identity.time.sleep") + @patch("products.tasks.backend.temporal.process_task.sandbox_identity.send_refresh_session") + @patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_user_mcp_server_configs") + @patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_sandbox_ph_mcp_configs") + @patch("products.tasks.backend.temporal.process_task.sandbox_identity.create_oauth_access_token_for_run") + def test_two_failures_are_non_fatal( + self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh, _mock_sleep + ): + mock_oauth.return_value = "fresh-token" + mock_ph_configs.return_value = [_make_mcp_config()] + mock_user_configs.return_value = [] + mock_send_refresh.return_value = CommandResult(success=False, status_code=502, error="down") + + # Must not raise. + _ensure(_make_task_run_mock()) + + assert mock_send_refresh.call_count == 2 + # Cache stays empty so the next follow-up retries the dispatch. + assert cache.get(_mcp_token_issued_cache_key("run-1", 42)) is None + assert cache.get(_sandbox_identity_cache_key("run-1", "mcp")) is None + + @patch("products.tasks.backend.temporal.process_task.sandbox_identity.send_refresh_session") + @patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_user_mcp_server_configs") + @patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_sandbox_ph_mcp_configs") + @patch("products.tasks.backend.temporal.process_task.sandbox_identity.create_oauth_access_token_for_run") + def test_token_mint_failure_is_non_fatal_and_skips_send( + self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh + ): + mock_oauth.side_effect = RuntimeError("oauth service down") + + _ensure(_make_task_run_mock()) + + mock_ph_configs.assert_not_called() + mock_user_configs.assert_not_called() + mock_send_refresh.assert_not_called() + + @patch("products.tasks.backend.temporal.process_task.sandbox_identity.send_refresh_session") + @patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_user_mcp_server_configs") + @patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_sandbox_ph_mcp_configs") + @patch("products.tasks.backend.temporal.process_task.sandbox_identity.create_oauth_access_token_for_run") + def test_skips_send_when_no_mcp_configs_resolved( + self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh + ): + mock_oauth.return_value = "fresh-token" + mock_ph_configs.return_value = [] + mock_user_configs.return_value = [] + + _ensure(_make_task_run_mock()) + + mock_send_refresh.assert_not_called() + + @patch("products.tasks.backend.temporal.process_task.sandbox_identity.send_refresh_session") + @patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_user_mcp_server_configs") + @patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_sandbox_ph_mcp_configs") + @patch("products.tasks.backend.temporal.process_task.sandbox_identity.create_oauth_access_token_for_run") + def test_user_mcp_configs_skipped_when_no_creator( + self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh + ): + mock_oauth.return_value = "fresh-token" + mock_ph_configs.return_value = [_make_mcp_config()] + mock_send_refresh.return_value = CommandResult(success=True, status_code=200) + + _ensure(_make_task_run_mock(created_by_id=None)) + + mock_user_configs.assert_not_called() + mock_send_refresh.assert_called_once() + + @patch("products.tasks.backend.temporal.process_task.sandbox_identity.send_refresh_session") + @patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_user_mcp_server_configs") + @patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_sandbox_ph_mcp_configs") + @patch("products.tasks.backend.temporal.process_task.sandbox_identity.create_oauth_access_token_for_run") + def test_scopes_propagate_to_oauth_and_configs( + self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh + ): + mock_oauth.return_value = "fresh-token" + mock_ph_configs.return_value = [_make_mcp_config()] + mock_user_configs.return_value = [] + mock_send_refresh.return_value = CommandResult(success=True, status_code=200) + + _ensure(_make_task_run_mock(), scopes="full") + + mock_oauth.assert_called_once_with(mock_oauth.call_args.args[0], None, scopes="full") + mock_ph_configs.assert_called_once_with( + token="fresh-token", project_id=7, scopes="full", interaction_origin=None, task_id="task-1" + ) + + +class TestMcpRefreshIntervalGate: + """Refreshes within MCP_TOKEN_REFRESH_INTERVAL_SECONDS of a previous + successful issuance for the same actor must be skipped without minting a + new token or contacting the sandbox.""" + + @patch("products.tasks.backend.temporal.process_task.sandbox_identity.send_refresh_session") + @patch("products.tasks.backend.temporal.process_task.sandbox_identity.create_oauth_access_token_for_run") + def test_skipped_when_token_recently_issued(self, mock_oauth, mock_send_refresh): + mark_mcp_token_issued("run-1", 42) + + _ensure(_make_task_run_mock()) + + mock_oauth.assert_not_called() + mock_send_refresh.assert_not_called() + + @patch("products.tasks.backend.temporal.process_task.sandbox_identity.send_refresh_session") + @patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_user_mcp_server_configs") + @patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_sandbox_ph_mcp_configs") + @patch("products.tasks.backend.temporal.process_task.sandbox_identity.create_oauth_access_token_for_run") + def test_marks_after_successful_refresh(self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh): + mock_oauth.return_value = "fresh-token" + mock_ph_configs.return_value = [_make_mcp_config()] + mock_user_configs.return_value = [] + mock_send_refresh.return_value = CommandResult(success=True, status_code=200) + + _ensure(_make_task_run_mock()) + + # Cache entries now exist → next refresh for this actor within the + # interval is gated, and the session identity is recorded. + assert cache.get(_mcp_token_issued_cache_key("run-1", 42)) is True + assert cache.get(_sandbox_identity_cache_key("run-1", "mcp")) == 42 + + @patch("products.tasks.backend.temporal.process_task.sandbox_identity.time.sleep") + @patch("products.tasks.backend.temporal.process_task.sandbox_identity.send_refresh_session") + @patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_user_mcp_server_configs") + @patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_sandbox_ph_mcp_configs") + @patch("products.tasks.backend.temporal.process_task.sandbox_identity.create_oauth_access_token_for_run") + def test_marks_after_successful_retry( + self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh, _mock_sleep + ): + mock_oauth.return_value = "fresh-token" + mock_ph_configs.return_value = [_make_mcp_config()] + mock_user_configs.return_value = [] + mock_send_refresh.side_effect = [ + CommandResult(success=False, status_code=502, error="transient"), + CommandResult(success=True, status_code=200), + ] + + _ensure(_make_task_run_mock()) + + assert cache.get(_mcp_token_issued_cache_key("run-1", 42)) is True + + +@patch("products.tasks.backend.temporal.process_task.sandbox_identity.send_refresh_session") +@patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_user_mcp_server_configs") +@patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_sandbox_ph_mcp_configs") +@patch("products.tasks.backend.temporal.process_task.sandbox_identity.create_oauth_access_token_for_run") +class TestMcpIdentityTransitionGate: + """An actor transition (run-state actor differs from the identity last + pushed to the sandbox) must bypass the freshness window and rebind the + session; the marks are keyed per sandbox so a replacement sandbox starts + unmarked.""" + + def _arm_success(self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh): + mock_oauth.return_value = "fresh-token" + mock_ph_configs.return_value = [_make_mcp_config(token="fresh-token")] + mock_user_configs.return_value = [] + mock_send_refresh.return_value = CommandResult(success=True, status_code=200) + + def test_actor_change_bypasses_freshness_window( + self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh + ): + self._arm_success(mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh) + # The creator's token is fresh and the session is bound to them… + mark_mcp_token_issued("run-1", 42) + mark_sandbox_identity("run-1", "mcp", 42) + + # …but the next message comes from a different actor. + with _patch_actor(99): + _ensure(_make_task_run_mock()) + + mock_send_refresh.assert_called_once() + assert cache.get(_sandbox_identity_cache_key("run-1", "mcp")) == 99 + assert cache.get(_mcp_token_issued_cache_key("run-1", 99)) is True + + def test_switch_back_to_creator_refreshes_despite_fresh_window( + self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh + ): + self._arm_success(mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh) + # The creator's window is still warm, but the session was last bound + # to another user — the creator speaking again is a transition. + mark_mcp_token_issued("run-1", 42) + mark_sandbox_identity("run-1", "mcp", 99) + + with _patch_actor(42): + _ensure(_make_task_run_mock()) + + mock_send_refresh.assert_called_once() + assert cache.get(_sandbox_identity_cache_key("run-1", "mcp")) == 42 + + def test_same_actor_within_window_is_skipped( + self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh + ): + self._arm_success(mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh) + mark_mcp_token_issued("run-1", 99) + mark_sandbox_identity("run-1", "mcp", 99) + + with _patch_actor(99): + _ensure(_make_task_run_mock()) + + mock_oauth.assert_not_called() + mock_send_refresh.assert_not_called() + + def test_replacement_sandbox_starts_unmarked( + self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh + ): + self._arm_success(mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh) + # Marks recorded against the run id (legacy scope) must not gate a run + # whose state now points at a concrete sandbox. + mark_mcp_token_issued("run-1", 42) + mark_sandbox_identity("run-1", "mcp", 42) + + with _patch_actor(42): + _ensure(_make_task_run_mock(state={"sandbox_id": "sb-2"})) + + mock_send_refresh.assert_called_once() + assert cache.get(_mcp_token_issued_cache_key("sb-2", 42)) is True + assert cache.get(_sandbox_identity_cache_key("sb-2", "mcp")) == 42 + + +@patch("products.tasks.backend.temporal.process_task.sandbox_identity.send_refresh_session") +@patch("products.tasks.backend.temporal.process_task.sandbox_identity.update_sandbox_env_file") +@patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_git_identity_env_vars") +@patch("products.tasks.backend.temporal.process_task.sandbox_identity.GitHubSandboxCredential") +@patch("products.tasks.backend.temporal.process_task.sandbox_identity.Sandbox") +class TestGithubIdentityTransitionGate: + """The GitHub surface rebinds on actor transitions only — token TTL + between messages is owned by the credential-refresh loop.""" + + def _quiet_mcp(self, user_id: int, scope: str = "run-1") -> None: + """Pre-mark the MCP surface so `ensure_sandbox_identity` exercises only + the GitHub gate (no oauth/config mocks needed).""" + mark_mcp_token_issued(scope, user_id) + mark_sandbox_identity(scope, "mcp", user_id) + + def _arm_refresh(self, mock_credential_cls, mock_git_env, refreshed: bool = True): + mock_credential_cls.return_value.refresh.return_value = CredentialRefreshOutcome( + "github", refreshed=refreshed, next_refresh_seconds=60 + ) + mock_git_env.return_value = {"GIT_AUTHOR_NAME": "New Actor", "GIT_AUTHOR_EMAIL": "actor@example.com"} + + def test_actor_transition_rebinds_credentials_and_author( + self, mock_sandbox_cls, mock_credential_cls, mock_git_env, mock_env_file, mock_send_refresh + ): + self._arm_refresh(mock_credential_cls, mock_git_env) + mock_send_refresh.return_value = CommandResult(success=True, status_code=200) + self._quiet_mcp(99, scope="sb-2") + mark_sandbox_identity("sb-2", "github", 42) + state = {"sandbox_id": "sb-2", "pr_authorship_mode": "user"} + task_run = _make_task_run_mock(state=state) + + with _patch_actor(99): + _ensure(task_run, auth_token="jwt", processing_context=_make_processing_context()) + + mock_sandbox_cls.get_by_id.assert_called_once_with("sb-2") + refresh_call = mock_credential_cls.return_value.refresh.call_args + live_ctx = refresh_call.args[1] + assert live_ctx.state == state # boot snapshot replaced with live run state + mock_env_file.assert_called_once() + _, notify_kwargs = mock_send_refresh.call_args + assert notify_kwargs["refreshed_credentials"] == ["github"] + assert notify_kwargs["authorship"] == "user" + assert mock_send_refresh.call_args.args[1] == [] + assert cache.get(_sandbox_identity_cache_key("sb-2", "github")) == 99 + + def test_same_actor_does_not_touch_sandbox( + self, mock_sandbox_cls, mock_credential_cls, mock_git_env, mock_env_file, mock_send_refresh + ): + self._quiet_mcp(99) + mark_sandbox_identity("run-1", "github", 99) + + with _patch_actor(99): + _ensure(_make_task_run_mock(), processing_context=_make_processing_context()) + + mock_sandbox_cls.get_by_id.assert_not_called() + mock_send_refresh.assert_not_called() + + def test_skipped_without_processing_context( + self, mock_sandbox_cls, mock_credential_cls, mock_git_env, mock_env_file, mock_send_refresh + ): + self._quiet_mcp(99) + mark_sandbox_identity("run-1", "github", 42) + + with _patch_actor(99): + _ensure(_make_task_run_mock(), processing_context=None) + + mock_sandbox_cls.get_by_id.assert_not_called() + + def test_skipped_without_github_credentials( + self, mock_sandbox_cls, mock_credential_cls, mock_git_env, mock_env_file, mock_send_refresh + ): + self._quiet_mcp(99) + mark_sandbox_identity("run-1", "github", 42) + context = _make_processing_context(github_integration_id=None) + + with _patch_actor(99): + _ensure(_make_task_run_mock(), processing_context=context) + + mock_sandbox_cls.get_by_id.assert_not_called() + + def test_rebind_failure_leaves_identity_unmarked_for_retry( + self, mock_sandbox_cls, mock_credential_cls, mock_git_env, mock_env_file, mock_send_refresh + ): + mock_credential_cls.return_value.refresh.side_effect = RuntimeError("sandbox unreachable") + self._quiet_mcp(99) + mark_sandbox_identity("run-1", "github", 42) + + with _patch_actor(99): + _ensure(_make_task_run_mock(), processing_context=_make_processing_context()) + + assert cache.get(_sandbox_identity_cache_key("run-1", "github")) == 42 + mock_send_refresh.assert_not_called() + + def test_missing_sandbox_id_leaves_identity_unmarked( + self, mock_sandbox_cls, mock_credential_cls, mock_git_env, mock_env_file, mock_send_refresh + ): + self._quiet_mcp(99) + mark_sandbox_identity("run-1", "github", 42) + + with _patch_actor(99): + _ensure(_make_task_run_mock(state=None), processing_context=_make_processing_context()) + + assert cache.get(_sandbox_identity_cache_key("run-1", "github")) == 42 + mock_sandbox_cls.get_by_id.assert_not_called() + + def test_unrefreshable_credential_marks_without_notify( + self, mock_sandbox_cls, mock_credential_cls, mock_git_env, mock_env_file, mock_send_refresh + ): + # e.g. a caller-token run: nothing we manage, so nothing can diverge — + # mark to avoid re-probing the sandbox on every message. + self._arm_refresh(mock_credential_cls, mock_git_env, refreshed=False) + self._quiet_mcp(99, scope="sb-2") + mark_sandbox_identity("sb-2", "github", 42) + + with _patch_actor(99): + _ensure(_make_task_run_mock(state={"sandbox_id": "sb-2"}), processing_context=_make_processing_context()) + + assert cache.get(_sandbox_identity_cache_key("sb-2", "github")) == 99 + mock_env_file.assert_not_called() + mock_send_refresh.assert_not_called() 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 a0f7a1b4a656..31c100cb0290 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 @@ -1,54 +1,17 @@ import pytest from unittest.mock import MagicMock, patch -from django.core.cache import cache - from temporalio.exceptions import ApplicationError from products.tasks.backend.logic.services.agent_command import CommandResult from products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox import ( - REFRESH_RETRY_DELAY_SECONDS, SEND_FOLLOWUP_MAX_ATTEMPTS, SendFollowupToSandboxInput, - _refresh_sandbox_mcp, send_followup_to_sandbox, ) -from products.tasks.backend.temporal.process_task.utils import ( - McpServerConfig, - _mcp_token_issued_cache_key, - _sandbox_identity_cache_key, - mark_mcp_token_issued, - mark_sandbox_identity, -) pytestmark = pytest.mark.django_db -_GATE_CACHE_KEYS = [ - _mcp_token_issued_cache_key("run-1", 42), - _mcp_token_issued_cache_key("run-1", 99), - _mcp_token_issued_cache_key("sb-2", 42), - _sandbox_identity_cache_key("run-1", "mcp"), - _sandbox_identity_cache_key("sb-2", "mcp"), -] - - -@pytest.fixture(autouse=True) -def _clear_mcp_token_cache(): - """Ensure each test starts with no recorded token issuances or session - identities so the refresh gate doesn't carry state between tests.""" - cache.delete_many(_GATE_CACHE_KEYS) - yield - cache.delete_many(_GATE_CACHE_KEYS) - - -def _make_mcp_config(name: str = "posthog", token: str = "tok") -> McpServerConfig: - return McpServerConfig( - type="http", - name=name, - url="https://mcp.posthog.com/mcp", - headers=[{"name": "Authorization", "value": f"Bearer {token}"}], - ) - def _make_task_run_mock(team_id: int = 7, created_by_id: int | None = 42, state: dict | None = None) -> MagicMock: task = MagicMock() @@ -69,360 +32,6 @@ def _make_task_run_mock(team_id: int = 7, created_by_id: int | None = 42, state: return task_run -class TestRefreshSandboxMcp: - @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.send_refresh_session") - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_user_mcp_server_configs" - ) - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_sandbox_ph_mcp_configs" - ) - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.create_oauth_access_token_for_run" - ) - def test_success_path_single_call(self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh): - mock_oauth.return_value = "fresh-token" - mock_ph_configs.return_value = [_make_mcp_config(token="fresh-token")] - mock_user_configs.return_value = [] - mock_send_refresh.return_value = CommandResult(success=True, status_code=200) - - task_run = _make_task_run_mock() - _refresh_sandbox_mcp(task_run, "read_only", auth_token="jwt") - - mock_oauth.assert_called_once_with(task_run.task, task_run.state, scopes="read_only") - mock_ph_configs.assert_called_once_with( - token="fresh-token", project_id=7, scopes="read_only", interaction_origin=None, task_id="task-1" - ) - mock_user_configs.assert_called_once_with(token="fresh-token", team_id=7, user_id=42, interaction_origin=None) - mock_send_refresh.assert_called_once() - _, kwargs = mock_send_refresh.call_args - assert kwargs["auth_token"] == "jwt" - assert mock_send_refresh.call_args.args[0] is task_run - # mcpServers payload is serialized McpServerConfig shape - mcp_servers = mock_send_refresh.call_args.args[1] - assert mcp_servers == [_make_mcp_config(token="fresh-token").to_dict()] - - @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.time.sleep") - @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.send_refresh_session") - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_user_mcp_server_configs" - ) - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_sandbox_ph_mcp_configs" - ) - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.create_oauth_access_token_for_run" - ) - def test_retries_once_on_first_failure( - self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh, mock_sleep - ): - mock_oauth.return_value = "fresh-token" - mock_ph_configs.return_value = [_make_mcp_config()] - mock_user_configs.return_value = [] - mock_send_refresh.side_effect = [ - CommandResult(success=False, status_code=502, error="transient", retryable=True), - CommandResult(success=True, status_code=200), - ] - - _refresh_sandbox_mcp(_make_task_run_mock(), "read_only", auth_token=None) - - assert mock_send_refresh.call_count == 2 - mock_sleep.assert_called_once_with(REFRESH_RETRY_DELAY_SECONDS) - - @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.time.sleep") - @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.send_refresh_session") - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_user_mcp_server_configs" - ) - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_sandbox_ph_mcp_configs" - ) - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.create_oauth_access_token_for_run" - ) - def test_two_failures_are_non_fatal( - self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh, _mock_sleep - ): - mock_oauth.return_value = "fresh-token" - mock_ph_configs.return_value = [_make_mcp_config()] - mock_user_configs.return_value = [] - mock_send_refresh.return_value = CommandResult(success=False, status_code=502, error="down") - - # Must not raise. - _refresh_sandbox_mcp(_make_task_run_mock(), "read_only", auth_token=None) - - assert mock_send_refresh.call_count == 2 - - @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.send_refresh_session") - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_user_mcp_server_configs" - ) - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_sandbox_ph_mcp_configs" - ) - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.create_oauth_access_token_for_run" - ) - def test_token_mint_failure_is_non_fatal_and_skips_send( - self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh - ): - mock_oauth.side_effect = RuntimeError("oauth service down") - - _refresh_sandbox_mcp(_make_task_run_mock(), "read_only", auth_token=None) - - mock_ph_configs.assert_not_called() - mock_user_configs.assert_not_called() - mock_send_refresh.assert_not_called() - - @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.send_refresh_session") - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_user_mcp_server_configs" - ) - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_sandbox_ph_mcp_configs" - ) - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.create_oauth_access_token_for_run" - ) - def test_skips_send_when_no_mcp_configs_resolved( - self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh - ): - mock_oauth.return_value = "fresh-token" - mock_ph_configs.return_value = [] - mock_user_configs.return_value = [] - - _refresh_sandbox_mcp(_make_task_run_mock(), "read_only", auth_token=None) - - mock_send_refresh.assert_not_called() - - @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.send_refresh_session") - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_user_mcp_server_configs" - ) - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_sandbox_ph_mcp_configs" - ) - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.create_oauth_access_token_for_run" - ) - def test_user_mcp_configs_skipped_when_no_creator( - self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh - ): - mock_oauth.return_value = "fresh-token" - mock_ph_configs.return_value = [_make_mcp_config()] - mock_send_refresh.return_value = CommandResult(success=True, status_code=200) - - _refresh_sandbox_mcp(_make_task_run_mock(created_by_id=None), "read_only", auth_token=None) - - mock_user_configs.assert_not_called() - mock_send_refresh.assert_called_once() - - @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.send_refresh_session") - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_user_mcp_server_configs" - ) - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_sandbox_ph_mcp_configs" - ) - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.create_oauth_access_token_for_run" - ) - def test_scopes_propagate_to_oauth_and_configs( - self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh - ): - mock_oauth.return_value = "fresh-token" - mock_ph_configs.return_value = [_make_mcp_config()] - mock_user_configs.return_value = [] - mock_send_refresh.return_value = CommandResult(success=True, status_code=200) - - _refresh_sandbox_mcp(_make_task_run_mock(), "full", auth_token=None) - - mock_oauth.assert_called_once_with(mock_oauth.call_args.args[0], None, scopes="full") - mock_ph_configs.assert_called_once_with( - token="fresh-token", project_id=7, scopes="full", interaction_origin=None, task_id="task-1" - ) - - -class TestRefreshIntervalGate: - """Refreshes within MCP_TOKEN_REFRESH_INTERVAL_SECONDS of a previous - successful issuance for the same actor must be skipped without minting a - new token or contacting the sandbox.""" - - @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.send_refresh_session") - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.create_oauth_access_token_for_run" - ) - def test_skipped_when_token_recently_issued(self, mock_oauth, mock_send_refresh): - mark_mcp_token_issued("run-1", 42) - - _refresh_sandbox_mcp(_make_task_run_mock(), "read_only", auth_token=None) - - mock_oauth.assert_not_called() - mock_send_refresh.assert_not_called() - - @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.send_refresh_session") - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_user_mcp_server_configs" - ) - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_sandbox_ph_mcp_configs" - ) - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.create_oauth_access_token_for_run" - ) - def test_marks_after_successful_refresh(self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh): - mock_oauth.return_value = "fresh-token" - mock_ph_configs.return_value = [_make_mcp_config()] - mock_user_configs.return_value = [] - mock_send_refresh.return_value = CommandResult(success=True, status_code=200) - - _refresh_sandbox_mcp(_make_task_run_mock(), "read_only", auth_token=None) - - # Cache entries now exist → next refresh for this actor within the - # interval is gated, and the session identity is recorded. - assert cache.get(_mcp_token_issued_cache_key("run-1", 42)) is True - assert cache.get(_sandbox_identity_cache_key("run-1", "mcp")) == 42 - - @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.time.sleep") - @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.send_refresh_session") - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_user_mcp_server_configs" - ) - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_sandbox_ph_mcp_configs" - ) - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.create_oauth_access_token_for_run" - ) - def test_marks_after_successful_retry( - self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh, _mock_sleep - ): - mock_oauth.return_value = "fresh-token" - mock_ph_configs.return_value = [_make_mcp_config()] - mock_user_configs.return_value = [] - mock_send_refresh.side_effect = [ - CommandResult(success=False, status_code=502, error="transient"), - CommandResult(success=True, status_code=200), - ] - - _refresh_sandbox_mcp(_make_task_run_mock(), "read_only", auth_token=None) - - assert cache.get(_mcp_token_issued_cache_key("run-1", 42)) is True - - @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.time.sleep") - @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.send_refresh_session") - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_user_mcp_server_configs" - ) - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_sandbox_ph_mcp_configs" - ) - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.create_oauth_access_token_for_run" - ) - def test_does_not_mark_after_two_failures( - self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh, _mock_sleep - ): - mock_oauth.return_value = "fresh-token" - mock_ph_configs.return_value = [_make_mcp_config()] - mock_user_configs.return_value = [] - mock_send_refresh.return_value = CommandResult(success=False, status_code=502, error="down") - - _refresh_sandbox_mcp(_make_task_run_mock(), "read_only", auth_token=None) - - # Cache stays empty so the next follow-up retries the dispatch. - assert cache.get(_mcp_token_issued_cache_key("run-1", 42)) is None - assert cache.get(_sandbox_identity_cache_key("run-1", "mcp")) is None - - -def _patch_actor(user_id: int): - """Pin the resolved credential user so tests can drive actor transitions - without building real run state.""" - return patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_task_run_credential_user", - return_value=MagicMock(id=user_id), - ) - - -@patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.send_refresh_session") -@patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_user_mcp_server_configs") -@patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_sandbox_ph_mcp_configs") -@patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.create_oauth_access_token_for_run" -) -class TestIdentityTransitionGate: - """An actor transition (run-state actor differs from the identity last - pushed to the sandbox) must bypass the freshness window and rebind the - session; the marks are keyed per sandbox so a replacement sandbox starts - unmarked.""" - - def _arm_success(self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh): - mock_oauth.return_value = "fresh-token" - mock_ph_configs.return_value = [_make_mcp_config(token="fresh-token")] - mock_user_configs.return_value = [] - mock_send_refresh.return_value = CommandResult(success=True, status_code=200) - - def test_actor_change_bypasses_freshness_window( - self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh - ): - self._arm_success(mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh) - # The creator's token is fresh and the session is bound to them… - mark_mcp_token_issued("run-1", 42) - mark_sandbox_identity("run-1", "mcp", 42) - - # …but the next message comes from a different actor. - with _patch_actor(99): - _refresh_sandbox_mcp(_make_task_run_mock(), "read_only", auth_token=None) - - mock_send_refresh.assert_called_once() - assert cache.get(_sandbox_identity_cache_key("run-1", "mcp")) == 99 - assert cache.get(_mcp_token_issued_cache_key("run-1", 99)) is True - - def test_switch_back_to_creator_refreshes_despite_fresh_window( - self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh - ): - self._arm_success(mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh) - # The creator's window is still warm, but the session was last bound - # to another user — the creator speaking again is a transition. - mark_mcp_token_issued("run-1", 42) - mark_sandbox_identity("run-1", "mcp", 99) - - with _patch_actor(42): - _refresh_sandbox_mcp(_make_task_run_mock(), "read_only", auth_token=None) - - mock_send_refresh.assert_called_once() - assert cache.get(_sandbox_identity_cache_key("run-1", "mcp")) == 42 - - def test_same_actor_within_window_is_skipped( - self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh - ): - self._arm_success(mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh) - mark_mcp_token_issued("run-1", 99) - mark_sandbox_identity("run-1", "mcp", 99) - - with _patch_actor(99): - _refresh_sandbox_mcp(_make_task_run_mock(), "read_only", auth_token=None) - - mock_oauth.assert_not_called() - mock_send_refresh.assert_not_called() - - def test_replacement_sandbox_starts_unmarked( - self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh - ): - self._arm_success(mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh) - # Marks recorded against the run id (legacy scope) must not gate a run - # whose state now points at a concrete sandbox. - mark_mcp_token_issued("run-1", 42) - mark_sandbox_identity("run-1", "mcp", 42) - - with _patch_actor(42): - _refresh_sandbox_mcp(_make_task_run_mock(state={"sandbox_id": "sb-2"}), "read_only", auth_token=None) - - mock_send_refresh.assert_called_once() - assert cache.get(_mcp_token_issued_cache_key("sb-2", 42)) is True - assert cache.get(_sandbox_identity_cache_key("sb-2", "mcp")) == 42 - - class TestSendFollowupActivityRefreshOrdering: """Refresh call must precede user_message, and the activity must succeed when refresh fails (non-fatal) as long as user_message succeeds.""" @@ -438,7 +47,7 @@ def _patches(self): "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.create_sandbox_connection_token" ) as mock_conn_token, patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox._refresh_sandbox_mcp" + "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.ensure_sandbox_identity" ) as mock_refresh, patch( "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.send_user_message" @@ -485,18 +94,18 @@ def test_scopes_flow_from_input_to_refresh(self, _patches): send_followup_to_sandbox(SendFollowupToSandboxInput(run_id="run-1", message="hi", posthog_mcp_scopes="full")) _patches["refresh"].assert_called_once() - args, _kwargs = _patches["refresh"].call_args + args, kwargs = _patches["refresh"].call_args assert args[0] is _patches["task_run"] - assert args[1] == "full" - assert args[2] == "jwt" + assert kwargs["posthog_mcp_scopes"] == "full" + assert kwargs["auth_token"] == "jwt" def test_default_scope_is_read_only(self, _patches): _patches["user_msg"].return_value = CommandResult(success=True, status_code=200) send_followup_to_sandbox(SendFollowupToSandboxInput(run_id="run-1", message="hi")) - args, _kwargs = _patches["refresh"].call_args - assert args[1] == "read_only" + _args, kwargs = _patches["refresh"].call_args + assert kwargs["posthog_mcp_scopes"] == "read_only" class TestSendFollowupTurnTimeout: @@ -515,7 +124,7 @@ def _patches(self): "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.create_sandbox_connection_token" ) as mock_conn_token, patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox._refresh_sandbox_mcp" + "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.ensure_sandbox_identity" ), patch( "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.send_user_message" diff --git a/products/tasks/backend/temporal/process_task/utils.py b/products/tasks/backend/temporal/process_task/utils.py index 0980d32aac41..ca8d739a3e6a 100644 --- a/products/tasks/backend/temporal/process_task/utils.py +++ b/products/tasks/backend/temporal/process_task/utils.py @@ -394,7 +394,7 @@ def should_refresh_mcp_token(scope: str, user_id: int) -> bool: # boot-time one (the task creator). SANDBOX_IDENTITY_TTL_SECONDS = 7 * 24 * 60 * 60 -SandboxIdentityKind = Literal["mcp"] +SandboxIdentityKind = Literal["mcp", "github"] def _sandbox_identity_cache_key(scope: str, kind: SandboxIdentityKind) -> str: @@ -404,7 +404,8 @@ def _sandbox_identity_cache_key(scope: str, kind: SandboxIdentityKind) -> str: def mark_sandbox_identity(scope: str, kind: SandboxIdentityKind, value: int | str) -> None: """Record which identity the sandbox's live session currently holds. - ``mcp`` stores the user id the last-pushed OAuth token was minted for. + ``mcp`` stores the user id the last-pushed OAuth token was minted for; + ``github`` the user id the git/gh credentials were last rebound to. The run-state actor (``get_task_run_credential_user``) says who *should* be acting; this mark says whose token the session actually has — the gap between the two is an identity transition that must bypass the freshness diff --git a/products/tasks/backend/temporal/process_task/workflow.py b/products/tasks/backend/temporal/process_task/workflow.py index 5c19826e7a3f..f84f070fe4c4 100644 --- a/products/tasks/backend/temporal/process_task/workflow.py +++ b/products/tasks/backend/temporal/process_task/workflow.py @@ -1687,6 +1687,7 @@ async def _send_followup_to_sandbox(self, message: str | None, artifact_ids: lis posthog_mcp_scopes=self._posthog_mcp_scopes, artifact_ids=artifact_ids, message_id=str(workflow.uuid4()), + context=self.context, ), start_to_close_timeout=timedelta(minutes=35), # The activity heartbeats while blocked on the sync delivery From bb2c977fbb3922a877d06f789588463e07742704 Mon Sep 17 00:00:00 2001 From: Vojta Bartos Date: Tue, 14 Jul 2026 12:50:20 +0200 Subject: [PATCH 5/5] refactor(tasks): dedupe identity gate plumbing Applies review findings across the identity-gate stack: - Pass the caller's already-resolved actor into ensure_sandbox_identity instead of re-resolving it (2-3 redundant DB queries per follow-up). - Fold the git-author write into GitHubSandboxCredential so every credential write (identity gate and refresh loop alike) applies token and author together, and share one agent-server notify helper between the loop and the gate. - Add TaskProcessingContext.with_state as the single owner of the boot-snapshot-to-live-state rebase, and fetch the TaskRun once per refresh tick instead of twice. - Collapse the duplicated send/retry success blocks into a loop, share the mark-pair helper (utils.mark_mcp_session) with the boot path, and name the creator-default assumption once (_last_bound_identity). - Guard the GitHub gate on sandbox_id before logging a transition. - Tests: shared mock helpers module, class-level patch stacks, cache.clear() fixture instead of a hand-enumerated key list. --- .../activities/get_task_processing_context.py | 15 +- .../activities/refresh_sandbox_credentials.py | 29 +- .../activities/send_followup_to_sandbox.py | 1 + .../activities/start_agent_server.py | 5 +- .../process_task/sandbox_credentials.py | 37 ++- .../temporal/process_task/sandbox_identity.py | 152 ++++----- .../temporal/process_task/tests/helpers.py | 31 ++ .../tests/test_sandbox_credentials.py | 20 ++ .../tests/test_sandbox_identity.py | 312 ++++++------------ .../tests/test_send_followup_to_sandbox.py | 24 +- .../backend/temporal/process_task/utils.py | 8 + 11 files changed, 292 insertions(+), 342 deletions(-) create mode 100644 products/tasks/backend/temporal/process_task/tests/helpers.py 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 acc6bf603274..94152fc8d2a7 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 @@ -1,4 +1,4 @@ -from dataclasses import dataclass +from dataclasses import dataclass, replace from datetime import timedelta from django.conf import settings @@ -94,6 +94,19 @@ class TaskProcessingContext: # (the user's toggle) applies. Captured at workflow start so it's stable across retries. rtk_enabled: bool = True + def with_state(self, state: dict | None) -> "TaskProcessingContext": + """Copy of this context rebased onto the given (typically live) run state. + + The context is captured at workflow start, so its ``state`` is a + boot-time snapshot while the run's actual state moves on (multiplayer + follow-ups update the actor mid-run). Anything resolving the *current* + actor must go through a rebased copy. The other fields intentionally + stay boot-time: they are either immutable for the run (ids, repository) + or last-resort fallbacks (``github_user_integration_id``, + ``distinct_id``) that actor-aware resolution overrides via state. + """ + return replace(self, state=state) + @property def mode(self) -> str: """Get the execution mode from state. Defaults to 'background'.""" diff --git a/products/tasks/backend/temporal/process_task/activities/refresh_sandbox_credentials.py b/products/tasks/backend/temporal/process_task/activities/refresh_sandbox_credentials.py index 4279d6938923..e592cb0007fa 100644 --- a/products/tasks/backend/temporal/process_task/activities/refresh_sandbox_credentials.py +++ b/products/tasks/backend/temporal/process_task/activities/refresh_sandbox_credentials.py @@ -1,4 +1,4 @@ -from dataclasses import dataclass, field, replace +from dataclasses import dataclass, field from temporalio import activity @@ -11,7 +11,6 @@ SandboxNotRunningError, TaskNotFoundError, ) -from products.tasks.backend.logic.services.agent_command import send_refresh_session from products.tasks.backend.logic.services.connection_token import create_sandbox_connection_token from products.tasks.backend.logic.services.sandbox import Sandbox from products.tasks.backend.models import Task, TaskRun @@ -20,6 +19,7 @@ from products.tasks.backend.temporal.process_task.sandbox_credentials import ( DEFAULT_REFRESH_INTERVAL_SECONDS, build_sandbox_credentials, + notify_sandbox_credentials_refreshed, ) from products.tasks.backend.temporal.process_task.utils import ( get_actor_distinct_id, @@ -32,12 +32,13 @@ logger = get_logger(__name__) -def _notify_agent_server_of_refresh(ctx: TaskProcessingContext, task: Task, refreshed_kinds: list[str]) -> None: +def _notify_agent_server_of_refresh( + ctx: TaskProcessingContext, task: Task, task_run: TaskRun, refreshed_kinds: list[str] +) -> None: """Tell the running agent-server which credentials were re-injected so it logs them. This is best-effort since the sandbox may be unreachable, so a failure here never fails the refresh itself. """ try: - task_run = TaskRun.objects.get(id=ctx.run_id) auth_token = None actor_user = get_task_run_credential_user(task, ctx.state) if is_slack_interaction_state(ctx.state) and actor_user is None: @@ -47,10 +48,7 @@ def _notify_agent_server_of_refresh(ctx: TaskProcessingContext, task: Task, refr auth_token = create_sandbox_connection_token( task_run, user_id=actor_user.id, distinct_id=get_actor_distinct_id(actor_user) ) - authorship = (ctx.state or {}).get("pr_authorship_mode") - send_refresh_session( - task_run, [], auth_token=auth_token, refreshed_credentials=refreshed_kinds, authorship=authorship - ) + notify_sandbox_credentials_refreshed(task_run, refreshed_kinds, auth_token=auth_token) except Exception: logger.warning("sandbox_credentials_refresh_notify_failed", run_id=ctx.run_id, exc_info=True) @@ -98,12 +96,13 @@ def refresh_sandbox_credentials(input: RefreshSandboxCredentialsInput) -> Refres except Task.DoesNotExist as e: raise TaskNotFoundError(f"Task {ctx.task_id} not found", {"task_id": ctx.task_id}, cause=e) - # The workflow-start context carries a boot-time snapshot of the run - # state, but credential resolution must follow the *current* actor - # (multiplayer follow-ups update it mid-run) — re-read it live. + # Credential resolution must follow the *current* actor (multiplayer + # follow-ups update it mid-run), so rebase the boot-time context onto + # the live run state. + task_run = None try: - live_state = TaskRun.objects.values_list("state", flat=True).get(id=ctx.run_id) - ctx = replace(ctx, state=live_state) + task_run = TaskRun.objects.get(id=ctx.run_id) + ctx = ctx.with_state(task_run.state) except TaskRun.DoesNotExist: logger.warning("sandbox_credentials_refresh_run_missing", run_id=ctx.run_id) @@ -194,8 +193,8 @@ def refresh_sandbox_credentials(input: RefreshSandboxCredentialsInput) -> Refres if intervals: next_refresh = min(intervals) - if refreshed_kinds: - _notify_agent_server_of_refresh(ctx, task, refreshed_kinds) + if refreshed_kinds and task_run is not None: + _notify_agent_server_of_refresh(ctx, task, task_run, refreshed_kinds) track_event( "sandbox_credentials_refreshed", 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 38051ff5e6ae..2fa5457e05e6 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 @@ -126,6 +126,7 @@ def _deliver_followup(input: SendFollowupToSandboxInput) -> None: # still deliver the follow-up with the existing (possibly stale) creds. ensure_sandbox_identity( task_run, + actor_user, posthog_mcp_scopes=input.posthog_mcp_scopes, auth_token=auth_token, processing_context=input.context, 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 39cad165be42..735e5a03f1a4 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 @@ -28,7 +28,7 @@ get_sandbox_ph_mcp_configs, get_task_run_credential_user, get_user_mcp_server_configs, - mark_mcp_token_issued, + mark_mcp_session, mark_sandbox_identity, ) @@ -328,8 +328,7 @@ def _invoke_start_agent_server( # Keyed on the sandbox id: a replacement sandbox starts unmarked. if params.actor_user_id is not None: if params.mcp_configs: - mark_mcp_token_issued(sandbox.id, params.actor_user_id) - mark_sandbox_identity(sandbox.id, "mcp", params.actor_user_id) + mark_mcp_session(sandbox.id, params.actor_user_id) if ctx.has_github_credentials: mark_sandbox_identity(sandbox.id, "github", params.actor_user_id) diff --git a/products/tasks/backend/temporal/process_task/sandbox_credentials.py b/products/tasks/backend/temporal/process_task/sandbox_credentials.py index 6333a3fc391f..a15c780ec5c4 100644 --- a/products/tasks/backend/temporal/process_task/sandbox_credentials.py +++ b/products/tasks/backend/temporal/process_task/sandbox_credentials.py @@ -13,10 +13,12 @@ from posthog.redis import get_client from products.tasks.backend.exceptions import CredentialUnavailableError +from products.tasks.backend.logic.services.agent_command import CommandResult, send_refresh_session from products.tasks.backend.logic.services.agentsh import ENV_FILE from products.tasks.backend.models import Task, TaskRun from products.tasks.backend.temporal.process_task.utils import ( PrAuthorshipMode, + get_git_identity_env_vars, get_github_token, get_pr_authorship_mode, get_sandbox_github_token, @@ -247,6 +249,19 @@ class GitHubSandboxCredential: kind: str = "github" + def _apply(self, sandbox: "SandboxBase", ctx: "TaskProcessingContext", task: Task, token: str) -> None: + """Write the token and the matching git author into the sandbox. + + The author must track the same actor the token was resolved for, so + every credential write re-derives it from the same state — otherwise + an actor transition would switch pushes but keep attributing commits + to the previous speaker. + """ + apply_github_credentials_to_sandbox(sandbox, ctx.repository, token) + git_identity = get_git_identity_env_vars(task, ctx.state) + if git_identity: + update_sandbox_env_file(sandbox, git_identity) + def refresh(self, sandbox: "SandboxBase", ctx: "TaskProcessingContext", task: Task) -> CredentialRefreshOutcome: if not ctx.has_github_credentials: return CredentialRefreshOutcome( @@ -312,7 +327,7 @@ def refresh(self, sandbox: "SandboxBase", ctx: "TaskProcessingContext", task: Ta self.kind, refreshed=False, next_refresh_seconds=DEFAULT_REFRESH_INTERVAL_SECONDS ) - apply_github_credentials_to_sandbox(sandbox, ctx.repository, token) + self._apply(sandbox, ctx, task, token) return CredentialRefreshOutcome( self.kind, refreshed=True, next_refresh_seconds=github_refresh_interval_seconds(token) ) @@ -328,12 +343,12 @@ def _refresh_shared_user_integration( return CredentialRefreshOutcome( self.kind, refreshed=False, next_refresh_seconds=DEFAULT_REFRESH_INTERVAL_SECONDS ) - apply_github_credentials_to_sandbox(sandbox, ctx.repository, fallback) + self._apply(sandbox, ctx, task, fallback) return CredentialRefreshOutcome( self.kind, refreshed=True, next_refresh_seconds=github_refresh_interval_seconds(fallback) ) if token: - apply_github_credentials_to_sandbox(sandbox, ctx.repository, token) + self._apply(sandbox, ctx, task, token) return CredentialRefreshOutcome( self.kind, refreshed=bool(token), next_refresh_seconds=USER_TOKEN_REFRESH_INTERVAL_SECONDS ) @@ -355,6 +370,22 @@ def _installation_token_fallback(self, ctx: "TaskProcessingContext", task: Task, ) +def notify_sandbox_credentials_refreshed( + task_run: TaskRun, refreshed_kinds: list[str], *, auth_token: str | None +) -> CommandResult: + """Tell the running agent-server which credentials were just re-injected. + + Credentials-only ``refresh_session`` payload (empty ``mcpServers``): the + agent-server logs it and returns without rebuilding the session, so it is + safe to send mid-turn. ``authorship`` rides along so the agent-server + tracks the effective PR authorship for the run. + """ + authorship = (task_run.state or {}).get("pr_authorship_mode") + return send_refresh_session( + task_run, [], auth_token=auth_token, refreshed_credentials=refreshed_kinds, authorship=authorship + ) + + def build_sandbox_credentials(ctx: "TaskProcessingContext") -> list[SandboxCredential]: credentials: list[SandboxCredential] = [] if ctx.has_github_credentials: diff --git a/products/tasks/backend/temporal/process_task/sandbox_identity.py b/products/tasks/backend/temporal/process_task/sandbox_identity.py index 2f7bffa977c5..a39fbb46f075 100644 --- a/products/tasks/backend/temporal/process_task/sandbox_identity.py +++ b/products/tasks/backend/temporal/process_task/sandbox_identity.py @@ -15,32 +15,26 @@ """ import time -from dataclasses import replace from typing import TYPE_CHECKING import structlog from posthog.temporal.oauth import PosthogMcpScopes -from products.tasks.backend.logic.services.agent_command import ( - REFRESH_TIMEOUT_SECONDS, - CommandResult, - send_refresh_session, -) +from products.tasks.backend.logic.services.agent_command import REFRESH_TIMEOUT_SECONDS, send_refresh_session from products.tasks.backend.logic.services.sandbox import Sandbox from products.tasks.backend.models import TaskRun from products.tasks.backend.temporal.oauth import create_oauth_access_token_for_run from products.tasks.backend.temporal.process_task.sandbox_credentials import ( GitHubSandboxCredential, - update_sandbox_env_file, + notify_sandbox_credentials_refreshed, ) from products.tasks.backend.temporal.process_task.utils import ( - get_git_identity_env_vars, + SandboxIdentityKind, get_last_sandbox_identity, get_sandbox_ph_mcp_configs, - get_task_run_credential_user, get_user_mcp_server_configs, - mark_mcp_token_issued, + mark_mcp_session, mark_sandbox_identity, sandbox_identity_scope, should_refresh_mcp_token, @@ -60,6 +54,7 @@ def ensure_sandbox_identity( task_run: TaskRun, + actor_user: "User | None", *, posthog_mcp_scopes: PosthogMcpScopes, auth_token: str | None, @@ -67,6 +62,10 @@ def ensure_sandbox_identity( ) -> None: """Rebind the sandbox's live session to the run's current actor. + ``actor_user`` is the caller's already-resolved credential user + (``get_task_run_credential_user``); None means no valid actor (Slack + fail-closed), which the MCP mint surfaces as its standard warning. + Best-effort and never raises: a failed rebind must not block an otherwise-valid follow-up, and an unmarked failure is retried on the next message. The GitHub surface is only reconciled when the caller supplies @@ -74,12 +73,6 @@ def ensure_sandbox_identity( """ run_id = str(task_run.id) scope = sandbox_identity_scope(run_id, task_run.state) - try: - actor_user = get_task_run_credential_user(task_run.task, task_run.state) - except Exception as e: - logger.warning("sandbox_identity_actor_resolution_failed", run_id=run_id, error=str(e)) - return - try: _ensure_mcp_identity(task_run, actor_user, scope, posthog_mcp_scopes, auth_token) except Exception: @@ -91,6 +84,16 @@ def ensure_sandbox_identity( logger.warning("sandbox_identity_reconcile_failed", kind="github", run_id=run_id, exc_info=True) +def _last_bound_identity(task_run: TaskRun, scope: str, kind: SandboxIdentityKind) -> int | str | None: + """The identity the sandbox's session currently holds for a surface. + + Until a rebind records otherwise, the sandbox holds its boot-time + credentials, which were minted for the task creator — so an absent mark + (never written, evicted, or pre-rollout sandbox) defaults to the creator. + """ + return get_last_sandbox_identity(scope, kind) or task_run.task.created_by_id + + def _ensure_mcp_identity( task_run: TaskRun, actor_user: "User | None", @@ -104,9 +107,7 @@ def _ensure_mcp_identity( through to the mint, which surfaces the standard warning.""" run_id = str(task_run.id) if actor_user is not None: - # Until a refresh records otherwise, the sandbox holds the boot-time - # token, which was minted for the task creator. - last_identity = get_last_sandbox_identity(scope, "mcp") or task_run.task.created_by_id + last_identity = _last_bound_identity(task_run, scope, "mcp") identity_changed = actor_user.id != last_identity if not identity_changed and not should_refresh_mcp_token(scope, actor_user.id): logger.info("refresh_mcp_skipped_within_interval", run_id=run_id, user_id=actor_user.id) @@ -121,13 +122,6 @@ def _ensure_mcp_identity( _rebind_mcp(task_run, actor_user, scope, scopes, auth_token) -def _mark_mcp_session(scope: str, user_id: int) -> None: - """Record whose token the MCP session now holds and start that user's - freshness window.""" - mark_mcp_token_issued(scope, user_id) - mark_sandbox_identity(scope, "mcp", user_id) - - def _rebind_mcp( task_run: TaskRun, actor_user: "User | None", @@ -166,49 +160,39 @@ def _rebind_mcp( # Nothing to push means there is no MCP session to rebind — mark the # window anyway so we don't re-mint a token on every message. if actor_user is not None: - _mark_mcp_session(scope, actor_user.id) + mark_mcp_session(scope, actor_user.id) logger.info("refresh_mcp_skipped_no_configs", run_id=run_id) return mcp_servers = [config.to_dict() for config in mcp_configs] - result = send_refresh_session( - task_run, - mcp_servers, - auth_token=auth_token, - timeout=REFRESH_TIMEOUT_SECONDS, - ) - if result.success: - if actor_user is not None: - _mark_mcp_session(scope, actor_user.id) - logger.info("refresh_mcp_delivered", run_id=run_id, attempts=1) - return - - logger.info( - "refresh_mcp_retrying", - run_id=run_id, - error=result.error, - status_code=result.status_code, - ) - time.sleep(REFRESH_RETRY_DELAY_SECONDS) - retry: CommandResult = send_refresh_session( - task_run, - mcp_servers, - auth_token=auth_token, - timeout=REFRESH_TIMEOUT_SECONDS, - ) - if retry.success: - if actor_user is not None: - _mark_mcp_session(scope, actor_user.id) - logger.info("refresh_mcp_delivered", run_id=run_id, attempts=2) - return - - logger.warning( - "refresh_mcp_failed", - run_id=run_id, - error=retry.error, - status_code=retry.status_code, - ) + for attempt in (1, 2): + result = send_refresh_session( + task_run, + mcp_servers, + auth_token=auth_token, + timeout=REFRESH_TIMEOUT_SECONDS, + ) + if result.success: + if actor_user is not None: + mark_mcp_session(scope, actor_user.id) + logger.info("refresh_mcp_delivered", run_id=run_id, attempts=attempt) + return + if attempt == 1: + logger.info( + "refresh_mcp_retrying", + run_id=run_id, + error=result.error, + status_code=result.status_code, + ) + time.sleep(REFRESH_RETRY_DELAY_SECONDS) + else: + logger.warning( + "refresh_mcp_failed", + run_id=run_id, + error=result.error, + status_code=result.status_code, + ) def _ensure_github_identity( @@ -223,7 +207,11 @@ def _ensure_github_identity( run_id = str(task_run.id) if not processing_context.has_github_credentials: return - last_identity = get_last_sandbox_identity(scope, "github") or task_run.task.created_by_id + sandbox_id = (task_run.state or {}).get("sandbox_id") + if not sandbox_id: + # Nowhere to push yet; the boot path binds the identity itself. + return + last_identity = _last_bound_identity(task_run, scope, "github") if actor_user.id == last_identity: return logger.info( @@ -232,34 +220,25 @@ def _ensure_github_identity( previous_user_id=last_identity, user_id=actor_user.id, ) - if _rebind_github(task_run, processing_context, auth_token): + if _rebind_github(task_run, sandbox_id, processing_context, auth_token): mark_sandbox_identity(scope, "github", actor_user.id) def _rebind_github( task_run: TaskRun, + sandbox_id: str, processing_context: "TaskProcessingContext", auth_token: str | None, ) -> bool: - """Re-inject the actor's GitHub credentials and git author into the live - sandbox. Returns True when the sandbox now reflects the actor (or there is - nothing to rebind), False when the rebind should be retried next message.""" + """Re-inject the actor's GitHub credentials (token + git author) into the + live sandbox. Returns True when the sandbox now reflects the actor (or + there is nothing to rebind), False when the rebind should be retried on + the next message.""" run_id = str(task_run.id) - task = task_run.task - sandbox_id = (task_run.state or {}).get("sandbox_id") - if not sandbox_id: - return False - - # The workflow-start context carries a boot-time snapshot of the run - # state; credential resolution must see the current actor. - live_context = replace(processing_context, state=task_run.state) + live_context = processing_context.with_state(task_run.state) try: sandbox = Sandbox.get_by_id(sandbox_id) - outcome = GitHubSandboxCredential().refresh(sandbox, live_context, task) - if outcome.refreshed: - git_identity = get_git_identity_env_vars(task, task_run.state) - if git_identity: - update_sandbox_env_file(sandbox, git_identity) + outcome = GitHubSandboxCredential().refresh(sandbox, live_context, task_run.task) except Exception: logger.warning("refresh_github_identity_failed", run_id=run_id, exc_info=True) return False @@ -269,14 +248,7 @@ def _rebind_github( # nothing to diverge on; mark so we don't retry every message. return True - authorship = (task_run.state or {}).get("pr_authorship_mode") - notify = send_refresh_session( - task_run, - [], - auth_token=auth_token, - refreshed_credentials=["github"], - authorship=authorship, - ) + notify = notify_sandbox_credentials_refreshed(task_run, ["github"], auth_token=auth_token) if not notify.success: # Credentials already landed in the sandbox; the notification only # feeds the agent-server's debug log. diff --git a/products/tasks/backend/temporal/process_task/tests/helpers.py b/products/tasks/backend/temporal/process_task/tests/helpers.py new file mode 100644 index 000000000000..9d1c27f8204a --- /dev/null +++ b/products/tasks/backend/temporal/process_task/tests/helpers.py @@ -0,0 +1,31 @@ +from unittest.mock import MagicMock + +from products.tasks.backend.temporal.process_task.utils import McpServerConfig + + +def make_task_run_mock(team_id: int = 7, created_by_id: int | None = 42, state: dict | None = None) -> MagicMock: + task = MagicMock() + task.created_by_id = created_by_id + if created_by_id is not None: + task.created_by = MagicMock(id=created_by_id, distinct_id=f"user-{created_by_id}") + else: + task.created_by = None + task_run = MagicMock() + task_run.id = "run-1" + task_run.team_id = team_id + task_run.task = task + task_run.task_id = "task-1" + # Default to None so `(task_run.state or {}).get(...)` returns None cleanly. + # MagicMock auto-attributes would otherwise return further MagicMock objects + # and leak into kwargs passed to `get_sandbox_ph_mcp_configs`. + task_run.state = state + return task_run + + +def make_mcp_config(name: str = "posthog", token: str = "tok") -> McpServerConfig: + return McpServerConfig( + type="http", + name=name, + url="https://mcp.posthog.com/mcp", + headers=[{"name": "Authorization", "value": f"Bearer {token}"}], + ) diff --git a/products/tasks/backend/temporal/process_task/tests/test_sandbox_credentials.py b/products/tasks/backend/temporal/process_task/tests/test_sandbox_credentials.py index 014fd645ff43..bd95da42d93d 100644 --- a/products/tasks/backend/temporal/process_task/tests/test_sandbox_credentials.py +++ b/products/tasks/backend/temporal/process_task/tests/test_sandbox_credentials.py @@ -236,6 +236,26 @@ def test_refresh_applies_coordinated_token(self): resolve.assert_called_once() apply.assert_called_once_with(sandbox, "explore-science/paper-wizard-frontend", "ghu_fresh") + def test_refresh_applies_git_author_alongside_token(self): + # The author must track the same actor as the token: a credential + # write that switched identity but kept the old author would keep + # attributing commits to the previous speaker. + import contextlib + + with contextlib.ExitStack() as stack: + self._as_user_integration_run(stack) + stack.enter_context(patch(f"{MODULE}.resolve_user_github_integration_for_task", return_value=MagicMock())) + stack.enter_context(patch(f"{MODULE}.resolve_coordinated_user_token", return_value="ghu_fresh")) + stack.enter_context(patch(f"{MODULE}.apply_github_credentials_to_sandbox")) + git_env = {"GIT_AUTHOR_NAME": "New Actor", "GIT_AUTHOR_EMAIL": "actor@example.com"} + stack.enter_context(patch(f"{MODULE}.get_git_identity_env_vars", return_value=git_env)) + env_write = stack.enter_context(patch(f"{MODULE}.update_sandbox_env_file")) + + sandbox = MagicMock() + GitHubSandboxCredential().refresh(sandbox, _context(), MagicMock()) + + env_write.assert_called_once_with(sandbox, git_env) + def test_refresh_reports_not_refreshed_when_no_token(self): import contextlib diff --git a/products/tasks/backend/temporal/process_task/tests/test_sandbox_identity.py b/products/tasks/backend/temporal/process_task/tests/test_sandbox_identity.py index e311bfcdfbf7..7809737bdbda 100644 --- a/products/tasks/backend/temporal/process_task/tests/test_sandbox_identity.py +++ b/products/tasks/backend/temporal/process_task/tests/test_sandbox_identity.py @@ -10,8 +10,8 @@ REFRESH_RETRY_DELAY_SECONDS, ensure_sandbox_identity, ) +from products.tasks.backend.temporal.process_task.tests.helpers import make_mcp_config, make_task_run_mock from products.tasks.backend.temporal.process_task.utils import ( - McpServerConfig, _mcp_token_issued_cache_key, _sandbox_identity_cache_key, mark_mcp_token_issued, @@ -20,53 +20,14 @@ pytestmark = pytest.mark.django_db -_GATE_CACHE_KEYS = [ - _mcp_token_issued_cache_key("run-1", 42), - _mcp_token_issued_cache_key("run-1", 99), - _mcp_token_issued_cache_key("sb-2", 42), - _mcp_token_issued_cache_key("sb-2", 99), - _sandbox_identity_cache_key("run-1", "mcp"), - _sandbox_identity_cache_key("sb-2", "mcp"), - _sandbox_identity_cache_key("run-1", "github"), - _sandbox_identity_cache_key("sb-2", "github"), -] - @pytest.fixture(autouse=True) def _clear_identity_cache(): """Ensure each test starts with no recorded token issuances or session identities so the gates don't carry state between tests.""" - cache.delete_many(_GATE_CACHE_KEYS) + cache.clear() yield - cache.delete_many(_GATE_CACHE_KEYS) - - -def _make_mcp_config(name: str = "posthog", token: str = "tok") -> McpServerConfig: - return McpServerConfig( - type="http", - name=name, - url="https://mcp.posthog.com/mcp", - headers=[{"name": "Authorization", "value": f"Bearer {token}"}], - ) - - -def _make_task_run_mock(team_id: int = 7, created_by_id: int | None = 42, state: dict | None = None) -> MagicMock: - task = MagicMock() - task.created_by_id = created_by_id - if created_by_id is not None: - task.created_by = MagicMock(id=created_by_id, distinct_id=f"user-{created_by_id}") - else: - task.created_by = None - task_run = MagicMock() - task_run.id = "run-1" - task_run.team_id = team_id - task_run.task = task - task_run.task_id = "task-1" - # Default to None so `(task_run.state or {}).get(...)` returns None cleanly. - # MagicMock auto-attributes would otherwise return further MagicMock objects - # and leak into kwargs passed to `get_sandbox_ph_mcp_configs`. - task_run.state = state - return task_run + cache.clear() def _make_processing_context(**overrides: object) -> TaskProcessingContext: @@ -84,36 +45,36 @@ def _make_processing_context(**overrides: object) -> TaskProcessingContext: return TaskProcessingContext(**defaults) -def _ensure(task_run, scopes="read_only", auth_token=None, processing_context=None) -> None: +def _ensure(task_run, actor_id: int | None = 42, scopes="read_only", auth_token=None, processing_context=None) -> None: + actor = MagicMock(id=actor_id) if actor_id is not None else None ensure_sandbox_identity( task_run, + actor, posthog_mcp_scopes=scopes, auth_token=auth_token, processing_context=processing_context, ) -def _patch_actor(user_id: int): - """Pin the resolved credential user so tests can drive actor transitions - without building real run state.""" - return patch( - "products.tasks.backend.temporal.process_task.sandbox_identity.get_task_run_credential_user", - return_value=MagicMock(id=user_id), - ) +def _arm_mcp_success(mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh): + mock_oauth.return_value = "fresh-token" + mock_ph_configs.return_value = [make_mcp_config(token="fresh-token")] + mock_user_configs.return_value = [] + mock_send_refresh.return_value = CommandResult(success=True, status_code=200) +@patch("products.tasks.backend.temporal.process_task.sandbox_identity.time.sleep") +@patch("products.tasks.backend.temporal.process_task.sandbox_identity.send_refresh_session") +@patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_user_mcp_server_configs") +@patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_sandbox_ph_mcp_configs") +@patch("products.tasks.backend.temporal.process_task.sandbox_identity.create_oauth_access_token_for_run") class TestRebindMcp: - @patch("products.tasks.backend.temporal.process_task.sandbox_identity.send_refresh_session") - @patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_user_mcp_server_configs") - @patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_sandbox_ph_mcp_configs") - @patch("products.tasks.backend.temporal.process_task.sandbox_identity.create_oauth_access_token_for_run") - def test_success_path_single_call(self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh): - mock_oauth.return_value = "fresh-token" - mock_ph_configs.return_value = [_make_mcp_config(token="fresh-token")] - mock_user_configs.return_value = [] - mock_send_refresh.return_value = CommandResult(success=True, status_code=200) - - task_run = _make_task_run_mock() + def test_success_path_single_call( + self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh, _mock_sleep + ): + _arm_mcp_success(mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh) + + task_run = make_task_run_mock() _ensure(task_run, auth_token="jwt") mock_oauth.assert_called_once_with(task_run.task, task_run.state, scopes="read_only") @@ -127,109 +88,74 @@ def test_success_path_single_call(self, mock_oauth, mock_ph_configs, mock_user_c assert mock_send_refresh.call_args.args[0] is task_run # mcpServers payload is serialized McpServerConfig shape mcp_servers = mock_send_refresh.call_args.args[1] - assert mcp_servers == [_make_mcp_config(token="fresh-token").to_dict()] + assert mcp_servers == [make_mcp_config(token="fresh-token").to_dict()] - @patch("products.tasks.backend.temporal.process_task.sandbox_identity.time.sleep") - @patch("products.tasks.backend.temporal.process_task.sandbox_identity.send_refresh_session") - @patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_user_mcp_server_configs") - @patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_sandbox_ph_mcp_configs") - @patch("products.tasks.backend.temporal.process_task.sandbox_identity.create_oauth_access_token_for_run") def test_retries_once_on_first_failure( self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh, mock_sleep ): - mock_oauth.return_value = "fresh-token" - mock_ph_configs.return_value = [_make_mcp_config()] - mock_user_configs.return_value = [] + _arm_mcp_success(mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh) mock_send_refresh.side_effect = [ CommandResult(success=False, status_code=502, error="transient", retryable=True), CommandResult(success=True, status_code=200), ] - _ensure(_make_task_run_mock()) + _ensure(make_task_run_mock()) assert mock_send_refresh.call_count == 2 mock_sleep.assert_called_once_with(REFRESH_RETRY_DELAY_SECONDS) - @patch("products.tasks.backend.temporal.process_task.sandbox_identity.time.sleep") - @patch("products.tasks.backend.temporal.process_task.sandbox_identity.send_refresh_session") - @patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_user_mcp_server_configs") - @patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_sandbox_ph_mcp_configs") - @patch("products.tasks.backend.temporal.process_task.sandbox_identity.create_oauth_access_token_for_run") def test_two_failures_are_non_fatal( self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh, _mock_sleep ): - mock_oauth.return_value = "fresh-token" - mock_ph_configs.return_value = [_make_mcp_config()] - mock_user_configs.return_value = [] + _arm_mcp_success(mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh) mock_send_refresh.return_value = CommandResult(success=False, status_code=502, error="down") + mock_send_refresh.side_effect = None # Must not raise. - _ensure(_make_task_run_mock()) + _ensure(make_task_run_mock()) assert mock_send_refresh.call_count == 2 # Cache stays empty so the next follow-up retries the dispatch. assert cache.get(_mcp_token_issued_cache_key("run-1", 42)) is None assert cache.get(_sandbox_identity_cache_key("run-1", "mcp")) is None - @patch("products.tasks.backend.temporal.process_task.sandbox_identity.send_refresh_session") - @patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_user_mcp_server_configs") - @patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_sandbox_ph_mcp_configs") - @patch("products.tasks.backend.temporal.process_task.sandbox_identity.create_oauth_access_token_for_run") def test_token_mint_failure_is_non_fatal_and_skips_send( - self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh + self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh, _mock_sleep ): mock_oauth.side_effect = RuntimeError("oauth service down") - _ensure(_make_task_run_mock()) + _ensure(make_task_run_mock()) mock_ph_configs.assert_not_called() mock_user_configs.assert_not_called() mock_send_refresh.assert_not_called() - @patch("products.tasks.backend.temporal.process_task.sandbox_identity.send_refresh_session") - @patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_user_mcp_server_configs") - @patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_sandbox_ph_mcp_configs") - @patch("products.tasks.backend.temporal.process_task.sandbox_identity.create_oauth_access_token_for_run") def test_skips_send_when_no_mcp_configs_resolved( - self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh + self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh, _mock_sleep ): - mock_oauth.return_value = "fresh-token" + _arm_mcp_success(mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh) mock_ph_configs.return_value = [] - mock_user_configs.return_value = [] - _ensure(_make_task_run_mock()) + _ensure(make_task_run_mock()) mock_send_refresh.assert_not_called() - @patch("products.tasks.backend.temporal.process_task.sandbox_identity.send_refresh_session") - @patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_user_mcp_server_configs") - @patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_sandbox_ph_mcp_configs") - @patch("products.tasks.backend.temporal.process_task.sandbox_identity.create_oauth_access_token_for_run") - def test_user_mcp_configs_skipped_when_no_creator( - self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh + def test_user_mcp_configs_skipped_when_no_actor( + self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh, _mock_sleep ): - mock_oauth.return_value = "fresh-token" - mock_ph_configs.return_value = [_make_mcp_config()] - mock_send_refresh.return_value = CommandResult(success=True, status_code=200) + _arm_mcp_success(mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh) - _ensure(_make_task_run_mock(created_by_id=None)) + _ensure(make_task_run_mock(created_by_id=None), actor_id=None) mock_user_configs.assert_not_called() mock_send_refresh.assert_called_once() - @patch("products.tasks.backend.temporal.process_task.sandbox_identity.send_refresh_session") - @patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_user_mcp_server_configs") - @patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_sandbox_ph_mcp_configs") - @patch("products.tasks.backend.temporal.process_task.sandbox_identity.create_oauth_access_token_for_run") def test_scopes_propagate_to_oauth_and_configs( - self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh + self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh, _mock_sleep ): - mock_oauth.return_value = "fresh-token" - mock_ph_configs.return_value = [_make_mcp_config()] - mock_user_configs.return_value = [] - mock_send_refresh.return_value = CommandResult(success=True, status_code=200) + _arm_mcp_success(mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh) - _ensure(_make_task_run_mock(), scopes="full") + _ensure(make_task_run_mock(), scopes="full") mock_oauth.assert_called_once_with(mock_oauth.call_args.args[0], None, scopes="full") mock_ph_configs.assert_called_once_with( @@ -237,55 +163,48 @@ def test_scopes_propagate_to_oauth_and_configs( ) +@patch("products.tasks.backend.temporal.process_task.sandbox_identity.time.sleep") +@patch("products.tasks.backend.temporal.process_task.sandbox_identity.send_refresh_session") +@patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_user_mcp_server_configs") +@patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_sandbox_ph_mcp_configs") +@patch("products.tasks.backend.temporal.process_task.sandbox_identity.create_oauth_access_token_for_run") class TestMcpRefreshIntervalGate: """Refreshes within MCP_TOKEN_REFRESH_INTERVAL_SECONDS of a previous successful issuance for the same actor must be skipped without minting a new token or contacting the sandbox.""" - @patch("products.tasks.backend.temporal.process_task.sandbox_identity.send_refresh_session") - @patch("products.tasks.backend.temporal.process_task.sandbox_identity.create_oauth_access_token_for_run") - def test_skipped_when_token_recently_issued(self, mock_oauth, mock_send_refresh): + def test_skipped_when_token_recently_issued( + self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh, _mock_sleep + ): mark_mcp_token_issued("run-1", 42) - _ensure(_make_task_run_mock()) + _ensure(make_task_run_mock()) mock_oauth.assert_not_called() mock_send_refresh.assert_not_called() - @patch("products.tasks.backend.temporal.process_task.sandbox_identity.send_refresh_session") - @patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_user_mcp_server_configs") - @patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_sandbox_ph_mcp_configs") - @patch("products.tasks.backend.temporal.process_task.sandbox_identity.create_oauth_access_token_for_run") - def test_marks_after_successful_refresh(self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh): - mock_oauth.return_value = "fresh-token" - mock_ph_configs.return_value = [_make_mcp_config()] - mock_user_configs.return_value = [] - mock_send_refresh.return_value = CommandResult(success=True, status_code=200) + def test_marks_after_successful_refresh( + self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh, _mock_sleep + ): + _arm_mcp_success(mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh) - _ensure(_make_task_run_mock()) + _ensure(make_task_run_mock()) # Cache entries now exist → next refresh for this actor within the # interval is gated, and the session identity is recorded. assert cache.get(_mcp_token_issued_cache_key("run-1", 42)) is True assert cache.get(_sandbox_identity_cache_key("run-1", "mcp")) == 42 - @patch("products.tasks.backend.temporal.process_task.sandbox_identity.time.sleep") - @patch("products.tasks.backend.temporal.process_task.sandbox_identity.send_refresh_session") - @patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_user_mcp_server_configs") - @patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_sandbox_ph_mcp_configs") - @patch("products.tasks.backend.temporal.process_task.sandbox_identity.create_oauth_access_token_for_run") def test_marks_after_successful_retry( self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh, _mock_sleep ): - mock_oauth.return_value = "fresh-token" - mock_ph_configs.return_value = [_make_mcp_config()] - mock_user_configs.return_value = [] + _arm_mcp_success(mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh) mock_send_refresh.side_effect = [ CommandResult(success=False, status_code=502, error="transient"), CommandResult(success=True, status_code=200), ] - _ensure(_make_task_run_mock()) + _ensure(make_task_run_mock()) assert cache.get(_mcp_token_issued_cache_key("run-1", 42)) is True @@ -300,23 +219,16 @@ class TestMcpIdentityTransitionGate: session; the marks are keyed per sandbox so a replacement sandbox starts unmarked.""" - def _arm_success(self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh): - mock_oauth.return_value = "fresh-token" - mock_ph_configs.return_value = [_make_mcp_config(token="fresh-token")] - mock_user_configs.return_value = [] - mock_send_refresh.return_value = CommandResult(success=True, status_code=200) - def test_actor_change_bypasses_freshness_window( self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh ): - self._arm_success(mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh) + _arm_mcp_success(mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh) # The creator's token is fresh and the session is bound to them… mark_mcp_token_issued("run-1", 42) mark_sandbox_identity("run-1", "mcp", 42) # …but the next message comes from a different actor. - with _patch_actor(99): - _ensure(_make_task_run_mock()) + _ensure(make_task_run_mock(), actor_id=99) mock_send_refresh.assert_called_once() assert cache.get(_sandbox_identity_cache_key("run-1", "mcp")) == 99 @@ -325,14 +237,13 @@ def test_actor_change_bypasses_freshness_window( def test_switch_back_to_creator_refreshes_despite_fresh_window( self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh ): - self._arm_success(mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh) + _arm_mcp_success(mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh) # The creator's window is still warm, but the session was last bound # to another user — the creator speaking again is a transition. mark_mcp_token_issued("run-1", 42) mark_sandbox_identity("run-1", "mcp", 99) - with _patch_actor(42): - _ensure(_make_task_run_mock()) + _ensure(make_task_run_mock(), actor_id=42) mock_send_refresh.assert_called_once() assert cache.get(_sandbox_identity_cache_key("run-1", "mcp")) == 42 @@ -340,12 +251,11 @@ def test_switch_back_to_creator_refreshes_despite_fresh_window( def test_same_actor_within_window_is_skipped( self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh ): - self._arm_success(mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh) + _arm_mcp_success(mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh) mark_mcp_token_issued("run-1", 99) mark_sandbox_identity("run-1", "mcp", 99) - with _patch_actor(99): - _ensure(_make_task_run_mock()) + _ensure(make_task_run_mock(), actor_id=99) mock_oauth.assert_not_called() mock_send_refresh.assert_not_called() @@ -353,23 +263,20 @@ def test_same_actor_within_window_is_skipped( def test_replacement_sandbox_starts_unmarked( self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh ): - self._arm_success(mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh) + _arm_mcp_success(mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh) # Marks recorded against the run id (legacy scope) must not gate a run # whose state now points at a concrete sandbox. mark_mcp_token_issued("run-1", 42) mark_sandbox_identity("run-1", "mcp", 42) - with _patch_actor(42): - _ensure(_make_task_run_mock(state={"sandbox_id": "sb-2"})) + _ensure(make_task_run_mock(state={"sandbox_id": "sb-2"}), actor_id=42) mock_send_refresh.assert_called_once() assert cache.get(_mcp_token_issued_cache_key("sb-2", 42)) is True assert cache.get(_sandbox_identity_cache_key("sb-2", "mcp")) == 42 -@patch("products.tasks.backend.temporal.process_task.sandbox_identity.send_refresh_session") -@patch("products.tasks.backend.temporal.process_task.sandbox_identity.update_sandbox_env_file") -@patch("products.tasks.backend.temporal.process_task.sandbox_identity.get_git_identity_env_vars") +@patch("products.tasks.backend.temporal.process_task.sandbox_identity.notify_sandbox_credentials_refreshed") @patch("products.tasks.backend.temporal.process_task.sandbox_identity.GitHubSandboxCredential") @patch("products.tasks.backend.temporal.process_task.sandbox_identity.Sandbox") class TestGithubIdentityTransitionGate: @@ -382,108 +289,95 @@ def _quiet_mcp(self, user_id: int, scope: str = "run-1") -> None: mark_mcp_token_issued(scope, user_id) mark_sandbox_identity(scope, "mcp", user_id) - def _arm_refresh(self, mock_credential_cls, mock_git_env, refreshed: bool = True): + def _arm_refresh(self, mock_credential_cls, refreshed: bool = True): mock_credential_cls.return_value.refresh.return_value = CredentialRefreshOutcome( "github", refreshed=refreshed, next_refresh_seconds=60 ) - mock_git_env.return_value = {"GIT_AUTHOR_NAME": "New Actor", "GIT_AUTHOR_EMAIL": "actor@example.com"} - def test_actor_transition_rebinds_credentials_and_author( - self, mock_sandbox_cls, mock_credential_cls, mock_git_env, mock_env_file, mock_send_refresh - ): - self._arm_refresh(mock_credential_cls, mock_git_env) - mock_send_refresh.return_value = CommandResult(success=True, status_code=200) + def test_actor_transition_rebinds_credentials(self, mock_sandbox_cls, mock_credential_cls, mock_notify): + self._arm_refresh(mock_credential_cls) + mock_notify.return_value = CommandResult(success=True, status_code=200) self._quiet_mcp(99, scope="sb-2") mark_sandbox_identity("sb-2", "github", 42) state = {"sandbox_id": "sb-2", "pr_authorship_mode": "user"} - task_run = _make_task_run_mock(state=state) + task_run = make_task_run_mock(state=state) - with _patch_actor(99): - _ensure(task_run, auth_token="jwt", processing_context=_make_processing_context()) + _ensure(task_run, actor_id=99, auth_token="jwt", processing_context=_make_processing_context()) mock_sandbox_cls.get_by_id.assert_called_once_with("sb-2") refresh_call = mock_credential_cls.return_value.refresh.call_args live_ctx = refresh_call.args[1] assert live_ctx.state == state # boot snapshot replaced with live run state - mock_env_file.assert_called_once() - _, notify_kwargs = mock_send_refresh.call_args - assert notify_kwargs["refreshed_credentials"] == ["github"] - assert notify_kwargs["authorship"] == "user" - assert mock_send_refresh.call_args.args[1] == [] + mock_notify.assert_called_once_with(task_run, ["github"], auth_token="jwt") assert cache.get(_sandbox_identity_cache_key("sb-2", "github")) == 99 - def test_same_actor_does_not_touch_sandbox( - self, mock_sandbox_cls, mock_credential_cls, mock_git_env, mock_env_file, mock_send_refresh - ): + def test_same_actor_does_not_touch_sandbox(self, mock_sandbox_cls, mock_credential_cls, mock_notify): self._quiet_mcp(99) mark_sandbox_identity("run-1", "github", 99) - with _patch_actor(99): - _ensure(_make_task_run_mock(), processing_context=_make_processing_context()) + _ensure( + make_task_run_mock(state={"sandbox_id": "run-1"}), + actor_id=99, + processing_context=_make_processing_context(), + ) mock_sandbox_cls.get_by_id.assert_not_called() - mock_send_refresh.assert_not_called() + mock_notify.assert_not_called() - def test_skipped_without_processing_context( - self, mock_sandbox_cls, mock_credential_cls, mock_git_env, mock_env_file, mock_send_refresh - ): + def test_skipped_without_processing_context(self, mock_sandbox_cls, mock_credential_cls, mock_notify): self._quiet_mcp(99) mark_sandbox_identity("run-1", "github", 42) - with _patch_actor(99): - _ensure(_make_task_run_mock(), processing_context=None) + _ensure(make_task_run_mock(), actor_id=99, processing_context=None) mock_sandbox_cls.get_by_id.assert_not_called() - def test_skipped_without_github_credentials( - self, mock_sandbox_cls, mock_credential_cls, mock_git_env, mock_env_file, mock_send_refresh - ): + def test_skipped_without_github_credentials(self, mock_sandbox_cls, mock_credential_cls, mock_notify): self._quiet_mcp(99) mark_sandbox_identity("run-1", "github", 42) context = _make_processing_context(github_integration_id=None) - with _patch_actor(99): - _ensure(_make_task_run_mock(), processing_context=context) + _ensure(make_task_run_mock(), actor_id=99, processing_context=context) mock_sandbox_cls.get_by_id.assert_not_called() def test_rebind_failure_leaves_identity_unmarked_for_retry( - self, mock_sandbox_cls, mock_credential_cls, mock_git_env, mock_env_file, mock_send_refresh + self, mock_sandbox_cls, mock_credential_cls, mock_notify ): mock_credential_cls.return_value.refresh.side_effect = RuntimeError("sandbox unreachable") - self._quiet_mcp(99) - mark_sandbox_identity("run-1", "github", 42) + self._quiet_mcp(99, scope="sb-2") + mark_sandbox_identity("sb-2", "github", 42) - with _patch_actor(99): - _ensure(_make_task_run_mock(), processing_context=_make_processing_context()) + _ensure( + make_task_run_mock(state={"sandbox_id": "sb-2"}), + actor_id=99, + processing_context=_make_processing_context(), + ) - assert cache.get(_sandbox_identity_cache_key("run-1", "github")) == 42 - mock_send_refresh.assert_not_called() + assert cache.get(_sandbox_identity_cache_key("sb-2", "github")) == 42 + mock_notify.assert_not_called() - def test_missing_sandbox_id_leaves_identity_unmarked( - self, mock_sandbox_cls, mock_credential_cls, mock_git_env, mock_env_file, mock_send_refresh - ): + def test_missing_sandbox_id_leaves_identity_unmarked(self, mock_sandbox_cls, mock_credential_cls, mock_notify): self._quiet_mcp(99) mark_sandbox_identity("run-1", "github", 42) - with _patch_actor(99): - _ensure(_make_task_run_mock(state=None), processing_context=_make_processing_context()) + _ensure(make_task_run_mock(state=None), actor_id=99, processing_context=_make_processing_context()) assert cache.get(_sandbox_identity_cache_key("run-1", "github")) == 42 mock_sandbox_cls.get_by_id.assert_not_called() - def test_unrefreshable_credential_marks_without_notify( - self, mock_sandbox_cls, mock_credential_cls, mock_git_env, mock_env_file, mock_send_refresh - ): + def test_unrefreshable_credential_marks_without_notify(self, mock_sandbox_cls, mock_credential_cls, mock_notify): # e.g. a caller-token run: nothing we manage, so nothing can diverge — # mark to avoid re-probing the sandbox on every message. - self._arm_refresh(mock_credential_cls, mock_git_env, refreshed=False) + self._arm_refresh(mock_credential_cls, refreshed=False) self._quiet_mcp(99, scope="sb-2") mark_sandbox_identity("sb-2", "github", 42) - with _patch_actor(99): - _ensure(_make_task_run_mock(state={"sandbox_id": "sb-2"}), processing_context=_make_processing_context()) + _ensure( + make_task_run_mock(state={"sandbox_id": "sb-2"}), + actor_id=99, + processing_context=_make_processing_context(), + ) assert cache.get(_sandbox_identity_cache_key("sb-2", "github")) == 99 - mock_env_file.assert_not_called() - mock_send_refresh.assert_not_called() + mock_notify.assert_not_called() 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 31c100cb0290..fbc20b2d58be 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 @@ -9,29 +9,11 @@ SendFollowupToSandboxInput, send_followup_to_sandbox, ) +from products.tasks.backend.temporal.process_task.tests.helpers import make_task_run_mock pytestmark = pytest.mark.django_db -def _make_task_run_mock(team_id: int = 7, created_by_id: int | None = 42, state: dict | None = None) -> MagicMock: - task = MagicMock() - task.created_by_id = created_by_id - if created_by_id is not None: - task.created_by = MagicMock(id=created_by_id, distinct_id=f"user-{created_by_id}") - else: - task.created_by = None - task_run = MagicMock() - task_run.id = "run-1" - task_run.team_id = team_id - task_run.task = task - task_run.task_id = "task-1" - # Default to None so `(task_run.state or {}).get(...)` returns None cleanly. - # MagicMock auto-attributes would otherwise return further MagicMock objects - # and leak into kwargs passed to `get_sandbox_ph_mcp_configs`. - task_run.state = state - return task_run - - class TestSendFollowupActivityRefreshOrdering: """Refresh call must precede user_message, and the activity must succeed when refresh fails (non-fatal) as long as user_message succeeds.""" @@ -59,7 +41,7 @@ def _patches(self): "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox._write_error_and_complete" ), ): - task_run = _make_task_run_mock() + task_run = make_task_run_mock() task_run.task.created_by = MagicMock(id=42, distinct_id="u42") mock_task_run_cls.objects.select_related.return_value.get.return_value = task_run mock_conn_token.return_value = "jwt" @@ -136,7 +118,7 @@ def _patches(self): "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox._write_error_and_complete" ) as mock_error, ): - task_run = _make_task_run_mock() + task_run = make_task_run_mock() task_run.task.created_by = MagicMock(id=42, distinct_id="u42") mock_task_run_cls.objects.select_related.return_value.get.return_value = task_run mock_conn_token.return_value = "jwt" diff --git a/products/tasks/backend/temporal/process_task/utils.py b/products/tasks/backend/temporal/process_task/utils.py index ca8d739a3e6a..ae520c99dd30 100644 --- a/products/tasks/backend/temporal/process_task/utils.py +++ b/products/tasks/backend/temporal/process_task/utils.py @@ -389,6 +389,14 @@ def should_refresh_mcp_token(scope: str, user_id: int) -> bool: return get_tasks_cache().get(_mcp_token_issued_cache_key(scope, user_id)) is None +def mark_mcp_session(scope: str, user_id: int) -> None: + """Record whose token the sandbox's MCP session now holds and start that + user's freshness window. Issuing a token and recording the session identity + must always happen together, or the identity-transition gate mis-fires.""" + mark_mcp_token_issued(scope, user_id) + mark_sandbox_identity(scope, "mcp", user_id) + + # How long the sandbox's session identity is remembered — comfortably past any # plausible sandbox lifetime. On eviction the identity is assumed to be the # boot-time one (the task creator).