Skip to content

Commit 4e99506

Browse files
jjcampanajjcampana
committed
add mcp param and tests
1 parent 14d7d59 commit 4e99506

File tree

2 files changed

+235
-3
lines changed

2 files changed

+235
-3
lines changed

‎src/agents/mcp/server.py‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
from mcp import ClientSession, StdioServerParameters, Tool as MCPTool, stdio_client
1414
from mcp.client.sse import sse_client
1515
from mcp.client.streamable_http import GetSessionIdCallback, streamablehttp_client
16+
from mcp.shared._httpx_utils import McpHttpClientFactory
1617
from mcp.shared.message import SessionMessage
1718
from mcp.types import CallToolResult, GetPromptResult, InitializeResult, ListPromptsResult
1819
from typing_extensions import NotRequired, TypedDict
@@ -575,6 +576,9 @@ class MCPServerStreamableHttpParams(TypedDict):
575576
terminate_on_close: NotRequired[bool]
576577
"""Terminate on close"""
577578

579+
httpx_client_factory: NotRequired[McpHttpClientFactory]
580+
"""Custom HTTP client factory for configuring httpx.AsyncClient behavior."""
581+
578582

579583
class MCPServerStreamableHttp(_MCPServerWithClientSession):
580584
"""MCP server implementation that uses the Streamable HTTP transport. See the [spec]
@@ -597,9 +601,9 @@ def __init__(
597601
598602
Args:
599603
params: The params that configure the server. This includes the URL of the server,
600-
the headers to send to the server, the timeout for the HTTP request, and the
601-
timeout for the Streamable HTTP connection and whether we need to
602-
terminate on close.
604+
the headers to send to the server, the timeout for the HTTP request, the
605+
timeout for the Streamable HTTP connection, whether we need to
606+
terminate on close, and an optional custom HTTP client factory.
603607
604608
cache_tools_list: Whether to cache the tools list. If `True`, the tools list will be
605609
cached and only fetched from the server once. If `False`, the tools list will be
@@ -651,6 +655,7 @@ def create_streams(
651655
timeout=self.params.get("timeout", 5),
652656
sse_read_timeout=self.params.get("sse_read_timeout", 60 * 5),
653657
terminate_on_close=self.params.get("terminate_on_close", True),
658+
httpx_client_factory=self.params.get("httpx_client_factory", None),
654659
)
655660

656661
@property

‎tests/mcp/test_streamable_http_client_factory.py‎

Lines changed: 227 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,227 @@
1+
"""Tests for MCPServerStreamableHttp httpx_client_factory functionality."""
2+
3+
from unittest.mock import MagicMock, patch
4+
5+
import httpx
6+
import pytest
7+
8+
from agents.mcp import MCPServerStreamableHttp
9+
10+
11+
class TestMCPServerStreamableHttpClientFactory:
12+
"""Test cases for custom httpx_client_factory parameter."""
13+
14+
@pytest.mark.asyncio
15+
async def test_default_httpx_client_factory(self):
16+
"""Test that default behavior works when no custom factory is provided."""
17+
# Mock the streamablehttp_client to avoid actual network calls
18+
with patch("agents.mcp.server.streamablehttp_client") as mock_client:
19+
mock_client.return_value = MagicMock()
20+
21+
server = MCPServerStreamableHttp(
22+
params={
23+
"url": "http://localhost:8000/mcp",
24+
"headers": {"Authorization": "Bearer token"},
25+
"timeout": 10,
26+
}
27+
)
28+
29+
# Create streams should not pass httpx_client_factory when not provided
30+
server.create_streams()
31+
32+
# Verify streamablehttp_client was called with correct parameters
33+
mock_client.assert_called_once_with(
34+
url="http://localhost:8000/mcp",
35+
headers={"Authorization": "Bearer token"},
36+
timeout=10,
37+
sse_read_timeout=300, # Default value
38+
terminate_on_close=True, # Default value
39+
httpx_client_factory=None, # Should be None when not provided
40+
)
41+
42+
@pytest.mark.asyncio
43+
async def test_custom_httpx_client_factory(self):
44+
"""Test that custom httpx_client_factory is passed correctly."""
45+
46+
# Create a custom factory function
47+
def custom_factory() -> httpx.AsyncClient:
48+
return httpx.AsyncClient(
49+
verify=False, # Disable SSL verification for testing
50+
timeout=httpx.Timeout(60.0),
51+
headers={"X-Custom-Header": "test"},
52+
)
53+
54+
# Mock the streamablehttp_client to avoid actual network calls
55+
with patch("agents.mcp.server.streamablehttp_client") as mock_client:
56+
mock_client.return_value = MagicMock()
57+
58+
server = MCPServerStreamableHttp(
59+
params={
60+
"url": "http://localhost:8000/mcp",
61+
"headers": {"Authorization": "Bearer token"},
62+
"timeout": 10,
63+
"httpx_client_factory": custom_factory,
64+
}
65+
)
66+
67+
# Create streams should pass the custom factory
68+
server.create_streams()
69+
70+
# Verify streamablehttp_client was called with the custom factory
71+
mock_client.assert_called_once_with(
72+
url="http://localhost:8000/mcp",
73+
headers={"Authorization": "Bearer token"},
74+
timeout=10,
75+
sse_read_timeout=300, # Default value
76+
terminate_on_close=True, # Default value
77+
httpx_client_factory=custom_factory,
78+
)
79+
80+
@pytest.mark.asyncio
81+
async def test_custom_httpx_client_factory_with_ssl_cert(self):
82+
"""Test custom factory with SSL certificate configuration."""
83+
84+
def ssl_cert_factory() -> httpx.AsyncClient:
85+
return httpx.AsyncClient(
86+
verify="/path/to/cert.pem", # Custom SSL certificate
87+
timeout=httpx.Timeout(120.0),
88+
)
89+
90+
with patch("agents.mcp.server.streamablehttp_client") as mock_client:
91+
mock_client.return_value = MagicMock()
92+
93+
server = MCPServerStreamableHttp(
94+
params={
95+
"url": "https://secure-server.com/mcp",
96+
"timeout": 30,
97+
"httpx_client_factory": ssl_cert_factory,
98+
}
99+
)
100+
101+
server.create_streams()
102+
103+
mock_client.assert_called_once_with(
104+
url="https://secure-server.com/mcp",
105+
headers=None,
106+
timeout=30,
107+
sse_read_timeout=300,
108+
terminate_on_close=True,
109+
httpx_client_factory=ssl_cert_factory,
110+
)
111+
112+
@pytest.mark.asyncio
113+
async def test_custom_httpx_client_factory_with_proxy(self):
114+
"""Test custom factory with proxy configuration."""
115+
116+
def proxy_factory() -> httpx.AsyncClient:
117+
return httpx.AsyncClient(
118+
proxies="http://proxy.example.com:8080",
119+
timeout=httpx.Timeout(60.0),
120+
)
121+
122+
with patch("agents.mcp.server.streamablehttp_client") as mock_client:
123+
mock_client.return_value = MagicMock()
124+
125+
server = MCPServerStreamableHttp(
126+
params={
127+
"url": "http://localhost:8000/mcp",
128+
"httpx_client_factory": proxy_factory,
129+
}
130+
)
131+
132+
server.create_streams()
133+
134+
mock_client.assert_called_once_with(
135+
url="http://localhost:8000/mcp",
136+
headers=None,
137+
timeout=5, # Default value
138+
sse_read_timeout=300,
139+
terminate_on_close=True,
140+
httpx_client_factory=proxy_factory,
141+
)
142+
143+
@pytest.mark.asyncio
144+
async def test_custom_httpx_client_factory_with_retry_logic(self):
145+
"""Test custom factory with retry logic configuration."""
146+
147+
def retry_factory() -> httpx.AsyncClient:
148+
return httpx.AsyncClient(
149+
timeout=httpx.Timeout(30.0),
150+
# Note: httpx doesn't have built-in retry, but this shows how
151+
# a custom factory could be used to configure retry behavior
152+
# through middleware or other mechanisms
153+
)
154+
155+
with patch("agents.mcp.server.streamablehttp_client") as mock_client:
156+
mock_client.return_value = MagicMock()
157+
158+
server = MCPServerStreamableHttp(
159+
params={
160+
"url": "http://localhost:8000/mcp",
161+
"httpx_client_factory": retry_factory,
162+
}
163+
)
164+
165+
server.create_streams()
166+
167+
mock_client.assert_called_once_with(
168+
url="http://localhost:8000/mcp",
169+
headers=None,
170+
timeout=5,
171+
sse_read_timeout=300,
172+
terminate_on_close=True,
173+
httpx_client_factory=retry_factory,
174+
)
175+
176+
def test_httpx_client_factory_type_annotation(self):
177+
"""Test that the type annotation is correct for httpx_client_factory."""
178+
from agents.mcp.server import MCPServerStreamableHttpParams
179+
180+
# This test ensures the type annotation is properly set
181+
# We can't easily test the TypedDict at runtime, but we can verify
182+
# that the import works and the type is available
183+
assert hasattr(MCPServerStreamableHttpParams, "__annotations__")
184+
185+
# Verify that the httpx_client_factory parameter is in the annotations
186+
annotations = MCPServerStreamableHttpParams.__annotations__
187+
assert "httpx_client_factory" in annotations
188+
189+
# The annotation should contain the string representation of the type
190+
annotation_str = str(annotations["httpx_client_factory"])
191+
assert "McpHttpClientFactory" in annotation_str
192+
193+
@pytest.mark.asyncio
194+
async def test_all_parameters_with_custom_factory(self):
195+
"""Test that all parameters work together with custom factory."""
196+
197+
def comprehensive_factory() -> httpx.AsyncClient:
198+
return httpx.AsyncClient(
199+
verify=False,
200+
timeout=httpx.Timeout(90.0),
201+
headers={"X-Test": "value"},
202+
)
203+
204+
with patch("agents.mcp.server.streamablehttp_client") as mock_client:
205+
mock_client.return_value = MagicMock()
206+
207+
server = MCPServerStreamableHttp(
208+
params={
209+
"url": "https://api.example.com/mcp",
210+
"headers": {"Authorization": "Bearer token"},
211+
"timeout": 45,
212+
"sse_read_timeout": 600,
213+
"terminate_on_close": False,
214+
"httpx_client_factory": comprehensive_factory,
215+
}
216+
)
217+
218+
server.create_streams()
219+
220+
mock_client.assert_called_once_with(
221+
url="https://api.example.com/mcp",
222+
headers={"Authorization": "Bearer token"},
223+
timeout=45,
224+
sse_read_timeout=600,
225+
terminate_on_close=False,
226+
httpx_client_factory=comprehensive_factory,
227+
)

0 commit comments

Comments
 (0)