fix(sync): track only trailing assistant streaming (#1151)

* fix(sync): track only trailing assistant streaming

* test(sync): cover streaming completion paths

---------

Co-authored-by: Isaac Sanchez <isanchez-hawkins@arize.com>

packages/ui/src/sync/streaming.test.ts

@@ -0,0 +1,83 @@
+import { beforeEach, describe, expect, test } from "bun:test"
+import type { Message, SessionStatus } from "@opencode-ai/sdk/v2/client"
+import { INITIAL_STATE, type State } from "./types"
+import { updateStreamingState, useStreamingStore } from "./streaming"
+
+const message = (id: string, role: "user" | "assistant"): Message => ({
+  id,
+  role,
+} as unknown as Message)
+
+const stateWithMessages = (messages: Message[], status: SessionStatus = { type: "busy" } as SessionStatus): State => ({
+  ...INITIAL_STATE,
+  session_status: {
+    ses_1: status,
+  },
+  message: {
+    ses_1: messages,
+  },
+})
+
+describe("updateStreamingState", () => {
+  beforeEach(() => {
+    useStreamingStore.setState({
+      streamingMessageIds: new Map(),
+      messageStreamStates: new Map(),
+    })
+  })
+
+  test("does not mark a previous assistant message as streaming during a new user turn", () => {
+    updateStreamingState(stateWithMessages([
+      message("msg_user_1", "user"),
+      message("msg_assistant_1", "assistant"),
+    ]))
+    expect(useStreamingStore.getState().streamingMessageIds.get("ses_1")).toBe("msg_assistant_1")
+
+    updateStreamingState(stateWithMessages([
+      message("msg_user_1", "user"),
+      message("msg_assistant_1", "assistant"),
+      message("msg_user_2", "user"),
+    ]))
+
+    expect(useStreamingStore.getState().streamingMessageIds.get("ses_1")).toBeNull()
+    expect(useStreamingStore.getState().messageStreamStates.get("msg_assistant_1")?.phase).toBe("completed")
+  })
+
+  test("tracks the trailing assistant message once it appears", () => {
+    updateStreamingState(stateWithMessages([
+      message("msg_user_1", "user"),
+      message("msg_assistant_1", "assistant"),
+    ]))
+    updateStreamingState(stateWithMessages([
+      message("msg_user_1", "user"),
+      message("msg_assistant_1", "assistant"),
+      message("msg_user_2", "user"),
+    ]))
+    expect(useStreamingStore.getState().streamingMessageIds.get("ses_1")).toBeNull()
+
+    updateStreamingState(stateWithMessages([
+      message("msg_user_1", "user"),
+      message("msg_assistant_1", "assistant"),
+      message("msg_user_2", "user"),
+      message("msg_assistant_2", "assistant"),
+    ]))
+
+    expect(useStreamingStore.getState().streamingMessageIds.get("ses_1")).toBe("msg_assistant_2")
+  })
+
+  test("completes the streaming message when the session becomes idle", () => {
+    updateStreamingState(stateWithMessages([
+      message("msg_user_1", "user"),
+      message("msg_assistant_1", "assistant"),
+    ]))
+    expect(useStreamingStore.getState().streamingMessageIds.get("ses_1")).toBe("msg_assistant_1")
+
+    updateStreamingState(stateWithMessages([
+      message("msg_user_1", "user"),
+      message("msg_assistant_1", "assistant"),
+    ], { type: "idle" } as SessionStatus))
+
+    expect(useStreamingStore.getState().streamingMessageIds.get("ses_1")).toBeNull()
+    expect(useStreamingStore.getState().messageStreamStates.get("msg_assistant_1")?.phase).toBe("completed")
+  })
+})

packages/ui/src/sync/streaming.ts

@@ -57,20 +57,43 @@ export function updateStreamingState(state: State) {
     }
   }
 
+  const completeStreamingMessage = (sessionID: string, msgId: string) => {
+    nextStreamingIds.set(sessionID, null)
+    const existing = nextStreamStates.get(msgId)
+    if (existing && existing.phase === "streaming") {
+      nextStreamStates.set(msgId, {
+        ...existing,
+        phase: "completed",
+        completedAt: now,
+      })
+    }
+    changed = true
+  }
+
   for (const sessionID of busySessionIds) {
     const messages = state.message[sessionID]
     if (!messages || messages.length === 0) continue
 
-    // Find the last assistant message — that's the one streaming
+    // Only the trailing assistant turn can be streaming. If a new user turn is
+    // last, the next assistant message has not arrived yet.
     let streamingMsg: Message | null = null
     for (let i = messages.length - 1; i >= 0; i--) {
+      if (messages[i].role === "user") {
+        break
+      }
       if (messages[i].role === "assistant") {
         streamingMsg = messages[i]
         break
       }
     }
 
-    if (!streamingMsg) continue
+    if (!streamingMsg) {
+      const prevId = currentStreamingIds.get(sessionID)
+      if (prevId) {
+        completeStreamingMessage(sessionID, prevId)
+      }
+      continue
+    }
 
     const prevId = currentStreamingIds.get(sessionID)
     if (prevId !== streamingMsg.id) changed = true
@@ -100,16 +123,7 @@ export function updateStreamingState(state: State) {
     const isStillBusy = busySessionIds.has(sessionID)
     if (isStillBusy) continue
 
-    nextStreamingIds.set(sessionID, null)
-    const existing = nextStreamStates.get(msgId)
-    if (existing && existing.phase === "streaming") {
-      nextStreamStates.set(msgId, {
-        ...existing,
-        phase: "completed",
-        completedAt: now,
-      })
-      changed = true
-    }
+    completeStreamingMessage(sessionID, msgId)
   }
 
   if (changed) {