-
Notifications
You must be signed in to change notification settings - Fork 3k
refactor: extract auth components and streamable HTTP app helpers #1896
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Open
maxisbey
wants to merge
2
commits into
main
Choose a base branch
from
refactor-auth-components
base: main
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
+471
−158
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
- Add build_auth_components() in mcp.server.auth.components for reusable auth setup (middleware, endpoint wrapper, routes) - Refactor create_streamable_http_app() to take session_manager as first arg with keyword args for app config (removed StreamableHTTPAppConfig) - FastMCP now uses _build_auth_components() helper, reducing duplication between sse_app() and streamable_http_app() - Session manager is now created/owned by caller, passed to app creator - Add unit tests for build_auth_components() - Export AuthComponents, build_auth_components from mcp.server.auth - Export StreamableHTTPSessionManager, create_streamable_http_app from mcp.server Claude-Generated-By: Claude Code (cli/claude-opus-4-5=100%) Claude-Steers: 14 Claude-Permission-Prompts: 13 Claude-Escapes: 0 Claude-Plan: <claude-plan> # Plan: Extract Auth Helper and Make FastMCP a Thin Wrapper ## Summary Create a shared `build_auth_components()` helper in the auth module that both `sse_app()` and `streamable_http_app()` can use. This removes ~60 lines of duplicated auth logic from FastMCP and makes it a much thinner wrapper. ## Files to Modify 1. **`src/mcp/server/auth/routes.py`** - Add `AuthConfig` dataclass and `build_auth_components()` function 2. **`src/mcp/server/fastmcp/server.py`** - Refactor both `sse_app()` and `streamable_http_app()` to use the helper 3. **`src/mcp/server/__init__.py`** - Export new auth helper for low-level users ## Implementation ### Step 1: Add to `src/mcp/server/auth/routes.py` **Add new dataclass and helper function:** ```python from dataclasses import dataclass, field from starlette.middleware import Middleware from starlette.middleware.authentication import AuthenticationMiddleware from starlette.types import ASGIApp @DataClass class AuthConfig: """Configuration for auth components in Starlette apps.""" # Token verification (required) token_verifier: TokenVerifier # Auth settings issuer_url: AnyHttpUrl required_scopes: list[str] = field(default_factory=list) resource_server_url: AnyHttpUrl | None = None # Optional: Full OAuth AS provider for serving auth endpoints auth_server_provider: OAuthAuthorizationServerProvider[Any, Any, Any] | None = None service_documentation_url: AnyHttpUrl | None = None client_registration_options: ClientRegistrationOptions | None = None revocation_options: RevocationOptions | None = None @DataClass class AuthComponents: """Auth components ready to be used in a Starlette app.""" routes: list[Route] middleware: list[Middleware] endpoint_wrapper: Callable[[ASGIApp], ASGIApp] def build_auth_components(config: AuthConfig) -> AuthComponents: """ Build auth routes, middleware, and endpoint wrapper from config. Returns an AuthComponents with: - routes: OAuth AS routes (if provider set) + protected resource metadata - middleware: AuthenticationMiddleware + AuthContextMiddleware - endpoint_wrapper: RequireAuthMiddleware wrapper function """ routes: list[Route] = [] # Build middleware middleware = [ Middleware( AuthenticationMiddleware, backend=BearerAuthBackend(config.token_verifier), ), Middleware(AuthContextMiddleware), ] # Add OAuth AS routes if provider is configured if config.auth_server_provider: routes.extend( create_auth_routes( provider=config.auth_server_provider, issuer_url=config.issuer_url, service_documentation_url=config.service_documentation_url, client_registration_options=config.client_registration_options, revocation_options=config.revocation_options, ) ) # Add protected resource metadata routes if resource_server_url is set if config.resource_server_url: routes.extend( create_protected_resource_routes( resource_url=config.resource_server_url, authorization_servers=[config.issuer_url], scopes_supported=config.required_scopes or None, ) ) # Build endpoint wrapper resource_metadata_url = None if config.resource_server_url: resource_metadata_url = build_resource_metadata_url(config.resource_server_url) def endpoint_wrapper(app: ASGIApp) -> ASGIApp: return RequireAuthMiddleware(app, config.required_scopes, resource_metadata_url) return AuthComponents( routes=routes, middleware=middleware, endpoint_wrapper=endpoint_wrapper, ) ``` ### Step 2: Refactor `FastMCP.streamable_http_app()` Replace the ~50 lines of auth logic with: ```python def streamable_http_app(self) -> Starlette: """Return an instance of the StreamableHTTP server app.""" additional_routes: list[Route | Mount] = [] middleware: list[Middleware] = [] endpoint_wrapper: Callable[[ASGIApp], ASGIApp] | None = None # Build auth components if auth is configured if self.settings.auth and self._token_verifier: from mcp.server.auth.routes import AuthConfig, build_auth_components auth_config = AuthConfig( token_verifier=self._token_verifier, issuer_url=self.settings.auth.issuer_url, required_scopes=self.settings.auth.required_scopes or [], resource_server_url=self.settings.auth.resource_server_url, auth_server_provider=self._auth_server_provider, 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, ) auth_components = build_auth_components(auth_config) additional_routes.extend(auth_components.routes) middleware = auth_components.middleware endpoint_wrapper = auth_components.endpoint_wrapper # Add custom routes last additional_routes.extend(self._custom_starlette_routes) # Create config and call low-level function config = StreamableHTTPAppConfig( mcp_server=self._mcp_server, event_store=self._event_store, retry_interval=self._retry_interval, json_response=self.settings.json_response, stateless=self.settings.stateless_http, security_settings=self.settings.transport_security, endpoint_path=self.settings.streamable_http_path, debug=self.settings.debug, additional_routes=additional_routes, middleware=middleware, endpoint_wrapper=endpoint_wrapper, ) starlette_app, session_manager = create_streamable_http_app(config) self._session_manager = session_manager return starlette_app ``` ### Step 3: Refactor `FastMCP.sse_app()` Similar refactor - replace the auth logic with `build_auth_components()`. The SSE app has slightly different route structure (two endpoints: SSE and messages), so the wrapper is applied to each endpoint individually rather than via config: ```python def sse_app(self, mount_path: str | None = None) -> Starlette: """Return an instance of the SSE server app.""" if mount_path is not None: self.settings.mount_path = mount_path normalized_message_endpoint = self._normalize_path( self.settings.mount_path, self.settings.message_path ) sse = SseServerTransport( normalized_message_endpoint, security_settings=self.settings.transport_security, ) async def handle_sse(scope: Scope, receive: Receive, send: Send): async with sse.connect_sse(scope, receive, send) as streams: await self._mcp_server.run( streams[0], streams[1], self._mcp_server.create_initialization_options(), ) return Response() routes: list[Route | Mount] = [] middleware: list[Middleware] = [] # Build auth components if configured if self.settings.auth and self._token_verifier: from mcp.server.auth.routes import AuthConfig, build_auth_components auth_config = AuthConfig( token_verifier=self._token_verifier, issuer_url=self.settings.auth.issuer_url, required_scopes=self.settings.auth.required_scopes or [], resource_server_url=self.settings.auth.resource_server_url, auth_server_provider=self._auth_server_provider, 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, ) auth_components = build_auth_components(auth_config) routes.extend(auth_components.routes) middleware = auth_components.middleware # SSE has two endpoints that need wrapping routes.append(Route( self.settings.sse_path, endpoint=auth_components.endpoint_wrapper(handle_sse), methods=["GET"], )) routes.append(Mount( self.settings.message_path, app=auth_components.endpoint_wrapper(sse.handle_post_message), )) else: # No auth - add routes directly async def sse_endpoint(request: Request) -> Response: return await handle_sse(request.scope, request.receive, request._send) routes.append(Route(self.settings.sse_path, endpoint=sse_endpoint, methods=["GET"])) routes.append(Mount(self.settings.message_path, app=sse.handle_post_message)) routes.extend(self._custom_starlette_routes) return Starlette(debug=self.settings.debug, routes=routes, middleware=middleware) ``` ### Step 4: Update exports in `src/mcp/server/__init__.py` Add the new auth helper to exports: ```python from .auth.routes import AuthConfig, AuthComponents, build_auth_components ``` ## Verification 1. Run existing tests: ```bash PYTEST_DISABLE_PLUGIN_AUTOLOAD="" uv run --frozen pytest tests/server/fastmcp/ ``` 2. Run auth tests specifically: ```bash PYTEST_DISABLE_PLUGIN_AUTOLOAD="" uv run --frozen pytest tests/server/fastmcp/auth/ ``` 3. Run type checking: ```bash uv run --frozen pyright ``` 4. Test low-level usage with auth: ```python from mcp.server import Server, StreamableHTTPAppConfig, create_streamable_http_app from mcp.server.auth.routes import AuthConfig, build_auth_components server = Server('test') auth = build_auth_components(AuthConfig(...)) config = StreamableHTTPAppConfig( mcp_server=server, additional_routes=auth.routes, middleware=auth.middleware, endpoint_wrapper=auth.endpoint_wrapper, ) app, manager = create_streamable_http_app(config) ``` ## Result FastMCP's `streamable_http_app()` goes from ~90 lines to ~30 lines, and `sse_app()` similarly shrinks. The auth logic is now: - Reusable by low-level Server users - Testable in isolation - Shared between SSE and StreamableHTTP transports </claude-plan>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Summary
Refactors the auth and streamable HTTP setup to be more composable and reusable by low-level Server users.
Changes
New
build_auth_components()functionmcp.server.auth.componentsAuthComponentswith:middleware- Authentication middlewareendpoint_wrapper- Function to protect endpointsroutes- OAuth AS and/or protected resource metadata routesRefactored
create_streamable_http_app()session_manageras first argument (caller owns it)StreamableHTTPAppConfigdataclass)Starlette, not a tupleFastMCP improvements
_build_auth_components()private helpersse_app()andstreamable_http_app()use shared auth setupNew exports
mcp.server.auth:AuthComponents,build_auth_componentsmcp.server:StreamableHTTPSessionManager,create_streamable_http_appExample: Low-level usage with auth
Testing
build_auth_components()intests/server/auth/test_components.py