diff --git a/src/beaver_gateway/frontends/telegram/drafts.py b/src/beaver_gateway/frontends/telegram/drafts.py index f841c39..d61df6e 100644 --- a/src/beaver_gateway/frontends/telegram/drafts.py +++ b/src/beaver_gateway/frontends/telegram/drafts.py @@ -9,6 +9,7 @@ from __future__ import annotations import asyncio import contextlib +import html import logging import time import zlib @@ -26,6 +27,7 @@ __all__ = ["Draft"] _log = logging.getLogger("beaver_gateway.frontends.telegram.drafts") _TAIL = 3500 +_REASONING_TAIL = 400 _KEEPALIVE = 20.0 @@ -47,6 +49,7 @@ class Draft: self._interval = interval self.status = status self.text = "" + self.reasoning = "" self._paused = False self._dirty = True self._broken = False @@ -80,9 +83,20 @@ class Draft: self.text = self.text.rstrip() + "\n\n" + text.lstrip() else: self.text += text + self.reasoning = "" self._paused = False self._dirty = True + def think(self, text: str) -> None: + if text: + self.reasoning += text + self._dirty = True + + def rethink(self) -> None: + if self.reasoning: + self.reasoning = "" + self._dirty = True + async def stop(self) -> None: """Stop pushing and wait for the push in flight. @@ -146,8 +160,21 @@ class Draft: ) def _render(self) -> tuple[str, str]: + rendered, plain = [to_html(self.status)], [self.status] + thought = _tail(self.reasoning.strip(), _REASONING_TAIL) + if thought: + rendered.append(f"
{html.escape(thought)}") + plain.append(thought) tail = self.text[-_TAIL:] - if not tail.strip(): - return to_html(self.status), self.status - rendered = f"{to_html(self.status)}\n\n{to_html_tail(self.text, _TAIL)}" - return rendered, f"{self.status}\n\n{tail}" + if tail.strip(): + rendered.append(to_html_tail(self.text, _TAIL)) + plain.append(tail) + return "\n\n".join(rendered), "\n\n".join(plain) + + +def _tail(text: str, limit: int) -> str: + if len(text) <= limit: + return text + cut = text[-limit:] + _, _, rest = cut.partition(" ") + return f"… {rest or cut}" diff --git a/src/beaver_gateway/frontends/telegram/frontend.py b/src/beaver_gateway/frontends/telegram/frontend.py index 96a2fd2..a8f0c94 100644 --- a/src/beaver_gateway/frontends/telegram/frontend.py +++ b/src/beaver_gateway/frontends/telegram/frontend.py @@ -994,10 +994,13 @@ class TelegramFrontend(Frontend): draft.append(progress) else: draft.set_status(self.texts.thinking) + draft.think(progress) elif kind == "content_block_start": block = raw.get("content_block") or {} if block.get("type") == "tool_use": draft.set_status(self._tool_status(block.get("name") or "", None)) + elif block.get("type") == "thinking": + draft.rethink() def _tool_status(self, name: Any, tool_input: Any) -> str: label = status_label(str(name), tool_input, self.texts.tool_labels) diff --git a/tests/test_telegram.py b/tests/test_telegram.py index f5871a0..441c83c 100644 --- a/tests/test_telegram.py +++ b/tests/test_telegram.py @@ -1228,18 +1228,50 @@ def test_progress_updates_reach_the_draft_but_hidden_reasoning_does_not() -> Non assert draft.status == texts.writing -def test_reasoning_from_an_agent_whose_thinking_is_not_progress_stays_hidden() -> None: +def test_reasoning_rides_under_the_status_and_never_reaches_the_text() -> None: draft = _progress_draft() texts = TelegramTexts() stub = SimpleNamespace(_drafts={"k": draft}, texts=texts) - TelegramFrontend._feed_draft( - cast("Any", stub), - "k", - _thinking_delta("weighing it"), - thinking_is_progress=False, + feed = functools.partial( + TelegramFrontend._feed_draft, cast("Any", stub), "k", thinking_is_progress=False ) + + feed(_thinking_delta("weighing
weighing <it> against the note", + f"{texts.thinking}\n\nweighing