Adapt to PR #1177

Author: OS-anasantos

src/mcp/client/streamable_http.py

@@ -464,6 +464,7 @@ def get_session_id(self) -> str | None:
 @asynccontextmanager
 async def streamable_http_client(
     url: str,
+    extensions: dict[str, str] | None = None,
     *,
     http_client: httpx.AsyncClient | None = None,
     terminate_on_close: bool = True,
@@ -480,6 +481,7 @@ async def streamable_http_client(
 
     Args:
         url: The MCP server endpoint URL.
+        extensions: Optional extensions to include in requests.
         http_client: Optional pre-configured httpx.AsyncClient. If None, a default
             client with recommended MCP timeouts will be created. To configure headers,
             authentication, or other HTTP settings, create an httpx.AsyncClient and pass it here.
@@ -515,7 +517,7 @@ async def streamable_http_client(
     auth = client.auth
 
     # Create transport with extracted configuration
-    transport = StreamableHTTPTransport(url, headers_dict, timeout, sse_read_timeout, auth)
+    transport = StreamableHTTPTransport(url, headers_dict, extensions, timeout, sse_read_timeout, auth)
 
     async with anyio.create_task_group() as tg:
         try:

tests/shared/test_streamable_http.py

@@ -1799,7 +1799,7 @@ async def test_extensions_passed_to_streamablehttp_client(self, basic_server: No
             "custom_metadata": "custom_data",
         }
 
-        async with streamablehttp_client(f"{basic_server_url}/mcp", extensions=test_extensions) as (
+        async with streamable_http_client(f"{basic_server_url}/mcp", extensions=test_extensions) as (
             read_stream,
             write_stream,
             _,
@@ -1817,15 +1817,15 @@ async def test_extensions_passed_to_streamablehttp_client(self, basic_server: No
     @pytest.mark.anyio
     async def test_extensions_with_empty_dict(self, basic_server: None, basic_server_url: str):
         """Test streamablehttp_client with empty extensions dict."""
-        async with streamablehttp_client(f"{basic_server_url}/mcp", extensions={}) as (read_stream, write_stream, _):
+        async with streamable_http_client(f"{basic_server_url}/mcp", extensions={}) as (read_stream, write_stream, _):
             async with ClientSession(read_stream, write_stream) as session:
                 result = await session.initialize()
                 assert isinstance(result, InitializeResult)
 
     @pytest.mark.anyio
     async def test_extensions_with_none(self, basic_server: None, basic_server_url: str):
         """Test streamablehttp_client with None extensions."""
-        async with streamablehttp_client(f"{basic_server_url}/mcp", extensions=None) as (read_stream, write_stream, _):
+        async with streamable_http_client(f"{basic_server_url}/mcp", extensions=None) as (read_stream, write_stream, _):
             async with ClientSession(read_stream, write_stream) as session:
                 result = await session.initialize()
                 assert isinstance(result, InitializeResult)
@@ -1887,7 +1887,7 @@ async def test_extensions_isolation_between_clients(self, basic_server: None, ba
         # Create two clients with different extensions
         results: list[tuple[str, str]] = []
 
-        async with streamablehttp_client(f"{basic_server_url}/mcp", extensions=extensions_1) as (
+        async with streamable_http_client(f"{basic_server_url}/mcp", extensions=extensions_1) as (
             read_stream1,
             write_stream1,
             _,
@@ -1896,7 +1896,7 @@ async def test_extensions_isolation_between_clients(self, basic_server: None, ba
                 result1 = await session1.initialize()
                 results.append(("client1", result1.serverInfo.name))
 
-        async with streamablehttp_client(f"{basic_server_url}/mcp", extensions=extensions_2) as (
+        async with streamable_http_client(f"{basic_server_url}/mcp", extensions=extensions_2) as (
             read_stream2,
             write_stream2,
             _,
@@ -1950,18 +1950,11 @@ async def stream(self, *args: Any, **kwargs: Any):
                 async with super().stream(*args, **kwargs) as response:
                     yield response
 
-        # Custom client factory that returns our capturing client
-        def custom_client_factory(
-            headers: dict[str, str] | None = None, timeout: httpx.Timeout | None = None, auth: httpx.Auth | None = None
-        ) -> httpx.AsyncClient:
-            return ExtensionCapturingClient(
-                headers=headers,
-                timeout=timeout,
-                auth=auth,
-            )
-
-        async with streamablehttp_client(
-            f"{basic_server_url}/mcp/", extensions=test_extensions, httpx_client_factory=custom_client_factory
+        # Create the custom client that will capture extensions
+        custom_client = ExtensionCapturingClient()
+
+        async with streamable_http_client(
+            f"{basic_server_url}/mcp/", extensions=test_extensions, http_client=custom_client
         ) as (read_stream, write_stream, _):
             async with ClientSession(read_stream, write_stream) as session:
                 # Initialize - this should make a POST request with extensions
@@ -1970,6 +1963,9 @@ def custom_client_factory(
                 # Make another request to capture more extensions usage
                 await session.list_tools()
 
+        # Close the custom client
+        await custom_client.aclose()
+
         # Verify extensions were captured in requests
         assert len(captured_extensions) > 0
 
@@ -1986,7 +1982,7 @@ async def test_extensions_with_json_and_sse_responses(self, basic_server: None,
         test_extensions = {"response_test": "json_sse_test", "format": "both"}
 
         # Test with regular SSE response (default behavior)
-        async with streamablehttp_client(f"{basic_server_url}/mcp", extensions=test_extensions) as (
+        async with streamable_http_client(f"{basic_server_url}/mcp", extensions=test_extensions) as (
             read_stream,
             write_stream,
             _,
@@ -2010,7 +2006,7 @@ async def test_extensions_with_json_response_server(self, json_response_server:
         """Test extensions work with JSON response mode."""
         test_extensions = {"response_mode": "json_only", "test_id": "json_test_123"}
 
-        async with streamablehttp_client(f"{json_server_url}/mcp", extensions=test_extensions) as (
+        async with streamable_http_client(f"{json_server_url}/mcp", extensions=test_extensions) as (
             read_stream,
             write_stream,
             _,
@@ -2049,7 +2045,7 @@ async def test_extensions_with_special_characters(self, basic_server: None, basi
             "url_like": "https://example.com/path?param=value",
         }
 
-        async with streamablehttp_client(f"{basic_server_url}/mcp", extensions=test_extensions) as (
+        async with streamable_http_client(f"{basic_server_url}/mcp", extensions=test_extensions) as (
             read_stream,
             write_stream,
             _,
@@ -2061,4 +2057,4 @@ async def test_extensions_with_special_characters(self, basic_server: None, basi
 
                 # Should work normally with tools
                 tools = await session.list_tools()
-                assert len(tools.tools) == 6
+                assert len(tools.tools) == 6