From 8d7e4241419d589f4b5ec717a5a7ee0ddebace79 Mon Sep 17 00:00:00 2001 From: Copilot Date: Mon, 23 Mar 2026 10:10:42 +0000 Subject: [PATCH 01/11] Python: Add header_provider to MCPStreamableHTTPTool (#4808) Add a header_provider callback parameter to MCPStreamableHTTPTool that enables injecting dynamic per-request HTTP headers from runtime kwargs (originating from FunctionInvocationContext.kwargs set in agent middleware). The implementation uses contextvars and httpx event hooks to ensure headers are task-local and safe for concurrent tool calls: - header_provider receives the runtime kwargs dict and returns headers - call_tool sets a ContextVar before delegating to MCPTool.call_tool - An httpx request event hook reads from the ContextVar and injects headers Example usage: mcp_tool = MCPStreamableHTTPTool( name="web-api", url="https://api.example.com/mcp", header_provider=lambda kwargs: { "X-Auth-Token": kwargs.get("auth_token", ""), }, ) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- python/packages/core/agent_framework/_mcp.py | 58 ++++- python/packages/core/tests/core/test_mcp.py | 257 +++++++++++++++++++ 2 files changed, 309 insertions(+), 6 deletions(-) diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 0dab38c820..9bdfa7c4ea 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -4,6 +4,7 @@ import asyncio import base64 +import contextvars import json import logging import re @@ -59,6 +60,7 @@ class MCPSpecificApproval(TypedDict, total=False): _MCP_REMOTE_NAME_KEY = "_mcp_remote_name" _MCP_NORMALIZED_NAME_KEY = "_mcp_normalized_name" +_mcp_call_headers: contextvars.ContextVar[dict[str, str]] = contextvars.ContextVar("_mcp_call_headers") # region: Helpers @@ -1386,6 +1388,7 @@ def __init__( client: SupportsChatGetResponse | None = None, additional_properties: dict[str, Any] | None = None, http_client: AsyncClient | None = None, + header_provider: Callable[[dict[str, Any]], dict[str, str]] | None = None, **kwargs: Any, ) -> None: """Initialize the MCP streamable HTTP tool. @@ -1433,6 +1436,11 @@ def __init__( ``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 ``asyncClient`` instance. + header_provider: Optional callable that receives the runtime keyword arguments + (from ``FunctionInvocationContext.kwargs``) and returns a ``dict[str, str]`` + of HTTP headers to inject into every outbound request to the MCP server. + Use this to forward per-request context (e.g. authentication tokens set in + agent middleware) without creating a separate ``httpx.AsyncClient``. kwargs: Additional keyword arguments (accepted for backward compatibility but not used). """ super().__init__( @@ -1453,6 +1461,7 @@ def __init__( self.url = url self.terminate_on_close = terminate_on_close self._httpx_client: AsyncClient | None = http_client + self._header_provider = header_provider def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]: """Get an MCP streamable HTTP client. @@ -1460,18 +1469,55 @@ def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]: Returns: An async context manager for the streamable HTTP client transport. """ - try: - from mcp.client.streamable_http import streamable_http_client - except ModuleNotFoundError as ex: - raise ModuleNotFoundError("`mcp` is required to use `MCPStreamableHTTPTool`. Please install `mcp`.") from ex + from mcp.client.streamable_http import streamable_http_client + + http_client = self._httpx_client + if self._header_provider is not None: + if http_client is None: + from mcp.shared._httpx_utils import create_mcp_http_client + + http_client = create_mcp_http_client() + self._httpx_client = http_client + + async def _inject_headers(request: httpx.Request) -> None: # noqa: RUF029 + headers = _mcp_call_headers.get({}) + for key, value in headers.items(): + request.headers[key] = value + + http_client.event_hooks["request"].append(_inject_headers) - # Pass the http_client (which may be None) to streamable_http_client return streamable_http_client( url=self.url, - http_client=self._httpx_client, + http_client=http_client, terminate_on_close=self.terminate_on_close if self.terminate_on_close is not None else True, ) + async def call_tool(self, tool_name: str, **kwargs: Any) -> str | list[Content]: + """Call a tool, injecting headers from the header_provider if configured. + + When a ``header_provider`` was supplied at construction time, the runtime + *kwargs* (originating from ``FunctionInvocationContext.kwargs``) are passed + to the provider. The returned headers are attached to every HTTP request + made during this tool call via a ``contextvars.ContextVar``. + + Args: + tool_name: The name of the tool to call. + + Keyword Args: + kwargs: Arguments to pass to the tool. + + Returns: + A list of Content items representing the tool output. + """ + if self._header_provider is not None: + headers = self._header_provider(kwargs) + token = _mcp_call_headers.set(headers) + try: + return await super().call_tool(tool_name, **kwargs) + finally: + _mcp_call_headers.reset(token) + return await super().call_tool(tool_name, **kwargs) + class MCPWebsocketTool(MCPTool): """MCP tool for connecting to WebSocket-based MCP servers. diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index 09c036c704..64ee6ca690 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -3805,3 +3805,260 @@ def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]: # endregion + + +# region: MCPStreamableHTTPTool header_provider + + +async def test_mcp_streamable_http_tool_header_provider_injects_headers(): + """Test that header_provider injects per-request HTTP headers via runtime kwargs. + + When header_provider is configured, runtime kwargs from FunctionInvocationContext + are passed to the provider and the returned headers appear on outbound HTTP requests. + """ + + class _TestServer(MCPStreamableHTTPTool): + async def connect(self): + self.session = Mock(spec=ClientSession) + self.session.list_tools = AsyncMock( + return_value=types.ListToolsResult( + tools=[ + types.Tool( + name="greet", + description="Says hello", + inputSchema={ + "type": "object", + "properties": {"name": {"type": "string"}}, + "required": ["name"], + }, + ) + ] + ) + ) + self.session.call_tool = AsyncMock( + return_value=types.CallToolResult(content=[types.TextContent(type="text", text="Hello!")]) + ) + self.session.send_ping = AsyncMock() + self.is_connected = True + + def get_mcp_client(self): + return None + + def provider(kwargs): + return {"X-Some-Token": kwargs.get("some_token", "")} + + server = _TestServer( + name="test", + url="http://example.com/mcp", + header_provider=provider, + ) + async with server: + await server.load_tools() + + # Simulate the runtime kwargs that flow from FunctionInvocationContext.kwargs + await server.call_tool("greet", name="Alice", some_token="my-secret") + + # Verify the MCP session.call_tool was called + server.session.call_tool.assert_called_once() + + +async def test_mcp_streamable_http_tool_header_provider_sets_contextvar(): + """Test that call_tool sets the contextvar with headers from header_provider.""" + from agent_framework._mcp import _mcp_call_headers + + observed_headers: list[dict[str, str]] = [] + original_call_tool = MCPTool.call_tool + + async def spy_call_tool(self, tool_name, **kwargs): + # Capture the contextvar value during the super call + try: + observed_headers.append(_mcp_call_headers.get()) + except LookupError: + observed_headers.append({}) + return await original_call_tool(self, tool_name, **kwargs) + + class _TestServer(MCPStreamableHTTPTool): + async def connect(self): + self.session = Mock(spec=ClientSession) + self.session.list_tools = AsyncMock( + return_value=types.ListToolsResult( + tools=[ + types.Tool( + name="greet", + description="Says hello", + inputSchema={"type": "object", "properties": {"name": {"type": "string"}}}, + ) + ] + ) + ) + self.session.call_tool = AsyncMock( + return_value=types.CallToolResult(content=[types.TextContent(type="text", text="Hello!")]) + ) + self.session.send_ping = AsyncMock() + self.is_connected = True + + def get_mcp_client(self): + return None + + server = _TestServer( + name="test", + url="http://example.com/mcp", + header_provider=lambda kw: {"X-Auth": kw.get("auth_token", "")}, + ) + async with server: + await server.load_tools() + + with patch.object(MCPTool, "call_tool", spy_call_tool): + await server.call_tool("greet", name="Alice", auth_token="bearer-xyz") + + assert len(observed_headers) == 1 + assert observed_headers[0] == {"X-Auth": "bearer-xyz"} + + +async def test_mcp_streamable_http_tool_header_provider_contextvar_reset_after_call(): + """Test that the contextvar is properly reset after call_tool completes.""" + from agent_framework._mcp import _mcp_call_headers + + class _TestServer(MCPStreamableHTTPTool): + async def connect(self): + self.session = Mock(spec=ClientSession) + self.session.list_tools = AsyncMock( + return_value=types.ListToolsResult( + tools=[ + types.Tool( + name="greet", + description="Says hello", + inputSchema={"type": "object", "properties": {"name": {"type": "string"}}}, + ) + ] + ) + ) + self.session.call_tool = AsyncMock( + return_value=types.CallToolResult(content=[types.TextContent(type="text", text="Hello!")]) + ) + self.session.send_ping = AsyncMock() + self.is_connected = True + + def get_mcp_client(self): + return None + + server = _TestServer( + name="test", + url="http://example.com/mcp", + header_provider=lambda kw: {"X-Token": kw.get("token", "")}, + ) + async with server: + await server.load_tools() + await server.call_tool("greet", name="Alice", token="secret") + + # After call_tool, the contextvar should be unset (reset to no value) + with pytest.raises(LookupError): + _mcp_call_headers.get() + + +async def test_mcp_streamable_http_tool_without_header_provider(): + """Test that call_tool works normally when no header_provider is configured.""" + + class _TestServer(MCPStreamableHTTPTool): + async def connect(self): + self.session = Mock(spec=ClientSession) + self.session.list_tools = AsyncMock( + return_value=types.ListToolsResult( + tools=[ + types.Tool( + name="greet", + description="Says hello", + inputSchema={"type": "object", "properties": {"name": {"type": "string"}}}, + ) + ] + ) + ) + self.session.call_tool = AsyncMock( + return_value=types.CallToolResult(content=[types.TextContent(type="text", text="Hello!")]) + ) + self.session.send_ping = AsyncMock() + self.is_connected = True + + def get_mcp_client(self): + return None + + server = _TestServer( + name="test", + url="http://example.com/mcp", + ) + async with server: + await server.load_tools() + await server.call_tool("greet", name="Alice") + server.session.call_tool.assert_called_once() + + # Without header_provider, call_tool should delegate directly to MCPTool + assert server._header_provider is None + + +async def test_mcp_streamable_http_tool_header_provider_with_httpx_event_hook(): + """Test that the httpx event hook injects headers from the contextvar.""" + import httpx + + from agent_framework._mcp import _mcp_call_headers + + tool = MCPStreamableHTTPTool( + name="test", + url="http://example.com/mcp", + header_provider=lambda kw: {"X-Custom": kw.get("custom", "")}, + ) + + with patch("agent_framework._mcp.streamable_http_client"): + # Trigger get_mcp_client to set up the event hook + tool.get_mcp_client() + + # The tool should have created an httpx client with the event hook + assert tool._httpx_client is not None + hooks = tool._httpx_client.event_hooks.get("request", []) + assert len(hooks) == 1, "Expected one request event hook" + + # Simulate what happens during a call_tool: contextvar is set + token = _mcp_call_headers.set({"X-Custom": "test-value"}) + try: + request = httpx.Request("POST", "http://example.com/mcp") + await hooks[0](request) + assert request.headers.get("X-Custom") == "test-value" + finally: + _mcp_call_headers.reset(token) + + +async def test_mcp_streamable_http_tool_header_provider_with_user_httpx_client(): + """Test that header_provider works when the user provides their own httpx client.""" + import httpx + + from agent_framework._mcp import _mcp_call_headers + + user_client = httpx.AsyncClient(headers={"X-Base": "static"}) + + tool = MCPStreamableHTTPTool( + name="test", + url="http://example.com/mcp", + http_client=user_client, + header_provider=lambda kw: {"X-Dynamic": kw.get("dynamic", "")}, + ) + + with patch("agent_framework._mcp.streamable_http_client"): + tool.get_mcp_client() + + # The user's client should still be used + assert tool._httpx_client is user_client + hooks = user_client.event_hooks.get("request", []) + assert len(hooks) == 1 + + # Verify the hook injects headers + token = _mcp_call_headers.set({"X-Dynamic": "per-request"}) + try: + request = httpx.Request("POST", "http://example.com/mcp") + await hooks[0](request) + assert request.headers.get("X-Dynamic") == "per-request" + finally: + _mcp_call_headers.reset(token) + + await user_client.aclose() + + +# endregion From 4c08967185a59c842334d296db3be68688e48e1c Mon Sep 17 00:00:00 2001 From: Copilot Date: Mon, 23 Mar 2026 10:29:59 +0000 Subject: [PATCH 02/11] Address review feedback for #4808: Python: [Bug]: Unable to pass AgentContext to MCPStreamableHTTPTool --- python/packages/core/agent_framework/_mcp.py | 21 +++--- python/packages/core/tests/core/test_mcp.py | 70 ++++++++++++++------ 2 files changed, 63 insertions(+), 28 deletions(-) diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 9bdfa7c4ea..1f80c5e317 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -1469,22 +1469,27 @@ def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]: Returns: An async context manager for the streamable HTTP client transport. """ + from httpx import AsyncClient from mcp.client.streamable_http import streamable_http_client http_client = self._httpx_client if self._header_provider is not None: if http_client is None: - from mcp.shared._httpx_utils import create_mcp_http_client - - http_client = create_mcp_http_client() + http_client = AsyncClient( + follow_redirects=True, + timeout=httpx.Timeout(MCP_DEFAULT_TIMEOUT, read=MCP_DEFAULT_SSE_READ_TIMEOUT), + ) self._httpx_client = http_client - async def _inject_headers(request: httpx.Request) -> None: # noqa: RUF029 - headers = _mcp_call_headers.get({}) - for key, value in headers.items(): - request.headers[key] = value + if not hasattr(self, "_inject_headers_hook"): + + async def _inject_headers(request: httpx.Request) -> None: # noqa: RUF029 + headers = _mcp_call_headers.get({}) + for key, value in headers.items(): + request.headers[key] = value - http_client.event_hooks["request"].append(_inject_headers) + self._inject_headers_hook = _inject_headers # type: ignore[attr-defined] + http_client.event_hooks["request"].append(self._inject_headers_hook) # type: ignore[attr-defined] return streamable_http_client( url=self.url, diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index 64ee6ca690..cf53bd658f 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -3804,6 +3804,28 @@ def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]: assert meta is None +async def test_mcp_streamable_http_tool_hook_not_duplicated_on_repeated_get_mcp_client(): + """Test that calling get_mcp_client multiple times does not accumulate duplicate hooks.""" + tool = MCPStreamableHTTPTool( + name="test", + url="http://example.com/mcp", + header_provider=lambda kw: {"X-Token": kw.get("token", "")}, + ) + + try: + with patch("agent_framework._mcp.streamable_http_client"): + tool.get_mcp_client() + tool.get_mcp_client() + tool.get_mcp_client() + + assert tool._httpx_client is not None + hooks = tool._httpx_client.event_hooks.get("request", []) + assert len(hooks) == 1, f"Expected exactly one hook, got {len(hooks)}" + finally: + if getattr(tool, "_httpx_client", None) is not None: + await tool._httpx_client.aclose() + + # endregion @@ -3811,10 +3833,10 @@ def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]: async def test_mcp_streamable_http_tool_header_provider_injects_headers(): - """Test that header_provider injects per-request HTTP headers via runtime kwargs. + """Test that header_provider integrates with call_tool via runtime kwargs. When header_provider is configured, runtime kwargs from FunctionInvocationContext - are passed to the provider and the returned headers appear on outbound HTTP requests. + are passed to the provider and the MCP session.call_tool is invoked successfully. """ class _TestServer(MCPStreamableHTTPTool): @@ -3999,7 +4021,7 @@ async def test_mcp_streamable_http_tool_header_provider_with_httpx_event_hook(): """Test that the httpx event hook injects headers from the contextvar.""" import httpx - from agent_framework._mcp import _mcp_call_headers + from agent_framework._mcp import MCP_DEFAULT_SSE_READ_TIMEOUT, MCP_DEFAULT_TIMEOUT, _mcp_call_headers tool = MCPStreamableHTTPTool( name="test", @@ -4007,23 +4029,31 @@ async def test_mcp_streamable_http_tool_header_provider_with_httpx_event_hook(): header_provider=lambda kw: {"X-Custom": kw.get("custom", "")}, ) - with patch("agent_framework._mcp.streamable_http_client"): - # Trigger get_mcp_client to set up the event hook - tool.get_mcp_client() - - # The tool should have created an httpx client with the event hook - assert tool._httpx_client is not None - hooks = tool._httpx_client.event_hooks.get("request", []) - assert len(hooks) == 1, "Expected one request event hook" - - # Simulate what happens during a call_tool: contextvar is set - token = _mcp_call_headers.set({"X-Custom": "test-value"}) - try: - request = httpx.Request("POST", "http://example.com/mcp") - await hooks[0](request) - assert request.headers.get("X-Custom") == "test-value" - finally: - _mcp_call_headers.reset(token) + try: + with patch("agent_framework._mcp.streamable_http_client"): + # Trigger get_mcp_client to set up the event hook + tool.get_mcp_client() + + # The tool should have created an httpx client with the event hook + assert tool._httpx_client is not None + assert tool._httpx_client.follow_redirects is True + assert tool._httpx_client.timeout.connect == MCP_DEFAULT_TIMEOUT + assert tool._httpx_client.timeout.read == MCP_DEFAULT_SSE_READ_TIMEOUT + hooks = tool._httpx_client.event_hooks.get("request", []) + assert len(hooks) == 1, "Expected one request event hook" + + # Simulate what happens during a call_tool: contextvar is set + token = _mcp_call_headers.set({"X-Custom": "test-value"}) + try: + request = httpx.Request("POST", "http://example.com/mcp") + await hooks[0](request) + assert request.headers.get("X-Custom") == "test-value" + finally: + _mcp_call_headers.reset(token) + finally: + # Ensure any created httpx client is properly closed + if getattr(tool, "_httpx_client", None) is not None: + await tool._httpx_client.aclose() async def test_mcp_streamable_http_tool_header_provider_with_user_httpx_client(): From f344bfea4ff33b4ed10ab2aab4b83cbaaebba6eb Mon Sep 17 00:00:00 2001 From: Copilot Date: Mon, 23 Mar 2026 12:29:41 +0000 Subject: [PATCH 03/11] Add test for header_provider via FunctionTool.invoke with FunctionInvocationContext Addresses PR review comment: exercises the full pipeline from FunctionInvocationContext.kwargs through FunctionTool.invoke to MCPStreamableHTTPTool.call_tool and header_provider, rather than testing call_tool in isolation. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- python/packages/core/tests/core/test_mcp.py | 86 +++++++++++++++++++++ 1 file changed, 86 insertions(+) diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index cf53bd658f..287bf0b334 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -4091,4 +4091,90 @@ async def test_mcp_streamable_http_tool_header_provider_with_user_httpx_client() await user_client.aclose() +async def test_mcp_streamable_http_tool_header_provider_via_invoke_with_context(): + """Test that header_provider receives kwargs when invoked through FunctionTool.invoke with FunctionInvocationContext. + + This exercises the full pipeline: FunctionInvocationContext.kwargs -> FunctionTool.invoke + -> MCPStreamableHTTPTool.call_tool -> header_provider. + """ + from agent_framework._mcp import _mcp_call_headers + + observed_headers: list[dict[str, str]] = [] + original_call_tool = MCPStreamableHTTPTool.call_tool + + async def spy_call_tool(self, tool_name, **kwargs): + # Capture the contextvar value set by call_tool before delegating + result = await original_call_tool(self, tool_name, **kwargs) + try: + observed_headers.append(_mcp_call_headers.get()) + except LookupError: + observed_headers.append({}) + return result + + class _TestServer(MCPStreamableHTTPTool): + async def connect(self): + self.session = Mock(spec=ClientSession) + self.session.list_tools = AsyncMock( + return_value=types.ListToolsResult( + tools=[ + types.Tool( + name="greet", + description="Says hello", + inputSchema={ + "type": "object", + "properties": {"name": {"type": "string"}}, + "required": ["name"], + }, + ) + ] + ) + ) + self.session.call_tool = AsyncMock( + return_value=types.CallToolResult(content=[types.TextContent(type="text", text="Hello!")]) + ) + self.session.send_ping = AsyncMock() + self.is_connected = True + + def get_mcp_client(self): + return None + + provider_received: list[dict] = [] + + def provider(kwargs): + provider_received.append(dict(kwargs)) + return {"X-Some-Token": kwargs.get("some_token", "")} + + server = _TestServer( + name="test", + url="http://example.com/mcp", + header_provider=provider, + ) + async with server: + await server.load_tools() + func = server.functions[0] + + # Build a FunctionInvocationContext with runtime kwargs, as the agent framework would + context = FunctionInvocationContext( + function=func, + arguments={"name": "Alice"}, + kwargs={"some_token": "my-secret"}, + ) + + with patch.object(MCPStreamableHTTPTool, "call_tool", spy_call_tool): + result = await func.invoke(arguments={"name": "Alice"}, context=context) + + # Verify the invoke produced a result + assert isinstance(result, list) + assert result[0].text == "Hello!" + + # Verify header_provider was called with the runtime kwargs + assert len(provider_received) == 1 + assert provider_received[0]["some_token"] == "my-secret" + + # Verify session.call_tool was called with the tool arguments (not the runtime kwargs) + server.session.call_tool.assert_called_once() + call_args = server.session.call_tool.call_args + assert call_args.kwargs.get("arguments", {}).get("name") == "Alice" + + # endregion From ad8939bf8c57b59fbd0d311a5d91ceafe0e85a76 Mon Sep 17 00:00:00 2001 From: Copilot Date: Mon, 23 Mar 2026 12:31:39 +0000 Subject: [PATCH 04/11] Address review feedback for #4808: review comment fixes --- python/packages/core/tests/core/test_mcp.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index 287bf0b334..72b219b7a6 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -4092,7 +4092,7 @@ async def test_mcp_streamable_http_tool_header_provider_with_user_httpx_client() async def test_mcp_streamable_http_tool_header_provider_via_invoke_with_context(): - """Test that header_provider receives kwargs when invoked through FunctionTool.invoke with FunctionInvocationContext. + """Test that header_provider receives kwargs via FunctionTool.invoke with FunctionInvocationContext. This exercises the full pipeline: FunctionInvocationContext.kwargs -> FunctionTool.invoke -> MCPStreamableHTTPTool.call_tool -> header_provider. From b5ce44ad938bfb9eaaab9d57dd43dfc54b806b5e Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Mon, 30 Mar 2026 08:54:52 +0200 Subject: [PATCH 05/11] Fix streamable MCP transport defaults Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- python/packages/core/agent_framework/_mcp.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 1f80c5e317..5bee5c3b57 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -61,6 +61,8 @@ class MCPSpecificApproval(TypedDict, total=False): _MCP_REMOTE_NAME_KEY = "_mcp_remote_name" _MCP_NORMALIZED_NAME_KEY = "_mcp_normalized_name" _mcp_call_headers: contextvars.ContextVar[dict[str, str]] = contextvars.ContextVar("_mcp_call_headers") +MCP_DEFAULT_TIMEOUT = 30 +MCP_DEFAULT_SSE_READ_TIMEOUT = 60 * 5 # region: Helpers @@ -139,6 +141,13 @@ def _inject_otel_into_mcp_meta(meta: dict[str, Any] | None = None) -> dict[str, return meta +def streamable_http_client(*args: Any, **kwargs: Any) -> _AsyncGeneratorContextManager[Any, None]: + """Lazily import the MCP streamable HTTP transport.""" + from mcp.client.streamable_http import streamable_http_client as _streamable_http_client + + return cast(_AsyncGeneratorContextManager[Any, None], _streamable_http_client(*args, **kwargs)) + + # region: MCP Plugin @@ -1469,21 +1478,20 @@ def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]: Returns: An async context manager for the streamable HTTP client transport. """ - from httpx import AsyncClient - from mcp.client.streamable_http import streamable_http_client + from httpx import AsyncClient, Request, Timeout http_client = self._httpx_client if self._header_provider is not None: if http_client is None: http_client = AsyncClient( follow_redirects=True, - timeout=httpx.Timeout(MCP_DEFAULT_TIMEOUT, read=MCP_DEFAULT_SSE_READ_TIMEOUT), + timeout=Timeout(MCP_DEFAULT_TIMEOUT, read=MCP_DEFAULT_SSE_READ_TIMEOUT), ) self._httpx_client = http_client if not hasattr(self, "_inject_headers_hook"): - async def _inject_headers(request: httpx.Request) -> None: # noqa: RUF029 + async def _inject_headers(request: Request) -> None: # noqa: RUF029 headers = _mcp_call_headers.get({}) for key, value in headers.items(): request.headers[key] = value From 1ae713c60c2269ded0d00266fc1179290c45af19 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Mon, 30 Mar 2026 09:18:18 +0200 Subject: [PATCH 06/11] Fix Azure AI test client mocks Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- python/packages/azure-ai/tests/test_azure_ai_client.py | 2 +- python/packages/azure-ai/tests/test_provider.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/python/packages/azure-ai/tests/test_azure_ai_client.py b/python/packages/azure-ai/tests/test_azure_ai_client.py index f3e459d0a4..d54db07fab 100644 --- a/python/packages/azure-ai/tests/test_azure_ai_client.py +++ b/python/packages/azure-ai/tests/test_azure_ai_client.py @@ -77,7 +77,7 @@ def mock_project_client() -> MagicMock: mock_client.telemetry.get_application_insights_connection_string = AsyncMock() # Mock get_openai_client method - mock_client.get_openai_client = AsyncMock() + mock_client.get_openai_client = MagicMock() # Mock close method mock_client.close = AsyncMock() diff --git a/python/packages/azure-ai/tests/test_provider.py b/python/packages/azure-ai/tests/test_provider.py index bc85948fca..912b8625e7 100644 --- a/python/packages/azure-ai/tests/test_provider.py +++ b/python/packages/azure-ai/tests/test_provider.py @@ -34,7 +34,7 @@ def mock_project_client() -> MagicMock: mock_client.telemetry.get_application_insights_connection_string = AsyncMock() # Mock get_openai_client method - mock_client.get_openai_client = AsyncMock() + mock_client.get_openai_client = MagicMock() # Mock close method mock_client.close = AsyncMock() From db146347358fd3249f6cf3cbcc869e52d6c852c2 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Mon, 30 Mar 2026 09:25:05 +0200 Subject: [PATCH 07/11] Fix MCP runtime kwarg regressions Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- python/packages/core/agent_framework/_mcp.py | 23 ++++++++++++++++++-- 1 file changed, 21 insertions(+), 2 deletions(-) diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 5bee5c3b57..7a2833c391 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -143,7 +143,18 @@ def _inject_otel_into_mcp_meta(meta: dict[str, Any] | None = None) -> dict[str, def streamable_http_client(*args: Any, **kwargs: Any) -> _AsyncGeneratorContextManager[Any, None]: """Lazily import the MCP streamable HTTP transport.""" - from mcp.client.streamable_http import streamable_http_client as _streamable_http_client + try: + from mcp.client.streamable_http import streamable_http_client as _streamable_http_client + except ModuleNotFoundError as ex: + missing_name = ex.name or str(ex) + if missing_name == "mcp" or missing_name.startswith("mcp."): + raise ModuleNotFoundError("`MCPStreamableHTTPTool` requires `mcp`. Please install `mcp`.") from ex + if "mcp" in missing_name: + raise ModuleNotFoundError("`MCPStreamableHTTPTool` requires `mcp`. Please install `mcp`.") from ex + raise ModuleNotFoundError( + f"`MCPStreamableHTTPTool` requires streamable HTTP transport support. " + f"The optional dependency `{missing_name}` is not installed. Please update your dependencies." + ) from ex return cast(_AsyncGeneratorContextManager[Any, None], _streamable_http_client(*args, **kwargs)) @@ -962,9 +973,17 @@ async def load_tools(self) -> None: input_schema = dict(tool.inputSchema or {}) if input_schema.get("type") == "object" and "properties" not in input_schema: input_schema["properties"] = {} + + async def _call_tool_with_runtime_kwargs( + *, + _remote_tool_name: str = tool.name, + **kwargs: Any, + ) -> str | list[Content]: + return await self.call_tool(_remote_tool_name, **kwargs) + # Create FunctionTools out of each tool func: FunctionTool = FunctionTool( - func=partial(self.call_tool, tool.name), + func=_call_tool_with_runtime_kwargs, name=local_name, description=tool.description or "", approval_mode=approval_mode, From eb2884b8dd202edf3dda5cc6b523b5166e403221 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Mon, 30 Mar 2026 09:29:22 +0200 Subject: [PATCH 08/11] Stabilize MCP tool runtime kwargs Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- python/packages/core/agent_framework/_mcp.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 7a2833c391..a8f96f5b25 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -993,6 +993,7 @@ async def _call_tool_with_runtime_kwargs( _MCP_NORMALIZED_NAME_KEY: normalized_name, }, ) + func.__dict__["_forward_runtime_kwargs"] = True self._functions.append(func) existing_names.add(local_name) From 1e9cf83ee430cbe4385513bc15867a3d21187578 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Mon, 30 Mar 2026 09:36:06 +0200 Subject: [PATCH 09/11] Use context kwargs in MCP wrappers Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- python/packages/core/agent_framework/_mcp.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index a8f96f5b25..a9110e8749 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -39,6 +39,7 @@ from mcp.shared.session import RequestResponder from ._clients import SupportsChatGetResponse + from ._middleware import FunctionInvocationContext logger = logging.getLogger(__name__) @@ -975,11 +976,14 @@ async def load_tools(self) -> None: input_schema["properties"] = {} async def _call_tool_with_runtime_kwargs( + ctx: FunctionInvocationContext, *, _remote_tool_name: str = tool.name, **kwargs: Any, ) -> str | list[Content]: - return await self.call_tool(_remote_tool_name, **kwargs) + call_kwargs = dict(ctx.kwargs) + call_kwargs.update(kwargs) + return await self.call_tool(_remote_tool_name, **call_kwargs) # Create FunctionTools out of each tool func: FunctionTool = FunctionTool( @@ -993,7 +997,6 @@ async def _call_tool_with_runtime_kwargs( _MCP_NORMALIZED_NAME_KEY: normalized_name, }, ) - func.__dict__["_forward_runtime_kwargs"] = True self._functions.append(func) existing_names.add(local_name) From ed5f37240dc77883740492d0ce2bec656c93f3c3 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Mon, 30 Mar 2026 10:21:21 +0200 Subject: [PATCH 10/11] updated mcp samples --- python/samples/02-agents/mcp/README.md | 4 +- .../02-agents/mcp/agent_as_mcp_server.py | 4 +- .../samples/02-agents/mcp/mcp_api_key_auth.py | 73 +++++++++---------- .../samples/02-agents/mcp/mcp_github_pat.py | 4 +- 4 files changed, 41 insertions(+), 44 deletions(-) diff --git a/python/samples/02-agents/mcp/README.md b/python/samples/02-agents/mcp/README.md index e07d63ddbd..149ba19c0c 100644 --- a/python/samples/02-agents/mcp/README.md +++ b/python/samples/02-agents/mcp/README.md @@ -11,7 +11,7 @@ The Model Context Protocol (MCP) is an open standard for connecting AI agents to | Sample | File | Description | |--------|------|-------------| | **Agent as MCP Server** | [`agent_as_mcp_server.py`](agent_as_mcp_server.py) | Shows how to expose an Agent Framework agent as an MCP server that other AI applications can connect to | -| **API Key Authentication** | [`mcp_api_key_auth.py`](mcp_api_key_auth.py) | Demonstrates API key authentication with MCP servers | +| **API Key Authentication** | [`mcp_api_key_auth.py`](mcp_api_key_auth.py) | Demonstrates API key authentication with MCP servers using `header_provider`, runtime invocation kwargs, and a command-line API key argument | | **GitHub Integration with PAT** | [`mcp_github_pat.py`](mcp_github_pat.py) | Demonstrates connecting to GitHub's MCP server using Personal Access Token (PAT) authentication | ## Prerequisites @@ -19,5 +19,7 @@ The Model Context Protocol (MCP) is an open standard for connecting AI agents to - `OPENAI_API_KEY` environment variable - `OPENAI_RESPONSES_MODEL` environment variable +Run `mcp_api_key_auth.py` with the MCP API key as the first command-line argument. + For `mcp_github_pat.py`: - `GITHUB_PAT` - Your GitHub Personal Access Token (create at https://github.com/settings/tokens) diff --git a/python/samples/02-agents/mcp/agent_as_mcp_server.py b/python/samples/02-agents/mcp/agent_as_mcp_server.py index 97cbcf3f75..dd92c240fb 100644 --- a/python/samples/02-agents/mcp/agent_as_mcp_server.py +++ b/python/samples/02-agents/mcp/agent_as_mcp_server.py @@ -4,7 +4,7 @@ import anyio from agent_framework import Agent, tool -from agent_framework.openai import OpenAIResponsesClient +from agent_framework.openai import OpenAIChatClient from dotenv import load_dotenv # Load environment variables from .env file @@ -57,7 +57,7 @@ async def run() -> None: # Define an agent # Agent's name and description provide better context for AI model agent = Agent( - client=OpenAIResponsesClient(), + client=OpenAIChatClient(), name="RestaurantAgent", description="Answer questions about the menu.", tools=[get_specials, get_item_price], diff --git a/python/samples/02-agents/mcp/mcp_api_key_auth.py b/python/samples/02-agents/mcp/mcp_api_key_auth.py index 16748a593c..456db2878c 100644 --- a/python/samples/02-agents/mcp/mcp_api_key_auth.py +++ b/python/samples/02-agents/mcp/mcp_api_key_auth.py @@ -1,20 +1,31 @@ # Copyright (c) Microsoft. All rights reserved. import asyncio -import os +import sys from agent_framework import Agent, MCPStreamableHTTPTool -from agent_framework.openai import OpenAIResponsesClient +from agent_framework.openai import OpenAIChatClient from dotenv import load_dotenv -from httpx import AsyncClient # Load environment variables from .env file load_dotenv() """ -MCP Authentication Example +MCP API Key Authentication Example -This example demonstrates how to authenticate with MCP servers using API key headers. +This sample demonstrates the runtime ``header_provider`` pattern for +``MCPStreamableHTTPTool``. The MCP tool derives authentication headers from +``function_invocation_kwargs`` passed to ``Agent.run(...)`` so the API key stays +in runtime context instead of being baked into a shared ``httpx.AsyncClient``. + +Replace the ``url`` parameter in the ``MCPStreamableHTTPTool`` with your authenticated server URL and +run the sample with your API key as a command-line argument: + python mcp_api_key_auth.py + +The ``header_provider`` here is just a simple lambda, but it can be a more complex function that retrieves and +formats headers as needed, allowing for flexible authentication schemes. +For more complex scenarios, you could implement token refresh logic or support multiple authentication methods +within the header provider function. For more authentication examples including OAuth 2.0 flows, see: - https://github.com/modelcontextprotocol/python-sdk/tree/main/examples/clients/simple-auth-client @@ -22,44 +33,28 @@ """ -async def api_key_auth_example() -> None: - """Example of using API key authentication with MCP server.""" - # Configuration - mcp_server_url = os.getenv("MCP_SERVER_URL", "your-mcp-server-url") - api_key = os.getenv("MCP_API_KEY") - - # Create authentication headers - # Common patterns: - # - Bearer token: "Authorization": f"Bearer {api_key}" - # - API key header: "X-API-Key": api_key - # - Custom header: "Authorization": f"ApiKey {api_key}" - auth_headers = { - "Authorization": f"Bearer {api_key}", - } - - # Create HTTP client with authentication headers - http_client = AsyncClient(headers=auth_headers) +async def api_key_auth_example(api_key: str) -> None: + """Run an agent against an MCP server using runtime-provided API key headers.""" - # Create MCP tool with the configured HTTP client - async with ( - MCPStreamableHTTPTool( + async with Agent( + client=OpenAIChatClient(), + name="Agent", + instructions="You are a helpful assistant. Use your MCP tool when answering the user's question.", + tools=MCPStreamableHTTPTool( name="MCP tool", - description="MCP tool description", - url=mcp_server_url, - http_client=http_client, # Pass HTTP client with authentication headers - ) as mcp_tool, - Agent( - client=OpenAIResponsesClient(), - name="Agent", - instructions="You are a helpful assistant.", - tools=mcp_tool, - ) as agent, - ): - query = "What tools are available to you?" + description="MCP tool description.", + url="", + header_provider=lambda kwargs: {"Authorization": f"Bearer {kwargs['mcp_api_key']}"}, + ), + ) as agent: + query = "Use your MCP tool to tell me what tools are available to you." print(f"User: {query}") - result = await agent.run(query) + result = await agent.run( + query, + function_invocation_kwargs={"mcp_api_key": api_key}, + ) print(f"Agent: {result.text}") if __name__ == "__main__": - asyncio.run(api_key_auth_example()) + asyncio.run(api_key_auth_example(sys.argv[1])) diff --git a/python/samples/02-agents/mcp/mcp_github_pat.py b/python/samples/02-agents/mcp/mcp_github_pat.py index 63d70a344d..8c83d7c8e2 100644 --- a/python/samples/02-agents/mcp/mcp_github_pat.py +++ b/python/samples/02-agents/mcp/mcp_github_pat.py @@ -4,7 +4,7 @@ import os from agent_framework import Agent -from agent_framework.openai import OpenAIResponsesClient +from agent_framework.openai import OpenAIChatClient from dotenv import load_dotenv """ @@ -45,7 +45,7 @@ async def github_mcp_example() -> None: # 4. Create agent with the GitHub MCP tool using instance method # The MCP tool manages the connection to the MCP server and makes its tools available # Set approval_mode="never_require" to allow the MCP tool to execute without approval - client = OpenAIResponsesClient() + client = OpenAIChatClient() github_mcp_tool = client.get_mcp_tool( name="GitHub", url="https://api.githubcopilot.com/mcp/", From 2f3afadede7456ea354770c8aec15b175a3240a5 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Mon, 30 Mar 2026 10:30:40 +0200 Subject: [PATCH 11/11] fix link --- python/packages/core/agent_framework/_mcp.py | 6 ++---- .../samples/05-end-to-end/evaluation/red_teaming/README.md | 2 +- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index a9110e8749..e40bb46f00 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -148,16 +148,14 @@ def streamable_http_client(*args: Any, **kwargs: Any) -> _AsyncGeneratorContextM from mcp.client.streamable_http import streamable_http_client as _streamable_http_client except ModuleNotFoundError as ex: missing_name = ex.name or str(ex) - if missing_name == "mcp" or missing_name.startswith("mcp."): - raise ModuleNotFoundError("`MCPStreamableHTTPTool` requires `mcp`. Please install `mcp`.") from ex - if "mcp" in missing_name: + if missing_name == "mcp" or missing_name.startswith("mcp.") or "mcp" in missing_name: raise ModuleNotFoundError("`MCPStreamableHTTPTool` requires `mcp`. Please install `mcp`.") from ex raise ModuleNotFoundError( f"`MCPStreamableHTTPTool` requires streamable HTTP transport support. " f"The optional dependency `{missing_name}` is not installed. Please update your dependencies." ) from ex - return cast(_AsyncGeneratorContextManager[Any, None], _streamable_http_client(*args, **kwargs)) + return _streamable_http_client(*args, **kwargs) # type: ignore[return-value] # region: MCP Plugin diff --git a/python/samples/05-end-to-end/evaluation/red_teaming/README.md b/python/samples/05-end-to-end/evaluation/red_teaming/README.md index 39fda91ae4..9a9efaafe6 100644 --- a/python/samples/05-end-to-end/evaluation/red_teaming/README.md +++ b/python/samples/05-end-to-end/evaluation/red_teaming/README.md @@ -174,7 +174,7 @@ Overall Attack Success Rate: 7.2% - [Azure AI Evaluation SDK](https://learn.microsoft.com/azure/ai-foundry/how-to/develop/evaluate-sdk) - [Risk and Safety Evaluations](https://learn.microsoft.com/azure/ai-foundry/concepts/evaluation-metrics-built-in#risk-and-safety-evaluators) - [Azure AI Red Teaming Notebook](https://github.com/Azure-Samples/azureai-samples/blob/main/scenarios/evaluate/AI_RedTeaming/AI_RedTeaming.ipynb) -- [PyRIT - Python Risk Identification Toolkit](https://github.com/Azure/PyRIT) +- [PyRIT - Python Risk Identification Toolkit](https://github.com/microsoft/PyRIT) ## Troubleshooting