Python: fix(core): authenticate MCP connect with run-supplied header_provider kwargs

python/packages/core/agent_framework/_agents.py

@@ -1455,11 +1455,20 @@ async def _prepare_run_context(
         # Normalize tools
         normalized_tools = _normalize_tools(tools_)
 
+        # Extract additional function arguments
+        effective_function_invocation_kwargs = (
+            dict(function_invocation_kwargs) if function_invocation_kwargs is not None else {}
+        )
+        additional_function_arguments = {**effective_function_invocation_kwargs, **existing_additional_args}
+
         # Resolve final tool list (configured tools + runtime provided tools + local MCP server tools)
         final_tools = list(base_tools)
         for tool in normalized_tools:
             if isinstance(tool, MCPTool):
                 if not tool.is_connected:
+                    # The handshake and discovery requests are issued before any tool call, so the run's
+                    # kwargs must reach header_provider here or those requests go out unauthenticated.
+                    tool._seed_connection_kwargs(additional_function_arguments)  # pyright: ignore[reportPrivateUsage]
                     await self._async_exit_stack.enter_async_context(tool)
                 _append_unique_tools(
                     final_tools,
@@ -1471,18 +1480,14 @@ async def _prepare_run_context(
 
         for mcp_server in self.mcp_tools:
             if not mcp_server.is_connected:
+                mcp_server._seed_connection_kwargs(additional_function_arguments)  # pyright: ignore[reportPrivateUsage]
                 await self._async_exit_stack.enter_async_context(mcp_server)
             _append_unique_tools(
                 final_tools,
                 mcp_server.functions,
                 duplicate_error_message=mcp_duplicate_message,
             )
 
-        effective_function_invocation_kwargs = (
-            dict(function_invocation_kwargs) if function_invocation_kwargs is not None else {}
-        )
-        additional_function_arguments = {**effective_function_invocation_kwargs, **existing_additional_args}
-
         model = opts.pop("model", None)
 
         # Build options dict from run() options merged with provided options

python/packages/core/agent_framework/_mcp.py

@@ -1699,6 +1699,11 @@ async def _close_and_check_cancelled(self, ex: BaseException) -> tuple[bool, Bas
         error path holding only a bare cancellation can still describe it.
         """
         cleanup_error = await self._safe_close_exit_stack()
+        # Every abandoned connection attempt lands here, so a rejected handshake cannot
+        # leave one run's credentials visible to a later unseeded reconnect. Deliberately
+        # not in _safe_close_exit_stack: connect(reset=True) closes through that path and
+        # must keep its kwargs to re-authenticate the new connection.
+        self._release_connection_kwargs()
         return _should_propagate_cancelled_error(ex), cleanup_error
 
     def _reset_session_state(self) -> None:
@@ -1890,6 +1895,14 @@ async def _connect_on_owner(self, *, reset: bool = False, load_configured: bool
                 self._tool_param_names_by_name = param_names_before_discovery
             raise
 
+    def _seed_connection_kwargs(self, kwargs: Mapping[str, Any]) -> None:
+        """Offer run-scoped kwargs to connection-lifetime header resolution."""
+        return
+
+    def _release_connection_kwargs(self) -> None:
+        """Drop any run-scoped kwargs held for connection-lifetime header resolution."""
+        return
+
     async def _sampling_request_approved(self, params: types.CreateMessageRequestParams) -> bool:
         """Run the configured sampling approval gate.
 
@@ -3549,6 +3562,18 @@ def __init__(
                 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``.
+                Only tool calls carry a run's kwargs. Connection-lifetime requests - the
+                ``initialize`` handshake, tool and prompt discovery, and background pings -
+                belong to no call, so they reuse the kwargs of the run that established the
+                connection until the tool is closed; a later run's kwargs do not reach them.
+                A tool connected outside any run (eagerly via ``async with``, or standalone)
+                has no kwargs to reuse and the provider is called with an empty mapping, in
+                which case a ``KeyError`` from the provider is tolerated and the request is
+                sent without headers. Once a run has supplied kwargs, a ``KeyError`` is
+                raised instead, since a key missing there is a misconfiguration rather than
+                an unavoidable gap. A credential that must authenticate the handshake should
+                therefore come from somewhere the provider can read without a run - a closure
+                or a ``ContextVar`` - rather than from run kwargs alone.
                 The framework attaches these headers only to requests whose origin (scheme,
                 host, port) matches the configured ``url``, so they are not leaked to other
                 origins on cross-origin redirects; headers injected this way are also removed
@@ -3626,6 +3651,9 @@ def __init__(
         # when a header_provider is set: parallel invocations on the same instance would
         # otherwise overwrite each other's snapshot and attach the wrong per-call headers.
         self._active_call_headers: dict[str, str] | None = None
+        # None means no run seeded this connection, which an empty mapping cannot express:
+        # a run that supplies no kwargs still expects a missing provider key to be an error.
+        self._connection_kwargs: dict[str, Any] | None = None
         self._call_headers_lock = asyncio.Lock()
         self._header_request_owner = object()
         self._header_hook_client: AsyncClient | None = None
@@ -3689,22 +3717,24 @@ async def _inject_headers(request: Request) -> None:  # ruff:ignore[unused-async
                     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.
+                        # provider with the kwargs seeded by the run that established this
+                        # connection, so static providers and run-supplied credentials both
+                        # authenticate these requests. Provider failures propagate, matching the
+                        # call_tool path, except the one case below that no caller can avoid.
                         if self._header_provider is None:
                             raise RuntimeError("Header injection hook invoked without a header_provider.")
                         try:
-                            headers = self._header_provider({})
+                            headers = self._header_provider(self._connection_kwargs or {})
                         except KeyError:
-                            # A kwargs-dependent provider raises on every ambient request
-                            # (initialize, discovery, and recurring pings).
+                            # Unavoidable only when no run seeded this connection: the provider
+                            # wants per-call values a connection-lifetime request cannot have. Once
+                            # a run has seeded kwargs a missing key is a misconfiguration, and
+                            # silently dropping it would send the handshake unauthenticated.
+                            if self._connection_kwargs is not None:
+                                raise
                             logger.debug(
                                 "header_provider raised KeyError for MCP server %r on an ambient "
-                                "request (missing per-call kwargs); proceeding without headers.",
+                                "request (no connection kwargs available); proceeding without headers.",
                                 self.name,
                                 exc_info=True,
                             )
@@ -3760,8 +3790,22 @@ async def _close_on_owner(self) -> None:
         try:
             await super()._close_on_owner()
         finally:
+            self._release_connection_kwargs()
             self._remove_header_hook()
 
+    def _seed_connection_kwargs(self, kwargs: Mapping[str, Any]) -> None:
+        if self._header_provider is None or self.is_connected:
+            return
+        # is_connected stays false until initialize returns, so it alone would let a second
+        # concurrent run swap the credential out from under the first run's in-flight
+        # handshake. The claim is released when the connection closes or its setup fails.
+        if self._connection_kwargs is not None:
+            return
+        self._connection_kwargs = dict(kwargs)
+
+    def _release_connection_kwargs(self) -> None:
+        self._connection_kwargs = None
+
     async def call_tool(self, tool_name: str, **kwargs: Any) -> str | list[Content]:
         """Call a tool, injecting headers from the header_provider if configured.
 

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

@@ -21,6 +21,7 @@
 from pydantic import AnyUrl, BaseModel
 
 from agent_framework import (
+    Agent,
     ChatResponse,
     ChatResponseUpdate,
     Content,
@@ -31,6 +32,7 @@
     MCPStreamableHTTPTool,
     MCPWebsocketTool,
     Message,
+    SupportsChatGetResponse,
 )
 from agent_framework._feature_stage import _WARNED_FEATURES, ExperimentalFeature, ExperimentalWarning
 from agent_framework._mcp import (
@@ -7569,6 +7571,218 @@ def provider(kwargs):
         assert call_args.kwargs.get("arguments", {}).get("name") == "Alice"
 
 
+async def test_agent_run_supplies_mcp_connect_headers(
+    client: SupportsChatGetResponse,
+) -> None:
+    """Run-time credentials should authenticate implicit MCP initialization.
+
+    The agent receives function_invocation_kwargs before connecting the MCP tool
+    supplied to run(). This test exercises the real MCP transport against an
+    in-process mock HTTP endpoint and asserts that header_provider can use those
+    credentials on the initialize request, before any tool invocation occurs.
+    """
+    import httpx
+
+    captured_requests: list[tuple[str, str, dict[str, str]]] = []
+
+    async def handler(request: httpx.Request) -> httpx.Response:
+        if request.method == "DELETE":
+            return httpx.Response(200)
+        if request.method == "GET":
+            return httpx.Response(405)
+        body = json.loads(request.content.decode())
+        method = body.get("method", "")
+        captured_requests.append((request.method, method, {k.lower(): v for k, v in request.headers.items()}))
+        if method == "initialize":
+            result = {
+                "protocolVersion": body["params"]["protocolVersion"],
+                "capabilities": {"tools": {}},
+                "serverInfo": {"name": "mock-server", "version": "1.0.0"},
+            }
+            return httpx.Response(
+                200,
+                headers={"mcp-session-id": "test-session"},
+                json={"jsonrpc": "2.0", "id": body["id"], "result": result},
+            )
+        if method == "tools/list":
+            result = {
+                "tools": [
+                    {
+                        "name": "greet",
+                        "description": "Says hello",
+                        "inputSchema": {"type": "object", "properties": {"name": {"type": "string"}}},
+                    }
+                ]
+            }
+            return httpx.Response(200, json={"jsonrpc": "2.0", "id": body["id"], "result": result})
+        if method == "tools/call":
+            result = {"content": [{"type": "text", "text": "Hello!"}], "isError": False}
+            return httpx.Response(200, json={"jsonrpc": "2.0", "id": body["id"], "result": result})
+        if "id" in body:
+            # Any other request (e.g. ping) gets an empty result so the session doesn't block on it.
+            return httpx.Response(200, json={"jsonrpc": "2.0", "id": body["id"], "result": {}})
+        # Notifications (e.g. notifications/initialized)
+        return httpx.Response(202)
+
+    http_client = httpx.AsyncClient(transport=httpx.MockTransport(handler))
+    tool = MCPStreamableHTTPTool(
+        name="test",
+        url="http://127.0.0.1:8000/mcp",
+        load_prompts=False,
+        http_client=http_client,
+        header_provider=lambda kw: {"x-api-key": kw["api_key"]},  # failing scenario
+        # header_provider=lambda _: {"x-api-key": "connect-token"} # working scenario (no lambda)
+    )
+    try:
+        async with Agent(client=client) as agent:
+            # placement of tools, matters as we defer the resolution until the agent calls `run`
+            await agent.run("Hello", tools=[tool], function_invocation_kwargs={"api_key": "connect-token"})
+    finally:
+        await http_client.aclose()
+
+    initialize_headers = [headers for _, method, headers in captured_requests if method == "initialize"]
+    assert len(initialize_headers) == 1
+    assert initialize_headers[0].get("x-api-key") == "connect-token"
+
+
+async def test_agent_context_manager_authenticates_connect_with_closure_provider(
+    client: SupportsChatGetResponse,
+) -> None:
+    """A constructor-supplied MCP tool authenticates its eager handshake with no run involved.
+
+    Pins that ``header_provider`` already covers construction-time credentials: entering the
+    agent context connects before any run exists, and the server rejects unauthenticated calls.
+    """
+    import httpx
+
+    captured_requests: list[tuple[str, dict[str, str]]] = []
+
+    async def handler(request: httpx.Request) -> httpx.Response:
+        if request.method == "DELETE":
+            return httpx.Response(200)
+        if request.method == "GET":
+            return httpx.Response(405)
+        if request.headers.get("x-api-key") != "constructor-token":
+            return httpx.Response(401)
+        body = json.loads(request.content.decode())
+        method = body.get("method", "")
+        captured_requests.append((method, {k.lower(): v for k, v in request.headers.items()}))
+        if method == "initialize":
+            return httpx.Response(
+                200,
+                headers={"mcp-session-id": "test-session"},
+                json={
+                    "jsonrpc": "2.0",
+                    "id": body["id"],
+                    "result": {
+                        "protocolVersion": body["params"]["protocolVersion"],
+                        "capabilities": {"tools": {}},
+                        "serverInfo": {"name": "mock-server", "version": "1.0.0"},
+                    },
+                },
+            )
+        if method == "tools/list":
+            return httpx.Response(
+                200,
+                json={
+                    "jsonrpc": "2.0",
+                    "id": body["id"],
+                    "result": {"tools": [{"name": "greet", "inputSchema": {"type": "object", "properties": {}}}]},
+                },
+            )
+        if "id" in body:
+            return httpx.Response(200, json={"jsonrpc": "2.0", "id": body["id"], "result": {}})
+        return httpx.Response(202)
+
+    http_client = httpx.AsyncClient(transport=httpx.MockTransport(handler))
+    # The credential is known at construction, so the provider closes over it and ignores kwargs.
+    tool = MCPStreamableHTTPTool(
+        name="test",
+        url="http://127.0.0.1:8000/mcp",
+        load_prompts=False,
+        http_client=http_client,
+        header_provider=lambda _kwargs: {"x-api-key": "constructor-token"},
+    )
+    try:
+        async with Agent(client=client, tools=[tool]):
+            assert tool.is_connected
+    finally:
+        await http_client.aclose()
+
+    assert [method for method, _ in captured_requests].count("initialize") == 1
+    assert all(headers.get("x-api-key") == "constructor-token" for _, headers in captured_requests)
+
+
+async def test_constructor_supplied_mcp_tool_uses_run_credentials_on_lazy_connect(
+    client: SupportsChatGetResponse,
+) -> None:
+    """A constructor-supplied MCP tool connected at first run authenticates with that run's kwargs.
+
+    Without the agent context manager the handshake is deferred to ``run()``, so the run's
+    credentials are available and must reach ``header_provider``.
+    """
+    import httpx
+
+    captured_requests: list[tuple[str, dict[str, str]]] = []
+
+    async def handler(request: httpx.Request) -> httpx.Response:
+        if request.method == "DELETE":
+            return httpx.Response(200)
+        if request.method == "GET":
+            return httpx.Response(405)
+        if request.headers.get("x-api-key") != "run-token":
+            return httpx.Response(401)
+        body = json.loads(request.content.decode())
+        method = body.get("method", "")
+        captured_requests.append((method, {k.lower(): v for k, v in request.headers.items()}))
+        if method == "initialize":
+            return httpx.Response(
+                200,
+                headers={"mcp-session-id": "test-session"},
+                json={
+                    "jsonrpc": "2.0",
+                    "id": body["id"],
+                    "result": {
+                        "protocolVersion": body["params"]["protocolVersion"],
+                        "capabilities": {"tools": {}},
+                        "serverInfo": {"name": "mock-server", "version": "1.0.0"},
+                    },
+                },
+            )
+        if method == "tools/list":
+            return httpx.Response(
+                200,
+                json={
+                    "jsonrpc": "2.0",
+                    "id": body["id"],
+                    "result": {"tools": [{"name": "greet", "inputSchema": {"type": "object", "properties": {}}}]},
+                },
+            )
+        if "id" in body:
+            return httpx.Response(200, json={"jsonrpc": "2.0", "id": body["id"], "result": {}})
+        return httpx.Response(202)
+
+    http_client = httpx.AsyncClient(transport=httpx.MockTransport(handler))
+    tool = MCPStreamableHTTPTool(
+        name="test",
+        url="http://127.0.0.1:8000/mcp",
+        load_prompts=False,
+        http_client=http_client,
+        header_provider=lambda kw: {"x-api-key": kw["api_key"]},
+    )
+    agent = Agent(client=client, tools=[tool])
+    try:
+        # No agent context manager, so the tool is still unconnected when the run starts.
+        assert not tool.is_connected
+        await agent.run("Hello", function_invocation_kwargs={"api_key": "run-token"})
+    finally:
+        await tool.close()
+        await http_client.aclose()
+
+    assert [method for method, _ in captured_requests].count("initialize") == 1
+    assert all(headers.get("x-api-key") == "run-token" for _, headers in captured_requests)
+
+
 async def test_mcp_streamable_http_tool_header_provider_applies_across_transport_tasks():
     """Regression test for #7161: header_provider headers must reach tools/call requests.