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 856f7bb9e8ff..9751babef63b 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 @@ -29,13 +29,14 @@ from products.tasks.backend.temporal.process_task.utils import ( get_actor_distinct_id, get_imported_mcp_server_configs, + get_sandbox_mcp_session_user, 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_mcp_session, record_message_actor, - should_refresh_mcp_token, + sandbox_identity_scope, ) from ee.hogai.sandbox import STOP_REASON_END_TURN, TURN_COMPLETE_METHOD @@ -158,10 +159,13 @@ 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 - # still deliver the follow-up with the existing (possibly stale) creds. - _refresh_sandbox_mcp(task_run, input.posthog_mcp_scopes, auth_token, actor_user=actor_user, state=state) + # Rebind the sandbox's MCP session to this actor before the turn. On an + # actor transition this must rebind or clear the prior session; if it can't, + # fail closed rather than run the turn under the previous actor's creds. + # Same-actor and first-bind refreshes stay best-effort. + if not _refresh_sandbox_mcp(task_run, input.posthog_mcp_scopes, auth_token, actor_user=actor_user, state=state): + error_msg = "Could not rebind sandbox MCP credentials for the follow-up actor" + raise RuntimeError(f"send_followup failed: {error_msg}") artifacts = None artifact_ids = input.artifact_ids or [] if artifact_ids: @@ -256,28 +260,43 @@ def _refresh_sandbox_mcp( *, actor_user: Any, state: dict[str, Any] | None, -) -> None: - """Mint a fresh OAuth token for the actor and push updated MCP configs to - the sandbox. - - Best-effort: retries once on failure, then logs and returns — a failed - refresh must not block an otherwise-valid follow-up. Skipped when a token - was already issued for this run within MCP_TOKEN_REFRESH_INTERVAL_SECONDS. +) -> bool: + """Rebind the sandbox's MCP session to this message's actor. + + Returns ``True`` when the session is safe to use (unchanged actor or a + successful rebind) and ``False`` when a rebind could not be confirmed — the + caller then fails the follow-up closed. A rebind is unconfirmed whenever the + mint or refresh fails and the binding is not known to be this actor's, + including an *unknown* binding: the marker self-expires at half the token + lifetime, so an absent marker can mean the previous actor's session is still + live, not that the sandbox is fresh. Retries the refresh once before giving + up. """ 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 if actor_user is None: # Without a credential user the mint is guaranteed to fail; skip # quietly rather than warn on every message. - return + return True + + scope = sandbox_identity_scope(run_id, state) + bound_user_id = get_sandbox_mcp_session_user(scope) + if bound_user_id == actor_user.id: + logger.info("refresh_mcp_skipped_within_interval", run_id=run_id, user_id=actor_user.id) + return True + is_transition = bound_user_id is not None + if is_transition: + logger.info( + "refresh_mcp_identity_transition", + run_id=run_id, + previous_user_id=bound_user_id, + user_id=actor_user.id, + ) try: access_token = create_oauth_access_token_for_run(task_run.task, state, scopes=scopes) except Exception as e: logger.warning("refresh_mcp_token_mint_failed", run_id=run_id, error=str(e)) - return + return False # rebind unconfirmed → fail closed (unknown binding may hide a live session) mcp_configs = get_sandbox_ph_mcp_configs( token=access_token, @@ -302,8 +321,22 @@ def _refresh_sandbox_mcp( mcp_configs = mcp_configs + imported_mcp_configs if not mcp_configs: + if is_transition: + # A prior actor holds the live session and this actor resolves no MCP + # configs, so an empty-list refresh (a no-op on the agent-server) + # can neither rebind it nor tear it down. Fail closed rather than run + # the turn against the previous actor's retained session. + logger.info( + "refresh_mcp_no_configs_on_transition_fail_closed", run_id=run_id, previous_user_id=bound_user_id + ) + return False + # No recorded prior actor and no MCP configs to establish a session: + # there is nothing to leak, so let the turn run rather than block the + # agent just because MCP is unavailable. Record the binding so a later + # actor transition is still detected. + mark_sandbox_mcp_session(scope, actor_user.id) logger.info("refresh_mcp_skipped_no_configs", run_id=run_id) - return + return True mcp_servers = [config.to_dict() for config in mcp_configs] @@ -314,9 +347,9 @@ def _refresh_sandbox_mcp( timeout=REFRESH_TIMEOUT_SECONDS, ) if result.success: - mark_mcp_token_issued(run_id) + mark_sandbox_mcp_session(scope, actor_user.id) logger.info("refresh_mcp_delivered", run_id=run_id, attempts=1) - return + return True logger.info( "refresh_mcp_retrying", @@ -332,9 +365,9 @@ def _refresh_sandbox_mcp( timeout=REFRESH_TIMEOUT_SECONDS, ) if retry.success: - mark_mcp_token_issued(run_id) + mark_sandbox_mcp_session(scope, actor_user.id) logger.info("refresh_mcp_delivered", run_id=run_id, attempts=2) - return + return True logger.warning( "refresh_mcp_failed", @@ -342,6 +375,7 @@ def _refresh_sandbox_mcp( error=retry.error, status_code=retry.status_code, ) + return False # rebind never confirmed → fail closed (unknown binding may hide a live session) def _get_stop_reason(result_data: dict[str, Any] | None) -> str: 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 6cc63b2b83b7..10131c6319e8 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 @@ -30,7 +30,7 @@ get_sandbox_ph_mcp_configs, get_task_run_credential_user, get_user_mcp_server_configs, - mark_mcp_token_issued, + mark_sandbox_mcp_session, ) from .get_task_processing_context import TaskProcessingContext @@ -167,6 +167,9 @@ class StartAgentServerOutput: class _LaunchParams: mcp_configs: list[McpServerConfig] relayed_mcp_servers: list[str] + # The user the boot-time credentials were minted for, recorded as the + # sandbox's initial session identity. + actor_user_id: int | None agentsh_domains: list[str] | None protected_base_branch: str | None event_ingest_token: str | None @@ -298,6 +301,7 @@ def _prepare_launch(ctx: TaskProcessingContext, scopes: PosthogMcpScopes) -> _La return _LaunchParams( mcp_configs=mcp_configs, relayed_mcp_servers=relayed_names, + 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, @@ -340,10 +344,10 @@ 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 boot identity so same-actor follow-ups within the + # freshness window skip the redundant refresh. + if params.mcp_configs and params.actor_user_id is not None: + mark_sandbox_mcp_session(sandbox.id, 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/tests/test_send_followup_to_sandbox.py b/products/tasks/backend/temporal/process_task/tests/test_send_followup_to_sandbox.py index 300c995909f8..adca99b8e9ab 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 @@ -15,20 +15,21 @@ ) from products.tasks.backend.temporal.process_task.utils import ( McpServerConfig, - _mcp_token_issued_cache_key, - mark_mcp_token_issued, + _sandbox_mcp_session_cache_key, + get_sandbox_mcp_session_user, + mark_sandbox_mcp_session, ) 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 +def _clear_session_cache(): + """Ensure each test starts with no recorded session bindings so the refresh gate doesn't carry state between tests.""" - cache.delete(_mcp_token_issued_cache_key("run-1")) + cache.clear() yield - cache.delete(_mcp_token_issued_cache_key("run-1")) + cache.clear() def _make_mcp_config(name: str = "posthog", token: str = "tok") -> McpServerConfig: @@ -65,22 +66,23 @@ def _refresh(task_run, actor_id: int | None = 42, scopes="read_only", auth_token _refresh_sandbox_mcp(task_run, scopes, auth_token, actor_user=actor, state=task_run.state) +def _arm_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.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" +) 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) + def test_success_path_single_call(self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh, _sleep): + _arm_success(mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh) task_run = _make_task_run_mock() _refresh(task_run, auth_token="jwt") @@ -98,18 +100,8 @@ def test_success_path_single_call(self, mock_oauth, mock_ph_configs, mock_user_c 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.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_refresh_keeps_imported_mcp_servers( - self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh + self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh, _sleep ): """refresh_session replaces the session's server list wholesale; without this, the run's client-imported servers vanish at the first token refresh.""" @@ -141,23 +133,10 @@ def test_refresh_keeps_imported_mcp_servers( "headers": [{"name": "Authorization", "value": "Bearer x"}], } - @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 = [] + _arm_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), @@ -167,43 +146,24 @@ def test_retries_once_on_first_failure( assert mock_send_refresh.call_count == 2 mock_sleep.assert_called_once_with(REFRESH_RETRY_DELAY_SECONDS) + # Marked on the successful retry → next same-actor refresh is gated. + assert get_sandbox_mcp_session_user("run-1") == 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_two_failures_are_non_fatal( - self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh, _mock_sleep + def test_two_failures_are_non_fatal_and_unmarked( + self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh, _sleep ): - mock_oauth.return_value = "fresh-token" - mock_ph_configs.return_value = [_make_mcp_config()] - mock_user_configs.return_value = [] + _arm_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") # Must not raise. _refresh(_make_task_run_mock()) assert mock_send_refresh.call_count == 2 + # Cache stays empty so the next follow-up retries the dispatch. + assert get_sandbox_mcp_session_user("run-1") is None - @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 + self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh, _sleep ): mock_oauth.side_effect = RuntimeError("oauth service down") @@ -213,61 +173,30 @@ def test_token_mint_failure_is_non_fatal_and_skips_send( 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 + self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh, _sleep ): - mock_oauth.return_value = "fresh-token" + _arm_success(mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh) mock_ph_configs.return_value = [] - mock_user_configs.return_value = [] _refresh(_make_task_run_mock()) mock_send_refresh.assert_not_called() + # Marked anyway: with no session to rebind, don't re-mint per message. + assert get_sandbox_mcp_session_user("run-1") == 42 - @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_no_actor_skips_entirely(self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh): - # Without a credential user the mint can only fail — skip quietly. + def test_no_actor_skips_entirely(self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh, _sleep): + # Creator-less non-Slack run: the mint is guaranteed to fail, so the + # refresh must not attempt (and warn) on every message. _refresh(_make_task_run_mock(created_by_id=None), actor_id=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_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, _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_success(mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh) _refresh(_make_task_run_mock(), scopes="full") @@ -276,94 +205,145 @@ def test_scopes_propagate_to_oauth_and_configs( token="fresh-token", project_id=7, scopes="full", interaction_origin=None, task_id="task-1" ) + def test_transition_refresh_failure_reports_unsafe( + self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh, _sleep + ): + # A prior actor holds the session and the new actor's refresh fails both + # attempts: the rebind never happened, so the gate reports unsafe (the + # caller fails the follow-up closed) and the previous binding is left as is. + 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") + mark_sandbox_mcp_session("run-1", 99) -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.""" + actor = MagicMock(id=42) + safe = _refresh_sandbox_mcp(_make_task_run_mock(), "read_only", None, actor_user=actor, state=None) - @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") + assert safe is False + assert get_sandbox_mcp_session_user("run-1") == 99 - _refresh(_make_task_run_mock()) + def test_unknown_binding_refresh_failure_fails_closed( + self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh, _sleep + ): + # No marker for this scope: the marker self-expires before the OAuth + # session does, so an absent one may hide the previous actor's still-live + # session rather than a fresh sandbox. When the refresh can't confirm the + # rebind, the gate reports unsafe so the caller fails closed. + 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") + + actor = MagicMock(id=42) + safe = _refresh_sandbox_mcp(_make_task_run_mock(), "read_only", None, actor_user=actor, state=None) + + assert safe is False + + +@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 TestSessionIdentityGate: + """One cache entry per sandbox records who the live session was last bound + to and expires with the freshness window — so a same-actor repeat skips, + while a transition, an expired entry, or a replacement sandbox refreshes.""" + + def test_same_actor_within_window_is_skipped( + self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh + ): + mark_sandbox_mcp_session("run-1", 42) + + _refresh(_make_task_run_mock(), actor_id=42) 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) + def test_actor_change_bypasses_freshness_window( + self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh + ): + _arm_success(mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh) + # The session is freshly bound to the creator… + mark_sandbox_mcp_session("run-1", 42) - _refresh(_make_task_run_mock()) + # …but the next message comes from a different actor. + _refresh(_make_task_run_mock(), actor_id=99) + + mock_send_refresh.assert_called_once() + assert get_sandbox_mcp_session_user("run-1") == 99 - # Cache entry now exists → next refresh within the interval is gated. - assert cache.get(_mcp_token_issued_cache_key("run-1")) is True + def test_switch_back_refreshes(self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh): + _arm_success(mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh) + # The session was last bound to another user — the creator speaking + # again is a transition even though they spoke recently. + mark_sandbox_mcp_session("run-1", 99) - @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 + _refresh(_make_task_run_mock(), actor_id=42) + + mock_send_refresh.assert_called_once() + assert get_sandbox_mcp_session_user("run-1") == 42 + + def test_unknown_binding_refreshes(self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh): + # No entry (expired window, cache eviction, pre-rollout sandbox): + # fail safe by refreshing rather than guessing who the session holds. + _arm_success(mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh) + + _refresh(_make_task_run_mock(), actor_id=42) + + mock_send_refresh.assert_called_once() + assert get_sandbox_mcp_session_user("run-1") == 42 + + def test_replacement_sandbox_starts_unmarked( + self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh + ): + _arm_success(mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh) + # A binding recorded against the run id (legacy scope) must not gate a + # run whose state now points at a concrete sandbox. + mark_sandbox_mcp_session("run-1", 42) + + _refresh(_make_task_run_mock(state={"sandbox_id": "sb-2"}), actor_id=42) + + mock_send_refresh.assert_called_once() + assert get_sandbox_mcp_session_user("sb-2") == 42 + assert cache.get(_sandbox_mcp_session_cache_key("run-1")) == 42 # untouched + + def test_transition_with_no_configs_fails_closed( + self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh ): + # The prior actor holds the live session and the new actor resolves no MCP + # configs, so an empty-list refresh can neither rebind nor tear it down. + # Reject the turn rather than run it against the prior actor's session. mock_oauth.return_value = "fresh-token" - mock_ph_configs.return_value = [_make_mcp_config()] + mock_ph_configs.return_value = [] mock_user_configs.return_value = [] - mock_send_refresh.side_effect = [ - CommandResult(success=False, status_code=502, error="transient"), - CommandResult(success=True, status_code=200), - ] + mark_sandbox_mcp_session("run-1", 99) - _refresh(_make_task_run_mock()) + actor = MagicMock(id=42) + safe = _refresh_sandbox_mcp(_make_task_run_mock(), "read_only", None, actor_user=actor, state=None) - assert cache.get(_mcp_token_issued_cache_key("run-1")) is True + assert safe is False # fail closed: prior session may still be live + mock_send_refresh.assert_not_called() + assert get_sandbox_mcp_session_user("run-1") == 99 # binding unchanged - @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 + def test_unknown_binding_with_no_configs_runs( + self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh ): + # No recorded prior actor and no MCP configs to establish a session: there + # is nothing to leak, so the turn runs rather than being blocked just + # because MCP is unavailable. The binding is recorded for later transitions. mock_oauth.return_value = "fresh-token" - mock_ph_configs.return_value = [_make_mcp_config()] + mock_ph_configs.return_value = [] mock_user_configs.return_value = [] - mock_send_refresh.return_value = CommandResult(success=False, status_code=502, error="down") - _refresh(_make_task_run_mock()) + actor = MagicMock(id=42) + safe = _refresh_sandbox_mcp(_make_task_run_mock(), "read_only", None, actor_user=actor, state=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 safe is True + mock_send_refresh.assert_not_called() + assert get_sandbox_mcp_session_user("run-1") == 42 # binding recorded class TestSendFollowupActivityRefreshOrdering: @@ -411,6 +391,7 @@ def test_refresh_called_before_user_message(self, _patches): def _record_refresh(*a, **kw): call_order.append("refresh") + return True # refresh confirmed the session is safe; gate lets the turn proceed def _record_user_msg(*a, **kw): call_order.append("user_message") diff --git a/products/tasks/backend/temporal/process_task/utils.py b/products/tasks/backend/temporal/process_task/utils.py index 6e7d99dfc110..9dca99c83fb5 100644 --- a/products/tasks/backend/temporal/process_task/utils.py +++ b/products/tasks/backend/temporal/process_task/utils.py @@ -361,23 +361,33 @@ 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 marks describing a run's live sandbox. + + Keyed on the sandbox id so a replacement sandbox (fresh provision, + snapshot restore, workflow retry) starts unmarked — nothing ever needs + clearing. Falls back to the run id when state has no sandbox id yet. + """ + return (state or {}).get("sandbox_id") or run_id + + +def _sandbox_mcp_session_cache_key(scope: str) -> str: + return f"tasks:sandbox-mcp-session:{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_mcp_session(scope: str, user_id: int) -> None: + """Record whose OAuth token the sandbox's live MCP session holds. - The cache entry self-expires after MCP_TOKEN_REFRESH_INTERVAL_SECONDS, so - `should_refresh_mcp_token` returns True again past that window. + Self-expires after MCP_TOKEN_REFRESH_INTERVAL_SECONDS, so an absent + entry always reads as "must refresh". """ - get_tasks_cache().set(_mcp_token_issued_cache_key(run_id), True, timeout=MCP_TOKEN_REFRESH_INTERVAL_SECONDS) + get_tasks_cache().set(_sandbox_mcp_session_cache_key(scope), user_id, timeout=MCP_TOKEN_REFRESH_INTERVAL_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_sandbox_mcp_session_user(scope: str) -> int | None: + """User id the sandbox's MCP session was last bound to within the + freshness window, or None when unknown.""" + return get_tasks_cache().get(_sandbox_mcp_session_cache_key(scope)) @dataclass(frozen=True)