Change implementation to take in auth headers

Author: aravind-segu

src/mcp/client/streamable_http.py

@@ -78,13 +78,12 @@ class AuthClientProvider(Protocol):
     """Base class that can be extended to implement custom client-to-server
     authentication"""
 
-    async def get_token(self) -> str:
-        """Get a token for authenticating to an MCP server. The token is assumed to
-                be short-lived; clients may call this API multiple times per
-            request to an MCP server.
+    async def get_auth_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:
-            str: The authentication token.
+            dict[str, str]: The authentication headers.
         """
         ...
 
@@ -129,23 +128,22 @@ def _update_headers_with_session(
             headers[MCP_SESSION_ID] = self.session_id
         return headers
 
-    async def _update_headers_with_token(
+    async def _update_headers_with_auth_headers(
         self, base_headers: dict[str, str]
     ) -> dict[str, str]:
-        """Update headers with token if token provider is specified and authorization
-        header is not present."""
-        if self.auth_client_provider is None or "Authorization" in base_headers:
+        """Update headers with auth_headers if auth client provider is specified.
+        The headers are merged giving precedence to the base_headers to
+        avoid overwriting existing Authorization headers"""
+        if self.auth_client_provider is None:
             return base_headers
 
-        token = await self.auth_client_provider.get_token()
-        headers = base_headers.copy()
-        headers["Authorization"] = f"Bearer {token}"
-        return headers
+        auth_headers = await self.auth_client_provider.get_auth_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_token(headers)
+        headers = await self._update_headers_with_auth_headers(headers)
         return headers
 
     def _is_initialization_request(self, message: JSONRPCMessage) -> bool:
@@ -252,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",
@@ -275,6 +272,16 @@ async def _handle_resumption_request(self, ctx: RequestContext) -> None:
                 if is_complete:
                     break
 
+    async def _is_testing_header_capture(self, response: httpx.Response) -> str | None:
+        try:
+            content = await response.aread()
+            if content.decode().startswith("[TESTING_HEADER_CAPTURE]"):
+                return content.decode()
+        except Exception as _:
+            return None
+
+        return None
+
     async def _handle_post_request(self, ctx: RequestContext) -> None:
         """Handle a POST request with response processing."""
         headers = await self._update_headers(ctx.headers)
@@ -299,12 +306,24 @@ async def _handle_post_request(self, ctx: RequestContext) -> None:
                     )
                 return
 
+            if response.status_code == 418:
+                test_error_message = await self._is_testing_header_capture(response)
+                # If this is coming from the test case return the response content
+                if test_error_message and isinstance(message.root, JSONRPCRequest):
+                    jsonrpc_error = JSONRPCError(
+                        jsonrpc="2.0",
+                        id=message.root.id,
+                        error=ErrorData(code=32600, message=test_error_message),
+                    )
+                    session_message = SessionMessage(JSONRPCMessage(jsonrpc_error))
+                    await ctx.read_stream_writer.send(session_message)
+                    return
+
             response.raise_for_status()
             if is_initialization:
                 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):

tests/shared/test_streamable_http.py

@@ -4,6 +4,7 @@
 Contains tests for both server and client sides of the StreamableHTTP transport.
 """
 
+import json
 import multiprocessing
 import socket
 import time
@@ -18,6 +19,8 @@
 import uvicorn
 from pydantic import AnyUrl
 from starlette.applications import Starlette
+from starlette.requests import Request
+from starlette.responses import Response
 from starlette.routing import Mount
 
 import mcp.types as types
@@ -244,8 +247,46 @@ def create_app(
     return app
 
 
+def create_header_capture_app() -> Starlette:
+    """Implement a minimal Starlette app that intercepts every request,
+    extracts its headers, and responds with status 418 (Test Status code),
+    embedding the captured headers as the JSON response body.
+    We use this server solely to verify that the MCP Server is forwarding
+    headers correctly."""
+
+    # Create a wrapper that captures headers and returns them in error response
+    async def header_capture_wrapper(scope, receive, send):
+        # Capture headers
+        request = Request(scope, receive=receive)
+        headers = dict(request.headers)
+
+        # Return error response with headers in body
+        response = Response(
+            "[TESTING_HEADER_CAPTURE]:" + json.dumps({"headers": headers}),
+            status_code=418,
+        )
+        await response(scope, receive, send)
+
+    # Create an ASGI application that uses our wrapper
+    app = Starlette(
+        debug=True,
+        routes=[
+            Mount("/mcp", app=header_capture_wrapper),
+        ],
+    )
+
+    return app
+
+
+def _get_captured_headrs(str) -> dict[str, str]:
+    return json.loads(str.split("[TESTING_HEADER_CAPTURE]:")[1])["headers"]
+
+
 def run_server(
-    port: int, is_json_response_enabled=False, event_store: EventStore | None = None
+    port: int,
+    is_json_response_enabled=False,
+    event_store: EventStore | None = None,
+    testing_header_capture: bool = False,
 ) -> None:
     """Run the test server.
 
@@ -255,7 +296,11 @@ def run_server(
         event_store: Optional event store for testing resumability.
     """
 
-    app = create_app(is_json_response_enabled, event_store)
+    if testing_header_capture:
+        app = create_header_capture_app()
+    else:
+        app = create_app(is_json_response_enabled, event_store)
+
     # Configure server
     config = uvicorn.Config(
         app=app,
@@ -296,33 +341,48 @@ def json_server_port() -> int:
         return s.getsockname()[1]
 
 
-@pytest.fixture
-def basic_server(basic_server_port: int) -> Generator[None, None, None]:
-    """Start a basic server."""
+def _start_basic_server(
+    basic_server_port: int, testing_header_capture: bool
+) -> Generator[None, None, None]:
     proc = multiprocessing.Process(
-        target=run_server, kwargs={"port": basic_server_port}, daemon=True
+        target=run_server,
+        kwargs={
+            "port": basic_server_port,
+            "testing_header_capture": testing_header_capture,
+        },
+        daemon=True,
     )
     proc.start()
 
     # Wait for server to be running
     max_attempts = 20
-    attempt = 0
-    while attempt < max_attempts:
+    for attempt in range(max_attempts):
         try:
             with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                 s.connect(("127.0.0.1", basic_server_port))
                 break
         except ConnectionRefusedError:
             time.sleep(0.1)
-            attempt += 1
     else:
         raise RuntimeError(f"Server failed to start after {max_attempts} attempts")
 
-    yield
+    try:
+        yield
+    finally:
+        proc.kill()
+        proc.join(timeout=2)
 
-    # Clean up
-    proc.kill()
-    proc.join(timeout=2)
+
+@pytest.fixture
+def basic_server(basic_server_port: int) -> Generator[None, None, None]:
+    yield from _start_basic_server(basic_server_port, testing_header_capture=False)
+
+
+@pytest.fixture
+def basic_server_with_header_capture(
+    basic_server_port: int,
+) -> Generator[None, None, None]:
+    yield from _start_basic_server(basic_server_port, testing_header_capture=True)
 
 
 @pytest.fixture
@@ -1232,79 +1292,84 @@ class MockAuthClientProvider:
     def __init__(self, token: str):
         self.token = token
 
-    async def get_token(self) -> str:
-        return self.token
+    async def get_auth_headers(self) -> dict[str, str]:
+        return {"Authorization": f"Bearer {self.token}"}
 
 
 @pytest.mark.anyio
-async def test_auth_client_provider_headers(basic_server, basic_server_url):
+async def test_auth_client_provider_headers(
+    basic_server_with_header_capture, basic_server_url
+):
     """Test that auth token provider correctly sets Authorization header."""
     # Create a mock token provider
-    client_provider = MockAuthClientProvider("test-token-123")
-    client_provider.get_token = AsyncMock(return_value="test-token-123")
+    client_provider = MockAuthClientProvider("short-lived-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_token.assert_called()
+            with pytest.raises(McpError) as mcpError:
+                _ = await session.initialize()
+                assert (
+                    _get_captured_headrs(mcpError.value.error.message)["Authorization"]
+                    == "Bearer short-lived-token-123"
+                )
 
 
 @pytest.mark.anyio
-async def test_auth_client_provider_token_update(basic_server, basic_server_url):
+async def test_auth_client_provider_token_called_on_every_request(
+    basic_server_with_header_capture, basic_server_url
+):
     """Test that auth token provider can return different tokens."""
     # Create a dynamic token provider
-    client_provider = MockAuthClientProvider("test-token-123")
-    client_provider.get_token = AsyncMock(return_value="test-token-123")
+    client_provider = MockAuthClientProvider("short-lived-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
+            with pytest.raises(McpError) as mcpError:
+                _ = await session.initialize()
+                assert (
+                    _get_captured_headrs(mcpError.value.error.message)["Authorization"]
+                    == "Bearer short-lived-token-123"
+                )
 
-    client_provider.get_token.call_count > 1
+            # Mock a new token and ensure the new token is returned
+            client_provider.get_auth_headers = AsyncMock(
+                return_value={"Authorization": "Bearer short-lived-token-456"}
+            )
+            with pytest.raises(McpError) as mcpError:
+                _ = await session.initialize()
+                assert (
+                    _get_captured_headrs(mcpError.value.error.message)["Authorization"]
+                    == "Bearer short-lived-token-456"
+                )
 
 
 @pytest.mark.anyio
 async def test_auth_client_provider_headers_not_overridden(
-    basic_server, basic_server_url
+    basic_server_with_header_capture, basic_server_url
 ):
-    """Test that auth token provider correctly sets Authorization header."""
+    """Test that provided headers override auth client provider headers."""
     # Create a mock token provider
-    client_provider = MockAuthClientProvider("test-token-123")
-    client_provider.get_token = AsyncMock(return_value="test-token-123")
+    client_provider = MockAuthClientProvider("short-lived-token")
 
-    # Create client with token provider
+    # Create client with token provider and custom headers
+    custom_headers = {"Authorization": "Bearer original-long-lived-token"}
     async with streamablehttp_client(
         f"{basic_server_url}/mcp",
         auth_client_provider=client_provider,
-        headers={"Authorization": "test-token-123"},
+        headers=custom_headers,
     ) 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_token.assert_not_called()
+            # Original token is used and not short-lived-token from the provider
+            with pytest.raises(McpError) as mcpError:
+                _ = await session.initialize()
+                assert (
+                    _get_captured_headrs(mcpError.value.error.message)["Authorization"]
+                    == "Bearer original-long-lived-token"
+                )