Skip to content

Commit 4886fb5

Browse files
maxpetrusenkoagentmaxpetrusenkoagent
committed
Extract reusable OAuth authorization redirect builder
1 parent cf110e3 commit 4886fb5

3 files changed

Lines changed: 208 additions & 30 deletions

File tree

‎src/mcp/client/auth/__init__.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,16 +5,20 @@
55

66
from mcp.client.auth.exceptions import OAuthFlowError, OAuthRegistrationError, OAuthTokenError
77
from mcp.client.auth.oauth2 import (
8+
OAuthAuthorizationRedirect,
89
OAuthClientProvider,
910
PKCEParameters,
1011
TokenStorage,
12+
build_authorization_redirect,
1113
)
1214

1315
__all__ = [
16+
"OAuthAuthorizationRedirect",
1417
"OAuthClientProvider",
1518
"OAuthFlowError",
1619
"OAuthRegistrationError",
1720
"OAuthTokenError",
1821
"PKCEParameters",
1922
"TokenStorage",
23+
"build_authorization_redirect",
2024
]

‎src/mcp/client/auth/oauth2.py‎

Lines changed: 92 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
from collections.abc import AsyncGenerator, Awaitable, Callable
1313
from dataclasses import dataclass, field
1414
from typing import Any, Protocol
15-
from urllib.parse import quote, urlencode, urljoin, urlparse
15+
from urllib.parse import parse_qsl, quote, urlencode, urljoin, urlparse, urlunparse
1616

1717
import anyio
1818
import httpx
@@ -69,6 +69,85 @@ def generate(cls) -> "PKCEParameters":
6969
return cls(code_verifier=code_verifier, code_challenge=code_challenge)
7070

7171

72+
@dataclass(frozen=True)
73+
class OAuthAuthorizationRedirect:
74+
"""Resumable OAuth authorization redirect state.
75+
76+
Proxy and server-side callers can persist this value, send the authorization
77+
URL to a user, and later resume token exchange with the returned code plus
78+
the stored state and code verifier.
79+
"""
80+
81+
authorization_url: str
82+
state: str
83+
code_verifier: str = field(repr=False)
84+
85+
86+
def build_authorization_redirect(
87+
*,
88+
authorization_endpoint: str,
89+
client_info: OAuthClientInformationFull,
90+
client_metadata: OAuthClientMetadata,
91+
pkce_params: PKCEParameters | None = None,
92+
state: str | None = None,
93+
resource_url: str | None = None,
94+
) -> OAuthAuthorizationRedirect:
95+
"""Build an OAuth authorization URL and resumable state.
96+
97+
Args:
98+
authorization_endpoint: Authorization endpoint URL.
99+
client_info: Registered OAuth client information.
100+
client_metadata: Client metadata containing redirect URIs and scopes.
101+
pkce_params: Optional PKCE parameters. Generated when omitted.
102+
state: Optional OAuth state value. Generated when omitted.
103+
resource_url: Optional RFC 8707 resource value.
104+
105+
Returns:
106+
Authorization URL plus the state and code verifier needed to resume.
107+
108+
Raises:
109+
OAuthFlowError: If no client ID or redirect URI is available.
110+
"""
111+
if client_info.client_id is None:
112+
raise OAuthFlowError("No client ID provided for authorization code grant")
113+
114+
if client_metadata.redirect_uris is None:
115+
raise OAuthFlowError("No redirect URIs provided for authorization code grant")
116+
117+
pkce_params = pkce_params or PKCEParameters.generate()
118+
state = state or secrets.token_urlsafe(32)
119+
120+
auth_params = {
121+
"response_type": "code",
122+
"client_id": client_info.client_id,
123+
"redirect_uri": str(client_metadata.redirect_uris[0]),
124+
"state": state,
125+
"code_challenge": pkce_params.code_challenge,
126+
"code_challenge_method": "S256",
127+
}
128+
129+
if resource_url:
130+
auth_params["resource"] = resource_url
131+
132+
if client_metadata.scope:
133+
auth_params["scope"] = client_metadata.scope
134+
135+
# OIDC requires prompt=consent when offline_access is requested
136+
# https://openid.net/specs/openid-connect-core-1_0.html#OfflineAccess
137+
if "offline_access" in client_metadata.scope.split():
138+
auth_params["prompt"] = "consent"
139+
140+
parsed_endpoint = urlparse(authorization_endpoint)
141+
query_params = parse_qsl(parsed_endpoint.query, keep_blank_values=True)
142+
query_params.extend(auth_params.items())
143+
authorization_url = urlunparse(parsed_endpoint._replace(query=urlencode(query_params)))
144+
return OAuthAuthorizationRedirect(
145+
authorization_url=authorization_url,
146+
state=state,
147+
code_verifier=pkce_params.code_verifier,
148+
)
149+
150+
72151
class TokenStorage(Protocol):
73152
"""Protocol for token storage implementations."""
74153

@@ -327,45 +406,29 @@ async def _perform_authorization_code_grant(self) -> tuple[str, str]:
327406
if not self.context.client_info:
328407
raise OAuthFlowError("No client info available for authorization") # pragma: no cover
329408

330-
# Generate PKCE parameters
331-
pkce_params = PKCEParameters.generate()
332-
state = secrets.token_urlsafe(32)
333-
334-
auth_params = {
335-
"response_type": "code",
336-
"client_id": self.context.client_info.client_id,
337-
"redirect_uri": str(self.context.client_metadata.redirect_uris[0]),
338-
"state": state,
339-
"code_challenge": pkce_params.code_challenge,
340-
"code_challenge_method": "S256",
341-
}
342-
343-
# Only include resource param if conditions are met
409+
resource_url = None
344410
if self.context.should_include_resource_param(self.context.protocol_version):
345-
auth_params["resource"] = self.context.get_resource_url() # RFC 8707
411+
resource_url = self.context.get_resource_url() # RFC 8707
346412

347-
if self.context.client_metadata.scope: # pragma: no branch
348-
auth_params["scope"] = self.context.client_metadata.scope
349-
350-
# OIDC requires prompt=consent when offline_access is requested
351-
# https://openid.net/specs/openid-connect-core-1_0.html#OfflineAccess
352-
if "offline_access" in self.context.client_metadata.scope.split():
353-
auth_params["prompt"] = "consent"
354-
355-
authorization_url = f"{auth_endpoint}?{urlencode(auth_params)}"
356-
await self.context.redirect_handler(authorization_url)
413+
redirect = build_authorization_redirect(
414+
authorization_endpoint=auth_endpoint,
415+
client_info=self.context.client_info,
416+
client_metadata=self.context.client_metadata,
417+
resource_url=resource_url,
418+
)
419+
await self.context.redirect_handler(redirect.authorization_url)
357420

358421
# Wait for callback
359422
auth_code, returned_state = await self.context.callback_handler()
360423

361-
if returned_state is None or not secrets.compare_digest(returned_state, state):
362-
raise OAuthFlowError(f"State parameter mismatch: {returned_state} != {state}")
424+
if returned_state is None or not secrets.compare_digest(returned_state, redirect.state):
425+
raise OAuthFlowError(f"State parameter mismatch: {returned_state} != {redirect.state}")
363426

364427
if not auth_code:
365428
raise OAuthFlowError("No authorization code received")
366429

367430
# Return auth code and code verifier for token exchange
368-
return auth_code, pkce_params.code_verifier
431+
return auth_code, redirect.code_verifier
369432

370433
def _get_token_endpoint(self) -> str:
371434
if self.context.oauth_metadata and self.context.oauth_metadata.token_endpoint:

‎tests/client/test_auth.py‎

Lines changed: 112 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
from inline_snapshot import Is, snapshot
1111
from pydantic import AnyHttpUrl, AnyUrl
1212

13-
from mcp.client.auth import OAuthClientProvider, PKCEParameters
13+
from mcp.client.auth import OAuthClientProvider, PKCEParameters, build_authorization_redirect
1414
from mcp.client.auth.exceptions import OAuthFlowError
1515
from mcp.client.auth.utils import (
1616
build_oauth_authorization_server_metadata_discovery_urls,
@@ -72,6 +72,117 @@ def client_metadata():
7272
)
7373

7474

