Python: Apply header_provider headers to the MCP initialize handshake and other ambient requests

python/packages/core/agent_framework/_mcp.py

@@ -3124,8 +3124,36 @@ async def _inject_headers(request: Request) -> None:  # ruff:ignore[unused-async
                         return
                     # The transport may send this request from a task whose context was
                     # captured before call_tool set the ContextVar; fall back to the
-                    # instance-level snapshot of the active call's headers.
-                    headers = _mcp_call_headers.get({}) or self._active_call_headers or {}
+                    # instance-level snapshot of the active call's headers. Both are None
+                    # only when this is an ambient request outside call_tool; an active
+                    # call that legitimately produced no headers yields an empty dict and
+                    # must not trigger the ambient fallback below.
+                    headers = _mcp_call_headers.get(None)
+                    if headers is None:
+                        headers = self._active_call_headers
+                    if headers is None:
+                        # Ambient request made outside call_tool (the initialize handshake,
+                        # load_tools/load_prompts discovery, or background pings). Invoke the
+                        # provider with empty kwargs so static providers can authenticate these
+                        # requests too. A provider that indexes a required per-call kwarg (e.g.
+                        # kwargs["api_key"]) raises KeyError on the empty dict; that specific
+                        # case is tolerated so connect still succeeds. Any other error is a
+                        # genuine provider failure and is left to propagate, matching the
+                        # call_tool path which does not catch header_provider exceptions.
+                        if self._header_provider is None:
+                            raise RuntimeError("Header injection hook invoked without a header_provider.")
+                        try:
+                            headers = self._header_provider({})
+                        except KeyError:
+                            # A kwargs-dependent provider raises on every ambient request
+                            # (initialize, discovery, and recurring pings).
+                            logger.debug(
+                                "header_provider raised KeyError for MCP server %r on an ambient "
+                                "request (missing per-call kwargs); proceeding without headers.",
+                                self.name,
+                                exc_info=True,
+                            )
+                            headers = {}
                     for key, value in headers.items():
                         request.headers[key] = value
 

python/packages/core/tests/core/test_mcp.py

@@ -6027,6 +6027,187 @@ async def test_mcp_streamable_http_tool_header_provider_with_httpx_event_hook():
             await tool._httpx_client.aclose()  # type: ignore[union-attr]
 
 
