diff --git a/clients/terminal/tui/renderers/markdown_content.py b/clients/terminal/tui/renderers/markdown_content.py index f754ab0..17f85e4 100644 --- a/clients/terminal/tui/renderers/markdown_content.py +++ b/clients/terminal/tui/renderers/markdown_content.py @@ -47,16 +47,20 @@ theme=RichTheme(theme.rich_theme_styles()), force_terminal=True, ) - segments = list(themed_console.render(self._markdown)) link_color = Style.parse(theme.link.hex).color - for segment in segments: + # Stream segments straight through instead of materializing them into a + # list first — every paint (scroll, resize, reading-mode toggle) would + # otherwise allocate the full segment list + walk it twice (once to + # build, once to rewrite links). Inline the link rewrite in the loop. + for segment in themed_console.render(self._markdown): style = segment.style if style and style.link and link_color is not None: - segment = Segment( + yield Segment( segment.text, Style(color=link_color, underline=True, link=style.link), ) - yield segment + else: + yield segment class MarkdownRenderer(ContentRenderer): diff --git a/clients/terminal/tui/widgets/chat_panel.py b/clients/terminal/tui/widgets/chat_panel.py index 9d6a55d..1a02b7c 100644 --- a/clients/terminal/tui/widgets/chat_panel.py +++ b/clients/terminal/tui/widgets/chat_panel.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +import time from typing import Any from rich.console import RenderableType @@ -22,6 +23,14 @@ from clients.terminal.tui.settings import get_tui_settings from clients.terminal.tui.themes import get_active_theme +# Minimum gap between expensive rich-renderable rebuilds of a streaming +# assistant bubble. Assistant answers render as rich ``Markdown``; re-parsing +# the growing text on every ``stream_delta`` is O(N²) over the answer length. +# Coalesce rebuilds to at most ~one per window (the ChatItem content still +# updates per delta; only the Markdown parse + rich render is throttled). Plain +# renderables (thinking/planning are ``Text``) bypass this — they are cheap. +_STREAMING_RENDER_THROTTLE = 0.15 + def _item_msg(item: ChatItem) -> dict[str, Any]: """Map a ChatItem to the msg dict its renderer expects.""" @@ -122,23 +131,103 @@ # width: a repeated paint at the same width/ content reuses the Content, # and a resize or token update (new renderable) invalidates it. self._content_cache: tuple[tuple[int, int], Content] | None = None + # (id(renderable), width) -> height cache for get_content_height, so the + # layout pass does not re-render the rich renderable to measure it when + # neither the renderable nor the width changed (e.g. between throttled + # rebuilds, or repeated paints). Invalidated alongside _content_cache. + self._height_cache: tuple[tuple[int, int], int] | None = None + # Throttle state for streaming assistant rebuilds (see + # _STREAMING_RENDER_THROTTLE). ``_last_render_at`` starts at 0 so the + # first mutating delta after mount rebuilds immediately (snappy start), + # then rapid deltas coalesce into ~one rebuild per window. + self._render_timer = None + self._render_pending: bool = False + self._last_render_at: float = 0.0 + # Captured at maybe_update time for the deferred rebuild (the registry + # and reading-mode flag are owned by ChatPanel). + self._registry_ref = None + self._reading: bool = False def maybe_update(self, item: ChatItem, registry, reading: bool = False) -> bool: - """Re-render only if the item's signature changed. Returns True if updated.""" + """Re-render only if the item's signature changed. Returns True if updated. + + Assistant answers are rich Markdown — rebuilding the renderable on every + stream delta is O(N²) over the answer. Throttle those rebuilds to ~one + per ``_STREAMING_RENDER_THROTTLE`` window (coalesced); the ChatItem + content still updates per delta, only the Markdown parse/render is + deferred. Plain renderables (thinking/planning are Text) are cheap and + rebuild immediately. + """ msg = _item_msg(item) sig = _signature(msg) if sig == self._signature: return False self._item = item self._signature = sig + self._registry_ref = registry + self._reading = reading + if item.kind == "assistant_message": + self._schedule_rebuild() + else: + self._do_rebuild() + return True + + def _schedule_rebuild(self) -> None: + """Rebuild the rich renderable now, or coalesce into the pending one. + + A rapid burst of stream deltas should produce at most one rebuild per + throttle window, not one per token. If a rebuild is already scheduled, + it will pick up the freshest ``_item`` when it fires — so just return. + """ + if self._render_pending: + return + elapsed = time.monotonic() - self._last_render_at + delay = _STREAMING_RENDER_THROTTLE - elapsed + if delay <= 0: + self._do_rebuild() + else: + self._render_pending = True + self._render_timer = self.set_timer(delay, self._do_rebuild) + + def _do_rebuild(self) -> None: + """Build a fresh rich renderable for the current item and refresh.""" + self._render_pending = False + self._render_timer = None + self._last_render_at = time.monotonic() + registry = self._registry_ref + if registry is None: + return + msg = _item_msg(self._item) + # Always update the renderable (and invalidate the caches keyed on it): + # if the widget is mid-mount, the fresh renderable is what gets painted + # when mount completes. ``refresh`` only means something once mounted. self._rich_renderable = ( - registry.render_plain(msg) if reading else registry.render(msg) + registry.render_plain(msg) if self._reading else registry.render(msg) ) self._content_cache = None - # render() rebuilds the Content from _rich_renderable; layout=True so the - # widget re-measures (a longer/shorter message changes height: auto). + self._height_cache = None + if not self.is_mounted: + return self.refresh(layout=True) - return True + # A deferred rebuild grows the widget's height AFTER ChatPanel._sync + # already ran scroll_end against the old height. If the panel was + # sticking to the bottom, re-anchor so the freshly rendered streaming + # text stays in view instead of scrolling past the bottom. + parent = self.parent + if parent is not None and getattr(parent, "_stick_to_bottom", False): + self.call_after_refresh(parent.scroll_end, animate=False) + + def flush(self) -> None: + """Force a pending throttled rebuild to happen now. + + Called by ChatPanel on ``stream_end`` so the final chunk of streamed + text appears immediately rather than after the throttle window. + """ + if not self._render_pending: + return + if self._render_timer is not None: + self._render_timer.stop() + self._do_rebuild() # ── rendering + selection ──────────────────────────────────────────────── @@ -212,13 +301,23 @@ renderable = self._rich_renderable if not width or renderable is None: return 0 + # Height is a pure function of (renderable, width); cache it so the + # layout pass does not re-render the rich renderable to measure it when + # neither changed (between throttled rebuilds, or repeated paints). The + # cache key is content identity + width, mirroring _content_cache. + key = (id(renderable), width) + cache = self._height_cache + if cache is not None and cache[0] == key: + return cache[1] try: console = self.app.console options = console.options.update(width=width, highlight=False) segments = list(console.render(renderable, options)) - return sum(1 for _ in Segment.split_lines(segments)) + height = sum(1 for _ in Segment.split_lines(segments)) except Exception: return 0 + self._height_cache = (key, height) + return height @staticmethod def _denoise(lines: list[str]) -> tuple[list[str], int, bool]: @@ -348,6 +447,10 @@ # item -> view widget, keyed by id(item). Only holds visible items, so # it stays bounded by the window cap, not by total history length. self._widgets: dict[int, _ChatItemView] = {} + # Captured in _sync (before any DOM change) so a deferred throttled + # rebuild (see _ChatItemView._do_rebuild) knows whether to re-anchor the + # scroll to the bottom after the streaming bubble grows. + self._stick_to_bottom: bool = False def _render_for(self, item: ChatItem): """Renderable for an item, honouring the current reading mode. @@ -391,6 +494,15 @@ def handle_ws_event(self, msg: dict) -> None: self._model.handle_ws_event(msg) self._sync() + # stream_end closes the turn — flush any throttled assistant rebuild so + # the final chunk of streamed text appears immediately, not after the + # throttle window. (No-op for widgets without a pending rebuild.) + if msg.get("type") == "stream_end": + self._flush_pending_renders() + + def _flush_pending_renders(self) -> None: + for widget in self._widgets.values(): + widget.flush() def clear(self) -> None: """Reset the chat model and redraw an empty conversation.""" @@ -438,6 +550,10 @@ # false for the "content just grew" case too and we'd never follow. # Auto-follow resumes on its own when the user scrolls back to bottom. stick_to_bottom = self.is_vertical_scroll_end + # Record for _ChatItemView._do_rebuild: a deferred (throttled) rebuild + # grows the streaming bubble's height after this _sync already ran + # scroll_end; the rebuild re-anchors only if we were at the bottom here. + self._stick_to_bottom = stick_to_bottom items = self._model.items limit = self._max_visible_items() hidden = max(0, len(items) - limit) diff --git a/tests/clients/test_chat_panel.py b/tests/clients/test_chat_panel.py index 98594fb..2592380 100644 --- a/tests/clients/test_chat_panel.py +++ b/tests/clients/test_chat_panel.py @@ -59,9 +59,11 @@ @pytest.mark.anyio async def test_per_token_refresh_renders_only_the_streaming_item() -> None: - """A stream_delta changes only the assistant bubble's signature, so the - registry is called exactly once per token — every other visible item - widget is left untouched instead of being rebuilt/re-parsed.""" + """A stream_delta changes only the assistant bubble's signature, so only the + streaming widget is re-rendered. Rebuilds are throttled (~150 ms coalesce) + for the Markdown assistant bubble — a rapid burst produces a couple of + rebuilds, not one per token — but every rebuild is for the assistant item, + never the untouched user bubble.""" async with NaviCodeTui(new_session=True).run_test() as pilot: await pilot.pause() chat = pilot.app.query_one("ChatPanel") @@ -70,16 +72,22 @@ chat.add_user_message("hi") # renders the user bubble once chat.handle_ws_event({"type": "stream_start"}) chat.handle_ws_event({"type": "stream_delta", "delta": "a"}) # creates + renders assistant + await pilot.pause() # let the mount flush baseline = len(calls) - # Two more tokens — only the streaming assistant widget should re-render. - chat.handle_ws_event({"type": "stream_delta", "delta": "b"}) - after_one = len(calls) - chat.handle_ws_event({"type": "stream_delta", "delta": "c"}) - after_two = len(calls) + # A rapid burst of tokens — throttled, not one render per token. + for ch in "bcdefghijk": + chat.handle_ws_event({"type": "stream_delta", "delta": ch}) + await pilot.pause(0.2) # let the throttle window fire - assert after_one - baseline == 1 - assert after_two - baseline == 2 + after = len(calls) + # The first mutating delta rebuilds immediately (snappy start); the rest + # coalesce into one rebuild after the window — not ten rebuilds. + assert after > baseline + assert after - baseline <= 2 + # Every rebuild in the burst was for the streaming assistant item — the + # user bubble widget was never touched. + assert all(c.get("type") == "assistant_message" for c in calls[baseline:]) @pytest.mark.anyio @@ -779,3 +787,66 @@ w = list(chat._widgets.values())[0] plain = w._build_content().plain assert "╭" not in plain and "late message" in plain + + +@pytest.mark.anyio +async def test_stream_delta_throttle_coalesces_rebuilds() -> None: + """A rapid burst of stream deltas rebuilds the assistant bubble a couple of + times (throttle coalesces), not once per token — the O(N²) markdown-per-delta + cost is gone. The ChatItem content still accumulates every token.""" + async with NaviCodeTui(new_session=True).run_test() as pilot: + await pilot.pause() + chat = pilot.app.query_one("ChatPanel") + calls = _wrap_render(chat) + + chat.handle_ws_event({"type": "stream_start"}) + chat.handle_ws_event({"type": "stream_delta", "delta": "a"}) # mount + await pilot.pause() # mount flush + baseline = len(calls) + + # 20 rapid deltas with no pause between — coalesced by the throttle. + for _ in range(20): + chat.handle_ws_event({"type": "stream_delta", "delta": "x"}) + mid = len(calls) + await pilot.pause(0.2) # let the throttle window fire + after = len(calls) + + # At most the first mutating delta rebuilt immediately; the rest are + # coalesced into one rebuild after the window — far fewer than 20. + assert mid - baseline <= 1 + assert after - baseline <= 2 + + # The ChatItem content got every token regardless of throttle. + assistant = next(w for w in chat._widgets.values() if w._item.kind == "assistant_message") + assert assistant._item.content == "a" + "x" * 20 + + +@pytest.mark.anyio +async def test_stream_end_flushes_pending_throttled_rebuild() -> None: + """stream_end flushes a pending throttled rebuild so the final chunk of + streamed text appears immediately, not after the 150 ms window.""" + async with NaviCodeTui(new_session=True).run_test() as pilot: + await pilot.pause() + chat = pilot.app.query_one("ChatPanel") + calls = _wrap_render(chat) + + chat.handle_ws_event({"type": "stream_start"}) + chat.handle_ws_event({"type": "stream_delta", "delta": "a"}) # mount + await pilot.pause() # mount flush + baseline = len(calls) + + # "b" rebuilds immediately (first mutating delta); "c" is throttled + # (pending, window not elapsed). Then stream_end flushes — without a + # pause for the throttle timer to fire. + chat.handle_ws_event({"type": "stream_delta", "delta": "b"}) + chat.handle_ws_event({"type": "stream_delta", "delta": "c"}) + chat.handle_ws_event({"type": "stream_end"}) + + # The flush forced a rebuild of the final "abc" content despite the + # throttle window not elapsing. + assistant_renders = [ + c for c in calls[baseline:] if c.get("type") == "assistant_message" + ] + assert any(c.get("content") == "abc" for c in assistant_renders) + assistant = next(w for w in chat._widgets.values() if w._item.kind == "assistant_message") + assert assistant._item.content == "abc"