Skip to content

Commit 815ebb4

Browse files
ushelpwukath
authored andcommitted
fix: add httpx_client_factory support to SseConnectionParams
Merge #4842 ### Link to Issue or Description of Change **1. Link to an existing issue** Closes #4841 **2. Changes** This change aligns SSE transport behavior with Streamable HTTP transport for MCP tools, enabling users to customize the httpx.AsyncClient for SSE-based MCP connections (e.g., proxy, auth, TLS config). - Update `SseConnectionParams` model to include httpx_client_factory parameter (matching StreamableHTTPConnectionParams pattern) - Pass `httpx_client_factory` from `SseConnectionParams` to underlying `sse_client()` call - Add `model_config` to `SseConnectionParams` to allow arbitrary types for the factory function - Update docstrings to document the new `httpx_client_factory attribute` - Add unit test for new senarios ### Testing Plan **Unit Tests:** - [x] I have added or updated unit tests for my change. - `tests/unittests/tools/mcp_tool/test_mcp_session_manager.py` - [x] All unit tests pass locally. - All test cases have passed the test. **Manual End-to-End (E2E) Tests:** 1. Set custom httpx_client_factory parameter for SseConnectionParams ```Python from typing import Any import httpx from google.adk.agents.llm_agent import Agent from google.adk.models import LiteLlm from google.adk.tools.mcp_tool import SseConnectionParams, McpToolset # Cutom httpx client factory for MCP SSE transport def custom_httpx_client_factory( headers: dict[str, str] | None = None, timeout: httpx.Timeout | None = None, auth: httpx.Auth | None = None, ) -> httpx.AsyncClient: # Print message for debugging print(' === Custom httpx_client_factory running! ===') kwargs: dict[str, Any] = { "follow_redirects": True, } if timeout is None: kwargs["timeout"] = httpx.Timeout(30, read=300) else: kwargs["timeout"] = timeout if headers is not None: kwargs["headers"] = headers if auth is not None: kwargs["auth"] = auth return httpx.AsyncClient(**kwargs) connection_params = SseConnectionParams( url='http://127.0.0.1:9000/sse', httpx_client_factory=custom_httpx_client_factory, # Custom httpx client factory ) tool = McpToolset(connection_params=connection_params) root_agent = Agent( model=LiteLlm(model="ollama_chat/qwen3.5:35b-a3b-q4_K_M"), name='root_agent', description='A helpful assistant for user questions.', instruction='Answer user questions to the best of your knowledge', tools=[tool], ) ``` 2. Run the Agent, and display the debug log in console ``` 2026-03-15 23:54:14,304 - INFO - agent_loader.py:130 - Found root_agent in my_agent.agent INFO: 127.0.0.1:55642 - "POST /run_sse HTTP/1.1" 200 OK === Custom httpx_client_factory running! === 2026-03-15 23:54:14,345 - INFO - _client.py:1740 - HTTP Request: GET http://127.0.0.1:9000/sse "HTTP/1.1 200 OK" 2026-03-15 23:54:14,349 - INFO - _client.py:1740 - HTTP Request: POST http://127.0.0.1:9000/messages/?session_id=b5750c6e214048c28110dff11c412ecb "HTTP/1.1 202 Accepted" ``` ### Checklist - [x] I have read the [CONTRIBUTING.md](https://github.com/google/adk-python/blob/main/CONTRIBUTING.md) document. - [x] I have performed a self-review of my own code. - [x] I have commented my code, particularly in hard-to-understand areas. - [x] I have added tests that prove my fix is effective or that my feature works. - [x] New and existing unit tests pass locally with my changes. - [x] I have manually tested my changes end-to-end. - [x] Any dependent changes have been merged and published in downstream modules. ### Additional context Co-authored-by: Kathy Wu <wukathy@google.com> COPYBARA_INTEGRATE_REVIEW=#4842 from ushelp:main e3125fb PiperOrigin-RevId: 890708694
1 parent 3f3ed0f commit 815ebb4

File tree

2 files changed

+49
-0
lines changed

2 files changed

+49
-0
lines changed

src/google/adk/tools/mcp_tool/mcp_session_manager.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -99,12 +99,17 @@ class SseConnectionParams(BaseModel):
9999
server.
100100
sse_read_timeout: Timeout in seconds for reading data from the MCP SSE
101101
server.
102+
httpx_client_factory: Factory function to create a custom HTTPX client. If
103+
not provided, a default factory will be used.
102104
"""
103105

106+
model_config = ConfigDict(arbitrary_types_allowed=True)
107+
104108
url: str
105109
headers: dict[str, Any] | None = None
106110
timeout: float = 5.0
107111
sse_read_timeout: float = 60 * 5.0
112+
httpx_client_factory: CheckableMcpHttpClientFactory = create_mcp_http_client
108113

109114

110115
@runtime_checkable
@@ -398,6 +403,7 @@ def _create_client(self, merged_headers: Optional[Dict[str, str]] = None):
398403
headers=merged_headers,
399404
timeout=self._connection_params.timeout,
400405
sse_read_timeout=self._connection_params.sse_read_timeout,
406+
httpx_client_factory=self._connection_params.httpx_client_factory,
401407
)
402408
elif isinstance(self._connection_params, StreamableHTTPConnectionParams):
403409
client = streamablehttp_client(

tests/unittests/tools/mcp_tool/test_mcp_session_manager.py

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -134,6 +134,49 @@ def test_init_with_sse_connection_params(self):
134134

135135
assert manager._connection_params == sse_params
136136

137+
@patch("google.adk.tools.mcp_tool.mcp_session_manager.sse_client")
138+
def test_init_with_sse_custom_httpx_factory(self, mock_sse_client):
139+
"""Test that sse_client is called with custom httpx_client_factory."""
140+
custom_httpx_factory = Mock()
141+
142+
sse_params = SseConnectionParams(
143+
url="https://example.com/mcp",
144+
timeout=10.0,
145+
httpx_client_factory=custom_httpx_factory,
146+
)
147+
manager = MCPSessionManager(sse_params)
148+
149+
manager._create_client()
150+
151+
mock_sse_client.assert_called_once_with(
152+
url="https://example.com/mcp",
153+
headers=None,
154+
timeout=10.0,
155+
sse_read_timeout=300.0,
156+
httpx_client_factory=custom_httpx_factory,
157+
)
158+
159+
@patch("google.adk.tools.mcp_tool.mcp_session_manager.sse_client")
160+
def test_init_with_sse_default_httpx_factory(self, mock_sse_client):
161+
"""Test that sse_client is called with default httpx_client_factory."""
162+
sse_params = SseConnectionParams(
163+
url="https://example.com/mcp",
164+
timeout=10.0,
165+
)
166+
manager = MCPSessionManager(sse_params)
167+
168+
manager._create_client()
169+
170+
mock_sse_client.assert_called_once_with(
171+
url="https://example.com/mcp",
172+
headers=None,
173+
timeout=10.0,
174+
sse_read_timeout=300.0,
175+
httpx_client_factory=SseConnectionParams.model_fields[
176+
"httpx_client_factory"
177+
].get_default(),
178+
)
179+
137180
def test_init_with_streamable_http_params(self):
138181
"""Test initialization with StreamableHTTPConnectionParams."""
139182
http_params = StreamableHTTPConnectionParams(

0 commit comments

Comments
 (0)