diff --git a/libs/code/deepagents_code/tui/widgets/messages.py b/libs/code/deepagents_code/tui/widgets/messages.py index 5d8ecfc24a..21a77936bd 100644 --- a/libs/code/deepagents_code/tui/widgets/messages.py +++ b/libs/code/deepagents_code/tui/widgets/messages.py @@ -1377,11 +1377,12 @@ class AssistantMessage(Vertical): update) re-yields the stale value and wrapped fenced-code bodies vanish. A full re-parse rebuilds every fence with correct internal state. - Streamed tokens are coalesced in `_pending_append` and flushed to the - `MarkdownStream` on a throttled timer (`_STREAM_FLUSH_INTERVAL`). Writing - every token immediately forced a markdown re-parse per chunk on the UI - event loop, which starved keyboard input while the model streamed. - Batching the writes keeps the event loop free so typing stays responsive. + The first streamed fragment is written immediately so the response appears + without waiting for `_STREAM_FLUSH_INTERVAL`. Later tokens are coalesced in + `_pending_append` and flushed to the `MarkdownStream` on a throttled timer. + Writing every token immediately forced a markdown re-parse per chunk on the + UI event loop, which starved keyboard input while the model streamed; + batching subsequent writes keeps typing responsive. """ _STREAM_FLUSH_INTERVAL: ClassVar[float] = 0.1 @@ -1493,11 +1494,11 @@ def _ensure_stream(self) -> MarkdownStream: return self._stream async def append_content(self, text: str) -> None: - """Append streamed content, coalescing writes onto a throttled timer. + """Append streamed content, then coalesce later writes on a timer. - Tokens are buffered in `_pending_append` and written to the - `MarkdownStream` at most once per `_STREAM_FLUSH_INTERVAL` so the UI - event loop stays free to process keypresses while the model streams. + The first fragment is written immediately. Later fragments are buffered + and written at most once per `_STREAM_FLUSH_INTERVAL` so the UI event + loop stays free to process keypresses while the model streams. Args: text: Text to append @@ -1507,6 +1508,7 @@ async def append_content(self, text: str) -> None: self._content_parts.append(text) self._pending_append += text if self._flush_timer is None: + await self._flush_pending_append() self._flush_timer = self.set_interval( self._STREAM_FLUSH_INTERVAL, self._flush_pending_append ) diff --git a/libs/code/tests/unit_tests/tui/widgets/test_messages.py b/libs/code/tests/unit_tests/tui/widgets/test_messages.py index 9c67e9f0ca..8662905317 100644 --- a/libs/code/tests/unit_tests/tui/widgets/test_messages.py +++ b/libs/code/tests/unit_tests/tui/widgets/test_messages.py @@ -5,7 +5,7 @@ from time import time from types import SimpleNamespace from typing import ClassVar -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, call, patch import pytest from deepagents.backends.utils import ( @@ -540,8 +540,8 @@ def test_leave_before_mount_is_noop(self) -> None: class TestAssistantMessageStreamCoalescing: """Tests for the throttled streaming flush that keeps input responsive.""" - async def test_append_buffers_until_flush(self) -> None: - """Tokens accumulate in `_content` but defer the markdown write.""" + async def test_first_append_flushes_immediately(self) -> None: + """The first fragment renders immediately while later text is buffered.""" async with _AssistantMessageApp().run_test() as pilot: msg = pilot.app.query_one("#assistant", AssistantMessage) stream = MagicMock() @@ -549,16 +549,17 @@ async def test_append_buffers_until_flush(self) -> None: msg._stream = stream await msg.append_content("hello ") + stream.write.assert_awaited_once_with("hello ") + await msg.append_content("world") - # No immediate write — tokens are buffered for the timer. - stream.write.assert_not_awaited() + assert stream.write.await_count == 1 assert msg._content == "hello world" - assert msg._pending_append == "hello world" + assert msg._pending_append == "world" assert msg._flush_timer is not None - async def test_timer_flushes_coalesced_text_once(self) -> None: - """The throttled timer writes buffered tokens as a single fragment.""" + async def test_timer_flushes_later_text(self) -> None: + """The throttled timer coalesces fragments after the immediate first write.""" async with _AssistantMessageApp().run_test() as pilot: msg = pilot.app.query_one("#assistant", AssistantMessage) stream = MagicMock() @@ -570,7 +571,7 @@ async def test_timer_flushes_coalesced_text_once(self) -> None: await asyncio.sleep(msg._STREAM_FLUSH_INTERVAL * 2) await pilot.pause() - stream.write.assert_awaited_once_with("foobar") + assert stream.write.await_args_list[:2] == [call("foo"), call("bar")] assert msg._pending_append == "" async def test_stop_stream_flushes_and_cancels_timer(self) -> None: @@ -615,8 +616,9 @@ async def test_set_content_drains_and_cancels_active_timer(self) -> None: assert msg._flush_timer is None assert msg._pending_append == "" - # Buffered token must not bleed into the replacement render. - stream.write.assert_not_awaited() + # The initial fragment rendered immediately and must not bleed into + # the replacement render again. + stream.write.assert_awaited_once_with("buffered") markdown.update.assert_awaited_once_with("replacement") async def test_timer_created_once_across_appends(self) -> None: @@ -681,7 +683,8 @@ async def test_flush_restores_buffer_when_write_fails(self) -> None: # via the Textual timer's exception handler. await msg._flush_pending_append() - stream.write.assert_awaited_once_with("kept") + assert stream.write.await_count == 2 + stream.write.assert_awaited_with("kept") assert msg._pending_append == "kept" # Text arriving after the failure queues behind the retried fragment.