From bf6c5627b6945caef401d17814c30dd325df179b Mon Sep 17 00:00:00 2001 From: Merge Sim Date: Fri, 4 Sep 2026 15:45:55 -0700 Subject: [PATCH] fix(native-chat): retain reads through send handoff --- ...tructuredAgentSessionStatusBridge.test.tsx | 103 ++++++++++--- .../structured-agent-session-outbox-retry.ts | 39 +++++ ...d-agent-session-pending-turn-retentions.ts | 143 ++++++++++++++++++ .../structured-agent-session-read-owner.ts | 22 +++ ...e-structured-agent-session-outbox.test.tsx | 26 ++++ .../use-structured-agent-session-outbox.ts | 71 +++------ .../use-structured-agent-session-read.ts | 4 +- .../use-structured-agent-session.ts | 11 +- 8 files changed, 344 insertions(+), 75 deletions(-) create mode 100644 src/renderer/src/components/native-chat/structured-agent-session-outbox-retry.ts create mode 100644 src/renderer/src/components/native-chat/structured-agent-session-pending-turn-retentions.ts diff --git a/src/renderer/src/components/native-chat/StructuredAgentSessionStatusBridge.test.tsx b/src/renderer/src/components/native-chat/StructuredAgentSessionStatusBridge.test.tsx index 78cc8b75b04..fdad08e2f83 100644 --- a/src/renderer/src/components/native-chat/StructuredAgentSessionStatusBridge.test.tsx +++ b/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() - 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() 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() + await waitFor(() => expect(mocks.subscribe).toHaveBeenCalledOnce()) + act(() => { + retainPendingTurn?.('client-1') + }) + view.rerender() + 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() await waitFor(() => expect(mocks.setAgentStatus).toHaveBeenCalledOnce()) diff --git a/src/renderer/src/components/native-chat/structured-agent-session-outbox-retry.ts b/src/renderer/src/components/native-chat/structured-agent-session-outbox-retry.ts new file mode 100644 index 00000000000..fede9c64bc8 --- /dev/null +++ b/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 + } +} diff --git a/src/renderer/src/components/native-chat/structured-agent-session-pending-turn-retentions.ts b/src/renderer/src/components/native-chat/structured-agent-session-pending-turn-retentions.ts new file mode 100644 index 00000000000..200953663b3 --- /dev/null +++ b/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 & { + turnLifecycle: NonNullable< + Extract['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() + 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) + } + } +} diff --git a/src/renderer/src/components/native-chat/structured-agent-session-read-owner.ts b/src/renderer/src/components/native-chat/structured-agent-session-read-owner.ts index 4c43326acb1..5200c477f43 100644 --- a/src/renderer/src/components/native-chat/structured-agent-session-read-owner.ts +++ b/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 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() 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 () => { diff --git a/src/renderer/src/components/native-chat/use-structured-agent-session-outbox.test.tsx b/src/renderer/src/components/native-chat/use-structured-agent-session-outbox.test.tsx index 403769144fa..8f316275dd4 100644 --- a/src/renderer/src/components/native-chat/use-structured-agent-session-outbox.test.tsx +++ b/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>() const second = deferred>() diff --git a/src/renderer/src/components/native-chat/use-structured-agent-session-outbox.ts b/src/renderer/src/components/native-chat/use-structured-agent-session-outbox.ts index 79b1e2c95bb..a4dea180ec8 100644 --- a/src/renderer/src/components/native-chat/use-structured-agent-session-outbox.ts +++ b/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(() => readOutbox(sessionId) @@ -126,6 +129,7 @@ export function useStructuredAgentSessionOutbox(args: { } outboxRef.current = staged setOutbox(staged) + retainPendingTurn?.(next.clientMessageId) void callStructuredAgentSession>( 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 } } diff --git a/src/renderer/src/components/native-chat/use-structured-agent-session-read.ts b/src/renderer/src/components/native-chat/use-structured-agent-session-read.ts index df39ca3b21d..682db0b0321 100644 --- a/src/renderer/src/components/native-chat/use-structured-agent-session-read.ts +++ b/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 } } diff --git a/src/renderer/src/components/native-chat/use-structured-agent-session.ts b/src/renderer/src/components/native-chat/use-structured-agent-session.ts index 5bea1af8c50..1edaeaa6461 100644 --- a/src/renderer/src/components/native-chat/use-structured-agent-session.ts +++ b/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(null) const operationIds = useRef(new Map()) @@ -65,7 +62,9 @@ export function useStructuredAgentSession(args: { sessionId, target, fence: state.fence, - submissions: state.submissions + submissions: state.submissions, + releasePendingTurn, + retainPendingTurn }) useEffect(() => {