feat(pkg-py): segment-based message storage for mixed content type streaming

js/src/chat/ChatMessage.tsx

@@ -25,9 +25,7 @@ export const ChatMessage = memo(function ChatMessage({
 
   const roleClass = isUser ? "shiny-chat-user-message" : "shiny-chat-message"
 
-  const segments = message.segments ?? [
-    { content: message.content, contentType: message.contentType },
-  ]
+  const segments = message.segments
 
   return (
     <div className={roleClass}>

js/src/chat/chat-entry.ts

@@ -28,9 +28,9 @@ function parseInitialMessages(container: HTMLElement): ChatMessageData[] {
       id: uuid(),
       role,
       content,
-      contentType,
       streaming: false,
       icon,
+      segments: [{ content, contentType }],
     })
   })
 

js/src/chat/state.ts

@@ -1,6 +1,6 @@
 import type {
-  ContentType,
   ChatAction,
+  ContentType,
   MessagePayload,
 } from "../transport/types"
 import { uuid } from "../utils/uuid"
@@ -14,12 +14,11 @@ export interface ChatMessageData {
   id: string
   role: "user" | "assistant"
   content: string
-  contentType: ContentType
   streaming: boolean
   /** True for the empty placeholder message shown while waiting for the assistant to respond. */
   isPlaceholder?: boolean
   icon?: string
-  segments?: ContentSegment[]
+  segments: ContentSegment[]
 }
 
 export interface ChatInputState {
@@ -54,14 +53,18 @@ export const initialState: ChatState = {
 }
 
 function messagePayloadToData(msg: MessagePayload): ChatMessageData {
+  const segments: ContentSegment[] = msg.segments.map((s) => ({
+    content: s.content,
+    contentType: s.content_type,
+  }))
+
   return {
     id: msg.id ?? uuid(),
     role: msg.role,
-    content: msg.content,
-    contentType: msg.content_type,
+    content: segments.map((s) => s.content).join(""),
     streaming: false,
     icon: msg.icon,
-    segments: [{ content: msg.content, contentType: msg.content_type }],
+    segments,
   }
 }
 
@@ -76,16 +79,16 @@ export function chatReducer(state: ChatState, action: AnyAction): ChatState {
         id: uuid(),
         role: "user",
         content: action.content,
-        contentType: "markdown",
         streaming: false,
+        segments: [{ content: action.content, contentType: "markdown" }],
       }
       const loadingMsg: ChatMessageData = {
         id: uuid(),
         role: "assistant",
         content: "",
-        contentType: "markdown",
         streaming: false,
         isPlaceholder: true,
+        segments: [{ content: "", contentType: "markdown" }],
       }
       return {
         ...state,
@@ -120,9 +123,10 @@ export function chatReducer(state: ChatState, action: AnyAction): ChatState {
       const last = state.streamingMessage
       if (!last || !last.streaming) return state
 
+      // segments is always non-empty: initialized with at least one segment on chunk_start
       const chunkType =
         action.content_type ??
-        last.segments![last.segments!.length - 1]!.contentType
+        last.segments[last.segments.length - 1]!.contentType
 
       if (action.operation === "replace") {
         const segments = [{ content: action.content, contentType: chunkType }]
@@ -131,13 +135,12 @@ export function chatReducer(state: ChatState, action: AnyAction): ChatState {
           streamingMessage: {
             ...last,
             content: action.content,
-            contentType: chunkType,
             segments,
           },
         }
       }
 
-      const segments = [...last.segments!]
+      const segments = [...last.segments]
       const current = segments[segments.length - 1]!
 
       if (chunkType !== current.contentType) {
@@ -154,7 +157,6 @@ export function chatReducer(state: ChatState, action: AnyAction): ChatState {
         streamingMessage: {
           ...last,
           content,
-          contentType: chunkType,
           segments,
         },
       }

js/src/transport/types.ts

@@ -2,12 +2,16 @@ import type { HtmlDep } from "rstudio-shiny/srcts/types/src/shiny/render"
 
 export type ContentType = "markdown" | "html" | "text"
 
+export type MessagePayloadSegment = {
+  content: string
+  content_type: ContentType
+}
+
 export type MessagePayload = {
   id?: string
   role: "user" | "assistant"
-  content: string
-  content_type: ContentType
   icon?: string
+  segments: MessagePayloadSegment[]
   html_deps?: HtmlDep[]
 }
 

js/tests/chat/ChatApp.integration.test.tsx

@@ -58,7 +58,10 @@ describe("ChatApp integration: full message flow", () => {
     await act(async () => {
       transport.fire("test-chat", {
         type: "chunk_start",
-        message: { role: "assistant", content: "", content_type: "markdown" },
+        message: {
+          role: "assistant",
+          segments: [{ content: "", content_type: "markdown" }],
+        },
       })
     })
 
@@ -94,7 +97,10 @@ describe("ChatApp integration: full message flow", () => {
     await act(async () => {
       transport.fire("test-chat", {
         type: "chunk_start",
-        message: { role: "assistant", content: "", content_type: "markdown" },
+        message: {
+          role: "assistant",
+          segments: [{ content: "", content_type: "markdown" }],
+        },
       })
     })
 
@@ -134,8 +140,7 @@ describe("ChatApp integration: full message flow", () => {
         type: "message",
         message: {
           role: "assistant",
-          content: "Complete reply",
-          content_type: "markdown",
+          segments: [{ content: "Complete reply", content_type: "markdown" }],
         },
       })
     })
@@ -162,8 +167,7 @@ describe("ChatApp integration: full message flow", () => {
         type: "message",
         message: {
           role: "assistant",
-          content: "Reply",
-          content_type: "markdown",
+          segments: [{ content: "Reply", content_type: "markdown" }],
         },
       })
     })

js/tests/chat/ChatApp.test.tsx

@@ -47,9 +47,13 @@ describe("Issue #3: suggestion handlers work after re-renders", () => {
         type: "message",
         message: {
           role: "assistant",
-          content:
-            "Hello! <span class='suggestion' data-suggestion='click me'>click me</span>",
-          content_type: "html",
+          segments: [
+            {
+              content:
+                "Hello! <span class='suggestion' data-suggestion='click me'>click me</span>",
+              content_type: "html",
+            },
+          ],
         },
       })
     })
@@ -86,9 +90,13 @@ describe("Issue #3: suggestion handlers work after re-renders", () => {
         type: "message",
         message: {
           role: "assistant",
-          content:
-            "Hello! <span class='suggestion' data-suggestion='click me'>click me</span>",
-          content_type: "html",
+          segments: [
+            {
+              content:
+                "Hello! <span class='suggestion' data-suggestion='click me'>click me</span>",
+              content_type: "html",
+            },
+          ],
         },
       })
     })
@@ -274,8 +282,12 @@ describe("External link dialog", () => {
         type: "message",
         message: {
           role: "assistant",
-          content: "Visit [Example](https://example.com) for more info.",
-          content_type: "markdown",
+          segments: [
+            {
+              content: "Visit [Example](https://example.com) for more info.",
+              content_type: "markdown",
+            },
+          ],
         },
       })
     })