Improved error handling, generic types for provider

Author: praboud-ant

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

@@ -14,7 +14,9 @@
 )
 from mcp.server.auth.json_response import PydanticJSONResponse
 from mcp.server.auth.provider import (
+    AuthorizationErrorCode,
     AuthorizationParams,
+    AuthorizeError,
     OAuthServerProvider,
     construct_redirect_uri,
 )
@@ -49,20 +51,9 @@ 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_description: str | None
     error_uri: AnyUrl | None = None
     # must be set if provided in the request
     state: str | None = None
@@ -98,16 +89,14 @@ async def handle(self, request: Request) -> Response:
 
         async def error_response(
             error: AuthorizationErrorCode,
-            error_description: str,
+            error_description: str | None,
             attempt_load_client: bool = True,
         ):
             nonlocal client, redirect_uri, state
             if client is None and attempt_load_client:
                 # make last-ditch attempt to load the client
                 client_id = best_effort_extract_string("client_id", params)
-                client = client_id and await self.provider.clients_store.get_client(
-                    client_id
-                )
+                client = client_id and await self.provider.get_client(client_id)
             if redirect_uri is None and client:
                 # make last-ditch effort to load the redirect uri
                 if params is not None and "redirect_uri" not in params:
@@ -171,7 +160,7 @@ async def error_response(
                 )
 
             # Get client information
-            client = await self.provider.clients_store.get_client(
+            client = await self.provider.get_client(
                 auth_request.client_id,
             )
             if not client:
@@ -210,15 +199,22 @@ async def error_response(
                 redirect_uri=redirect_uri,
             )
 
-            # Let the provider pick the next URI to redirect to
-            return RedirectResponse(
-                url=await self.provider.authorize(
-                    client,
-                    auth_params,
-                ),
-                status_code=302,
-                headers={"Cache-Control": "no-store"},
-            )
+            try:
+                # Let the provider pick the next URI to redirect to
+                return RedirectResponse(
+                    url=await self.provider.authorize(
+                        client,
+                        auth_params,
+                    ),
+                    status_code=302,
+                    headers={"Cache-Control": "no-store"},
+                )
+            except AuthorizeError as e:
+                # Handle authorization errors as defined in RFC 6749 Section 4.1.2.1
+                return await error_response(
+                    error=e.error,
+                    error_description=e.error_description,
+                )
 
         except Exception as validation_error:
             # Catch-all for unexpected errors

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

@@ -1,7 +1,6 @@
 import secrets
 import time
 from dataclasses import dataclass
-from typing import Literal
 from uuid import uuid4
 
 from pydantic import BaseModel, RootModel, ValidationError
@@ -10,7 +9,11 @@
 
 from mcp.server.auth.errors import stringify_pydantic_error
 from mcp.server.auth.json_response import PydanticJSONResponse
-from mcp.server.auth.provider import OAuthRegisteredClientsStore
+from mcp.server.auth.provider import (
+    OAuthServerProvider,
+    RegistrationError,
+    RegistrationErrorCode,
+)
 from mcp.server.auth.settings import ClientRegistrationOptions
 from mcp.shared.auth import OAuthClientInformationFull, OAuthClientMetadata
 
@@ -22,18 +25,13 @@ class RegistrationRequest(RootModel):
 
 
 class RegistrationErrorResponse(BaseModel):
-    error: Literal[
-        "invalid_redirect_uri",
-        "invalid_client_metadata",
-        "invalid_software_statement",
-        "unapproved_software_statement",
-    ]
-    error_description: str
+    error: RegistrationErrorCode
+    error_description: str | None
 
 
 @dataclass
 class RegistrationHandler:
-    clients_store: OAuthRegisteredClientsStore
+    provider: OAuthServerProvider
     options: ClientRegistrationOptions
 
     async def handle(self, request: Request) -> Response:
@@ -116,8 +114,17 @@ async def handle(self, request: Request) -> Response:
             software_id=client_metadata.software_id,
             software_version=client_metadata.software_version,
         )
-        # Register client
-        await self.clients_store.register_client(client_info)
+        try:
+            # Register client
+            await self.provider.register_client(client_info)
 
-        # Return client information
-        return PydanticJSONResponse(content=client_info, status_code=201)
+            # Return client information
+            return PydanticJSONResponse(content=client_info, status_code=201)
+        except RegistrationError as e:
+            # Handle registration errors as defined in RFC 7591 Section 3.2.2
+            return PydanticJSONResponse(
+                content=RegistrationErrorResponse(
+                    error=e.error, error_description=e.error_description
+                ),
+                status_code=400,
+            )

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

@@ -81,8 +81,10 @@ async def handle(self, request: Request) -> Response:
             if token is not None:
                 break
 
+        # if token is not found, just return HTTP 200 per the RFC
         if token and token.client_id == client.client_id:
-            # Revoke token
+            # Revoke token; provider is not meant to be able to do validation
+            # at this point that would result in an error
             await self.provider.revoke_token(token)
 
         # Return successful empty response

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

@@ -16,7 +16,7 @@
     AuthenticationError,
     ClientAuthenticator,
 )
-from mcp.server.auth.provider import OAuthServerProvider
+from mcp.server.auth.provider import OAuthServerProvider, TokenError, TokenErrorCode
 from mcp.shared.auth import OAuthToken
 
 
@@ -56,14 +56,7 @@ 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: TokenErrorCode
     error_description: str | None = None
     error_uri: AnyHttpUrl | None = None
 
@@ -184,10 +177,18 @@ async def handle(self, request: Request):
                         )
                     )
 
-                # Exchange authorization code for tokens
-                tokens = await self.provider.exchange_authorization_code(
-                    client_info, auth_code
-                )
+                try:
+                    # 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,
+                        )
+                    )
 
             case RefreshTokenRequest():
                 refresh_token = await self.provider.load_refresh_token(
@@ -233,9 +234,17 @@ async def handle(self, request: Request):
                             )
                         )
 
-                # Exchange refresh token for new tokens
-                tokens = await self.provider.exchange_refresh_token(
-                    client_info, refresh_token, scopes
-                )
+                try:
+                    # 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(TokenSuccessResponse(root=tokens))

src/mcp/server/auth/middleware/client_auth.py

@@ -1,6 +1,6 @@
 import time
 
-from mcp.server.auth.provider import OAuthRegisteredClientsStore
+from mcp.server.auth.provider import OAuthServerProvider
 from mcp.shared.auth import OAuthClientInformationFull
 
 
@@ -20,20 +20,20 @@ class ClientAuthenticator:
     logic is skipped.
     """
 
-    def __init__(self, clients_store: OAuthRegisteredClientsStore):
+    def __init__(self, provider: OAuthServerProvider):
         """
         Initialize the dependency.
 
         Args:
-            clients_store: Store to look up client information
+            provider: Provider to look up client information
         """
-        self.clients_store = clients_store
+        self.provider = provider
 
     async def authenticate(
         self, client_id: str, client_secret: str | None
     ) -> OAuthClientInformationFull:
         # Look up client information
-        client = await self.clients_store.get_client(client_id)
+        client = await self.provider.get_client(client_id)
         if not client:
             raise AuthenticationError("Invalid client_id")