diff --git a/agent/agent_runtime_helpers.py b/agent/agent_runtime_helpers.py index e849b955d19..bf4d31b4f4d 100644 --- a/agent/agent_runtime_helpers.py +++ b/agent/agent_runtime_helpers.py @@ -3161,20 +3161,21 @@ def apply_pending_steer_to_tool_results(agent, messages: list, num_tool_msgs: in -def force_close_tcp_sockets(client: Any, *, release_fds: bool = False) -> int: - """Abort in-flight TCP I/O by shutting down pool sockets. +def force_close_tcp_sockets(client: Any) -> int: + """Abort in-flight TCP I/O by shutting down sockets WITHOUT closing FDs. When a provider drops a connection mid-stream — or the user issues an interrupt — we want to unblock httpx's reader/writer immediately rather than waiting for the kernel's per-connection timeout. ``shutdown(SHUT_RDWR)`` achieves that: it sends FIN, breaks any pending ``recv``/``send`` with EOF - or ``EPIPE``. + or ``EPIPE``, but does NOT release the file descriptor. - By default (``release_fds=False``) this helper does **not** call - ``socket.close()`` / release the FD. That default is load-bearing for - cross-thread abort paths (#29507): + Historically this helper also called ``socket.close()`` so the FD got + released immediately, but that's unsafe when (as is the case for both the + interrupt-abort path and stale-call kill path) the helper runs on a + different thread than the one driving the request: - * The Python ``socket.socket`` we close is the SAME object held by + * The Python ``socket.socket`` we close here is the SAME object held by httpx's pool, so closing it via Python sets its ``_fd`` to -1 and future operations on that Python object fail safely. * BUT the SSL wrapper (``ssl.SSLSocket``'s underlying OpenSSL ``BIO``) @@ -3186,20 +3187,15 @@ def force_close_tcp_sockets(client: Any, *, release_fds: bool = False) -> int: wrong file (issue #29507: 24-byte TLS application-data record clobbering SQLite header bytes 5..28). - ``shutdown()`` from any thread is FD-safe; ``close()`` is not when a - stranger thread still has the BIO holding the raw FD. + The fix is to let the owning thread own the close. ``shutdown()`` from any + thread is FD-safe; ``close()`` is not. The httpx connection's own close + path — which runs from the worker thread when it unwinds — will release + the FD via the same ``socket.socket`` object, and because Python's socket + close atomically swaps ``_fd`` to -1 *before* issuing ``os.close``, there + is no FD-aliasing window when only one thread closes. - When the **owning** thread is disposing of a client that is no longer - shared (``_close_openai_client`` after replace / request-complete), pass - ``release_fds=True``. httpx's own ``client.close()`` does not reliably - ``os.close()`` sockets that were already ``shutdown()``'d, so without an - explicit ``sock.close()`` those FDs stay in kernel CLOSED state forever - and accumulate under long-lived gateways (issue #61979 — ~1 CLOSED fd - per ~6 minutes through a local proxy path). - - Returns the number of sockets shut down (and optionally closed). Field - kept as ``tcp_force_closed=N`` in log lines for backwards-compatible - parsing. + Returns the number of sockets shut down. (Field kept as + ``tcp_force_closed=N`` in the log line for backwards-compatible parsing.) """ import socket as _socket @@ -3211,13 +3207,7 @@ def force_close_tcp_sockets(client: Any, *, release_fds: bool = False) -> int: except OSError: # Already shut down / not connected / FD invalid — all benign. pass - # IMPORTANT (#29507): never release FDs from stranger-thread - # abort paths. Only the owning-thread close path may opt in. - if release_fds: - try: - sock.close() - except OSError: - pass + # IMPORTANT (#29507): do NOT call sock.close() here. See docstring. shutdown_count += 1 except Exception as exc: _ra().logger.debug("Force-close TCP sockets sweep error: %s", exc) diff --git a/run_agent.py b/run_agent.py index 51448bb9cc4..89c746235b5 100644 --- a/run_agent.py +++ b/run_agent.py @@ -3992,18 +3992,17 @@ class AIAgent: return create_openai_client(self, client_kwargs, reason=reason, shared=shared) @staticmethod - def _force_close_tcp_sockets(client: Any, *, release_fds: bool = False) -> int: + def _force_close_tcp_sockets(client: Any) -> int: """Forwarder — see ``agent.agent_runtime_helpers.force_close_tcp_sockets``.""" from agent.agent_runtime_helpers import force_close_tcp_sockets - return force_close_tcp_sockets(client, release_fds=release_fds) + return force_close_tcp_sockets(client) def _close_openai_client(self, client: Any, *, reason: str, shared: bool) -> None: if client is None: return - # Owning-thread dispose: shutdown + release FDs so already-shutdown - # sockets don't linger in kernel CLOSED state (#61979). Cross-thread - # abort uses _abort_request_openai_client (release_fds=False) instead. - force_closed = self._force_close_tcp_sockets(client, release_fds=True) + # Force-close TCP sockets first to prevent CLOSE-WAIT accumulation, + # then do the graceful SDK-level close. + force_closed = self._force_close_tcp_sockets(client) try: client.close() logger.info( diff --git a/tests/run_agent/test_tls_fd_recycle_corruption.py b/tests/run_agent/test_tls_fd_recycle_corruption.py index 4c5f781b5e3..29c35612fdf 100644 --- a/tests/run_agent/test_tls_fd_recycle_corruption.py +++ b/tests/run_agent/test_tls_fd_recycle_corruption.py @@ -7,11 +7,9 @@ SQLite header. The fix has two prongs: -1. ``force_close_tcp_sockets`` defaults to ``shutdown(SHUT_RDWR)`` only — - no ``sock.close()``. Shutdown unblocks the worker's pending - ``recv``/``send`` without releasing the FD. Owning-thread dispose - opts into ``release_fds=True`` so already-shutdown sockets do not - accumulate as kernel CLOSED fds (#61979). +1. ``force_close_tcp_sockets`` no longer calls ``sock.close()`` — only + ``shutdown(SHUT_RDWR)``. Shutdown unblocks the worker's pending + ``recv``/``send`` without releasing the FD. 2. ``_close_request_client_once`` is thread-aware: a stranger thread (the interrupt-check / stale-call loop) only aborts the sockets and leaves @@ -63,10 +61,11 @@ def _build_fake_client(sock): def test_force_close_tcp_sockets_shutdown_only_no_close(): - """Default path: shutdown is called, close is NOT. + """The smoking-gun guarantee: shutdown is called, close is NOT. - Stranger-thread abort (#29507) relies on this default. Releasing the - FD from a non-owning thread re-opens the TLS→SQLite corruption race. + If a future refactor reintroduces ``sock.close()`` here, the + FD-recycling race that corrupted ``kanban.db`` (issue #29507) will + re-open. Pin the contract explicitly. """ from agent.agent_runtime_helpers import force_close_tcp_sockets @@ -78,30 +77,11 @@ def test_force_close_tcp_sockets_shutdown_only_no_close(): assert n == 1 assert sock.shutdown_calls == 1, "shutdown() must run — it's how we unblock the worker" assert sock.close_calls == 0, ( - "close() must NOT run by default — releasing the FD here is the " + "close() must NOT run from this helper — releasing the FD here is the " "race that wrote TLS bytes into kanban.db (#29507)" ) -def test_force_close_tcp_sockets_release_fds_closes_after_shutdown(): - """Owning-thread dispose path (#61979): shutdown then close. - - httpx does not reliably ``os.close()`` sockets that were already - shutdown, so owner-thread close must release FDs explicitly or they - accumulate in kernel CLOSED state under long-lived gateways. - """ - from agent.agent_runtime_helpers import force_close_tcp_sockets - - sock = _FakeSocket() - client = _build_fake_client(sock) - - n = force_close_tcp_sockets(client, release_fds=True) - - assert n == 1 - assert sock.shutdown_calls == 1 - assert sock.close_calls == 1 - - def test_force_close_tcp_sockets_uses_shut_rdwr(): """Both directions must be shut down so the SSL state machine fully unwinds. @@ -135,7 +115,7 @@ def test_force_close_tcp_sockets_swallows_oserror_on_shutdown(): def shutdown(self, _how): raise OSError("not connected") - def close(self): # pragma: no cover — must not run on default path + def close(self): # pragma: no cover — must not run raise AssertionError("close() must not be called") client = _build_fake_client(_AlreadyShut()) @@ -144,27 +124,6 @@ def test_force_close_tcp_sockets_swallows_oserror_on_shutdown(): assert force_close_tcp_sockets(client) == 1 -def test_force_close_tcp_sockets_release_fds_swallows_oserror_on_close(): - """Already-closed sockets must not abort the owning-thread sweep.""" - from agent.agent_runtime_helpers import force_close_tcp_sockets - - class _CloseRaises: - def __init__(self): - self.shutdown_calls = 0 - - def shutdown(self, _how): - self.shutdown_calls += 1 - - def close(self): - raise OSError("Bad file descriptor") - - sock = _CloseRaises() - client = _build_fake_client(sock) - - assert force_close_tcp_sockets(client, release_fds=True) == 1 - assert sock.shutdown_calls == 1 - - def test_force_close_tcp_sockets_handles_multiple_pool_entries(): """Walk every pool connection — the bug equally applies to all of them.""" from agent.agent_runtime_helpers import force_close_tcp_sockets @@ -184,20 +143,6 @@ def test_force_close_tcp_sockets_handles_multiple_pool_entries(): assert s.shutdown_calls == 1 assert s.close_calls == 0 - # Owning-thread dispose releases every pool entry's FD. - socks2 = [_FakeSocket(), _FakeSocket(), _FakeSocket()] - entries2 = [ - SimpleNamespace(_connection=SimpleNamespace(_network_stream=SimpleNamespace(_sock=s))) - for s in socks2 - ] - client2 = SimpleNamespace( - _client=SimpleNamespace(_transport=SimpleNamespace(_pool=SimpleNamespace(_connections=entries2))) - ) - assert force_close_tcp_sockets(client2, release_fds=True) == 3 - for s in socks2: - assert s.shutdown_calls == 1 - assert s.close_calls == 1 - # --------------------------------------------------------------------------- # Prong 2: _close_request_client_once is thread-aware. @@ -447,38 +392,6 @@ def test_agent_abort_request_openai_client_does_not_call_client_close(caplog): ), f"missing abort log line; got: {msgs!r}" -def test_agent_close_openai_client_releases_fds(caplog): - """Owning-thread ``_close_openai_client`` must release pool FDs (#61979). - - After #29507 the helper defaulted to shutdown-only; combined with - httpx not closing already-shutdown sockets, gateway processes leaked - CLOSED fds until rlimit. The owner dispose path must opt into - ``release_fds=True``. - """ - from run_agent import AIAgent - - sock = _FakeSocket() - client = _build_fake_client(sock) - client.close = MagicMock() - - agent = AIAgent.__new__(AIAgent) - agent._client_log_context = lambda: "provider=test" - - with caplog.at_level(logging.INFO, logger="run_agent"): - agent._close_openai_client(client, reason="replace:test", shared=True) - - assert sock.shutdown_calls == 1 - assert sock.close_calls == 1 - client.close.assert_called_once() - - msgs = [r.getMessage() for r in caplog.records] - assert any( - "OpenAI client closed (replace:test" in m - and "tcp_force_closed=1" in m - for m in msgs - ), f"missing close log line; got: {msgs!r}" - - def test_agent_abort_request_openai_client_null_client_is_noop(): """A ``None`` client must short-circuit cleanly (defensive).""" from run_agent import AIAgent