Fixups while integrating new auth capabilities

Author: praboud-ant

src/mcp/server/auth/provider.py

@@ -1,4 +1,4 @@
-from typing import Protocol
+from typing import Generic, Protocol, TypeVar
 from urllib.parse import parse_qs, urlencode, urlparse, urlunparse
 
 from pydantic import AnyHttpUrl, BaseModel
@@ -62,7 +62,16 @@ async def register_client(self, client_info: OAuthClientInformationFull) -> None
         ...
 
 
-class OAuthServerProvider(Protocol):
+# NOTE: FastMCP doesn't render any of these types in the user response, so it's
+# OK to add fields to subclasses which should not be exposed externally.
+AuthorizationCodeT = TypeVar("AuthorizationCodeT", bound=AuthorizationCode)
+RefreshTokenT = TypeVar("RefreshTokenT", bound=RefreshToken)
+AuthInfoT = TypeVar("AuthInfoT", bound=AuthInfo)
+
+
+class OAuthServerProvider(
+    Protocol, Generic[AuthorizationCodeT, RefreshTokenT, AuthInfoT]
+):
     @property
     def clients_store(self) -> OAuthRegisteredClientsStore:
         """
@@ -107,7 +116,7 @@ async def authorize(
 
     async def load_authorization_code(
         self, client: OAuthClientInformationFull, authorization_code: str
-    ) -> AuthorizationCode | None:
+    ) -> AuthorizationCodeT | None:
         """
         Loads metadata for the authorization code challenge.
 
@@ -121,7 +130,7 @@ async def load_authorization_code(
         ...
 
     async def exchange_authorization_code(
-        self, client: OAuthClientInformationFull, authorization_code: AuthorizationCode
+        self, client: OAuthClientInformationFull, authorization_code: AuthorizationCodeT
     ) -> OAuthToken:
         """
         Exchanges an authorization code for an access token.
@@ -137,12 +146,12 @@ async def exchange_authorization_code(
 
     async def load_refresh_token(
         self, client: OAuthClientInformationFull, refresh_token: str
-    ) -> RefreshToken | None: ...
+    ) -> RefreshTokenT | None: ...
 
     async def exchange_refresh_token(
         self,
         client: OAuthClientInformationFull,
-        refresh_token: RefreshToken,
+        refresh_token: RefreshTokenT,
         scopes: list[str],
     ) -> OAuthToken:
         """
@@ -158,7 +167,7 @@ async def exchange_refresh_token(
         """
         ...
 
-    async def load_access_token(self, token: str) -> AuthInfo | None:
+    async def load_access_token(self, token: str) -> AuthInfoT | None:
         """
         Verifies an access token and returns information about it.
 
@@ -172,7 +181,7 @@ async def load_access_token(self, token: str) -> AuthInfo | None:
 
     async def revoke_token(
         self,
-        token: AuthInfo | RefreshToken,
+        token: AuthInfoT | RefreshTokenT,
     ) -> None:
         """
         Revokes an access or refresh token.

src/mcp/server/auth/router.py

@@ -2,7 +2,7 @@
 from typing import Callable
 
 from pydantic import AnyHttpUrl
-from starlette.routing import Route, Router
+from starlette.routing import Route
 
 from mcp.server.auth.handlers.authorize import AuthorizationHandler
 from mcp.server.auth.handlers.metadata import MetadataHandler
@@ -57,27 +57,13 @@ def validate_issuer_url(url: AnyHttpUrl):
 REVOCATION_PATH = "/revoke"
 
 
-def create_auth_router(
+def create_auth_routes(
     provider: OAuthServerProvider,
     issuer_url: AnyHttpUrl,
     service_documentation_url: AnyHttpUrl | None = None,
     client_registration_options: ClientRegistrationOptions | None = None,
     revocation_options: RevocationOptions | None = None,
-) -> Router:
-    """
-    Create a Starlette router with standard MCP authorization endpoints.
-
-    Args:
-        provider: OAuth server provider
-        issuer_url: Issuer URL for the authorization server
-        service_documentation_url: Optional URL for service documentation
-        client_registration_options: Options for client registration
-        revocation_options: Options for token revocation
-
-    Returns:
-        Starlette router with authorization endpoints
-    """
-
+) -> list[Route]:
     validate_issuer_url(issuer_url)
 
     client_registration_options = (
@@ -93,32 +79,30 @@ def create_auth_router(
     client_authenticator = ClientAuthenticator(provider.clients_store)
 
     # Create routes
-    auth_router = Router(
-        routes=[
-            Route(
-                "/.well-known/oauth-authorization-server",
-                endpoint=MetadataHandler(metadata).handle,
-                methods=["GET"],
-            ),
-            Route(
-                AUTHORIZATION_PATH,
-                endpoint=AuthorizationHandler(provider).handle,
-                methods=["GET", "POST"],
-            ),
-            Route(
-                TOKEN_PATH,
-                endpoint=TokenHandler(provider, client_authenticator).handle,
-                methods=["POST"],
-            ),
-        ]
-    )
+    routes = [
+        Route(
+            "/.well-known/oauth-authorization-server",
+            endpoint=MetadataHandler(metadata).handle,
+            methods=["GET"],
+        ),
+        Route(
+            AUTHORIZATION_PATH,
+            endpoint=AuthorizationHandler(provider).handle,
+            methods=["GET", "POST"],
+        ),
+        Route(
+            TOKEN_PATH,
+            endpoint=TokenHandler(provider, client_authenticator).handle,
+            methods=["POST"],
+        ),
+    ]
 
     if client_registration_options.enabled:
         registration_handler = RegistrationHandler(
             provider.clients_store,
             client_secret_expiry_seconds=client_registration_options.client_secret_expiry_seconds,
         )
-        auth_router.routes.append(
+        routes.append(
             Route(
                 REGISTRATION_PATH,
                 endpoint=registration_handler.handle,
@@ -128,11 +112,11 @@ def create_auth_router(
 
     if revocation_options.enabled:
         revocation_handler = RevocationHandler(provider, client_authenticator)
-        auth_router.routes.append(
+        routes.append(
             Route(REVOCATION_PATH, endpoint=revocation_handler.handle, methods=["POST"])
         )
 
-    return auth_router
+    return routes
 
 
 def modify_url_path(url: AnyHttpUrl, path_mapper: Callable[[str], str]) -> AnyHttpUrl:

src/mcp/server/fastmcp/server.py

@@ -11,7 +11,7 @@
     asynccontextmanager,
 )
 from itertools import chain
-from typing import Any, Callable, Generic, Literal, Sequence
+from typing import Any, Awaitable, Callable, Generic, Literal, Sequence
 
 import anyio
 import pydantic_core
@@ -24,6 +24,7 @@
 from starlette.authentication import requires
 from starlette.middleware.authentication import AuthenticationMiddleware
 
+from mcp.server.auth.middleware.auth_context import AuthContextMiddleware
 from mcp.server.auth.middleware.bearer_auth import (
     BearerAuthBackend,
     RequireAuthMiddleware,
@@ -151,6 +152,7 @@ def __init__(
             warn_on_duplicate_prompts=self.settings.warn_on_duplicate_prompts
         )
         self._auth_provider = auth_provider
+        self._custom_starlette_routes = []
         self.dependencies = self.settings.dependencies
 
         # Set up MCP protocol handlers
@@ -477,6 +479,33 @@ def decorator(func: AnyFunction) -> AnyFunction:
 
         return decorator
 
+    def custom_route(
+        self,
+        path: str,
+        methods: list[str],
+        name: str | None = None,
+        include_in_schema: bool = True,
+    ):
+        from starlette.requests import Request
+        from starlette.responses import Response
+        from starlette.routing import Route
+
+        def decorator(
+            func: Callable[[Request], Awaitable[Response]],
+        ) -> Callable[[Request], Awaitable[Response]]:
+            self._custom_starlette_routes.append(
+                Route(
+                    path,
+                    endpoint=func,
+                    methods=methods,
+                    name=name,
+                    include_in_schema=include_in_schema,
+                )
+            )
+            return func
+
+        return decorator
+
     async def run_stdio_async(self) -> None:
         """Run the server using stdio transport."""
         async with stdio_server() as (read_stream, write_stream):
@@ -513,31 +542,33 @@ async def handle_sse(request) -> EventSourceResponse:
         routes = []
         middleware = []
         required_scopes = self.settings.auth_required_scopes or []
-        auth_router = None
 
         # Add auth endpoints if auth provider is configured
         if self._auth_provider and self.settings.auth_issuer_url:
-            from mcp.server.auth.router import create_auth_router
+            from mcp.server.auth.router import create_auth_routes
 
-            # Set up bearer auth middleware if auth is required
             middleware = [
+                # extract auth info from request (but do not require it)
                 Middleware(
                     AuthenticationMiddleware,
                     backend=BearerAuthBackend(
                         provider=self._auth_provider,
                     ),
-                )
+                ),
+                # Add the auth context middleware to store
+                # authenticated user in a contextvar
+                Middleware(AuthContextMiddleware),
             ]
-            auth_router = create_auth_router(
-                provider=self._auth_provider,
-                issuer_url=self.settings.auth_issuer_url,
-                service_documentation_url=self.settings.auth_service_documentation_url,
-                client_registration_options=self.settings.auth_client_registration_options,
-                revocation_options=self.settings.auth_revocation_options,
+            routes.extend(
+                create_auth_routes(
+                    provider=self._auth_provider,
+                    issuer_url=self.settings.auth_issuer_url,
+                    service_documentation_url=self.settings.auth_service_documentation_url,
+                    client_registration_options=self.settings.auth_client_registration_options,
+                    revocation_options=self.settings.auth_revocation_options,
+                )
             )
 
-            # Add the auth router as a mount
-
         routes.append(
             Route(
                 "/sse", endpoint=requires(required_scopes)(handle_sse), methods=["GET"]
@@ -549,8 +580,8 @@ async def handle_sse(request) -> EventSourceResponse:
                 app=RequireAuthMiddleware(sse.handle_post_message, required_scopes),
             )
         )
-        if auth_router:
-            routes.append(Mount("/", app=auth_router))
+        # mount these routes last, so they have the lowest route matching precedence
+        routes.extend(self._custom_starlette_routes)
 
         # Create Starlette app with routes and middleware
         return Starlette(

src/mcp/shared/auth.py

@@ -41,7 +41,8 @@ class OAuthClientMetadata(BaseModel):
     )
     # grant_types: this implementation only supports authorization_code & refresh_token
     grant_types: list[Literal["authorization_code", "refresh_token"]] = [
-        "authorization_code"
+        "authorization_code",
+        "refresh_token",
     ]
     # this implementation only supports code; ie: it does not support implicit grants
     response_types: list[Literal["code"]] = ["code"]

tests/server/fastmcp/auth/test_auth_integration.py

@@ -30,7 +30,7 @@
 from mcp.server.auth.router import (
     ClientRegistrationOptions,
     RevocationOptions,
-    create_auth_router,
+    create_auth_routes,
 )
 from mcp.server.fastmcp import FastMCP
 from mcp.shared.auth import (
@@ -222,7 +222,7 @@ def mock_oauth_provider():
 @pytest.fixture
 def auth_app(mock_oauth_provider):
     # Create auth router
-    auth_router = create_auth_router(
+    auth_routes = create_auth_routes(
         mock_oauth_provider,
         AnyHttpUrl("https://auth.example.com"),
         AnyHttpUrl("https://docs.example.com"),
@@ -231,7 +231,7 @@ def auth_app(mock_oauth_provider):
     )
 
     # Create Starlette app
-    app = Starlette(routes=[Mount("/", app=auth_router)])
+    app = Starlette(routes=auth_routes)
 
     return app