fastmcp 2.11.2__py3-none-any.whl → 2.12.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (77) hide show
  1. fastmcp/__init__.py +5 -4
  2. fastmcp/cli/claude.py +22 -18
  3. fastmcp/cli/cli.py +472 -136
  4. fastmcp/cli/install/claude_code.py +37 -40
  5. fastmcp/cli/install/claude_desktop.py +37 -42
  6. fastmcp/cli/install/cursor.py +148 -38
  7. fastmcp/cli/install/mcp_json.py +38 -43
  8. fastmcp/cli/install/shared.py +64 -7
  9. fastmcp/cli/run.py +122 -215
  10. fastmcp/client/auth/oauth.py +69 -13
  11. fastmcp/client/client.py +46 -9
  12. fastmcp/client/logging.py +25 -1
  13. fastmcp/client/oauth_callback.py +91 -91
  14. fastmcp/client/sampling.py +12 -4
  15. fastmcp/client/transports.py +143 -67
  16. fastmcp/experimental/sampling/__init__.py +0 -0
  17. fastmcp/experimental/sampling/handlers/__init__.py +3 -0
  18. fastmcp/experimental/sampling/handlers/base.py +21 -0
  19. fastmcp/experimental/sampling/handlers/openai.py +163 -0
  20. fastmcp/experimental/server/openapi/routing.py +1 -3
  21. fastmcp/experimental/server/openapi/server.py +10 -25
  22. fastmcp/experimental/utilities/openapi/__init__.py +2 -2
  23. fastmcp/experimental/utilities/openapi/formatters.py +34 -0
  24. fastmcp/experimental/utilities/openapi/models.py +5 -2
  25. fastmcp/experimental/utilities/openapi/parser.py +252 -70
  26. fastmcp/experimental/utilities/openapi/schemas.py +135 -106
  27. fastmcp/mcp_config.py +40 -20
  28. fastmcp/prompts/prompt_manager.py +4 -2
  29. fastmcp/resources/resource_manager.py +16 -6
  30. fastmcp/server/auth/__init__.py +11 -1
  31. fastmcp/server/auth/auth.py +19 -2
  32. fastmcp/server/auth/oauth_proxy.py +1047 -0
  33. fastmcp/server/auth/providers/azure.py +270 -0
  34. fastmcp/server/auth/providers/github.py +287 -0
  35. fastmcp/server/auth/providers/google.py +305 -0
  36. fastmcp/server/auth/providers/jwt.py +27 -16
  37. fastmcp/server/auth/providers/workos.py +256 -2
  38. fastmcp/server/auth/redirect_validation.py +65 -0
  39. fastmcp/server/auth/registry.py +1 -1
  40. fastmcp/server/context.py +91 -41
  41. fastmcp/server/dependencies.py +32 -2
  42. fastmcp/server/elicitation.py +60 -1
  43. fastmcp/server/http.py +44 -37
  44. fastmcp/server/middleware/logging.py +66 -28
  45. fastmcp/server/proxy.py +2 -0
  46. fastmcp/server/sampling/handler.py +19 -0
  47. fastmcp/server/server.py +85 -20
  48. fastmcp/settings.py +18 -3
  49. fastmcp/tools/tool.py +23 -10
  50. fastmcp/tools/tool_manager.py +5 -1
  51. fastmcp/tools/tool_transform.py +75 -32
  52. fastmcp/utilities/auth.py +34 -0
  53. fastmcp/utilities/cli.py +148 -15
  54. fastmcp/utilities/components.py +21 -5
  55. fastmcp/utilities/inspect.py +166 -37
  56. fastmcp/utilities/json_schema_type.py +4 -2
  57. fastmcp/utilities/logging.py +4 -1
  58. fastmcp/utilities/mcp_config.py +47 -18
  59. fastmcp/utilities/mcp_server_config/__init__.py +25 -0
  60. fastmcp/utilities/mcp_server_config/v1/__init__.py +0 -0
  61. fastmcp/utilities/mcp_server_config/v1/environments/__init__.py +6 -0
  62. fastmcp/utilities/mcp_server_config/v1/environments/base.py +30 -0
  63. fastmcp/utilities/mcp_server_config/v1/environments/uv.py +306 -0
  64. fastmcp/utilities/mcp_server_config/v1/mcp_server_config.py +446 -0
  65. fastmcp/utilities/mcp_server_config/v1/schema.json +361 -0
  66. fastmcp/utilities/mcp_server_config/v1/sources/__init__.py +0 -0
  67. fastmcp/utilities/mcp_server_config/v1/sources/base.py +30 -0
  68. fastmcp/utilities/mcp_server_config/v1/sources/filesystem.py +216 -0
  69. fastmcp/utilities/openapi.py +4 -4
  70. fastmcp/utilities/tests.py +7 -2
  71. fastmcp/utilities/types.py +15 -2
  72. {fastmcp-2.11.2.dist-info → fastmcp-2.12.0.dist-info}/METADATA +3 -2
  73. fastmcp-2.12.0.dist-info/RECORD +129 -0
  74. fastmcp-2.11.2.dist-info/RECORD +0 -108
  75. {fastmcp-2.11.2.dist-info → fastmcp-2.12.0.dist-info}/WHEEL +0 -0
  76. {fastmcp-2.11.2.dist-info → fastmcp-2.12.0.dist-info}/entry_points.txt +0 -0
  77. {fastmcp-2.11.2.dist-info → fastmcp-2.12.0.dist-info}/licenses/LICENSE +0 -0
@@ -0,0 +1,1047 @@
1
+ """OAuth Proxy Provider for FastMCP.
2
+
3
+ This provider acts as a transparent proxy to an upstream OAuth Authorization Server,
4
+ handling Dynamic Client Registration locally while forwarding all other OAuth flows.
5
+ This enables authentication with upstream providers that don't support DCR or have
6
+ restricted client registration policies.
7
+
8
+ Key features:
9
+ - Proxies authorization and token endpoints to upstream server
10
+ - Implements local Dynamic Client Registration with fixed upstream credentials
11
+ - Validates tokens using upstream JWKS
12
+ - Maintains minimal local state for bookkeeping
13
+ - Enhanced logging with request correlation
14
+
15
+ This implementation is based on the OAuth 2.1 specification and is designed for
16
+ production use with enterprise identity providers.
17
+ """
18
+
19
+ from __future__ import annotations
20
+
21
+ import secrets
22
+ import time
23
+ from typing import TYPE_CHECKING, Any, Final
24
+ from urllib.parse import urlencode
25
+
26
+ import httpx
27
+ from authlib.integrations.httpx_client import AsyncOAuth2Client
28
+ from mcp.server.auth.provider import (
29
+ AccessToken,
30
+ AuthorizationCode,
31
+ AuthorizationParams,
32
+ RefreshToken,
33
+ TokenError,
34
+ )
35
+ from mcp.server.auth.settings import (
36
+ ClientRegistrationOptions,
37
+ RevocationOptions,
38
+ )
39
+ from mcp.shared.auth import OAuthClientInformationFull, OAuthToken
40
+ from pydantic import AnyHttpUrl, AnyUrl, SecretStr
41
+ from starlette.requests import Request
42
+ from starlette.responses import JSONResponse, RedirectResponse
43
+ from starlette.routing import Route
44
+
45
+ from fastmcp.server.auth.auth import OAuthProvider, TokenVerifier
46
+ from fastmcp.server.auth.redirect_validation import validate_redirect_uri
47
+ from fastmcp.utilities.logging import get_logger
48
+
49
+ if TYPE_CHECKING:
50
+ pass
51
+
52
+ logger = get_logger(__name__)
53
+
54
+
55
+ class ProxyDCRClient(OAuthClientInformationFull):
56
+ """Client for DCR proxy with configurable redirect URI validation.
57
+
58
+ This special client class is critical for the OAuth proxy to work correctly
59
+ with Dynamic Client Registration (DCR). Here's why it exists:
60
+
61
+ Problem:
62
+ --------
63
+ When MCP clients use OAuth, they dynamically register with random localhost
64
+ ports (e.g., http://localhost:55454/callback). The OAuth proxy needs to:
65
+ 1. Accept these dynamic redirect URIs from clients based on configured patterns
66
+ 2. Use its own fixed redirect URI with the upstream provider (Google, GitHub, etc.)
67
+ 3. Forward the authorization code back to the client's dynamic URI
68
+
69
+ Solution:
70
+ ---------
71
+ This class validates redirect URIs against configurable patterns,
72
+ while the proxy internally uses its own fixed redirect URI with the upstream
73
+ provider. This allows the flow to work even when clients reconnect with
74
+ different ports or when tokens are cached.
75
+
76
+ Without proper validation, clients could get "Redirect URI not registered" errors
77
+ when trying to authenticate with cached tokens, or security vulnerabilities could
78
+ arise from accepting arbitrary redirect URIs.
79
+ """
80
+
81
+ def __init__(
82
+ self, *args, allowed_redirect_uri_patterns: list[str] | None = None, **kwargs
83
+ ):
84
+ """Initialize with allowed redirect URI patterns.
85
+
86
+ Args:
87
+ allowed_redirect_uri_patterns: List of allowed redirect URI patterns with wildcard support.
88
+ If None, defaults to localhost-only patterns.
89
+ If empty list, allows all redirect URIs.
90
+ """
91
+ super().__init__(*args, **kwargs)
92
+ self._allowed_redirect_uri_patterns = allowed_redirect_uri_patterns
93
+
94
+ def validate_redirect_uri(self, redirect_uri: AnyUrl | None) -> AnyUrl:
95
+ """Validate redirect URI against allowed patterns.
96
+
97
+ Since we're acting as a proxy and clients register dynamically,
98
+ we validate their redirect URIs against configurable patterns.
99
+ This is essential for cached token scenarios where the client may
100
+ reconnect with a different port.
101
+ """
102
+ if redirect_uri is not None:
103
+ # Validate against allowed patterns
104
+ if validate_redirect_uri(redirect_uri, self._allowed_redirect_uri_patterns):
105
+ return redirect_uri
106
+ # Fall back to normal validation if not in allowed patterns
107
+ return super().validate_redirect_uri(redirect_uri)
108
+ # If no redirect_uri provided, use default behavior
109
+ return super().validate_redirect_uri(redirect_uri)
110
+
111
+
112
+ # Default token expiration times
113
+ DEFAULT_ACCESS_TOKEN_EXPIRY_SECONDS: Final[int] = 60 * 60 # 1 hour
114
+ DEFAULT_AUTH_CODE_EXPIRY_SECONDS: Final[int] = 5 * 60 # 5 minutes
115
+
116
+ # HTTP client timeout
117
+ HTTP_TIMEOUT_SECONDS: Final[int] = 30
118
+
119
+
120
+ class OAuthProxy(OAuthProvider):
121
+ """OAuth provider that presents a DCR-compliant interface while proxying to non-DCR IDPs.
122
+
123
+ Purpose
124
+ -------
125
+ MCP clients expect OAuth providers to support Dynamic Client Registration (DCR),
126
+ where clients can register themselves dynamically and receive unique credentials.
127
+ Most enterprise IDPs (Google, GitHub, Azure AD, etc.) don't support DCR and require
128
+ pre-registered OAuth applications with fixed credentials.
129
+
130
+ This proxy bridges that gap by:
131
+ - Presenting a full DCR-compliant OAuth interface to MCP clients
132
+ - Translating DCR registration requests to use pre-configured upstream credentials
133
+ - Proxying all OAuth flows to the upstream IDP with appropriate translations
134
+ - Managing the state and security requirements of both protocols
135
+
136
+ Architecture Overview
137
+ --------------------
138
+ The proxy maintains a single OAuth app registration with the upstream provider
139
+ while allowing unlimited MCP clients to register and authenticate dynamically.
140
+ It implements the complete OAuth 2.1 + DCR specification for clients while
141
+ translating to whatever OAuth variant the upstream provider requires.
142
+
143
+ Key Translation Challenges Solved
144
+ ---------------------------------
145
+ 1. Dynamic Client Registration:
146
+ - MCP clients expect to register dynamically and get unique credentials
147
+ - Upstream IDPs require pre-registered apps with fixed credentials
148
+ - Solution: Accept DCR requests, return shared upstream credentials
149
+
150
+ 2. Dynamic Redirect URIs:
151
+ - MCP clients use random localhost ports that change between sessions
152
+ - Upstream IDPs require fixed, pre-registered redirect URIs
153
+ - Solution: Use proxy's fixed callback URL with upstream, forward to client's dynamic URI
154
+
155
+ 3. Authorization Code Mapping:
156
+ - Upstream returns codes for the proxy's redirect URI
157
+ - Clients expect codes for their own redirect URIs
158
+ - Solution: Exchange upstream code server-side, issue new code to client
159
+
160
+ 4. State Parameter Collision:
161
+ - Both client and proxy need to maintain state through the flow
162
+ - Only one state parameter available in OAuth
163
+ - Solution: Use transaction ID as state with upstream, preserve client's state
164
+
165
+ 5. Token Management:
166
+ - Clients may expect different token formats/claims than upstream provides
167
+ - Need to track tokens for revocation and refresh
168
+ - Solution: Store token relationships, forward upstream tokens transparently
169
+
170
+ OAuth Flow Implementation
171
+ ------------------------
172
+ 1. Client Registration (DCR):
173
+ - Accept any client registration request
174
+ - Store ProxyDCRClient that accepts dynamic redirect URIs
175
+ - Return shared upstream credentials to all clients
176
+
177
+ 2. Authorization:
178
+ - Store transaction mapping client details to proxy flow
179
+ - Redirect to upstream with proxy's fixed redirect URI
180
+ - Use transaction ID as state parameter with upstream
181
+
182
+ 3. Upstream Callback:
183
+ - Exchange upstream authorization code for tokens (server-side)
184
+ - Generate new authorization code bound to client's PKCE challenge
185
+ - Redirect to client's original dynamic redirect URI
186
+
187
+ 4. Token Exchange:
188
+ - Validate client's code and PKCE verifier
189
+ - Return previously obtained upstream tokens
190
+ - Clean up one-time use authorization code
191
+
192
+ 5. Token Refresh:
193
+ - Forward refresh requests to upstream using authlib
194
+ - Handle token rotation if upstream issues new refresh token
195
+ - Update local token mappings
196
+
197
+ State Management
198
+ ---------------
199
+ The proxy maintains minimal but crucial state:
200
+ - _clients: DCR registrations (all use ProxyDCRClient for flexibility)
201
+ - _oauth_transactions: Active authorization flows with client context
202
+ - _client_codes: Authorization codes with PKCE challenges and upstream tokens
203
+ - _access_tokens, _refresh_tokens: Token storage for revocation
204
+ - Token relationship mappings for cleanup and rotation
205
+
206
+ Security Considerations
207
+ ----------------------
208
+ - PKCE enforced end-to-end (client to proxy, proxy to upstream)
209
+ - Authorization codes are single-use with short expiry
210
+ - Transaction IDs are cryptographically random
211
+ - All state is cleaned up after use to prevent replay
212
+ - Token validation delegates to upstream provider
213
+
214
+ Provider Compatibility
215
+ ---------------------
216
+ Works with any OAuth 2.0 provider that supports:
217
+ - Authorization code flow
218
+ - Fixed redirect URI (configured in provider's app settings)
219
+ - Standard token endpoint
220
+
221
+ Handles provider-specific requirements:
222
+ - Google: Ensures minimum scope requirements
223
+ - GitHub: Compatible with OAuth Apps and GitHub Apps
224
+ - Azure AD: Handles tenant-specific endpoints
225
+ - Generic: Works with any spec-compliant provider
226
+ """
227
+
228
+ def __init__(
229
+ self,
230
+ *,
231
+ # Upstream server configuration
232
+ upstream_authorization_endpoint: str,
233
+ upstream_token_endpoint: str,
234
+ upstream_client_id: str,
235
+ upstream_client_secret: str,
236
+ upstream_revocation_endpoint: str | None = None,
237
+ # Token validation
238
+ token_verifier: TokenVerifier,
239
+ # FastMCP server configuration
240
+ base_url: AnyHttpUrl | str,
241
+ redirect_path: str = "/auth/callback",
242
+ issuer_url: AnyHttpUrl | str | None = None,
243
+ service_documentation_url: AnyHttpUrl | str | None = None,
244
+ resource_server_url: AnyHttpUrl | str | None = None,
245
+ # Client redirect URI validation
246
+ allowed_client_redirect_uris: list[str] | None = None,
247
+ ):
248
+ """Initialize the OAuth proxy provider.
249
+
250
+ Args:
251
+ upstream_authorization_endpoint: URL of upstream authorization endpoint
252
+ upstream_token_endpoint: URL of upstream token endpoint
253
+ upstream_client_id: Client ID registered with upstream server
254
+ upstream_client_secret: Client secret for upstream server
255
+ upstream_revocation_endpoint: Optional upstream revocation endpoint
256
+ token_verifier: Token verifier for validating access tokens
257
+ base_url: Public URL of the server that exposes this FastMCP server; redirect path is
258
+ relative to this URL
259
+ redirect_path: Redirect path configured in upstream OAuth app (defaults to "/auth/callback")
260
+ issuer_url: Issuer URL for OAuth metadata (defaults to base_url)
261
+ service_documentation_url: Optional service documentation URL
262
+ resource_server_url: Path of the FastMCP server. If None, FastMCP will
263
+ attempt to overwrite this with the correct path to the server
264
+ e.g. {base_url}/mcp
265
+ allowed_client_redirect_uris: List of allowed redirect URI patterns for MCP clients.
266
+ Patterns support wildcards (e.g., "http://localhost:*", "https://*.example.com/*").
267
+ If None (default), only localhost redirect URIs are allowed.
268
+ If empty list, all redirect URIs are allowed (not recommended for production).
269
+ These are for MCP clients performing loopback redirects, NOT for the upstream OAuth app.
270
+ """
271
+ # Always enable DCR since we implement it locally for MCP clients
272
+ client_registration_options = ClientRegistrationOptions(enabled=True)
273
+
274
+ # Enable revocation only if upstream endpoint provided
275
+ revocation_options = (
276
+ RevocationOptions(enabled=True) if upstream_revocation_endpoint else None
277
+ )
278
+
279
+ super().__init__(
280
+ base_url=base_url,
281
+ issuer_url=issuer_url,
282
+ service_documentation_url=service_documentation_url,
283
+ client_registration_options=client_registration_options,
284
+ revocation_options=revocation_options,
285
+ required_scopes=token_verifier.required_scopes,
286
+ resource_server_url=resource_server_url,
287
+ )
288
+
289
+ # Store upstream configuration
290
+ self._upstream_authorization_endpoint = upstream_authorization_endpoint
291
+ self._upstream_token_endpoint = upstream_token_endpoint
292
+ self._upstream_client_id = upstream_client_id
293
+ self._upstream_client_secret = SecretStr(upstream_client_secret)
294
+ self._upstream_revocation_endpoint = upstream_revocation_endpoint
295
+ self._default_scope_str = " ".join(self.required_scopes or [])
296
+
297
+ # Store redirect configuration
298
+ self._redirect_path = (
299
+ redirect_path if redirect_path.startswith("/") else f"/{redirect_path}"
300
+ )
301
+ self._allowed_client_redirect_uris = allowed_client_redirect_uris
302
+
303
+ # Local state for DCR and token bookkeeping
304
+ self._clients: dict[str, OAuthClientInformationFull] = {}
305
+ self._access_tokens: dict[str, AccessToken] = {}
306
+ self._refresh_tokens: dict[str, RefreshToken] = {}
307
+
308
+ # Token relation mappings for cleanup
309
+ self._access_to_refresh: dict[str, str] = {}
310
+ self._refresh_to_access: dict[str, str] = {}
311
+
312
+ # OAuth transaction storage for IdP callback forwarding
313
+ self._oauth_transactions: dict[
314
+ str, dict[str, Any]
315
+ ] = {} # txn_id -> transaction_data
316
+ self._client_codes: dict[str, dict[str, Any]] = {} # client_code -> code_data
317
+
318
+ # Use the provided token validator
319
+ self._token_validator = token_verifier
320
+
321
+ logger.debug(
322
+ "Initialized OAuth proxy provider with upstream server %s",
323
+ self._upstream_authorization_endpoint,
324
+ )
325
+
326
+ # -------------------------------------------------------------------------
327
+ # Client Registration (Local Implementation)
328
+ # -------------------------------------------------------------------------
329
+
330
+ async def get_client(self, client_id: str) -> OAuthClientInformationFull | None:
331
+ """Get client information by ID.
332
+
333
+ For unregistered clients, returns a ProxyDCRClient that accepts
334
+ any localhost redirect URI for DCR clients.
335
+
336
+ Even registered clients use ProxyDCRClient to ensure they can
337
+ authenticate with different dynamic ports on reconnection. This
338
+ handles the case where a client with cached tokens reconnects
339
+ on a different port.
340
+ """
341
+ client = self._clients.get(client_id)
342
+
343
+ if client is None:
344
+ # For unregistered DCR clients, create a permissive client
345
+ # that will accept any localhost redirect URI
346
+ # We need at least one URI for Pydantic validation, but our custom
347
+ # validate_redirect_uri will accept any localhost URI
348
+ client = ProxyDCRClient(
349
+ client_id=client_id,
350
+ client_secret=None,
351
+ redirect_uris=[
352
+ AnyUrl("http://localhost")
353
+ ], # Placeholder, validation uses allowed_patterns
354
+ grant_types=["authorization_code", "refresh_token"],
355
+ scope=self._default_scope_str,
356
+ token_endpoint_auth_method="none",
357
+ allowed_redirect_uri_patterns=self._allowed_client_redirect_uris,
358
+ )
359
+ logger.debug("Created ProxyDCRClient for unregistered client %s", client_id)
360
+
361
+ return client
362
+
363
+ async def register_client(self, client_info: OAuthClientInformationFull) -> None:
364
+ """Register a client locally using fixed upstream credentials.
365
+
366
+ This implementation always uses the upstream client_id and client_secret
367
+ regardless of what the client requests. It modifies the client_info object
368
+ in place since the MCP framework ignores return values.
369
+
370
+ This ensures all clients use the same credentials that are registered
371
+ with the upstream server.
372
+
373
+ Implementation Detail:
374
+ We store a ProxyDCRClient (not the original client_info) to ensure
375
+ the client can reconnect with different dynamic redirect URIs. This is
376
+ essential for cached token scenarios where the client port changes.
377
+
378
+ The flow:
379
+ 1. Client provides its desired redirect URIs (dynamic localhost ports)
380
+ 2. We create a ProxyDCRClient that will accept ANY localhost URI
381
+ 3. We store this flexible client for future authentications
382
+ 4. When client reconnects with a different port, ProxyDCRClient accepts it
383
+ """
384
+ # Always use the upstream credentials
385
+ upstream_id = self._upstream_client_id
386
+ upstream_secret = self._upstream_client_secret.get_secret_value()
387
+
388
+ # Create a ProxyDCRClient with configured redirect URI validation
389
+ proxy_client = ProxyDCRClient(
390
+ client_id=upstream_id,
391
+ client_secret=upstream_secret,
392
+ redirect_uris=client_info.redirect_uris or [AnyUrl("http://localhost")],
393
+ grant_types=client_info.grant_types
394
+ or ["authorization_code", "refresh_token"],
395
+ scope=self._default_scope_str,
396
+ token_endpoint_auth_method="none",
397
+ allowed_redirect_uri_patterns=self._allowed_client_redirect_uris,
398
+ )
399
+
400
+ # Store the ProxyDCRClient using the upstream ID
401
+ self._clients[upstream_id] = proxy_client
402
+
403
+ # Log redirect URIs to help users discover what patterns they might need
404
+ if client_info.redirect_uris:
405
+ for uri in client_info.redirect_uris:
406
+ logger.debug(
407
+ "Client registered with redirect_uri: %s - if restricting redirect URIs, "
408
+ "ensure this pattern is allowed in allowed_client_redirect_uris",
409
+ uri,
410
+ )
411
+
412
+ logger.debug(
413
+ "Registered client %s with %d redirect URIs",
414
+ upstream_id,
415
+ len(proxy_client.redirect_uris),
416
+ )
417
+
418
+ # -------------------------------------------------------------------------
419
+ # Authorization Flow (Proxy to Upstream)
420
+ # -------------------------------------------------------------------------
421
+
422
+ async def authorize(
423
+ self,
424
+ client: OAuthClientInformationFull,
425
+ params: AuthorizationParams,
426
+ ) -> str:
427
+ """Start OAuth transaction and redirect to upstream IdP.
428
+
429
+ This implements the DCR-compliant proxy pattern:
430
+ 1. Store transaction with client details and PKCE challenge
431
+ 2. Use transaction ID as state for IdP
432
+ 3. Redirect to IdP with our fixed callback URL
433
+ """
434
+ # Generate transaction ID for this authorization request
435
+ txn_id = secrets.token_urlsafe(32)
436
+
437
+ # Store transaction data for IdP callback processing
438
+ self._oauth_transactions[txn_id] = {
439
+ "client_id": client.client_id,
440
+ "client_redirect_uri": str(params.redirect_uri),
441
+ "client_state": params.state,
442
+ "code_challenge": params.code_challenge,
443
+ "code_challenge_method": getattr(params, "code_challenge_method", "S256"),
444
+ "scopes": params.scopes or [],
445
+ "created_at": time.time(),
446
+ }
447
+
448
+ # Build query parameters for upstream IdP authorization request
449
+ # Use our fixed IdP callback and transaction ID as state
450
+ query_params: dict[str, Any] = {
451
+ "response_type": "code",
452
+ "client_id": self._upstream_client_id,
453
+ "redirect_uri": f"{str(self.base_url).rstrip('/')}{self._redirect_path}",
454
+ "state": txn_id, # Use txn_id as IdP state
455
+ }
456
+
457
+ # Add scopes - use client scopes or fallback to required scopes
458
+ scopes_to_use = params.scopes or self.required_scopes or []
459
+
460
+ if scopes_to_use:
461
+ query_params["scope"] = " ".join(scopes_to_use)
462
+
463
+ # Build the upstream authorization URL
464
+ separator = "&" if "?" in self._upstream_authorization_endpoint else "?"
465
+ upstream_url = f"{self._upstream_authorization_endpoint}{separator}{urlencode(query_params)}"
466
+
467
+ logger.debug(
468
+ "Starting OAuth transaction %s for client %s, redirecting to IdP",
469
+ txn_id,
470
+ client.client_id,
471
+ )
472
+ return upstream_url
473
+
474
+ # -------------------------------------------------------------------------
475
+ # Authorization Code Handling
476
+ # -------------------------------------------------------------------------
477
+
478
+ async def load_authorization_code(
479
+ self,
480
+ client: OAuthClientInformationFull,
481
+ authorization_code: str,
482
+ ) -> AuthorizationCode | None:
483
+ """Load authorization code for validation.
484
+
485
+ Look up our client code and return authorization code object
486
+ with PKCE challenge for validation.
487
+ """
488
+ # Look up client code data
489
+ code_data = self._client_codes.get(authorization_code)
490
+ if not code_data:
491
+ logger.debug("Authorization code not found: %s", authorization_code)
492
+ return None
493
+
494
+ # Check if code expired
495
+ if time.time() > code_data["expires_at"]:
496
+ logger.debug("Authorization code expired: %s", authorization_code)
497
+ self._client_codes.pop(authorization_code, None)
498
+ return None
499
+
500
+ # Verify client ID matches
501
+ if code_data["client_id"] != client.client_id:
502
+ logger.debug(
503
+ "Authorization code client ID mismatch: %s vs %s",
504
+ code_data["client_id"],
505
+ client.client_id,
506
+ )
507
+ return None
508
+
509
+ # Create authorization code object with PKCE challenge
510
+ return AuthorizationCode(
511
+ code=authorization_code,
512
+ client_id=client.client_id,
513
+ redirect_uri=code_data["redirect_uri"],
514
+ redirect_uri_provided_explicitly=True,
515
+ scopes=code_data["scopes"],
516
+ expires_at=code_data["expires_at"],
517
+ code_challenge=code_data.get("code_challenge", ""),
518
+ )
519
+
520
+ async def exchange_authorization_code(
521
+ self,
522
+ client: OAuthClientInformationFull,
523
+ authorization_code: AuthorizationCode,
524
+ ) -> OAuthToken:
525
+ """Exchange authorization code for stored IdP tokens.
526
+
527
+ For the DCR-compliant proxy flow, we return the IdP tokens that were obtained
528
+ during the IdP callback exchange. PKCE validation is handled by the MCP framework.
529
+ """
530
+ # Look up stored code data
531
+ code_data = self._client_codes.get(authorization_code.code)
532
+ if not code_data:
533
+ logger.error(
534
+ "Authorization code not found in client codes: %s",
535
+ authorization_code.code,
536
+ )
537
+ raise TokenError("invalid_grant", "Authorization code not found")
538
+
539
+ # Get stored IdP tokens
540
+ idp_tokens = code_data["idp_tokens"]
541
+
542
+ # Clean up client code (one-time use)
543
+ self._client_codes.pop(authorization_code.code, None)
544
+
545
+ # Extract token information for local tracking
546
+ access_token_value = idp_tokens["access_token"]
547
+ refresh_token_value = idp_tokens.get("refresh_token")
548
+ expires_in = int(
549
+ idp_tokens.get("expires_in", DEFAULT_ACCESS_TOKEN_EXPIRY_SECONDS)
550
+ )
551
+ expires_at = int(time.time() + expires_in)
552
+
553
+ # Store access token locally for tracking
554
+ access_token = AccessToken(
555
+ token=access_token_value,
556
+ client_id=client.client_id,
557
+ scopes=authorization_code.scopes,
558
+ expires_at=expires_at,
559
+ )
560
+ self._access_tokens[access_token_value] = access_token
561
+
562
+ # Store refresh token if provided
563
+ if refresh_token_value:
564
+ refresh_token = RefreshToken(
565
+ token=refresh_token_value,
566
+ client_id=client.client_id,
567
+ scopes=authorization_code.scopes,
568
+ expires_at=None, # Refresh tokens typically don't expire
569
+ )
570
+ self._refresh_tokens[refresh_token_value] = refresh_token
571
+
572
+ # Maintain token relationships for cleanup
573
+ self._access_to_refresh[access_token_value] = refresh_token_value
574
+ self._refresh_to_access[refresh_token_value] = access_token_value
575
+
576
+ logger.debug(
577
+ "Successfully exchanged client code for stored IdP tokens (client: %s)",
578
+ client.client_id,
579
+ )
580
+
581
+ return OAuthToken(**idp_tokens) # type: ignore[arg-type]
582
+
583
+ # -------------------------------------------------------------------------
584
+ # Refresh Token Flow
585
+ # -------------------------------------------------------------------------
586
+
587
+ async def load_refresh_token(
588
+ self,
589
+ client: OAuthClientInformationFull,
590
+ refresh_token: str,
591
+ ) -> RefreshToken | None:
592
+ """Load refresh token from local storage."""
593
+ return self._refresh_tokens.get(refresh_token)
594
+
595
+ async def exchange_refresh_token(
596
+ self,
597
+ client: OAuthClientInformationFull,
598
+ refresh_token: RefreshToken,
599
+ scopes: list[str],
600
+ ) -> OAuthToken:
601
+ """Exchange refresh token for new access token using authlib."""
602
+
603
+ # Use authlib's AsyncOAuth2Client for refresh token exchange
604
+ oauth_client = AsyncOAuth2Client(
605
+ client_id=self._upstream_client_id,
606
+ client_secret=self._upstream_client_secret.get_secret_value(),
607
+ timeout=HTTP_TIMEOUT_SECONDS,
608
+ )
609
+
610
+ try:
611
+ logger.debug("Using authlib to refresh token from upstream")
612
+
613
+ # Let authlib handle the refresh token exchange
614
+ token_response: dict[str, Any] = await oauth_client.refresh_token( # type: ignore[misc]
615
+ url=self._upstream_token_endpoint,
616
+ refresh_token=refresh_token.token,
617
+ scope=" ".join(scopes) if scopes else None,
618
+ )
619
+
620
+ logger.debug(
621
+ "Successfully refreshed access token via authlib (client: %s)",
622
+ client.client_id,
623
+ )
624
+
625
+ except Exception as e:
626
+ logger.error("Authlib refresh token exchange failed: %s", e)
627
+ raise TokenError(
628
+ "invalid_grant", f"Upstream refresh token exchange failed: {e}"
629
+ ) from e
630
+
631
+ # Update local token storage
632
+ new_access_token = token_response["access_token"]
633
+ expires_in = int(
634
+ token_response.get("expires_in", DEFAULT_ACCESS_TOKEN_EXPIRY_SECONDS)
635
+ )
636
+
637
+ self._access_tokens[new_access_token] = AccessToken(
638
+ token=new_access_token,
639
+ client_id=client.client_id,
640
+ scopes=scopes,
641
+ expires_at=int(time.time() + expires_in),
642
+ )
643
+
644
+ # Handle refresh token rotation if new one provided
645
+ if "refresh_token" in token_response:
646
+ new_refresh_token = token_response["refresh_token"]
647
+ if new_refresh_token != refresh_token.token:
648
+ # Remove old refresh token
649
+ self._refresh_tokens.pop(refresh_token.token, None)
650
+ old_access = self._refresh_to_access.pop(refresh_token.token, None)
651
+ if old_access:
652
+ self._access_to_refresh.pop(old_access, None)
653
+
654
+ # Store new refresh token
655
+ self._refresh_tokens[new_refresh_token] = RefreshToken(
656
+ token=new_refresh_token,
657
+ client_id=client.client_id,
658
+ scopes=scopes,
659
+ expires_at=None,
660
+ )
661
+ self._access_to_refresh[new_access_token] = new_refresh_token
662
+ self._refresh_to_access[new_refresh_token] = new_access_token
663
+
664
+ return OAuthToken(**token_response) # type: ignore[arg-type]
665
+
666
+ # -------------------------------------------------------------------------
667
+ # Token Validation
668
+ # -------------------------------------------------------------------------
669
+
670
+ async def load_access_token(self, token: str) -> AccessToken | None:
671
+ """Validate access token using upstream JWKS.
672
+
673
+ Delegates to the JWT verifier which handles signature validation,
674
+ expiration checking, and claims validation using the upstream JWKS.
675
+ """
676
+ result = await self._token_validator.verify_token(token)
677
+ if result:
678
+ logger.debug("Token validated successfully")
679
+ else:
680
+ logger.debug("Token validation failed")
681
+ return result
682
+
683
+ # -------------------------------------------------------------------------
684
+ # Token Revocation
685
+ # -------------------------------------------------------------------------
686
+
687
+ async def revoke_token(self, token: AccessToken | RefreshToken) -> None:
688
+ """Revoke token locally and with upstream server if supported.
689
+
690
+ Removes tokens from local storage and attempts to revoke them with
691
+ the upstream server if a revocation endpoint is configured.
692
+ """
693
+ # Clean up local token storage
694
+ if isinstance(token, AccessToken):
695
+ self._access_tokens.pop(token.token, None)
696
+ # Also remove associated refresh token
697
+ paired_refresh = self._access_to_refresh.pop(token.token, None)
698
+ if paired_refresh:
699
+ self._refresh_tokens.pop(paired_refresh, None)
700
+ self._refresh_to_access.pop(paired_refresh, None)
701
+ else: # RefreshToken
702
+ self._refresh_tokens.pop(token.token, None)
703
+ # Also remove associated access token
704
+ paired_access = self._refresh_to_access.pop(token.token, None)
705
+ if paired_access:
706
+ self._access_tokens.pop(paired_access, None)
707
+ self._access_to_refresh.pop(paired_access, None)
708
+
709
+ # Attempt upstream revocation if endpoint is configured
710
+ if self._upstream_revocation_endpoint:
711
+ try:
712
+ async with httpx.AsyncClient(
713
+ timeout=HTTP_TIMEOUT_SECONDS
714
+ ) as http_client:
715
+ await http_client.post(
716
+ self._upstream_revocation_endpoint,
717
+ data={"token": token.token},
718
+ auth=(
719
+ self._upstream_client_id,
720
+ self._upstream_client_secret.get_secret_value(),
721
+ ),
722
+ )
723
+ logger.debug("Successfully revoked token with upstream server")
724
+ except Exception as e:
725
+ logger.warning("Failed to revoke token with upstream server: %s", e)
726
+ else:
727
+ logger.debug("No upstream revocation endpoint configured")
728
+
729
+ logger.debug("Token revoked successfully")
730
+
731
+ # -------------------------------------------------------------------------
732
+ # Custom Route Handling
733
+ # -------------------------------------------------------------------------
734
+
735
+ async def _handle_proxy_token_request(self, request: Request) -> JSONResponse:
736
+ """Custom token endpoint using authlib for upstream requests.
737
+
738
+ This handler uses authlib's OAuth2Client to forward token requests to the
739
+ upstream OAuth server, automatically handling response format differences.
740
+ """
741
+ try:
742
+ # Parse the incoming request form data
743
+ form_data = await request.form()
744
+
745
+ # Log the incoming request (with sensitive data redacted)
746
+ redacted_form = {
747
+ k: (
748
+ str(v)[:8] + "..."
749
+ if k in {"code", "code_verifier", "client_secret", "refresh_token"}
750
+ and v
751
+ else str(v)
752
+ )
753
+ for k, v in form_data.items()
754
+ }
755
+ logger.debug("Proxy token request form data: %s", redacted_form)
756
+
757
+ # Create authlib OAuth2 client
758
+ oauth_client = AsyncOAuth2Client(
759
+ client_id=self._upstream_client_id,
760
+ client_secret=self._upstream_client_secret.get_secret_value(),
761
+ timeout=HTTP_TIMEOUT_SECONDS,
762
+ )
763
+
764
+ grant_type = str(form_data.get("grant_type", ""))
765
+
766
+ if grant_type == "authorization_code":
767
+ # Authorization code grant
768
+ try:
769
+ token_data: dict[str, Any] = await oauth_client.fetch_token( # type: ignore[misc]
770
+ url=self._upstream_token_endpoint,
771
+ code=str(form_data.get("code", "")),
772
+ redirect_uri=str(form_data.get("redirect_uri", "")),
773
+ code_verifier=str(form_data.get("code_verifier"))
774
+ if "code_verifier" in form_data
775
+ else None,
776
+ )
777
+
778
+ # Store tokens locally for tracking
779
+ if "access_token" in token_data:
780
+ self._store_tokens_from_response(token_data)
781
+
782
+ logger.debug(
783
+ "Successfully proxied authorization code exchange via authlib"
784
+ )
785
+
786
+ except Exception as e:
787
+ logger.error("Authlib authorization code exchange failed: %s", e)
788
+ return JSONResponse(
789
+ content={
790
+ "error": "invalid_grant",
791
+ "error_description": f"Authorization code exchange failed: {e}",
792
+ },
793
+ status_code=400,
794
+ )
795
+
796
+ elif grant_type == "refresh_token":
797
+ # Refresh token grant
798
+ try:
799
+ token_data: dict[str, Any] = await oauth_client.refresh_token( # type: ignore[misc]
800
+ url=self._upstream_token_endpoint,
801
+ refresh_token=str(form_data.get("refresh_token", "")),
802
+ scope=str(form_data.get("scope"))
803
+ if "scope" in form_data
804
+ else None,
805
+ )
806
+
807
+ logger.debug(
808
+ "Successfully proxied refresh token exchange via authlib"
809
+ )
810
+
811
+ except Exception as e:
812
+ logger.error("Authlib refresh token exchange failed: %s", e)
813
+ return JSONResponse(
814
+ content={
815
+ "error": "invalid_grant",
816
+ "error_description": f"Refresh token exchange failed: {e}",
817
+ },
818
+ status_code=400,
819
+ )
820
+ else:
821
+ # Unsupported grant type
822
+ logger.error("Unsupported grant type: %s", grant_type)
823
+ return JSONResponse(
824
+ content={
825
+ "error": "unsupported_grant_type",
826
+ "error_description": f"Grant type '{grant_type}' not supported by proxy",
827
+ },
828
+ status_code=400,
829
+ )
830
+
831
+ return JSONResponse(content=token_data)
832
+
833
+ except Exception as e:
834
+ logger.error("Error in proxy token handler: %s", e, exc_info=True)
835
+ return JSONResponse(
836
+ content={
837
+ "error": "server_error",
838
+ "error_description": "Internal server error",
839
+ },
840
+ status_code=500,
841
+ )
842
+
843
+ def _store_tokens_from_response(self, token_data: dict[str, Any]) -> None:
844
+ """Store tokens from upstream response for local tracking."""
845
+ try:
846
+ access_token_value = token_data.get("access_token")
847
+ refresh_token_value = token_data.get("refresh_token")
848
+ expires_in = int(
849
+ token_data.get("expires_in", DEFAULT_ACCESS_TOKEN_EXPIRY_SECONDS)
850
+ )
851
+ expires_at = int(time.time() + expires_in)
852
+
853
+ if access_token_value:
854
+ access_token = AccessToken(
855
+ token=access_token_value,
856
+ client_id=self._upstream_client_id,
857
+ scopes=[], # Will be determined by token validation
858
+ expires_at=expires_at,
859
+ )
860
+ self._access_tokens[access_token_value] = access_token
861
+
862
+ if refresh_token_value:
863
+ refresh_token = RefreshToken(
864
+ token=refresh_token_value,
865
+ client_id=self._upstream_client_id,
866
+ scopes=[],
867
+ expires_at=None,
868
+ )
869
+ self._refresh_tokens[refresh_token_value] = refresh_token
870
+
871
+ # Maintain token relationships
872
+ self._access_to_refresh[access_token_value] = refresh_token_value
873
+ self._refresh_to_access[refresh_token_value] = access_token_value
874
+
875
+ logger.debug("Stored tokens from upstream response for tracking")
876
+
877
+ except Exception as e:
878
+ logger.warning("Failed to store tokens from upstream response: %s", e)
879
+
880
+ def get_routes(self) -> list[Route]:
881
+ """Get OAuth routes with custom proxy token handler.
882
+
883
+ This method creates standard OAuth routes and replaces the token endpoint
884
+ with our proxy handler that forwards requests to the upstream OAuth server.
885
+ """
886
+ # Get standard OAuth routes from parent class
887
+ routes = super().get_routes()
888
+ custom_routes = []
889
+ token_route_found = False
890
+
891
+ logger.debug(
892
+ f"get_routes called - configuring OAuth routes in {len(routes)} routes"
893
+ )
894
+
895
+ for i, route in enumerate(routes):
896
+ logger.debug(
897
+ f"Route {i}: {route} - path: {getattr(route, 'path', 'N/A')}, methods: {getattr(route, 'methods', 'N/A')}"
898
+ )
899
+
900
+ # Keep all standard OAuth routes unchanged - our DCR-compliant flow handles everything
901
+ custom_routes.append(route)
902
+
903
+ if (
904
+ isinstance(route, Route)
905
+ and route.path == "/token"
906
+ and route.methods is not None
907
+ and "POST" in route.methods
908
+ ):
909
+ token_route_found = True
910
+
911
+ # Add OAuth callback endpoint for forwarding to client callbacks
912
+ custom_routes.append(
913
+ Route(
914
+ path=self._redirect_path,
915
+ endpoint=self._handle_idp_callback,
916
+ methods=["GET"],
917
+ )
918
+ )
919
+
920
+ logger.debug(
921
+ f"✅ OAuth routes configured: token_endpoint={token_route_found}, total routes={len(custom_routes)} (includes OAuth callback)"
922
+ )
923
+ return custom_routes
924
+
925
+ # -------------------------------------------------------------------------
926
+ # IdP Callback Forwarding
927
+ # -------------------------------------------------------------------------
928
+
929
+ async def _handle_idp_callback(self, request: Request) -> RedirectResponse:
930
+ """Handle callback from upstream IdP and forward to client.
931
+
932
+ This implements the DCR-compliant callback forwarding:
933
+ 1. Receive IdP callback with code and txn_id as state
934
+ 2. Exchange IdP code for tokens (server-side)
935
+ 3. Generate our own client code bound to PKCE challenge
936
+ 4. Redirect to client's callback with client code and original state
937
+ """
938
+ try:
939
+ idp_code = request.query_params.get("code")
940
+ txn_id = request.query_params.get("state")
941
+ error = request.query_params.get("error")
942
+
943
+ if error:
944
+ logger.error(
945
+ "IdP callback error: %s - %s",
946
+ error,
947
+ request.query_params.get("error_description"),
948
+ )
949
+ # TODO: Forward error to client callback
950
+ return RedirectResponse(
951
+ url=f"data:text/html,<h1>OAuth Error</h1><p>{error}: {request.query_params.get('error_description', 'Unknown error')}</p>",
952
+ status_code=302,
953
+ )
954
+
955
+ if not idp_code or not txn_id:
956
+ logger.error("IdP callback missing code or transaction ID")
957
+ return RedirectResponse(
958
+ url="data:text/html,<h1>OAuth Error</h1><p>Missing authorization code or transaction ID</p>",
959
+ status_code=302,
960
+ )
961
+
962
+ # Look up transaction data
963
+ transaction = self._oauth_transactions.get(txn_id)
964
+ if not transaction:
965
+ logger.error("IdP callback with invalid transaction ID: %s", txn_id)
966
+ return RedirectResponse(
967
+ url="data:text/html,<h1>OAuth Error</h1><p>Invalid or expired transaction</p>",
968
+ status_code=302,
969
+ )
970
+
971
+ # Exchange IdP code for tokens (server-side)
972
+ oauth_client = AsyncOAuth2Client(
973
+ client_id=self._upstream_client_id,
974
+ client_secret=self._upstream_client_secret.get_secret_value(),
975
+ timeout=HTTP_TIMEOUT_SECONDS,
976
+ )
977
+
978
+ try:
979
+ idp_redirect_uri = (
980
+ f"{str(self.base_url).rstrip('/')}{self._redirect_path}"
981
+ )
982
+ logger.debug(
983
+ f"Exchanging IdP code for tokens with redirect_uri: {idp_redirect_uri}"
984
+ )
985
+
986
+ idp_tokens: dict[str, Any] = await oauth_client.fetch_token( # type: ignore[misc]
987
+ url=self._upstream_token_endpoint,
988
+ code=idp_code,
989
+ redirect_uri=idp_redirect_uri,
990
+ )
991
+
992
+ logger.debug(
993
+ f"Successfully exchanged IdP code for tokens (transaction: {txn_id})"
994
+ )
995
+
996
+ except Exception as e:
997
+ logger.error("IdP token exchange failed: %s", e)
998
+ # TODO: Forward error to client callback
999
+ return RedirectResponse(
1000
+ url=f"data:text/html,<h1>OAuth Error</h1><p>Token exchange failed: {e}</p>",
1001
+ status_code=302,
1002
+ )
1003
+
1004
+ # Generate our own authorization code for the client
1005
+ client_code = secrets.token_urlsafe(32)
1006
+ code_expires_at = int(time.time() + DEFAULT_AUTH_CODE_EXPIRY_SECONDS)
1007
+
1008
+ # Store client code with PKCE challenge and IdP tokens
1009
+ self._client_codes[client_code] = {
1010
+ "client_id": transaction["client_id"],
1011
+ "redirect_uri": transaction["client_redirect_uri"],
1012
+ "code_challenge": transaction["code_challenge"],
1013
+ "code_challenge_method": transaction["code_challenge_method"],
1014
+ "scopes": transaction["scopes"],
1015
+ "idp_tokens": idp_tokens,
1016
+ "expires_at": code_expires_at,
1017
+ "created_at": time.time(),
1018
+ }
1019
+
1020
+ # Clean up transaction
1021
+ self._oauth_transactions.pop(txn_id, None)
1022
+
1023
+ # Build client callback URL with our code and original state
1024
+ client_redirect_uri = transaction["client_redirect_uri"]
1025
+ client_state = transaction["client_state"]
1026
+
1027
+ callback_params = {
1028
+ "code": client_code,
1029
+ "state": client_state,
1030
+ }
1031
+
1032
+ # Add query parameters to client redirect URI
1033
+ separator = "&" if "?" in client_redirect_uri else "?"
1034
+ client_callback_url = (
1035
+ f"{client_redirect_uri}{separator}{urlencode(callback_params)}"
1036
+ )
1037
+
1038
+ logger.debug(f"Forwarding to client callback for transaction {txn_id}")
1039
+
1040
+ return RedirectResponse(url=client_callback_url, status_code=302)
1041
+
1042
+ except Exception as e:
1043
+ logger.error("Error in IdP callback handler: %s", e, exc_info=True)
1044
+ return RedirectResponse(
1045
+ url="data:text/html,<h1>OAuth Error</h1><p>Internal server error during IdP callback</p>",
1046
+ status_code=302,
1047
+ )