diff --git a/backend/src/agents/main_agent/base_agent.py b/backend/src/agents/main_agent/base_agent.py index 8d795276..75c41b74 100644 --- a/backend/src/agents/main_agent/base_agent.py +++ b/backend/src/agents/main_agent/base_agent.py @@ -340,18 +340,6 @@ async def disconnected_lookup(provider_id: str) -> bool: self.user_id, provider_id ) - async def mark_disconnected(provider_id: str) -> None: - # Persist a disconnect from the AfterToolCallEvent 401-retry - # path so subsequent requests (potentially on other replicas) - # also force a fresh consent. - from apis.shared.oauth.disconnect_repository import ( - get_disconnect_repository, - ) - - await get_disconnect_repository().mark_disconnected( - self.user_id, provider_id - ) - return OAuthConsentHook( user_id=self.user_id, provider_lookup=provider_lookup, @@ -359,7 +347,6 @@ async def mark_disconnected(provider_id: str) -> None: provider_type_lookup=provider_type_lookup, custom_parameters_lookup=custom_parameters_lookup, disconnected_lookup=disconnected_lookup, - mark_disconnected=mark_disconnected, ) def _build_filtered_tools(self) -> List: diff --git a/backend/src/agents/main_agent/integrations/oauth_token_cache.py b/backend/src/agents/main_agent/integrations/oauth_token_cache.py index 5feb18aa..0ce2481f 100644 --- a/backend/src/agents/main_agent/integrations/oauth_token_cache.py +++ b/backend/src/agents/main_agent/integrations/oauth_token_cache.py @@ -9,33 +9,66 @@ * `OAuthConsentHook` warms the cache after a successful vault lookup, or * the resume path re-fetches a token after the user completes consent. -Tokens are evicted explicitly via `clear_user_provider` when consent is -revoked or expires; we don't track expiry locally because AgentCore -Identity owns refresh. +Each entry has a TTL — by default a little under the upstream token's +lifetime (Google's access tokens are 3600s; we expire locally at 3000s). +On expiry `get` returns None so `OAuthConsentHook._gate` re-asks AgentCore +Identity, which transparently refreshes the access token using the +refresh_token in the vault. Without a TTL the cache would hand out an +expired token until the upstream MCP server returned 401, and the 401 +retry path would force a full re-consent — i.e. the user reconnects +hourly even though refresh would have worked. -Disconnect intent ("user pressed Disconnect" / "tool returned 401") is *not* -held here — it lives in the DDB-backed `OAuthDisconnectRepository` so it's -visible across replicas. The cache only holds tokens. +Tokens are evicted explicitly via `clear_user_provider` when consent is +revoked. Disconnect intent ("user pressed Disconnect") is *not* held here +— it lives in the DDB-backed `OAuthDisconnectRepository` so it's visible +across replicas. The cache only holds tokens. """ from __future__ import annotations import threading +import time from typing import Optional +# Default TTL for cached access tokens. Google access tokens last 3600s; +# expiring locally at 3000s gives us a 10-minute safety margin so we +# refresh before the upstream rejects the token. Most other providers +# (Microsoft, GitHub) use the same or longer lifetimes, so this is a safe +# floor across the board. +DEFAULT_TTL_SECONDS = 3000 _lock = threading.Lock() -_cache: dict[tuple[str, str], str] = {} +# (user_id, provider_id) -> (token, expires_at_monotonic) +_cache: dict[tuple[str, str], tuple[str, float]] = {} def get(user_id: str, provider_id: str) -> Optional[str]: - with _lock: - return _cache.get((user_id, provider_id)) + """Return the cached token if present and not expired, else None. - -def set(user_id: str, provider_id: str, token: str) -> None: + Expired entries are evicted on access; we don't run a sweeper because + the cache is small (one entry per (user, provider) pair) and a stale + entry that's never read costs nothing. + """ + with _lock: + entry = _cache.get((user_id, provider_id)) + if entry is None: + return None + token, expires_at = entry + if time.monotonic() >= expires_at: + del _cache[(user_id, provider_id)] + return None + return token + + +def set( + user_id: str, + provider_id: str, + token: str, + *, + ttl_seconds: float = DEFAULT_TTL_SECONDS, +) -> None: with _lock: - _cache[(user_id, provider_id)] = token + _cache[(user_id, provider_id)] = (token, time.monotonic() + ttl_seconds) def clear_user_provider(user_id: str, provider_id: str) -> None: diff --git a/backend/src/agents/main_agent/session/hooks/oauth_consent.py b/backend/src/agents/main_agent/session/hooks/oauth_consent.py index f4ed4bd0..d4009011 100644 --- a/backend/src/agents/main_agent/session/hooks/oauth_consent.py +++ b/backend/src/agents/main_agent/session/hooks/oauth_consent.py @@ -130,12 +130,6 @@ def _looks_like_auth_failure(tool_result: Any) -> bool: # fresh consent URL with `force_authentication=True`. DisconnectedLookup = Callable[[str], Union[bool, Awaitable[bool]]] -# Records a disconnect for the caller — invoked from the AfterToolCallEvent -# path when a tool returns a 401 against the cached vault token. Called -# instead of mutating per-process state so the intent is durable across -# replicas. -MarkDisconnected = Callable[[str], Union[None, Awaitable[None]]] - class OAuthConsentHook(HookProvider): """Pause the agent if a tool needs OAuth and we don't have a token yet.""" @@ -148,7 +142,6 @@ def __init__( provider_type_lookup: Optional[ProviderTypeLookup] = None, custom_parameters_lookup: Optional[CustomParametersLookup] = None, disconnected_lookup: Optional[DisconnectedLookup] = None, - mark_disconnected: Optional[MarkDisconnected] = None, ): """Initialize. @@ -169,11 +162,6 @@ def __init__( effectively assumes the user has not disconnected. Wire this to the durable disconnect repository in production so a /disconnect on one replica is visible from any other. - mark_disconnected: See `MarkDisconnected`. Optional. Invoked - from the 401-retry path; without it, a 401 still flips - `event.retry = True` but leaves no durable record, so the - next BeforeToolCallEvent on a different replica won't know - to force a fresh consent. """ self._user_id = user_id self._provider_lookup = provider_lookup @@ -181,7 +169,6 @@ def __init__( self._provider_type_lookup = provider_type_lookup self._custom_parameters_lookup = custom_parameters_lookup self._disconnected_lookup = disconnected_lookup - self._mark_disconnected = mark_disconnected # Cache scopes per provider for the lifetime of this hook (one agent # invocation). Avoids repeated DB hits if the same provider is used # across multiple tool calls in a single turn. @@ -196,10 +183,9 @@ def __init__( # Providers that already burned their one 401-retry in the current # turn. The agent instance is cached across turns by `get_agent`, so # this set must be reset on `BeforeInvocationEvent`. Without the cap, - # a misconfigured provider (wrong scope, perma-401) would surface a - # consent prompt on every tool call in the turn — `_record_disconnect` - # forces fresh consent on the next BeforeToolCallEvent, the user - # consents, the tool 401s again, and the loop repeats per tool use. + # a misconfigured provider (wrong scope, perma-401) would loop: + # the cache is cleared, AgentCore returns the same expired/invalid + # token from the vault, the tool 401s again, and so on. self._reauth_attempted_providers: set[str] = set() def register_hooks(self, registry: HookRegistry, **kwargs: Any) -> None: @@ -298,7 +284,11 @@ async def _fetch_token_or_url( scopes=scopes, user_id=self._user_id, force_authentication=force_authentication, - custom_parameters=custom_parameters_for(provider_type, admin_extras), + custom_parameters=custom_parameters_for( + provider_type, + admin_extras, + force_authentication=force_authentication, + ), ) except WorkloadTokenUnavailableError: logger.error( @@ -328,19 +318,28 @@ async def _fetch_token_or_url( } async def _handle_auth_failure(self, event: AfterToolCallEvent) -> None: - """Detect a 401 from an OAuth-gated MCP tool and retry with fresh consent. - - AgentCore Identity has no revoke API, so when the user revokes our - app at the provider (or the refresh token expires), AgentCore's - vault keeps serving the now-stale token. The MCP server rejects it - with a 401 — and that's where the staleness first becomes visible. + """Detect a 401 from an OAuth-gated MCP tool and retry. - We detect the 401 in the tool result, mark the (user, provider) - for forced re-consent in the cache, and set `event.retry = True`. + We clear the local hot-path cache and set `event.retry = True`. Strands' tool executor then re-fires `BeforeToolCallEvent`, our - `_gate` callback sees the force-reauth flag, asks AgentCore for a - fresh consent URL with `force_authentication=True`, and raises an - interrupt — same path as a first-time consent. + `_gate` callback misses the cache, and re-asks AgentCore Identity. + AgentCore handles refresh transparently: + * If the access_token was just expired, it uses the vault's + refresh_token to mint a new one and returns it — the user + never sees a prompt. + * If the refresh_token itself is dead (user revoked our app at + the provider, or it lapsed), AgentCore returns an + authorization URL instead, which `_gate` surfaces as the + standard `oauth_required` interrupt. + + We deliberately do NOT write the durable disconnect flag here. A + 401 is most commonly an expired access token, and writing the + flag would force `_gate` to call AgentCore with + `force_authentication=True` — which bypasses the vault entirely, + ignores the still-valid refresh_token, and prompts the user to + re-consent unnecessarily. The disconnect flag is reserved for + explicit user intent (the "Disconnect" button in the settings + page). """ provider_id = self._provider_lookup(event.selected_tool) if not provider_id: @@ -370,11 +369,9 @@ async def _handle_auth_failure(self, event: AfterToolCallEvent) -> None: provider_id, ) # Drop the local hot-path token so the BeforeToolCallEvent retry - # doesn't short-circuit to it, and record the intent durably so - # other replicas (and subsequent requests on this one) also force a - # fresh consent. + # doesn't short-circuit to it. The retry will re-fetch from + # AgentCore Identity, which handles refresh internally. oauth_token_cache.clear_user_provider(self._user_id, provider_id) - await self._record_disconnect(provider_id) event.retry = True async def _resolve_scopes(self, provider_id: str) -> list[str]: @@ -434,10 +431,3 @@ async def _is_disconnected(self, provider_id: str) -> bool: return bool(await result) return bool(result) - async def _record_disconnect(self, provider_id: str) -> None: - """Persist a disconnect from the AfterToolCallEvent retry path.""" - if self._mark_disconnected is None: - return - result = self._mark_disconnected(provider_id) - if inspect.isawaitable(result): - await result diff --git a/backend/src/apis/app_api/connectors/routes.py b/backend/src/apis/app_api/connectors/routes.py index 21bab2af..fc606096 100644 --- a/backend/src/apis/app_api/connectors/routes.py +++ b/backend/src/apis/app_api/connectors/routes.py @@ -229,7 +229,9 @@ async def initiate_consent( user_id=current_user.user_id, force_authentication=force_auth, custom_parameters=custom_parameters_for( - provider.provider_type.value, provider.custom_parameters + provider.provider_type.value, + provider.custom_parameters, + force_authentication=force_auth, ), # No custom_state: AgentCore appears to treat its presence as a # signal to start a fresh flow, never short-circuiting to the diff --git a/backend/src/apis/shared/oauth/agentcore_identity.py b/backend/src/apis/shared/oauth/agentcore_identity.py index 3ab96418..ba694a6f 100644 --- a/backend/src/apis/shared/oauth/agentcore_identity.py +++ b/backend/src/apis/shared/oauth/agentcore_identity.py @@ -88,7 +88,9 @@ async def poll_for_token(self) -> str: _CALLBACK_URL_ENV = "AGENTCORE_LOCAL_OAUTH_CALLBACK_URL" -def _vendor_baseline_params(provider_type: Optional[str]) -> Dict[str, str]: +def _vendor_baseline_params( + provider_type: Optional[str], *, force_authentication: bool = False +) -> Dict[str, str]: """Hardcoded params AgentCore Identity *requires* for a given vendor. Per the AgentCore Identity authentication docs @@ -97,30 +99,51 @@ def _vendor_baseline_params(provider_type: Optional[str]) -> Dict[str, str]: without it the vault entry expires after ~1 hour with no refresh path. This is non-negotiable: it always wins over admin-supplied extras to prevent an admin from accidentally turning it off. + + When `force_authentication=True` we additionally send `prompt=consent` + for Google. Google only re-issues a refresh token on subsequent + authorizations if the user is shown the consent screen again — so a + user who clicks "Disconnect" then "Reconnect" would otherwise get a + vault entry with an access token but no refresh token, putting them + right back in the hourly-reconsent loop on the next access-token + expiry. We only set this on the explicit re-consent path; first-time + consent and silent refreshes don't trigger it. """ if not provider_type: return {} if provider_type.lower() == "google": - return {"access_type": "offline"} + params = {"access_type": "offline"} + if force_authentication: + params["prompt"] = "consent" + return params return {} def custom_parameters_for( provider_type: Optional[str], admin_extras: Optional[Dict[str, str]] = None, + *, + force_authentication: bool = False, ) -> Optional[Dict[str, str]]: """Build the `customParameters` payload AgentCore Identity wants forwarded. Merges admin-supplied extras (e.g. Google `hd=mycorp.com` for domain - restriction, `prompt=consent` for stricter UX) with the hardcoded - vendor baseline. Baseline keys win on conflict — admins cannot turn - off a documented requirement. + restriction) with the hardcoded vendor baseline. Baseline keys win on + conflict — admins cannot turn off a documented requirement. + + `force_authentication=True` mirrors the same flag on + `get_token_for_user`: when the caller is forcing AgentCore to bypass + the vault and walk the user through consent again, this enables any + vendor-specific extras needed for that path (e.g. Google's + `prompt=consent` so a refresh token is re-issued). Returns None when the merged result would be empty, so callers can pass the value through to the SDK unconditionally without sending an empty `customParameters` map. """ - baseline = _vendor_baseline_params(provider_type) + baseline = _vendor_baseline_params( + provider_type, force_authentication=force_authentication + ) merged = {**(admin_extras or {}), **baseline} return merged or None diff --git a/backend/tests/agents/main_agent/integrations/test_oauth_token_cache.py b/backend/tests/agents/main_agent/integrations/test_oauth_token_cache.py index 939ec006..2fa3e0d0 100644 --- a/backend/tests/agents/main_agent/integrations/test_oauth_token_cache.py +++ b/backend/tests/agents/main_agent/integrations/test_oauth_token_cache.py @@ -1,5 +1,8 @@ """Tests for the in-process OAuth token cache.""" +import time +from unittest.mock import patch + from agents.main_agent.integrations import oauth_token_cache @@ -59,3 +62,91 @@ def test_clear_user_provider_drops_only_that_pair(): assert oauth_token_cache.get("tester", "google") is None assert oauth_token_cache.get("tester", "github") == "gh" + + +class TestExpiry: + """The TTL is the load-bearing piece that keeps the cache from + handing out tokens past their upstream lifetime. Without it, the 401 + retry path used to mark the user as disconnected and force re-consent + every ~1 hour. + """ + + def test_get_returns_none_after_ttl_elapses(self): + _isolate() + + with patch("agents.main_agent.integrations.oauth_token_cache.time") as mock_time: + mock_time.monotonic.return_value = 1000.0 + oauth_token_cache.set("tester", "google", "tok", ttl_seconds=60) + + mock_time.monotonic.return_value = 1059.999 + assert oauth_token_cache.get("tester", "google") == "tok" + + mock_time.monotonic.return_value = 1060.0 + assert oauth_token_cache.get("tester", "google") is None + + def test_expired_entry_is_evicted_on_read(self): + """An expired entry must not linger — a subsequent set with the + same key should not be merged with stale state.""" + _isolate() + + with patch("agents.main_agent.integrations.oauth_token_cache.time") as mock_time: + mock_time.monotonic.return_value = 0.0 + oauth_token_cache.set("tester", "google", "old", ttl_seconds=10) + + mock_time.monotonic.return_value = 100.0 + assert oauth_token_cache.get("tester", "google") is None + # Internal dict no longer contains the stale entry. + assert ("tester", "google") not in oauth_token_cache._cache + + def test_default_ttl_is_under_googles_access_token_lifetime(self): + """3000s gives ~10 min of safety margin under Google's 3600s + access token lifetime — we want to refresh proactively, not + wait for the upstream 401.""" + assert oauth_token_cache.DEFAULT_TTL_SECONDS <= 3000 + # Sanity floor: a too-aggressive TTL would defeat the cache. + assert oauth_token_cache.DEFAULT_TTL_SECONDS >= 60 + + def test_set_uses_default_ttl_when_unspecified(self): + _isolate() + + with patch("agents.main_agent.integrations.oauth_token_cache.time") as mock_time: + mock_time.monotonic.return_value = 0.0 + oauth_token_cache.set("tester", "google", "tok") + + # Just under default TTL — still cached. + mock_time.monotonic.return_value = ( + oauth_token_cache.DEFAULT_TTL_SECONDS - 1 + ) + assert oauth_token_cache.get("tester", "google") == "tok" + + # Past default TTL — gone. + mock_time.monotonic.return_value = ( + oauth_token_cache.DEFAULT_TTL_SECONDS + 1 + ) + assert oauth_token_cache.get("tester", "google") is None + + def test_set_resets_expiry(self): + """A fresh set after expiry stores a new entry with a new clock, + so the old expiry doesn't leak through.""" + _isolate() + + with patch("agents.main_agent.integrations.oauth_token_cache.time") as mock_time: + mock_time.monotonic.return_value = 0.0 + oauth_token_cache.set("tester", "google", "old", ttl_seconds=10) + + mock_time.monotonic.return_value = 100.0 + assert oauth_token_cache.get("tester", "google") is None + + # Re-set should not be subject to the old expiry. + oauth_token_cache.set("tester", "google", "new", ttl_seconds=10) + mock_time.monotonic.return_value = 105.0 + assert oauth_token_cache.get("tester", "google") == "new" + + def test_real_clock_short_ttl_expires(self): + """End-to-end check without clock mocking — guard against the + mock hiding a real-world bug.""" + _isolate() + oauth_token_cache.set("tester", "google", "tok", ttl_seconds=0.05) + assert oauth_token_cache.get("tester", "google") == "tok" + time.sleep(0.1) + assert oauth_token_cache.get("tester", "google") is None diff --git a/backend/tests/agents/main_agent/session/test_oauth_consent_hook.py b/backend/tests/agents/main_agent/session/test_oauth_consent_hook.py index 14367c6f..144b1db5 100644 --- a/backend/tests/agents/main_agent/session/test_oauth_consent_hook.py +++ b/backend/tests/agents/main_agent/session/test_oauth_consent_hook.py @@ -113,6 +113,71 @@ async def test_disconnected_lookup_bypasses_token_cache(self): kwargs = identity.get_token_for_user.call_args.kwargs assert kwargs["force_authentication"] is True + @pytest.mark.asyncio + async def test_force_authentication_sends_prompt_consent_for_google(self): + """Google only re-issues a refresh token on subsequent grants if + the user is shown the consent screen — so the explicit re-consent + path (force_authentication=True) must propagate prompt=consent + through to AgentCore Identity. Without it, a Disconnect/Reconnect + cycle leaves the vault with an access token but no refresh token, + putting the user back in the hourly-reconsent loop.""" + identity = MagicMock() + identity.get_token_for_user = AsyncMock( + return_value=TokenResult(authorization_url="https://accounts/consent") + ) + + hook = OAuthConsentHook( + user_id="alice", + provider_lookup=lambda _tool: "google", + scopes_lookup=lambda _: ["openid"], + provider_type_lookup=lambda _: "google", + disconnected_lookup=lambda _pid: True, + ) + event = _make_event(provider_id="google") + + with patch( + "agents.main_agent.session.hooks.oauth_consent.get_agentcore_identity_client", + return_value=identity, + ): + with pytest.raises(InterruptException): + await hook._gate(event) + + kwargs = identity.get_token_for_user.call_args.kwargs + assert kwargs["force_authentication"] is True + assert kwargs["custom_parameters"] == { + "access_type": "offline", + "prompt": "consent", + } + + @pytest.mark.asyncio + async def test_silent_refresh_does_not_send_prompt_consent(self): + """The refresh path (force_authentication=False) must not send + prompt=consent — that would force the consent screen on every + silent refresh, which is the exact UX we're trying to avoid.""" + identity = MagicMock() + identity.get_token_for_user = AsyncMock( + return_value=TokenResult(access_token="refreshed-token") + ) + + hook = OAuthConsentHook( + user_id="alice", + provider_lookup=lambda _tool: "google", + scopes_lookup=lambda _: ["openid"], + provider_type_lookup=lambda _: "google", + ) + event = _make_event(provider_id="google") + + with patch( + "agents.main_agent.session.hooks.oauth_consent.get_agentcore_identity_client", + return_value=identity, + ): + await hook._gate(event) + + kwargs = identity.get_token_for_user.call_args.kwargs + assert kwargs["force_authentication"] is False + assert kwargs["custom_parameters"] == {"access_type": "offline"} + assert "prompt" not in kwargs["custom_parameters"] + @pytest.mark.asyncio async def test_uses_cached_token_without_calling_identity(self): oauth_token_cache.set("alice", "google", "cached-token") @@ -316,8 +381,13 @@ def test_parallel_tool_calls_same_provider_produce_distinct_interrupt_ids(self): class TestOAuthConsentHookAuthFailureRetry: - """The AfterToolCallEvent handler turns a 401-style tool error into - a retry that forces re-consent at AgentCore Identity.""" + """The AfterToolCallEvent handler turns a 401-style tool error into a + retry that re-fetches the token from AgentCore Identity. The retry + deliberately does NOT write the durable disconnect flag — most 401s + are just an expired access token, and AgentCore can refresh + transparently using the refresh_token in the vault. Forcing a + disconnect on every 401 caused users to reconnect hourly even though + refresh would have worked.""" def _after_event( self, @@ -341,17 +411,11 @@ def _after_event( return event @pytest.mark.asyncio - async def test_401_records_disconnect_and_retries(self): - recorded: list[str] = [] - - async def mark_disconnected(pid: str) -> None: - recorded.append(pid) - + async def test_401_clears_cache_and_retries(self): hook = OAuthConsentHook( user_id="alice", provider_lookup=lambda _tool: "google", scopes_lookup=lambda _: [], - mark_disconnected=mark_disconnected, ) oauth_token_cache.set("alice", "google", "stale-token") event = self._after_event( @@ -362,13 +426,29 @@ async def mark_disconnected(pid: str) -> None: await hook._handle_auth_failure(event) assert event.retry is True - # Durable record of the disconnect intent so other replicas force - # fresh consent on the next request, too. - assert recorded == ["google"] # Local cache cleared so the BeforeToolCallEvent retry doesn't - # short-circuit on this replica. + # short-circuit on this replica — it'll re-fetch from AgentCore + # Identity, which can refresh transparently. assert oauth_token_cache.get("alice", "google") is None + @pytest.mark.asyncio + async def test_401_does_not_write_durable_disconnect_flag(self): + """Regression guard for the hourly-reconnect bug: writing the + disconnect flag here would force `_gate` to call AgentCore with + `force_authentication=True`, bypassing the vault and prompting + the user to reconsent even though the refresh_token was valid. + The flag is reserved for the explicit Disconnect button in the + settings page.""" + # The hook no longer accepts a mark_disconnected hook at all — + # the parameter and the `_record_disconnect` method are gone. + # If either reappears, the tests for the consent flow itself + # should explicitly opt back in. + import inspect + + sig = inspect.signature(OAuthConsentHook.__init__) + assert "mark_disconnected" not in sig.parameters + assert not hasattr(OAuthConsentHook, "_record_disconnect") + @pytest.mark.asyncio async def test_non_oauth_tool_is_ignored(self): hook = OAuthConsentHook( @@ -384,24 +464,19 @@ async def test_non_oauth_tool_is_ignored(self): @pytest.mark.asyncio async def test_non_auth_error_is_ignored(self): - recorded: list[str] = [] - - async def mark_disconnected(pid: str) -> None: - recorded.append(pid) - hook = OAuthConsentHook( user_id="alice", provider_lookup=lambda _tool: "google", scopes_lookup=lambda _: [], - mark_disconnected=mark_disconnected, ) + oauth_token_cache.set("alice", "google", "good-token") event = self._after_event("google", "Network unreachable") await hook._handle_auth_failure(event) assert event.retry is False - # No disconnect persisted — the failure wasn't auth-related. - assert recorded == [] + # Cache untouched — the failure wasn't auth-related. + assert oauth_token_cache.get("alice", "google") == "good-token" @pytest.mark.asyncio async def test_does_not_retry_twice_for_same_tool_use(self): @@ -427,16 +502,10 @@ async def test_caps_retry_across_tool_calls_in_same_turn(self): on every tool call in a turn. Cap at one retry per provider per turn so subsequent 401s for the same provider just surface to the model instead of triggering another consent flow.""" - recorded: list[str] = [] - - async def mark_disconnected(pid: str) -> None: - recorded.append(pid) - hook = OAuthConsentHook( user_id="alice", provider_lookup=lambda _tool: "google", scopes_lookup=lambda _: [], - mark_disconnected=mark_disconnected, ) # First tool call 401s — retry path fires. @@ -454,17 +523,12 @@ async def mark_disconnected(pid: str) -> None: ) await hook._handle_auth_failure(event2) assert event2.retry is False - # Disconnect was already recorded on the first 401 — don't write - # again. - assert recorded == ["google"] @pytest.mark.asyncio async def test_before_invocation_event_resets_per_turn_budget(self): """The agent instance is cached across turns by `get_agent`, so the per-provider retry budget on the hook must be reset whenever a new agent invocation begins (fresh turn or resume).""" - from unittest.mock import MagicMock - hook = OAuthConsentHook( user_id="alice", provider_lookup=lambda _tool: "google", diff --git a/backend/tests/apis/shared/oauth/test_agentcore_identity.py b/backend/tests/apis/shared/oauth/test_agentcore_identity.py index 6dafd8c6..c1b7e89d 100644 --- a/backend/tests/apis/shared/oauth/test_agentcore_identity.py +++ b/backend/tests/apis/shared/oauth/test_agentcore_identity.py @@ -67,6 +67,38 @@ def test_empty_admin_extras_treated_as_none(self) -> None: assert custom_parameters_for("microsoft", {}) is None assert custom_parameters_for("microsoft", None) is None + def test_force_authentication_adds_prompt_consent_for_google(self) -> None: + # Google only re-issues a refresh token on subsequent grants when + # the consent screen is shown — so the explicit re-consent path + # must add prompt=consent on top of the baseline. + result = custom_parameters_for("google", force_authentication=True) + assert result == {"access_type": "offline", "prompt": "consent"} + + def test_force_authentication_does_not_set_prompt_for_other_vendors( + self, + ) -> None: + # Other vendors don't have the same constraint — leaving prompt + # unset means we don't accidentally annoy a Microsoft / GitHub + # user with an unnecessary consent screen on every reconnect. + assert custom_parameters_for("microsoft", force_authentication=True) is None + assert custom_parameters_for("github", force_authentication=True) is None + + def test_force_authentication_baseline_still_wins_over_admin(self) -> None: + # Even when force_authentication adds prompt=consent, an admin + # supplying prompt=login can't override it — the re-consent path + # needs the consent screen specifically. + result = custom_parameters_for( + "google", {"prompt": "login"}, force_authentication=True + ) + assert result == {"access_type": "offline", "prompt": "consent"} + + def test_default_force_authentication_is_false(self) -> None: + # Silent refresh path must not get prompt=consent — otherwise + # every refresh would force the consent screen. + result = custom_parameters_for("google") + assert result == {"access_type": "offline"} + assert "prompt" not in result + class TestTokenResult: def test_access_token_only_is_valid(self) -> None: