diff --git a/plugins/web/parallel/provider.py b/plugins/web/parallel/provider.py index 028f5df3fc37b..597da8208fe3c 100644 --- a/plugins/web/parallel/provider.py +++ b/plugins/web/parallel/provider.py @@ -3,8 +3,8 @@ Subclasses :class:`agent.web_search_provider.WebSearchProvider`. Uses two distinct Parallel SDK clients: -- ``Parallel`` (sync) — for :meth:`search` -- ``AsyncParallel`` (async) — for :meth:`extract` +- ``Parallel`` (sync, cached) — for :meth:`search` +- ``AsyncParallel`` (async, request-scoped) — for :meth:`extract` This is the first plugin to exercise the **async-extract** code path in the ABC: :meth:`extract` is declared ``async def``, and the dispatcher @@ -36,11 +36,11 @@ logger = logging.getLogger(__name__) -# Module-level note: the canonical cache slots ``_parallel_client`` and -# ``_async_parallel_client`` live on :mod:`tools.web_tools` so tests that do -# ``tools.web_tools._parallel_client = None`` between cases see fresh state. -# The plugin reads/writes through that public module (see -# :func:`_get_sync_client` / :func:`_get_async_client`). +# Module-level note: the canonical sync cache slot ``_parallel_client`` lives +# on :mod:`tools.web_tools` so tests that reset it between cases see fresh +# state. Async clients are deliberately request-scoped: httpx transports are +# bound to the event loop that first uses them and cannot be shared safely by +# the per-thread loops used for concurrent tool execution. def _ensure_parallel_sdk_installed() -> None: @@ -91,16 +91,13 @@ def _get_sync_client() -> Any: def _get_async_client() -> Any: - """Lazy-load + cache the async Parallel client. + """Create an async Parallel client owned by the current extraction. - Cache lives on :mod:`tools.web_tools` (as ``_async_parallel_client``). + The caller must close the client on the same event loop that uses it. + Caching this process-wide lets concurrent tool workers race to publish + loop-affine clients; losing clients can then be finalized from + prompt_toolkit's loop after their worker loops have died. """ - import tools.web_tools as _wt - - cached = getattr(_wt, "_async_parallel_client", None) - if cached is not None: - return cached - from agent.web_search_provider import get_provider_env api_key = get_provider_env("PARALLEL_API_KEY") @@ -113,21 +110,17 @@ def _get_async_client() -> Any: _ensure_parallel_sdk_installed() from parallel import AsyncParallel # noqa: WPS433 — deliberately lazy - client = AsyncParallel(api_key=api_key) - _wt._async_parallel_client = client - return client + return AsyncParallel(api_key=api_key) def _reset_clients_for_tests() -> None: - """Drop both cached clients so tests can re-instantiate cleanly. + """Drop the cached sync client so tests can re-instantiate cleanly. - Clears the canonical slots on :mod:`tools.web_tools` (where - :func:`_get_sync_client` / :func:`_get_async_client` read/write them). + The async client is request-scoped and therefore has no cache to reset. """ import tools.web_tools as _wt _wt._parallel_client = None - _wt._async_parallel_client = None # Backward-compatible aliases for the names that lived in tools.web_tools @@ -234,10 +227,41 @@ async def extract( ] logger.info("Parallel extract: %d URL(s)", len(urls)) - response = await _get_async_client().beta.extract( - urls=urls, - full_content=True, - ) + client = _get_async_client() + try: + response = await client.beta.extract( + urls=urls, + full_content=True, + ) + finally: + # AsyncParallel owns an httpx connection pool whose transports + # are tied to this running loop. Drain it here, before the + # concurrent-tool worker and its loop go out of scope. + # + # Never let a teardown failure destroy an otherwise successful + # extraction: ``close()`` funnels into ``httpx.aclose()`` -> + # ``transport.aclose()``, which can raise (a mid-shutdown TLS + # error, or the very RuntimeError this cleanup exists to + # prevent). Without this guard that exception propagates out + # of ``finally``, past the response handling below, and the + # outer ``except Exception`` rewrites every URL into an error + # result even though the content was already fetched. + # + # Masked from the caller, but NOT from operators: a failed + # close can leave partial resources behind, and an "Event loop + # is closed" here would mean this ownership fix regressed, so + # it must stay visible. ``except Exception`` deliberately lets + # ``CancelledError`` (BaseException since 3.8) propagate so + # cancellation semantics are preserved. + try: + await client.close() + except Exception as close_exc: # noqa: BLE001 — cleanup is best-effort + logger.warning( + "Parallel async client close failed; preserving " + "extraction result: %s: %s", + type(close_exc).__name__, + close_exc, + ) results: List[Dict[str, Any]] = [] for result in response.results or []: diff --git a/tests/tools/test_parallel_async_client_lifecycle.py b/tests/tools/test_parallel_async_client_lifecycle.py new file mode 100644 index 0000000000000..ce1cd6df74e5f --- /dev/null +++ b/tests/tools/test_parallel_async_client_lifecycle.py @@ -0,0 +1,189 @@ +"""Regression coverage for Parallel's loop-affine async client lifecycle.""" + +import asyncio +import logging +import sys +import threading +import types +from types import SimpleNamespace + +from tools.daemon_pool import DaemonThreadPoolExecutor + + +def test_concurrent_extract_closes_each_client_on_its_owner_loop(monkeypatch): + """Concurrent workers must never publish an async client process-wide. + + A constructor barrier forces the old cache implementation's three workers + to observe an empty slot before any client can be published. Two clients + then lose the publication race and become cyclic garbage after their + worker threads exit, reproducing the lifecycle that surfaced as + ``RuntimeError('Event loop is closed')`` in prompt_toolkit. + """ + from model_tools import _run_async + from plugins.web.parallel import provider as parallel_provider + import tools.web_tools as web_tools + + constructor_barrier = threading.Barrier(3, timeout=15) + instances = [] + + class FakeBeta: + def __init__(self, client): + self.client = client + + async def extract(self, *, urls, full_content): + assert full_content is True + self.client.use_loop_id = id(asyncio.get_running_loop()) + return SimpleNamespace( + results=[ + SimpleNamespace( + url=urls[0], + title="Example", + full_content="content", + excerpts=[], + ) + ], + errors=[], + ) + + class FakeAsyncParallel: + def __init__(self, *, api_key): + assert api_key == "parallel-test-key" + self.beta = FakeBeta(self) + self.closed = False + self.use_loop_id = None + self.close_loop_id = None + instances.append(self) + constructor_barrier.wait() + + async def close(self): + self.close_loop_id = id(asyncio.get_running_loop()) + self.closed = True + + fake_parallel = types.ModuleType("parallel") + fake_parallel.AsyncParallel = FakeAsyncParallel + monkeypatch.setitem(sys.modules, "parallel", fake_parallel) + monkeypatch.setenv("PARALLEL_API_KEY", "parallel-test-key") + monkeypatch.setattr( + parallel_provider, + "_ensure_parallel_sdk_installed", + lambda: None, + ) + # Preserve compatibility with callers/tests that still create this legacy + # attribute dynamically. The provider must neither read nor populate it. + monkeypatch.setattr(web_tools, "_async_parallel_client", None, raising=False) + + provider = parallel_provider.ParallelWebSearchProvider() + + def extract(index): + return _run_async(provider.extract([f"https://example.test/{index}"])) + + with DaemonThreadPoolExecutor(max_workers=3) as executor: + results = list(executor.map(extract, range(3))) + + assert len(instances) == 3 + assert all(client.closed for client in instances) + assert all(client.close_loop_id == client.use_loop_id for client in instances) + assert web_tools._async_parallel_client is None + assert [result[0]["content"] for result in results] == ["content"] * 3 + + +def test_close_failure_does_not_discard_a_successful_extraction(monkeypatch, caplog): + """A teardown error must not rewrite fetched content into an error result. + + ``AsyncParallel.close()`` funnels into ``httpx.aclose()`` -> + ``transport.aclose()``, which can raise while the pool is being drained. + That happens *after* the content is already in hand, so the extraction + must still succeed; the cleanup failure is masked from the caller but + logged at warning level so a regressed ownership fix stays visible. + """ + from plugins.web.parallel import provider as parallel_provider + + class FakeBeta: + def __init__(self, client): + self.client = client + + async def extract(self, *, urls, full_content): + assert full_content is True + return SimpleNamespace( + results=[ + SimpleNamespace( + url=urls[0], + title="Example", + full_content="content", + excerpts=[], + ) + ], + errors=[], + ) + + class ExplodingOnCloseParallel: + def __init__(self, *, api_key): + self.beta = FakeBeta(self) + self.close_calls = 0 + instances.append(self) + + async def close(self): + self.close_calls += 1 + raise RuntimeError("Event loop is closed") + + instances = [] + fake_parallel = types.ModuleType("parallel") + fake_parallel.AsyncParallel = ExplodingOnCloseParallel + monkeypatch.setitem(sys.modules, "parallel", fake_parallel) + monkeypatch.setenv("PARALLEL_API_KEY", "parallel-test-key") + monkeypatch.setattr( + parallel_provider, + "_ensure_parallel_sdk_installed", + lambda: None, + ) + + provider = parallel_provider.ParallelWebSearchProvider() + with caplog.at_level(logging.WARNING, logger=parallel_provider.__name__): + results = asyncio.run(provider.extract(["https://example.test/1"])) + + assert results[0]["content"] == "content" + assert "error" not in results[0] + # Pin this test's own trigger: deleting the close entirely must not pass. + assert [client.close_calls for client in instances] == [1] + close_warnings = [ + record + for record in caplog.records + if record.levelno == logging.WARNING and "close failed" in record.message + ] + assert len(close_warnings) == 1 + + +def test_extraction_failure_outranks_a_close_failure(monkeypatch): + """When both the request and the cleanup fail, report the request error. + + The primary exception must win: a secondary cleanup failure must not + overwrite the reason the extraction actually failed. + """ + from plugins.web.parallel import provider as parallel_provider + + class FakeBeta: + async def extract(self, *, urls, full_content): + raise RuntimeError("extract failed") + + class ExplodingBothWaysParallel: + def __init__(self, *, api_key): + self.beta = FakeBeta() + + async def close(self): + raise RuntimeError("close failed") + + fake_parallel = types.ModuleType("parallel") + fake_parallel.AsyncParallel = ExplodingBothWaysParallel + monkeypatch.setitem(sys.modules, "parallel", fake_parallel) + monkeypatch.setenv("PARALLEL_API_KEY", "parallel-test-key") + monkeypatch.setattr( + parallel_provider, + "_ensure_parallel_sdk_installed", + lambda: None, + ) + + provider = parallel_provider.ParallelWebSearchProvider() + results = asyncio.run(provider.extract(["https://example.test/1"])) + + assert "extract failed" in results[0]["error"] + assert "close failed" not in results[0]["error"] diff --git a/tools/web_tools.py b/tools/web_tools.py index e8c62142af8b9..f412cb157fe00 100644 --- a/tools/web_tools.py +++ b/tools/web_tools.py @@ -65,22 +65,22 @@ _normalize_tavily_search_results, _tavily_request, ) -# Parallel + Exa clients re-exported for backward-compat with existing +# Parallel + Exa client helpers re-exported for backward-compat with existing # unit tests (tests/tools/test_web_tools_config.py imports _get_parallel_client -# / _get_async_parallel_client / _get_exa_client directly). +# / _get_async_parallel_client / _get_exa_client directly). The async Parallel +# helper is a factory; its caller owns and closes each returned client. from plugins.web.parallel.provider import ( # noqa: F401 — backward-compat names _get_async_parallel_client, _get_parallel_client, ) from plugins.web.exa.provider import _get_exa_client # noqa: F401 -# Module-level cache slots for the per-vendor clients. The plugins read/write -# these via tools.web_tools so unit tests that reset +# Module-level cache slots for reusable synchronous vendor clients. The +# plugins read/write these via tools.web_tools so unit tests that reset # ``tools.web_tools.__client = None`` between cases keep working. _firecrawl_client: Optional[Any] = None _firecrawl_client_config: Optional[Any] = None _parallel_client: Optional[Any] = None -_async_parallel_client: Optional[Any] = None _exa_client: Optional[Any] = None from tools.debug_helpers import DebugSession