75+
def test_build_authorization_redirect_returns_resumable_oauth_state(client_metadata: OAuthClientMetadata):
76+
"""The authorization step can be built without browser/callback handlers."""
77+
redirect = build_authorization_redirect(
78+
authorization_endpoint="https://auth.example.com/authorize",
79+
client_info=OAuthClientInformationFull(
80+
client_id="test_client",
81+
redirect_uris=[AnyUrl("http://localhost:3030/callback")],
82+
),
83+
client_metadata=client_metadata,
84+
pkce_params=PKCEParameters(code_verifier="v" * 64, code_challenge="c" * 64),
85+
state="stored-state",
86+
resource_url="https://api.example.com/v1/mcp",
87+
)
88+
89+
assert redirect.state == "stored-state"
90+
assert redirect.code_verifier == "v" * 64
91+
92+
parsed = urlparse(redirect.authorization_url)
93+
assert parsed.scheme == "https"
94+
assert parsed.netloc == "auth.example.com"
95+
assert parsed.path == "/authorize"
96+
97+
params = parse_qs(parsed.query)
98+
assert params == {
99+
"response_type": ["code"],
100+
"client_id": ["test_client"],
101+
"redirect_uri": ["http://localhost:3030/callback"],
102+
"state": ["stored-state"],
103+
"code_challenge": ["c" * 64],
104+
"code_challenge_method": ["S256"],
105+
"resource": ["https://api.example.com/v1/mcp"],
106+
"scope": ["read write"],
107+
}
108+
109+
110+
def test_build_authorization_redirect_requires_redirect_uri():
111+
with pytest.raises(OAuthFlowError, match="No redirect URIs provided"):
112+
build_authorization_redirect(
113+
authorization_endpoint="https://auth.example.com/authorize",
114+
client_info=OAuthClientInformationFull(
115+
client_id="test_client",
116+
redirect_uris=[AnyUrl("http://localhost:3030/callback")],
117+
),
118+
client_metadata=OAuthClientMetadata(redirect_uris=None),
119+
)
120+
121+
122+
def test_build_authorization_redirect_requires_client_id():
123+
with pytest.raises(OAuthFlowError, match="No client ID provided"):
124+
build_authorization_redirect(
125+
authorization_endpoint="https://auth.example.com/authorize",
126+
client_info=OAuthClientInformationFull(
127+
client_id=None,
128+
redirect_uris=[AnyUrl("http://localhost:3030/callback")],
129+
),
130+
client_metadata=OAuthClientMetadata(redirect_uris=[AnyUrl("http://localhost:3030/callback")]),
131+
)
132+
133+
134+
def test_build_authorization_redirect_prompts_for_offline_access():
135+
redirect = build_authorization_redirect(
136+
authorization_endpoint="https://auth.example.com/authorize",
137+
client_info=OAuthClientInformationFull(
138+
client_id="test_client",
139+
redirect_uris=[AnyUrl("http://localhost:3030/callback")],
140+
),
141+
client_metadata=OAuthClientMetadata(
142+
redirect_uris=[AnyUrl("http://localhost:3030/callback")],
143+
scope="read offline_access",
144+
),
145+
pkce_params=PKCEParameters(code_verifier="v" * 64, code_challenge="c" * 64),
146+
state="stored-state",
147+
)
148+
149+
assert parse_qs(urlparse(redirect.authorization_url).query)["prompt"] == ["consent"]
150+
151+
152+
def test_build_authorization_redirect_preserves_existing_authorization_endpoint_query():
153+
redirect = build_authorization_redirect(
154+
authorization_endpoint="https://auth.example.com/authorize?audience=api",
155+
client_info=OAuthClientInformationFull(
156+
client_id="test_client",
157+
redirect_uris=[AnyUrl("http://localhost:3030/callback")],
158+
),
159+
client_metadata=OAuthClientMetadata(redirect_uris=[AnyUrl("http://localhost:3030/callback")]),
160+
pkce_params=PKCEParameters(code_verifier="v" * 64, code_challenge="c" * 64),
161+
state="stored-state",
162+
)
163+
164+
parsed = urlparse(redirect.authorization_url)
165+
assert parsed.path == "/authorize"
166+
assert parse_qs(parsed.query)["audience"] == ["api"]
167+
assert parse_qs(parsed.query)["response_type"] == ["code"]
168+
169+
170+
def test_authorization_redirect_repr_hides_code_verifier():
171+
redirect = build_authorization_redirect(
172+
authorization_endpoint="https://auth.example.com/authorize",
173+
client_info=OAuthClientInformationFull(
174+
client_id="test_client",
175+
redirect_uris=[AnyUrl("http://localhost:3030/callback")],
176+
),
177+
client_metadata=OAuthClientMetadata(redirect_uris=[AnyUrl("http://localhost:3030/callback")]),
178+
pkce_params=PKCEParameters(code_verifier="v" * 64, code_challenge="c" * 64),
179+
state="stored-state",
180+
)
181+
182+
assert "code_verifier" not in repr(redirect)
183+
assert "v" * 64 not in repr(redirect)
184+
185+
75186
@pytest.fixture
76187
def valid_tokens():
77188
return OAuthToken(

0 commit comments

Comments
 (0)