Skip to content

Commit e63c148

Browse files
Python: Fix MCPStreamableHTTPTool to use new streamable_http_client API (microsoft#3088)
* 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>
1 parent c7cb5be commit e63c148

6 files changed

Lines changed: 99 additions & 55 deletions

File tree

python/packages/core/agent_framework/_mcp.py

Lines changed: 19 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -10,10 +10,11 @@
1010
from functools import partial
1111
from typing import TYPE_CHECKING, Any, Literal
1212

13+
import httpx
1314
from mcp import types
1415
from mcp.client.session import ClientSession
1516
from mcp.client.stdio import StdioServerParameters, stdio_client
16-
from mcp.client.streamable_http import streamablehttp_client
17+
from mcp.client.streamable_http import streamable_http_client
1718
from mcp.client.websocket import websocket_client
1819
from mcp.shared.context import RequestContext
1920
from mcp.shared.exceptions import McpError
@@ -897,7 +898,6 @@ class MCPStreamableHTTPTool(MCPTool):
897898
mcp_tool = MCPStreamableHTTPTool(
898899
name="web-api",
899900
url="https://api.example.com/mcp",
900-
headers={"Authorization": "Bearer token"},
901901
description="Web API operations",
902902
)
903903
@@ -919,21 +919,19 @@ def __init__(
919919
description: str | None = None,
920920
approval_mode: (Literal["always_require", "never_require"] | HostedMCPSpecificApproval | None) = None,
921921
allowed_tools: Collection[str] | None = None,
922-
headers: dict[str, Any] | None = None,
923-
timeout: float | None = None,
924-
sse_read_timeout: float | None = None,
925922
terminate_on_close: bool | None = None,
926923
chat_client: "ChatClientProtocol | None" = None,
927924
additional_properties: dict[str, Any] | None = None,
925+
http_client: httpx.AsyncClient | None = None,
928926
**kwargs: Any,
929927
) -> None:
930928
"""Initialize the MCP streamable HTTP tool.
931929
932930
Note:
933-
The arguments are used to create a streamable HTTP client.
934-
See ``mcp.client.streamable_http.streamablehttp_client`` for more details.
935-
Any extra arguments passed to the constructor will be passed to the
936-
streamable HTTP client constructor.
931+
The arguments are used to create a streamable HTTP client using the
932+
new ``mcp.client.streamable_http.streamable_http_client`` API.
933+
If an httpx.AsyncClient is provided via ``http_client``, it will be used directly.
934+
Otherwise, the ``streamable_http_client`` API will create and manage a default client.
937935
938936
Args:
939937
name: The name of the tool.
@@ -953,12 +951,13 @@ def __init__(
953951
A tool should not be listed in both, if so, it will require approval.
954952
allowed_tools: A list of tools that are allowed to use this tool.
955953
additional_properties: Additional properties.
956-
headers: The headers to send with the request.
957-
timeout: The timeout for the request.
958-
sse_read_timeout: The timeout for reading from the SSE stream.
959954
terminate_on_close: Close the transport when the MCP client is terminated.
960955
chat_client: The chat client to use for sampling.
961-
kwargs: Any extra arguments to pass to the SSE client.
956+
http_client: Optional httpx.AsyncClient to use. If not provided, the
957+
``streamable_http_client`` API will create and manage a default client.
958+
To configure headers, timeouts, or other HTTP client settings, create
959+
and pass your own ``httpx.AsyncClient`` instance.
960+
kwargs: Additional keyword arguments (accepted for backward compatibility but not used).
962961
"""
963962
super().__init__(
964963
name=name,
@@ -973,32 +972,21 @@ def __init__(
973972
request_timeout=request_timeout,
974973
)
975974
self.url = url
976-
self.headers = headers or {}
977-
self.timeout = timeout
978-
self.sse_read_timeout = sse_read_timeout
979975
self.terminate_on_close = terminate_on_close
980-
self._client_kwargs = kwargs
976+
self._httpx_client: httpx.AsyncClient | None = http_client
981977

982978
def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]:
983979
"""Get an MCP streamable HTTP client.
984980
985981
Returns:
986982
An async context manager for the streamable HTTP client transport.
987983
"""
988-
args: dict[str, Any] = {
989-
"url": self.url,
990-
}
991-
if self.headers:
992-
args["headers"] = self.headers
993-
if self.timeout is not None:
994-
args["timeout"] = self.timeout
995-
if self.sse_read_timeout is not None:
996-
args["sse_read_timeout"] = self.sse_read_timeout
997-
if self.terminate_on_close is not None:
998-
args["terminate_on_close"] = self.terminate_on_close
999-
if self._client_kwargs:
1000-
args.update(self._client_kwargs)
1001-
return streamablehttp_client(**args)
984+
# Pass the http_client (which may be None) to streamable_http_client
985+
return streamable_http_client(
986+
url=self.url,
987+
http_client=self._httpx_client,
988+
terminate_on_close=self.terminate_on_close if self.terminate_on_close is not None else True,
989+
)
1002990

1003991

1004992
class MCPWebsocketTool(MCPTool):

python/packages/core/pyproject.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@ dependencies = [
3434
# connectors and functions
3535
"openai>=1.99.0",
3636
"azure-identity>=1,<2",
37-
"mcp[ws]>=1.23",
37+
"mcp[ws]>=1.24.0,<2",
3838
"packaging>=24.1",
3939
]
4040

python/packages/core/tests/core/test_mcp.py

Lines changed: 62 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1512,24 +1512,18 @@ def test_mcp_streamable_http_tool_get_mcp_client_all_params():
15121512
tool = MCPStreamableHTTPTool(
15131513
name="test",
15141514
url="http://example.com",
1515-
headers={"Auth": "token"},
1516-
timeout=30.0,
1517-
sse_read_timeout=10.0,
15181515
terminate_on_close=True,
1519-
custom_param="test",
15201516
)
15211517

1522-
with patch("agent_framework._mcp.streamablehttp_client") as mock_http_client:
1518+
with patch("agent_framework._mcp.streamable_http_client") as mock_http_client:
15231519
tool.get_mcp_client()
15241520

1525-
# Verify all parameters were passed
1521+
# Verify streamable_http_client was called with None for http_client
1522+
# (since we didn't provide one, the API will create its own)
15261523
mock_http_client.assert_called_once_with(
15271524
url="http://example.com",
1528-
headers={"Auth": "token"},
1529-
timeout=30.0,
1530-
sse_read_timeout=10.0,
1525+
http_client=None,
15311526
terminate_on_close=True,
1532-
custom_param="test",
15331527
)
15341528

15351529

@@ -1692,3 +1686,61 @@ async def test_load_prompts_prevents_multiple_calls():
16921686
tool._prompts_loaded = True
16931687

16941688
assert mock_session.list_prompts.call_count == 1 # Still 1, not incremented
1689+
1690+
1691+
@pytest.mark.asyncio
1692+
async def test_mcp_streamable_http_tool_httpx_client_cleanup():
1693+
"""Test that MCPStreamableHTTPTool properly passes through httpx clients."""
1694+
from unittest.mock import AsyncMock, Mock, patch
1695+
1696+
from agent_framework import MCPStreamableHTTPTool
1697+
1698+
# Mock the streamable_http_client to avoid actual connections
1699+
with (
1700+
patch("agent_framework._mcp.streamable_http_client") as mock_client,
1701+
patch("agent_framework._mcp.ClientSession") as mock_session_class,
1702+
):
1703+
# Setup mock context manager for streamable_http_client
1704+
mock_transport = (Mock(), Mock())
1705+
mock_context_manager = Mock()
1706+
mock_context_manager.__aenter__ = AsyncMock(return_value=mock_transport)
1707+
mock_context_manager.__aexit__ = AsyncMock(return_value=None)
1708+
mock_client.return_value = mock_context_manager
1709+
1710+
# Setup mock session
1711+
mock_session = Mock()
1712+
mock_session.initialize = AsyncMock()
1713+
mock_session_class.return_value.__aenter__ = AsyncMock(return_value=mock_session)
1714+
mock_session_class.return_value.__aexit__ = AsyncMock(return_value=None)
1715+
1716+
# Test 1: Tool without provided client (passes None to streamable_http_client)
1717+
tool1 = MCPStreamableHTTPTool(
1718+
name="test",
1719+
url="http://localhost:8081/mcp",
1720+
load_tools=False,
1721+
load_prompts=False,
1722+
terminate_on_close=False,
1723+
)
1724+
await tool1.connect()
1725+
# When no client is provided, _httpx_client should be None
1726+
assert tool1._httpx_client is None, "httpx client should be None when not provided"
1727+
1728+
# Test 2: Tool with user-provided client
1729+
user_client = Mock()
1730+
tool2 = MCPStreamableHTTPTool(
1731+
name="test",
1732+
url="http://localhost:8081/mcp",
1733+
load_tools=False,
1734+
load_prompts=False,
1735+
terminate_on_close=False,
1736+
http_client=user_client,
1737+
)
1738+
await tool2.connect()
1739+
1740+
# Verify the user-provided client was stored
1741+
assert tool2._httpx_client is user_client, "User-provided client should be stored"
1742+
1743+
# Verify streamable_http_client was called with the user's client
1744+
# Get the last call (should be from tool2.connect())
1745+
call_args = mock_client.call_args
1746+
assert call_args.kwargs["http_client"] is user_client, "User's client should be passed through"

python/pyproject.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -146,7 +146,7 @@ ignore = [
146146
"TD003", # allow missing link to todo issue
147147
"FIX002", # allow todo
148148
"B027", # allow empty non-abstract method in ABC
149-
"RUF067", # allow version detection in __init__.py
149+
"RUF067" # Allow version in __init__.py
150150
]
151151

152152
[tool.ruff.lint.per-file-ignores]

python/samples/getting_started/mcp/mcp_api_key_auth.py

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44

55
from agent_framework import ChatAgent, MCPStreamableHTTPTool
66
from agent_framework.openai import OpenAIResponsesClient
7+
from httpx import AsyncClient
78

89
"""
910
MCP Authentication Example
@@ -31,13 +32,16 @@ async def api_key_auth_example() -> None:
3132
"Authorization": f"Bearer {api_key}",
3233
}
3334

34-
# Create MCP tool with authentication headers
35+
# Create HTTP client with authentication headers
36+
http_client = AsyncClient(headers=auth_headers)
37+
38+
# Create MCP tool with the configured HTTP client
3539
async with (
3640
MCPStreamableHTTPTool(
3741
name="MCP tool",
3842
description="MCP tool description",
3943
url=mcp_server_url,
40-
headers=auth_headers, # Authentication headers
44+
http_client=http_client, # Pass HTTP client with authentication headers
4145
) as mcp_tool,
4246
ChatAgent(
4347
chat_client=OpenAIResponsesClient(),

python/uv.lock

Lines changed: 10 additions & 10 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)