 Max Isbey
|
34bcf2d415
|
refactor: extract auth components and streamable HTTP app helpers
- 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>
|
2026-01-16 15:32:33 +01:00 |
|