fix(agent): retire replaced shared OpenAI clients instead of cross-thread pool close
Widen the #70773 fix beyond the three in-request cleanup sites removed in the cherry-picked commit: every remaining path that swaps out the shared OpenAI client could still hard-close its pool from a thread that doesn't own the in-flight sockets (credential rotation/refresh on the turn thread, dead-connection cleanup, gateway cache eviction, transport recovery) — the same FD-recycle corruption vector, just rarer. Add AIAgent._retire_shared_openai_client(): shutdown(SHUT_RDWR) all pooled sockets (FD-safe from any thread, unblocks in-flight readers) but never call client.close() — FD release is deferred to GC, which cannot run until every borrowing thread has unwound its SSL BIO. Refcounting is the ownership handshake; with no borrowers the FDs are released immediately. Wired into: - _replace_primary_openai_client (rotation/refresh/dead-conn cleanup) - try_recover_primary_transport (primary_recovery) - release_clients (gateway cache_evict) agent.close() keeps the hard close: full teardown is a real session boundary where no request may be in flight. Tests: new tests/run_agent/test_70773_shared_client_fd_corruption.py covers the three watchdog/retry sites plus retire semantics; existing close-assertions updated to pin retire-not-close.
This commit is contained in:
teknium1
@@ -1225,11 +1225,17 @@ def try_recover_primary_transport(
|
||||
return False
|
||||
|
||||
try:
|
||||
# Close existing client to release stale connections
|
||||
# Retire the existing client to release stale connections. #70773:
|
||||
# never hard-close the shared client here — this runs on the
|
||||
# conversation-loop thread while workers from stale-killed streaming
|
||||
# attempts may still be unwinding their SSL BIOs on the old pool.
|
||||
# ``_retire_shared_openai_client`` shuts the sockets down (FD-safe
|
||||
# from any thread) and defers the FD release to GC, which cannot
|
||||
# complete until every borrowing thread has unwound.
|
||||
if getattr(agent, "client", None) is not None:
|
||||
try:
|
||||
agent._close_openai_client(
|
||||
agent.client, reason="primary_recovery", shared=True,
|
||||
agent._retire_shared_openai_client(
|
||||
agent.client, reason="primary_recovery",
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
+55
-3
@@ -3817,11 +3817,16 @@ class AIAgent:
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Close the OpenAI/httpx client to release sockets immediately.
|
||||
# Retire the OpenAI/httpx client to release sockets immediately.
|
||||
# #70773: eviction runs on the gateway's memory-manager thread — a
|
||||
# cross-thread hard close of the shared client can release TLS FDs
|
||||
# under a still-unwinding worker (FD-recycle → SQLite corruption).
|
||||
# Retirement shuts the pooled sockets down (the memory/socket win we
|
||||
# want here) and lets GC release the FDs once no thread holds them.
|
||||
try:
|
||||
client = getattr(self, "client", None)
|
||||
if client is not None:
|
||||
self._close_openai_client(client, reason="cache_evict", shared=True)
|
||||
self._retire_shared_openai_client(client, reason="cache_evict")
|
||||
self.client = None
|
||||
except Exception:
|
||||
pass
|
||||
@@ -4385,6 +4390,47 @@ class AIAgent:
|
||||
exc,
|
||||
)
|
||||
|
||||
def _retire_shared_openai_client(self, client: Any, *, reason: str) -> None:
|
||||
"""Ownership-safe retirement of a replaced shared OpenAI client.
|
||||
|
||||
#70773 / #67142 / #29507: ``client.close()`` releases the pool's raw
|
||||
FDs from the *calling* thread. The shared primary client has no single
|
||||
owning thread — worker threads from stale-killed attempts may still be
|
||||
unwinding their SSL BIOs, and the codex-direct / MoA paths stream on
|
||||
the shared client itself. If we release an FD while another thread's
|
||||
SSL layer still caches the raw integer fd, the kernel can recycle it
|
||||
into an unrelated ``open()`` (e.g. ``kanban.db``) and the unwinding
|
||||
TLS flush then writes an application-data record into that file — the
|
||||
SQLite-header corruption documented in #29507/#70773.
|
||||
|
||||
Only an owner may release FDs, and a replaced shared client has none.
|
||||
So nobody calls ``close()``: we ``shutdown()`` the pooled sockets
|
||||
(FD-safe from any thread; unblocks in-flight readers with EOF/EPIPE)
|
||||
and defer the actual FD release to garbage collection. Refcounting
|
||||
guarantees the underlying sockets are only collected once every
|
||||
thread that borrowed the client has unwound — GC *is* the ownership
|
||||
handshake. In the common case (no borrower) the refcount hits zero on
|
||||
this line and the FDs are released immediately anyway.
|
||||
"""
|
||||
if client is None:
|
||||
return
|
||||
try:
|
||||
shutdown_count = self._force_close_tcp_sockets(client)
|
||||
logger.info(
|
||||
"Shared OpenAI client retired (%s, tcp_shutdown=%d, "
|
||||
"fd_release=deferred_to_gc) %s",
|
||||
reason,
|
||||
shutdown_count,
|
||||
self._client_log_context(),
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.debug(
|
||||
"Shared OpenAI client retire failed (%s) %s error=%s",
|
||||
reason,
|
||||
self._client_log_context(),
|
||||
exc,
|
||||
)
|
||||
|
||||
def _replace_primary_openai_client(self, *, reason: str) -> bool:
|
||||
with self._openai_client_lock():
|
||||
old_client = getattr(self, "client", None)
|
||||
@@ -4399,7 +4445,13 @@ class AIAgent:
|
||||
)
|
||||
return False
|
||||
self.client = new_client
|
||||
self._close_openai_client(old_client, reason=f"replace:{reason}", shared=True)
|
||||
# #70773: never hard-close the replaced shared client from here — the
|
||||
# caller may not be the thread whose request is still unwinding on the
|
||||
# old pool (credential rotation and dead-connection cleanup run on the
|
||||
# turn thread while stale-killed workers unwind; the codex-direct path
|
||||
# streams on the shared client itself). Retire it instead: sockets are
|
||||
# shut down (FD-safe), FD release deferred to GC.
|
||||
self._retire_shared_openai_client(old_client, reason=f"replace:{reason}")
|
||||
return True
|
||||
|
||||
def _ensure_primary_openai_client(self, *, reason: str) -> Any:
|
||||
|
||||
@@ -0,0 +1,243 @@
|
||||
"""#70773: the shared OpenAI-wire client must never be pool-closed from a
|
||||
non-owner thread.
|
||||
|
||||
A custom OpenAI-compatible provider reproduced the #29507/#67142 TLS-FD →
|
||||
SQLite corruption shape on v0.18.2: the streaming stale watchdog (poll
|
||||
thread) called ``_replace_primary_openai_client(reason=
|
||||
"stale_stream_pool_cleanup")``, which hard-closed the old shared client's
|
||||
connection pool while worker threads from previous stale-killed attempts
|
||||
were still unwinding their SSL BIOs. The kernel recycled a just-released
|
||||
TLS FD onto ``kanban.db`` and the unwinding TLS flush wrote a 24-byte
|
||||
application-data record over the SQLite header.
|
||||
|
||||
Fix (mirrors the #67142 Anthropic ownership discipline):
|
||||
|
||||
1. The three in-request cleanup sites (stale watchdog, mid-tool retry,
|
||||
stream retry) no longer touch the shared client at all — the
|
||||
request-local client is closed via ``_close_request_client_once`` and
|
||||
the shared client is replaced lazily by ``_ensure_primary_openai_client``
|
||||
on the next request, on the requesting thread.
|
||||
2. ``_replace_primary_openai_client`` (still used by credential rotation,
|
||||
refresh, and dead-connection cleanup) retires the old client instead of
|
||||
hard-closing it: sockets are ``shutdown()`` (FD-safe from any thread),
|
||||
FD release is deferred to GC so it cannot happen under a borrowing
|
||||
thread's live SSL BIO.
|
||||
"""
|
||||
import threading
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from tests.run_agent.test_streaming import _make_stream_chunk
|
||||
|
||||
|
||||
def _make_agent():
|
||||
from run_agent import AIAgent
|
||||
|
||||
agent = AIAgent(
|
||||
api_key="test-key",
|
||||
base_url="https://custom.example.com/v1",
|
||||
provider="custom",
|
||||
model="deepseek-v4-pro",
|
||||
quiet_mode=True,
|
||||
skip_context_files=True,
|
||||
skip_memory=True,
|
||||
)
|
||||
agent.api_mode = "chat_completions"
|
||||
agent._interrupt_requested = False
|
||||
return agent
|
||||
|
||||
|
||||
class TestStaleWatchdogNeverClosesSharedClient:
|
||||
"""The stale watchdog / retry cleanups must not rebuild (and therefore
|
||||
must not close) the shared OpenAI client from inside a request."""
|
||||
|
||||
@pytest.mark.filterwarnings(
|
||||
"ignore::pytest.PytestUnhandledThreadExceptionWarning"
|
||||
)
|
||||
@patch("run_agent.AIAgent._replace_primary_openai_client")
|
||||
@patch("run_agent.AIAgent._abort_request_openai_client")
|
||||
@patch("run_agent.AIAgent._create_request_openai_client")
|
||||
@patch("run_agent.AIAgent._close_request_openai_client")
|
||||
def test_stale_stream_kill_does_not_replace_primary(
|
||||
self, mock_close, mock_create, mock_abort, mock_replace, monkeypatch,
|
||||
):
|
||||
"""Stale-stream watchdog fires → request-local client is aborted, the
|
||||
shared primary is never replaced (poll thread must not close it)."""
|
||||
monkeypatch.setenv("HERMES_STREAM_STALE_TIMEOUT", "0.05")
|
||||
monkeypatch.setenv("HERMES_STREAM_RETRIES", "1")
|
||||
|
||||
unblock = threading.Event()
|
||||
|
||||
class StaleThenDeadStream:
|
||||
response = SimpleNamespace(headers={})
|
||||
|
||||
def __iter__(self):
|
||||
# Yield nothing; block until the watchdog aborts the request
|
||||
# client, then surface the transport error.
|
||||
unblock.wait(timeout=5.0)
|
||||
raise httpx.ConnectError("connection dropped after abort")
|
||||
yield # pragma: no cover — make this a generator
|
||||
|
||||
retry_chunks = [
|
||||
_make_stream_chunk(content="recovered"),
|
||||
_make_stream_chunk(finish_reason="stop", model="test/model"),
|
||||
]
|
||||
|
||||
class RetryStream:
|
||||
response = SimpleNamespace(headers={})
|
||||
|
||||
def __iter__(self):
|
||||
return iter(retry_chunks)
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.chat.completions.create.side_effect = [
|
||||
StaleThenDeadStream(),
|
||||
RetryStream(),
|
||||
]
|
||||
mock_create.return_value = mock_client
|
||||
mock_abort.side_effect = lambda *a, **k: unblock.set()
|
||||
|
||||
agent = _make_agent()
|
||||
response = agent._interruptible_streaming_api_call({})
|
||||
|
||||
assert response.choices[0].message.content == "recovered"
|
||||
# The watchdog aborted the request-local client from the poll thread…
|
||||
assert mock_abort.call_count >= 1
|
||||
# …and NEVER replaced/closed the shared primary client (#70773).
|
||||
mock_replace.assert_not_called()
|
||||
|
||||
@pytest.mark.filterwarnings(
|
||||
"ignore::pytest.PytestUnhandledThreadExceptionWarning"
|
||||
)
|
||||
@patch("run_agent.AIAgent._replace_primary_openai_client")
|
||||
@patch("run_agent.AIAgent._create_request_openai_client")
|
||||
@patch("run_agent.AIAgent._close_request_openai_client")
|
||||
def test_stream_retry_cleanup_does_not_replace_primary(
|
||||
self, mock_close, mock_create, mock_replace, monkeypatch,
|
||||
):
|
||||
"""Connection drop before first delta → retry path closes only the
|
||||
request-local client; the shared primary is left alone."""
|
||||
monkeypatch.setenv("HERMES_STREAM_RETRIES", "1")
|
||||
|
||||
attempts = {"n": 0}
|
||||
|
||||
def _pick_stream(*a, **kw):
|
||||
attempts["n"] += 1
|
||||
if attempts["n"] == 1:
|
||||
def _dead():
|
||||
raise httpx.ConnectError("connection reset by peer")
|
||||
yield # pragma: no cover
|
||||
return _dead()
|
||||
return iter([
|
||||
_make_stream_chunk(content="ok"),
|
||||
_make_stream_chunk(finish_reason="stop", model="test/model"),
|
||||
])
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.chat.completions.create.side_effect = _pick_stream
|
||||
mock_create.return_value = mock_client
|
||||
|
||||
agent = _make_agent()
|
||||
response = agent._interruptible_streaming_api_call({})
|
||||
|
||||
assert attempts["n"] == 2
|
||||
assert response.choices[0].message.content == "ok"
|
||||
# Request-local cleanup ran; shared client untouched (#70773).
|
||||
assert mock_close.call_count >= 1
|
||||
mock_replace.assert_not_called()
|
||||
|
||||
@pytest.mark.filterwarnings(
|
||||
"ignore::pytest.PytestUnhandledThreadExceptionWarning"
|
||||
)
|
||||
@patch("run_agent.AIAgent._replace_primary_openai_client")
|
||||
@patch("run_agent.AIAgent._create_request_openai_client")
|
||||
@patch("run_agent.AIAgent._close_request_openai_client")
|
||||
def test_mid_tool_retry_cleanup_does_not_replace_primary(
|
||||
self, mock_close, mock_create, mock_replace, monkeypatch,
|
||||
):
|
||||
"""Connection drop mid tool-call → silent retry closes only the
|
||||
request-local client; the shared primary is left alone."""
|
||||
from tests.run_agent.test_streaming import _make_tool_call_delta
|
||||
|
||||
monkeypatch.setenv("HERMES_STREAM_RETRIES", "2")
|
||||
|
||||
attempts = {"n": 0}
|
||||
|
||||
def _first_stream():
|
||||
yield _make_stream_chunk(content="Working: ")
|
||||
yield _make_stream_chunk(tool_calls=[
|
||||
_make_tool_call_delta(index=0, tc_id="call_1", name="terminal"),
|
||||
])
|
||||
raise httpx.RemoteProtocolError("peer closed connection")
|
||||
|
||||
def _second_stream():
|
||||
yield _make_stream_chunk(tool_calls=[
|
||||
_make_tool_call_delta(index=0, tc_id="call_1", name="terminal"),
|
||||
])
|
||||
yield _make_stream_chunk(tool_calls=[
|
||||
_make_tool_call_delta(index=0, arguments='{"command": "ls"}'),
|
||||
])
|
||||
yield _make_stream_chunk(finish_reason="tool_calls")
|
||||
|
||||
def _pick_stream(*a, **kw):
|
||||
attempts["n"] += 1
|
||||
return _first_stream() if attempts["n"] == 1 else _second_stream()
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.chat.completions.create.side_effect = _pick_stream
|
||||
mock_create.return_value = mock_client
|
||||
|
||||
agent = _make_agent()
|
||||
response = agent._interruptible_streaming_api_call({})
|
||||
|
||||
assert attempts["n"] == 2
|
||||
assert response.choices[0].message.tool_calls
|
||||
mock_replace.assert_not_called()
|
||||
|
||||
|
||||
class TestReplacePrimaryRetiresInsteadOfClosing:
|
||||
"""_replace_primary_openai_client (credential rotation / refresh /
|
||||
dead-connection cleanup) must retire the old shared client — shutdown
|
||||
sockets, defer FD release to GC — never hard-close its pool."""
|
||||
|
||||
def test_replace_retires_old_client_without_close(self):
|
||||
agent = _make_agent()
|
||||
|
||||
old_client = MagicMock()
|
||||
agent.client = old_client
|
||||
agent._client_kwargs = {
|
||||
"api_key": "test-key",
|
||||
"base_url": "https://custom.example.com/v1",
|
||||
}
|
||||
|
||||
shutdown_calls = []
|
||||
with patch.object(
|
||||
agent, "_force_close_tcp_sockets",
|
||||
side_effect=lambda c: shutdown_calls.append(c) or 1,
|
||||
):
|
||||
with patch("run_agent.OpenAI", MagicMock()):
|
||||
ok = agent._replace_primary_openai_client(reason="test_rotate")
|
||||
|
||||
assert ok
|
||||
# Sockets were shut down (FD-safe from any thread)…
|
||||
assert shutdown_calls == [old_client]
|
||||
# …but the old client's pool was NOT hard-closed: no thread owns a
|
||||
# replaced shared client, so nobody may release its FDs (#70773).
|
||||
old_client.close.assert_not_called()
|
||||
|
||||
def test_retire_helper_never_calls_close(self):
|
||||
agent = _make_agent()
|
||||
client = MagicMock()
|
||||
|
||||
with patch.object(agent, "_force_close_tcp_sockets", return_value=2):
|
||||
agent._retire_shared_openai_client(client, reason="unit_test")
|
||||
|
||||
client.close.assert_not_called()
|
||||
|
||||
def test_retire_helper_none_is_noop(self):
|
||||
agent = _make_agent()
|
||||
# Must not raise.
|
||||
agent._retire_shared_openai_client(None, reason="unit_test")
|
||||
@@ -577,20 +577,26 @@ class TestTryRecoverPrimaryTransport:
|
||||
# wait_time = min(3 + 10, 8) = 8
|
||||
mock_sleep.assert_called_once_with(8)
|
||||
|
||||
def test_closes_existing_client_before_rebuild(self):
|
||||
def test_retires_existing_client_before_rebuild(self):
|
||||
"""#70773: the old shared client is retired (sockets shutdown, FD
|
||||
release deferred to GC), never hard-closed — this path runs on the
|
||||
conversation-loop thread while stale-killed workers may still be
|
||||
unwinding on the old pool."""
|
||||
agent = _make_agent(provider="custom")
|
||||
old_client = agent.client
|
||||
error = _make_transport_error("ReadTimeout")
|
||||
|
||||
with patch("run_agent.OpenAI", return_value=MagicMock()), \
|
||||
patch("time.sleep"), \
|
||||
patch.object(agent, "_close_openai_client") as mock_close:
|
||||
patch.object(agent, "_close_openai_client") as mock_close, \
|
||||
patch.object(agent, "_retire_shared_openai_client") as mock_retire:
|
||||
agent._try_recover_primary_transport(
|
||||
error, retry_count=3, max_retries=3,
|
||||
)
|
||||
mock_close.assert_called_once_with(
|
||||
old_client, reason="primary_recovery", shared=True,
|
||||
mock_retire.assert_called_once_with(
|
||||
old_client, reason="primary_recovery",
|
||||
)
|
||||
mock_close.assert_not_called()
|
||||
|
||||
def test_survives_rebuild_failure(self):
|
||||
"""If client rebuild fails, returns False gracefully."""
|
||||
|
||||
@@ -6550,6 +6550,7 @@ class TestNousCredentialRefresh:
|
||||
agent.api_mode = "chat_completions"
|
||||
|
||||
closed = {"value": False}
|
||||
retired = {"value": False}
|
||||
rebuilt = {"kwargs": None}
|
||||
captured = {}
|
||||
|
||||
@@ -6575,12 +6576,28 @@ class TestNousCredentialRefresh:
|
||||
"hermes_cli.auth.resolve_nous_runtime_credentials", _fake_resolve
|
||||
)
|
||||
|
||||
agent.client = _ExistingClient()
|
||||
existing = _ExistingClient()
|
||||
agent.client = existing
|
||||
|
||||
_orig_retire = agent._retire_shared_openai_client
|
||||
|
||||
def _spy_retire(client, *, reason):
|
||||
if client is existing:
|
||||
retired["value"] = True
|
||||
return _orig_retire(client, reason=reason)
|
||||
|
||||
monkeypatch.setattr(agent, "_retire_shared_openai_client", _spy_retire)
|
||||
|
||||
with patch("run_agent.OpenAI", side_effect=_fake_openai):
|
||||
ok = agent._try_refresh_nous_client_credentials(force=True)
|
||||
|
||||
assert ok is True
|
||||
assert closed["value"] is True
|
||||
# #70773: the replaced shared client is RETIRED (sockets shutdown,
|
||||
# FD release deferred to GC), never hard-closed from the refreshing
|
||||
# thread — close() releasing pool FDs cross-thread was the
|
||||
# TLS-FD→SQLite corruption vector.
|
||||
assert retired["value"] is True
|
||||
assert closed["value"] is False
|
||||
assert captured["force_refresh"] is True
|
||||
assert rebuilt["kwargs"]["api_key"] == "new-nous-key"
|
||||
assert (
|
||||
|
||||
Reference in New Issue
Block a user