Skip to content

feat(auth): persist oauth_metadata via TokenStorage to fix refresh after restart#2492

Open
mangibaud33 wants to merge 1 commit intomodelcontextprotocol:mainfrom
mangibaud33:fix/auth-persist-oauth-metadata
Open

feat(auth): persist oauth_metadata via TokenStorage to fix refresh after restart#2492
mangibaud33 wants to merge 1 commit intomodelcontextprotocol:mainfrom
mangibaud33:fix/auth-persist-oauth-metadata

Conversation

@mangibaud33
Copy link
Copy Markdown

Summary

OAuthClientProvider._initialize() restores current_tokens and client_info from storage but not oauth_metadata. After a restart with cached tokens, _refresh_token() loses the authoritative token endpoint discovered during the prior session and falls back to <base_url>/token — wrong for servers whose token endpoint sits on a non-standard path.

This PR makes metadata persistence symmetric with tokens/client_info: optional get_oauth_metadata / set_oauth_metadata methods on TokenStorage, loaded in _initialize, saved after OASM discovery.

Closes part of #1318 (Bug 2 — the _refresh_token fallback URL issue). Bug 1 from #1318 (token_expiry_time not restored) is tracked separately in #1784.

Repro

Concrete case: HubSpot MCP Auth App (mcp.hubspot.com). Its discovery document reports "token_endpoint": "https://mcp.hubspot.com/oauth/v3/token". After completing the initial OAuth flow and storing tokens to disk, a restart causes the next refresh attempt to hit https://mcp.hubspot.com/token (the <base_url>/token fallback), which returns 404 HTML. The httpx.Response.json() validation then fails, _handle_refresh_response returns False, and the SDK cascades into a full interactive OAuth flow. That flow cannot complete in non-interactive environments (daemons, containers, CI).

Fix approach

Alternative considered: fetch discovery inline in _initialize via a private httpx.AsyncClient. Rejected — inconsistent with the yield-based flow, complicates testing (would need to mock HTTP in _initialize tests), and can't easily use the same TLS / proxy configuration the host httpx client is using.

Chosen: persist oauth_metadata alongside tokens/client_info. Symmetric with existing storage pattern, zero extra HTTP call on hot path, cache survives restarts.

Changes

src/mcp/client/auth/oauth2.py

  • TokenStorage Protocol: add two methods with default no-op bodies (return None / return). Marked as optional in docstrings.
  • _initialize: load oauth_metadata via getattr(storage, 'get_oauth_metadata', None) — preserves backward compatibility with duck-typed storage implementations predating this API (they return None and the refresh path falls back to the legacy behaviour).
  • async_auth_flow 401 branch: after successful OASM discovery, persist via getattr(storage, 'set_oauth_metadata', None).
  • Removed three obsolete # pragma: no cover markers on lines now exercised by the new tests (_initialize, async_auth_flow._initialize call, async_auth_flow except Exception), per AGENTS.md.

tests/client/test_auth.py

5 new tests:

Test Purpose
test_initialize_restores_oauth_metadata Happy path: _initialize reads metadata from storage into context
test_refresh_token_uses_persisted_metadata_endpoint The actual bug: after restart with persisted metadata, _refresh_token uses the correct non-standard endpoint
test_initialize_backward_compat_without_metadata_methods Duck-typed legacy storage (no get_oauth_metadata) doesn't raise AttributeError on init
test_token_storage_protocol_default_metadata_methods Protocol defaults return None / no-op for subclasses choosing not to override
test_auth_flow_discovery_with_legacy_storage_skips_metadata_persistence Full OAuth flow with legacy storage: OASM discovery completes, getattr fallback skips persistence silently

MockTokenStorage also updated to implement the two new methods.

Backward compatibility

  • Protocol Subclasses: inherit new default no-op implementations automatically. No action required.
  • Duck-typed implementations (matching the 4 required methods without inheriting the Protocol class — the existing MockTokenStorage pattern): continue to work thanks to the getattr fallback in the provider. They will not persist metadata and will re-trigger discovery on each restart, identical to pre-PR behaviour.

Zero breaking change for any existing TokenStorage implementation.

Test plan

  • uv run --frozen pytest tests/client/test_auth.py — 101 passed
  • uv run --frozen pyright src/mcp/client/auth/oauth2.py tests/client/test_auth.py — 0 errors
  • uv run --frozen ruff format + check — clean
  • coverage report --include='src/mcp/client/auth/oauth2.py' — 99.38% (parity with main baseline; the only missing line is the pre-existing return data, headers in prepare_token_auth, covered by other test files in CI)
  • UV_FROZEN=1 uv run --frozen strict-no-cover — ✅ (removed 3 obsolete pragmas)
  • Verified locally against a running HubSpot MCP Auth App server: gateway reconnect no longer triggers browser OAuth flow.

Notes for reviewers

…ter restart

Closes part of modelcontextprotocol#1318.

Problem: OAuthClientProvider._initialize() restores tokens and client_info
from storage but not oauth_metadata. After a restart with cached tokens,
_refresh_token() falls back to <base_url>/token (via get_authorization_base_url
+ urljoin), which is incorrect for servers whose token endpoint is at a
non-standard path (e.g. HubSpot's MCP Auth App uses /oauth/v3/token).
Refresh requests return 404, cascading into a full interactive OAuth flow
that cannot complete in non-interactive environments (daemons, containers,
long-running services).

Fix: add optional get_oauth_metadata / set_oauth_metadata methods to the
TokenStorage protocol. _initialize now restores metadata alongside tokens
and client_info via a getattr fallback that preserves backward compatibility
with storage implementations predating this API; they return None and the
refresh path falls back to the legacy behaviour as before. Discovered
metadata is persisted after OASM discovery in the 401 flow so subsequent
restarts can resolve the correct token endpoint without rediscovery.

Coverage: added 5 tests (3 feature + 2 to cover the getattr fallback
branches and the Protocol default bodies). Also removed three "# pragma:
no cover" markers on lines now exercised by the new tests per AGENTS.md.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant