From 55f8deeeb88b06f06ed4cb5802909e9bf6f1f8a6 Mon Sep 17 00:00:00 2001 From: Jose Alvarez Date: Fri, 16 Oct 2026 11:21:19 +0200 Subject: [PATCH 1/3] Isolate MCP lifecycle task context --- python/packages/core/agent_framework/_mcp.py | 5 +++-- python/packages/core/tests/core/test_mcp.py | 23 ++++++++++++++++++++ 2 files changed, 26 insertions(+), 2 deletions(-) diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index d5b613efc15..511976dca93 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -12,7 +12,7 @@ from abc import abstractmethod from collections.abc import Callable, Collection, Coroutine, Mapping, Sequence from contextlib import AsyncExitStack, _AsyncGeneratorContextManager # type: ignore -from contextvars import ContextVar +from contextvars import Context, ContextVar from copy import copy from dataclasses import dataclass from datetime import date, datetime @@ -2010,7 +2010,8 @@ async def _ensure_lifecycle_owner(self) -> None: return self._lifecycle_queue = asyncio.Queue() - self._lifecycle_owner_task = asyncio.create_task( + self._lifecycle_owner_task = Context().run( + asyncio.create_task, self._run_lifecycle_owner(), name=f"mcp-lifecycle:{self.name}", ) diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index fc836e17b3f..3bcf2b218b2 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -10740,6 +10740,29 @@ 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] = [] + + async def record_owner_context(**_: Any) -> None: + observed_call_headers.append(_mcp_call_headers.get(None)) + tool.is_connected = True + + token = _mcp_call_headers.set((tool._header_request_owner, {"Authorization": "principal-a"})) + try: + with patch.object(tool, "_connect_on_owner", new=record_owner_context): + await tool.connect() + + assert observed_call_headers == [None] + finally: + _mcp_call_headers.reset(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() From a75404d088ae5701bc6daf109a307862c79846ff Mon Sep 17 00:00:00 2001 From: Jose Alvarez Date: Fri, 16 Oct 2026 11:31:42 +0200 Subject: [PATCH 2/3] Narrow MCP lifecycle context isolation --- python/packages/core/agent_framework/_mcp.py | 8 +++++--- python/packages/core/tests/core/test_mcp.py | 10 ++++++++-- 2 files changed, 13 insertions(+), 5 deletions(-) diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 511976dca93..151f8284a40 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -12,7 +12,7 @@ from abc import abstractmethod from collections.abc import Callable, Collection, Coroutine, Mapping, Sequence from contextlib import AsyncExitStack, _AsyncGeneratorContextManager # type: ignore -from contextvars import Context, ContextVar +from contextvars import ContextVar, copy_context from copy import copy from dataclasses import dataclass from datetime import date, datetime @@ -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,7 +2010,9 @@ async def _ensure_lifecycle_owner(self) -> None: return self._lifecycle_queue = asyncio.Queue() - self._lifecycle_owner_task = Context().run( + owner_context = copy_context() + owner_context.run(_mcp_call_headers.set, None) + self._lifecycle_owner_task = owner_context.run( asyncio.create_task, self._run_lifecycle_owner(), name=f"mcp-lifecycle:{self.name}", diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index 3bcf2b218b2..d2233286829 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -10746,19 +10746,25 @@ async def test_lifecycle_owner_does_not_inherit_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 - token = _mcp_call_headers.set((tool._header_request_owner, {"Authorization": "principal-a"})) + 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: - _mcp_call_headers.reset(token) + request_context.reset(request_context_token) + _mcp_call_headers.reset(call_headers_token) tool.is_connected = False await tool.close() From 080bc112d4291cb0074d45cdc1665cdf82be69c9 Mon Sep 17 00:00:00 2001 From: Jose Alvarez Date: Fri, 16 Oct 2026 11:34:49 +0200 Subject: [PATCH 3/3] Simplify MCP lifecycle context masking --- python/packages/core/agent_framework/_mcp.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 151f8284a40..21edadb96aa 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -12,7 +12,7 @@ from abc import abstractmethod from collections.abc import Callable, Collection, Coroutine, Mapping, Sequence from contextlib import AsyncExitStack, _AsyncGeneratorContextManager # type: ignore -from contextvars import ContextVar, copy_context +from contextvars import ContextVar from copy import copy from dataclasses import dataclass from datetime import date, datetime @@ -2010,13 +2010,14 @@ async def _ensure_lifecycle_owner(self) -> None: return self._lifecycle_queue = asyncio.Queue() - owner_context = copy_context() - owner_context.run(_mcp_call_headers.set, None) - self._lifecycle_owner_task = owner_context.run( - 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