Add headers, timeout, and sse_read_timeout to MCPServerHTTP (#1218)

Co-authored-by: John Uiterwyk <[email protected]>
Co-authored-by: Marcelo Trylesinski <[email protected]>

Author: JohnUiterwyk

pydantic_ai_slim/pydantic_ai/mcp.py

@@ -188,11 +188,35 @@ async def main():
     For example for a server running locally, this might be `http://localhost:3001/sse`.
     """
 
+    headers: dict[str, Any] | None = None
+    """Optional HTTP headers to be sent with each request to the SSE endpoint.
+
+    These headers will be passed directly to the underlying `httpx.AsyncClient`.
+    Useful for authentication, custom headers, or other HTTP-specific configurations.
+    """
+
+    timeout: float = 5
+    """Initial connection timeout in seconds for establishing the SSE connection.
+
+    This timeout applies to the initial connection setup and handshake.
+    If the connection cannot be established within this time, the operation will fail.
+    """
+
+    sse_read_timeout: float = 60 * 5
+    """Maximum time in seconds to wait for new SSE messages before timing out.
+
+    This timeout applies to the long-lived SSE connection after it's established.
+    If no new messages are received within this time, the connection will be considered stale
+    and may be closed. Defaults to 5 minutes (300 seconds).
+    """
+
     @asynccontextmanager
     async def client_streams(
         self,
     ) -> AsyncIterator[
         tuple[MemoryObjectReceiveStream[JSONRPCMessage | Exception], MemoryObjectSendStream[JSONRPCMessage]]
     ]:  # pragma: no cover
-        async with sse_client(url=self.url) as (read_stream, write_stream):
+        async with sse_client(
+            url=self.url, headers=self.headers, timeout=self.timeout, sse_read_timeout=self.sse_read_timeout
+        ) as (read_stream, write_stream):
             yield read_stream, write_stream

tests/test_mcp.py

@@ -43,6 +43,19 @@ def test_sse_server():
     assert sse_server.url == 'http://localhost:8000/sse'
 
 
+def test_sse_server_with_header_and_timeout():
+    sse_server = MCPServerHTTP(
+        url='http://localhost:8000/sse',
+        headers={'my-custom-header': 'my-header-value'},
+        timeout=10,
+        sse_read_timeout=100,
+    )
+    assert sse_server.url == 'http://localhost:8000/sse'
+    assert sse_server.headers is not None and sse_server.headers['my-custom-header'] == 'my-header-value'
+    assert sse_server.timeout == 10
+    assert sse_server.sse_read_timeout == 100
+
+
 async def test_agent_with_stdio_server(allow_model_requests: None, openai_api_key: str):
     server = MCPServerStdio('python', ['-m', 'tests.mcp_server'])
     model = OpenAIModel('gpt-4o', provider=OpenAIProvider(api_key=openai_api_key))