diff --git a/anton/chat_ui.py b/anton/chat_ui.py index 4c5af725..dacb895f 100644 --- a/anton/chat_ui.py +++ b/anton/chat_ui.py @@ -202,11 +202,7 @@ def __init__(self, console: Console, toolbar: dict | None = None) -> None: self._live: Live | None = None self._toolbar = toolbar self._activities: list[_ToolActivity] = [] - self._buffer = "" # answer text accumulated during streaming - self._in_tool_phase = False - self._last_was_tool = False - self._initial_text = "" - self._initial_printed = False + self._pending = "" # assistant text accumulated during streaming self._active = False # 3-line footer state self._line1_fun: str = "" # Line 1: Esc to cancel — fun message @@ -280,11 +276,7 @@ def start(self) -> None: self._line3_peek = "" self._set_status(self._line1_fun) self._activities = [] - self._buffer = "" - self._initial_text = "" - self._initial_printed = False - self._in_tool_phase = False - self._last_was_tool = False + self._pending = "" self._cancel_msg = "" self._active = True self._start_spinner() @@ -292,15 +284,9 @@ def start(self) -> None: def append_text(self, delta: str) -> None: if not self._active: return - if self._in_tool_phase: - self._buffer += delta - self._last_was_tool = False - self._line3_peek = self._extract_peek(self._buffer) - self._update_spinner() - else: - self._initial_text += delta - self._line3_peek = self._extract_peek(self._initial_text) - self._update_spinner() + self._pending += delta + self._line3_peek = self._extract_peek(self._pending) + self._update_spinner() def show_tool_result(self, content: str) -> None: """Print a tool result permanently (immediately scrollable).""" @@ -308,7 +294,6 @@ def show_tool_result(self, content: str) -> None: return self._stop_spinner() self._console.print(Markdown(content)) - self._last_was_tool = True self._start_spinner() def show_tool_execution(self, task: str) -> None: @@ -316,13 +301,25 @@ def show_tool_execution(self, task: str) -> None: self.on_tool_use_start(f"_compat_{id(task)}", task) def on_tool_use_start(self, tool_id: str, name: str) -> None: - """Track a new tool use.""" + """Track a new tool use. + + Any text accumulated since the last flush is the current round's + preamble (inner-speech). Print it dimmed now and clear the + accumulator, so it is visually separated from other rounds' + preambles and from the final answer. + """ import time as _time if not self._active: return - self._in_tool_phase = True - self._last_was_tool = True + preamble = self._pending.rstrip() + if preamble: + self._stop_spinner() + self._console.print(Text(preamble, style="anton.muted")) + self._start_spinner() + self._pending = "" + self._line3_peek = "" + self._update_spinner() # refresh footer even when preamble was empty activity = _ToolActivity( tool_id=tool_id, name=name, start_time=_time.monotonic() ) @@ -467,23 +464,12 @@ def finish(self) -> None: act.done_line_printed = True self._print_done_line(act, act.work_elapsed) - # Print initial text as muted "inner speech" (if not already printed) - if self._initial_text and not self._initial_printed: - if self._activities: - self._console.print( - Text(self._initial_text.rstrip(), style="anton.muted") - ) - - # Print answer - if self._activities: - if self._buffer: - self._console.print(Text("anton> ", style="anton.prompt"), end="") - self._console.print(Markdown(self._buffer)) - else: - all_text = self._initial_text + self._buffer - if all_text: - self._console.print(Text("anton> ", style="anton.prompt"), end="") - self._console.print(Markdown(all_text)) + # Whatever remains in _pending is the text after the last tool — the + # true final answer. All preambles were already flushed dimmed at + # each tool start, so there is nothing to disambiguate here. + if self._pending: + self._console.print(Text("anton> ", style="anton.prompt"), end="") + self._console.print(Markdown(self._pending)) self._active = False self._console.print() diff --git a/tests/test_chat_ui.py b/tests/test_chat_ui.py index 4ed6d141..415e2bd2 100644 --- a/tests/test_chat_ui.py +++ b/tests/test_chat_ui.py @@ -20,7 +20,7 @@ def test_start_creates_live(self, MockLive): MockLive.return_value.start.assert_called_once() @patch("anton.chat_ui.Live") - def test_append_text_updates_buffer(self, MockLive): + def test_append_text_accumulates_in_pending(self, MockLive): display, console = self._make_display() display.start() live = MockLive.return_value @@ -28,8 +28,8 @@ def test_append_text_updates_buffer(self, MockLive): display.append_text("Hello ") display.append_text("world!") - # Before any tool use, text goes to _initial_text - assert display._initial_text == "Hello world!" + # All streamed text accumulates in the single _pending buffer + assert display._pending == "Hello world!" assert live.update.call_count == 2 @patch("anton.chat_ui.Live") @@ -115,22 +115,50 @@ def test_json_delta_accumulation(self, MockLive): @patch("anton.chat_ui.Live") def test_finish_prints_activity_summary(self, MockLive): + from rich.markdown import Markdown as RichMarkdown + from rich.text import Text as RichText + display, console = self._make_display() display.start() - # Initial text before tools + # Preamble before the tool — must flush dimmed AT on_tool_use_start display.append_text("Let me check...") - display.on_tool_use_start("tool_1", "scratchpad") + + muted_before_finish = [ + c.args[0].plain + for c in console.print.call_args_list + if c.args + and isinstance(c.args[0], RichText) + and c.args[0].style == "anton.muted" + ] + assert muted_before_finish == ["Let me check..."] + display.on_tool_use_delta("tool_1", '{"action": "exec", "name": "pad"}') display.on_tool_use_end("tool_1") - # Answer text after tools + # Answer text after the tool display.append_text("Here's what I found...") - display.finish() - # finish should print: muted initial, activity tree, anton> + answer markdown, trailing newline - assert console.print.call_count >= 4 + calls_before_finish = len(console.print.call_args_list) + display.finish() + finish_calls = console.print.call_args_list[calls_before_finish:] + + # finish() prints NO muted inner-speech (it was flushed earlier) … + assert not [ + c + for c in finish_calls + if c.args + and isinstance(c.args[0], RichText) + and c.args[0].style == "anton.muted" + ] + # … and prints the final answer as a single Markdown block. + markdowns = [ + c.args[0].markup + for c in finish_calls + if c.args and isinstance(c.args[0], RichMarkdown) + ] + assert markdowns == ["Here's what I found..."] @patch("anton.chat_ui.Live") def test_no_activities_no_tree(self, MockLive): @@ -203,24 +231,130 @@ def test_scratchpad_display_falls_back_to_action(self): assert result == "exec" @patch("anton.chat_ui.Live") - def test_text_routes_to_initial_before_tools(self, MockLive): - display, _ = self._make_display() - display.start() - - display.append_text("Let me check...") - assert display._initial_text == "Let me check..." - assert display._buffer == "" - assert not display._in_tool_phase + def test_preamble_flushed_dimmed_at_tool_start(self, MockLive): + from rich.text import Text as RichText - @patch("anton.chat_ui.Live") - def test_text_routes_to_buffer_after_tools(self, MockLive): - display, _ = self._make_display() + display, console = self._make_display() display.start() display.append_text("Initial text") display.on_tool_use_start("tool_1", "scratchpad") + + # Preamble printed dimmed at the tool boundary, accumulator cleared + muted = [ + c.args[0].plain + for c in console.print.call_args_list + if c.args + and isinstance(c.args[0], RichText) + and c.args[0].style == "anton.muted" + ] + assert muted == ["Initial text"] + assert display._pending == "" + + # Subsequent text accumulates fresh display.append_text("Answer text") + assert display._pending == "Answer text" + + @patch("anton.chat_ui.Live") + def test_multiround_preambles_flushed_separately(self, MockLive): + from rich.markdown import Markdown as RichMarkdown + from rich.text import Text as RichText + + display, console = self._make_display() + display.start() + + # Round 1: preamble → tool + display.append_text("Now launching the backend:") + display.on_tool_use_start("t1", "scratchpad") + display.on_tool_use_delta("t1", '{"action": "exec", "name": "p"}') + display.on_tool_use_end("t1") + + # Round 2: preamble → tool + display.append_text("Launched! Checking the API:") + display.on_tool_use_start("t2", "scratchpad") + display.on_tool_use_delta("t2", '{"action": "exec", "name": "p"}') + display.on_tool_use_end("t2") + + # Trailing text after the last tool = the real final answer + display.append_text("Everything works.") + display.finish() + + muted = [ + c.args[0].plain + for c in console.print.call_args_list + if c.args + and isinstance(c.args[0], RichText) + and c.args[0].style == "anton.muted" + ] + # Both preambles printed live, in order, each on its own line + assert muted == [ + "Now launching the backend:", + "Launched! Checking the API:", + ] + + markdowns = [ + c.args[0].markup + for c in console.print.call_args_list + if c.args and isinstance(c.args[0], RichMarkdown) + ] + # Final answer is a single block — NOT concatenated with the preambles + assert markdowns == ["Everything works."] + + @patch("anton.chat_ui.Live") + def test_consecutive_tools_no_text_no_flush(self, MockLive): + from rich.text import Text as RichText + + display, console = self._make_display() + display.start() + + display.on_tool_use_start("t1", "scratchpad") + display.on_tool_use_end("t1") + display.on_tool_use_start("t2", "scratchpad") + display.on_tool_use_end("t2") + + muted = [ + c + for c in console.print.call_args_list + if c.args + and isinstance(c.args[0], RichText) + and getattr(c.args[0], "style", None) == "anton.muted" + ] + assert muted == [] + + @patch("anton.chat_ui.Live") + def test_turn_ending_with_tool_prints_no_answer(self, MockLive): + from rich.markdown import Markdown as RichMarkdown + + display, console = self._make_display() + display.start() + + display.append_text("Preamble") + display.on_tool_use_start("t1", "scratchpad") + display.on_tool_use_end("t1") + display.finish() + + markdowns = [ + c + for c in console.print.call_args_list + if c.args and isinstance(c.args[0], RichMarkdown) + ] + # No trailing text → no anton> answer block + assert markdowns == [] + + @patch("anton.chat_ui.Live") + def test_no_tools_single_markdown_answer(self, MockLive): + from rich.markdown import Markdown as RichMarkdown + + display, console = self._make_display() + display.start() + + display.append_text("Hello ") + display.append_text("world!") + display.finish() - assert display._initial_text == "Initial text" - assert display._buffer == "Answer text" - assert display._in_tool_phase + markdowns = [ + c.args[0].markup + for c in console.print.call_args_list + if c.args and isinstance(c.args[0], RichMarkdown) + ] + assert markdowns == ["Hello world!"]