From 34ada2c7adbbdf0f575c51b8056f439e22cb2323 Mon Sep 17 00:00:00 2001 From: Taylor Wilsdon Date: Sat, 28 Feb 2026 11:50:09 -0400 Subject: [PATCH] better cachce management --- core/server.py | 55 ++++++++++++++++++++++---------------------------- 1 file changed, 24 insertions(+), 31 deletions(-) diff --git a/core/server.py b/core/server.py index 7b68848..87e32ea 100644 --- a/core/server.py +++ b/core/server.py @@ -1,13 +1,13 @@ import hashlib import logging import os -from typing import Callable, List, Optional +from typing import List, Optional from importlib import metadata from fastapi.responses import HTMLResponse, JSONResponse, FileResponse from starlette.applications import Starlette -from starlette.middleware.base import BaseHTTPMiddleware from starlette.requests import Request +from starlette.responses import Response from starlette.middleware import Middleware from fastmcp import FastMCP @@ -46,31 +46,27 @@ def _compute_scope_fingerprint() -> str: return hashlib.sha256(scopes_str.encode()).hexdigest()[:12] -class OAuthMetadataCacheBustMiddleware(BaseHTTPMiddleware): - """Override the upstream 1-hour Cache-Control on OAuth discovery endpoints. +def _wrap_well_known_endpoint(endpoint, etag: str): + """Wrap a well-known metadata endpoint to prevent browser caching. - The MCP SDK sets ``Cache-Control: public, max-age=3600`` on the - ``.well-known`` metadata responses. When the server is restarted with a - different ``--permissions`` or ``--read-only`` configuration, browsers / - MCP clients can serve stale discovery docs that advertise the wrong - scopes, causing the OAuth flow to silently fail. + The MCP SDK hardcodes ``Cache-Control: public, max-age=3600`` on discovery + responses. When the server restarts with different ``--permissions`` or + ``--read-only`` flags, browsers / MCP clients serve stale metadata that + advertises the wrong scopes, causing OAuth to silently fail. - This middleware replaces the cache header with ``no-store`` and adds an - ``ETag`` derived from the current scope set so that intermediary caches - that *do* store the response will still invalidate on config change. + The wrapper overrides the header to ``no-store`` and adds an ``ETag`` + derived from the current scope set so intermediary caches that ignore + ``no-store`` still see a fingerprint change. """ - def __init__(self, app: Starlette, scope_fingerprint: str) -> None: - super().__init__(app) - self._etag = f'"{scope_fingerprint}"' - - async def dispatch(self, request: Request, call_next: Callable): - response = await call_next(request) - if request.url.path.startswith("/.well-known/"): - response.headers["Cache-Control"] = "no-store, must-revalidate" - response.headers["ETag"] = self._etag + async def _no_cache_endpoint(request: Request) -> Response: + response = await endpoint(request) + response.headers["Cache-Control"] = "no-store, must-revalidate" + response.headers["ETag"] = etag return response + return _no_cache_endpoint + # Custom FastMCP that adds secure middleware stack for OAuth 2.1 class SecureFastMCP(FastMCP): @@ -82,13 +78,6 @@ class SecureFastMCP(FastMCP): # Session Management - extracts session info for MCP context app.user_middleware.insert(0, session_middleware) - # Prevent browser caching of OAuth discovery endpoints across config changes - fingerprint = _compute_scope_fingerprint() - app.user_middleware.insert( - 0, - Middleware(OAuthMetadataCacheBustMiddleware, scope_fingerprint=fingerprint), - ) - # Rebuild middleware stack app.middleware_stack = app.build_middleware_stack() logger.info("Added middleware stack: Session Management") @@ -428,14 +417,18 @@ def configure_server_for_http(): "OAuth 2.1 enabled using FastMCP GoogleProvider with protocol-level auth" ) - # Explicitly mount well-known routes from the OAuth provider - # These should be auto-mounted but we ensure they're available + # Mount well-known routes with cache-busting headers. + # The MCP SDK hardcodes Cache-Control: public, max-age=3600 + # on discovery responses which causes stale-scope bugs when + # the server is restarted with a different --permissions config. try: + scope_etag = f'"{_compute_scope_fingerprint()}"' well_known_routes = provider.get_well_known_routes() for route in well_known_routes: logger.info(f"Mounting OAuth well-known route: {route.path}") + wrapped = _wrap_well_known_endpoint(route.endpoint, scope_etag) server.custom_route(route.path, methods=list(route.methods))( - route.endpoint + wrapped ) except Exception as e: logger.warning(f"Could not mount well-known routes: {e}")