Delay sending of chat UI messages until reactive graph is flushed (#1593)

Author: cpsievert

CHANGELOG.md

@@ -22,6 +22,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * A handful of fixes for `ui.Chat()`, including:
   * A fix for use inside Shiny modules. (#1582)
   * `.messages(format="google")` now returns the correct role. (#1622)
+  * `ui.Chat(messages)` are no longer dropped when dynamically rendered. (#1593)
   * `transform_assistant_response` can now return `None` and correctly handles change of content on the last chunk. (#1641)
 
 * An empty `ui.input_date()` value no longer crashes Shiny. (#1528)

shiny/ui/_chat.py

@@ -564,7 +564,7 @@ async def append_message_stream(self, message: Iterable[Any] | AsyncIterable[Any
         async def _stream_task():
             await self._append_message_stream(message)
 
-        _stream_task()
+        self._session.on_flushed(_stream_task, once=True)
 
         # Since the task runs in the background (outside/beyond the current context,
         # if any), we need to manually raise any exceptions that occur
@@ -642,7 +642,9 @@ async def _send_append_message(
 
         # print(msg)
 
-        await self._send_custom_message(msg_type, msg)
+        # When streaming (i.e., chunk is truthy), we can send messages immediately
+        # since we already waited for the flush in order to start the stream
+        await self._send_custom_message(msg_type, msg, on_flushed=chunk is False)
         # TODO: Joe said it's a good idea to yield here, but I'm not sure why?
         # await asyncio.sleep(0)
 
@@ -994,15 +996,23 @@ def destroy(self):
     async def _remove_loading_message(self):
         await self._send_custom_message("shiny-chat-remove-loading-message", None)
 
-    async def _send_custom_message(self, handler: str, obj: ClientMessage | None):
-        await self._session.send_custom_message(
-            "shinyChatMessage",
-            {
-                "id": self.id,
-                "handler": handler,
-                "obj": obj,
-            },
-        )
+    async def _send_custom_message(
+        self, handler: str, obj: ClientMessage | None, on_flushed: bool = True
+    ):
+        async def _do_send():
+            await self._session.send_custom_message(
+                "shinyChatMessage",
+                {
+                    "id": self.id,
+                    "handler": handler,
+                    "obj": obj,
+                },
+            )
+
+        if on_flushed:
+            self._session.on_flushed(_do_send, once=True)
+        else:
+            await _do_send()
 
 
 @add_example(ex_dir="../api-examples/chat")

tests/playwright/shiny/components/chat/dynamic_ui/app.py

@@ -0,0 +1,8 @@
+from shiny.express import render, ui
+
+chat = ui.Chat(id="chat", messages=["A starting message"])
+
+
+@render.ui
+def chat_output():
+    return chat.ui()

tests/playwright/shiny/components/chat/dynamic_ui/test_chat_dynamic_ui.py

@@ -0,0 +1,15 @@
+from playwright.sync_api import Page, expect
+from utils.deploy_utils import skip_on_webkit
+
+from shiny.playwright import controller
+from shiny.run import ShinyAppProc
+
+
+@skip_on_webkit
+def test_validate_chat_basic(page: Page, local_app: ShinyAppProc) -> None:
+    page.goto(local_app.url)
+
+    chat = controller.Chat(page, "chat")
+
+    expect(chat.loc).to_be_visible(timeout=30 * 1000)
+    chat.expect_latest_message("A starting message", timeout=30 * 1000)

tests/playwright/shiny/components/chat/stream/test_chat_stream.py

@@ -21,10 +21,10 @@ def test_validate_chat(page: Page, local_app: ShinyAppProc) -> None:
     expect(chat.loc_input_button).to_be_disabled()
 
     messages = [
-        "FIRST FIRST FIRST",
         "SECOND SECOND SECOND",
-        "THIRD THIRD THIRD",
         "FOURTH FOURTH FOURTH",
+        "FIRST FIRST FIRST",
+        "THIRD THIRD THIRD",
         "FIFTH FIFTH FIFTH",
     ]
     # Allow for any whitespace between messages