mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: Fix MCPStreamableHTTPTool to use new streamable_http_client API (#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>
This commit is contained in:
Copilot
committed by
GitHub
Unverified
GitHub
parent
c7cb5be231
commit
e63c148fc7
@@ -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):
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -146,7 +146,7 @@ ignore = [
|
||||
"TD003", # allow missing link to todo issue
|
||||
"FIX002", # allow todo
|
||||
"B027", # allow empty non-abstract method in ABC
|
||||
"RUF067", # allow version detection in __init__.py
|
||||
"RUF067" # Allow version in __init__.py
|
||||
]
|
||||
|
||||
[tool.ruff.lint.per-file-ignores]
|
||||
|
||||
@@ -4,6 +4,7 @@ import os
|
||||
|
||||
from agent_framework import ChatAgent, MCPStreamableHTTPTool
|
||||
from agent_framework.openai import OpenAIResponsesClient
|
||||
from httpx import AsyncClient
|
||||
|
||||
"""
|
||||
MCP Authentication Example
|
||||
@@ -31,13 +32,16 @@ async def api_key_auth_example() -> None:
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
}
|
||||
|
||||
# Create MCP tool with authentication headers
|
||||
# Create HTTP client with authentication headers
|
||||
http_client = AsyncClient(headers=auth_headers)
|
||||
|
||||
# Create MCP tool with the configured HTTP client
|
||||
async with (
|
||||
MCPStreamableHTTPTool(
|
||||
name="MCP tool",
|
||||
description="MCP tool description",
|
||||
url=mcp_server_url,
|
||||
headers=auth_headers, # Authentication headers
|
||||
http_client=http_client, # Pass HTTP client with authentication headers
|
||||
) as mcp_tool,
|
||||
ChatAgent(
|
||||
chat_client=OpenAIResponsesClient(),
|
||||
|
||||
Generated
+10
-10
@@ -377,7 +377,7 @@ requires-dist = [
|
||||
{ name = "agent-framework-purview", marker = "extra == 'all'", editable = "packages/purview" },
|
||||
{ name = "agent-framework-redis", marker = "extra == 'all'", editable = "packages/redis" },
|
||||
{ name = "azure-identity", specifier = ">=1,<2" },
|
||||
{ name = "mcp", extras = ["ws"], specifier = ">=1.23" },
|
||||
{ name = "mcp", extras = ["ws"], specifier = ">=1.24.0,<2" },
|
||||
{ name = "openai", specifier = ">=1.99.0" },
|
||||
{ name = "opentelemetry-api", specifier = ">=1.39.0" },
|
||||
{ name = "opentelemetry-sdk", specifier = ">=1.39.0" },
|
||||
@@ -2595,7 +2595,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "huggingface-hub"
|
||||
version = "1.2.4"
|
||||
version = "1.3.1"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "filelock", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
|
||||
@@ -2609,9 +2609,9 @@ dependencies = [
|
||||
{ name = "typer-slim", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
|
||||
{ name = "typing-extensions", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/fb/94/42ed2ff780f4bc58acbe4b8cb98eb4574310ad6feba12f76a820e7546120/huggingface_hub-1.2.4.tar.gz", hash = "sha256:7a1d9ec4802e64372d1d152d69fb8e26d943f15a2289096fbc8e09e7b90c21a5", size = 614771, upload-time = "2026-01-06T11:01:29.828Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/dd/dd/1cc985c5dda36298b152f75e82a1c81f52243b78fb7e9cad637a29561ad1/huggingface_hub-1.3.1.tar.gz", hash = "sha256:e80e0cfb4a75557c51ab20d575bdea6bb6106c2f97b7c75d8490642f1efb6df5", size = 622356, upload-time = "2026-01-09T14:08:16.888Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/dd/b0/113c4a688e7af9f0b92f5585cb425e71134e04c83a0a4a1e62db90edee20/huggingface_hub-1.2.4-py3-none-any.whl", hash = "sha256:2db69b91877d9d34825f5cd2a63b94f259011a77dcf761b437bf510fbe9522e9", size = 520980, upload-time = "2026-01-06T11:01:27.789Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/90/fb/cb8fe5f71d5622427f20bcab9e06a696a5aaf21bfe7bd0a8a0c63c88abf5/huggingface_hub-1.3.1-py3-none-any.whl", hash = "sha256:efbc7f3153cb84e2bb69b62ed90985e21ecc9343d15647a419fc0ee4b85f0ac3", size = 533351, upload-time = "2026-01-09T14:08:14.519Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -3068,7 +3068,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "litellm"
|
||||
version = "1.80.12"
|
||||
version = "1.80.13"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "aiohttp", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
|
||||
@@ -3086,9 +3086,9 @@ dependencies = [
|
||||
{ name = "tiktoken", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
|
||||
{ name = "tokenizers", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/d1/fa/29cea38f88a024e835eced4630ede0d65a0d9e0d951b1e17c6921dfce393/litellm-1.80.12.tar.gz", hash = "sha256:75eca39cb00f3ae00f3b5282e34fb304ffc0d36f7d4cfd189c9191b3a17b5815", size = 13264051, upload-time = "2026-01-07T17:50:33.965Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/e1/b4/ef75838159aabee15ad93d65ee0e91d04ba0e310784b7b0d3f490cca270c/litellm-1.80.13.tar.gz", hash = "sha256:61ed22dfad633ada3b97dd8a50d8e8d804da0115105006d2f9d77ba3fb247a0b", size = 13277620, upload-time = "2026-01-09T04:37:08.529Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/2b/34/57b06609b0b637bd1b2d68ed5282a4dfa1d2c35535202de70d36352ea4d4/litellm-1.80.12-py3-none-any.whl", hash = "sha256:d602add4486cc32f30ebdf1a05c58e8074ef1e87838fb5f1306a862968496ff7", size = 11544878, upload-time = "2026-01-07T17:50:30.776Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/a6/90/e8e0ad7f57d3a56c3411b3867e02768f9722b5975a263c8aaaaba6693d91/litellm-1.80.13-py3-none-any.whl", hash = "sha256:43dcdbca010961f17d7a5a6a995a38d1a46101350959b0e8269576cfe913cf0b", size = 11562501, upload-time = "2026-01-09T04:37:05.551Z" },
|
||||
]
|
||||
|
||||
[package.optional-dependencies]
|
||||
@@ -3130,11 +3130,11 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "litellm-proxy-extras"
|
||||
version = "0.4.18"
|
||||
version = "0.4.20"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/b8/96/ed23a3c3af1913b4a9b4afbdbd5176b8b1e06138f26a9584427f09f652c1/litellm_proxy_extras-0.4.18.tar.gz", hash = "sha256:898b28e3e74acdc29142906b84787ab05a90e30aa3c0c8aee849915e3a16adb3", size = 20676, upload-time = "2026-01-07T09:08:13.903Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/44/ab/df85ce715ebf488cacba338607f67d46c4e4db0b717c9d2f904b8dc7de12/litellm_proxy_extras-0.4.20.tar.gz", hash = "sha256:4fcc95db25cc8b75abbc3f00bb79fd6b94edd1b838ad7bb12cf839b39c67923a", size = 21044, upload-time = "2026-01-07T19:11:32.562Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/af/60/2465d2d58625af9ffc80c3933b43cf63e94f1722078b5da197398b0705de/litellm_proxy_extras-0.4.18-py3-none-any.whl", hash = "sha256:c3edee68bf8eb073c6158dcf7df05727dfc829e63c03a617fcb48853d11490df", size = 45338, upload-time = "2026-01-07T09:08:12.886Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/d4/f5/eb350c49e7cf09db5b335aaeef410c2094e19c84bfe51733cab8470dc011/litellm_proxy_extras-0.4.20-py3-none-any.whl", hash = "sha256:7737cd693dd1aa0bd25ada6d300b37f42c8c18d1820535aceb0ed38ed21f68f5", size = 46565, upload-time = "2026-01-07T19:11:29.728Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
||||
Reference in New Issue
Block a user