Dynamic Authorization in Streamable HTTP Client

src/mcp/client/streamable_http.py

@@ -11,7 +11,7 @@
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
 from datetime import timedelta
-from typing import Any
+from typing import Any, Protocol
 
 import anyio
 import httpx
@@ -74,6 +74,20 @@ class RequestContext:
     sse_read_timeout: timedelta
 
 
+class AuthClientProvider(Protocol):
+    """Base class that can be extended to implement custom client-to-server
+    authentication"""
+
+    async def get_headers(self) -> dict[str, str]:
+        """Gets auth headers for authenticating to an MCP server.
+        Clients may call this API multiple times per request to an MCP server.
+
+        Returns:
+            dict[str, str]: The authentication headers.
+        """
+        ...
+
+
 class StreamableHTTPTransport:
     """StreamableHTTP client transport implementation."""
 
@@ -83,6 +97,7 @@ def __init__(
         headers: dict[str, Any] | None = None,
         timeout: timedelta = timedelta(seconds=30),
         sse_read_timeout: timedelta = timedelta(seconds=60 * 5),
+        auth_client_provider: AuthClientProvider | None = None,
     ) -> None:
         """Initialize the StreamableHTTP transport.
 
@@ -102,6 +117,7 @@ def __init__(
             CONTENT_TYPE: JSON,
             **self.headers,
         }
+        self.auth_client_provider = auth_client_provider
 
     def _update_headers_with_session(
         self, base_headers: dict[str, str]
@@ -112,6 +128,24 @@ def _update_headers_with_session(
             headers[MCP_SESSION_ID] = self.session_id
         return headers
 
+    async def _update_headers_with_auth_headers(
+        self, base_headers: dict[str, str]
+    ) -> dict[str, str]:
+        """Update headers with auth_headers if auth client provider is specified.
+        The headers are merged, giving precedence to any headers already
+        specified in base_headers"""
+        if self.auth_client_provider is None:
+            return base_headers
+
+        auth_headers = await self.auth_client_provider.get_headers()
+        return {**auth_headers, **base_headers}
+
+    async def _update_headers(self, base_headers: dict[str, str]) -> dict[str, str]:
+        """Update headers with session ID and token if available."""
+        headers = self._update_headers_with_session(base_headers)
+        headers = await self._update_headers_with_auth_headers(headers)
+        return headers
+
     def _is_initialization_request(self, message: JSONRPCMessage) -> bool:
         """Check if the message is an initialization request."""
         return (
@@ -184,7 +218,7 @@ async def handle_get_stream(
             if not self.session_id:
                 return
 
-            headers = self._update_headers_with_session(self.request_headers)
+            headers = await self._update_headers(self.request_headers)
 
             async with aconnect_sse(
                 client,
@@ -206,7 +240,7 @@ async def handle_get_stream(
 
     async def _handle_resumption_request(self, ctx: RequestContext) -> None:
         """Handle a resumption request using GET with SSE."""
-        headers = self._update_headers_with_session(ctx.headers)
+        headers = await self._update_headers(ctx.headers)
         if ctx.metadata and ctx.metadata.resumption_token:
             headers[LAST_EVENT_ID] = ctx.metadata.resumption_token
         else:
@@ -216,7 +250,6 @@ async def _handle_resumption_request(self, ctx: RequestContext) -> None:
         original_request_id = None
         if isinstance(ctx.session_message.message.root, JSONRPCRequest):
             original_request_id = ctx.session_message.message.root.id
-
         async with aconnect_sse(
             ctx.client,
             "GET",
@@ -241,7 +274,7 @@ async def _handle_resumption_request(self, ctx: RequestContext) -> None:
 
     async def _handle_post_request(self, ctx: RequestContext) -> None:
         """Handle a POST request with response processing."""
-        headers = self._update_headers_with_session(ctx.headers)
+        headers = await self._update_headers(ctx.headers)
         message = ctx.session_message.message
         is_initialization = self._is_initialization_request(message)
 
@@ -268,7 +301,6 @@ async def _handle_post_request(self, ctx: RequestContext) -> None:
                 self._maybe_extract_session_id_from_response(response)
 
             content_type = response.headers.get(CONTENT_TYPE, "").lower()
-
             if content_type.startswith(JSON):
                 await self._handle_json_response(response, ctx.read_stream_writer)
             elif content_type.startswith(SSE):
@@ -405,7 +437,7 @@ async def terminate_session(self, client: httpx.AsyncClient) -> None:
             return
 
         try:
-            headers = self._update_headers_with_session(self.request_headers)
+            headers = await self._update_headers(self.request_headers)
             response = await client.delete(self.url, headers=headers)
 
             if response.status_code == 405:
@@ -427,6 +459,7 @@ async def streamablehttp_client(
     timeout: timedelta = timedelta(seconds=30),
     sse_read_timeout: timedelta = timedelta(seconds=60 * 5),
     terminate_on_close: bool = True,
+    auth_client_provider: AuthClientProvider | None = None,
 ) -> AsyncGenerator[
     tuple[
         MemoryObjectReceiveStream[SessionMessage | Exception],
@@ -441,13 +474,21 @@ async def streamablehttp_client(
     `sse_read_timeout` determines how long (in seconds) the client will wait for a new
     event before disconnecting. All other HTTP operations are controlled by `timeout`.
 
+    `auth_client_provider` instance of `AuthClientProvider` that can be passed to
+    support client-to-server authentication. Before each request to the MCP Server,
+    the auth_client_provider.get_headers() method is invoked to retrieve headers
+    for authentication. Note that any headers already specified in `headers`
+    will take precedence over headers returned by auth_client_provider.get_headers()
+
     Yields:
         Tuple containing:
             - read_stream: Stream for reading messages from the server
             - write_stream: Stream for sending messages to the server
             - get_session_id_callback: Function to retrieve the current session ID
     """
