Address comments

Author: aravind-segu

src/mcp/client/streamable_http.py

@@ -78,7 +78,7 @@ class AuthClientProvider(Protocol):
     """Base class that can be extended to implement custom client-to-server
     authentication"""
 
-    async def get_auth_headers(self) -> dict[str, str]:
+    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.
 
@@ -132,12 +132,12 @@ 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 the base_headers to
-        avoid overwriting existing Authorization headers"""
+        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_auth_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]:
@@ -476,10 +476,9 @@ async def streamablehttp_client(
 
     `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_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, that header will take precedence
-    over any tokens generated by this provider.
+    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:

tests/shared/test_streamable_http.py

@@ -1226,24 +1226,14 @@ async def sampling_callback(
             )
 
 
-class MockAuthClientProvider:
-    """Mock implementation of AuthClientProvider for testing."""
-
-    def __init__(self, token: str):
-        self.token = token
-
-    async def get_auth_headers(self) -> dict[str, str]:
-        return {"Authorization": "Bearer " + self.token}
-
-
 @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 = MockAuthClientProvider("test-token-123")
-    client_provider.get_auth_headers = AsyncMock(
-        return_value={"Authorization": "Bearer test-token-123"}
-    )
+    client_provider = AsyncMock()
+    client_provider.get_headers.return_value = {
+        "Authorization": "Bearer test-token-123"
+    }
 
     # Create client with token provider
     async with streamablehttp_client(
@@ -1258,17 +1248,17 @@ async def test_auth_client_provider_headers(basic_server, basic_server_url):
             tools = await session.list_tools()
             assert len(tools.tools) == 4
 
-    client_provider.get_auth_headers.assert_called()
+    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 = MockAuthClientProvider("test-token-123")
-    client_provider.get_auth_headers = AsyncMock(
-        return_value={"Authorization": "Bearer test-token-123"}
-    )
+    client_provider = AsyncMock()
+    client_provider.get_headers.return_value = {
+        "Authorization": "Bearer test-token-123"
+    }
 
     # Create client with dynamic token provider
     async with streamablehttp_client(
@@ -1284,4 +1274,6 @@ async def test_auth_client_provider_called_per_request(basic_server, basic_serve
                 tools = await session.list_tools()
                 assert len(tools.tools) == 4
 
-    client_provider.get_auth_headers.call_count > 1
+    # 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