From 5b540ddbe681fb39e5cd56edceb4db3383387619 Mon Sep 17 00:00:00 2001 From: Teknium <127238744+teknium1@users.noreply.github.com> Date: Wed, 2 Sep 2026 19:14:36 -0700 Subject: [PATCH] refactor(agent/think_scrubber,thread_scoped_output): single partial-suffix scan, forwarded stream methods --- agent/think_scrubber.py | 51 +++++++++++------------------------ agent/thread_scoped_output.py | 20 +++++++------- 2 files changed, 25 insertions(+), 46 deletions(-) diff --git a/agent/think_scrubber.py b/agent/think_scrubber.py index bb41604e2d..b5ff98e1be 100644 --- a/agent/think_scrubber.py +++ b/agent/think_scrubber.py @@ -20,10 +20,9 @@ __all__ = ["StreamingThinkScrubber"] class StreamingThinkScrubber: """Stateful scrubber for streaming reasoning/thinking blocks. - State: ``_in_block`` (inside an open block; text discarded), ``_buf`` - (held-back partial-tag tail), ``_last_emitted_ended_newline`` (True iff the - last emission ended with ``\\n`` or nothing has been emitted yet — decides - whether an open tag at buffer position 0 sits at a block boundary). + State: ``_in_block`` (inside an open block; text discarded), ``_buf`` (held-back partial-tag + tail), ``_last_emitted_ended_newline`` (True iff the last emission ended with ``\\n`` or nothing + was emitted yet — decides whether an open tag at buffer position 0 sits at a block boundary). """ _OPEN_TAG_NAMES: Tuple[str, ...] = ("think", "thinking", "reasoning", "thought", "REASONING_SCRATCHPAD") @@ -31,7 +30,8 @@ class StreamingThinkScrubber: # Lowercased literal tags so the hot path does string ops, not regex per feed(). _OPEN_TAGS: Tuple[str, ...] = tuple(f"<{name.lower()}>" for name in _OPEN_TAG_NAMES) _CLOSE_TAGS: Tuple[str, ...] = tuple(f"" for name in _OPEN_TAG_NAMES) - _MAX_TAG_LEN: int = max(len(tag) for tag in _OPEN_TAGS + _CLOSE_TAGS) + _ALL_TAGS: Tuple[str, ...] = _OPEN_TAGS + _CLOSE_TAGS + _MAX_TAG_LEN: int = max(len(tag) for tag in _ALL_TAGS) # Orphan close tag plus trailing whitespace (matches _strip_think_blocks case 3). _ORPHAN_CLOSE_RE = re.compile( "(?:" + "|".join(re.escape(t) for t in _CLOSE_TAGS) + r")[ \t\n\r]*", re.IGNORECASE @@ -54,11 +54,7 @@ class StreamingThinkScrubber: self._last_emitted_ended_newline = text.endswith("\n") def feed(self, text: str) -> str: - """Feed one delta; return the scrubbed visible portion. - - Returns "" when the whole delta is reasoning content or is held back - pending resolution of a partial tag at the boundary. - """ + """Feed one delta; return the scrubbed visible portion ("" when it is all reasoning or held back).""" if not text: return "" buf = self._buf + text @@ -94,10 +90,7 @@ class StreamingThinkScrubber: # No resolvable tag: hold back any partial-tag prefix at the tail # so a tag split across deltas isn't missed, then emit the rest. - held = max( - self._max_partial_suffix(buf, self._OPEN_TAGS), - self._max_partial_suffix(buf, self._CLOSE_TAGS), - ) + held = self._max_partial_suffix(buf, self._ALL_TAGS) self._emit(out, buf[:-held] if held else buf) self._buf = buf[-held:] if held else "" break @@ -105,13 +98,10 @@ class StreamingThinkScrubber: return "".join(out) def flush(self) -> str: - """End-of-stream flush. - - Inside an unterminated block the held-back content is discarded (leaking partial - reasoning is worse than a truncated answer); otherwise the tail is emitted verbatim. - Always resets the boundary flag: intra-turn retries flush then stream again without - ``reset()``, and a stale False flag made the new stream's opening ```` look mid-line. - """ + """End-of-stream flush: inside an unterminated block the held-back content is discarded (leaking + partial reasoning is worse than a truncated answer), otherwise the tail is emitted verbatim. + Always resets the boundary flag — intra-turn retries flush then stream again without ``reset()``, + and a stale False flag made the new stream's opening ```` look mid-line.""" tail = "" if self._in_block else self._buf self._buf = "" self._in_block = False @@ -128,10 +118,7 @@ class StreamingThinkScrubber: return min(hits) if hits else (-1, 0) def _find_earliest_closed_pair(self, buf: str): - """Return (start_idx, end_idx) of the earliest ``...`` pair, else None. - - Case-insensitive and non-greedy (closest close after the open wins); the earliest open tag wins. - """ + """(start_idx, end_idx) of the earliest ``...`` pair (non-greedy, case-insensitive), else None.""" buf_lower = buf.lower() best: "tuple[int, int] | None" = None for open_tag, close_tag in zip(self._OPEN_TAGS, self._CLOSE_TAGS): @@ -156,12 +143,9 @@ class StreamingThinkScrubber: return best_idx, best_len def _is_block_boundary(self, buf: str, idx: int, already_emitted: list[str]) -> bool: - """True iff position *idx* in *buf* is a block boundary. - - Boundary = position 0 with the prior emission ending in a newline (or nothing emitted - yet), or any position whose preceding text on the current line is whitespace-only - (when no newline precedes it in *buf*, the prior emission must also have ended with a newline). - """ + """True iff *idx* is a block boundary: position 0 after a newline-terminated (or no) prior emission, + or any position whose preceding text on the current line is whitespace-only (when no newline + precedes it in *buf*, the prior emission must also have ended with a newline).""" prior_newline = ( already_emitted[-1].endswith("\n") if already_emitted else self._last_emitted_ended_newline ) @@ -175,10 +159,7 @@ class StreamingThinkScrubber: @classmethod def _max_partial_suffix(cls, buf: str, tags: Tuple[str, ...]) -> int: - """Longest buf-suffix that is a strict prefix of any tag (case-insensitive). - - Full-length matches are real tags handled elsewhere, not held-back partials. - """ + """Longest buf-suffix that is a strict prefix of any tag (full matches are real tags, handled elsewhere).""" buf_lower = buf.lower() for i in range(min(len(buf_lower), cls._MAX_TAG_LEN - 1), 0, -1): suffix = buf_lower[-i:] diff --git a/agent/thread_scoped_output.py b/agent/thread_scoped_output.py index 56322a9f1d..1e5c061034 100644 --- a/agent/thread_scoped_output.py +++ b/agent/thread_scoped_output.py @@ -60,23 +60,21 @@ class _ThreadRoutingStream: else: self._state.silenced.pop(ident, None) - def write(self, data): # type: ignore[no-untyped-def] + def _forward(self, name: str, fallback, *args): # type: ignore[no-untyped-def] + """Call ``name`` on the current target; a dead target yields ``fallback(*args)`` instead of raising.""" try: - return self._target().write(data) + return getattr(self._target(), name)(*args) except Exception: - return len(data) if isinstance(data, str) else 0 + return fallback(*args) + + def write(self, data): # type: ignore[no-untyped-def] + return self._forward("write", lambda d: len(d) if isinstance(d, str) else 0, data) def flush(self): # type: ignore[no-untyped-def] - try: - return self._target().flush() - except Exception: - return None + return self._forward("flush", lambda: None) def writelines(self, lines): # type: ignore[no-untyped-def] - try: - return self._target().writelines(lines) - except Exception: - return None + return self._forward("writelines", lambda _l: None, lines) def isatty(self) -> bool: try: