Merge branch 'main' into switch-kvstore

pyproject.toml

@@ -16,6 +16,7 @@ dependencies = [
     "pyperclip>=1.9.0",
     "openapi-core>=0.19.5",
     "py-key-value-aio[disk,memory]>=0.2.0",
+    "websockets>=15.0.1",
 ]
 
 requires-python = ">=3.10"
@@ -42,7 +43,6 @@ classifiers = [
 ]
 
 [project.optional-dependencies]
-websockets = ["websockets>=15.0.1"]
 openai = ["openai>=1.102.0"]
 
 [dependency-groups]

src/fastmcp/client/oauth_callback.py

@@ -289,6 +289,7 @@ def create_oauth_callback_server(
             port=port,
             lifespan="off",
             log_level="warning",
+            ws="websockets-sansio",
         )
     )
 

src/fastmcp/server/auth/oauth_proxy.py

@@ -28,9 +28,13 @@ from urllib.parse import urlencode
 
 import httpx
 from authlib.common.security import generate_token
-from authlib.integrations.httpx_client import AsyncOAuth2Client
+from authlib.integrations.httpx_client import AsyncOAuth2Cli
 from key_value.aio.adapters.pydantic import PydanticAdapter
 from key_value.aio.protocols import AsyncKeyValue
+from mcp.server.auth.handlers.token import TokenErrorResponse, TokenSuccessResponse
+from mcp.server.auth.handlers.token import TokenHandler as _SDKTokenHandler
+from mcp.server.auth.json_response import PydanticJSONResponse
+from mcp.server.auth.middleware.client_auth import ClientAuthenticator
 from mcp.server.auth.provider import (
     AccessToken,
     AuthorizationCode,
@@ -38,6 +42,7 @@ from mcp.server.auth.provider import (
     RefreshToken,
     TokenError,
 )
+from mcp.server.auth.routes import cors_middleware
 from mcp.server.auth.settings import (
     ClientRegistrationOptions,
     RevocationOptions,
@@ -116,10 +121,53 @@ DEFAULT_AUTH_CODE_EXPIRY_SECONDS: Final[int] = 5 * 60  # 5 minutes
 HTTP_TIMEOUT_SECONDS: Final[int] = 30
 
 
-@dataclass
-class RelatedTokens(BaseModel):
-    access_token: str
-    refresh_token: str
+class TokenHandler(_SDKTokenHandler):
+    """TokenHandler that returns OAuth 2.1 compliant error responses.
+
+    The MCP SDK always returns HTTP 400 for all client authentication issues.
+    However, OAuth 2.1 Section 5.3 and the MCP specification require that
+    invalid or expired tokens MUST receive a HTTP 401 response.
+
+    This handler extends the base MCP SDK TokenHandler to transform client
+    authentication failures into OAuth 2.1 compliant responses:
+    - Changes 'unauthorized_client' to 'invalid_client' error code
+    - Returns HTTP 401 status code instead of 400 for client auth failures
+
+    Per OAuth 2.1 Section 5.3: "The authorization server MAY return an HTTP 401
+    (Unauthorized) status code to indicate which HTTP authentication schemes
+    are supported."
+
+    Per MCP spec: "Invalid or expired tokens MUST receive a HTTP 401 response."
+    """
+
+    def response(self, obj: TokenSuccessResponse | TokenErrorResponse):
+        """Override response method to provide OAuth 2.1 compliant error handling."""
+        # Check if this is a client authentication failure (not just unauthorized for grant type)
+        # unauthorized_client can mean two things:
+        # 1. Client authentication failed (client_id not found or wrong credentials) -> invalid_client 401
+        # 2. Client not authorized for this grant type -> unauthorized_client 400 (correct per spec)
+        if (
+            isinstance(obj, TokenErrorResponse)
+            and obj.error == "unauthorized_client"
+            and obj.error_description
+            and "Invalid client_id" in obj.error_description
+        ):
+            # Transform client auth failure to OAuth 2.1 compliant response
+            return PydanticJSONResponse(
+                content=TokenErrorResponse(
+                    error="invalid_client",
+                    error_description=obj.error_description,
+                    error_uri=obj.error_uri,
+                ),
+                status_code=401,
+                headers={
+                    "Cache-Control": "no-store",
+                    "Pragma": "no-cache",
+                },
+            )
+
+        # Otherwise use default behavior from parent class
+        return super().response(obj)
 
 
 class OAuthProxy(OAuthProvider):
@@ -835,9 +883,7 @@ class OAuthProxy(OAuthProvider):
                 f"Route {i}: {route} - path: {getattr(route, 'path', 'N/A')}, methods: {getattr(route, 'methods', 'N/A')}"
             )
 
-            # Keep all standard OAuth routes unchanged - our DCR-compliant flow handles everything
-            custom_routes.append(route)
-
+            # Replace the token endpoint with our custom handler that returns proper OAuth 2.1 error codes
             if (
                 isinstance(route, Route)
                 and route.path == "/token"
@@ -845,6 +891,22 @@ class OAuthProxy(OAuthProvider):
                 and "POST" in route.methods
             ):
                 token_route_found = True
+                # Replace with our OAuth 2.1 compliant token handler
+                token_handler = TokenHandler(
+                    provider=self, client_authenticator=ClientAuthenticator(self)
+                )
+                custom_routes.append(
+                    Route(
+                        path="/token",
+                        endpoint=cors_middleware(
+                            token_handler.handle, ["POST", "OPTIONS"]
+                        ),
+                        methods=["POST", "OPTIONS"],
+                    )
+                )
+            else:
+                # Keep all other standard OAuth routes unchanged
+                custom_routes.append(route)
 
         # Add OAuth callback endpoint for forwarding to client callbacks
         custom_routes.append(

src/fastmcp/server/server.py

@@ -1566,6 +1566,7 @@ class FastMCP(Generic[LifespanResultT]):
         config_kwargs: dict[str, Any] = {
             "timeout_graceful_shutdown": 0,
             "lifespan": "on",
+            "ws": "websockets-sansio",
         }
         config_kwargs.update(_uvicorn_config_from_user)
 

src/fastmcp/utilities/tests.py

@@ -66,6 +66,7 @@ def _run_server(mcp_server: FastMCP, transport: Literal["sse"], port: int) -> No
             host="127.0.0.1",
             port=port,
             log_level="error",
+            ws="websockets-sansio",
         )
     )
     uvicorn_server.run()

