Python: Apply header_provider headers to the MCP initialize handshake

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.

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."""