Add back authorization to the /revoke endpoint, simplify revoke

Author: praboud-ant

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

@@ -1,18 +1,20 @@
 from dataclasses import dataclass
+from functools import partial
 from typing import Literal
 
 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,
 )
-from mcp.server.auth.provider import OAuthServerProvider
+from mcp.server.auth.provider import AuthInfo, OAuthServerProvider, RefreshToken
 
 
 class RevocationRequest(BaseModel):
@@ -22,6 +24,8 @@ class RevocationRequest(BaseModel):
 
     token: str
     token_type_hint: Literal["access_token", "refresh_token"] | None = None
+    client_id: str
+    client_secret: str | None
 
 
 class RevocationErrorResponse(BaseModel):
@@ -50,10 +54,30 @@ async def handle(self, request: Request) -> Response:
                 ),
             )
 
-        # Revoke token
-        await self.provider.revoke_token(
-            revocation_request.token, revocation_request.token_type_hint
-        )
+        # Authenticate client
+        try:
+            client = await self.client_authenticator.authenticate(
+                revocation_request.client_id, revocation_request.client_secret
+            )
+        except InvalidClientError as e:
+            return PydanticJSONResponse(status_code=401, content=e.error_response())
+
+        loaders = [
+            self.provider.load_access_token,
+            partial(self.provider.load_refresh_token, client),
+        ]
+        if revocation_request.token_type_hint == "refresh_token":
+            loaders = reversed(loaders)
+
+        token: None | AuthInfo | RefreshToken = None
+        for loader in loaders:
+            token = await loader(revocation_request.token)
+            if token is not None:
+                break
+
+        if token and token.client_id == client.client_id:
+            # Revoke token
+            await self.provider.revoke_token(token)
 
         # Return successful empty response
         return Response(

src/mcp/server/auth/provider.py

@@ -1,4 +1,4 @@
-from typing import Literal, Protocol
+from typing import Protocol
 from urllib.parse import parse_qs, urlencode, urlparse, urlunparse
 
 from pydantic import AnyHttpUrl, BaseModel
@@ -172,8 +172,7 @@ async def load_access_token(self, token: str) -> AuthInfo | None:
 
     async def revoke_token(
         self,
-        token: str,
-        token_type_hint: Literal["access_token", "refresh_token"] | None = None,
+        token: AuthInfo | RefreshToken,
     ) -> None:
         """
         Revokes an access or refresh token.

tests/server/fastmcp/auth/test_auth_integration.py

@@ -8,7 +8,6 @@
 import secrets
 import time
 import unittest.mock
-from typing import Literal
 from urllib.parse import parse_qs, urlparse
 
 import anyio
@@ -164,11 +163,12 @@ async def exchange_refresh_token(
         new_refresh_token = f"refresh_{secrets.token_hex(32)}"
 
         # Store the new tokens
-        self.tokens[new_access_token] = {
-            "client_id": client.client_id,
-            "scopes": scopes or token_info.scopes,
-            "expires_at": int(time.time()) + 3600,
-        }
+        self.tokens[new_access_token] = AuthInfo(
+            token=new_access_token,
+            client_id=client.client_id,
+            scopes=scopes or token_info.scopes,
+            expires_at=int(time.time()) + 3600,
+        )
 
         self.refresh_tokens[new_refresh_token] = new_access_token
 
@@ -198,25 +198,20 @@ async def load_access_token(self, token: str) -> AuthInfo | None:
             expires_at=token_info.expires_at,
         )
 
-    async def revoke_token(
-        self,
-        token: str,
-        token_type_hint: Literal["access_token", "refresh_token"] | None = None,
-    ) -> None:
-        # Check if it's a refresh token
-        if token in self.refresh_tokens:
-            # Remove the refresh token
-            del self.refresh_tokens[token]
-
-        # Check if it's an access token
-        elif token in self.tokens:
-            # Remove the access token
-            del self.tokens[token]
-
-            # Also remove any refresh tokens that point to this access token
-            for refresh_token, access_token in list(self.refresh_tokens.items()):
-                if access_token == token:
-                    del self.refresh_tokens[refresh_token]
+    async def revoke_token(self, token: OAuthToken | RefreshToken) -> None:
+        match token:
+            case RefreshToken():
+                # Remove the refresh token
+                del self.refresh_tokens[token.token]
+
+            case AuthInfo():
+                # Remove the access token
+                del self.tokens[token.token]
+
+                # Also remove any refresh tokens that point to this access token
+                for refresh_token, access_token in list(self.refresh_tokens.items()):
+                    if access_token == token.token:
+                        del self.refresh_tokens[refresh_token]
 
 
 @pytest.fixture