Move around the response models to be closer to the handlers

Author: praboud-ant

src/mcp/server/auth/handlers/authorize.py

@@ -53,6 +53,25 @@ class AuthorizationRequest(BaseModel):
     )
 
 
+AuthorizationErrorCode = Literal[
+    "invalid_request",
+    "unauthorized_client",
+    "access_denied",
+    "unsupported_response_type",
+    "invalid_scope",
+    "server_error",
+    "temporarily_unavailable",
+]
+
+
+class AuthorizationErrorResponse(BaseModel):
+    error: AuthorizationErrorCode
+    error_description: str
+    error_uri: AnyUrl | None = None
+    # must be set if provided in the request
+    state: str | None = None
+
+
 def validate_scope(
     requested_scope: str | None, client: OAuthClientInformationFull
 ) -> list[str] | None:
@@ -84,25 +103,6 @@ def validate_redirect_uri(
         )
 
 
-ErrorCode = Literal[
-    "invalid_request",
-    "unauthorized_client",
-    "access_denied",
-    "unsupported_response_type",
-    "invalid_scope",
-    "server_error",
-    "temporarily_unavailable",
-]
-
-
-class ErrorResponse(BaseModel):
-    error: ErrorCode
-    error_description: str
-    error_uri: AnyUrl | None = None
-    # must be set if provided in the request
-    state: str | None = None
-
-
 def best_effort_extract_string(
     key: str, params: None | FormData | QueryParams
 ) -> str | None:
@@ -132,7 +132,9 @@ async def handle(self, request: Request) -> Response:
         params = None
 
         async def error_response(
-            error: ErrorCode, error_description: str, attempt_load_client: bool = True
+            error: AuthorizationErrorCode,
+            error_description: str,
+            attempt_load_client: bool = True,
         ):
             nonlocal client, redirect_uri, state
             if client is None and attempt_load_client:
