diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index abdb0aa5d0..d58ea81080 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -2984,7 +2984,9 @@ def __init__( (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``. + agent middleware) without creating a separate ``httpx.AsyncClient``. Also + invoked with ``{}`` before the ``initialize`` handshake, since no per-call + kwargs exist yet at that point. task_options: Options for tools that advertise ``execution.taskSupport == "required"``. See :class:`MCPTaskOptions`. additional_tool_argument_names: Extra argument names to forward to the MCP server in @@ -3110,6 +3112,40 @@ async def call_tool(self, tool_name: str, **kwargs: Any) -> str | list[Content]: _mcp_call_headers.reset(token) return await super().call_tool(tool_name, **kwargs) + async def _connect_on_owner(self, *, reset: bool = False, load_configured: bool = True) -> None: + """Connect to the MCP server, applying header_provider headers to the initialize handshake. + + ``initialize`` (and any ``load_tools``/``load_prompts`` in the same connect pass) runs + before ``call_tool()`` ever does, so servers that require auth on ``initialize`` need + headers seeded here too, or they reject the connection before ``header_provider`` runs. + + Keyword Args: + reset: If True, forces a reconnection even if already connected. + load_configured: If True, loads tools and prompts according to the constructor flags. + """ + if self._header_provider is not None: + try: + # No per-call kwargs exist yet, so header_provider is invoked with {}; an + # implementation written only against call_tool()'s contract may not expect that. + headers = self._header_provider({}) + except Exception: + logger.warning( + "header_provider raised for %r during connect; proceeding without headers.", + self.name, + exc_info=True, + ) + headers = {} + # Must set/reset here, not in the connect()/__aenter__() caller: this can run in a + # different asyncio task (the MCP lifecycle-owner task), and a Token can only be + # reset in the context that created it. + token = _mcp_call_headers.set(headers) + try: + await super()._connect_on_owner(reset=reset, load_configured=load_configured) + finally: + _mcp_call_headers.reset(token) + return + await super()._connect_on_owner(reset=reset, load_configured=load_configured) + 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 4b61cd2451..39f556b725 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -5890,6 +5890,106 @@ def get_mcp_client(self): # pyrefly: ignore[bad-override] _mcp_call_headers.get() +async def test_mcp_streamable_http_tool_header_provider_applied_during_connect(): + """Test that header_provider headers are seeded before session.initialize() runs. + + Regression test: header_provider was previously only invoked from call_tool(), so the + initialize handshake issued during connect() went out with no headers, causing an + unconditional 401 against MCP servers that require auth on initialize. + """ + from agent_framework._mcp import _mcp_call_headers + + observed_headers: list[dict[str, str]] = [] + + tool = MCPStreamableHTTPTool( + name="test", + url="http://example.com/mcp", + header_provider=lambda kw: {"Authorization": "Bearer token-123"}, + ) + + 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) + tool.get_mcp_client = Mock(return_value=mock_context_manager) # type: ignore[method-assign] + + mock_session = Mock() + + async def fake_initialize(): + try: + observed_headers.append(_mcp_call_headers.get()) + except LookupError: + observed_headers.append({}) + return Mock(protocolVersion="2025-06-18", capabilities=None) + + mock_session.initialize = AsyncMock(side_effect=fake_initialize) + + with patch("mcp.client.session.ClientSession") as mock_session_class: + mock_session_class.return_value.__aenter__ = AsyncMock(return_value=mock_session) + mock_session_class.return_value.__aexit__ = AsyncMock(return_value=None) + + await tool.connect() + + # header_provider's output must have been visible to the initialize() call + assert observed_headers == [{"Authorization": "Bearer token-123"}] + + # contextvar must not leak past connect() + with pytest.raises(LookupError): + _mcp_call_headers.get() + + +async def test_mcp_streamable_http_tool_header_provider_error_during_connect_does_not_fail_connect(): + """Test that a header_provider written only for call_tool()'s contract doesn't break connect(). + + header_provider is now also invoked with ``{}`` at connect time. An implementation that + assumes kwargs is always populated (e.g. indexing a key only present on real tool calls) + would raise KeyError when called this way. That must be caught and logged, not allowed to + fail the connection. A genuine auth failure still surfaces via session.initialize() itself. + """ + from agent_framework._mcp import _mcp_call_headers + + observed_headers: list[dict[str, str]] = [] + + def unsafe_header_provider(kwargs: dict[str, object]) -> dict[str, str]: + return {"Authorization": f"Bearer {kwargs['mcp_api_key']}"} # type: ignore[index] + + tool = MCPStreamableHTTPTool( + name="test", + url="http://example.com/mcp", + header_provider=unsafe_header_provider, + ) + + 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) + tool.get_mcp_client = Mock(return_value=mock_context_manager) # type: ignore[method-assign] + + mock_session = Mock() + + async def fake_initialize(): + try: + observed_headers.append(_mcp_call_headers.get()) + except LookupError: + observed_headers.append({}) + return Mock(protocolVersion="2025-06-18", capabilities=None) + + mock_session.initialize = AsyncMock(side_effect=fake_initialize) + + with patch("mcp.client.session.ClientSession") as mock_session_class: + mock_session_class.return_value.__aenter__ = AsyncMock(return_value=mock_session) + mock_session_class.return_value.__aexit__ = AsyncMock(return_value=None) + + # Must not raise KeyError from unsafe_header_provider. + await tool.connect() + + # initialize() proceeded with no headers rather than the connection failing. + assert observed_headers == [{}] + + 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."""