tests/client/test_sse.py

@@ -96,7 +96,9 @@ def run_nested_server(host: str, port: int) -> None:
     mount = Starlette(routes=[Mount("/nest-inner", app=app)])
     mount2 = Starlette(routes=[Mount("/nest-outer", app=mount)])
     server = uvicorn.Server(
-        config=uvicorn.Config(app=mount2, host=host, port=port, log_level="error")
+        config=uvicorn.Config(
+            app=mount2, host=host, port=port, log_level="error", ws="websockets-sansio"
+        )
     )
     server.run()
 

tests/client/test_streamable_http.py

@@ -103,6 +103,7 @@ def run_nested_server(host: str, port: int) -> None:
             port=port,
             log_level="error",
             lifespan="on",
+            ws="websockets-sansio",
         )
     )
     server.run()

tests/server/auth/test_oauth_proxy.py

@@ -243,7 +243,13 @@ class MockOAuthProvider:
                 self.port = s.getsockname()[1]
 
         self.base_url = f"http://localhost:{self.port}"
-        config = Config(self.app, host="localhost", port=self.port, log_level="error")
+        config = Config(
+            self.app,
+            host="localhost",
+            port=self.port,
+            log_level="error",
+            ws="websockets-sansio",
+        )
         self.server = Server(config)
 
         # Start server in background
@@ -973,3 +979,126 @@ class TestParameterForwarding:
         assert query_params["audience"][0] == "https://api.example.com"
         assert query_params["prompt"][0] == "consent"
         assert query_params["max_age"][0] == "3600"
