diff --git a/src/mcp/client/auth/oauth2.py b/src/mcp/client/auth/oauth2.py index 7dc62b52b9..424cc33512 100644 --- a/src/mcp/client/auth/oauth2.py +++ b/src/mcp/client/auth/oauth2.py @@ -5,6 +5,7 @@ import base64 import hashlib +import inspect import logging import secrets import string @@ -551,6 +552,38 @@ async def _initialize(self) -> None: """Load stored tokens and client info.""" self.context.current_tokens = await self.context.storage.get_tokens() self.context.client_info = await self.context.storage.get_client_info() + # Compute the absolute expiry time from the loaded token's relative + # `expires_in`. Without this, `is_token_valid()` short-circuits to True + # when `token_expiry_time` is None and the refresh-on-expiry path in + # `async_auth_flow` never fires — every process restart sends an + # expired access_token, gets a 401, and lands in the full re-auth + # branch. Mirrors what `set_tokens` does on the write path. + if self.context.current_tokens is not None: + self.context.update_token_expiry(self.context.current_tokens) + # Optionally load OAuth metadata from storage. Some downstream storage + # implementations (e.g. ones persisting `.meta.json` from server + # discovery) implement this; SDK-provided storages do not, so the + # `getattr` guard keeps this a no-op when absent. + # + # Without this, `_refresh_token` falls back to `urljoin(base, "/token")` + # which gives the wrong endpoint for any IdP that mounts its token + # endpoint at a non-root path (Hydra/Ory-style: /oauth/token, Auth0: + # /oauth/token, Keycloak: /protocol/openid-connect/token, etc.) and + # silently 404s the refresh grant. + loader = getattr(self.context.storage, "load_oauth_metadata", None) + if callable(loader): + try: + meta = loader() + if inspect.iscoroutine(meta): + meta = await meta + if meta is not None: + self.context.oauth_metadata = meta # type: ignore[assignment] + except Exception: + # Storage implementations are user-provided; a misbehaving + # metadata loader must not break auth initialization. The + # 401-handling path will populate `oauth_metadata` via server + # discovery as a fallback. + pass self._initialized = True def _add_auth_header(self, request: httpx2.Request) -> None: diff --git a/tests/client/test_auth.py b/tests/client/test_auth.py index be96cc8eec..1de396e536 100644 --- a/tests/client/test_auth.py +++ b/tests/client/test_auth.py @@ -636,6 +636,114 @@ async def test_refresh_token_request(self, oauth_provider: OAuthClientProvider, assert "client_id=test_client" in content assert "client_secret=test_secret" in content + @pytest.mark.anyio + async def test_initialize_computes_token_expiry_time(self, oauth_provider: OAuthClientProvider, valid_tokens: OAuthToken): + """`_initialize` must compute `context.token_expiry_time` from the loaded + token's relative `expires_in` — otherwise `is_token_valid()` short- + circuits to True and refresh-on-expiry never fires after a process + restart (see https://github.com/modelcontextprotocol/python-sdk/issues/3250).""" + inner_storage: MockTokenStorage = oauth_provider.context.storage # type: ignore[assignment] + await inner_storage.set_tokens(valid_tokens) + await inner_storage.set_client_info( + OAuthClientInformationFull( + client_id="test_client", + redirect_uris=[AnyUrl("http://localhost:3030/callback")], + token_endpoint_auth_method="none", + ) + ) + + # token_expiry_time should be None before init, populated after init. + assert oauth_provider.context.token_expiry_time is None + await oauth_provider._initialize() + # valid_tokens fixture has expires_in=3600 + assert valid_tokens.expires_in is not None + expected = time.time() + valid_tokens.expires_in + assert oauth_provider.context.token_expiry_time is not None + assert abs(oauth_provider.context.token_expiry_time - expected) < 5 + assert oauth_provider.context.is_token_valid() + + # An immediately-expired token (clamped `expires_in=0` from + # `HermesTokenStorage`-style storage rewrite of a stale on-disk token) + # must produce an invalid state — not short-circuit to True. + class _ExpiredStorage(MockTokenStorage): + async def get_tokens(self) -> OAuthToken | None: + token = await super().get_tokens() + if token is not None: + return token.model_copy(update={"expires_in": 0}) + return None + + wrapped = _ExpiredStorage() + await wrapped.set_tokens(valid_tokens) + client_info = await inner_storage.get_client_info() + assert client_info is not None + await wrapped.set_client_info(client_info) + oauth_provider.context.storage = wrapped # type: ignore[assignment] + oauth_provider._initialized = False + oauth_provider.context.token_expiry_time = None + await oauth_provider._initialize() + assert oauth_provider.context.token_expiry_time is not None + assert oauth_provider.context.token_expiry_time <= time.time() + assert not oauth_provider.context.is_token_valid() + + @pytest.mark.anyio + async def test_initialize_loads_oauth_metadata_from_storage(self, oauth_provider: OAuthClientProvider, valid_tokens: OAuthToken): + """`_initialize` must call `storage.load_oauth_metadata()` if the storage + implements it — otherwise `_refresh_token` falls back to + `urljoin(base, "/token")` which silently 404s against any IdP that + mounts its token endpoint at a non-root path (Hydra, Auth0, Keycloak).""" + canonical_meta = OAuthMetadata( + issuer=AnyHttpUrl("https://hydra.example.com/"), + authorization_endpoint=AnyHttpUrl("https://hydra.example.com/oauth2/auth"), + token_endpoint=AnyHttpUrl("https://hydra.example.com/oauth2/token"), + token_endpoint_auth_methods_supported=["none"], + response_types_supported=["code"], + grant_types_supported=["authorization_code", "refresh_token"], + code_challenge_methods_supported=["S256"], + ) + + class _MetadataStorage(MockTokenStorage): + def load_oauth_metadata(self) -> OAuthMetadata: + return canonical_meta + + meta_storage = _MetadataStorage() + await meta_storage.set_tokens(valid_tokens) + await meta_storage.set_client_info( + OAuthClientInformationFull( + client_id="test_client", + redirect_uris=[AnyUrl("http://localhost:3030/callback")], + token_endpoint_auth_method="none", + ) + ) + oauth_provider.context.storage = meta_storage # type: ignore[assignment] + oauth_provider._initialized = False + + await oauth_provider._initialize() + assert oauth_provider.context.oauth_metadata is canonical_meta + + # Verify the metadata is actually used by _refresh_token — without the + # patch, this would build `https://api.example.com/token` (404 for Hydra). + oauth_provider.context.current_tokens = valid_tokens + request = await oauth_provider._refresh_token() + assert str(request.url) == "https://hydra.example.com/oauth2/token" + + @pytest.mark.anyio + async def test_initialize_metadata_loader_failure_is_non_fatal(self, oauth_provider: OAuthClientProvider, valid_tokens: OAuthToken): + """A misbehaving `load_oauth_metadata` must not break auth init — the + 401-handling path will populate metadata via server discovery as a fallback.""" + class _BrokenMetadataStorage(MockTokenStorage): + def load_oauth_metadata(self) -> OAuthMetadata: + raise RuntimeError("simulated storage failure") + + broken = _BrokenMetadataStorage() + await broken.set_tokens(valid_tokens) + oauth_provider.context.storage = broken # type: ignore[assignment] + + # Should not raise — the try/except in _initialize swallows loader errors. + await oauth_provider._initialize() + assert oauth_provider.context.current_tokens is not None + assert oauth_provider.context.token_expiry_time is not None + assert oauth_provider.context.oauth_metadata is None + @pytest.mark.anyio async def test_basic_auth_token_exchange(self, oauth_provider: OAuthClientProvider): """Test token exchange with client_secret_basic authentication."""