fix(native-chat): retain reads through send handoff

src/renderer/src/components/native-chat/StructuredAgentSessionStatusBridge.test.tsx

@@ -85,6 +85,8 @@ import {
 import { resetStructuredAgentSessionReadOwnersForTests } from './structured-agent-session-read-owner'
 import { useStructuredAgentSessionRead } from './use-structured-agent-session-read'
 
+let retainPendingTurn: ((clientMessageId: string) => void) | null = null
+
 const structuredTab = {
   id: 'structured-tab-1',
   worktreeId: 'wt-1',
@@ -131,11 +133,12 @@ const historyResult = {
 }
 
 function ActiveSessionRead({ isVisible = true }: { isVisible?: boolean }): null {
-  useStructuredAgentSessionRead({
+  const read = useStructuredAgentSessionRead({
     sessionId: structuredTab.entityId,
     target: { kind: 'local' },
     isVisible
   })
+  retainPendingTurn = read.retainPendingTurn
   return null
 }
 
@@ -152,6 +155,7 @@ describe('StructuredAgentSessionStatusBridge', () => {
   beforeEach(() => {
     vi.clearAllMocks()
     resetStructuredAgentSessionReadOwnersForTests()
+    retainPendingTurn = null
     mocks.call.mockResolvedValue(historyResult)
     mocks.subscribe.mockResolvedValue({ unsubscribe: mocks.unsubscribe })
     mocks.store?.setState({
@@ -215,7 +219,7 @@ describe('StructuredAgentSessionStatusBridge', () => {
     })
   })
 
-  it('keeps observing a hidden active turn until its settlement removes the running item', async () => {
+  it('hands an immediate hidden send to the status observer before its running event', async () => {
     const runningItem = {
       itemId: 'turn-running',
       revision: 1,
@@ -227,33 +231,38 @@ describe('StructuredAgentSessionStatusBridge', () => {
         turnLifecycle: { turnId: 'turn-1', state: 'running' }
       }
     } as const
-    mocks.call.mockResolvedValue({
-      ...historyResult,
-      page: {
-        ...historyResult.page,
-        items: [userItem, runningItem],
-        window: {
-          ...historyResult.page.window,
-          newest: { epoch: 'epoch-1', sequence: 2 },
-          nextCursor: { epoch: 'epoch-1', sequence: 2 }
-        },
-        liveCursor: { epoch: 'epoch-1', sequence: 2 }
-      }
-    })
     const view = render(<ActiveComposition />)
 
-    await waitFor(() =>
-      expect(Object.values(mocks.store?.getState().agentStatusByPaneKey ?? {})).toEqual([
-        expect.objectContaining({ state: 'working' })
-      ])
-    )
-    expect(mocks.subscribe).toHaveBeenCalledOnce()
+    await waitFor(() => expect(mocks.subscribe).toHaveBeenCalledOnce())
+    act(() => {
+      expect(retainPendingTurn).not.toBeNull()
+      retainPendingTurn?.('client-1')
+    })
 
     view.rerender(<ActiveComposition isVisible={false} />)
     await act(() => Promise.resolve())
     expect(mocks.unsubscribe).not.toHaveBeenCalled()
 
     const onEvent = mocks.subscribe.mock.calls[0]?.[2] as (event: unknown) => void
+    act(() =>
+      onEvent({
+        type: 'batch',
+        sessionId: 'session-1',
+        batch: {
+          cursor: { epoch: 'epoch-1', sequence: 2 },
+          items: [runningItem],
+          removedItemIds: [],
+          submissions: []
+        }
+      })
+    )
+
+    await waitFor(() =>
+      expect(Object.values(mocks.store?.getState().agentStatusByPaneKey ?? {})).toEqual([
+        expect.objectContaining({ state: 'working' })
+      ])
+    )
+
     act(() =>
       onEvent({
         type: 'batch',
@@ -275,6 +284,58 @@ describe('StructuredAgentSessionStatusBridge', () => {
     await waitFor(() => expect(mocks.unsubscribe).toHaveBeenCalledOnce())
   })
 
+  it('releases a hidden send whose running and settled lifecycle arrive together', async () => {
+    const view = render(<ActiveComposition />)
+    await waitFor(() => expect(mocks.subscribe).toHaveBeenCalledOnce())
+    act(() => {
+      retainPendingTurn?.('client-1')
+    })
+    view.rerender(<ActiveComposition isVisible={false} />)
+    await act(() => Promise.resolve())
+
+    const onEvent = mocks.subscribe.mock.calls[0]?.[2] as (event: unknown) => void
+    act(() =>
+      onEvent({
+        type: 'batch',
+        sessionId: 'session-1',
+        batch: {
+          cursor: { epoch: 'epoch-1', sequence: 3 },
+          items: [
+            {
+              itemId: 'turn-running',
+              revision: 1,
+              sequence: 2,
+              observedAt: 2,
+              body: {
+                kind: 'status',
+                text: 'Working',
+                turnLifecycle: { turnId: 'turn-1', state: 'running' }
+              }
+            },
+            {
+              itemId: 'turn-complete',
+              revision: 1,
+              sequence: 3,
+              observedAt: 3,
+              body: {
+                kind: 'status',
+                text: 'Done',
+                turnLifecycle: { turnId: 'turn-1', state: 'completed' }
+              }
+            }
+          ],
+          removedItemIds: [],
+          submissions: []
+        }
+      })
+    )
+
+    await waitFor(() => expect(mocks.unsubscribe).toHaveBeenCalledOnce())
+    expect(Object.values(mocks.store?.getState().agentStatusByPaneKey ?? {})).toEqual([
+      expect.objectContaining({ state: 'done' })
+    ])
+  })
+
   it('keeps the status map reference stable for coalesced assistant deltas', async () => {
     render(<ActiveComposition />)
     await waitFor(() => expect(mocks.setAgentStatus).toHaveBeenCalledOnce())

src/renderer/src/components/native-chat/structured-agent-session-outbox-retry.ts

@@ -0,0 +1,39 @@
+import type { AgentJournalSubmission } from '../../../../shared/agent-session-journal-types'
+import type { StructuredAgentSessionOutboxEntry } from '../../../../shared/structured-agent-session-outbox'
+
+export function prepareStructuredAgentSessionOutboxRetry(args: {
+  clientMessageId: string
+  createOperationId: () => string
+  entries: readonly StructuredAgentSessionOutboxEntry[]
+  submissions: readonly AgentJournalSubmission[]
+}): {
+  entries: StructuredAgentSessionOutboxEntry[]
+  nextClientMessageId: string
+} {
+  const { clientMessageId, createOperationId, entries, submissions } = args
+  const submission = submissions.find((candidate) => candidate.clientMessageId === clientMessageId)
+  const current = entries.find((entry) => entry.clientMessageId === clientMessageId)
+  // A settled rejection must use a fresh operation id or the host replays it forever.
+  const nextClientMessageId =
+    current && submission?.dispatchState === 'rejected' ? createOperationId() : clientMessageId
+  const retryAfterUnknownSubmittedAt =
+    submission?.dispatchState === 'unknown'
+      ? submission.submittedAt
+      : current?.state === 'unconfirmed'
+        ? -1
+        : null
+  return {
+    entries: entries.map((entry) =>
+      entry.clientMessageId === clientMessageId
+        ? {
+            ...entry,
+            clientMessageId: nextClientMessageId,
+            state: 'queued' as const,
+            retryAfterUnknownSubmittedAt:
+              nextClientMessageId === clientMessageId ? retryAfterUnknownSubmittedAt : null
+          }
+        : entry
+    ),
+    nextClientMessageId
+  }
+}

src/renderer/src/components/native-chat/structured-agent-session-pending-turn-retentions.ts

@@ -0,0 +1,143 @@
+import type {
+  AgentJournalRenderItem,
+  AgentJournalSubmission
+} from '../../../../shared/agent-session-journal-types'
+import { activeStructuredAgentSessionTurnId } from '../../../../shared/structured-agent-session-projection'
+
+type TurnLifecycleItem = AgentJournalRenderItem & {
+  body: Extract<AgentJournalRenderItem['body'], { kind: 'status' }> & {
+    turnLifecycle: NonNullable<
+      Extract<AgentJournalRenderItem['body'], { kind: 'status' }>['turnLifecycle']
+    >
+  }
+}
+
+function turnLifecycleItems(items: readonly AgentJournalRenderItem[]): TurnLifecycleItem[] {
+  return items.filter(
+    (item): item is TurnLifecycleItem =>
+      item.body.kind === 'status' && Boolean(item.body.turnLifecycle)
+  )
+}
+
+export function createStructuredAgentSessionPendingTurnRetentions(
+  getState: () => {
+    items: readonly AgentJournalRenderItem[]
+    submissions: readonly AgentJournalSubmission[]
+  }
+): {
+  dispose: () => void
+  hasRetentions: () => boolean
+  observe: () => void
+  releaseObservedTurn: () => void
+  release: (clientMessageId: string) => void
+  retain: (clientMessageId: string, acquireActivation: () => () => void) => void
+} {
+  const retentions = new Map<
+    string,
+    {
+      afterSequence: number
+      release: () => void
+      turnId: string | null
+    }
+  >()
+  let observedActiveTurnId: string | null = null
+
+  const observe = (): void => {
+    if (retentions.size === 0) {
+      return
+    }
+    const { items, submissions } = getState()
+    for (const submission of submissions) {
+      if (submission.dispatchState === 'rejected') {
+        retentions.get(submission.clientMessageId)?.release()
+      }
+    }
+    const lifecycleItems = turnLifecycleItems(items)
+    const activeTurnId = activeStructuredAgentSessionTurnId(lifecycleItems)
+    if (observedActiveTurnId && observedActiveTurnId !== activeTurnId) {
+      for (const retention of retentions.values()) {
+        if (retention.turnId === observedActiveTurnId) {
+          retention.release()
+          break
+        }
+      }
+    }
+    if (activeTurnId && activeTurnId !== observedActiveTurnId) {
+      const runningSequence = lifecycleItems.findLast(
+        (item) => item.body.turnLifecycle.turnId === activeTurnId
+      )?.sequence
+      for (const retention of retentions.values()) {
+        if (
+          retention.turnId === null &&
+          runningSequence !== undefined &&
+          runningSequence > retention.afterSequence
+        ) {
+          retention.turnId = activeTurnId
+          break
+        }
+      }
+    }
+    if (!activeTurnId) {
+      const claimedTurnIds = new Set<string>()
+      for (const retention of retentions.values()) {
+        if (retention.turnId !== null) {
+          continue
+        }
+        const settled = lifecycleItems.find(
+          (item) =>
+            item.sequence > retention.afterSequence &&
+            item.body.turnLifecycle.state !== 'running' &&
+            !claimedTurnIds.has(item.body.turnLifecycle.turnId)
+        )
+        if (settled) {
+          claimedTurnIds.add(settled.body.turnLifecycle.turnId)
+          retention.release()
+        }
+      }
+    }
+    observedActiveTurnId = activeTurnId
+  }
+
+  return {
+    dispose: () => {
+      for (const retention of retentions.values()) {
+        retention.release()
+      }
+    },
+    hasRetentions: () => retentions.size > 0,
+    observe,
+    releaseObservedTurn: () => {
+      if (!observedActiveTurnId) {
+        return
+      }
+      for (const retention of retentions.values()) {
+        if (retention.turnId === observedActiveTurnId) {
+          retention.release()
+          return
+        }
+      }
+    },
+    release: (clientMessageId) => retentions.get(clientMessageId)?.release(),
+    retain: (clientMessageId, acquireActivation) => {
+      if (retentions.has(clientMessageId)) {
+        return
+      }
+      const lifecycleItems = turnLifecycleItems(getState().items)
+      if (retentions.size === 0) {
+        observedActiveTurnId = activeStructuredAgentSessionTurnId(lifecycleItems)
+      }
+      const releaseActivation = acquireActivation()
+      const retention = {
+        afterSequence: lifecycleItems.at(-1)?.sequence ?? -1,
+        turnId: null as string | null,
+        release: (): void => {
+          if (!retentions.delete(clientMessageId)) {
+            return
+          }
+          releaseActivation()
+        }
+      }
+      retentions.set(clientMessageId, retention)
+    }
+  }
+}

src/renderer/src/components/native-chat/structured-agent-session-read-owner.ts

@@ -15,6 +15,7 @@ import {
 import type { RuntimeClientTarget } from '@/runtime/runtime-rpc-client'
 import { callStructuredAgentSession } from '@/runtime/structured-agent-session-client'
 import { NATIVE_CHAT_INITIAL_LIMIT } from './native-chat-pagination'
+import { createStructuredAgentSessionPendingTurnRetentions } from './structured-agent-session-pending-turn-retentions'
 import { startStructuredAgentSessionReadTransport } from './structured-agent-session-read-transport'
 
 export type StructuredAgentSessionReadSnapshot = {
@@ -29,6 +30,8 @@ export type StructuredAgentSessionReadOwner = {
   getSnapshot: () => StructuredAgentSessionReadSnapshot
   loadOlder: () => Promise<void>
   refresh: () => void
+  releasePendingTurn: (clientMessageId: string) => void
+  retainPendingTurn: (clientMessageId: string) => void
   subscribe: (listener: () => void) => () => void
 }
 
@@ -58,6 +61,9 @@ function createReadOwner(
   let captureActiveHistoryReadGuard = (): (() => boolean) => retiredHistoryRead
   const activations = new Set<symbol>()
   const listeners = new Set<() => void>()
+  const pendingTurnRetentions = createStructuredAgentSessionPendingTurnRetentions(
+    () => snapshot.state
+  )
 
   const emit = (): void => {
     for (const listener of listeners) {
@@ -70,6 +76,7 @@ function createReadOwner(
     }
     snapshot = next
     emit()
+    pendingTurnRetentions.observe()
   }
   const apply = (action: StructuredAgentSessionAction): void => {
     const state = reduceStructuredAgentSession(snapshot.state, action)
@@ -178,7 +185,17 @@ function createReadOwner(
   }
 
   let owner: StructuredAgentSessionReadOwner
+  let releasingUnusedRetentions = false
   const deleteIfUnused = (): void => {
+    if (
+      listeners.size === 0 &&
+      pendingTurnRetentions.hasRetentions() &&
+      !releasingUnusedRetentions
+    ) {
+      releasingUnusedRetentions = true
+      pendingTurnRetentions.dispose()
+      releasingUnusedRetentions = false
+    }
     if (activations.size === 0 && listeners.size === 0 && owners.get(key) === owner) {
       owners.delete(key)
     }
@@ -190,6 +207,7 @@ function createReadOwner(
       if (activations.size === 1) {
         start()
       }
+      pendingTurnRetentions.releaseObservedTurn()
       return () => {
         activations.delete(token)
         if (activations.size === 0) {
@@ -199,6 +217,7 @@ function createReadOwner(
       }
     },
     dispose: () => {
+      pendingTurnRetentions.dispose()
       activations.clear()
       listeners.clear()
       stopActiveRun?.()
@@ -240,6 +259,9 @@ function createReadOwner(
       }
     },
     refresh: () => refreshActiveRun(),
+    releasePendingTurn: pendingTurnRetentions.release,
+    retainPendingTurn: (clientMessageId) =>
+      pendingTurnRetentions.retain(clientMessageId, owner.activate),
     subscribe: (listener) => {
       listeners.add(listener)
       return () => {

src/renderer/src/components/native-chat/use-structured-agent-session-outbox.test.tsx

@@ -106,6 +106,32 @@ describe('useStructuredAgentSessionOutbox', () => {
     )
   })
 
+  it('retains the session read from queueing until lifecycle settlement owns release', async () => {
+    const releasePendingTurn = vi.fn()
+    const retainPendingTurn = vi.fn()
+    mocks.call.mockResolvedValue(acceptedResult(1))
+    const view = renderHook(() =>
+      useStructuredAgentSessionOutbox({
+        sessionId: 'session-1',
+        target: LOCAL_TARGET,
+        fence: 1,
+        submissions: [],
+        releasePendingTurn,
+        retainPendingTurn
+      })
+    )
+
+    act(() => expect(view.result.current.send('hello')).toBe(true))
+
+    expect(retainPendingTurn).toHaveBeenCalledTimes(2)
+    expect(new Set(retainPendingTurn.mock.calls.flat()).size).toBe(1)
+    await waitFor(() => expect(view.result.current.outbox).toHaveLength(0))
+    expect(releasePendingTurn).not.toHaveBeenCalled()
+
+    view.unmount()
+    expect(releasePendingTurn).not.toHaveBeenCalled()
+  })
+
   it('requeues across a fence change and ignores the stale settlement', async () => {
     const first = deferred<ReturnType<typeof acceptedResult>>()
     const second = deferred<ReturnType<typeof acceptedResult>>()

src/renderer/src/components/native-chat/use-structured-agent-session-outbox.ts

@@ -16,6 +16,7 @@ import {
 import type { RuntimeClientTarget } from '@/runtime/runtime-rpc-client'
 import { callStructuredAgentSession } from '@/runtime/structured-agent-session-client'
 import { readOutbox, writeOutbox } from './structured-agent-session-outbox-storage'
+import { prepareStructuredAgentSessionOutboxRetry } from './structured-agent-session-outbox-retry'
 
 export function structuredSessionOperationId(): string {
   return createStructuredAgentSessionOperationId(() => crypto.randomUUID())
@@ -38,8 +39,10 @@ export function useStructuredAgentSessionOutbox(args: {
   target: RuntimeClientTarget
   fence: number | null
   submissions: readonly AgentJournalSubmission[]
+  releasePendingTurn?: (clientMessageId: string) => void
+  retainPendingTurn?: (clientMessageId: string) => void
 }) {
-  const { fence, sessionId, submissions, target } = args
+  const { fence, releasePendingTurn, retainPendingTurn, sessionId, submissions, target } = args
   const targetKey = target.kind === 'local' ? 'local' : `environment:${target.environmentId}`
   const [outbox, setOutbox] = useState<StructuredAgentSessionOutboxEntry[]>(() =>
     readOutbox(sessionId)
@@ -126,6 +129,7 @@ export function useStructuredAgentSessionOutbox(args: {
     }
     outboxRef.current = staged
     setOutbox(staged)
+    retainPendingTurn?.(next.clientMessageId)
     void callStructuredAgentSession<AgentSessionMutationResult<AgentSessionSendResult>>(
       target,
       'agentSession.send',
@@ -136,6 +140,7 @@ export function useStructuredAgentSessionOutbox(args: {
           return
         }
         if (!result.ok) {
+          releasePendingTurn?.(next.clientMessageId)
           setError(result.refusal.message)
           const updated = outboxRef.current.map((entry) =>
             entry.clientMessageId === next.clientMessageId
@@ -154,6 +159,7 @@ export function useStructuredAgentSessionOutbox(args: {
         }
         const submission = result.value.submission
         if (submission.dispatchState === 'rejected') {
+          releasePendingTurn?.(next.clientMessageId)
           blockedIdRef.current = next.clientMessageId
           setError(submission.reason ?? 'Message was not accepted')
         } else {
@@ -184,6 +190,7 @@ export function useStructuredAgentSessionOutbox(args: {
         }
         const failure = classifyStructuredAgentSessionSendFailure(caught, isDesktopDeliveryUnknown)
         if (failure === 'failed') {
+          releasePendingTurn?.(next.clientMessageId)
           blockedIdRef.current = next.clientMessageId
         }
         const updated = outboxRef.current.map((entry) =>
@@ -207,7 +214,7 @@ export function useStructuredAgentSessionOutbox(args: {
           dispatchingRef.current = false
         }
       })
-  }, [fence, outbox, sessionId, target])
+  }, [fence, outbox, releasePendingTurn, retainPendingTurn, sessionId, target])
 
   // A transport-side unknown may never have reached the host, and nothing else
   // moves it out of `unconfirmed`, so one wedges the whole FIFO queue. Re-issuing
@@ -270,61 +277,31 @@ export function useStructuredAgentSessionOutbox(args: {
       outboxRef.current = next
       setOutbox(next)
       setError(null)
+      retainPendingTurn?.(entry.clientMessageId)
       return true
     },
-    [sessionId]
+    [retainPendingTurn, sessionId]
   )
 
   const retry = (clientMessageId: string): void => {
     blockedIdRef.current = null
     setError(null)
-    const submission = submissions.find(
-      (candidate) => candidate.clientMessageId === clientMessageId
-    )
-    const current = outboxRef.current.find((entry) => entry.clientMessageId === clientMessageId)
-    // A provider-history reconciliation can settle an earlier unknown as
-    // rejected before the user presses Retry. Reusing that operation id only
-    // replays the settled rejection forever, so rotate the id for a safe resend.
-    if (current && submission?.dispatchState === 'rejected') {
-      const rotated = outboxRef.current.map((entry) =>
-        entry.clientMessageId === clientMessageId
-          ? {
-              ...entry,
-              clientMessageId: structuredSessionOperationId(),
-              state: 'queued' as const,
-              retryAfterUnknownSubmittedAt: null
-            }
-          : entry
-      )
-      if (!writeOutbox(sessionId, rotated)) {
-        setError('Message could not be saved to the outbox')
-        return
-      }
-      outboxRef.current = rotated
-      setOutbox(rotated)
-      return
-    }
-    const retryAfterUnknownSubmittedAt =
-      submission?.dispatchState === 'unknown'
-        ? submission.submittedAt
-        : current?.state === 'unconfirmed'
-          ? -1
-          : null
-    const next = outboxRef.current.map((entry) =>
-      entry.clientMessageId === clientMessageId
-        ? {
-            ...entry,
-            state: 'queued' as const,
-            retryAfterUnknownSubmittedAt
-          }
-        : entry
-    )
-    if (!writeOutbox(sessionId, next)) {
+    const prepared = prepareStructuredAgentSessionOutboxRetry({
+      clientMessageId,
+      createOperationId: structuredSessionOperationId,
+      entries: outboxRef.current,
+      submissions
+    })
+    if (!writeOutbox(sessionId, prepared.entries)) {
       setError('Message could not be saved to the outbox')
       return
     }
-    outboxRef.current = next
-    setOutbox(next)
+    outboxRef.current = prepared.entries
+    setOutbox(prepared.entries)
+    if (prepared.nextClientMessageId !== clientMessageId) {
+      releasePendingTurn?.(clientMessageId)
+    }
+    retainPendingTurn?.(prepared.nextClientMessageId)
   }
   return { outbox, error, blockedClientMessageId: blockedIdRef.current, send, retry }
 }

src/renderer/src/components/native-chat/use-structured-agent-session-read.ts

@@ -61,6 +61,8 @@ export function useStructuredAgentSessionRead(args: {
     state: snapshot.state,
     loadingOlder: snapshot.loadingOlder,
     loadOlder: owner.loadOlder,
-    providerSession: snapshot.providerSession
+    providerSession: snapshot.providerSession,
+    releasePendingTurn: owner.releasePendingTurn,
+    retainPendingTurn: owner.retainPendingTurn
   }
 }

src/renderer/src/components/native-chat/use-structured-agent-session.ts

@@ -48,11 +48,8 @@ export function useStructuredAgentSession(args: {
     surface: 'desktop-chat',
     enabled: isVisible
   })
-  const { state, loadingOlder, loadOlder } = useStructuredAgentSessionRead({
-    sessionId,
-    target,
-    isVisible
-  })
+  const { state, loadingOlder, loadOlder, releasePendingTurn, retainPendingTurn } =
+    useStructuredAgentSessionRead({ sessionId, target, isVisible })
   const stateRef = useRef(state)
   const [writeError, setWriteError] = useState<string | null>(null)
   const operationIds = useRef(new Map<string, string>())
@@ -65,7 +62,9 @@ export function useStructuredAgentSession(args: {
     sessionId,
     target,
     fence: state.fence,
-    submissions: state.submissions
+    submissions: state.submissions,
+    releasePendingTurn,
+    retainPendingTurn
   })
 
   useEffect(() => {