From 82e2c9ce40a7f0d96cd55073da09ddef9abaeb1b Mon Sep 17 00:00:00 2001 From: Soju06 Date: Tue, 14 Jul 2026 03:57:59 +0000 Subject: [PATCH] perf(agent): reuse the per-request OpenAI wire client across sequential LLM calls MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every LLM call built a fresh openai.OpenAI wire client (new httpx pool, TCP+TLS handshake, measured 19.2ms p50 / 35.5ms p95 per call at ~5 calls per tool-loop turn) and closed it when the request finished. Cache ONE reusable wire client on the agent, keyed by the effective client kwargs: - _create_request_openai_client hands back the cached client when the effective kwargs are identical; any change (credential rotation, provider failover, vision default_headers) evicts and rebuilds. - Only a request that produced a response reports a reuse close reason (request_complete / stream_request_complete); error unwinds report *_error_cleanup and really close, so a retry after a request error always builds a fresh pool. - Cross-thread aborts poison the slot: a pool whose sockets were shutdown(SHUT_RDWR) from a stranger thread is never reused (#29507) — the owner-thread close discards it and the next create rebuilds. The holder read and the abort are atomic (under the holder lock) at all three abort sites, so a late abort can never poison the NEXT request's checked-out client. - Worker-side interrupt breaks close the half-read SSE stream on the owning thread before building the partial response; a failed close poisons the slot (otherwise each interrupt leaked one checked-out connection until the pool hit PoolTimeout). run_codex_stream gets the same poison-on-close-failure handling. - Single checked-out slot (in_use): a concurrent call gets an untracked client with the old per-request lifecycle. - release_clients() / close() really close the cached client when idle; if a worker has it checked out they abort the sockets and detach the slot, deferring the FD release to the worker's own close. - MoA facade and Mock passthroughs never enter the cache; max_retries=0 is preserved on all request clients. --- agent/chat_completion_helpers.py | 113 ++++-- agent/codex_runtime.py | 15 +- run_agent.py | 131 ++++++- tests/agent/test_request_client_reuse.py | 289 ++++++++++++++ .../run_agent/test_openai_client_lifecycle.py | 12 +- .../test_request_client_reuse_abort_races.py | 369 ++++++++++++++++++ 6 files changed, 899 insertions(+), 30 deletions(-) create mode 100644 tests/agent/test_request_client_reuse.py create mode 100644 tests/run_agent/test_request_client_reuse_abort_races.py diff --git a/agent/chat_completion_helpers.py b/agent/chat_completion_helpers.py index 5fecd0e26cfe..f4bca4fa1755 100644 --- a/agent/chat_completion_helpers.py +++ b/agent/chat_completion_helpers.py @@ -519,10 +519,16 @@ def direct_api_call(agent, api_kwargs: dict): def _abort_active_request(reason: str) -> None: """Abort the inline request from a watchdog/interrupt thread.""" + # Abort while still holding the holder lock: the instant it is + # released, the inline finally may pop + cache the client for reuse + # and the NEXT call check it out — a late abort would then poison + # the slot and shut down an innocent in-flight request's sockets + # (same atomicity contract as _close_request_client_once in the + # interruptible variants; the abort itself never blocks). with request_client_lock: request_client = request_client_holder["client"] - if request_client is not None: - agent._abort_request_openai_client(request_client, reason=reason) + if request_client is not None: + agent._abort_request_openai_client(request_client, reason=reason) def _make_client(reason: str, kind: str = "openai"): # direct_api_call only runs for OpenAI-wire chat_completions cron @@ -535,6 +541,10 @@ def direct_api_call(agent, api_kwargs: dict): agent._active_request_abort = _abort_active_request return client + # Only a clean return may report the reuse reason (request_complete): + # after an error or interrupt the wire client is really closed so the + # retry builds a fresh pool (see _REQUEST_CLIENT_REUSE_REASONS). + succeeded = False try: response = _dispatch_nonstreaming_api_request( agent, api_kwargs, make_client=_make_client @@ -547,6 +557,7 @@ def direct_api_call(agent, api_kwargs: dict): if getattr(agent, "_interrupt_requested", False): raise InterruptedError("Agent interrupted during API call") _reset_stale_streak(agent) + succeeded = True return response finally: if getattr(agent, "_active_request_abort", None) is _abort_active_request: @@ -555,7 +566,10 @@ def direct_api_call(agent, api_kwargs: dict): request_client = request_client_holder["client"] request_client_holder["client"] = None if request_client is not None: - agent._close_request_openai_client(request_client, reason="request_complete") + agent._close_request_openai_client( + request_client, + reason="request_complete" if succeeded else "request_error_cleanup", + ) def interruptible_api_call(agent, api_kwargs: dict): @@ -633,20 +647,28 @@ def interruptible_api_call(agent, api_kwargs: dict): and owner_tid is not None and owner_tid != threading.get_ident() ) - if not stranger_thread: - # Owning thread (or no recorded owner) → pop and fully close. - request_client_holder["client"] = None - request_client_holder["owner_tid"] = None + if stranger_thread: + # Abort while still holding the holder lock: the instant it + # is released, the worker's finally may pop + cache the client + # for reuse and the NEXT call check it out — an abort landing + # after that would poison the slot and shut down an innocent + # in-flight request's sockets. The abort itself never blocks + # (socket shutdown + slot poison), so holding the lock across + # it only delays the racing pop, never the data path. + if request_client_kind.get("value", "openai") == "anthropic_messages": + agent._abort_request_anthropic_client( + request_client, reason=reason + ) + else: + agent._abort_request_openai_client(request_client, reason=reason) + return + # Owning thread (or no recorded owner) → pop and fully close. + request_client_holder["client"] = None + request_client_holder["owner_tid"] = None if request_client is None: return - kind = request_client_kind.get("value", "openai") - if kind == "anthropic_messages": - if stranger_thread: - agent._abort_request_anthropic_client(request_client, reason=reason) - else: - agent._close_request_anthropic_client(request_client, reason=reason) - elif stranger_thread: - agent._abort_request_openai_client(request_client, reason=reason) + if request_client_kind.get("value", "openai") == "anthropic_messages": + agent._close_request_anthropic_client(request_client, reason=reason) else: agent._close_request_openai_client(request_client, reason=reason) @@ -683,7 +705,15 @@ def interruptible_api_call(agent, api_kwargs: dict): return result["error"] = e finally: - _close_request_client_once("request_complete") + # Reuse reason only on a clean response; any other outcome — + # error, or the cancel-swallow return above (which leaves both + # result slots None) — really closes so the next attempt builds + # a fresh pool (see _REQUEST_CLIENT_REUSE_REASONS). + _close_request_client_once( + "request_complete" + if result["response"] is not None + else "request_error_cleanup" + ) # ── Stale-call timeout (mirrors streaming stale detector) ──────── # Non-streaming calls return nothing until the full response is @@ -2642,20 +2672,27 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= and owner_tid is not None and owner_tid != threading.get_ident() ) - if not stranger_thread: - request_client_holder["client"] = None - request_client_holder["owner_tid"] = None + if stranger_thread: + # Abort under the holder lock — see the non-streaming variant + # for why the holder read and the abort must be atomic (a late + # abort would otherwise hit the NEXT request's checkout). + if request_client_kind.get("value", "openai") == "anthropic_messages": + agent._abort_request_anthropic_client( + request_client, reason=reason + ) + else: + agent._abort_request_openai_client(request_client, reason=reason) + return + request_client_holder["client"] = None + request_client_holder["owner_tid"] = None if request_client is None: return + # Stranger threads returned under the lock above, so only the owner + # (or an any-thread-safe stream handle) reaches the close dispatch. if request_kind == "stream": _close_request_stream_handle(request_client, reason) elif request_kind == "anthropic_messages": - if stranger_thread: - agent._abort_request_anthropic_client(request_client, reason=reason) - else: - agent._close_request_anthropic_client(request_client, reason=reason) - elif stranger_thread: - agent._abort_request_openai_client(request_client, reason=reason) + agent._close_request_anthropic_client(request_client, reason=reason) else: agent._close_request_openai_client(request_client, reason=reason) @@ -2944,6 +2981,23 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= pass if agent._interrupt_requested: + # Abandoning a half-read SSE response leaves its connection + # permanently checked out of the httpx pool — and the partial + # response built below makes the worker's finally report a + # reuse-reason close, which would cache the client together + # with the leaked connection (each interrupt leaking one more + # until the pool exhausts). Close the stream here, on the + # owning thread, so the connection is released first. + try: + stream.close() + except Exception: + # Connection may still be checked out — poison the slot so + # the finally's close really closes the pool instead of + # caching it (owner-thread abort: shutdown is safe, and the + # FD release still happens in the finally below). + agent._abort_request_openai_client( + request_client, reason="interrupt_stream_close_failed" + ) break if not _stream_attempt_is_active(stream_attempt_id): @@ -3695,7 +3749,14 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= result["error"] = e return finally: - _close_request_client_once("stream_request_complete") + # Reuse reason only on a clean stream; any other outcome (error, + # cancel-swallow) really closes so the next attempt builds a + # fresh pool (see _REQUEST_CLIENT_REUSE_REASONS). + _close_request_client_once( + "stream_request_complete" + if result["response"] is not None + else "stream_error_cleanup" + ) # Provider-configured stale timeout takes priority over env default. _cfg_stale = get_provider_stale_timeout(agent.provider, agent.model) diff --git a/agent/codex_runtime.py b/agent/codex_runtime.py index ed4dae1b57d7..f7d3d3c9f181 100644 --- a/agent/codex_runtime.py +++ b/agent/codex_runtime.py @@ -1351,7 +1351,20 @@ def run_codex_stream(agent, api_kwargs: dict, client: Any = None, on_first_delta try: close_fn() except Exception: - pass + # A failed close can leave this response's connection + # checked out of the httpx pool while the caller's finally + # reports a reuse-reason close (e.g. interrupt_check broke + # the event loop with collected output) — caching the + # client with the leaked connection. Poison the slot so + # that close really closes the pool (owner-thread abort; + # mirrors the chat-streaming interrupt-break handling). + # ``client is None`` means the shared primary client, + # which is never reuse-cached and must not have its + # sockets force-shut here. + if client is not None: + agent._abort_request_openai_client( + active_client, reason="codex_stream_close_failed" + ) def run_codex_create_stream_fallback(agent, api_kwargs: dict, client: Any = None): diff --git a/run_agent.py b/run_agent.py index 745e7c942286..b1f0008d8909 100644 --- a/run_agent.py +++ b/run_agent.py @@ -3880,6 +3880,13 @@ class AIAgent: except Exception: pass + # Also drop the cached per-request wire client (reused across + # sequential LLM calls) — same socket/memory rationale as above. + try: + self._close_cached_request_openai_client(reason="cache_evict") + except Exception: + pass + def close(self) -> None: """Release all resources held by this agent instance. @@ -3936,6 +3943,13 @@ class AIAgent: except Exception: pass + # 5b. Close the cached per-request wire client (reused across + # sequential LLM calls; see _create_request_openai_client). + try: + self._close_cached_request_openai_client(reason="agent_close") + except Exception: + pass + # 6. Free conversation history. Mirrors _release_evicted_agent_soft's # soft-eviction clear — close() is the hard teardown for true session # boundaries (/new, /reset, session expiry), so the message list won't @@ -4645,6 +4659,27 @@ class AIAgent: return copilot_request_headers(is_agent_turn=True, is_vision=is_vision) + # Close reasons the request workers' own ``finally`` unwind reports for + # a request that produced a response — the only closes that both come + # from the thread that owns the pool's FDs AND attest a healthy pool. + # Only these may keep the wire client for the next call, and poisoning + # still wins: a cross-thread abort (#29507) marks the slot so even a + # worker-finally close discards it. Every other reason (error cleanups, + # stale/interrupt kills, retry cleanups) gets a real close, so a retry + # after a request error always builds a fresh pool. + _REQUEST_CLIENT_REUSE_REASONS = frozenset({ + "request_complete", + "stream_request_complete", + }) + + def _request_client_cache_ref(self) -> dict: + # Lazy init — tests build agents via AIAgent.__new__ without __init__. + cache = getattr(self, "_request_client_cache", None) + if cache is None: + cache = {"client": None, "kwargs": None, "poisoned": False, "in_use": False} + self._request_client_cache = cache + return cache + def _create_request_openai_client(self, *, reason: str, api_kwargs: Optional[dict] = None) -> Any: from unittest.mock import Mock @@ -4671,9 +4706,96 @@ class AIAgent: and self._api_kwargs_have_image_parts(api_kwargs or {}) ): request_kwargs["default_headers"] = self._copilot_headers_for_request(is_vision=True) - return self._create_openai_client(request_kwargs, reason=reason, shared=False) + # Reuse the cached wire client while the effective kwargs are + # unchanged: constructing openai.OpenAI + its httpx pool costs + # ~19-35ms per LLM call (fresh TCP+TLS handshake), ~5x per turn. + # The cache is a single checked-out slot: `in_use` prevents two + # concurrent calls from sharing one pool's close/abort lifecycle + # (a second concurrent call gets a fresh untracked client with + # the old build-per-request behavior). + stale = None + with self._openai_client_lock(): + cache = self._request_client_cache_ref() + cached = cache["client"] + if cached is not None and not cache["in_use"]: + if ( + not cache["poisoned"] + and cache["kwargs"] == request_kwargs + and not self._is_openai_client_closed(cached) + ): + cache["in_use"] = True + return cached + # kwargs changed (credential rotation, provider failover), + # poisoned by a cross-thread abort (#29507), or externally + # closed — never reuse; discard and rebuild below. + stale = cached + cache["client"] = None + cache["kwargs"] = None + cache["poisoned"] = False + if stale is not None: + # Safe to close from this thread: in_use was False, so no + # worker thread owns the pool's FDs (#29507 concerns clients + # with an in-flight request on another thread). + self._close_openai_client(stale, reason=f"reuse_evict:{reason}", shared=False) + client = self._create_openai_client(request_kwargs, reason=reason, shared=False) + with self._openai_client_lock(): + cache = self._request_client_cache_ref() + if cache["client"] is None: + cache["client"] = client + # Snapshot nested dicts (default_headers): rotation sites + # assign fresh inner dicts today, but an aliased inner + # object would compare equal even after in-place mutation. + cache["kwargs"] = { + k: dict(v) if isinstance(v, dict) else v + for k, v in request_kwargs.items() + } + cache["poisoned"] = False + cache["in_use"] = True + # else: a concurrent call holds the slot — hand this client + # out untracked; _close_request_openai_client fully closes + # untracked clients, preserving the per-request lifecycle. + return client def _close_request_openai_client(self, client: Any, *, reason: str) -> None: + with self._openai_client_lock(): + cache = self._request_client_cache_ref() + if cache["client"] is client: + if reason in self._REQUEST_CLIENT_REUSE_REASONS and not cache["poisoned"]: + # Clean finish on the owning thread — keep the wire client + # (and its warm httpx pool) for the next sequential call. + cache["in_use"] = False + return + # Failure / kill / abort outcome: drop the slot and fall + # through to a real close. This runs on the owning worker + # thread, which is where the FD release belongs (#29507). + cache["client"] = None + cache["kwargs"] = None + cache["poisoned"] = False + cache["in_use"] = False + self._close_openai_client(client, reason=reason, shared=False) + + def _close_cached_request_openai_client(self, *, reason: str) -> None: + """Teardown hook: really close the cached per-request wire client.""" + with self._openai_client_lock(): + cache = getattr(self, "_request_client_cache", None) + client = cache["client"] if cache else None + in_use = bool(cache["in_use"]) if cache else False + if cache is not None: + cache["client"] = None + cache["kwargs"] = None + cache["poisoned"] = False + cache["in_use"] = False + if client is None: + return + if in_use: + # A worker thread has this client checked out for an in-flight + # request (workers can outlive turns — see interruptible_api_call). + # client.close() here would release its FDs from a stranger thread, + # the #29507 race teardown must not reintroduce. Abort the sockets + # instead; the slot is already cleared, so the worker's own finally + # sees an untracked client and does the real close on its thread. + self._abort_request_openai_client(client, reason=f"{reason}_in_flight") + return self._close_openai_client(client, reason=reason, shared=False) def _abort_request_openai_client(self, client: Any, *, reason: str) -> None: @@ -4692,6 +4814,13 @@ class AIAgent: """ if client is None: return + # A pool whose sockets were shut down from a stranger thread must + # never be reused: poison the cache slot so the owner-thread close + # discards it and the next create builds a fresh client. + with self._openai_client_lock(): + cache = self._request_client_cache_ref() + if cache["client"] is client: + cache["poisoned"] = True try: shutdown_count = self._force_close_tcp_sockets(client) # tcp_force_closed=0 means the stranger-thread abort found no diff --git a/tests/agent/test_request_client_reuse.py b/tests/agent/test_request_client_reuse.py new file mode 100644 index 000000000000..cda3c36d20b8 --- /dev/null +++ b/tests/agent/test_request_client_reuse.py @@ -0,0 +1,289 @@ +"""Per-request OpenAI wire client reuse across sequential LLM calls. + +Building a fresh ``openai.OpenAI`` client per LLM call costs ~19-35ms (new +httpx pool, TCP+TLS handshake), so ``_create_request_openai_client`` caches +ONE reusable wire client on the agent, keyed by the effective client kwargs: + +- identical kwargs → same client object handed back (the reuse win); +- kwargs change (credential rotation, provider failover) → evict + rebuild; +- cross-thread abort (#29507) poisons the slot → the owner-thread close does + a real close and the next create rebuilds; +- non-reuse close reasons (error cleanups, stale/interrupt kills, retry + cleanups) discard — only a request that produced a response reports a + reuse reason (request_complete / stream_request_complete); +- vision-header copilot variant is a distinct kwargs key; +- teardown (release_clients / close) really closes the cached client, or + detaches it to the in-flight worker's own close when checked out (#29507); +- MoA facade and Mock passthroughs never enter the cache. +""" + +from unittest.mock import MagicMock, patch + +from run_agent import AIAgent + + +class _StubClient: + """Minimal non-Mock client: _is_openai_client_closed reads ``is_closed``.""" + + def __init__(self): + self.is_closed = False + + def close(self): + self.is_closed = True + + +def _make_agent(provider="openai", base_url="https://api.openai.com/v1", model="gpt-5.4"): + with patch("run_agent.OpenAI") as mock_openai: + mock_openai.return_value = MagicMock() + agent = AIAgent( + api_key="sk-test", + base_url=base_url, + provider=provider, + model=model, + quiet_mode=True, + skip_context_files=True, + skip_memory=True, + ) + # Non-Mock shared client so the Mock passthrough branch doesn't trigger. + agent.client = _StubClient() + return agent + + +class _Harness: + """Patch the wire-client build/close/socket seams and record calls.""" + + def __init__(self, agent): + self.agent = agent + self.built = [] # (kwargs_copy, reason) + self.closed = [] # (client, reason) + self._patchers = [] + + def __enter__(self): + def _fake_create(kwargs, *, reason, shared): + # Only record per-request wire clients; teardown tests can also + # trigger a shared-client rebuild via _ensure_primary_openai_client. + if not shared: + self.built.append((dict(kwargs), reason)) + return _StubClient() + + def _fake_close(client, *, reason, shared): + self.closed.append((client, reason)) + + self._patchers = [ + patch.object(self.agent, "_create_openai_client", side_effect=_fake_create), + patch.object(self.agent, "_close_openai_client", side_effect=_fake_close), + patch.object(self.agent, "_force_close_tcp_sockets", return_value=0), + ] + for p in self._patchers: + p.start() + return self + + def __exit__(self, *exc): + for p in self._patchers: + p.stop() + + def closed_clients(self): + return [client for client, _reason in self.closed] + + +def test_reuse_on_identical_kwargs_same_object(): + agent = _make_agent() + with _Harness(agent) as h: + a = agent._create_request_openai_client(reason="chat_completion_request") + # The agent's outer loop owns retries — the SDK loop must stay off. + assert h.built[-1][0]["max_retries"] == 0 + agent._close_request_openai_client(a, reason="request_complete") + assert h.closed == [] # kept for reuse, not really closed + + b = agent._create_request_openai_client(reason="chat_completion_request") + assert b is a + assert len(h.built) == 1 + + +def test_reuse_after_streaming_clean_finish(): + agent = _make_agent() + with _Harness(agent) as h: + a = agent._create_request_openai_client(reason="chat_completion_stream_request") + agent._close_request_openai_client(a, reason="stream_request_complete") + b = agent._create_request_openai_client(reason="chat_completion_stream_request") + assert b is a + assert h.closed == [] + + +def test_rebuild_on_client_kwargs_change(): + agent = _make_agent() + with _Harness(agent) as h: + a = agent._create_request_openai_client(reason="r") + agent._close_request_openai_client(a, reason="request_complete") + + # Credential rotation / provider failover mutate _client_kwargs. + agent._client_kwargs["api_key"] = "sk-rotated" + b = agent._create_request_openai_client(reason="r") + assert b is not a + # The stale cached client was really closed on eviction. + assert a in h.closed_clients() + assert h.built[-1][0]["api_key"] == "sk-rotated" + + # And the rotated client is itself reusable. + agent._close_request_openai_client(b, reason="request_complete") + c = agent._create_request_openai_client(reason="r") + assert c is b + + +def test_rebuild_after_cross_thread_abort_poisons_cache(): + agent = _make_agent() + with _Harness(agent) as h: + a = agent._create_request_openai_client(reason="r") + # Stranger-thread abort (#29507): stale-call detector / interrupt loop + # shutdown(SHUT_RDWR) the pool's sockets. This client must never be + # reused even though the worker's own finally reports a clean reason. + agent._abort_request_openai_client(a, reason="stale_call_kill") + agent._close_request_openai_client(a, reason="request_complete") + assert a in h.closed_clients() # poisoned → real close, not kept + + b = agent._create_request_openai_client(reason="r") + assert b is not a + + +def test_kill_reason_close_discards_cached_client(): + agent = _make_agent() + with _Harness(agent) as h: + a = agent._create_request_openai_client(reason="r") + agent._close_request_openai_client(a, reason="stream_retry_cleanup") + assert a in h.closed_clients() + + b = agent._create_request_openai_client(reason="r") + assert b is not a + + +def test_externally_closed_cached_client_rebuilds(): + agent = _make_agent() + with _Harness(agent) as h: + a = agent._create_request_openai_client(reason="r") + agent._close_request_openai_client(a, reason="request_complete") + a.is_closed = True # e.g. closed behind our back + + b = agent._create_request_openai_client(reason="r") + assert b is not a + assert len(h.built) == 2 + + +def test_copilot_vision_variant_gets_own_client(): + agent = _make_agent(provider="copilot", base_url="https://api.githubcopilot.com") + text_kwargs = {"model": "gpt-5.4", "messages": [{"role": "user", "content": "hi"}]} + image_kwargs = { + "model": "gpt-5.4", + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is in this image?"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,abc"}}, + ], + } + ], + } + with _Harness(agent) as h: + a = agent._create_request_openai_client(reason="r", api_kwargs=text_kwargs) + agent._close_request_openai_client(a, reason="request_complete") + + # Vision request: different default_headers → must not reuse the + # text-variant client. + b = agent._create_request_openai_client(reason="r", api_kwargs=image_kwargs) + assert b is not a + assert h.built[-1][0]["default_headers"]["Copilot-Vision-Request"] == "true" + agent._close_request_openai_client(b, reason="request_complete") + + # Consecutive vision requests reuse the vision-variant client. + c = agent._create_request_openai_client(reason="r", api_kwargs=image_kwargs) + assert c is b + + +def test_release_clients_closes_cached_request_client(): + agent = _make_agent() + with _Harness(agent) as h: + a = agent._create_request_openai_client(reason="r") + agent._close_request_openai_client(a, reason="request_complete") + + agent.release_clients() + assert (a, "cache_evict") in h.closed + + # Next create rebuilds instead of handing back the closed client. + b = agent._create_request_openai_client(reason="r") + assert b is not a + + +def test_agent_close_closes_cached_request_client(): + agent = _make_agent() + with _Harness(agent) as h: + a = agent._create_request_openai_client(reason="r") + agent._close_request_openai_client(a, reason="request_complete") + + agent.close() + assert (a, "agent_close") in h.closed + + # Idempotent: a second teardown must not double-close. + before = len(h.closed) + agent._close_cached_request_openai_client(reason="agent_close") + assert len(h.closed) == before + + +def test_teardown_while_checked_out_defers_close_to_worker(): + agent = _make_agent() + with _Harness(agent) as h: + a = agent._create_request_openai_client(reason="r") + # Checked out by an in-flight worker (workers can outlive turns): + # teardown must not client.close() from a stranger thread (#29507) — + # it aborts the sockets and detaches the slot instead. + agent._close_cached_request_openai_client(reason="agent_close") + assert a not in h.closed_clients() + + # The worker's own finally sees an untracked client → real close on + # the owning thread, even with a clean-finish reason. + agent._close_request_openai_client(a, reason="request_complete") + assert a in h.closed_clients() + + b = agent._create_request_openai_client(reason="r") + assert b is not a + + +def test_concurrent_checkout_gets_untracked_client(): + agent = _make_agent() + with _Harness(agent) as h: + a = agent._create_request_openai_client(reason="r") + # Slot checked out by the first (still in-flight) call: a concurrent + # call gets a fresh client with the old per-request lifecycle. + b = agent._create_request_openai_client(reason="r") + assert b is not a + + agent._close_request_openai_client(b, reason="request_complete") + assert b in h.closed_clients() # untracked → really closed + + agent._close_request_openai_client(a, reason="request_complete") + assert a not in h.closed_clients() # cached → kept + + c = agent._create_request_openai_client(reason="r") + assert c is a + + +def test_moa_passthrough_unaffected(): + agent = _make_agent() + agent.provider = "moa" + facade = agent.client + with _Harness(agent) as h: + a = agent._create_request_openai_client(reason="r") + assert a is facade + assert h.built == [] # facade handed back, no wire client built + + # Close behaves exactly as before the cache existed. + agent._close_request_openai_client(facade, reason="request_complete") + assert facade in h.closed_clients() + + +def test_mock_client_passthrough_unaffected(): + agent = _make_agent() + agent.client = MagicMock() + with _Harness(agent) as h: + a = agent._create_request_openai_client(reason="r") + assert a is agent.client + assert h.built == [] diff --git a/tests/run_agent/test_openai_client_lifecycle.py b/tests/run_agent/test_openai_client_lifecycle.py index e38c1f726e4b..dafe4484307d 100644 --- a/tests/run_agent/test_openai_client_lifecycle.py +++ b/tests/run_agent/test_openai_client_lifecycle.py @@ -92,7 +92,11 @@ def test_retry_after_api_connection_error_recreates_request_client(monkeypatch): assert result == {"ok": True} assert len(factory.calls) == 2 assert first_request.close_calls >= 1 - assert second_request.close_calls >= 1 + # The successful request's wire client is cached for reuse across + # sequential calls (not closed at request end); teardown really closes it. + assert second_request.close_calls == 0 + agent._close_cached_request_openai_client(reason="agent_close") + assert second_request.close_calls == 1 def test_stale_non_stream_close_is_single_owner(monkeypatch): @@ -205,5 +209,9 @@ def test_streaming_call_recreates_closed_shared_client_before_request(monkeypatc assert response.choices[0].message.content == "Hello world" assert agent.client is replacement_shared assert stale_shared.close_calls >= 1 - assert request_client.close_calls >= 1 + # The clean stream's wire client is cached for reuse across sequential + # calls (not closed at request end); teardown really closes it. + assert request_client.close_calls == 0 + agent._close_cached_request_openai_client(reason="agent_close") + assert request_client.close_calls == 1 assert len(factory.calls) == 2 diff --git a/tests/run_agent/test_request_client_reuse_abort_races.py b/tests/run_agent/test_request_client_reuse_abort_races.py new file mode 100644 index 000000000000..73a9c41887cc --- /dev/null +++ b/tests/run_agent/test_request_client_reuse_abort_races.py @@ -0,0 +1,369 @@ +"""Races between the request-client reuse cache and the abort machinery. + +Invariants pinned here: + +1. A worker-side interrupt that breaks out of the SSE chunk loop must close + the half-read stream first (owner thread). Abandoning it leaves its + connection permanently checked out of the httpx pool while the partial + response makes the worker's finally report a reuse-reason close, caching + the client with the leaked connection — one more per interrupt until the + pool hits ``max_connections`` and every request dies with PoolTimeout. + If the close fails, the slot must be poisoned so the finally really + closes the pool. + +2. The stranger-thread abort (stale detector / interrupt loop) must read + the holder and fire the abort atomically, under ``request_client_lock``. + An abort landing after the lock is released races the worker's finally, + which can pop + cache the client and let the NEXT call check it out — + the late abort would then poison the slot and shut down an innocent + in-flight request's sockets. + +3. The same atomicity contract holds at the third holder-abort site: + ``direct_api_call``'s ``_abort_active_request`` (the cron inline path, + reached cross-thread via ``AIAgent.interrupt()``). + +4. ``run_codex_stream``'s finally must poison the reuse slot when + ``event_stream.close()`` fails; otherwise a failed close caches the + client with a connection still checked out of the pool. +""" +import threading +import time +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + + +def _make_agent(): + from run_agent import AIAgent + + agent = AIAgent( + api_key="test-key", + base_url="https://example.com/v1", + model="test/model", + quiet_mode=True, + skip_context_files=True, + skip_memory=True, + ) + agent.api_mode = "chat_completions" + return agent + + +def _chunk(content=None, finish_reason=None): + return SimpleNamespace( + choices=[ + SimpleNamespace( + index=0, + delta=SimpleNamespace( + content=content, + tool_calls=None, + reasoning_content=None, + reasoning=None, + ), + finish_reason=finish_reason, + ) + ], + model="test/model", + usage=None, + ) + + +class _FakeStream: + """SSE stream stand-in. + + Deliberately has NO ``choices`` attribute so ``_call_chat_completions`` + treats it as a genuine token stream (a MagicMock would auto-create + ``choices`` and get misread as a completed response object). + """ + + response = None + + def __init__(self, chunk_iter_factory, close_raises=False): + self._factory = chunk_iter_factory + self._close_raises = close_raises + self.close_calls = 0 + + def __iter__(self): + return self._factory() + + def close(self): + self.close_calls += 1 + if self._close_raises: + raise RuntimeError("close failed") + + +def _mock_wire_client(stream): + client = MagicMock() + client.chat.completions.create.return_value = stream + return client + + +@pytest.mark.filterwarnings("ignore::pytest.PytestUnhandledThreadExceptionWarning") +def test_worker_interrupt_break_closes_stream(): + """Interrupt noticed between chunks must close the half-read stream. + + Without the close, the connection stays checked out of the httpx pool + while the partial-response finally caches the client for reuse — the + leak that eventually exhausted the pool (PoolTimeout on every request). + """ + agent = _make_agent() + + def chunks(): + yield _chunk(content="partial ") + # /stop arrives while the provider is still streaming. + agent._interrupt_requested = True + yield _chunk(content="never processed") + + stream = _FakeStream(chunks) + + with patch.object( + agent, "_create_request_openai_client", return_value=_mock_wire_client(stream) + ), patch.object(agent, "_close_request_openai_client"): + with pytest.raises(InterruptedError): + agent._interruptible_streaming_api_call({}) + + assert stream.close_calls == 1 + + +@pytest.mark.filterwarnings("ignore::pytest.PytestUnhandledThreadExceptionWarning") +def test_worker_interrupt_break_poisons_slot_when_stream_close_fails(): + """If the half-read stream can't be released, the client must not be + cached: the owner-thread abort poisons the slot so the worker's finally + really closes the pool (leaked connection and all).""" + agent = _make_agent() + + def chunks(): + yield _chunk(content="partial ") + agent._interrupt_requested = True + yield _chunk(content="never processed") + + stream = _FakeStream(chunks, close_raises=True) + abort_reasons = [] + + with patch.object( + agent, "_create_request_openai_client", return_value=_mock_wire_client(stream) + ), patch.object(agent, "_close_request_openai_client"), patch.object( + agent, + "_abort_request_openai_client", + side_effect=lambda client, *, reason: abort_reasons.append(reason), + ): + with pytest.raises(InterruptedError): + agent._interruptible_streaming_api_call({}) + + assert stream.close_calls == 1 + assert "interrupt_stream_close_failed" in abort_reasons + + +def test_stale_abort_is_atomic_with_holder_read(monkeypatch): + """The stranger-thread abort must complete before the worker's finally + can pop + cache the client. + + The abort runs under ``request_client_lock``; a worker that finishes + while the abort is in flight must block in its finally until the abort + returns. Without that, the worker could cache the client (and the next + call check it out) between the holder read and the abort — the abort + then killed an innocent request's sockets. + """ + monkeypatch.setenv("HERMES_STREAM_STALE_TIMEOUT", "0.05") + agent = _make_agent() + + allow_finish = threading.Event() + worker_close_reasons = [] + abort_reasons = [] + observed = {"worker_finished_during_abort": None} + + def chunks(): + yield _chunk(content="hello") + # Stall long enough for the stale detector to fire, then finish + # cleanly the moment the abort (below) unblocks us. + allow_finish.wait(timeout=5.0) + yield _chunk(finish_reason="stop") + + stream = _FakeStream(chunks) + + def fake_abort(client, *, reason): + abort_reasons.append(reason) + # Let the worker race toward its finally while the abort is still + # in flight. Under the fix it must block on the holder lock, so the + # owner-side close cannot land until this abort returns. + allow_finish.set() + deadline = time.time() + 0.6 + while time.time() < deadline and not worker_close_reasons: + time.sleep(0.02) + observed["worker_finished_during_abort"] = bool(worker_close_reasons) + + with patch.object( + agent, "_create_request_openai_client", return_value=_mock_wire_client(stream) + ), patch.object( + agent, + "_close_request_openai_client", + side_effect=lambda client, *, reason: worker_close_reasons.append(reason), + ), patch.object( + agent, "_abort_request_openai_client", side_effect=fake_abort + ), patch.object( + agent, "_replace_primary_openai_client" + ): + response = agent._interruptible_streaming_api_call({}) + + assert response is not None + assert "stale_stream_kill" in abort_reasons + # The atomicity contract: no owner-side close slipped in mid-abort. + assert observed["worker_finished_during_abort"] is False + # ...and the worker's own finally still performed its close afterwards. + # The stale kill cancels the stream attempt before aborting, so the + # worker's late clean finish is treated as a superseded stream and the + # finally reports the error-cleanup reason — really closing the + # socket-aborted (poisoned) client instead of caching it for reuse. + assert worker_close_reasons == ["stream_error_cleanup"] + + +def test_direct_api_call_abort_is_atomic_with_holder_read(): + """Same atomicity contract for the cron inline path's holder abort. + + ``direct_api_call``'s ``_abort_active_request`` (fired cross-thread via + ``AIAgent.interrupt()`` -> ``_active_request_abort``) must complete + before the inline finally can pop + cache the client — otherwise the + delayed abort poisons the slot and kills the NEXT request's sockets. + """ + from agent.chat_completion_helpers import direct_api_call + + agent = _make_agent() + + allow_finish = threading.Event() + close_reasons = [] + abort_reasons = [] + observed = {"close_during_abort": None} + + client = MagicMock() + + def blocking_create(**kwargs): + # Block until the aborting thread lets us race toward the finally. + allow_finish.wait(timeout=5.0) + return SimpleNamespace(choices=[]) + + client.chat.completions.create.side_effect = blocking_create + + def fake_abort(aborted_client, *, reason): + abort_reasons.append(reason) + # Unblock the owner mid-abort. Under the fix it must block in its + # finally on the holder lock until this abort returns. + allow_finish.set() + deadline = time.time() + 0.6 + while time.time() < deadline and not close_reasons: + time.sleep(0.02) + observed["close_during_abort"] = bool(close_reasons) + + with patch.object( + agent, "_create_request_openai_client", return_value=client + ), patch.object( + agent, + "_close_request_openai_client", + side_effect=lambda c, *, reason: close_reasons.append(reason), + ), patch.object( + agent, "_abort_request_openai_client", side_effect=fake_abort + ): + result = {} + + def owner(): + result["response"] = direct_api_call(agent, {}) + + owner_thread = threading.Thread(target=owner, daemon=True) + owner_thread.start() + deadline = time.time() + 5.0 + while time.time() < deadline and getattr(agent, "_active_request_abort", None) is None: + time.sleep(0.01) + abort_fn = getattr(agent, "_active_request_abort", None) + assert abort_fn is not None + # Stranger-thread abort (this thread) racing the inline finally. + abort_fn("test_stranger_abort") + owner_thread.join(timeout=5.0) + assert not owner_thread.is_alive() + + assert abort_reasons == ["test_stranger_abort"] + # The atomicity contract: no owner-side close slipped in mid-abort. + assert observed["close_during_abort"] is False + # ...and the inline finally still performed its close afterwards. + assert close_reasons == ["request_complete"] + assert result["response"] is not None + + +class _FakeCodexEventStream: + """Codex Responses SSE stream stand-in (iterable, no ``output`` attr).""" + + def __init__(self, events, close_raises=False): + self._events = events + self._close_raises = close_raises + self.close_calls = 0 + + def __iter__(self): + return iter(self._events) + + def close(self): + self.close_calls += 1 + if self._close_raises: + raise RuntimeError("close failed") + + +def _codex_delta_event(text): + return SimpleNamespace(type="response.output_text.delta", delta=text) + + +def _codex_completed_event(): + return SimpleNamespace( + type="response.completed", + response=SimpleNamespace(usage=None, id="resp_1", status="completed"), + ) + + +def test_codex_stream_close_failure_poisons_slot(): + """A failed ``event_stream.close()`` must poison the reuse slot. + + Otherwise the connection stays checked out of the pool while the + worker's finally reports ``request_complete`` and caches the client + with the leaked connection (the codex sibling of the chat-streaming + interrupt-break leak). + """ + from agent.codex_runtime import run_codex_stream + + agent = _make_agent() + stream = _FakeCodexEventStream( + [_codex_delta_event("partial "), _codex_completed_event()], + close_raises=True, + ) + client = MagicMock() + client.responses.create.return_value = stream + abort_reasons = [] + + with patch.object( + agent, + "_abort_request_openai_client", + side_effect=lambda c, *, reason: abort_reasons.append(reason), + ): + final = run_codex_stream(agent, {"model": "test/model"}, client=client) + + assert final.output_text == "partial " + assert stream.close_calls == 1 + assert abort_reasons == ["codex_stream_close_failed"] + + +def test_codex_stream_close_failure_on_primary_client_does_not_abort(): + """``client=None`` means the shared primary client — never reuse-cached, + and its sockets must not be force-shut by the close-failure handler.""" + from agent.codex_runtime import run_codex_stream + + agent = _make_agent() + stream = _FakeCodexEventStream( + [_codex_delta_event("hello"), _codex_completed_event()], + close_raises=True, + ) + primary = MagicMock() + primary.responses.create.return_value = stream + + with patch.object( + agent, "_ensure_primary_openai_client", return_value=primary + ), patch.object(agent, "_abort_request_openai_client") as abort_mock: + final = run_codex_stream(agent, {"model": "test/model"}, client=None) + + assert final.output_text == "hello" + assert stream.close_calls == 1 + abort_mock.assert_not_called()