Address comments

Author: aravind-segu

src/mcp/client/streamable_http.py

@@ -75,7 +75,9 @@ class RequestContext:
 
 
 class AuthTokenProvider(Protocol):
-    """Protocol for providers that supply authentication tokens."""
+    """Protocol that can be extended to implement custom client-to-server authentication
+    The get_token method is invoked before each request to the MCP Server to retrieve a
+    fresh authentication token and update the request headers."""
 
     async def get_token(self) -> str:
         """Get an authentication token.
@@ -129,8 +131,9 @@ def _update_headers_with_session(
     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:
+        """Update headers with token if token provider is specified and authorization
+        header is not present."""
+        if self.auth_token_provider is None or "Authorization" in base_headers:
             return base_headers
 
         token = await self.auth_token_provider.get_token()
@@ -474,6 +477,12 @@ 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_token_provider` is an optional protocol that can be extended to implement
+    custom client-to-server authentication. Before each request to the MCP Server,
+    the get_token method is invoked to retrieve a fresh authentication token and
+    update the request headers. Note that if the passed in headers already
+    contain an authorization header, this provider will not be called.
+
     Yields:
         Tuple containing:
             - read_stream: Stream for reading messages from the server

tests/shared/test_streamable_http.py

@@ -1279,6 +1279,32 @@ async def test_auth_token_provider_token_update(basic_server, basic_server_url):
             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
+
+
+@pytest.mark.anyio
+async def test_auth_token_provider_headers_not_overridden(
+    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,
+        headers={"Authorization": "test-token-123"},
+    ) 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_not_called()