fix(tools): close cached transports when async worker threads exit

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.
 

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)

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)