diff --git a/pyproject.toml b/pyproject.toml index 97feb117f..130d85c57 100644 --- a/pyproject.toml +++ b/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] diff --git a/src/fastmcp/client/oauth_callback.py b/src/fastmcp/client/oauth_callback.py index 8a92b9b08..c6da794e8 100644 --- a/src/fastmcp/client/oauth_callback.py +++ b/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", ) ) diff --git a/src/fastmcp/server/auth/oauth_proxy.py b/src/fastmcp/server/auth/oauth_proxy.py index 599defe0e..6cab473d7 100644 --- a/src/fastmcp/server/auth/oauth_proxy.py +++ b/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( diff --git a/src/fastmcp/server/server.py b/src/fastmcp/server/server.py index 2341f8634..489561166 100644 --- a/src/fastmcp/server/server.py +++ b/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) diff --git a/src/fastmcp/utilities/tests.py b/src/fastmcp/utilities/tests.py index 309a4bf9d..9d0efeaca 100644 --- a/src/fastmcp/utilities/tests.py +++ b/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() diff --git a/tests/client/test_sse.py b/tests/client/test_sse.py index 5bc35d50b..a60b975ac 100644 --- a/tests/client/test_sse.py +++ b/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() diff --git a/tests/client/test_streamable_http.py b/tests/client/test_streamable_http.py index 43923239a..9c7896330 100644 --- a/tests/client/test_streamable_http.py +++ b/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() diff --git a/tests/server/auth/test_oauth_proxy.py b/tests/server/auth/test_oauth_proxy.py index 3b52d7e47..c171547eb 100644 --- a/tests/server/auth/test_oauth_proxy.py +++ b/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 diff --git a/uv.lock b/uv.lock index 87136c3b2..918781ab3 100644 --- a/uv.lock +++ b/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 = [