mirrors/hermes-agent
1
0
Fork 0
mirror of https://github.com/NousResearch/hermes-agent.git synced 2026-07-29 18:46:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 7

perf(agent): reuse the per-request OpenAI wire client across sequential LLM calls

Browse source 
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.
This commit is contained in:
Soju06 2026-07-14 03:57:59 +00:00 • committed by kshitij
parent a41dc65ba7
commit 82e2c9ce40
6 changed files with 899 additions and 30 deletions
Download patch file Download diff file Expand all files Collapse all files

113
agent/chat_completion_helpers.py
View file

@ -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)

15
agent/codex_runtime.py
View file

@ -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):

131
run_agent.py
View file

@ -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

289
tests/agent/test_request_client_reuse.py Normal file
View file

@ -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 == []

12
tests/run_agent/test_openai_client_lifecycle.py
View file

@ -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

369
tests/run_agent/test_request_client_reuse_abort_races.py Normal file
View file

@ -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()