-    transport = StreamableHTTPTransport(url, headers, timeout, sse_read_timeout)
+    transport = StreamableHTTPTransport(
+        url, headers, timeout, sse_read_timeout, auth_client_provider
+    )
 
     read_stream_writer, read_stream = anyio.create_memory_object_stream[
         SessionMessage | Exception

tests/shared/test_streamable_http.py

@@ -9,6 +9,7 @@
 import time
 from collections.abc import Generator
 from typing import Any
+from unittest.mock import AsyncMock
 
 import anyio
 import httpx
@@ -1223,3 +1224,56 @@ async def sampling_callback(
                 captured_message_params.messages[0].content.text
                 == "Server needs client sampling"
             )
+
+
+@pytest.mark.anyio
+async def test_auth_client_provider_headers(basic_server, basic_server_url):
+    """Test that auth token provider correctly sets Authorization header."""
+    # Create a mock token provider
+    client_provider = AsyncMock()
+    client_provider.get_headers.return_value = {
+        "Authorization": "Bearer test-token-123"
+    }
+
+    # Create client with token provider
+    async with streamablehttp_client(
+        f"{basic_server_url}/mcp", auth_client_provider=client_provider
+    ) as (read_stream, write_stream, _):
+        async with ClientSession(read_stream, write_stream) as session:
+            # Initialize the session
+            result = await session.initialize()
+            assert isinstance(result, InitializeResult)
+
+            # Make a request to verify headers
+            tools = await session.list_tools()
+            assert len(tools.tools) == 4
+
+    client_provider.get_headers.assert_called()
+
+
+@pytest.mark.anyio
+async def test_auth_client_provider_called_per_request(basic_server, basic_server_url):
+    """Test that auth token provider can return different tokens."""
+    # Create a dynamic token provider
+    client_provider = AsyncMock()
+    client_provider.get_headers.return_value = {
+        "Authorization": "Bearer test-token-123"
+    }
+
+    # Create client with dynamic token provider
+    async with streamablehttp_client(
+        f"{basic_server_url}/mcp", auth_client_provider=client_provider
+    ) as (read_stream, write_stream, _):
+        async with ClientSession(read_stream, write_stream) as session:
+            # Initialize the session
+            result = await session.initialize()
+            assert isinstance(result, InitializeResult)
+
+            # Make multiple requests to verify token updates
+            for i in range(3):
+                tools = await session.list_tools()
+                assert len(tools.tools) == 4
+
+    # list_tools is called 3 times, but get_auth_headers is also used during
+    # session initialization and setup. Verify it's called at least 3 times.
+    assert client_provider.get_headers.call_count > 3