Move TokenHandler to OAuthProvider for consistent error codes (#2538)

All OAuth providers now return correct invalid_client error codes
instead of unauthorized_client for auth failures. Previously only
OAuthProxy had this fix; now OAuthProvider (and InMemoryOAuthProvider)
also benefit.

Author: jlowin

src/fastmcp/server/auth/auth.py

@@ -1,10 +1,15 @@
 from __future__ import annotations
 
+import json
 from typing import Any, cast
 from urllib.parse import urlparse
 
+from mcp.server.auth.handlers.token import TokenErrorResponse
+from mcp.server.auth.handlers.token import TokenHandler as _SDKTokenHandler
+from mcp.server.auth.json_response import PydanticJSONResponse
 from mcp.server.auth.middleware.auth_context import AuthContextMiddleware
 from mcp.server.auth.middleware.bearer_auth import BearerAuthBackend
+from mcp.server.auth.middleware.client_auth import ClientAuthenticator
 from mcp.server.auth.provider import (
     AccessToken as _SDKAccessToken,
 )
@@ -17,6 +22,7 @@
     TokenVerifier as TokenVerifierProtocol,
 )
 from mcp.server.auth.routes import (
+    cors_middleware,
     create_auth_routes,
     create_protected_resource_routes,
 )
@@ -40,6 +46,48 @@ class AccessToken(_SDKAccessToken):
     claims: dict[str, Any] = Field(default_factory=dict)
 
 
+class TokenHandler(_SDKTokenHandler):
+    """TokenHandler that returns OAuth 2.1 compliant error responses.
+
+    The MCP SDK returns `unauthorized_client` for client authentication failures.
+    However, per RFC 6749 Section 5.2, authentication failures should return
+    `invalid_client` with HTTP 401, not `unauthorized_client`.
+
+    This distinction matters: `unauthorized_client` means "client exists but
+    can't do this", while `invalid_client` means "client doesn't exist or
+    credentials are wrong". Claude's OAuth client uses this to decide whether
+    to re-register.
+
+    This handler transforms 401 responses with `unauthorized_client` to use
+    `invalid_client` instead, making the error semantics correct per OAuth spec.
+    """
+
+    async def handle(self, request: Any):
+        """Wrap SDK handle() and transform auth error responses."""
+        response = await super().handle(request)
+
+        # Transform 401 unauthorized_client -> invalid_client
+        if response.status_code == 401:
+            try:
+                body = json.loads(response.body)
+                if body.get("error") == "unauthorized_client":
+                    return PydanticJSONResponse(
+                        content=TokenErrorResponse(
+                            error="invalid_client",
+                            error_description=body.get("error_description"),
+                        ),
+                        status_code=401,
+                        headers={
+                            "Cache-Control": "no-store",
+                            "Pragma": "no-cache",
+                        },
+                    )
+            except (json.JSONDecodeError, AttributeError):
+                pass  # Not JSON or unexpected format, return as-is
+
+        return response
+
+
 class AuthProvider(TokenVerifierProtocol):
     """Base class for all FastMCP authentication providers.
 
@@ -368,14 +416,40 @@ def get_routes(
             self.issuer_url is not None
         )  # typing check (issuer_url defaults to base_url)
 
-        oauth_routes = create_auth_routes(
+        sdk_routes = create_auth_routes(
             provider=self,
             issuer_url=self.base_url,
             service_documentation_url=self.service_documentation_url,
             client_registration_options=self.client_registration_options,
             revocation_options=self.revocation_options,
         )
 
+        # Replace the token endpoint with our custom handler that returns
+        # proper OAuth 2.1 error codes (invalid_client instead of unauthorized_client)
+        oauth_routes: list[Route] = []
+        for route in sdk_routes:
+            if (
+                isinstance(route, Route)
+                and route.path == "/token"
+                and route.methods is not None
+                and "POST" in route.methods
+            ):
+                # Replace with our OAuth 2.1 compliant token handler
+                token_handler = TokenHandler(
+                    provider=self, client_authenticator=ClientAuthenticator(self)
+                )
+                oauth_routes.append(
+                    Route(
+                        path="/token",
+                        endpoint=cors_middleware(
+                            token_handler.handle, ["POST", "OPTIONS"]
+                        ),
+                        methods=["POST", "OPTIONS"],
+                    )
+                )
+            else:
+                oauth_routes.append(route)
+
         # Get the resource URL based on the MCP path
         resource_url = self._get_resource_url(mcp_path)
 

src/fastmcp/server/auth/oauth_proxy.py

@@ -36,10 +36,6 @@
 from key_value.aio.protocols import AsyncKeyValue
 from key_value.aio.stores.disk import DiskStore
 from key_value.aio.wrappers.encryption import FernetEncryptionWrapper
-from mcp.server.auth.handlers.token import TokenErrorResponse
-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,
@@ -48,7 +44,6 @@
     RefreshToken,
     TokenError,
 )
-from mcp.server.auth.routes import cors_middleware
 from mcp.server.auth.settings import (
     ClientRegistrationOptions,
     RevocationOptions,
@@ -514,53 +509,6 @@ def create_error_html(
     )
 
 
-# -------------------------------------------------------------------------
-# Handler Classes
-# -------------------------------------------------------------------------
-
-
-class TokenHandler(_SDKTokenHandler):
-    """TokenHandler that returns OAuth 2.1 compliant error responses.
-
-    The MCP SDK returns `unauthorized_client` for client authentication failures.
-    However, per RFC 6749 Section 5.2, authentication failures should return
-    `invalid_client` with HTTP 401, not `unauthorized_client`.
-
-    This distinction matters: `unauthorized_client` means "client exists but
-    can't do this", while `invalid_client` means "client doesn't exist or
-    credentials are wrong". Claude's OAuth client uses this to decide whether
-    to re-register.
-
-    This handler transforms 401 responses with `unauthorized_client` to use
-    `invalid_client` instead, making the error semantics correct per OAuth spec.
-    """
-
-    async def handle(self, request: Any):
-        """Wrap SDK handle() and transform auth error responses."""
-        response = await super().handle(request)
-
-        # Transform 401 unauthorized_client -> invalid_client
-        if response.status_code == 401:
-            try:
-                body = json.loads(response.body)
-                if body.get("error") == "unauthorized_client":
-                    return PydanticJSONResponse(
-                        content=TokenErrorResponse(
-                            error="invalid_client",
-                            error_description=body.get("error_description"),
-                        ),
-                        status_code=401,
-                        headers={
-                            "Cache-Control": "no-store",
-                            "Pragma": "no-cache",
-                        },
-                    )
-            except (json.JSONDecodeError, AttributeError):
-                pass  # Not JSON or unexpected format, return as-is
-
-        return response
-
-
 class OAuthProxy(OAuthProvider):
     """OAuth provider that presents a DCR-compliant interface while proxying to non-DCR IDPs.
 
@@ -1668,10 +1616,9 @@ def get_routes(
                 This is used to advertise the resource URL in metadata.
         """
         # Get standard OAuth routes from parent class
+        # Note: parent already replaces /token with TokenHandler for proper error codes
         routes = super().get_routes(mcp_path)
         custom_routes = []
-        token_route_found = False
-        authorize_route_found = False
 
         logger.debug(
             f"get_routes called - configuring OAuth routes in {len(routes)} routes"
@@ -1689,7 +1636,6 @@ def get_routes(
                 and route.methods is not None
                 and ("GET" in route.methods or "POST" in route.methods)
             ):
-                authorize_route_found = True
                 # Replace with our enhanced authorization handler
                 # Note: self.base_url is guaranteed to be set in parent __init__
                 authorize_handler = AuthorizationHandler(
@@ -1705,27 +1651,6 @@ def get_routes(
                         methods=["GET", "POST"],
                     )
                 )
-            # Replace the token endpoint with our custom handler that returns proper OAuth 2.1 error codes
-            elif (
-                isinstance(route, Route)
-                and route.path == "/token"
-                and route.methods is not None
-                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)
@@ -1747,9 +1672,6 @@ def get_routes(
             )
         )
 
-        logger.debug(
-            f"✅ OAuth routes configured: authorize_endpoint={authorize_route_found}, token_endpoint={token_route_found}, total routes={len(custom_routes)} (includes OAuth callback + consent)"
-        )
         return custom_routes
 
     # -------------------------------------------------------------------------

tests/server/auth/test_oauth_proxy.py

@@ -1254,7 +1254,7 @@ async def test_transforms_client_auth_failure_to_invalid_client_401(self):
 
         from mcp.server.auth.handlers.token import TokenHandler as SDKTokenHandler
 
-        from fastmcp.server.auth.oauth_proxy import TokenHandler
+        from fastmcp.server.auth.auth import TokenHandler
 
         handler = TokenHandler(provider=Mock(), client_authenticator=Mock())
 
@@ -1283,7 +1283,7 @@ 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
+        from fastmcp.server.auth.auth import TokenHandler
 
         handler = TokenHandler(provider=Mock(), client_authenticator=Mock())
 
@@ -1303,7 +1303,7 @@ 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
+        from fastmcp.server.auth.auth import TokenHandler
 
         handler = TokenHandler(provider=Mock(), client_authenticator=Mock())