mirror of
https://github.com/NousResearch/hermes-agent.git
synced 2026-11-06 05:17:15 +00:00
Brooklyn Nicholson
82732650a1
Keeps agent_runtime_helpers.py within its code-health line cap and logs the best-effort probe/cleanup failures with their traceback (BLE001).
94 lines
3.8 KiB
Python
94 lines
3.8 KiB
Python
"""Pre-turn dead-connection sweep over an agent's OpenAI client pools (primary + cached request
|
|
client). Sibling of ``agent.agent_runtime_helpers``, which owns the pool-socket walk."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import contextlib
|
|
import logging
|
|
from typing import Any
|
|
|
|
logger = logging.getLogger("run_agent") # origin module's logger name: log records / caplog filters unchanged
|
|
|
|
|
|
def _socket_is_dead(sock) -> bool:
|
|
"""Probe socket health with a non-blocking recv peek."""
|
|
import socket as _socket
|
|
try:
|
|
sock.setblocking(False)
|
|
return sock.recv(1, _socket.MSG_PEEK | _socket.MSG_DONTWAIT) == b""
|
|
except BlockingIOError:
|
|
return False # no data available: socket is healthy
|
|
except OSError:
|
|
return True
|
|
finally:
|
|
with contextlib.suppress(OSError):
|
|
sock.setblocking(True)
|
|
|
|
|
|
def _count_dead_pool_sockets(client: Any) -> int:
|
|
"""Return the number of sockets in ``client``'s pool that have reached EOF (CLOSE-WAIT)."""
|
|
from agent.agent_runtime_helpers import _iter_pool_sockets
|
|
return sum(1 for sock in _iter_pool_sockets(client) if _socket_is_dead(sock))
|
|
|
|
|
|
def cleanup_dead_connections(agent) -> bool:
|
|
"""Detect and clean up dead TCP connections owned by an agent.
|
|
|
|
Long-lived chat-completion turns use a cached per-request OpenAI client
|
|
for sequential calls, while ``agent.client`` remains the primary client
|
|
used for runtime rebuilds and provider management. Inspect both pools:
|
|
a provider FIN on the cached client otherwise stays in CLOSE-WAIT until
|
|
session teardown and can be handed back to the next request.
|
|
|
|
Dead sockets on the primary client rebuild that client. Dead sockets on
|
|
the request cache go through the cache's ownership-aware teardown hook so
|
|
an in-flight worker is aborted safely and an idle client is closed by its
|
|
owner. Returns True if dead connections were found and cleaned up.
|
|
"""
|
|
primary = getattr(agent, "client", None)
|
|
cache = getattr(agent, "_request_client_cache", None)
|
|
request_client = cache.get("client") if isinstance(cache, dict) else None
|
|
|
|
clients: list[tuple[str, Any]] = []
|
|
if primary is not None:
|
|
clients.append(("primary", primary))
|
|
if request_client is not None and request_client is not primary:
|
|
clients.append(("request", request_client))
|
|
if not clients:
|
|
return False
|
|
|
|
dead_clients: list[tuple[str, Any, int]] = []
|
|
for label, client in clients:
|
|
try:
|
|
dead_count = _count_dead_pool_sockets(client)
|
|
if dead_count:
|
|
dead_clients.append((label, client, dead_count))
|
|
except Exception:
|
|
# Best-effort probe over private httpcore internals: one client's failure must not
|
|
# skip the other's check or abort the turn.
|
|
logger.debug("Dead connection check error for %s client", label, exc_info=True)
|
|
|
|
if not dead_clients:
|
|
return False
|
|
|
|
total_dead = sum(count for _label, _client, count in dead_clients)
|
|
labels = ", ".join(label for label, _client, _count in dead_clients)
|
|
logger.warning(
|
|
"Found %d dead connection(s) in %s client pool(s) — cleaning up",
|
|
total_dead,
|
|
labels,
|
|
)
|
|
|
|
for label, _client, _count in dead_clients:
|
|
try:
|
|
if label == "primary":
|
|
agent._replace_primary_openai_client(reason="dead_connection_cleanup")
|
|
else:
|
|
# This hook clears the cache and chooses close() vs socket
|
|
# shutdown based on whether another worker owns the client.
|
|
agent._close_cached_request_openai_client(reason="dead_connection_cleanup")
|
|
except Exception:
|
|
# Cleaning one pool must not stop the other pool's cleanup.
|
|
logger.debug("Dead connection cleanup error for %s client", label, exc_info=True)
|
|
return True
|