From 9a0fa659af0a581d022327cb4a6fbdad7dd070f7 Mon Sep 17 00:00:00 2001
From: Teknium <127238744+teknium1@users.noreply.github.com>
Date: Wed, 2 Sep 2026 19:42:28 -0700
Subject: [PATCH] 
 refactor(agent/thread_scoped_output,stream_diag,stream_single_writer): shared
 fence-call helper, compact installs

---
 agent/stream_diag.py          |  3 +--
 agent/stream_single_writer.py | 25 +++++++++++++------------
 agent/thread_scoped_output.py | 18 +++++-------------
 3 files changed, 19 insertions(+), 27 deletions(-)

diff --git a/agent/stream_diag.py b/agent/stream_diag.py
index f52deea9f9..c9f83c89df 100644
--- a/agent/stream_diag.py
+++ b/agent/stream_diag.py
@@ -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
diff --git a/agent/stream_single_writer.py b/agent/stream_single_writer.py
index 67fc06220a..fd6ee1dc73 100644
--- a/agent/stream_single_writer.py
+++ b/agent/stream_single_writer.py
@@ -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
diff --git a/agent/thread_scoped_output.py b/agent/thread_scoped_output.py
index 1e5c061034..02b7e273d0 100644
--- a/agent/thread_scoped_output.py
+++ b/agent/thread_scoped_output.py
@@ -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: