Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions docs_src/identity_assertion/tutorial001.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,13 +20,13 @@ def __init__(self) -> None:
async def get_tokens(self) -> OAuthToken | None:
return self.tokens

async def set_tokens(self, tokens: OAuthToken) -> None:
async def set_tokens(self, tokens: OAuthToken | None) -> None:
self.tokens = tokens

async def get_client_info(self) -> OAuthClientInformationFull | None:
return self.client_info

async def set_client_info(self, client_info: OAuthClientInformationFull) -> None:
async def set_client_info(self, client_info: OAuthClientInformationFull | None) -> None:
self.client_info = client_info


Expand Down
4 changes: 2 additions & 2 deletions docs_src/oauth_clients/tutorial001.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,13 +17,13 @@ def __init__(self) -> None:
async def get_tokens(self) -> OAuthToken | None:
return self.tokens

async def set_tokens(self, tokens: OAuthToken) -> None:
async def set_tokens(self, tokens: OAuthToken | None) -> None:
self.tokens = tokens

async def get_client_info(self) -> OAuthClientInformationFull | None:
return self.client_info

async def set_client_info(self, client_info: OAuthClientInformationFull) -> None:
async def set_client_info(self, client_info: OAuthClientInformationFull | None) -> None:
self.client_info = client_info


Expand Down
4 changes: 2 additions & 2 deletions docs_src/oauth_clients/tutorial002.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,13 +14,13 @@ def __init__(self) -> None:
async def get_tokens(self) -> OAuthToken | None:
return self.tokens

async def set_tokens(self, tokens: OAuthToken) -> None:
async def set_tokens(self, tokens: OAuthToken | None) -> None:
self.tokens = tokens

async def get_client_info(self) -> OAuthClientInformationFull | None:
return self.client_info

async def set_client_info(self, client_info: OAuthClientInformationFull) -> None:
async def set_client_info(self, client_info: OAuthClientInformationFull | None) -> None:
self.client_info = client_info


Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -37,13 +37,13 @@ def __init__(self):
async def get_tokens(self) -> OAuthToken | None:
return self._tokens

async def set_tokens(self, tokens: OAuthToken) -> None:
async def set_tokens(self, tokens: OAuthToken | None) -> None:
self._tokens = tokens

async def get_client_info(self) -> OAuthClientInformationFull | None:
return self._client_info

async def set_client_info(self, client_info: OAuthClientInformationFull) -> None:
async def set_client_info(self, client_info: OAuthClientInformationFull | None) -> None:
self._client_info = client_info


Expand Down
4 changes: 2 additions & 2 deletions examples/snippets/clients/identity_assertion_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,13 +34,13 @@ def __init__(self) -> None:
async def get_tokens(self) -> OAuthToken | None:
return self.tokens

async def set_tokens(self, tokens: OAuthToken) -> None:
async def set_tokens(self, tokens: OAuthToken | None) -> None:
self.tokens = tokens

async def get_client_info(self) -> OAuthClientInformationFull | None:
return self.client_info

async def set_client_info(self, client_info: OAuthClientInformationFull) -> None:
async def set_client_info(self, client_info: OAuthClientInformationFull | None) -> None:
self.client_info = client_info


Expand Down
4 changes: 2 additions & 2 deletions examples/snippets/clients/oauth_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,15 +29,15 @@ async def get_tokens(self) -> OAuthToken | None:
"""Get stored tokens."""
return self.tokens

async def set_tokens(self, tokens: OAuthToken) -> None:
async def set_tokens(self, tokens: OAuthToken | None) -> None:
"""Store tokens."""
self.tokens = tokens

async def get_client_info(self) -> OAuthClientInformationFull | None:
"""Get stored client information."""
return self.client_info

async def set_client_info(self, client_info: OAuthClientInformationFull) -> None:
async def set_client_info(self, client_info: OAuthClientInformationFull | None) -> None:
"""Store client information."""
self.client_info = client_info

Expand Down
4 changes: 2 additions & 2 deletions examples/stories/_shared/auth.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,13 +38,13 @@ class InMemoryTokenStorage:
async def get_tokens(self) -> OAuthToken | None:
return self.tokens

async def set_tokens(self, tokens: OAuthToken) -> None:
async def set_tokens(self, tokens: OAuthToken | None) -> None:
self.tokens = tokens

async def get_client_info(self) -> OAuthClientInformationFull | None:
return self.client_info

async def set_client_info(self, client_info: OAuthClientInformationFull) -> None:
async def set_client_info(self, client_info: OAuthClientInformationFull | None) -> None:
self.client_info = client_info


Expand Down
44 changes: 38 additions & 6 deletions src/mcp/client/auth/oauth2.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,13 +5,14 @@

import base64
import hashlib
import json
import logging
import secrets
import string
import time
from collections.abc import AsyncGenerator, Awaitable, Callable
from dataclasses import dataclass, field
from typing import Any, Protocol, get_args
from typing import Any, Protocol, cast, get_args
from urllib.parse import quote, urlencode, urljoin, urlparse

import anyio
Expand Down Expand Up @@ -109,6 +110,21 @@ def check_registration_usable(client_info: OAuthClientInformationFull) -> None:
)


def _is_expired_client_secret(client_info: OAuthClientInformationFull) -> bool:
"""Return whether a stored registration reports an expired client secret."""
expires_at = client_info.client_secret_expires_at
return expires_at is not None and expires_at != 0 and expires_at <= int(time.time())


def _is_invalid_client_response(body: bytes) -> bool:
"""Identify RFC 6749 ``invalid_client`` responses independent of HTTP status."""
try:
payload: Any = json.loads(body)
except (json.JSONDecodeError, UnicodeDecodeError):
return False
return isinstance(payload, dict) and cast(dict[str, Any], payload).get("error") == "invalid_client"


class PKCEParameters(BaseModel):
"""PKCE (Proof Key for Code Exchange) parameters."""

Expand All @@ -131,16 +147,16 @@ async def get_tokens(self) -> OAuthToken | None:
"""Get stored tokens."""
...

async def set_tokens(self, tokens: OAuthToken) -> None:
"""Store tokens."""
async def set_tokens(self, tokens: OAuthToken | None) -> None:
"""Store tokens, or clear them when ``tokens`` is ``None``."""
...

async def get_client_info(self) -> OAuthClientInformationFull | None:
"""Get stored client information."""
...

async def set_client_info(self, client_info: OAuthClientInformationFull) -> None:
"""Store client information."""
async def set_client_info(self, client_info: OAuthClientInformationFull | None) -> None:
"""Store client information, or clear it when ``client_info`` is ``None``."""
...


Expand Down Expand Up @@ -468,6 +484,8 @@ async def _handle_token_response(self, response: httpx2.Response) -> None:
"""Handle token exchange response."""
if response.status_code not in {200, 201}:
body = await response.aread()
if _is_invalid_client_response(body):
await self._clear_stored_credentials()
body_text = body.decode("utf-8")
raise OAuthTokenError(f"Token exchange failed ({response.status_code}): {body_text}")

Expand Down Expand Up @@ -519,8 +537,12 @@ async def _refresh_token(self) -> httpx2.Request:
async def _handle_refresh_response(self, response: httpx2.Response) -> bool:
"""Handle token refresh response. Returns True if successful."""
if response.status_code != 200:
body = await response.aread()
logger.warning(f"Token refresh failed: {response.status_code}")
self.context.clear_tokens()
if _is_invalid_client_response(body):
await self._clear_stored_credentials()
else:
self.context.clear_tokens()
return False

try:
Expand Down Expand Up @@ -551,8 +573,18 @@ 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()
if self.context.client_info and _is_expired_client_secret(self.context.client_info):
logger.info("Stored client registration has expired; clearing credentials and re-registering")
await self._clear_stored_credentials()
self._initialized = True

async def _clear_stored_credentials(self) -> None:
"""Clear the in-memory and persisted credentials bound to a client registration."""
self.context.client_info = None
self.context.clear_tokens()
await self.context.storage.set_client_info(None)
await self.context.storage.set_tokens(None)

def _add_auth_header(self, request: httpx2.Request) -> None:
"""Add authorization header to request if we have valid tokens."""
if self.context.current_tokens and self.context.current_tokens.access_token: # pragma: no branch
Expand Down
4 changes: 2 additions & 2 deletions tests/client/auth/extensions/test_client_credentials.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,13 +27,13 @@ def __init__(self):
async def get_tokens(self) -> OAuthToken | None:
return self._tokens

async def set_tokens(self, tokens: OAuthToken) -> None: # pragma: no cover
async def set_tokens(self, tokens: OAuthToken | None) -> None: # pragma: no cover
self._tokens = tokens

async def get_client_info(self) -> OAuthClientInformationFull | None: # pragma: no cover
return self._client_info

async def set_client_info(self, client_info: OAuthClientInformationFull) -> None: # pragma: no cover
async def set_client_info(self, client_info: OAuthClientInformationFull | None) -> None: # pragma: no cover
self._client_info = client_info


Expand Down
4 changes: 2 additions & 2 deletions tests/client/auth/extensions/test_identity_assertion.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,13 +29,13 @@ def __init__(self, tokens: OAuthToken | None = None) -> None:
async def get_tokens(self) -> OAuthToken | None:
return self.tokens

async def set_tokens(self, tokens: OAuthToken) -> None:
async def set_tokens(self, tokens: OAuthToken | None) -> None:
self.tokens = tokens

async def get_client_info(self) -> OAuthClientInformationFull | None:
raise NotImplementedError

async def set_client_info(self, client_info: OAuthClientInformationFull) -> None:
async def set_client_info(self, client_info: OAuthClientInformationFull | None) -> None:
raise NotImplementedError


Expand Down
95 changes: 93 additions & 2 deletions tests/client/test_auth.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@

from mcp.client.auth import OAuthClientProvider, PKCEParameters
from mcp.client.auth.exceptions import OAuthFlowError, OAuthRegistrationError, OAuthTokenError
from mcp.client.auth.oauth2 import _is_invalid_client_response
from mcp.client.auth.utils import (
build_oauth_authorization_server_metadata_discovery_urls,
build_protected_resource_metadata_discovery_urls,
Expand Down Expand Up @@ -53,13 +54,13 @@ def __init__(self):
async def get_tokens(self) -> OAuthToken | None:
return self._tokens

async def set_tokens(self, tokens: OAuthToken) -> None:
async def set_tokens(self, tokens: OAuthToken | None) -> None:
self._tokens = tokens

async def get_client_info(self) -> OAuthClientInformationFull | None:
return self._client_info

async def set_client_info(self, client_info: OAuthClientInformationFull) -> None:
async def set_client_info(self, client_info: OAuthClientInformationFull | None) -> None:
self._client_info = client_info


Expand All @@ -78,6 +79,20 @@ def client_metadata():
)


@pytest.mark.parametrize(
("body", "expected"),
[
(b"not json", False),
(b"\xff", False),
(b"[]", False),
(b'{"error": "invalid_grant"}', False),
(b'{"error": "invalid_client"}', True),
],
)
def test_invalid_client_response_detection(body: bytes, expected: bool) -> None:
assert _is_invalid_client_response(body) is expected


@pytest.fixture
def valid_tokens():
return OAuthToken(
Expand Down Expand Up @@ -264,6 +279,26 @@ def test_clear_tokens(self, oauth_provider: OAuthClientProvider, valid_tokens: O
assert context.current_tokens is None
assert context.token_expiry_time is None

@pytest.mark.anyio
async def test_initialize_discards_expired_client_registration(
self, oauth_provider: OAuthClientProvider, mock_storage: MockTokenStorage, valid_tokens: OAuthToken
):
expired_client = OAuthClientInformationFull(
client_id="expired-client",
client_secret="expired-secret",
client_secret_expires_at=int(time.time()) - 1,
redirect_uris=[AnyUrl("http://localhost:3030/callback")],
)
await mock_storage.set_client_info(expired_client)
await mock_storage.set_tokens(valid_tokens)

await oauth_provider._initialize()

assert oauth_provider.context.client_info is None
assert oauth_provider.context.current_tokens is None
assert await mock_storage.get_client_info() is None
assert await mock_storage.get_tokens() is None


class TestOAuthFlow:
"""Test OAuth flow methods."""
Expand Down Expand Up @@ -2955,6 +2990,34 @@ async def test_handle_token_response_raises_on_non_2xx_with_body(oauth_provider:
await oauth_provider._handle_token_response(response)


@pytest.mark.anyio
async def test_handle_token_response_invalid_client_clears_stored_credentials(
oauth_provider: OAuthClientProvider, mock_storage: MockTokenStorage, valid_tokens: OAuthToken
):
client_info = OAuthClientInformationFull(
client_id="stale-client",
client_secret="stale-secret",
redirect_uris=[AnyUrl("http://localhost:3030/callback")],
)
oauth_provider.context.client_info = client_info
oauth_provider.context.current_tokens = valid_tokens
await mock_storage.set_client_info(client_info)
await mock_storage.set_tokens(valid_tokens)
response = httpx2.Response(
401,
json={"error": "invalid_client"},
request=httpx2.Request("POST", "https://auth.example.com/token"),
)

with pytest.raises(OAuthTokenError, match=r"Token exchange failed \(401\).*invalid_client"):
await oauth_provider._handle_token_response(response)

assert oauth_provider.context.client_info is None
assert oauth_provider.context.current_tokens is None
assert await mock_storage.get_client_info() is None
assert await mock_storage.get_tokens() is None


@pytest.mark.anyio
async def test_handle_refresh_response_carries_prior_scope_and_refresh_token_when_omitted(
oauth_provider: OAuthClientProvider, mock_storage: MockTokenStorage
Expand Down Expand Up @@ -3007,6 +3070,34 @@ async def test_handle_refresh_response_adopts_rotated_refresh_token_when_returne
assert stored.refresh_token == "rotated"


@pytest.mark.anyio
async def test_handle_refresh_response_invalid_client_clears_stored_credentials(
oauth_provider: OAuthClientProvider, mock_storage: MockTokenStorage, valid_tokens: OAuthToken
):
client_info = OAuthClientInformationFull(
client_id="stale-client",
client_secret="stale-secret",
redirect_uris=[AnyUrl("http://localhost:3030/callback")],
)
oauth_provider.context.client_info = client_info
oauth_provider.context.current_tokens = valid_tokens
await mock_storage.set_client_info(client_info)
await mock_storage.set_tokens(valid_tokens)
response = httpx2.Response(
400,
json={"error": "invalid_client"},
request=httpx2.Request("POST", "https://auth.example.com/token"),
)

ok = await oauth_provider._handle_refresh_response(response)

assert ok is False
assert oauth_provider.context.client_info is None
assert oauth_provider.context.current_tokens is None
assert await mock_storage.get_client_info() is None
assert await mock_storage.get_tokens() is None


@pytest.mark.anyio
async def test_issuer_binding_re_evaluated_after_asm_when_prm_discovery_failed(
oauth_provider: OAuthClientProvider,
Expand Down
Loading
Loading