diff --git a/agent/auxiliary_client.py b/agent/auxiliary_client.py index 4844545125079..8de08691f68bc 100644 --- a/agent/auxiliary_client.py +++ b/agent/auxiliary_client.py @@ -24,6 +24,7 @@ from typing import Any, Callable, Dict, List, NamedTuple, Optional, Tuple, TYPE_CHECKING, Union from urllib.parse import urlparse, parse_qs, urlunparse +from agent.auxiliary_client_lifecycle import _close_cached_client from agent.error_classifier import ( _BILLING_PATTERNS, _OVERLOADED_PATTERNS, @@ -5900,73 +5901,6 @@ def neuter_async_httpx_del() -> None: pass # Graceful degradation if the SDK changes its internals -def _force_close_async_httpx(client: Any) -> None: - """Mark the httpx AsyncClient inside an AsyncOpenAI client as closed so ``__del__`` won't - schedule ``aclose()`` on a dead loop. Skips the full async close — the OS drops connections.""" - with contextlib.suppress(Exception): - from httpx._client import ClientState - inner = getattr(client, "_client", None) - if inner is not None and not getattr(inner, "is_closed", True): - inner._state = ClientState.CLOSED - - -def _schedule_async_close(close_result: Any, client: Any) -> None: - """Finish an async close without leaking an unawaited coroutine.""" - async def _await_close() -> None: - try: - await close_result - except Exception: - pass - finally: - _force_close_async_httpx(client) - runner = _await_close() - try: - import asyncio as _aio - try: - loop = _aio.get_running_loop() - except RuntimeError: - _aio.run(runner) - else: - task = loop.create_task(runner) - - def _consume(completed_task) -> None: - with contextlib.suppress(BaseException): - completed_task.exception() - task.add_done_callback(_consume) - runner = None - except Exception: - if runner is not None: - with contextlib.suppress(Exception): - runner.close() - _force_close_async_httpx(client) - - -def _close_cached_client(client: Any, *, close_async: bool = False) -> None: - """Close one cached client, awaiting async transports only when safe.""" - if client is None: - return - close_fn = getattr(client, "close", None) - if not callable(close_fn): - _force_close_async_httpx(client) - return - try: - close_result = close_fn() - except Exception: - _force_close_async_httpx(client) - return - if inspect.isawaitable(close_result): - if close_async: - _schedule_async_close(close_result, client) - else: - # Never await a client owned by another live loop; close the coroutine (no - # unawaited warning) and neuter the transport. - with contextlib.suppress(Exception): - close_result.close() - _force_close_async_httpx(client) - return - _force_close_async_httpx(client) - - def shutdown_cached_clients() -> None: """Close all cached clients; call at CLI shutdown *before* the loop closes. diff --git a/agent/auxiliary_client_lifecycle.py b/agent/auxiliary_client_lifecycle.py new file mode 100644 index 0000000000000..4db6685422f5f --- /dev/null +++ b/agent/auxiliary_client_lifecycle.py @@ -0,0 +1,103 @@ +"""Teardown of cached auxiliary-client transports.""" + +import contextlib +import inspect +from typing import Any, Optional, TYPE_CHECKING + +if TYPE_CHECKING: + from asyncio import Task + + +def _force_close_async_httpx(client: Any) -> None: + """Mark the httpx AsyncClient inside an AsyncOpenAI client as closed so ``__del__`` won't + schedule ``aclose()`` on a dead loop. Skips the full async close — the OS drops connections.""" + with contextlib.suppress(Exception): + from httpx._client import ClientState + + inner = getattr(client, "_client", None) + if inner is not None and not getattr(inner, "is_closed", True): + inner._state = ClientState.CLOSED + + +def _schedule_async_close(close_result: Any, client: Any) -> Optional["Task[None]"]: + """Finish an async close without leaking an unawaited coroutine.""" + + async def _await_close() -> None: + try: + await close_result + except Exception: + pass + finally: + _force_close_async_httpx(client) + + runner = _await_close() + try: + import asyncio as _aio + + try: + loop = _aio.get_running_loop() + except RuntimeError: + _aio.run(runner) + else: + task = loop.create_task(runner) + + def _consume(completed_task) -> None: + with contextlib.suppress(BaseException): + completed_task.exception() + + task.add_done_callback(_consume) + runner = None + return task + except Exception: + if runner is not None: + with contextlib.suppress(Exception): + runner.close() + _force_close_async_httpx(client) + + +def _close_cached_client( + client: Any, *, close_async: bool = False +) -> Optional["Task[None]"]: + """Close one cached client, awaiting async transports only when safe.""" + if client is None: + return + close_fn = getattr(client, "close", None) + if not callable(close_fn): + _force_close_async_httpx(client) + return + try: + close_result = close_fn() + except Exception: + _force_close_async_httpx(client) + return + if inspect.isawaitable(close_result): + if close_async: + return _schedule_async_close(close_result, client) + else: + # Never await a client owned by another live loop; close the coroutine (no + # unawaited warning) and neuter the transport. + with contextlib.suppress(Exception): + close_result.close() + _force_close_async_httpx(client) + return + _force_close_async_httpx(client) + + +async def shutdown_cached_clients_for_current_loop() -> None: + """Close cached transports owned by the current event loop.""" + import asyncio + from agent import auxiliary_client + + loop = asyncio.get_running_loop() + with auxiliary_client._client_cache_lock: + keys = [ + key + for key, entry in auxiliary_client._client_cache.items() + if entry[2] is loop + ] + clients = [auxiliary_client._client_cache.pop(key)[0] for key in keys] + + tasks = [_close_cached_client(client, close_async=True) for client in clients] + pending = [task for task in tasks if task is not None] + if pending: + await asyncio.gather(*pending) diff --git a/model_tools.py b/model_tools.py index a4ef18942889d..125073686984d 100644 --- a/model_tools.py +++ b/model_tools.py @@ -16,6 +16,8 @@ import logging import threading import time +import sys +import weakref from typing import Dict, Any, List, Optional, Tuple from tools.registry import CHECK_FN_CACHE_BYPASS, check_fn_cache_scope, discover_builtin_tools, registry, tool_error @@ -80,14 +82,46 @@ def _get_tool_loop(): return _tool_loop -def _get_worker_loop(): +async def _finish_worker_loop() -> None: + current = asyncio.current_task() + pending = [task for task in asyncio.all_tasks() if task is not current] + for task in pending: + task.cancel() + if pending: + await asyncio.gather(*pending, return_exceptions=True) + + await asyncio.get_running_loop().shutdown_asyncgens() + if "agent.auxiliary_client" in sys.modules: + from agent.auxiliary_client_lifecycle import shutdown_cached_clients_for_current_loop + + await shutdown_cached_clients_for_current_loop() + + +def _close_worker_loop(runner: asyncio.Runner) -> None: + if runner.get_loop().is_closed(): + return + + try: + runner.run(_finish_worker_loop()) + finally: + runner.close() + + +class _WorkerLoopOwner: + def __init__(self) -> None: + self._runner = asyncio.Runner() + self.loop = self._runner.get_loop() + # A bound owner method would keep this owner alive after thread exit. + self._finalizer = weakref.finalize(self, _close_worker_loop, self._runner) + + +def _get_worker_loop() -> asyncio.AbstractEventLoop: """Persistent event loop for the current worker thread (thread-local).""" - loop = getattr(_worker_thread_local, 'loop', None) - if loop is None or loop.is_closed(): - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) - _worker_thread_local.loop = loop - return loop + owner = getattr(_worker_thread_local, "owner", None) + if owner is None or owner.loop.is_closed(): + owner = _WorkerLoopOwner() + _worker_thread_local.owner = owner + return owner.loop def _run_async(coro): @@ -101,30 +135,24 @@ def _run_async(coro): # reference to, so on timeout we can cancel the task inside it # (ThreadPoolExecutor.cancel() is a no-op on a running worker). import concurrent.futures + worker_loop: Optional[asyncio.AbstractEventLoop] = None loop_ready = threading.Event() def _run_in_worker(): nonlocal worker_loop - worker_loop = asyncio.new_event_loop() + runner = asyncio.Runner() + worker_loop = runner.get_loop() loop_ready.set() try: - asyncio.set_event_loop(worker_loop) return worker_loop.run_until_complete(coro) finally: - try: # drain tasks still pending after an external cancel - pending = asyncio.all_tasks(worker_loop) - for t in pending: - t.cancel() - if pending: - worker_loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True)) - except Exception: - pass - worker_loop.close() + _close_worker_loop(runner) pool = concurrent.futures.ThreadPoolExecutor(max_workers=1) # Carry profile + approval/sudo context so get_hermes_home() resolves correctly. from tools.thread_context import propagate_context_to_thread + future = pool.submit(propagate_context_to_thread(_run_in_worker)) try: return future.result(timeout=300) diff --git a/tests/test_model_tools_async_bridge.py b/tests/test_model_tools_async_bridge.py index e1d22268282c6..6cc614953cdc8 100644 --- a/tests/test_model_tools_async_bridge.py +++ b/tests/test_model_tools_async_bridge.py @@ -78,6 +78,123 @@ def worker(): results = [future.result(timeout=15) for future in futures] assert len({thread for _, thread in results}) == 3 assert len({main, *(loop for loop, _ in results)}) == 4 + try: + assert [loop.is_closed() for loop, _ in results] == [True, True, True] + assert not main.is_closed() + finally: + for loop, _ in results: + if not loop.is_closed(): + loop.close() + + +@pytest.mark.parametrize("worker_mode", ["persistent", "running-loop"]) +def test_worker_exit_closes_its_cached_client_before_the_loop(monkeypatch, worker_mode): + from concurrent.futures import ThreadPoolExecutor + from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + + from openai import AsyncOpenAI + + from agent import auxiliary_client as aux + from model_tools import _run_async + + class Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def do_GET(self): + self.send_response(200) + self.send_header("Content-Length", "2") + self.end_headers() + self.wfile.write(b"ok") + + def log_message(self, format: str, *args: object) -> None: + pass + + monkeypatch.setattr(aux, "_client_cache", {}) + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + serving = threading.Thread(target=server.serve_forever, daemon=True) + serving.start() + url = f"http://127.0.0.1:{server.server_port}" + + async def connect(key): + loop = asyncio.get_running_loop() + client = AsyncOpenAI(api_key="test", base_url=url, max_retries=0) + aux._store_cached_client(key, client, None, bound_loop=loop) + assert (await client._client.get(url)).text == "ok" + return loop, client + + main_key = ("main-lifetime-control",) + worker_key = ("worker-lifetime",) + stream_closed = threading.Event() + + async def open_stream(client): + async def stream(): + try: + yield client + finally: + assert not client.is_closed() + stream_closed.set() + + generator = stream() + await anext(generator) + return generator + + async def read(client): + return (await client._client.get(url)).text + + main, main_client = _run_async(connect(main_key)) + try: + + def worker(): + loop, client = _run_async(connect(worker_key)) + assert _run_async(_get_current_loop()) is loop + assert not client.is_closed() + assert _run_async(read(client)) == "ok" + return loop, client, _run_async(open_stream(client)) + + if worker_mode == "persistent": + with ThreadPoolExecutor(max_workers=1) as pool: + loop, client, generator = pool.submit(worker).result(timeout=15) + else: + + async def operation(): + loop, client = await connect(worker_key) + assert await _get_current_loop() is loop + assert not client.is_closed() + assert await read(client) == "ok" + return loop, client, await open_stream(client) + + async def running_caller(): + return _run_async(operation()) + + loop, client, generator = asyncio.run(running_caller()) + + assert ( + loop.is_closed(), + client.is_closed(), + main.is_closed(), + main_client.is_closed(), + set(aux._client_cache), + stream_closed.is_set(), + ) == (True, True, False, False, {main_key}, True) + finally: + if not stream_closed.is_set(): + if loop.is_closed(): + asyncio.run(generator.aclose()) + else: + loop.run_until_complete(generator.aclose()) + if not client.is_closed(): + if loop.is_closed(): + aux._close_cached_client(client, close_async=True) + else: + loop.run_until_complete(client.close()) + _run_async(main_client.close()) + aux._client_cache.clear() + if not loop.is_closed(): + loop.close() + server.shutdown() + server.server_close() + serving.join(timeout=10) + class TestRunAsyncWithRunningLoop: @@ -154,8 +271,12 @@ async def _never_finishes(): FakeExecutor, ) - with pytest.raises(concurrent.futures.TimeoutError): - _run_async(_never_finishes()) + coroutine = _never_finishes() + try: + with pytest.raises(concurrent.futures.TimeoutError): + _run_async(coroutine) + finally: + coroutine.close() assert events["result_timeout"] == 300 # The worker wrapper creates its own event loop so _run_async can diff --git a/tests/test_model_tools_worker_shutdown.py b/tests/test_model_tools_worker_shutdown.py new file mode 100644 index 0000000000000..95d09b8755c67 --- /dev/null +++ b/tests/test_model_tools_worker_shutdown.py @@ -0,0 +1,76 @@ +"""Worker teardown must cancel background tasks and join their executor threads.""" + +import asyncio +from concurrent.futures import ThreadPoolExecutor +import threading + +import pytest + + +@pytest.mark.parametrize("mode", ["persistent", "running-loop"]) +def test_worker_exit_drains_background_task_and_default_executor(mode): + from model_tools import _run_async + + started = threading.Event() + release = threading.Event() + finished = threading.Event() + cleaned = threading.Event() + threads = [] + + def blocking_work(): + threads.append(threading.current_thread()) + started.set() + try: + assert release.wait(timeout=15) + finally: + finished.set() + + async def background(): + try: + await asyncio.to_thread(blocking_work) + finally: + release.set() + cleaned.set() + + async def operation(): + task = asyncio.create_task(background()) + await asyncio.sleep(0) + assert started.wait(timeout=5) + return asyncio.get_running_loop(), task + + async def running_caller(): + return _run_async(operation()) + + loop = None + try: + if mode == "persistent": + with ThreadPoolExecutor(max_workers=1) as pool: + loop, task = pool.submit(_run_async, operation()).result(timeout=10) + else: + loop, task = asyncio.run(running_caller()) + + assert ( + loop.is_closed(), + task.done(), + task.cancelled(), + cleaned.is_set(), + finished.is_set(), + any(thread.is_alive() for thread in threads), + ) == (True, True, True, True, True, False) + finally: + release.set() + if loop is not None and not loop.is_closed(): + + async def cleanup(): + pending = asyncio.all_tasks() + pending.discard(asyncio.current_task()) + for pending_task in pending: + pending_task.cancel() + await asyncio.gather(*pending, return_exceptions=True) + await asyncio.get_running_loop().shutdown_default_executor() + + loop.run_until_complete(cleanup()) + loop.close() + for thread in threads: + thread.join(timeout=10) + assert not thread.is_alive()