refactor: Fix markdown streaming race conditions

- Add asyncio locks to CopyableMarkdown update methods
- Serialize fragment appends to prevent stale state
- Handle fenced code blocks across async boundaries
- Update streaming logic to ignore stale snapshots
- Add tests for fenced block serialization

Author: ibeex

src/mother/runtime_coordinator.py

@@ -249,7 +249,7 @@ def _show_error(self, response: Response, error_text: str) -> None:
             self.callbacks.update_response_output, response, error_text
         )
         _ = self.callbacks.call_from_thread(response.stop_stream)
-        response.reset_state(error_text)
+        _ = self.callbacks.call_from_thread(response.reset_state, error_text)
 
     async def run_runtime_request(
         self,
@@ -342,7 +342,7 @@ async def run_runtime_request(
 
         _ = self.callbacks.call_from_thread(response.stop_stream)
         session.conversation_state.message_history = list(runtime_response.all_messages)
-        response.reset_state(stream_state.visible_text)
+        _ = self.callbacks.call_from_thread(response.reset_state, stream_state.visible_text)
         if runtime_response.tool_limit_recovery_used:
             session.record_session_event(
                 "tool_limit_recovery",
@@ -511,7 +511,7 @@ def on_progress(message: str) -> None:
                 stream_state.visible_text,
             )
         _ = self.callbacks.call_from_thread(response.stop_stream)
-        response.reset_state(stream_state.visible_text)
+        _ = self.callbacks.call_from_thread(response.reset_state, stream_state.visible_text)
         session.record_session_message("assistant", stream_state.visible_text)
         return stream_state.visible_text
 
@@ -589,7 +589,7 @@ def on_progress(update: CouncilProgressUpdate) -> None:
             )
         _ = self.callbacks.call_from_thread(self.callbacks.show_council_trace, result)
         _ = self.callbacks.call_from_thread(response.stop_stream)
-        response.reset_state(final_text)
+        _ = self.callbacks.call_from_thread(response.reset_state, final_text)
         session.conversation_state.append_synthetic_turn(user_question, final_text)
         session.record_session_message("assistant", final_text)
         session.record_session_event("council_completed", result.to_event_details())

src/mother/runtime_presentation.py

@@ -271,13 +271,7 @@ async def update_response_output(self, response: Response, text: str) -> None:
         had_waiting_animation = self.has_waiting_animation(response)
         self.clear_response_waiting_animation(response)
         should_follow = self.host.should_follow_chat_updates()
-        current_text = response.raw_markdown
-        if had_waiting_animation:
-            await response.replace_markdown(text)
-        elif text.startswith(current_text):
-            await response.append_fragment(text[len(current_text) :])
-        else:
-            await response.replace_markdown(text)
+        await response.update_streamed_markdown(text, force_replace=had_waiting_animation)
         if should_follow:
             self.host.scroll_chat_to_end(force=True)
 

src/mother/widgets.py

@@ -1,5 +1,6 @@
 """Reusable TUI widget classes for Mother."""
 
+import asyncio
 import re
 from dataclasses import dataclass
 from itertools import pairwise
@@ -835,6 +836,7 @@ def __init__(self, markdown: str = "") -> None:
         self._cursor: int = 0
         self._raw: str = markdown
         self._stream: MarkdownStream | None = None
+        self._update_lock: asyncio.Lock = asyncio.Lock()
 
     @property
     def raw_markdown(self) -> str:
@@ -850,31 +852,61 @@ def stream(self) -> MarkdownStream:
 
     async def append_fragment(self, fragment: str) -> None:
         """Append a streamed markdown fragment without reparsing the full response."""
+        async with self._update_lock:
+            await self._append_fragment_locked(fragment)
+
+    async def _append_fragment_locked(self, fragment: str) -> None:
+        """Append a fragment while the markdown update lock is held."""
         if not fragment:
             return
         previous_raw = self._raw
         self._raw += fragment
         if (
-            _ends_inside_fenced_block(previous_raw)
+            _contains_fenced_block_marker(previous_raw)
+            or _ends_inside_fenced_block(previous_raw)
             or _ends_inside_fenced_block(self._raw)
             or _contains_fenced_block_marker(fragment)
         ):
-            # Textual's incremental Markdown.append can lose fenced code contents
-            # or produce stale block layout when streamed fragments contain fences.
-            # Reparse the full response whenever a fence is active or appears.
-            await self.stop_stream()
+            # Textual's incremental Markdown.append can lose fenced code contents,
+            # append following paragraphs into the preceding fence, or produce stale
+            # block layout when streamed fragments contain fences. Once a response
+            # has seen a fence, reparse the full response for all later fragments.
+            await self._stop_stream_locked()
             await self.update(self._raw)
             return
         await self.stream.write(fragment)
 
     async def replace_markdown(self, markdown: str) -> None:
         """Replace the rendered markdown, stopping any active incremental stream first."""
-        await self.stop_stream()
+        async with self._update_lock:
+            await self._replace_markdown_locked(markdown)
+
+    async def _replace_markdown_locked(self, markdown: str) -> None:
+        """Replace markdown while the markdown update lock is held."""
+        await self._stop_stream_locked()
         self._raw = markdown
         await self.update(markdown)
 
+    async def update_streamed_markdown(self, markdown: str, *, force_replace: bool = False) -> None:
+        """Render a streamed full-text snapshot in order, ignoring stale prefix snapshots."""
+        async with self._update_lock:
+            current_text = self._raw
+            if not force_replace and current_text.startswith(markdown) and markdown != current_text:
+                return
+            if force_replace:
+                await self._replace_markdown_locked(markdown)
+            elif markdown.startswith(current_text):
+                await self._append_fragment_locked(markdown[len(current_text) :])
+            else:
+                await self._replace_markdown_locked(markdown)
+
     async def stop_stream(self) -> None:
         """Stop the active incremental markdown stream, flushing any pending fragments."""
+        async with self._update_lock:
+            await self._stop_stream_locked()
+
+    async def _stop_stream_locked(self) -> None:
+        """Stop the active incremental markdown stream while the update lock is held."""
         if self._stream is None:
             return
         stream = self._stream

tests/test_markdown_streaming.py

@@ -5,6 +5,7 @@
 from unittest.mock import AsyncMock
 
 from textual.app import App, ComposeResult
+from textual.await_complete import AwaitComplete
 from textual.widgets.markdown import MarkdownFence
 
 from mother.widgets import Response
@@ -25,6 +26,53 @@ def compose(self) -> ComposeResult:
         yield Response("")
 
 
+@final
+class _RecordingResponse(Response):
+    """Response test double that records rendered markdown with configurable delay."""
+
+    def __init__(self) -> None:
+        super().__init__("")
+        self.rendered_markdown: list[str] = []
+
+    @override
+    def update(self, markdown: str) -> AwaitComplete:
+        async def record() -> None:
+            if markdown == "```bash\n":
+                await asyncio.sleep(0.02)
+            self.rendered_markdown.append(markdown)
+
+        return AwaitComplete(record())
+
+
+def test_response_streaming_serializes_fenced_code_reparses() -> None:
+    async def run() -> None:
+        response = _RecordingResponse()
+
+        first_update = asyncio.create_task(response.append_fragment("```bash\n"))
+        await asyncio.sleep(0)
+        final_update = asyncio.create_task(response.append_fragment("echo ok\n```\n"))
+        _ = await asyncio.gather(first_update, final_update)
+
+        assert response.raw_markdown == "```bash\necho ok\n```\n"
+        assert response.rendered_markdown[-1] == response.raw_markdown
+
+    asyncio.run(run())
+
+
+def test_response_streaming_ignores_stale_prefix_snapshots() -> None:
+    async def run() -> None:
+        response = _RecordingResponse()
+        final_markdown = "```bash\necho ok\n```\n"
+
+        await response.update_streamed_markdown(final_markdown)
+        await response.update_streamed_markdown("```bash\n")
+
+        assert response.raw_markdown == final_markdown
+        assert response.rendered_markdown[-1] == final_markdown
+
+    asyncio.run(run())
+
+
 def test_response_streaming_reparses_fragments_that_contain_complete_fences() -> None:
     async def run() -> None:
         app = _MarkdownStreamingApp()
@@ -47,6 +95,25 @@ async def run() -> None:
     asyncio.run(run())
 
 
+def test_response_streaming_keeps_text_after_closed_fence_out_of_code_block() -> None:
+    async def run() -> None:
+        app = _MarkdownStreamingApp()
+
+        async with app.run_test() as pilot:
+            response = app.query_one(Response)
+
+            await response.append_fragment("```bash\n")
+            await response.append_fragment("echo ok\n")
+            await response.append_fragment("```\n")
+            await response.append_fragment("\nAfter\n")
+            await pilot.pause()
+
+            assert response.query_one(MarkdownFence).code == "echo ok"
+            assert response.raw_markdown.endswith("After\n")
+
+    asyncio.run(run())
+
+
 def test_response_streaming_preserves_fenced_code_blocks_across_fragments() -> None:
     async def run() -> None:
         app = _MarkdownStreamingApp()

tests/test_mother_send_prompt.py

@@ -67,6 +67,14 @@ async def replace_markdown(self, markdown: str) -> None:
         self.raw_markdown = markdown
         self.updated_texts.append(markdown)
 
+    async def update_streamed_markdown(self, markdown: str, *, force_replace: bool = False) -> None:
+        if force_replace:
+            await self.replace_markdown(markdown)
+        elif markdown.startswith(self.raw_markdown):
+            await self.append_fragment(markdown[len(self.raw_markdown) :])
+        else:
+            await self.replace_markdown(markdown)
+
     async def stop_stream(self) -> None:
         self.stop_stream_calls += 1