refactor(agent/relay_runtime): _Lazy singleton for executor/segments config; _fail_open decorator replaces *_unguarded pairs; dead context reset dropped

This commit is contained in:
Teknium
2026-09-02 22:22:18 -07:00
parent eb92eb3a38
commit 763cee5c20
+86 -112
View File
@@ -6,6 +6,7 @@ import atexit
import asyncio
import contextlib
import contextvars
import functools
import importlib
import inspect
import logging
@@ -37,8 +38,34 @@ _PROFILE_KEY_CACHE: dict[str, str] = {}
# lost span, never a blocked agent.
_SCOPE_OP_TIMEOUT = 10.0
_SCOPE_OP_EXECUTOR: Any = None
_SCOPE_OP_EXECUTOR_LOCK = threading.Lock()
class _Lazy:
"""Double-checked, thread-safe once-only factory; ``reset()`` re-arms it (tests)."""
def __init__(self, factory: Callable[[], Any]) -> None:
self._factory, self._value, self._lock = factory, None, threading.Lock()
def get(self) -> Any:
if self._value is None:
with self._lock:
if self._value is None:
self._value = self._factory()
return self._value
def reset(self) -> None:
self._value = None
def _new_scope_op_executor() -> Any:
# Daemon workers: a wedged call abandoned at timeout cannot block interpreter exit;
# ``Future.result(timeout=...)`` still bounds callers when every worker is wedged.
from tools.daemon_pool import DaemonThreadPoolExecutor
return DaemonThreadPoolExecutor(max_workers=8, thread_name_prefix="relay-scope-op")
_SCOPE_OP_EXECUTOR = _Lazy(_new_scope_op_executor)
_scope_op_executor = _SCOPE_OP_EXECUTOR.get
def runtime_metadata(runtime_id: str, **extra: Any) -> dict[str, Any]:
@@ -46,25 +73,11 @@ def runtime_metadata(runtime_id: str, **extra: Any) -> dict[str, Any]:
return {RUNTIME_SCHEMA_KEY: RUNTIME_SCHEMA_VERSION, RUNTIME_INSTANCE_KEY: runtime_id, **extra}
def _scope_op_executor():
"""Shared daemon executor for bounded native scope ops.
Daemon workers so a wedged call abandoned at timeout cannot block interpreter exit;
``Future.result(timeout=...)`` still bounds callers when every worker is wedged."""
global _SCOPE_OP_EXECUTOR
if _SCOPE_OP_EXECUTOR is None:
with _SCOPE_OP_EXECUTOR_LOCK:
if _SCOPE_OP_EXECUTOR is None:
from tools.daemon_pool import DaemonThreadPoolExecutor
_SCOPE_OP_EXECUTOR = DaemonThreadPoolExecutor(max_workers=8, thread_name_prefix="relay-scope-op")
return _SCOPE_OP_EXECUTOR
def _run_on_daemon_thread(
fn: Callable[[], Any], *, name: str, timeout: float | None = None, timeout_message: str = ""
) -> Any:
"""Run ``fn`` on a fresh daemon thread; re-raise its error or return its result.
With ``timeout`` a still-running worker is abandoned with ``TimeoutError`` (daemon:
cannot block interpreter exit)."""
With ``timeout`` a still-running worker is abandoned with ``TimeoutError`` (daemon: cannot block exit)."""
outcome: dict[str, Any] = {}
def _target() -> None:
@@ -143,23 +156,16 @@ class RelaySession:
closing: bool = False
handle: Any = None
context: contextvars.Context | None = None
# Session-span segmentation: rotation closes the current session scope and pushes
# segment N+1 at a turn boundary (the only LIFO-safe point).
# Segmentation: rotation closes the session scope and pushes segment N+1 at a turn
# boundary (the only LIFO-safe point). Flags are set by compaction, consumed at that boundary.
segment: int = 0 # index of the CURRENT session scope (0 = first)
segment_turns: int = 0 # turns completed within the current segment
rotate_pending: bool = False # set by compaction; consumed at next begin_turn
# Rotating compaction landed while a turn was live here; closing now would pop the
# session scope under the live turn, so end_turn consumes this instead.
close_pending: bool = False
# gateway.telemetry.session_segments, cached at first read. Both defaults OFF =>
# rotation never fires and the scope lifecycle is unchanged.
_SEGMENTS_CONFIG: dict[str, Any] | None = None
_SEGMENTS_CONFIG_LOCK = threading.Lock()
rotate_pending: bool = False # consumed at next begin_turn
close_pending: bool = False # rotating compaction hit a live turn; end_turn consumes it
def _load_segments_config() -> dict[str, Any]:
"""gateway.telemetry.session_segments; both defaults OFF => rotation never fires."""
segments: dict[str, Any] = {}
with contextlib.suppress(Exception): # config absence must not crash
from gateway.run import _load_gateway_config # late import
@@ -172,19 +178,9 @@ def _load_segments_config() -> dict[str, Any]:
return {"on_compaction": bool(segments.get("on_compaction", False)), "max_turns": max_turns}
def _segments_config() -> dict[str, Any]:
"""Resolve session-segmentation settings; inert defaults when unset."""
global _SEGMENTS_CONFIG
if _SEGMENTS_CONFIG is None:
with _SEGMENTS_CONFIG_LOCK:
if _SEGMENTS_CONFIG is None:
_SEGMENTS_CONFIG = _load_segments_config()
return _SEGMENTS_CONFIG
def _reset_segments_config_for_tests() -> None:
global _SEGMENTS_CONFIG
_SEGMENTS_CONFIG = None
_SEGMENTS_CONFIG = _Lazy(_load_segments_config) # cached at first read
_segments_config = _SEGMENTS_CONFIG.get
_reset_segments_config_for_tests = _SEGMENTS_CONFIG.reset
class RelayOperationLease:
@@ -291,10 +287,9 @@ class _ProcessRelayPluginConfiguration:
def release(self, owner: Any) -> None:
"""Release one host and clear Relay after the final host exits."""
with self._lock:
if id(owner) not in self._owners:
return
self._owners.remove(id(owner))
self.retry_pending_cleanup()
if id(owner) in self._owners:
self._owners.remove(id(owner))
self.retry_pending_cleanup()
def reset_for_tests(self) -> None:
"""Clear process-global state left by directly constructed test hosts."""
@@ -385,9 +380,8 @@ class RelayRuntime:
exit_fallback: bool = False, **push_kwargs: Any,
) -> None:
"""Push a fresh SESSION_SCOPE for ``session`` (bounded by ``_SCOPE_OP_TIMEOUT``); record handle + context.
Subagents parent under their spawning turn/session handle; ``resolve_parent`` creates the parent
session when its handle is unknown. ``exit_fallback``: at interpreter shutdown the executor refuses
new futures; push synchronously instead (no agent turn waits at exit)."""
Subagents parent under their spawning turn/session handle (``resolve_parent`` creates the parent when
unknown). ``exit_fallback``: at interpreter shutdown the executor refuses futures; push synchronously."""
parent_handle = None
if session.parent_session_id:
with self._sessions_lock:
@@ -428,21 +422,15 @@ class RelayRuntime:
if session.closing:
return None
if session.handle is None:
try:
self._open_session_scope(
session, {**(metadata or {}), **runtime_metadata(self.runtime_id)},
resolve_parent=True, data=data, exit_fallback=True,
)
except Exception:
session.context = None
raise
self._open_session_scope(
session, {**(metadata or {}), **runtime_metadata(self.runtime_id)},
resolve_parent=True, data=data, exit_fallback=True,
)
return session
def rotate_session_scope(self, session: RelaySession, *, reason: str) -> None:
"""Close the current session scope and open the next segment.
Called ONLY at a turn boundary: the stack is LIFO and rotating under a live child
would close a parent out of order. Bookkeeping advances even when a native call
fails so a degraded rotation cannot retry on every turn."""
"""Close the current session scope and open the next segment (turn boundary ONLY: LIFO).
Bookkeeping advances even when a native call fails so a degraded rotation cannot retry every turn."""
with session.lock:
if session.closing or session.handle is None:
return
@@ -542,9 +530,8 @@ class RelayRuntime:
allow_closing: bool = False, timeout: float | None = None, **kwargs: Any,
) -> Any:
"""Run a Relay operation against a session's isolated scope stack.
``timeout`` bounds the native call on the daemon executor (``TimeoutError`` on
breach); ``None`` runs synchronously. Lifecycle ops gating turn/session completion
pass ``_SCOPE_OP_TIMEOUT``: a wedged pipeline must cost one span, never the agent."""
``timeout`` bounds the native call on the daemon executor (``TimeoutError`` on breach); ``None``
runs synchronously. Lifecycle ops pass ``_SCOPE_OP_TIMEOUT``: a wedged pipeline costs one span."""
with self._operation():
return self._run_in_session_untracked(
session, callback, *args, allow_closing=allow_closing, timeout=timeout, **kwargs
@@ -639,9 +626,8 @@ class RelayRuntime:
def _pop_with_drain(
self, handle: Any, *, output: dict[str, Any], metadata: dict[str, Any], session_root: Any, drain_limit: int,
) -> BaseException | None:
"""Pop ``handle``; if that fails, drain orphans above it and retry once.
Returns the retry's error (None on success). Must run inside ONE ``run_in_session``
callback so ContextVar stack views stay consistent."""
"""Pop ``handle``; if that fails, drain orphans above it and retry once; return the retry's error.
Must run inside ONE ``run_in_session`` callback so ContextVar stack views stay consistent."""
with contextlib.suppress(Exception):
pop_relay_scope(self.relay, handle, output=output, metadata=metadata)
return None
@@ -672,10 +658,9 @@ class RelayRuntime:
self, session: RelaySession, handle: Any, *, output: dict[str, Any] | None = None, allow_closing: bool = False,
failure_label: str = "scope close failed", drain_limit: int = 32, operation_already_held: bool = False,
) -> str | None:
"""Pop ``handle``, draining orphaned children in the same session context.
Relay scopes are strict LIFO; empty-stream retries + interrupt can abandon a
physical LLM scope above TURN/SESSION. Drain+close is bounded so a wedged pipeline
never blocks turn/session completion. Returns a failure string or None."""
"""Pop ``handle``, draining orphaned children in the same session context; failure string or None.
Relay scopes are strict LIFO; empty-stream retries + interrupt can abandon a physical LLM scope
above TURN/SESSION. Drain+close is bounded so a wedged pipeline never blocks completion."""
if handle is None:
return None
run_in_session = (self._run_in_session_untracked if operation_already_held else self.run_in_session)
@@ -837,8 +822,7 @@ class ConversationLease:
def live_runtime(self) -> RelayRuntime | None:
"""Return the real Relay host when this lease owns an open session."""
host = self.host
return host if isinstance(host, RelayRuntime) and self.session is not None else None
return self.host if isinstance(self.host, RelayRuntime) and self.session is not None else None
@dataclass
@@ -862,21 +846,18 @@ _CURRENT_TURN: contextvars.ContextVar[RelayTurnContext | None] = contextvars.Con
"hermes_relay_turn", default=None
)
# Depth of managed Relay callbacks on the current call path (>0 while the native pipeline
# is mid-dispatch of a Hermes tool/LLM callback). Nested managed execution there is
# structurally broken: the pipeline binds its Futures to the OUTER call's event loop, which
# is blocked inside the synchronous callback (wrong loop / deadlock / "Event loop is
# closed"), so resolve_execution_context() bypasses Relay while set. A ContextVar so the
# marker follows copy_context() into worker threads / per-thread loops.
# >0 while the native pipeline is mid-dispatch of a Hermes tool/LLM callback. Nested managed
# execution there is structurally broken (the pipeline binds its Futures to the OUTER call's
# loop, blocked inside the synchronous callback), so resolve_execution_context() bypasses Relay.
# A ContextVar so the marker follows copy_context() into worker threads / per-thread loops.
_MANAGED_CALLBACK_DEPTH: contextvars.ContextVar[int] = contextvars.ContextVar(
"hermes_relay_managed_callback_depth", default=0
)
class managed_callback_guard:
"""Mark the current context as inside a managed Relay callback.
Wrap the ``invoke()`` callbacks handed to the native pipeline; everything they
transitively call (incl. work forwarded via copy_context()) runs unmanaged."""
"""Mark the current context as inside a managed Relay callback: everything the wrapped ``invoke()``
transitively calls (incl. work forwarded via copy_context()) runs unmanaged."""
def __enter__(self) -> "managed_callback_guard":
self._token = _MANAGED_CALLBACK_DEPTH.set(_MANAGED_CALLBACK_DEPTH.get() + 1)
@@ -895,6 +876,16 @@ def _warn_on_error(what: str, callback: Callable[..., Any], *args: Any, **kwargs
return None
def _fail_open(what: str) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
"""Decorator form of ``_warn_on_error`` for telemetry hooks that must never block the caller."""
def wrap(fn: Callable[..., Any]) -> Callable[..., Any]:
@functools.wraps(fn)
def guarded(*args: Any, **kwargs: Any) -> Any:
return _warn_on_error(what, fn, *args, **kwargs)
return guarded
return wrap
def _flag_open_session(session: RelaySession, flag: str) -> None:
"""Set a pending-rotation/close flag unless the session is already closing."""
with session.lock:
@@ -1032,15 +1023,10 @@ class RelaySessionCoordinator:
if failure:
logger.warning("Hermes Relay turn finalization failed: %s", failure)
def _consume_deferred_close(self, lease: Any) -> None:
"""Close a session whose rotating-compaction close was deferred.
``notify_session_compacted`` sets ``close_pending`` when the old session had a live
turn (closing then breaks LIFO). The last live turn consumes it here after its own
scope popped and it left the active-turn table."""
# Telemetry must never block end_turn.
_warn_on_error("deferred session close", self._consume_deferred_close_unguarded, lease)
def _consume_deferred_close_unguarded(self, lease: ConversationLease) -> None:
@_fail_open("deferred session close")
def _consume_deferred_close(self, lease: ConversationLease) -> None:
"""Close a session whose rotating-compaction close was deferred (``close_pending``).
The last live turn consumes it after its own scope popped and it left the active-turn table."""
host = lease.live_runtime()
if host is None:
return
@@ -1049,19 +1035,12 @@ class RelaySessionCoordinator:
if pending and not self.has_active_turn(profile_key=lease.profile_key, session_id=lease.session_id):
host.close_session({"session_id": lease.session_id})
@_fail_open("compaction notification")
def notify_session_compacted(self, *, profile_key: str, session_id: str, old_session_id: str = "") -> None:
"""React to a completed compaction, per compaction mode.
In-place (``old_session_id`` empty/equal): flag rotation for the next turn boundary
— never rotate immediately, a turn may be live and rotating under it breaks LIFO.
Rotating (ids differ): the next turn gets a fresh session under the new id, so close
the OLD session now or its scope stays an unexported orphan. Unknown sessions and
disabled config are silent no-ops."""
# Telemetry must never block compaction.
_warn_on_error(
"compaction notification", self._notify_session_compacted_unguarded, profile_key, session_id, old_session_id
)
def _notify_session_compacted_unguarded(self, profile_key: str, session_id: str, old_session_id: str) -> None:
"""React to a completed compaction, per compaction mode; unknown sessions / disabled config are no-ops.
In-place (``old_session_id`` empty/equal): flag rotation for the next turn boundary — never rotate
immediately, a live turn under it breaks LIFO. Rotating (ids differ): the next turn gets a fresh
session under the new id, so close the OLD session now or its scope stays an unexported orphan."""
if not _segments_config()["on_compaction"]:
return
host = self.registry.for_profile(profile_key)
@@ -1205,9 +1184,7 @@ def resolve_execution_context(session_id: str) -> tuple[RelayRuntime | None, Rel
def apply_tool_request_intercepts(*, session_id: str, tool_name: str, args: dict[str, Any]) -> dict[str, Any]:
"""Return Relay-rewritten arguments at Hermes's authorization boundary."""
if not session_id:
return args
runtime = get_runtime(create=False)
runtime = get_runtime(create=False) if session_id else None
if runtime is None:
return args
return runtime.apply_tool_request_intercepts(session_id=session_id, tool_name=tool_name, args=args)
@@ -1219,12 +1196,9 @@ def _is_relay_wrapped_callback_error(relay_error: BaseException, callback_error:
return True
if not isinstance(relay_error, RuntimeError):
return False
callback_type = callback_error.__class__
type_names = {
callback_type.__name__, callback_type.__qualname__, f"{callback_type.__module__}.{callback_type.__qualname__}",
}
message = str(relay_error)
return any(message.startswith(f"internal error: {type_name}: {callback_error}") for type_name in type_names)
kind = type(callback_error)
type_names = {kind.__name__, kind.__qualname__, f"{kind.__module__}.{kind.__qualname__}"}
return any(str(relay_error).startswith(f"internal error: {name}: {callback_error}") for name in type_names)
def get_runtime(*, create: bool = True, profile_key: str | None = None) -> RelayRuntime | None: