refactor(agent/thread_scoped_output,stream_diag,stream_single_writer): shared fence-call helper, compact installs
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user