Merge pull request #36 from jmolz/fix/sse-header-forwarding

fix: forward headers for SSE MCP clients 🤖🤖🤖

Author: rdasilveiracabral

src/nooa/mcp/client.py

@@ -67,15 +67,18 @@ class MCPSSEClient(MCPBaseClient):
     Args:
         url: The URL of the MCP server
         tool_call_timeout: Timeout for tool calls
+        headers: Optional custom HTTP headers to include in requests (e.g., for authentication)
     """
 
     def __init__(
         self,
         url: str,
         tool_call_timeout: timedelta = timedelta(seconds=60),
+        headers: dict[str, str] | None = None,
     ):
         super().__init__(tool_call_timeout=tool_call_timeout)
         self._url = url
+        self._headers = headers or {}
 
     @property
     def transport(self) -> Literal["sse", "stdio", "streamable-http"]:
@@ -87,6 +90,7 @@ def server_config(self) -> dict[str, Any]:
         """Return the server configuration."""
         return {
             "url": self._url,
+            "headers": self._headers,
             "transport": self.transport,
         }
 
@@ -95,6 +99,11 @@ def url(self) -> str:
         """Return the server URL."""
         return self._url
 
+    @property
+    def headers(self) -> dict[str, str]:
+        """Return the custom headers configured for this client."""
+        return self._headers
+
     @asynccontextmanager
     @override
     async def connect_to_server(self):
@@ -107,7 +116,10 @@ async def connect_to_server(self):
             RuntimeError: If session initialization fails (MCP protocol error)
         """
         async with (
-            sse_client(url=self._url) as (read, write),
+            sse_client(
+                url=self._url,
+                headers=self._headers if self._headers else None,
+            ) as (read, write),
             ClientSession(read, write) as session,
         ):
             await session.initialize()
@@ -320,7 +332,7 @@ def create_mcp_client(
         case "sse":
             if url is None:
                 raise ValueError("url must be provided for sse transport")
-            return MCPSSEClient(url=url)
+            return MCPSSEClient(url=url, headers=headers)
         case "streamable-http":
             if url is None:
                 raise ValueError("url must be provided for streamable-http transport")

tests/test_mcp/test_client.py

@@ -45,6 +45,7 @@ def sse_client() -> MCPSSEClient:
     """Create an SSE client for testing."""
     return MCPSSEClient(
         url="http://localhost:8000",
+        headers={"Authorization": "Bearer token"},
         tool_call_timeout=timedelta(seconds=45),
     )
 
@@ -253,11 +254,24 @@ def test_sse_client_properties(sse_client: MCPSSEClient):
     """MCPSSEClient properties return correct values."""
     assert sse_client.transport == "sse"
     assert sse_client.url == "http://localhost:8000"
+    assert sse_client.headers == {"Authorization": "Bearer token"}
     assert sse_client.tool_call_timeout == timedelta(seconds=45)
 
     config = sse_client.server_config
     assert config["transport"] == "sse"
     assert config["url"] == "http://localhost:8000"
+    assert config["headers"] == {"Authorization": "Bearer token"}
+
+
+def test_sse_client_preserves_positional_timeout():
+    """The existing second positional argument remains the tool-call timeout."""
+    client = MCPSSEClient(
+        "http://localhost:8000",
+        timedelta(seconds=7),
+    )
+
+    assert client.tool_call_timeout == timedelta(seconds=7)
+    assert client.headers == {}
 
 
 def test_streamable_http_client_properties(streamable_http_client: MCPStreamableHTTPClient):
@@ -328,6 +342,35 @@ async def test_connect_context_manager(
                 mock_client_session.initialize.assert_awaited_once()
 
 
+@pytest.mark.asyncio
+async def test_sse_headers_passed_to_transport(
+    mock_mcp_transport: tuple[MagicMock, MagicMock],
+    mock_client_session: AsyncMock,
+):
+    """SSE clients forward custom headers to the transport."""
+    client = create_mcp_client(
+        "sse",
+        url="https://example.test/sse",
+        headers={"Authorization": "Bearer token"},
+    )
+    mock_read, mock_write = mock_mcp_transport
+
+    with (
+        patch("nooa.mcp.client.sse_client") as mock_sse,
+        patch("nooa.mcp.client.ClientSession") as mock_session_class,
+    ):
+        mock_sse.return_value.__aenter__.return_value = (mock_read, mock_write)
+        mock_session_class.return_value.__aenter__.return_value = mock_client_session
+
+        async with client.connect_to_server():
+            pass
+
+    mock_sse.assert_called_once_with(
+        url="https://example.test/sse",
+        headers={"Authorization": "Bearer token"},
+    )
+
+
 @pytest.mark.asyncio
 async def test_streamable_http_connect_context_manager(
     streamable_http_client: MCPStreamableHTTPClient,