clean up

Author: ihrpr

src/mcp/shared/httpx_utils.py

@@ -1,32 +1,40 @@
 """Utilities for creating standardized httpx AsyncClient instances."""
 
+from __future__ import annotations
+
 from typing import Any
 
 import httpx
 
+__all__ = ["create_mcp_http_client"]
+
 
 def create_mcp_http_client(
     *,
     headers: dict[str, Any] | None = None,
-    timeout: httpx.Timeout | float | None = None,
+    timeout: httpx.Timeout | None = None,
     **kwargs: Any,
 ) -> httpx.AsyncClient:
     """Create a standardized httpx AsyncClient with MCP defaults.
 
     This function provides common defaults used throughout the MCP codebase:
     - follow_redirects=True (always enabled)
     - Default timeout of 30 seconds if not specified
-    - Header will be merged
+    - Headers will be merged with any existing headers in kwargs
 
     Args:
         headers: Optional headers to include with all requests.
-        timeout: Request timeout in seconds (float) or httpx.Timeout object.
+        timeout: Request timeout as httpx.Timeout object.
             Defaults to 30 seconds if not specified.
         **kwargs: Additional keyword arguments to pass to AsyncClient.
 
     Returns:
         Configured httpx.AsyncClient instance with MCP defaults.
 
+    Note:
+        The returned AsyncClient must be used as a context manager to ensure
+        proper cleanup of connections.
+
     Examples:
         # Basic usage with MCP defaults
         async with create_mcp_http_client() as client:
@@ -50,16 +58,16 @@ def create_mcp_http_client(
     # Handle timeout
     if timeout is None:
         defaults["timeout"] = httpx.Timeout(30.0)
-    elif isinstance(timeout, int | float):
-        defaults["timeout"] = httpx.Timeout(timeout)
     else:
         defaults["timeout"] = timeout
 
-    # Handle headers
+    # Handle headers with proper merging
     if headers is not None:
-        kwargs["headers"] = headers
+        existing_headers = kwargs.get("headers", {})
+        merged_headers = {**existing_headers, **headers}
+        kwargs["headers"] = merged_headers
 
-    # Merge defaults with provided kwargs
-    kwargs = {**defaults, **kwargs}
+    # Merge kwargs with defaults (defaults take precedence)
+    kwargs = {**kwargs, **defaults}
 
     return httpx.AsyncClient(**kwargs)

tests/shared/test_httpx_utils.py

@@ -0,0 +1,46 @@
+"""Tests for httpx utility functions."""
+
+import httpx
+
+from mcp.shared.httpx_utils import create_mcp_http_client
+
+
+class TestCreateMcpHttpClient:
+    """Test create_mcp_http_client function."""
+
+    def test_default_settings(self):
+        """Test that default settings are applied correctly."""
+        client = create_mcp_http_client()
+        
+        # Check follow_redirects is True
+        assert client.follow_redirects is True
+        
+        # Check default timeout is 30 seconds
+        assert client.timeout.connect == 30.0
+        assert client.timeout.read == 30.0
+        assert client.timeout.write == 30.0
+        assert client.timeout.pool == 30.0
+
+    def test_custom_parameters(self):
+        """Test custom headers and timeout are set correctly."""
+        headers = {"Authorization": "Bearer token", "X-Custom": "value"}
+        timeout = httpx.Timeout(connect=5.0, read=10.0, write=15.0, pool=20.0)
+        
+        client = create_mcp_http_client(headers=headers, timeout=timeout)
+        
+        # Check headers
+        assert client.headers["Authorization"] == "Bearer token"
+        assert client.headers["X-Custom"] == "value"
+        
+        # Check custom timeout
+        assert client.timeout.connect == 5.0
+        assert client.timeout.read == 10.0
+        assert client.timeout.write == 15.0
+        assert client.timeout.pool == 20.0
+
+    def test_follow_redirects_enforced(self):
+        """Test follow_redirects is always True even if False is passed."""
+        client = create_mcp_http_client(follow_redirects=False)
+        
+        # Should still be True because our defaults override user input
+        assert client.follow_redirects is True