diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index d5b613efc15..21edadb96aa 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -227,7 +227,7 @@ class MCPSpecificApproval(TypedDict, total=False): "_meta", }) # object is a reference to the owner of the headers, so we keep track which headers belong to which tool -_mcp_call_headers: ContextVar[tuple[object, dict[str, str]]] = ContextVar("_mcp_call_headers") +_mcp_call_headers: ContextVar[tuple[object, dict[str, str]] | None] = ContextVar("_mcp_call_headers") _mcp_tool_runtime_context: ContextVar[tuple[object, Mapping[str, Any]] | None] = ContextVar( "_mcp_tool_runtime_context", default=None ) @@ -2010,10 +2010,14 @@ async def _ensure_lifecycle_owner(self) -> None: return self._lifecycle_queue = asyncio.Queue() - self._lifecycle_owner_task = asyncio.create_task( - self._run_lifecycle_owner(), - name=f"mcp-lifecycle:{self.name}", - ) + token = _mcp_call_headers.set(None) + try: + self._lifecycle_owner_task = asyncio.create_task( + self._run_lifecycle_owner(), + name=f"mcp-lifecycle:{self.name}", + ) + finally: + _mcp_call_headers.reset(token) async def _run_lifecycle_owner(self) -> None: queue = self._lifecycle_queue diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index fc836e17b3f..d2233286829 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -10740,6 +10740,35 @@ async def _pending_lifecycle_task_names() -> list[str]: return names +async def test_lifecycle_owner_does_not_inherit_call_headers(): + """Connection-lifetime work must not inherit headers from the call that starts it.""" + from agent_framework._mcp import _mcp_call_headers + + tool = MCPStreamableHTTPTool(name="context-check", url="https://mcp.example/mcp") + observed_call_headers: list[Any] = [] + observed_request_context: list[str | None] = [] + request_context = ContextVar[str]("request_context") + + async def record_owner_context(**_: Any) -> None: + observed_call_headers.append(_mcp_call_headers.get(None)) + observed_request_context.append(request_context.get(None)) + tool.is_connected = True + + call_headers_token = _mcp_call_headers.set((tool._header_request_owner, {"Authorization": "principal-a"})) + request_context_token = request_context.set("request-a") + try: + with patch.object(tool, "_connect_on_owner", new=record_owner_context): + await tool.connect() + + assert observed_call_headers == [None] + assert observed_request_context == ["request-a"] + finally: + request_context.reset(request_context_token) + _mcp_call_headers.reset(call_headers_token) + tool.is_connected = False + await tool.close() + + async def test_failed_connect_stops_the_lifecycle_owner(): """A connect that fails must not leave its owner task blocked on the queue forever.""" client = _unauthorized_http_client()