Python: Apply header_provider headers to the MCP initialize handshake#7080
Python: Apply header_provider headers to the MCP initialize handshake#7080efranken wants to merge 2 commits into
Conversation
MCPStreamableHTTPTool.header_provider was only invoked from call_tool(), so
servers that require auth on initialize() (e.g. Azure AI Search MCP endpoints)
rejected the connection with a 401 before header_provider ever ran. The same
gap affected load_tools()/load_prompts() calls issued during the same connect
pass.
Seed the _mcp_call_headers contextvar in an override of _connect_on_owner()
before delegating to the base implementation, mirroring the existing
call_tool() pattern. header_provider is invoked with {} at connect time since
no per-call kwargs exist yet; if that raises, the exception is logged and the
connection proceeds without headers rather than failing outright, since a
genuine auth failure still surfaces via the server's own rejection of
initialize().
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
@microsoft-github-policy-service agree |
There was a problem hiding this comment.
Pull request overview
This PR fixes an authentication gap in the Python MCP HTTP tool by ensuring MCPStreamableHTTPTool.header_provider is applied during the MCP initialize handshake (and related connect-time loading), not only during call_tool(). This enables connecting to MCP servers that require auth headers on initialize (e.g., Azure AI Search MCP endpoints) without failing early with 401s.
Changes:
- Seed the
_mcp_call_headersContextVarfromheader_provider({})for the duration of_connect_on_owner(), soinitialize()(and connect-timeload_tools/load_prompts) can use the same header injection mechanism ascall_tool(). - Make connect-time
header_provider({})best-effort: exceptions are caught/logged and the connection proceeds without injected headers. - Add regression tests verifying headers are visible during
initialize(), no contextvar leakage occurs afterconnect(), and a raisingheader_providerdoes not breakconnect().
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated no comments.
| File | Description |
|---|---|
| python/packages/core/agent_framework/_mcp.py | Apply header_provider-derived headers during connect/initialize by seeding _mcp_call_headers inside _connect_on_owner(). |
| python/packages/core/tests/core/test_mcp.py | Add regression tests covering connect-time header seeding, contextvar cleanup, and error-tolerant header_provider behavior. |
| # 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) |
There was a problem hiding this comment.
Could we avoid keeping the initialize headers in the transport task's inherited context? streamable_http_client() starts a long-lived post_writer during this await, so it retains header_provider({}); later per-call values set by call_tool() live in another task and never reach the HTTP hooks. With a shared tool, every caller can therefore execute under the connect-time principal instead of its own runtime credentials, crossing tenant authorization boundaries. Could initialize auth be supplied without pinning it into the writer while passing each call's headers explicitly to that task?
Motivation & Context
MCPStreamableHTTPTool.header_provideronly fires insidecall_tool(). But the MCPinitializehandshake happens earlier, duringconnect()/__aenter__()(_connect_on_owner), so servers that require auth oninitialize(e.g. Azure AI Search MCP endpoints) get an immediate 401 beforeheader_providerever runs. The same gap affectsload_tools()/load_prompts()when they run during the same connect pass.Description & Review Guide
Major changes
_connect_on_owneronMCPStreamableHTTPToolto seed the_mcp_call_headersContextVar fromheader_provider({})before delegating tosuper()._connect_on_owner(...), the same patterncall_tool()already uses.header_providerraises on an empty dict (e.g. an implementation written only againstcall_tool()'s contract, like the existingmcp_api_key_auth.pysample), the exception is caught and logged rather than failing the connection. A real auth failure still surfaces via the server's own rejection ofinitialize, this only guards the opportunistic best-effort call.initialize(), the ContextVar doesn't leak pastconnect(), and a raisingheader_providerdoesn't breakconnect().Impact
header_providernow also applies toinitialize,load_tools, andload_promptsduring connect, not justcall_tool(). No public API changes; existingheader_providerimplementations keep working unmodified.Outstanding question: Is seeding headers for the whole
_connect_on_ownercall correct, or should it be scoped tighter to justsession.initialize()?AI assistance was used after several hours of research to ensure fix is inline with existing patterns.
Related Issue
Fixes #7079
Contribution Checklist
breaking changelabel (or add "[BREAKING]" to the title prefix, before or after any language prefix) — a workflow keeps the label and title prefix in sync automatically.