@@ -157,7 +159,7 @@ async def error_response(
                 # make last-ditch effort to load state
                 state = best_effort_extract_string("state", params)
 
-            error_resp = ErrorResponse(
+            error_resp = AuthorizationErrorResponse(
                 error=error,
                 error_description=error_description,
                 state=state,
@@ -194,7 +196,7 @@ async def error_response(
                 auth_request = AuthorizationRequest.model_validate(params)
                 state = auth_request.state  # Update with validated state
             except ValidationError as validation_error:
-                error: ErrorCode = "invalid_request"
+                error: AuthorizationErrorCode = "invalid_request"
                 for e in validation_error.errors():
                     if e["loc"] == ("response_type",) and e["type"] == "literal_error":
                         error = "unsupported_response_type"
@@ -264,11 +266,11 @@ async def error_response(
 
 
 def create_error_redirect(
-    redirect_uri: AnyUrl, error: Exception | ErrorResponse
+    redirect_uri: AnyUrl, error: Exception | AuthorizationErrorResponse
 ) -> str:
     parsed_uri = urlparse(str(redirect_uri))
 
-    if isinstance(error, ErrorResponse):
+    if isinstance(error, AuthorizationErrorResponse):
         # Convert ErrorResponse to dict
         error_dict = error.model_dump(exclude_none=True)
         query_params = {}

src/mcp/server/auth/handlers/register.py

@@ -10,7 +10,7 @@
 from typing import Literal
 from uuid import uuid4
 
-from pydantic import BaseModel, ValidationError
+from pydantic import BaseModel, RootModel, ValidationError
 from starlette.requests import Request
 from starlette.responses import Response
 
@@ -20,7 +20,13 @@
 from mcp.shared.auth import OAuthClientInformationFull, OAuthClientMetadata
 
 
-class ErrorResponse(BaseModel):
+class RegistrationRequest(RootModel):
+    # this wrapper is a no-op; it's just to separate out the types exposed to the
+    # provider from what we use in the HTTP handler
+    root: OAuthClientMetadata
+
+
+class RegistrationErrorResponse(BaseModel):
     error: Literal[
         "invalid_redirect_uri",
         "invalid_client_metadata",
@@ -43,7 +49,7 @@ async def handle(self, request: Request) -> Response:
             client_metadata = OAuthClientMetadata.model_validate(body)
         except ValidationError as validation_error:
             return PydanticJSONResponse(
-                content=ErrorResponse(
+                content=RegistrationErrorResponse(
                     error="invalid_client_metadata",
                     error_description=stringify_pydantic_error(validation_error),
                 ),

src/mcp/server/auth/handlers/revoke.py

@@ -5,26 +5,34 @@
 """
 
 from dataclasses import dataclass
+from typing import Literal
 
-from pydantic import ValidationError
+from pydantic import BaseModel, ValidationError
 from starlette.requests import Request
 from starlette.responses import Response
 
 from mcp.server.auth.errors import (
-    InvalidClientError,
     stringify_pydantic_error,
 )
 from mcp.server.auth.json_response import PydanticJSONResponse
 from mcp.server.auth.middleware.client_auth import (
     ClientAuthenticator,
-    ClientAuthRequest,
 )
-from mcp.server.auth.provider import OAuthServerProvider, OAuthTokenRevocationRequest
-from mcp.shared.auth import TokenErrorResponse
+from mcp.server.auth.provider import OAuthServerProvider
 
 
-class RevocationRequest(OAuthTokenRevocationRequest, ClientAuthRequest):
-    pass
+class RevocationRequest(BaseModel):
+    """
+    # See https://datatracker.ietf.org/doc/html/rfc7009#section-2.1
+    """
+
+    token: str
+    token_type_hint: Literal["access_token", "refresh_token"] | None = None
+
+
+class RevocationErrorResponse(BaseModel):
+    error: Literal["invalid_request",]
+    error_description: str | None = None
 
 
 @dataclass
@@ -42,21 +50,16 @@ async def handle(self, request: Request) -> Response:
         except ValidationError as e:
             return PydanticJSONResponse(
                 status_code=400,
-                content=TokenErrorResponse(
+                content=RevocationErrorResponse(
                     error="invalid_request",
                     error_description=stringify_pydantic_error(e),
                 ),
             )
 
-        # Authenticate client
-        try:
-            client_auth_result = await self.client_authenticator(revocation_request)
-        except InvalidClientError as e:
-            return PydanticJSONResponse(status_code=401, content=e.error_response())
-
         # Revoke token
-        if self.provider.revoke_token:
-            await self.provider.revoke_token(client_auth_result, revocation_request)
+        await self.provider.revoke_token(
+            revocation_request.token, revocation_request.token_type_hint
+        )
 
         # Return successful empty response
         return Response(

src/mcp/server/auth/handlers/token.py

@@ -10,7 +10,7 @@
 from dataclasses import dataclass
 from typing import Annotated, Literal
 
-from pydantic import AnyHttpUrl, Field, RootModel, ValidationError
+from pydantic import AnyHttpUrl, BaseModel, Field, RootModel, ValidationError
 from starlette.requests import Request
 
 from mcp.server.auth.errors import (
@@ -24,7 +24,7 @@
     ClientAuthRequest,
 )
 from mcp.server.auth.provider import OAuthServerProvider
-from mcp.shared.auth import TokenErrorResponse, TokenSuccessResponse
+from mcp.shared.auth import OAuthToken
 
 
 class AuthorizationCodeRequest(ClientAuthRequest):
@@ -53,6 +53,30 @@ class TokenRequest(RootModel):
     ]
 
 
+class TokenErrorResponse(BaseModel):
+    """
+    See https://datatracker.ietf.org/doc/html/rfc6749#section-5.2
+    """
+
+    error: Literal[
+        "invalid_request",
+        "invalid_client",
+        "invalid_grant",
+        "unauthorized_client",
+        "unsupported_grant_type",
+        "invalid_scope",
+    ]
+    error_description: str | None = None
+    error_uri: AnyHttpUrl | None = None
+
+
+class TokenSuccessResponse(RootModel):
+    # this is just a wrapper over OAuthToken; the only reason we do this
+    # is to have some separation between the HTTP response type, and the
+    # type returned by the provider
+    root: OAuthToken
+
+
 @dataclass
 class TokenHandler:
     provider: OAuthServerProvider
@@ -100,7 +124,7 @@ async def handle(self, request: Request):
                 )
             )
 
-        tokens: TokenSuccessResponse
+        tokens: OAuthToken
 
         match token_request:
             case AuthorizationCodeRequest():
@@ -208,4 +232,4 @@ async def handle(self, request: Request):
                     client_info, refresh_token, scopes
                 )
 
-        return self.response(tokens)
+        return self.response(TokenSuccessResponse(root=tokens))

src/mcp/server/auth/provider.py

@@ -12,7 +12,7 @@
 from mcp.server.auth.types import AuthInfo
 from mcp.shared.auth import (
     OAuthClientInformationFull,
-    TokenSuccessResponse,
+    OAuthToken,
 )
 
 
@@ -45,15 +45,6 @@ class RefreshToken(BaseModel):
     expires_at: int | None = None
 
 
-class OAuthTokenRevocationRequest(BaseModel):
-    """
-    # See https://datatracker.ietf.org/doc/html/rfc7009#section-2.1
-    """
-
-    token: str
-    token_type_hint: Literal["access_token", "refresh_token"] | None = None
-
-
 class OAuthRegisteredClientsStore(Protocol):
     """
     Interface for storing and retrieving registered OAuth clients.
@@ -149,7 +140,7 @@ async def load_authorization_code(
 
     async def exchange_authorization_code(
         self, client: OAuthClientInformationFull, authorization_code: AuthorizationCode
-    ) -> TokenSuccessResponse:
+    ) -> OAuthToken:
         """
         Exchanges an authorization code for an access token.
 
@@ -171,7 +162,7 @@ async def exchange_refresh_token(
         client: OAuthClientInformationFull,
         refresh_token: RefreshToken,
         scopes: list[str],
-    ) -> TokenSuccessResponse:
+    ) -> OAuthToken:
         """
         Exchanges a refresh token for an access token.
 
@@ -198,16 +189,20 @@ async def load_access_token(self, token: str) -> AuthInfo | None:
         ...
 
     async def revoke_token(
-        self, client: OAuthClientInformationFull, request: OAuthTokenRevocationRequest
+        self,
+        token: str,
+        token_type_hint: Literal["access_token", "refresh_token"] | None = None,
     ) -> None:
         """
         Revokes an access or refresh token.
 
         If the given token is invalid or already revoked, this method should do nothing.
 
         Args:
-            client: The client revoking the token.
-            request: The token revocation request.
+            token: the token to revoke
+            token_type_hint: hint about the type of token to revoke; optional. if the
+            token cannot be located using this hint, the provider MUST extend its search
+            to include all tokens.
         """
         ...
 

src/mcp/shared/auth.py

@@ -9,24 +9,7 @@
 from pydantic import AnyHttpUrl, BaseModel, Field
 
 
-class TokenErrorResponse(BaseModel):
-    """
-    See https://datatracker.ietf.org/doc/html/rfc6749#section-5.2
-    """
-
-    error: Literal[
-        "invalid_request",
-        "invalid_client",
-        "invalid_grant",
-        "unauthorized_client",
-        "unsupported_grant_type",
-        "invalid_scope",
-    ]
-    error_description: Optional[str] = None
-    error_uri: Optional[AnyHttpUrl] = None
-
-
-class TokenSuccessResponse(BaseModel):
+class OAuthToken(BaseModel):
     """
     See https://datatracker.ietf.org/doc/html/rfc6749#section-5.1
     """