OAuth Update: Adding the Client Credentials & Token Exchange Grant Types

docs/api.md

@@ -1 +1,5 @@
+The Python SDK exposes the entire `mcp` package for use in your own projects.
+It includes an OAuth server implementation with support for the RFC 8693
+`token_exchange` grant type.
+
 ::: mcp

docs/index.md

@@ -3,3 +3,7 @@
 This is the MCP Server implementation in Python.
 
 It only contains the [API Reference](api.md) for the time being.
+
+The built-in OAuth server supports the RFC 8693 `token_exchange` grant type,
+allowing clients to exchange user tokens from external providers for MCP
+access tokens.

examples/servers/simple-auth/mcp_simple_auth/simple_auth_provider.py

@@ -238,6 +238,52 @@ async def exchange_authorization_code(
             scope=" ".join(authorization_code.scopes),
         )
 
+    async def exchange_client_credentials(self, client: OAuthClientInformationFull, scopes: list[str]) -> OAuthToken:
+        """Exchange client credentials for an MCP access token."""
+        mcp_token = f"mcp_{secrets.token_hex(32)}"
+        self.tokens[mcp_token] = AccessToken(
+            token=mcp_token,
+            client_id=client.client_id,
+            scopes=scopes,
+            expires_at=int(time.time()) + 3600,
+        )
+        return OAuthToken(
+            access_token=mcp_token,
+            token_type="Bearer",
+            expires_in=3600,
+            scope=" ".join(scopes),
+        )
+
+    async def exchange_token(
+        self,
+        client: OAuthClientInformationFull,
+        subject_token: str,
+        subject_token_type: str,
+        actor_token: str | None,
+        actor_token_type: str | None,
+        scope: list[str] | None,
+        audience: str | None,
+        resource: str | None,
+    ) -> OAuthToken:
+        """Exchange an external token for an MCP access token."""
+        if not subject_token:
+            raise ValueError("Invalid subject token")
+
+        mcp_token = f"mcp_{secrets.token_hex(32)}"
+        self.tokens[mcp_token] = AccessToken(
+            token=mcp_token,
+            client_id=client.client_id,
+            scopes=scope or [self.settings.mcp_scope],
+            expires_at=int(time.time()) + 3600,
+            resource=resource,
+        )
+        return OAuthToken(
+            access_token=mcp_token,
+            token_type="Bearer",
+            expires_in=3600,
+            scope=" ".join(scope or [self.settings.mcp_scope]),
+        )
+
     async def load_access_token(self, token: str) -> AccessToken | None:
         """Load and validate an access token."""
         access_token = self.tokens.get(token)

src/mcp/client/auth.py

@@ -125,15 +125,15 @@ def update_token_expiry(self, token: OAuthToken) -> None:
             self.token_expiry_time = None
 
     def is_token_valid(self) -> bool:
-        """Check if current token is valid."""
+        """Check if the current token is valid."""
         return bool(
             self.current_tokens
             and self.current_tokens.access_token
             and (not self.token_expiry_time or time.time() <= self.token_expiry_time)
         )
 
     def can_refresh_token(self) -> bool:
-        """Check if token can be refreshed."""
+        """Check if the token can be refreshed."""
         return bool(self.current_tokens and self.current_tokens.refresh_token and self.client_info)
 
     def clear_tokens(self) -> None:
@@ -546,6 +546,286 @@ async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.
                     logger.exception("OAuth flow error")
                     raise
 
