refactor(agent/thread_scoped_output,stream_diag,stream_single_writer): shared fence-call helper, compact installs

This commit is contained in:
Teknium
2026-09-02 19:42:28 -07:00
parent 401269014a
commit 9a0fa659af
3 changed files with 19 additions and 27 deletions
+1 -2
View File
@@ -40,8 +40,7 @@ def stream_diag_capture_response(agent: Any, diag: Dict[str, Any], http_response
captured: Dict[str, str] = {}
for name in getattr(agent, "_STREAM_DIAG_HEADERS", STREAM_DIAG_HEADERS): # per-agent override (back-compat)
try:
val = headers.get(name)
if val:
if val := headers.get(name):
captured[name] = str(val)[:120] # keep log lines bounded
except Exception:
continue
+13 -12
View File
@@ -17,21 +17,22 @@ logger = logging.getLogger(__name__)
def claim_stream_writer(agent: Any) -> int:
"""Claim the delta sink for this stream attempt; ``0`` (never fenced) when the agent lacks the fence or the claim raised."""
claim = getattr(agent, "_claim_stream_writer", None)
if callable(claim):
try:
return int(claim())
except Exception:
logger.debug("stream single-writer: claim failed; proceeding unfenced", exc_info=True)
return 0
return _fence_call(agent, "_claim_stream_writer", int, 0, "claim failed; proceeding unfenced")
def stream_writer_is_current(agent: Any, token: int) -> bool:
"""True when ``token`` is still the active writer; a falsy token or a fence-less agent cannot prove supersession, so True."""
is_current = getattr(agent, "_stream_writer_is_current", None) if token else None
if callable(is_current):
if not token:
return True
return _fence_call(agent, "_stream_writer_is_current", bool, True, "is_current check failed; treating as current", token)
def _fence_call(agent: Any, name: str, cast, fallback, failure_note: str, *args):
"""Call ``agent.<name>(*args)`` when it exists; ``fallback`` when missing or raising (logged at debug)."""
fn = getattr(agent, name, None)
if callable(fn):
try:
return bool(is_current(token))
return cast(fn(*args))
except Exception:
logger.debug("stream single-writer: is_current check failed; treating as current", exc_info=True)
return True
logger.debug("stream single-writer: %s", failure_note, exc_info=True)
return fallback
+5 -13
View File
@@ -44,9 +44,7 @@ class _ThreadRoutingStream:
self._state = state
def _target(self) -> TextIO:
if self._state.silenced.get(threading.get_ident(), 0) > 0:
return self._state.sink
return self._passthrough
return self._state.sink if self._state.silenced.get(threading.get_ident(), 0) > 0 else self._passthrough
def silence(self, ident: int) -> None:
with self._state.lock:
@@ -107,15 +105,12 @@ def _ensure_installed(attr: str, passthrough: TextIO) -> "_ThreadRoutingStream":
passthrough = current if current is not None else passthrough
sink = _sinks.get(attr)
if sink is None or sink.closed:
sink = open(os.devnull, "w", encoding="utf-8")
_sinks[attr] = sink
sink = _sinks[attr] = open(os.devnull, "w", encoding="utf-8")
state = _routing_states.get(attr)
if state is None or state.sink is not sink:
state = _RoutingState(sink)
_routing_states[attr] = state
proxy = _ThreadRoutingStream(passthrough, state)
state = _routing_states[attr] = _RoutingState(sink)
proxy = _installed[attr] = _ThreadRoutingStream(passthrough, state)
setattr(sys, attr, proxy)
_installed[attr] = proxy
return proxy
@@ -123,10 +118,7 @@ def _ensure_installed(attr: str, passthrough: TextIO) -> "_ThreadRoutingStream":
def thread_scoped_silence() -> Iterator[None]:
"""Silence ``stdout``/``stderr`` for the *current thread only*."""
ident = threading.get_ident()
proxies = (
_ensure_installed("stdout", sys.__stdout__ or sys.stdout),
_ensure_installed("stderr", sys.__stderr__ or sys.stderr),
)
proxies = (_ensure_installed("stdout", sys.__stdout__ or sys.stdout), _ensure_installed("stderr", sys.__stderr__ or sys.stderr))
for proxy in proxies:
proxy.silence(ident)
try: