diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 9071deafcf3..d0dc62e91cd 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/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 diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index 4890c0eccb4..028cac027b3 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/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