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/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 9adab94feea9..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 @@ -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,6 +96,16 @@ 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) + # 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: + 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) + refreshed_kinds: list[str] = [] orphaned_kinds: list[str] = [] next_refresh = DEFAULT_REFRESH_INTERVAL_SECONDS @@ -185,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 b770a7d03207..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 @@ -1,5 +1,4 @@ import json -import time import threading import contextvars from dataclasses import dataclass @@ -12,35 +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_sandbox_ph_mcp_configs, get_task_run_credential_user, - get_user_mcp_server_configs, is_slack_interaction_state, - mark_mcp_token_issued, - 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. @@ -57,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 @@ -128,10 +121,16 @@ 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, + actor_user, + 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: @@ -216,92 +215,6 @@ def _deliver_followup(input: SendFollowupToSandboxInput) -> None: raise ApplicationError(f"send_followup failed: {error_msg}", non_retryable=True) -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 entirely if a token was issued for this run within the last - MCP_TOKEN_REFRESH_INTERVAL_SECONDS — the in-sandbox token is still fresh. - """ - 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 - try: - actor_user = get_task_run_credential_user(task, task_run.state) - 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: - 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: - mark_mcp_token_issued(run_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: - mark_mcp_token_issued(run_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 19c6581b8c71..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,8 @@ 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, ) from .get_task_processing_context import TaskProcessingContext @@ -164,6 +165,10 @@ class StartAgentServerOutput: @dataclass class _LaunchParams: mcp_configs: list[McpServerConfig] + # 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 event_ingest_token: str | None @@ -275,6 +280,7 @@ def _prepare_launch(ctx: TaskProcessingContext, scopes: PosthogMcpScopes) -> _La return _LaunchParams( mcp_configs=mcp_configs, + 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 +322,15 @@ 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 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.actor_user_id is not None: + if params.mcp_configs: + mark_mcp_session(sandbox.id, 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_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 new file mode 100644 index 000000000000..a39fbb46f075 --- /dev/null +++ b/products/tasks/backend/temporal/process_task/sandbox_identity.py @@ -0,0 +1,257 @@ +"""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 typing import TYPE_CHECKING + +import structlog + +from posthog.temporal.oauth import PosthogMcpScopes + +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, + notify_sandbox_credentials_refreshed, +) +from products.tasks.backend.temporal.process_task.utils import ( + SandboxIdentityKind, + get_last_sandbox_identity, + get_sandbox_ph_mcp_configs, + get_user_mcp_server_configs, + mark_mcp_session, + 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, + actor_user: "User | None", + *, + 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. + + ``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 + 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: + _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 _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", + 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: + 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) + 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 _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] + + 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( + 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 + 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( + "refresh_github_identity_transition", + run_id=run_id, + previous_user_id=last_identity, + user_id=actor_user.id, + ) + 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 (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) + live_context = processing_context.with_state(task_run.state) + try: + sandbox = Sandbox.get_by_id(sandbox_id) + 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 + + 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 + + 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. + 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/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 new file mode 100644 index 000000000000..7809737bdbda --- /dev/null +++ b/products/tasks/backend/temporal/process_task/tests/test_sandbox_identity.py @@ -0,0 +1,383 @@ +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.tests.helpers import make_mcp_config, make_task_run_mock +from products.tasks.backend.temporal.process_task.utils import ( + _mcp_token_issued_cache_key, + _sandbox_identity_cache_key, + mark_mcp_token_issued, + mark_sandbox_identity, +) + +pytestmark = pytest.mark.django_db + + +@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.clear() + yield + cache.clear() + + +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, 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 _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: + 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") + 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()] + + def test_retries_once_on_first_failure( + 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) + 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) + + def test_two_failures_are_non_fatal( + 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) + 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()) + + 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 + + def test_token_mint_failure_is_non_fatal_and_skips_send( + 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()) + + mock_ph_configs.assert_not_called() + mock_user_configs.assert_not_called() + mock_send_refresh.assert_not_called() + + def test_skips_send_when_no_mcp_configs_resolved( + 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) + mock_ph_configs.return_value = [] + + _ensure(make_task_run_mock()) + + mock_send_refresh.assert_not_called() + + def test_user_mcp_configs_skipped_when_no_actor( + 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(created_by_id=None), actor_id=None) + + mock_user_configs.assert_not_called() + mock_send_refresh.assert_called_once() + + def test_scopes_propagate_to_oauth_and_configs( + 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(), 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" + ) + + +@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.""" + + 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()) + + mock_oauth.assert_not_called() + mock_send_refresh.assert_not_called() + + 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()) + + # 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 + + def test_marks_after_successful_retry( + 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) + 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 test_actor_change_bypasses_freshness_window( + self, 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. + _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 + 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 + ): + _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) + + _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 + + def test_same_actor_within_window_is_skipped( + self, 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) + + _ensure(make_task_run_mock(), actor_id=99) + + 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 + ): + _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) + + _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.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: + """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, refreshed: bool = True): + mock_credential_cls.return_value.refresh.return_value = CredentialRefreshOutcome( + "github", refreshed=refreshed, next_refresh_seconds=60 + ) + + 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) + + _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_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_notify): + self._quiet_mcp(99) + mark_sandbox_identity("run-1", "github", 99) + + _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_notify.assert_not_called() + + 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) + + _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_notify): + self._quiet_mcp(99) + mark_sandbox_identity("run-1", "github", 42) + context = _make_processing_context(github_integration_id=None) + + _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_notify + ): + mock_credential_cls.return_value.refresh.side_effect = RuntimeError("sandbox unreachable") + self._quiet_mcp(99, scope="sb-2") + mark_sandbox_identity("sb-2", "github", 42) + + _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")) == 42 + mock_notify.assert_not_called() + + 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) + + _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_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, refreshed=False) + self._quiet_mcp(99, scope="sb-2") + mark_sandbox_identity("sb-2", "github", 42) + + _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_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 5ff23e9eaf7c..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 @@ -1,327 +1,19 @@ 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, - mark_mcp_token_issued, -) +from products.tasks.backend.temporal.process_task.tests.helpers import make_task_run_mock pytestmark = pytest.mark.django_db -@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")) - yield - cache.delete(_mcp_token_issued_cache_key("run-1")) - - -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 - - -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 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") - - _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 entry now exists → next refresh within the interval is gated. - assert cache.get(_mcp_token_issued_cache_key("run-1")) 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_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")) 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")) is None - - class TestSendFollowupActivityRefreshOrdering: """Refresh call must precede user_message, and the activity must succeed when refresh fails (non-fatal) as long as user_message succeeds.""" @@ -337,7 +29,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" @@ -349,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" @@ -384,18 +76,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: @@ -414,7 +106,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" @@ -426,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 c7f3632f69f8..ae520c99dd30 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,77 @@ 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"tasks: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 + + +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). +SANDBOX_IDENTITY_TTL_SECONDS = 7 * 24 * 60 * 60 + +SandboxIdentityKind = Literal["mcp", "github"] + + +def _sandbox_identity_cache_key(scope: str, kind: SandboxIdentityKind) -> str: + return f"tasks: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; + ``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 + 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) 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