+async def test_mcp_streamable_http_tool_header_provider_injects_on_ambient_request():
+    """Regression test for #7079: static header_provider must authenticate ambient requests.
+
+    The initialize handshake, load_tools/load_prompts discovery, and background pings run
+    outside call_tool, so the contextvar/snapshot are unset. A static header_provider should
+    still be invoked (with empty kwargs) so these requests carry auth headers.
+    """
+    import httpx
+
+    tool = MCPStreamableHTTPTool(
+        name="test",
+        url="http://example.com/mcp",
+        header_provider=lambda kw: {"Authorization": "******"},
+    )
+
+    try:
+        with patch("agent_framework._mcp.streamable_http_client"):
+            tool.get_mcp_client()
+
+            assert tool._httpx_client is not None
+            hooks = tool._httpx_client.event_hooks.get("request", [])
+            assert len(hooks) == 1
+
+            # No contextvar set and no active call snapshot: simulates the initialize handshake.
+            request = httpx.Request("POST", "http://example.com/mcp")
+            await hooks[0](request)
+            assert request.headers.get("Authorization") == "******"
+    finally:
+        if getattr(tool, "_httpx_client", None) is not None:
+            await tool._httpx_client.aclose()  # type: ignore[union-attr]
+
+
+async def test_mcp_streamable_http_tool_header_provider_ambient_request_tolerates_kwargs_provider():
+    """A header_provider that requires per-call kwargs must not crash ambient requests.
+
+    Providers that index runtime kwargs (e.g. kw["mcp_api_key"]) which are absent at connect
+    time raise KeyError. The hook should swallow that specific error and proceed without headers
+    rather than failing the initialize handshake.
+    """
+    import httpx
+
+    tool = MCPStreamableHTTPTool(
+        name="test",
+        url="http://example.com/mcp",
+        header_provider=lambda kw: {"Authorization": f"Bearer {kw['mcp_api_key']}"},
+    )
+
+    try:
+        with patch("agent_framework._mcp.streamable_http_client"):
+            tool.get_mcp_client()
+
+            assert tool._httpx_client is not None
+            hooks = tool._httpx_client.event_hooks.get("request", [])
+            assert len(hooks) == 1
+
+            # No kwargs available at connect time -> provider raises KeyError -> hook swallows it.
+            request = httpx.Request("POST", "http://example.com/mcp")
+            await hooks[0](request)
+            assert "Authorization" not in request.headers
+    finally:
+        if getattr(tool, "_httpx_client", None) is not None:
+            await tool._httpx_client.aclose()  # type: ignore[union-attr]
+
+
+async def test_mcp_streamable_http_tool_header_provider_empty_active_call_skips_ambient_fallback():
+    """Regression test: an active call that yields no headers must not trigger the ambient fallback.
+
+    When header_provider legitimately returns {} for a real call_tool invocation, the empty dict
+    is a valid "set" value. The hook must treat it as set-but-empty (leave the request unchanged)
+    rather than as "unset", which would re-invoke header_provider({}) mid-call and inject headers
+    the caller deliberately omitted.
+    """
+    import httpx
+
+    from agent_framework._mcp import _mcp_call_headers
+
+    call_count = 0
+
+    def provider(kw: dict[str, Any]) -> dict[str, str]:
+        nonlocal call_count
+        call_count += 1
+        # Non-empty for ambient ({}), empty for the active call below.
+        return {} if kw.get("suppress") else {"Authorization": "******"}
+
+    tool = MCPStreamableHTTPTool(name="test", url="http://example.com/mcp", header_provider=provider)
+
+    try:
+        with patch("agent_framework._mcp.streamable_http_client"):
+            tool.get_mcp_client()
+
+            assert tool._httpx_client is not None
+            hooks = tool._httpx_client.event_hooks.get("request", [])
+            assert len(hooks) == 1
+
+            # Simulate an active call whose provider returned {} (both ContextVar and snapshot set).
+            token = _mcp_call_headers.set({})
+            tool._active_call_headers = {}
+            try:
+                call_count = 0
+                request = httpx.Request("POST", "http://example.com/mcp")
+                await hooks[0](request)
+                assert "Authorization" not in request.headers
+                assert call_count == 0, "ambient fallback must not run during a set-but-empty call"
+            finally:
+                tool._active_call_headers = None
+                _mcp_call_headers.reset(token)
+    finally:
+        if getattr(tool, "_httpx_client", None) is not None:
+            await tool._httpx_client.aclose()  # type: ignore[union-attr]
+
+
+async def test_mcp_streamable_http_tool_header_provider_ambient_kwarg_error_is_benign(caplog):
+    """A kwargs-dependent provider that raises KeyError on ambient requests must not fail them.
+
+    Ambient requests (initialize/discovery plus recurring pings) call the provider with empty
+    kwargs, so a provider indexing a required per-call kwarg raises KeyError. That is tolerated:
+    the request proceeds without headers and only a DEBUG line is logged (no WARNING spam).
+    """
+    import logging
+
+    import httpx
+
+    tool = MCPStreamableHTTPTool(
+        name="test",
+        url="http://example.com/mcp",
+        header_provider=lambda kw: {"Authorization": "Bearer " + kw["mcp_api_key"]},
+    )
+
+    try:
+        with patch("agent_framework._mcp.streamable_http_client"):
+            tool.get_mcp_client()
+
+            assert tool._httpx_client is not None
+            hooks = tool._httpx_client.event_hooks.get("request", [])
+            assert len(hooks) == 1
+
+            with caplog.at_level(logging.DEBUG, logger="agent_framework._mcp"):
+                for _ in range(3):
+                    request = httpx.Request("POST", "http://example.com/mcp")
+                    await hooks[0](request)
+                    assert "Authorization" not in request.headers
+
+            # The benign kwargs case must not escalate to WARNING/ERROR.
+            assert not [r for r in caplog.records if r.levelno >= logging.WARNING]
+    finally:
+        if getattr(tool, "_httpx_client", None) is not None:
+            await tool._httpx_client.aclose()  # type: ignore[union-attr]
+
+
+async def test_mcp_streamable_http_tool_header_provider_ambient_non_keyerror_propagates():
+    """A genuine header_provider failure on an ambient request must surface, not be swallowed.
+
+    Only the missing-kwargs case (KeyError) is tolerated. Any other error - e.g. a token-refresh
+    failure or a provider bug - propagates so it is not silently converted into unauthenticated
+    traffic, matching the call_tool path which does not catch header_provider exceptions.
+    """
+    import httpx
+
+    class TokenRefreshError(RuntimeError):
+        pass
+
+    def failing_provider(kw: dict[str, Any]) -> dict[str, str]:
+        raise TokenRefreshError("token endpoint unreachable")
+
+    tool = MCPStreamableHTTPTool(name="test", url="http://example.com/mcp", header_provider=failing_provider)
+
+    try:
+        with patch("agent_framework._mcp.streamable_http_client"):
+            tool.get_mcp_client()
+
+            assert tool._httpx_client is not None
+            hooks = tool._httpx_client.event_hooks.get("request", [])
+            assert len(hooks) == 1
+
+            with pytest.raises(TokenRefreshError):
+                await hooks[0](httpx.Request("POST", "http://example.com/mcp"))
+    finally:
+        if getattr(tool, "_httpx_client", None) is not None:
+            await tool._httpx_client.aclose()  # type: ignore[union-attr]
+
+
 async def test_mcp_streamable_http_tool_header_provider_skips_cross_origin_redirect():
     """The request hook must not re-add caller headers after a cross-origin redirect."""
     import httpx