Completely drop RootModel from types module (#1910)

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

@@ -2,7 +2,8 @@ import logging
 from dataclasses import dataclass
 from typing import Any, Literal
 
-from pydantic import AnyUrl, BaseModel, Field, RootModel, ValidationError
+# TODO(Marcelo): We should drop the `RootModel`.
+from pydantic import AnyUrl, BaseModel, Field, RootModel, ValidationError  # noqa: TID251
 from starlette.datastructures import FormData, QueryParams
 from starlette.requests import Request
 from starlette.responses import RedirectResponse, Response

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

@@ -4,7 +4,7 @@ from dataclasses import dataclass
 from typing import Any
 from uuid import uuid4
 
-from pydantic import BaseModel, RootModel, ValidationError
+from pydantic import BaseModel, ValidationError
 from starlette.requests import Request
 from starlette.responses import Response
 
@@ -14,11 +14,9 @@ from mcp.server.auth.provider import OAuthAuthorizationServerProvider, Registrat
 from mcp.server.auth.settings import ClientRegistrationOptions
 from mcp.shared.auth import OAuthClientInformationFull, OAuthClientMetadata
 
-
-class RegistrationRequest(RootModel[OAuthClientMetadata]):
-    # 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
+# this alias is a no-op; it's just to separate out the types exposed to the
+# provider from what we use in the HTTP handler
+RegistrationRequest = OAuthClientMetadata
 
 
 class RegistrationErrorResponse(BaseModel):
@@ -35,6 +33,7 @@ class RegistrationHandler:
         # Implements dynamic client registration as defined in https://datatracker.ietf.org/doc/html/rfc7591#section-3.1
         try:
             # Parse request body as JSON
+            # TODO(Marcelo): This is unnecessary. We should use `request.body()`.
             body = await request.json()
             client_metadata = OAuthClientMetadata.model_validate(body)
 

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

@@ -4,7 +4,7 @@ import time
 from dataclasses import dataclass
 from typing import Annotated, Any, Literal
 
-from pydantic import AnyHttpUrl, AnyUrl, BaseModel, Field, RootModel, ValidationError
+from pydantic import AnyHttpUrl, AnyUrl, BaseModel, Field, TypeAdapter, ValidationError
 from starlette.requests import Request
 
 from mcp.server.auth.errors import stringify_pydantic_error
@@ -40,18 +40,8 @@ class RefreshTokenRequest(BaseModel):
     resource: str | None = Field(None, description="Resource indicator for the token")
 
 
-class TokenRequest(
-    RootModel[
-        Annotated[
-            AuthorizationCodeRequest | RefreshTokenRequest,
-            Field(discriminator="grant_type"),
-        ]
-    ]
-):
-    root: Annotated[
-        AuthorizationCodeRequest | RefreshTokenRequest,
-        Field(discriminator="grant_type"),
-    ]
+TokenRequest = Annotated[AuthorizationCodeRequest | RefreshTokenRequest, Field(discriminator="grant_type")]
+token_request_adapter = TypeAdapter[TokenRequest](TokenRequest)
 
 
 class TokenErrorResponse(BaseModel):
@@ -62,11 +52,10 @@ class TokenErrorResponse(BaseModel):
     error_uri: AnyHttpUrl | None = None
 
 
-class TokenSuccessResponse(RootModel[OAuthToken]):
-    # 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
+# this is just an alias 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
+TokenSuccessResponse = OAuthToken
 
 
 @dataclass
@@ -107,7 +96,8 @@ class TokenHandler:
 
         try:
             form_data = await request.form()
-            token_request = TokenRequest.model_validate(dict(form_data)).root
+            # TODO(Marcelo): Can someone check if this `dict()` wrapper is necessary?
+            token_request = token_request_adapter.validate_python(dict(form_data))
         except ValidationError as validation_error:  # pragma: no cover
             return self.response(
                 TokenErrorResponse(
@@ -186,12 +176,7 @@ class TokenHandler:
                     # Exchange authorization code for tokens
                     tokens = await self.provider.exchange_authorization_code(client_info, auth_code)
                 except TokenError as e:
-                    return self.response(
-                        TokenErrorResponse(
-                            error=e.error,
-                            error_description=e.error_description,
-                        )
-                    )
+                    return self.response(TokenErrorResponse(error=e.error, error_description=e.error_description))
 
             case RefreshTokenRequest():  # pragma: no cover
                 refresh_token = await self.provider.load_refresh_token(client_info, token_request.refresh_token)
@@ -229,11 +214,6 @@ class TokenHandler:
                     # Exchange refresh token for new tokens
                     tokens = await self.provider.exchange_refresh_token(client_info, refresh_token, scopes)
                 except TokenError as e:
-                    return self.response(
-                        TokenErrorResponse(
-                            error=e.error,
-                            error_description=e.error_description,
-                        )
-                    )
+                    return self.response(TokenErrorResponse(error=e.error, error_description=e.error_description))
 
-        return self.response(TokenSuccessResponse(root=tokens))
+        return self.response(tokens)