-        # Retry with new tokens
-        self._add_auth_header(request)
-        yield request
+                    # Retry with new tokens
+                    self._add_auth_header(request)
+                    yield request
+
+
+class ClientCredentialsProvider(httpx.Auth):
+    """HTTPX auth using the OAuth2 client credentials grant."""
+
+    def __init__(
+        self,
+        server_url: str,
+        client_metadata: OAuthClientMetadata,
+        storage: TokenStorage,
+        resource: str | None = None,
+        timeout: float = 300.0,
+    ):
+        self.server_url = server_url
+        self.client_metadata = client_metadata
+        self.storage = storage
+        self.timeout = timeout
+        self.resource = resource or resource_url_from_server_url(server_url)
+
+        self._current_tokens: OAuthToken | None = None
+        self._metadata: OAuthMetadata | None = None
+        self._client_info: OAuthClientInformationFull | None = None
+        self._token_expiry_time: float | None = None
+
+        self._token_lock = anyio.Lock()
+
+    def _get_authorization_base_url(self, server_url: str) -> str:
+        """Return base authorization server URL without path."""
+        parsed = urlparse(server_url)
+        return f"{parsed.scheme}://{parsed.netloc}"
+
+    async def _discover_oauth_metadata(self, server_url: str) -> OAuthMetadata | None:
+        """Discover OAuth server metadata for client credentials."""
+        auth_base_url = self._get_authorization_base_url(server_url)
+        url = urljoin(auth_base_url, "/.well-known/oauth-authorization-server")
+        headers = {"MCP-Protocol-Version": LATEST_PROTOCOL_VERSION}
+
+        async with httpx.AsyncClient() as client:
+            try:
+                response = await client.get(url, headers=headers)
+                if response.status_code == 404:
+                    return None
+                response.raise_for_status()
+                return OAuthMetadata.model_validate(response.json())
+            except Exception:
+                try:
+                    response = await client.get(url)
+                    if response.status_code == 404:
+                        return None
+                    response.raise_for_status()
+                    return OAuthMetadata.model_validate(response.json())
+                except Exception:
+                    logger.exception("Failed to discover OAuth metadata")
+                    return None
+
+    async def _register_oauth_client(
+        self,
+        server_url: str,
+        client_metadata: OAuthClientMetadata,
+        metadata: OAuthMetadata | None = None,
+    ) -> OAuthClientInformationFull:
+        if not metadata:
+            metadata = await self._discover_oauth_metadata(server_url)
+
+        if metadata and metadata.registration_endpoint:
+            registration_url = str(metadata.registration_endpoint)
+        else:
+            auth_base_url = self._get_authorization_base_url(server_url)
+            registration_url = urljoin(auth_base_url, "/register")
+
+        if client_metadata.scope is None and metadata and metadata.scopes_supported is not None:
+            client_metadata.scope = " ".join(metadata.scopes_supported)
+
+        registration_data = client_metadata.model_dump(by_alias=True, mode="json", exclude_none=True)
+
+        async with httpx.AsyncClient() as client:
+            response = await client.post(
+                registration_url,
+                json=registration_data,
+                headers={"Content-Type": "application/json"},
+            )
+
+            if response.status_code not in (200, 201):
+                raise httpx.HTTPStatusError(
+                    f"Registration failed: {response.status_code}",
+                    request=response.request,
+                    response=response,
+                )
+
+            return OAuthClientInformationFull.model_validate(response.json())
+
+    def _has_valid_token(self) -> bool:
+        if not self._current_tokens or not self._current_tokens.access_token:
+            return False
+
+        if self._token_expiry_time and time.time() > self._token_expiry_time:
+            return False
+        return True
+
+    async def _validate_token_scopes(self, token_response: OAuthToken) -> None:
+        if not token_response.scope:
+            return
+
+        requested_scopes: set[str] = set()
+        if self.client_metadata.scope:
+            requested_scopes = set(self.client_metadata.scope.split())
+            returned_scopes = set(token_response.scope.split())
+            unauthorized_scopes = returned_scopes - requested_scopes
+            if unauthorized_scopes:
+                raise Exception(f"Server granted unauthorized scopes: {unauthorized_scopes}.")
+        else:
+            granted = set(token_response.scope.split())
+            logger.debug(
+                "No explicit scopes requested, accepting server-granted scopes: %s",
+                granted,
+            )
+
+    async def initialize(self) -> None:
+        self._current_tokens = await self.storage.get_tokens()
+        self._client_info = await self.storage.get_client_info()
+
+    async def _get_or_register_client(self) -> OAuthClientInformationFull:
+        if not self._client_info:
+            self._client_info = await self._register_oauth_client(self.server_url, self.client_metadata, self._metadata)
+            await self.storage.set_client_info(self._client_info)
+        return self._client_info
+
+    async def _request_token(self) -> None:
+        if not self._metadata:
+            self._metadata = await self._discover_oauth_metadata(self.server_url)
+
+        client_info = await self._get_or_register_client()
+
+        if self._metadata and self._metadata.token_endpoint:
+            token_url = str(self._metadata.token_endpoint)
+        else:
+            auth_base_url = self._get_authorization_base_url(self.server_url)
+            token_url = urljoin(auth_base_url, "/token")
+
+        token_data = {
+            "grant_type": "client_credentials",
+            "client_id": client_info.client_id,
+            "resource": self.resource,
+        }
+
+        if client_info.client_secret:
+            token_data["client_secret"] = client_info.client_secret
+
+        if self.client_metadata.scope:
+            token_data["scope"] = self.client_metadata.scope
+
+        async with httpx.AsyncClient() as client:
+            response = await client.post(
+                token_url,
+                data=token_data,
+                headers={"Content-Type": "application/x-www-form-urlencoded"},
+                timeout=30.0,
+            )
+
+        if response.status_code != 200:
+            raise Exception(f"Token request failed: {response.status_code} {response.text}")
+
+        token_response = OAuthToken.model_validate(response.json())
+        await self._validate_token_scopes(token_response)
+
+        if token_response.expires_in:
+            self._token_expiry_time = time.time() + token_response.expires_in
+        else:
+            self._token_expiry_time = None
+
+        await self.storage.set_tokens(token_response)
+        self._current_tokens = token_response
+
+    async def ensure_token(self) -> None:
+        async with self._token_lock:
+            if self._has_valid_token():
+                return
+            await self._request_token()
+
+    async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.Request, httpx.Response]:
+        if not self._has_valid_token():
+            await self.initialize()
+            await self.ensure_token()
+
+        if self._current_tokens and self._current_tokens.access_token:
+            request.headers["Authorization"] = f"Bearer {self._current_tokens.access_token}"
+
+        response = yield request
+
+        if response.status_code == 401:
+            self._current_tokens = None
+
+
+class TokenExchangeProvider(ClientCredentialsProvider):
+    """OAuth2 token exchange based on RFC 8693."""
+
+    def __init__(
+        self,
+        server_url: str,
+        client_metadata: OAuthClientMetadata,
+        storage: TokenStorage,
+        subject_token_supplier: Callable[[], Awaitable[str]],
+        subject_token_type: str = "access_token",
+        actor_token_supplier: Callable[[], Awaitable[str]] | None = None,
+        actor_token_type: str | None = None,
+        audience: str | None = None,
+        resource: str | None = None,
+        timeout: float = 300.0,
+    ):
+        """Create a new token exchange provider.
+
+        Parameters are forwarded to ClientCredentialsProvider for
+        client authentication. The resource parameter binds issued tokens to
+        the target resource, as defined by RFC 8707.
+        """
+
+        super().__init__(server_url, client_metadata, storage, resource, timeout)
+        self.subject_token_supplier = subject_token_supplier
+        self.subject_token_type = subject_token_type
+        self.actor_token_supplier = actor_token_supplier
+        self.actor_token_type = actor_token_type
+        self.audience = audience
+
+    async def _request_token(self) -> None:
+        if not self._metadata:
+            self._metadata = await self._discover_oauth_metadata(self.server_url)
+
+        client_info = await self._get_or_register_client()
+
+        if self._metadata and self._metadata.token_endpoint:
+            token_url = str(self._metadata.token_endpoint)
+        else:
+            auth_base_url = self._get_authorization_base_url(self.server_url)
+            token_url = urljoin(auth_base_url, "/token")
+
+        subject_token = await self.subject_token_supplier()
+        actor_token = await self.actor_token_supplier() if self.actor_token_supplier else None
+
+        token_data = {
+            "grant_type": "token_exchange",
+            "client_id": client_info.client_id,
+            "subject_token": subject_token,
+            "subject_token_type": self.subject_token_type,
+        }
+
+        if client_info.client_secret:
+            token_data["client_secret"] = client_info.client_secret
+
+        if actor_token:
+            token_data["actor_token"] = actor_token
+        if self.actor_token_type:
+            token_data["actor_token_type"] = self.actor_token_type
+        if self.audience:
+            token_data["audience"] = self.audience
+        if self.resource:
+            token_data["resource"] = self.resource
+        if self.client_metadata.scope:
+            token_data["scope"] = self.client_metadata.scope
+
+        async with httpx.AsyncClient() as client:
+            response = await client.post(
+                token_url,
+                data=token_data,
+                headers={"Content-Type": "application/x-www-form-urlencoded"},
+                timeout=30.0,
+            )
+
+        if response.status_code != 200:
+            raise Exception(f"Token request failed: {response.status_code} {response.text}")
+
+        token_response = OAuthToken.model_validate(response.json())
+        await self._validate_token_scopes(token_response)
+
+        if token_response.expires_in:
+            self._token_expiry_time = time.time() + token_response.expires_in
+        else:
+            self._token_expiry_time = None
+
+        await self.storage.set_tokens(token_response)
+        self._current_tokens = token_response

src/mcp/client/streamable_http.py

@@ -173,9 +173,11 @@ async def _handle_sse_event(
                 session_message = SessionMessage(message)
                 await read_stream_writer.send(session_message)
 
-                # Call resumption token callback if we have an ID
-                if sse.id and resumption_callback:
-                    await resumption_callback(sse.id)
+                # Call resumption token callback if we have an ID. Only update
+                # the resumption token on notifications to avoid overwriting it
+                # with the token from the final response.
+                if sse.id and resumption_callback and not isinstance(message.root, JSONRPCResponse | JSONRPCError):
+                    await resumption_callback(sse.id.strip())
 
                 # If this is a response or error return True indicating completion
                 # Otherwise, return False to continue listening
@@ -221,7 +223,7 @@ async def _handle_resumption_request(self, ctx: RequestContext) -> None:
         """Handle a resumption request using GET with SSE."""
         headers = self._prepare_request_headers(ctx.headers)
         if ctx.metadata and ctx.metadata.resumption_token:
-            headers[LAST_EVENT_ID] = ctx.metadata.resumption_token
+            headers[LAST_EVENT_ID] = ctx.metadata.resumption_token.strip()
         else:
             raise ResumptionError("Resumption request requires a resumption token")