diff --git a/mobile/src/session/use-mobile-structured-agent-state.test.tsx b/mobile/src/session/use-mobile-structured-agent-state.test.tsx index eaaf5610dc0..a141ec47d2c 100644 --- a/mobile/src/session/use-mobile-structured-agent-state.test.tsx +++ b/mobile/src/session/use-mobile-structured-agent-state.test.tsx @@ -6,7 +6,11 @@ import type { AgentJournalRenderItem, AgentJournalResetReason } from '../../../src/shared/agent-session-journal-types' -import type { AgentSessionHistoryPage } from '../../../src/shared/agent-session-wire' +import { STRUCTURED_AGENT_SESSION_CLIENT_COALESCE_MS } from '../../../src/shared/structured-agent-session-coalescer' +import type { + AgentSessionSubscribeEvent, + AgentSessionHistoryPage +} from '../../../src/shared/agent-session-wire' import { publishConnectionRevival } from '../transport/connection-revival-signal' import type { RpcClient } from '../transport/rpc-client' import { useMobileStructuredAgentState } from './use-mobile-structured-agent-state' @@ -122,9 +126,39 @@ describe('useMobileStructuredAgentState older history', () => { }) }) await vi.waitFor(() => expect(hook.current!.state.status).toBe('ready')) - return { hook, historyCursors } + return { hook, historyCursors, emit: (event: AgentSessionSubscribeEvent) => onFrame?.(event) } } + it('keeps a row removed when its upsert and removal arrive within one streaming window', async () => { + const { hook, emit } = await mountWithOlderPages({}, []) + vi.useFakeTimers() + try { + const batch = { + cursor: cursorAt(1001), + items: [said(1001)], + removedItemIds: [], + submissions: [] + } + act(() => emit({ type: 'batch', sessionId: 'session-1', batch })) + act(() => + emit({ + type: 'batch', + sessionId: 'session-1', + batch: { ...batch, cursor: cursorAt(1002), items: [], removedItemIds: ['item-1001'] } + }) + ) + await act(async () => { + await vi.advanceTimersByTimeAsync(STRUCTURED_AGENT_SESSION_CLIENT_COALESCE_MS) + }) + expect(hook.current!.state.items).toEqual([]) + expect(hook.current!.state.cursor?.sequence).toBe(1002) + } finally { + act(() => renderer?.unmount()) + renderer = null + vi.useRealTimers() + } + }) + it("reads past pages of a subagent's rows to the session's own older rows", async () => { const { hook, historyCursors } = await mountWithOlderPages({ 1000: page(range(800, 1_000, 'task-1'), true), diff --git a/src/renderer/src/components/native-chat/use-structured-agent-session-read.test.tsx b/src/renderer/src/components/native-chat/use-structured-agent-session-read.test.tsx index f6262411d6f..0513383f67e 100644 --- a/src/renderer/src/components/native-chat/use-structured-agent-session-read.test.tsx +++ b/src/renderer/src/components/native-chat/use-structured-agent-session-read.test.tsx @@ -100,6 +100,48 @@ describe('useStructuredAgentSessionRead history window', () => { mocks.subscribe.mockResolvedValue({ unsubscribe: vi.fn() }) }) + it('keeps a row removed when its upsert and removal arrive within one streaming window', async () => { + mocks.call.mockResolvedValueOnce({ ok: true, page: page('tail', [], false) }) + const { result, unmount } = renderHook(() => + useStructuredAgentSessionRead({ sessionId: 'session-a', target: LOCAL_TARGET }) + ) + await waitFor(() => expect(result.current.state.status).toBe('ready')) + const emit: (event: AgentSessionSubscribeEvent) => void = mocks.subscribe.mock.calls[0][2] + act(() => + emit({ type: 'snapshot', sessionId: 'session-a', page: page('tail', [], false), fence: 1 }) + ) + vi.useFakeTimers() + try { + const batch = { + cursor: { epoch: 'epoch-a', sequence: 501 }, + items: [message('answer', 501, 'assistant')], + removedItemIds: [], + submissions: [] + } + act(() => emit({ type: 'batch', sessionId: 'session-a', batch })) + act(() => + emit({ + type: 'batch', + sessionId: 'session-a', + batch: { + ...batch, + cursor: { epoch: 'epoch-a', sequence: 502 }, + items: [], + removedItemIds: ['answer'] + } + }) + ) + await act(async () => { + await vi.advanceTimersByTimeAsync(STRUCTURED_AGENT_SESSION_CLIENT_COALESCE_MS) + }) + expect(result.current.state.items).toEqual([]) + expect(result.current.state.cursor?.sequence).toBe(502) + } finally { + unmount() + vi.useRealTimers() + } + }) + it('restores a realistic 21-turn window across the wire-safe bridge-sized read', async () => { const items = Array.from({ length: 21 }, (_, turn) => [ message(`user-${turn}`, turn * 2 + 1, 'user'), diff --git a/src/shared/structured-agent-session-coalescer.test.ts b/src/shared/structured-agent-session-coalescer.test.ts index 26030fa56be..5a10cb3a552 100644 --- a/src/shared/structured-agent-session-coalescer.test.ts +++ b/src/shared/structured-agent-session-coalescer.test.ts @@ -1,6 +1,17 @@ -import { describe, expect, it } from 'vitest' +import { describe, expect, it, vi } from 'vitest' +import type { AgentJournalRenderItem } from './agent-session-journal-types' +import { projectNativeChatTranscriptMessages } from './native-chat-transcript-projection' +import { projectStructuredAgentSessionMessages } from './structured-agent-session-message-projection' +import { + EMPTY_STRUCTURED_AGENT_SESSION, + reduceStructuredAgentSession, + type StructuredAgentSessionState +} from './structured-agent-session-reducer' import type { AgentSessionSubscribeEvent } from './agent-session-wire' -import { createStructuredAgentSessionEventCoalescer } from './structured-agent-session-coalescer' +import { + STRUCTURED_AGENT_SESSION_CLIENT_COALESCE_MS, + createStructuredAgentSessionEventCoalescer +} from './structured-agent-session-coalescer' function batch( sequence: number, @@ -215,3 +226,192 @@ describe('structured agent session event coalescer', () => { expect(events[0]).not.toHaveProperty('latestTurn') }) }) + +describe('coalesced history removals', () => { + const roster = { + type: 'subagent-group' as const, + groupId: 'group', + agents: [{ id: 'helper', label: 'review', state: 'working' as const }] + } + const item = (withRoster = false): AgentJournalRenderItem => ({ + itemId: 'answer', + revision: 1, + sequence: 1, + observedAt: 1, + body: { + kind: 'message', + role: 'assistant', + blocks: withRoster ? [roster] : [{ type: 'text', text: 'answer' }] + } + }) + function history() { + let state: StructuredAgentSessionState = { ...EMPTY_STRUCTURED_AGENT_SESSION, epoch: 'epoch-1' } + const events: AgentSessionSubscribeEvent[] = [] + const coalescer = createStructuredAgentSessionEventCoalescer((event) => { + events.push(event) + state = reduceStructuredAgentSession(state, { type: 'event', event }, 0) + }) + const push = ( + sequence: number, + items: AgentJournalRenderItem[], + removedItemIds: string[] = [] + ) => { + const event = { + ...batch(sequence), + batch: { ...batch(sequence).batch, items, removedItemIds } + } + coalescer.push(event) + return event + } + return { coalescer, events, push, state: () => state } + } + + it.each([ + { name: 'assistant text', withRoster: false, replace: false }, + { name: 'assistant roster', withRoster: true, replace: false }, + { name: 'same-frame replacement', withRoster: false, replace: true } + ])( + 'preserves removal order for $name, then permits a later reintroduction', + ({ withRoster, replace }) => { + vi.useFakeTimers() + const h = history() + const visibleBlocks = () => + projectNativeChatTranscriptMessages( + projectStructuredAgentSessionMessages(h.state().items, [], h.state().submissions) + ).map((message) => message.blocks) + try { + const original = item(withRoster) + h.push(1, [original]) + h.push(2, [{ ...original, revision: 2 }]) + expect(h.events).toHaveLength(0) + const replacement: AgentJournalRenderItem = { + ...item(), + revision: 3, + body: { + kind: 'message', + role: 'assistant', + blocks: [{ type: 'text', text: 'replacement' }] + } + } + const removal = h.push(3, replace ? [replacement] : [], ['answer']) + const immediateCursor = h.state().cursor + vi.advanceTimersByTime(STRUCTURED_AGENT_SESSION_CLIENT_COALESCE_MS) + expect(h.state().items).toEqual(replace ? [replacement] : []) + expect(immediateCursor).toEqual(removal.batch.cursor) + expect(h.events).toHaveLength(2) + expect(h.events[0]).toMatchObject({ batch: { items: [{ itemId: 'answer', revision: 2 }] } }) + expect(h.events[1]).toBe(removal) + expect(h.state().items).toEqual(replace ? [replacement] : []) + expect(h.state().subagentRoster?.size).toBe(0) + expect(h.state().subagentRoster?.has('helper')).toBe(false) + expect(h.state().cursor).toEqual(removal.batch.cursor) + vi.advanceTimersByTime(STRUCTURED_AGENT_SESSION_CLIENT_COALESCE_MS) + expect(h.state().items).toEqual(replace ? [replacement] : []) + expect(visibleBlocks()).toEqual(replace ? [[{ type: 'text', text: 'replacement' }]] : []) + const reintroduced = { ...item(), revision: 4, sequence: 4 } + h.push(4, [reintroduced]) + vi.advanceTimersByTime(STRUCTURED_AGENT_SESSION_CLIENT_COALESCE_MS) + expect(h.state().items).toEqual([reintroduced]) + expect(visibleBlocks()).toEqual([[{ type: 'text', text: 'answer' }]]) + } finally { + h.coalescer.dispose() + vi.useRealTimers() + } + } + ) + + it('removes a system roster immediately after flushing an unrelated assistant update', () => { + const h = history() + try { + h.push(1, [ + { ...item(), itemId: 'roster', body: { kind: 'message', role: 'system', blocks: [roster] } } + ]) + expect(h.state().subagentRoster?.has('helper')).toBe(true) + const answer = { ...item(), sequence: 2 } + h.push(2, [answer]) + const removal = h.push(3, [], ['roster']) + expect(h.state().cursor).toEqual(removal.batch.cursor) + expect(h.state().items).toEqual([answer]) + expect(h.state().subagentRoster?.size).toBe(0) + } finally { + h.coalescer.dispose() + } + }) + + it('bounds destructive burst work at N and 10N while ordinary updates still aggregate', () => { + const population = 1_000 + const retained = Array.from({ length: population }, (_, index) => ({ + ...item(), + itemId: `retained-${index}`, + sequence: index + 1 + })) + function run(count: number, remove: boolean) { + let state: StructuredAgentSessionState = { + ...EMPTY_STRUCTURED_AGENT_SESSION, + epoch: 'epoch-1', + status: 'ready', + items: retained, + retainedOwnItemLimit: population + 1, + retainedItemCap: population + 1 + } + let reductions = 0 + let inputVisits = 0 + const coalescer = createStructuredAgentSessionEventCoalescer((event) => { + reductions++ + // Count retained input reads, including head lookup, without instrumenting production. + const items = [...state.items] + items.forEach((row, index) => { + Object.defineProperty(items, index, { + get() { + inputVisits++ + return row + } + }) + }) + state = reduceStructuredAgentSession({ ...state, items }, { type: 'event', event }, 0) + }) + try { + for (let index = 0; index < count; index++) { + const update = batch(population + index * 2 + 1) + const row: AgentJournalRenderItem = { + ...item(), + itemId: remove ? `answer-${index}` : 'answer', + sequence: remove ? update.batch.cursor.sequence : population + 1, + revision: index + 1 + } + coalescer.push({ ...update, batch: { ...update.batch, items: [row] } }) + if (remove) { + const removal = batch(update.batch.cursor.sequence + 1) + coalescer.push({ + ...removal, + batch: { ...removal.batch, removedItemIds: [row.itemId] } + }) + } + } + expect(reductions).toBe(remove ? 2 * count : 0) + vi.advanceTimersByTime(STRUCTURED_AGENT_SESSION_CLIENT_COALESCE_MS) + expect(reductions).toBe(remove ? 2 * count : 1) + expect(inputVisits).toBeGreaterThan(0) + expect(inputVisits).toBeLessThanOrEqual(reductions * (population + 2)) + expect(state.items.slice(0, population)).toEqual(retained) + expect(state.items.slice(population)).toEqual( + remove ? [] : [expect.objectContaining({ itemId: 'answer', revision: count })] + ) + return { reductions, inputVisits } + } finally { + coalescer.dispose() + } + } + vi.useFakeTimers() + try { + const small = run(10, true) + const large = run(100, true) + expect(large.reductions).toBe(10 * small.reductions) + expect(large.inputVisits).toBeLessThanOrEqual(10 * small.inputVisits) + run(10, false) + run(100, false) + } finally { + vi.useRealTimers() + } + }) +}) diff --git a/src/shared/structured-agent-session-coalescer.ts b/src/shared/structured-agent-session-coalescer.ts index 398335509fc..bcf2366a2ed 100644 --- a/src/shared/structured-agent-session-coalescer.ts +++ b/src/shared/structured-agent-session-coalescer.ts @@ -13,6 +13,7 @@ function bypassCoalescing(event: AgentSessionSubscribeEvent): boolean { // Deliver queue lists immediately or the queued response removes Sending before its replacement card arrives. event.queuedMessages !== undefined || event.currentRequests !== undefined || + event.batch.removedItemIds.length > 0 || event.batch.items.some((item) => item.body.kind !== 'message' || item.body.role !== 'assistant') ) } @@ -47,7 +48,7 @@ function mergeBatch( batch: { cursor: right.batch.cursor, items: [...items.values()], - removedItemIds: [...new Set([...left.batch.removedItemIds, ...right.batch.removedItemIds])], + removedItemIds: [], submissions: [...submissions.values()] }, ...(right.fence !== undefined || left.fence !== undefined