From 1b9fdbbdff69e51dbb1e4750f3a01a152f613519 Mon Sep 17 00:00:00 2001 From: garvitkaushik-123 Date: Thu, 20 Aug 2026 18:53:02 +0530 Subject: [PATCH] fix: hold the MCP session open in one background task, don't split enter/exit (issue #60 part 2) OpenDirect21Client.connect()/disconnect() used to split streamablehttp_client's and ClientSession's __aenter__/__aexit__ across two separate method calls. That's fine as long as the connection succeeds -- but the moment the attempt itself failed, __aenter__ never returned, so __aexit__ was never called by anything. The transport was abandoned to Python's async-generator GC finalizer, which runs in whatever task the garbage collector happens to be executing in at finalization time -- not necessarily the task that opened the connection. anyio requires a cancel scope to be entered and exited by the same task, so that mismatch crashed with "Attempted to exit cancel scope in a different task than it was entered in", and the crash propagated all the way up through kickoff_async(), taking down the whole ExecutionActivationFlow. Confirmed the root cause in isolation before touching anything: a bare `streamablehttp_client(url).__aenter__()` / manual `__aexit__()` later against an unreachable host reproduces the crash on its own, with nothing else from this codebase involved. A plain, unbroken nested `async with streamablehttp_client(...) as (r, w, _): async with ClientSession(r, w) as session: ...` against the same unreachable host fails cleanly with an ordinary ExceptionGroup -- no crash. The fix has to make connect()/disconnect() behave like that nested block, not like two independent calls. connect() now starts a background task that owns the MCP session's entire lifetime -- opening AND closing streamablehttp_client and ClientSession within one unbroken `async with`, in that one task, from start to finish. The task blocks on an asyncio.Event in between; disconnect() sets the event and awaits the task to let it unwind naturally, in the same task it opened in. This mirrors the pattern deals_api_mcp_client.py already uses for exactly this reason, cited in its own docstring: "Satisfies anyio's cancel-scope invariant by running the full streamablehttp_client lifecycle inside a single background asyncio Task." Verified: the original ExecutionActivationFlow repro from #60 no longer raises anything -- kickoff_async() now completes and the connection failure is recorded as a plain state warning, exactly the graceful-degrade behavior create_execution_order's own (now reachable) try/except was always meant to provide. Verified both execution_type paths (deal_id and io_order). Verified 5 repeated connect/disconnect cycles against an unreachable host in a row. Found along the way, not fixed here (separate bug, out of scope for this crash): with the crash no longer masking it, the resulting warning reads "'OpenDirect21Client' object has no attribute 'create_execution_order'" -- UnifiedClient.create_execution_order() calls a method OpenDirect21Client never actually defines. Worth its own issue. Tests: new tests/unit/test_opendirect21_client.py -- connection failure degrades cleanly with no raise (the regression), REST fallback via _call_tool when MCP never connected, 5 repeated cycles, disconnect without a prior connect is a no-op, and a mocked successful connection still sets up the session/tools and disconnects both context managers correctly (the happy path the rewrite must not break). Full suite: 1484 passed, 28 skipped (pre-existing, unrelated), no regressions. ruff check / format clean. Closes #60 (part 2, alongside #63 for part 1). --- src/ad_seller/clients/opendirect21_client.py | 94 ++++++++++++--- tests/unit/test_opendirect21_client.py | 118 +++++++++++++++++++ 2 files changed, 194 insertions(+), 18 deletions(-) create mode 100644 tests/unit/test_opendirect21_client.py diff --git a/src/ad_seller/clients/opendirect21_client.py b/src/ad_seller/clients/opendirect21_client.py index 2e20801e..96dd4be5 100644 --- a/src/ad_seller/clients/opendirect21_client.py +++ b/src/ad_seller/clients/opendirect21_client.py @@ -7,6 +7,7 @@ OpenDirect 2.1 specification via MCP (Model Context Protocol). """ +import asyncio from typing import Any, Optional import httpx @@ -40,6 +41,13 @@ def __init__(self, base_url: Optional[str] = None): self._session: Optional[Any] = None self._tools: dict[str, Any] = {} + # The MCP session is held open by a dedicated background task -- + # see connect()/_run_mcp_session() for why. + self._session_task: Optional[asyncio.Task] = None + self._session_ready: Optional[asyncio.Event] = None + self._session_done: Optional[asyncio.Event] = None + self._session_error: Optional[BaseException] = None + async def __aenter__(self) -> "OpenDirect21Client": """Async context manager entry.""" await self.connect() @@ -66,29 +74,79 @@ async def connect(self) -> None: timeout=30.0, ) - # Try MCP connection + # The MCP session runs entirely inside one background task: the + # streamablehttp_client/ClientSession context managers are entered + # AND exited within one unbroken `async with` block, in one task, + # start to finish -- the task just blocks on _session_done.wait() + # in between. This is deliberate, not incidental: splitting + # __aenter__/__aexit__ across separate connect()/disconnect() calls + # (the previous shape of this method) breaks anyio's same-task + # cancel-scope invariant the moment the connection attempt fails, + # and crashes with "Attempted to exit cancel scope in a different + # task than it was entered in" -- reproduced with a bare + # streamablehttp_client() call against an unreachable URL, + # independent of anything else this class does (issue #60 part 2). + # Mirrors the already-proven-safe pattern in deals_api_mcp_client.py. + self._session_ready = asyncio.Event() + self._session_done = asyncio.Event() + self._session_error = None + self._session_task = asyncio.create_task(self._run_mcp_session()) + await self._session_ready.wait() + # self._session_error set means the MCP attempt failed (or was + # cancelled, in which case it was re-raised out of the background + # task already) -- self._session stays None and callers fall back + # to REST via _call_tool(). + + async def _run_mcp_session(self) -> None: + """Own the MCP session's full lifetime in this one task. + + Runs until _session_done is set (normal disconnect()) or the + connection attempt itself fails, in which case _session_error is + recorded and _session_ready is released so connect() doesn't hang. + """ + from mcp import ClientSession + from mcp.client.streamable_http import streamablehttp_client + + assert self._session_ready is not None + assert self._session_done is not None try: - from mcp import ClientSession - from mcp.client.streamable_http import streamablehttp_client - - transport = await streamablehttp_client(self.mcp_url).__aenter__() - read_stream, write_stream, _ = transport - self._session = ClientSession(read_stream, write_stream) - await self._session.__aenter__() - await self._session.initialize() - - # Cache available tools - tools_result = await self._session.list_tools() - self._tools = {tool.name: tool for tool in tools_result.tools} - except Exception: - # MCP not available, fall back to REST + async with streamablehttp_client(self.mcp_url) as (read_stream, write_stream, _): + async with ClientSession(read_stream, write_stream) as session: + await session.initialize() + tools_result = await session.list_tools() + self._tools = {tool.name: tool for tool in tools_result.tools} + self._session = session + self._session_ready.set() + await self._session_done.wait() + except BaseException as exc: + self._session_error = exc + if not self._session_ready.is_set(): + self._session_ready.set() + # A cancellation means something upstream wants this task to + # stop; swallowing it here would hide that from the task's own + # cancellation machinery. + if isinstance(exc, (asyncio.CancelledError, KeyboardInterrupt, SystemExit)): + raise + finally: self._session = None async def disconnect(self) -> None: """Disconnect from the server.""" - if self._session: - await self._session.__aexit__(None, None, None) - self._session = None + if self._session_task is not None: + if self._session_done is not None: + self._session_done.set() + try: + await asyncio.wait_for(self._session_task, timeout=5.0) + except (TimeoutError, asyncio.CancelledError): + self._session_task.cancel() + try: + await self._session_task + except (asyncio.CancelledError, Exception): + pass + except Exception: + pass + self._session_task = None + self._session = None if self._http_client: await self._http_client.aclose() diff --git a/tests/unit/test_opendirect21_client.py b/tests/unit/test_opendirect21_client.py new file mode 100644 index 00000000..2db97bcd --- /dev/null +++ b/tests/unit/test_opendirect21_client.py @@ -0,0 +1,118 @@ +# Author: Green Mountain Systems AI Inc. +# Donated to IAB Tech Lab + +"""Unit tests for OpenDirect21Client's MCP connection lifecycle (issue #60 part 2). + +connect()/disconnect() used to split streamablehttp_client's/ClientSession's +__aenter__ and __aexit__ across separate method calls. That's fine on the +happy path, but the moment the connection attempt itself failed, __aenter__ +never returned so __aexit__ was never called -- the transport was abandoned +to Python's async-generator GC finalizer, which runs in whatever task the +garbage collector happens to be executing in, not the task that opened the +connection. anyio requires a cancel scope to be entered and exited by the +same task, so that mismatch crashed with "Attempted to exit cancel scope in +a different task than it was entered in" -- and that crash propagated all +the way up through kickoff_async(), taking down the whole flow. + +connect() now runs the MCP session's entire lifetime -- open through close +-- inside one background task via a proper nested `async with`, mirroring +the pattern already proven safe in deals_api_mcp_client.py. +""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from ad_seller.clients.opendirect21_client import OpenDirect21Client + + +@pytest.fixture +def settings_stub(): + with patch("ad_seller.clients.opendirect21_client.get_settings") as mock_settings: + mock_settings.return_value = MagicMock( + opendirect_base_url="http://127.0.0.1:1", + opendirect_api_key=None, + opendirect_token=None, + ) + yield mock_settings + + +class TestConnectionFailure: + """The regression: a connection attempt that fails must degrade to REST + fallback cleanly, never crash, and never leak the background task.""" + + async def test_unreachable_server_does_not_raise(self, settings_stub): + client = OpenDirect21Client(base_url="http://127.0.0.1:1") + await client.connect() # must not raise + assert client._session is None + assert client._session_error is not None + await client.disconnect() + + async def test_unreachable_server_records_error_for_rest_fallback(self, settings_stub): + client = OpenDirect21Client(base_url="http://127.0.0.1:1") + await client.connect() + assert client._session_error is not None + assert client._tools == {} + await client.disconnect() + + async def test_call_tool_falls_back_to_rest_when_mcp_unreachable(self, settings_stub): + """_call_tool must route through _rest_call, not blow up, when + the MCP session never established.""" + client = OpenDirect21Client(base_url="http://127.0.0.1:1") + await client.connect() + with patch.object( + client, "_rest_call", new=AsyncMock(return_value=[{"id": "p1"}]) + ) as mock_rest: + result = await client._call_tool("list_products", {}) + mock_rest.assert_awaited_once() + assert result == [{"id": "p1"}] + await client.disconnect() + + async def test_repeated_connect_disconnect_cycles_do_not_crash(self, settings_stub): + """Regression guard: the original bug reproduced reliably on a + single failed connection attempt -- this exercises several in a + row to catch anything that only shows up on reuse.""" + for _ in range(5): + client = OpenDirect21Client(base_url="http://127.0.0.1:1") + await client.connect() + assert client._session is None + await client.disconnect() + + async def test_disconnect_without_connect_is_a_noop(self, settings_stub): + client = OpenDirect21Client(base_url="http://127.0.0.1:1") + await client.disconnect() # must not raise + + +class TestConnectionSuccess: + """The happy path must still work after the rewrite.""" + + async def test_successful_connection_sets_session_and_tools(self, settings_stub): + mock_tool = MagicMock(name="list_products") + mock_tool.name = "list_products" + mock_session = AsyncMock() + mock_session.initialize = AsyncMock() + mock_session.list_tools = AsyncMock(return_value=MagicMock(tools=[mock_tool])) + mock_session.__aenter__ = AsyncMock(return_value=mock_session) + mock_session.__aexit__ = AsyncMock(return_value=None) + + mock_transport_cm = AsyncMock() + mock_transport_cm.__aenter__ = AsyncMock(return_value=(AsyncMock(), AsyncMock(), None)) + mock_transport_cm.__aexit__ = AsyncMock(return_value=None) + + with ( + patch( + "mcp.client.streamable_http.streamablehttp_client", + return_value=mock_transport_cm, + ), + patch("mcp.ClientSession", return_value=mock_session), + ): + client = OpenDirect21Client(base_url="http://127.0.0.1:1") + await client.connect() + + assert client._session is mock_session + assert client._session_error is None + assert "list_products" in client._tools + + await client.disconnect() + mock_session.__aexit__.assert_awaited() + mock_transport_cm.__aexit__.assert_awaited()