fix(sync): settle completed turns and finished messages promptly (#2667)

fix(sync): settle completed turns and finished messages promptly

packages/ui/src/sync/DOCUMENTATION.md

@@ -220,6 +220,10 @@ The event pipeline delivers each ordered per-directory flush as one reducer batc
 
 Streaming lifecycle derivation has two paths. Directory attach, switch, bootstrap, and reconnect may perform a full reconciliation. Normal store publications reconcile only sessions whose `session_status` or `message` bucket changed; part-only events update the affected streaming message heartbeat directly and must not rescan all busy sessions.
 
+A trailing assistant message that the server stamped `time.completed` is never marked as streaming: the stamp means the whole response (text plus every tool call) finished, so even while the session stays busy for the next step of the turn, the typing indicator and the streaming part-update suspension must not linger on finished content. The message-level streaming state (`streamingMessageIds` / `messageStreamStates`) is therefore a *message* lifecycle, not a turn lifecycle — it is completed by an explicit `time.completed`, by a newer trailing message, or by the session leaving `busy`.
+
+When an assistant `message.updated` event carries `time.completed` and the store still believes the session busy, sync schedules one deferred status check (`maybePollStatusAfterMessageCompletion`, ~750ms). The status is re-read when the timer fires, so a normal turn whose `session.idle` lands inside that window issues no request at all; only a still-busy session spends a directory status poll, sharing the watchdog's one-in-flight-per-directory guard. The invariant is unchanged from the watchdog escalation: the monotonic pass confirms or raises active status and never lowers it, and an authoritative resync runs only when the snapshot disagrees with a store that still believes the session busy. This narrows the stuck-spinner window after a lost `session.idle` from a watchdog interval to one round-trip; the 5s watchdog poll remains the backstop.
+
 Incomplete-session materialization is deduplicated by runtime, directory, and session for the full cooldown window, including after a fast success or failure. A settled-running-tool recovery may supersede a different request in that window so an earlier pre-settlement refresh cannot consume the only terminal recovery signal. Deferred recovery is dropped if its captured runtime is no longer active. If recovery requests a tail refresh while an older load is in flight, one refresh runs after that load instead of losing the newer authority demand. Completion retains the cooldown marker until expiry, and an older completion cannot clear a newer request marker. Recovery starts after the current ordered event batch and rechecks whether local state already contains the requested entity before starting HTTP. An explicit empty part bucket is authoritative fetched-empty state, not a missing snapshot. This prevents repeated orphan/missing-part events from creating message-tail and status request storms while preserving later recovery.
 
 When `session.idle` or `session.error` settles a session but the trailing assistant message still contains a `pending` or `running` tool, sync refreshes that session tail. This narrowly reconciles a missed terminal tool-part event without refetching normally completed turns or stale tools from older turns. A stale refresh or delayed part event cannot regress a locally observed terminal tool to an active status.

packages/ui/src/sync/__tests__/message-completion-status-poll.test.ts

@@ -0,0 +1,141 @@
+/**
+ * Tests for the deferred status poll fired when an assistant message completes
+ * (issue OPE-193): the busy spinner must not linger for up to a full watchdog
+ * poll interval after a turn completed when the session.idle event was delayed
+ * or lost — and a normal turn, whose session.idle arrives promptly, must not
+ * cost a single extra request.
+ */
+import { beforeEach, describe, expect, mock, test } from "bun:test"
+import { create, type StoreApi } from "zustand"
+import type { SessionStatus } from "@opencode-ai/sdk/v2/client"
+import { INITIAL_STATE } from "../types"
+import type { DirectoryStore } from "../child-store"
+
+type StatusSnapshot = Record<string, SessionStatus | undefined>
+
+let respondWithSnapshot: () => Promise<StatusSnapshot | null> = () => Promise.resolve({ ses_1: { type: "idle" } })
+const statusSnapshotCalls: string[] = []
+
+mock.module("@/lib/opencode/client", () => ({
+  opencodeClient: {
+    getSessionStatusForDirectory: mock((directory: string) => {
+      statusSnapshotCalls.push(directory)
+      return respondWithSnapshot()
+    }),
+  },
+}))
+
+mock.module("@/lib/runtime-switch", () => ({
+  getRuntimeKey: () => "test-runtime",
+}))
+
+import { maybePollStatusAfterMessageCompletion, MESSAGE_COMPLETION_STATUS_POLL_DELAY_MS } from "../sync-context"
+
+const createStore = (status: SessionStatus): StoreApi<DirectoryStore> => {
+  return create<DirectoryStore>()((set) => ({
+    ...INITIAL_STATE,
+    session_status: { ses_1: status },
+    patch: (partial) => set(partial),
+    replace: (next) => set(next),
+  }))
+}
+
+const sleep = (ms: number): Promise<void> => new Promise((resolve) => setTimeout(resolve, ms))
+
+/** Past the deferral, plus room for the background-network task chain. */
+const waitForPollSettled = async (): Promise<void> => {
+  await sleep(MESSAGE_COMPLETION_STATUS_POLL_DELAY_MS + 50)
+  await sleep(50)
+}
+
+describe("maybePollStatusAfterMessageCompletion (issue OPE-193)", () => {
+  beforeEach(() => {
+    respondWithSnapshot = () => Promise.resolve({ ses_1: { type: "idle" } })
+    statusSnapshotCalls.length = 0
+  })
+
+  test("does not poll when the store believes the session is already idle", async () => {
+    const store = createStore({ type: "idle" })
+
+    maybePollStatusAfterMessageCompletion("/test/project", store, "ses_1")
+    await waitForPollSettled()
+
+    expect(statusSnapshotCalls).toEqual([])
+    expect(store.getState().session_status?.ses_1?.type).toBe("idle")
+  })
+
+  test("does not poll without a directory or session id", async () => {
+    const store = createStore({ type: "busy" })
+
+    maybePollStatusAfterMessageCompletion("", store, "ses_1")
+    maybePollStatusAfterMessageCompletion("global", store, "ses_1")
+    maybePollStatusAfterMessageCompletion("/test/project", store, "")
+    await waitForPollSettled()
+
+    expect(statusSnapshotCalls).toEqual([])
+  })
+
+  test("issues no request when session.idle arrives inside the deferral window", async () => {
+    const store = createStore({ type: "busy" })
+
+    maybePollStatusAfterMessageCompletion("/test/project", store, "ses_1")
+    // The turn's own session.idle event lands well before the timer fires.
+    await sleep(50)
+    store.getState().patch({ session_status: { ses_1: { type: "idle" } } })
+    await waitForPollSettled()
+
+    expect(statusSnapshotCalls).toEqual([])
+  })
+
+  test("settles a busy session to idle when the idle event never arrives", async () => {
+    const store = createStore({ type: "busy" })
+
+    maybePollStatusAfterMessageCompletion("/test/project", store, "ses_1")
+    // Nothing settles the session inside the window; the poll must run.
+    expect(statusSnapshotCalls).toEqual([])
+
+    await waitForPollSettled()
+
+    expect(statusSnapshotCalls).toEqual(["/test/project", "/test/project"])
+    expect(store.getState().session_status?.ses_1?.type).toBe("idle")
+  })
+
+  test("keeps the session busy when the snapshot confirms it is still active", async () => {
+    const store = createStore({ type: "busy" })
+    respondWithSnapshot = () => Promise.resolve({ ses_1: { type: "busy" } })
+
+    maybePollStatusAfterMessageCompletion("/test/project", store, "ses_1")
+    await waitForPollSettled()
+
+    // Monotonic poll confirms busy; the snapshot is not idle, so no
+    // authoritative escalation runs.
+    expect(statusSnapshotCalls).toEqual(["/test/project"])
+    expect(store.getState().session_status?.ses_1?.type).toBe("busy")
+  })
+
+  test("preserves the busy status when the status fetch fails", async () => {
+    const store = createStore({ type: "busy" })
+    respondWithSnapshot = () => Promise.resolve(null)
+
+    maybePollStatusAfterMessageCompletion("/test/project", store, "ses_1")
+    await waitForPollSettled()
+
+    expect(statusSnapshotCalls).toEqual(["/test/project"])
+    // Failure is not treated as authoritative empty: the busy status stays
+    // until the watchdog poll (or a live event) corrects it.
+    expect(store.getState().session_status?.ses_1?.type).toBe("busy")
+  })
+
+  test("schedules one check for a burst of completions on the same session", async () => {
+    const store = createStore({ type: "busy" })
+
+    maybePollStatusAfterMessageCompletion("/test/project", store, "ses_1")
+    maybePollStatusAfterMessageCompletion("/test/project", store, "ses_1")
+    maybePollStatusAfterMessageCompletion("/test/project", store, "ses_1")
+    await waitForPollSettled()
+
+    // One monotonic poll plus its authoritative escalation, not three.
+    expect(statusSnapshotCalls).toEqual(["/test/project", "/test/project"])
+    expect(store.getState().session_status?.ses_1?.type).toBe("idle")
+  })
+})

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

@@ -16,8 +16,16 @@ import {
 const message = (id: string, role: "user" | "assistant"): Message => ({
   id,
   role,
+  time: { created: 1 },
 } as unknown as Message)
 
+const completedAssistantMessage = (id: string): Message => {
+  const base = message(id, "assistant")
+  // SAFETY: test fixture — the streaming reducers read only `id`, `role`, and
+  // `time.completed`, which this literal provides.
+  return { ...base, time: { created: 1, completed: 100 } } as Message
+}
+
 const stateWithMessages = (messages: Message[], status: SessionStatus = { type: "busy" } as SessionStatus): State => ({
   ...INITIAL_STATE,
   session_status: {
@@ -163,4 +171,72 @@ describe("updateStreamingState", () => {
     expect(streaming.messageStreamStates.get("msg_assistant_1")?.phase).toBe("completed")
     expect(streaming.messageStreamStates.get("msg_assistant_2")?.phase).toBe("streaming")
   })
+
+  test("completes a streaming message when the trailing assistant message finishes while the session stays busy", () => {
+    updateStreamingState(stateWithMessages([
+      message("msg_user_1", "user"),
+      message("msg_assistant_1", "assistant"),
+    ]))
+    expect(useStreamingStore.getState().streamingMessageIds.get("ses_1")).toBe("msg_assistant_1")
+
+    // The message completed (time.completed) but the turn keeps running
+    // (next step / tool phase) — the finished message must not stay marked
+    // as streaming with the typing indicator and part-update suspension on it.
+    updateStreamingState(stateWithMessages([
+      message("msg_user_1", "user"),
+      completedAssistantMessage("msg_assistant_1"),
+    ]))
+
+    const streaming = useStreamingStore.getState()
+    expect(streaming.streamingMessageIds.get("ses_1")).toBeNull()
+    expect(streaming.messageStreamStates.get("msg_assistant_1")?.phase).toBe("completed")
+  })
+
+  test("does not mark an already-completed trailing assistant message as streaming", () => {
+    updateStreamingState(stateWithMessages([
+      message("msg_user_1", "user"),
+      completedAssistantMessage("msg_assistant_1"),
+    ]))
+
+    const streaming = useStreamingStore.getState()
+    expect(streaming.streamingMessageIds.get("ses_1") ?? null).toBeNull()
+    expect(streaming.messageStreamStates.has("msg_assistant_1")).toBe(false)
+  })
+
+  test("incrementally clears the streaming marker when the trailing message completes while busy", () => {
+    const previous = stateWithMessages([
+      message("msg_user_1", "user"),
+      message("msg_assistant_1", "assistant"),
+    ])
+    updateStreamingState(previous, 10)
+    expect(useStreamingStore.getState().streamingMessageIds.get("ses_1")).toBe("msg_assistant_1")
+
+    const next = stateWithMessages([
+      message("msg_user_1", "user"),
+      completedAssistantMessage("msg_assistant_1"),
+    ])
+    updateChangedStreamingSessions(next, previous, 20)
+
+    const streaming = useStreamingStore.getState()
+    expect(streaming.streamingMessageIds.get("ses_1")).toBeNull()
+    expect(streaming.messageStreamStates.get("msg_assistant_1")?.phase).toBe("completed")
+  })
+
+  test("keeps the next assistant message streaming after an intermediate message completed while busy", () => {
+    const previous = stateWithMessages([
+      message("msg_user_1", "user"),
+      completedAssistantMessage("msg_assistant_1"),
+    ])
+    updateStreamingState(previous, 10)
+    expect(useStreamingStore.getState().streamingMessageIds.get("ses_1") ?? null).toBeNull()
+
+    const next = stateWithMessages([
+      message("msg_user_1", "user"),
+      completedAssistantMessage("msg_assistant_1"),
+      message("msg_assistant_2", "assistant"),
+    ])
+    updateChangedStreamingSessions(next, previous, 20)
+
+    expect(useStreamingStore.getState().streamingMessageIds.get("ses_1")).toBe("msg_assistant_2")
+  })
 })

packages/ui/src/sync/streaming.ts

@@ -58,6 +58,18 @@ const findTrailingAssistantMessage = (messages: Message[] | undefined): Message
   return null
 }
 
+/**
+ * The server stamps `time.completed` on an assistant message only after its
+ * whole response (text + every tool call) finished. A completed trailing
+ * message therefore means the message itself is done even when the turn keeps
+ * running (next step, follow-up tool phase) — it must not stay marked as
+ * streaming, or the typing indicator and the part-update suspension linger on
+ * finished content until the session settles.
+ */
+const isTrailingMessageComplete = (message: Message): boolean => {
+  return message.role === "assistant" && message.time.completed !== undefined
+}
+
 export function updateStreamingState(state: State, now = Date.now()) {
   countSyncPerformance("streamingFullReconciliations")
   const currentStore = useStreamingStore.getState()
@@ -108,6 +120,18 @@ export function updateStreamingState(state: State, now = Date.now()) {
       continue
     }
 
+    // The trailing assistant message already finished (time.completed), so
+    // nothing is streaming right now even though the session stays busy for
+    // the rest of the turn. Complete any previously streaming message instead
+    // of re-marking the finished one as streaming.
+    if (isTrailingMessageComplete(streamingMsg)) {
+      const prevId = currentStreamingIds.get(sessionID)
+      if (prevId) {
+        completeStreamingMessage(sessionID, prevId)
+      }
+      continue
+    }
+
     const prevId = currentStreamingIds.get(sessionID)
     if (prevId !== streamingMsg.id) changed = true
     nextStreamingIds.set(sessionID, streamingMsg.id)
@@ -222,6 +246,14 @@ export function updateChangedStreamingSessions(state: State, previous: State, no
       continue
     }
 
+    // Completed trailing message while the turn keeps running: nothing is
+    // streaming — clear the marker and any previous streaming message instead
+    // of keeping the finished message flagged as streaming.
+    if (isTrailingMessageComplete(streamingMessage)) {
+      if (previousMessageID) complete(sessionID, previousMessageID)
+      continue
+    }
+
     if (previousMessageID && previousMessageID !== streamingMessage.id) {
       complete(sessionID, previousMessageID)
     }

packages/ui/src/sync/sync-context.tsx

@@ -341,6 +341,18 @@ type PendingSessionMaterialization = {
 const SESSION_MATERIALIZATION_COOLDOWN_MS = 5_000
 const pendingSessionMaterializations = new Map<string, PendingSessionMaterialization>()
 
+// One in-flight directory status fetch at a time, shared by the active-session
+// watchdog poll and the deferred completion poll so the two cannot overlap on
+// the same directory.
+const statusPollingDirectories = new Set<string>()
+
+// Deferred completion polls awaiting their delay, keyed by directory+session so
+// a burst of completing messages schedules one check.
+const pendingMessageCompletionPolls = new Map<string, ReturnType<typeof setTimeout>>()
+
+// How long to wait for the turn's own `session.idle` before spending a request.
+export const MESSAGE_COMPLETION_STATUS_POLL_DELAY_MS = 750
+
 function enqueueSessionMaterialization(
   directory: string,
   sessionID: string,
@@ -731,6 +743,63 @@ async function resyncDirectorySessionStatuses(
   return nextStatuses
 }
 
+/**
+ * Re-check the session status shortly after an assistant message completes.
+ * The turn-ending `session.idle` event can be delayed or lost; left alone, the
+ * busy spinner keeps showing until the next watchdog poll tick (up to ~5s) and
+ * its escalation (up to ~10s).
+ *
+ * The check is deferred by `MESSAGE_COMPLETION_STATUS_POLL_DELAY_MS`, and the
+ * status is read again when the timer fires: a normal turn whose `session.idle`
+ * arrives inside that window settles on its own and issues no request at all.
+ * Only a session the store still believes busy costs one status fetch, which
+ * mirrors the watchdog escalation — the monotonic pass confirms/raises busy but
+ * never lowers it, and when the snapshot reports the session idle while the
+ * store still believes it busy, an authoritative resync settles the status.
+ *
+ * Bounded: one scheduled check per session, one in-flight status fetch per
+ * directory (shared with the watchdog poll), best-effort — the watchdog poll
+ * remains the backstop.
+ */
+export function maybePollStatusAfterMessageCompletion(
+  directory: string,
+  store: StoreApi<DirectoryStore>,
+  sessionID: string,
+): void {
+  if (!directory || directory === "global" || !sessionID) return
+  const current = store.getState().session_status?.[sessionID]
+  if (!current || current.type === "idle") return
+
+  const pendingKey = `${directory}\u0000${sessionID}`
+  if (pendingMessageCompletionPolls.has(pendingKey)) return
+
+  const timer = setTimeout(() => {
+    pendingMessageCompletionPolls.delete(pendingKey)
+    const latest = store.getState().session_status?.[sessionID]
+    if (!latest || latest.type === "idle") return
+    if (statusPollingDirectories.has(directory)) return
+
+    statusPollingDirectories.add(directory)
+    void (async () => {
+      try {
+        const statuses = await runBackgroundNetworkTask(() =>
+          resyncDirectorySessionStatuses(directory, store, [sessionID], "monotonic"))
+        if (!statuses) return
+        if (needsSnapshotAfterStatusPoll(store.getState(), sessionID, statuses[sessionID])) {
+          await runBackgroundNetworkTask(() =>
+            resyncDirectorySessionStatuses(directory, store, [sessionID], "authoritative"))
+        }
+      } catch {
+        // Best-effort — the watchdog poll retries on its own cadence.
+      } finally {
+        statusPollingDirectories.delete(directory)
+      }
+    })()
+  }, MESSAGE_COMPLETION_STATUS_POLL_DELAY_MS)
+
+  pendingMessageCompletionPolls.set(pendingKey, timer)
+}
+
 // After a monotonic poll, decide whether to escalate to a full authoritative
 // resync: the store believes the session is active but the snapshot reports it
 // idle/absent — a suspected missed idle that the monotonic poll deliberately
@@ -1854,6 +1923,12 @@ export function handleEvent(
           messageID,
         })
       }
+      // An assistant message that finished is strong evidence the turn may
+      // have ended; if the session.idle event was delayed or lost, settle the
+      // busy status immediately instead of waiting for the next watchdog poll.
+      if (info.role === "assistant" && typeof info.time?.completed === "number") {
+        maybePollStatusAfterMessageCompletion(resolvedDirectory, store, sessionID)
+      }
     }
   } else {
     const sessionID = getSessionIdFromPayload(payload) ?? undefined
@@ -2069,7 +2144,6 @@ export function SyncProvider(props: {
   const lastChildDiscoveryAtByDirectoryRef = useRef(new Map<string, number>())
   const resyncingDirectoriesRef = useRef(new Set<string>())
   const blockingRequestResyncingDirectoriesRef = useRef(new Set<string>())
-  const statusPollingDirectoriesRef = useRef(new Set<string>())
   const pipelineReconnectRef = useRef<((reason?: string) => void) | null>(null)
   const pipelineHasConnectedRef = useRef(false)
   const pipelineDisconnectedBeforeFirstConnectRef = useRef(false)
@@ -2432,7 +2506,7 @@ export function SyncProvider(props: {
       store: StoreApi<DirectoryStore>,
       candidateSessionIds: string[],
     ) => {
-      const polling = statusPollingDirectoriesRef.current
+      const polling = statusPollingDirectories
       if (polling.has(directory)) return
       polling.add(directory)
       try {
@@ -2490,7 +2564,7 @@ export function SyncProvider(props: {
         .finally(() => {
           running = false
           if (stopped) {
-            statusPollingDirectoriesRef.current.clear()
+            statusPollingDirectories.clear()
           }
         })
     }