diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index a25f359a59..034ed5e9ff 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -10,10 +10,11 @@ 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 @@ def __init__( 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 @@ def __init__( 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 @@ def __init__( 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 @@ def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]: 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): diff --git a/python/packages/core/pyproject.toml b/python/packages/core/pyproject.toml index 1b245277c0..942f9d8919 100644 --- a/python/packages/core/pyproject.toml +++ b/python/packages/core/pyproject.toml @@ -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", ] diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index 18c90d64b3..89c3c520fe 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -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" diff --git a/python/pyproject.toml b/python/pyproject.toml index cbdf3f0d75..b8ea697142 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -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] diff --git a/python/samples/getting_started/mcp/mcp_api_key_auth.py b/python/samples/getting_started/mcp/mcp_api_key_auth.py index f3ec1777e6..d80d92d4fa 100644 --- a/python/samples/getting_started/mcp/mcp_api_key_auth.py +++ b/python/samples/getting_started/mcp/mcp_api_key_auth.py @@ -4,6 +4,7 @@ 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(), diff --git a/python/uv.lock b/python/uv.lock index 15d1fd855e..26fb2ed57e 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -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]]