agent/agent-framework
2
0
Fork 0
mirror of https://github.com/microsoft/agent-framework.git synced 2026-06-16 21:04:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

Python: Fix MCPStreamableHTTPTool to use new streamable_http_client API (#3088)

Browse Source 
* Fix MCPStreamableHTTPTool to use new streamable_http_client API with proper httpx client cleanup

Co-authored-by: eavanvalkenburg <13749212+eavanvalkenburg@users.noreply.github.com>

* Update docstring to reflect new streamable_http_client API usage

Co-authored-by: eavanvalkenburg <13749212+eavanvalkenburg@users.noreply.github.com>

* Refactor MCPStreamableHTTPTool to accept optional http_client parameter and delegate client creation to streamable_http_client

Co-authored-by: eavanvalkenburg <13749212+eavanvalkenburg@users.noreply.github.com>

* Update mcp package minimum version to 1.24.0 for streamable_http_client API support

Co-authored-by: eavanvalkenburg <13749212+eavanvalkenburg@users.noreply.github.com>

* Fix critical bugs: apply headers/timeout/sse_read_timeout when creating httpx client, add version constraint <2, and properly manage client lifecycle

Co-authored-by: eavanvalkenburg <13749212+eavanvalkenburg@users.noreply.github.com>

* Simplify implementation: remove headers/timeout/sse_read_timeout params, remove kwargs, remove close() override per feedback

Co-authored-by: eavanvalkenburg <13749212+eavanvalkenburg@users.noreply.github.com>

* Add back **kwargs parameter for backward compatibility (accepted but not used)

Co-authored-by: eavanvalkenburg <13749212+eavanvalkenburg@users.noreply.github.com>

* Remove unused httpx import from test file

Note: The uv.lock file needs to be updated with 'uv sync' to reflect the mcp version constraint change (>=1.24.0,<2)

Co-authored-by: eavanvalkenburg <13749212+eavanvalkenburg@users.noreply.github.com>

* cicd fixes

* udpated samples with headers examples

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: eavanvalkenburg <13749212+eavanvalkenburg@users.noreply.github.com>
Co-authored-by: eavanvalkenburg <github@vanvalkenburg.eu>
This commit is contained in:
Copilot
2026-01-12 09:26:53 -08:00
committed by GitHub
Unverified
parent c7cb5be231
commit e63c148fc7
6 changed files with 99 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/core/agent_framework/_mcp.py
+19 -31
View File
@@ -10,10 +10,11 @@ from datetime import timedelta
from functools import partial
from typing import TYPE_CHECKING, Any, Literal
import httpx
from mcp import types
from mcp.client.session import ClientSession
from mcp.client.stdio import StdioServerParameters, stdio_client
from mcp.client.streamable_http import streamablehttp_client
from mcp.client.streamable_http import streamable_http_client
from mcp.client.websocket import websocket_client
from mcp.shared.context import RequestContext
from mcp.shared.exceptions import McpError
@@ -897,7 +898,6 @@ class MCPStreamableHTTPTool(MCPTool):
mcp_tool = MCPStreamableHTTPTool(
name="web-api",
url="https://api.example.com/mcp",
headers={"Authorization": "Bearer token"},
description="Web API operations",
)
@@ -919,21 +919,19 @@ class MCPStreamableHTTPTool(MCPTool):
description: str | None = None,
approval_mode: (Literal["always_require", "never_require"] | HostedMCPSpecificApproval | None) = None,
allowed_tools: Collection[str] | None = None,
headers: dict[str, Any] | None = None,
timeout: float | None = None,
sse_read_timeout: float | None = None,
terminate_on_close: bool | None = None,
chat_client: "ChatClientProtocol | None" = None,
additional_properties: dict[str, Any] | None = None,
http_client: httpx.AsyncClient | None = None,
**kwargs: Any,
) -> None:
"""Initialize the MCP streamable HTTP tool.
Note:
The arguments are used to create a streamable HTTP client.
See ``mcp.client.streamable_http.streamablehttp_client`` for more details.
Any extra arguments passed to the constructor will be passed to the
streamable HTTP client constructor.
The arguments are used to create a streamable HTTP client using the
new ``mcp.client.streamable_http.streamable_http_client`` API.
If an httpx.AsyncClient is provided via ``http_client``, it will be used directly.
Otherwise, the ``streamable_http_client`` API will create and manage a default client.
Args:
name: The name of the tool.
@@ -953,12 +951,13 @@ class MCPStreamableHTTPTool(MCPTool):
A tool should not be listed in both, if so, it will require approval.
allowed_tools: A list of tools that are allowed to use this tool.
additional_properties: Additional properties.
headers: The headers to send with the request.
timeout: The timeout for the request.
sse_read_timeout: The timeout for reading from the SSE stream.
terminate_on_close: Close the transport when the MCP client is terminated.
chat_client: The chat client to use for sampling.
kwargs: Any extra arguments to pass to the SSE client.
http_client: Optional httpx.AsyncClient to use. If not provided, the
``streamable_http_client`` API will create and manage a default client.
To configure headers, timeouts, or other HTTP client settings, create
and pass your own ``httpx.AsyncClient`` instance.
kwargs: Additional keyword arguments (accepted for backward compatibility but not used).
"""
super().__init__(
name=name,
@@ -973,11 +972,8 @@ class MCPStreamableHTTPTool(MCPTool):
request_timeout=request_timeout,
)
self.url = url
self.headers = headers or {}
self.timeout = timeout
self.sse_read_timeout = sse_read_timeout
self.terminate_on_close = terminate_on_close
self._client_kwargs = kwargs
self._httpx_client: httpx.AsyncClient | None = http_client
def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]:
"""Get an MCP streamable HTTP client.
@@ -985,20 +981,12 @@ class MCPStreamableHTTPTool(MCPTool):
Returns:
An async context manager for the streamable HTTP client transport.
"""
args: dict[str, Any] = {
"url": self.url,
}
if self.headers:
args["headers"] = self.headers
if self.timeout is not None:
args["timeout"] = self.timeout
if self.sse_read_timeout is not None:
args["sse_read_timeout"] = self.sse_read_timeout
if self.terminate_on_close is not None:
args["terminate_on_close"] = self.terminate_on_close
if self._client_kwargs:
args.update(self._client_kwargs)
return streamablehttp_client(**args)
# Pass the http_client (which may be None) to streamable_http_client
return streamable_http_client(
url=self.url,
http_client=self._httpx_client,
terminate_on_close=self.terminate_on_close if self.terminate_on_close is not None else True,
)
class MCPWebsocketTool(MCPTool):
python/packages/core/pyproject.toml
+1 -1
View File
@@ -34,7 +34,7 @@ dependencies = [
# connectors and functions
"openai>=1.99.0",
"azure-identity>=1,<2",
"mcp[ws]>=1.23",
"mcp[ws]>=1.24.0,<2",
"packaging>=24.1",
]
python/packages/core/tests/core/test_mcp.py
+62 -10
View File
@@ -1512,24 +1512,18 @@ def test_mcp_streamable_http_tool_get_mcp_client_all_params():
tool = MCPStreamableHTTPTool(
name="test",
url="http://example.com",
headers={"Auth": "token"},
timeout=30.0,
sse_read_timeout=10.0,
terminate_on_close=True,
custom_param="test",
)
with patch("agent_framework._mcp.streamablehttp_client") as mock_http_client:
with patch("agent_framework._mcp.streamable_http_client") as mock_http_client:
tool.get_mcp_client()
# Verify all parameters were passed
# Verify streamable_http_client was called with None for http_client
# (since we didn't provide one, the API will create its own)
mock_http_client.assert_called_once_with(
url="http://example.com",
headers={"Auth": "token"},
timeout=30.0,
sse_read_timeout=10.0,
http_client=None,
terminate_on_close=True,
custom_param="test",
)
@@ -1692,3 +1686,61 @@ async def test_load_prompts_prevents_multiple_calls():
tool._prompts_loaded = True
assert mock_session.list_prompts.call_count == 1 # Still 1, not incremented
@pytest.mark.asyncio
async def test_mcp_streamable_http_tool_httpx_client_cleanup():
"""Test that MCPStreamableHTTPTool properly passes through httpx clients."""
from unittest.mock import AsyncMock, Mock, patch
from agent_framework import MCPStreamableHTTPTool
# Mock the streamable_http_client to avoid actual connections
with (
patch("agent_framework._mcp.streamable_http_client") as mock_client,
patch("agent_framework._mcp.ClientSession") as mock_session_class,
):
# Setup mock context manager for streamable_http_client
mock_transport = (Mock(), Mock())
mock_context_manager = Mock()
mock_context_manager.__aenter__ = AsyncMock(return_value=mock_transport)
mock_context_manager.__aexit__ = AsyncMock(return_value=None)
mock_client.return_value = mock_context_manager
# Setup mock session
mock_session = Mock()
mock_session.initialize = AsyncMock()
mock_session_class.return_value.__aenter__ = AsyncMock(return_value=mock_session)
mock_session_class.return_value.__aexit__ = AsyncMock(return_value=None)
# Test 1: Tool without provided client (passes None to streamable_http_client)
tool1 = MCPStreamableHTTPTool(
name="test",
url="http://localhost:8081/mcp",
load_tools=False,
load_prompts=False,
terminate_on_close=False,
)
await tool1.connect()
# When no client is provided, _httpx_client should be None
assert tool1._httpx_client is None, "httpx client should be None when not provided"
# Test 2: Tool with user-provided client
user_client = Mock()
tool2 = MCPStreamableHTTPTool(
name="test",
url="http://localhost:8081/mcp",
load_tools=False,
load_prompts=False,
terminate_on_close=False,
http_client=user_client,
)
await tool2.connect()
# Verify the user-provided client was stored
assert tool2._httpx_client is user_client, "User-provided client should be stored"
# Verify streamable_http_client was called with the user's client
# Get the last call (should be from tool2.connect())
call_args = mock_client.call_args
assert call_args.kwargs["http_client"] is user_client, "User's client should be passed through"