Add support for get tokens method

Author: aravind-segu

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,18 @@ class RequestContext:
     sse_read_timeout: timedelta
 
 
+class AuthTokenProvider(Protocol):
+    """Protocol for providers that supply authentication tokens."""
+
+    async def get_token(self) -> str:
+        """Get an authentication token.
+
+        Returns:
+            str: The authentication token.
+        """
+        ...
+
+
 class StreamableHTTPTransport:
     """StreamableHTTP client transport implementation."""
 
@@ -83,6 +95,7 @@ def __init__(
         headers: dict[str, Any] | None = None,
         timeout: timedelta = timedelta(seconds=30),
         sse_read_timeout: timedelta = timedelta(seconds=60 * 5),
+        auth_token_provider: AuthTokenProvider | None = None,
     ) -> None:
         """Initialize the StreamableHTTP transport.
 
@@ -102,6 +115,7 @@ def __init__(
             CONTENT_TYPE: JSON,
             **self.headers,
         }
+        self.auth_token_provider = auth_token_provider
 
     def _update_headers_with_session(
         self, base_headers: dict[str, str]
@@ -112,6 +126,24 @@ def _update_headers_with_session(
             headers[MCP_SESSION_ID] = self.session_id
         return headers
 
+    async def _update_headers_with_token(
+        self, base_headers: dict[str, str]
+    ) -> dict[str, str]:
+        """Update headers with token if token provider is specified."""
+        if self.auth_token_provider is None:
+            return base_headers
+
+        token = await self.auth_token_provider.get_token()
+        headers = base_headers.copy()
+        headers["Authorization"] = f"Bearer {token}"
+        return 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_token(headers)
+        return headers
+
     def _is_initialization_request(self, message: JSONRPCMessage) -> bool:
         """Check if the message is an initialization request."""
         return (
@@ -184,7 +216,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 +238,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:
@@ -241,7 +273,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)
 
@@ -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_token_provider: AuthTokenProvider | None = None,
 ) -> AsyncGenerator[
     tuple[
         MemoryObjectReceiveStream[SessionMessage | Exception],
@@ -447,7 +480,9 @@ async def streamablehttp_client(
             - 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_token_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,61 @@ async def sampling_callback(
                 captured_message_params.messages[0].content.text
                 == "Server needs client sampling"
             )
+
+
+class MockAuthTokenProvider:
+    """Mock implementation of AuthTokenProvider for testing."""
+
+    def __init__(self, token: str):
+        self.token = token
+
+    async def get_token(self) -> str:
+        return self.token
+
+
+@pytest.mark.anyio
+async def test_auth_token_provider_headers(basic_server, basic_server_url):
+    """Test that auth token provider correctly sets Authorization header."""
+    # Create a mock token provider
+    token_provider = MockAuthTokenProvider("test-token-123")
+    token_provider.get_token = AsyncMock(return_value="test-token-123")
+
+    # Create client with token provider
+    async with streamablehttp_client(
+        f"{basic_server_url}/mcp", auth_token_provider=token_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
+
+    token_provider.get_token.assert_called()
+
+
+@pytest.mark.anyio
+async def test_auth_token_provider_token_update(basic_server, basic_server_url):
+    """Test that auth token provider can return different tokens."""
+    # Create a dynamic token provider
+    token_provider = MockAuthTokenProvider("test-token-123")
+    token_provider.get_token = AsyncMock(return_value="test-token-123")
+
+    # Create client with dynamic token provider
+    async with streamablehttp_client(
+        f"{basic_server_url}/mcp", auth_token_provider=token_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
+                await anyio.sleep(0.1)  # Small delay to ensure token updates
+
+    token_provider.get_token.call_count > 1