+
+    @pytest.mark.asyncio
+    async def test_token_endpoint_invalid_client_error(self, jwt_verifier):
+        """Test that invalid client_id returns OAuth 2.1 compliant error response.
+
+        When a client ID is not found during token exchange, the proxy should:
+        1. Return HTTP 401 status code
+        2. Use 'invalid_client' error code instead of 'unauthorized_client'
+
+        This aligns with OAuth 2.1 spec and enables Claude's automatic client re-registration.
+        """
+        from starlette.applications import Starlette
+        from starlette.testclient import TestClient
+
+        proxy = OAuthProxy(
+            upstream_authorization_endpoint="https://oauth.example.com/authorize",
+            upstream_token_endpoint="https://oauth.example.com/token",
+            upstream_client_id="upstream-client",
+            upstream_client_secret="upstream-secret",
+            token_verifier=jwt_verifier,
+            base_url="https://proxy.example.com",
+        )
+
+        # Create a test app with OAuth routes
+        app = Starlette(routes=proxy.get_routes())
+
+        # Test the token endpoint with an invalid (non-existent) client_id
+        with TestClient(app) as client:
+            response = client.post(
+                "/token",
+                data={
+                    "grant_type": "authorization_code",
+                    "code": "test-auth-code",
+                    "client_id": "non-existent-client-id",
+                    "code_verifier": "test-code-verifier",
+                    "redirect_uri": "http://localhost:12345/callback",
+                },
+                headers={
+                    "Content-Type": "application/x-www-form-urlencoded",
+                },
+            )
+
+            # Verify OAuth 2.1 compliant error response
+            assert response.status_code == 401, (
+                f"Expected 401 but got {response.status_code}"
+            )
+
+            error_data = response.json()
+            assert error_data["error"] == "invalid_client", (
+                f"Expected 'invalid_client' but got '{error_data.get('error')}'"
+            )
+            assert "Invalid client_id" in error_data["error_description"]
+
+            # Verify proper cache headers are set
+            assert response.headers.get("Cache-Control") == "no-store"
+            assert response.headers.get("Pragma") == "no-cache"
+
+
+class TestTokenHandlerErrorTransformation:
+    """Tests for TokenHandler's OAuth 2.1 compliant error transformation."""
+
+    def test_transforms_client_auth_failure_to_invalid_client_401(self):
+        """Test that client authentication failures return invalid_client with 401."""
+        from mcp.server.auth.handlers.token import TokenErrorResponse
+
+        from fastmcp.server.auth.oauth_proxy import TokenHandler
+
+        handler = TokenHandler(provider=Mock(), client_authenticator=Mock())
+
+        # Simulate error from ClientAuthenticator.authenticate() failure
+        error_response = TokenErrorResponse(
+            error="unauthorized_client",
+            error_description="Invalid client_id 'test-client-id'",
+        )
+
+        response = handler.response(error_response)
+
+        # Should transform to OAuth 2.1 compliant response
+        assert response.status_code == 401
+        assert b'"error":"invalid_client"' in response.body
+        assert (
+            b'"error_description":"Invalid client_id \'test-client-id\'"'
+            in response.body
+        )
+
+    def test_does_not_transform_grant_type_unauthorized_to_invalid_client(self):
+        """Test that grant type authorization errors stay as unauthorized_client with 400."""
+        from mcp.server.auth.handlers.token import TokenErrorResponse
+
+        from fastmcp.server.auth.oauth_proxy import TokenHandler
+
+        handler = TokenHandler(provider=Mock(), client_authenticator=Mock())
+
+        # Simulate error from grant_type not in client_info.grant_types
+        error_response = TokenErrorResponse(
+            error="unauthorized_client",
+            error_description="Client not authorized for this grant type",
+        )
+
+        response = handler.response(error_response)
+
+        # Should NOT transform - keep as 400 unauthorized_client
+        assert response.status_code == 400
+        assert b'"error":"unauthorized_client"' in response.body
+
+    def test_does_not_transform_other_errors(self):
+        """Test that other error types pass through unchanged."""
+        from mcp.server.auth.handlers.token import TokenErrorResponse
+
+        from fastmcp.server.auth.oauth_proxy import TokenHandler
+
+        handler = TokenHandler(provider=Mock(), client_authenticator=Mock())
+
+        error_response = TokenErrorResponse(
+            error="invalid_grant",
+            error_description="Authorization code has expired",
+        )
+
+        response = handler.response(error_response)
+
+        # Should pass through unchanged
+        assert response.status_code == 400
+        assert b'"error":"invalid_grant"' in response.body

uv.lock

@@ -1,5 +1,5 @@
 version = 1
-revision = 2
+revision = 3
 requires-python = ">=3.10"
 resolution-markers = [
     "python_full_version >= '3.11'",
@@ -549,15 +549,13 @@ dependencies = [
     { name = "pyperclip" },
     { name = "python-dotenv" },
     { name = "rich" },
+    { name = "websockets" },
 ]
 
 [package.optional-dependencies]
 openai = [
     { name = "openai" },
 ]
-websockets = [
-    { name = "websockets" },
-]
 
 [package.dev-dependencies]
 dev = [
@@ -600,9 +598,9 @@ requires-dist = [
     { name = "pyperclip", specifier = ">=1.9.0" },
     { name = "python-dotenv", specifier = ">=1.1.0" },
     { name = "rich", specifier = ">=13.9.4" },
-    { name = "websockets", marker = "extra == 'websockets'", specifier = ">=15.0.1" },
+    { name = "websockets", specifier = ">=15.0.1" },
 ]
-provides-extras = ["openai", "websockets"]
+provides-extras = ["openai"]
 
 [package.metadata.requires-dev]
 dev = [