fix structured outbox and tab restore races

This commit is contained in:
Brennan Benson
2026-08-12 20:27:38 -07:00
parent e444d4bedd
commit b8d906228f
6 changed files with 258 additions and 47 deletions
@@ -0,0 +1,143 @@
import {
useCallback,
useMemo,
useRef,
type Dispatch,
type MutableRefObject,
type SetStateAction
} from 'react'
import type { AgentJournalSubmission } from '../../../src/shared/agent-session-journal-types'
import {
reconcileStructuredAgentSessionOutbox,
updateStructuredAgentSessionOutboxEntry
} from '../../../src/shared/structured-agent-session-outbox'
import type { MobileStructuredOutboxHydration } from './mobile-structured-outbox-hydration'
import type { MobileStructuredOutboxEntry } from './mobile-structured-outbox-store'
type PersistOutbox = (sessionId: string, entries: MobileStructuredOutboxEntry[]) => Promise<void>
export function useMobileStructuredOutboxMutations(args: {
activeSessionRef: MutableRefObject<string | null>
hydrationRef: MutableRefObject<MobileStructuredOutboxHydration | null>
outboxRef: MutableRefObject<MobileStructuredOutboxEntry[]>
setOutbox: Dispatch<SetStateAction<MobileStructuredOutboxEntry[]>>
persist: PersistOutbox
}) {
const { activeSessionRef, hydrationRef, outboxRef, persist, setOutbox } = args
const tailRef = useRef<Promise<void>>(Promise.resolve())
const sessionGenerationRef = useRef(0)
const serialize = useCallback(<T>(mutation: () => Promise<T>): Promise<T> => {
const operation = tailRef.current.then(mutation, mutation)
tailRef.current = operation.then(
() => undefined,
() => undefined
)
return operation
}, [])
const beginSession = useCallback(() => {
sessionGenerationRef.current += 1
}, [])
const isCurrent = useCallback(
(sessionId: string, generation: number) =>
activeSessionRef.current === sessionId && sessionGenerationRef.current === generation,
[activeSessionRef]
)
const commit = useCallback(
async (sessionId: string, generation: number, next: MobileStructuredOutboxEntry[]) => {
await persist(sessionId, next)
if (!isCurrent(sessionId, generation)) {
return false
}
outboxRef.current = next
setOutbox(next)
return true
},
[isCurrent, outboxRef, persist, setOutbox]
)
const reconcile = useCallback(
(sessionId: string, submissions: readonly AgentJournalSubmission[]) => {
const generation = sessionGenerationRef.current
return serialize(async () => {
if (!isCurrent(sessionId, generation)) {
return
}
const current = outboxRef.current
const next = reconcileStructuredAgentSessionOutbox(current, submissions)
if (
next.length === current.length &&
next.every((entry, index) => entry === current[index])
) {
return
}
await commit(sessionId, generation, next)
})
},
[commit, isCurrent, outboxRef, serialize]
)
const normalizeDispatching = useCallback(
(sessionId: string) => {
const generation = sessionGenerationRef.current
return serialize(async () => {
if (!isCurrent(sessionId, generation)) {
return
}
const current = outboxRef.current
const next = current.map((entry) =>
entry.state === 'dispatching' ? { ...entry, state: 'queued' as const } : entry
)
if (next.every((entry, index) => entry === current[index])) {
return
}
await commit(sessionId, generation, next)
})
},
[commit, isCurrent, outboxRef, serialize]
)
const replaceEntry = useCallback(
(
sessionId: string,
id: string,
update: (entry: MobileStructuredOutboxEntry) => MobileStructuredOutboxEntry | null
) => {
const generation = sessionGenerationRef.current
return serialize(async () => {
if (!isCurrent(sessionId, generation)) {
return
}
const next = updateStructuredAgentSessionOutboxEntry(outboxRef.current, id, update)
await commit(sessionId, generation, next)
})
},
[commit, isCurrent, outboxRef, serialize]
)
const enqueue = useCallback(
(
sessionId: string,
hydration: MobileStructuredOutboxHydration,
entry: MobileStructuredOutboxEntry
) => {
const generation = sessionGenerationRef.current
return serialize(async () => {
if (!isCurrent(sessionId, generation) || hydrationRef.current !== hydration) {
return false
}
const next = [...outboxRef.current, entry]
return commit(sessionId, generation, next)
})
},
[commit, hydrationRef, isCurrent, outboxRef, serialize]
)
return useMemo(
() => ({ beginSession, enqueue, normalizeDispatching, reconcile, replaceEntry }),
[beginSession, enqueue, normalizeDispatching, reconcile, replaceEntry]
)
}
@@ -182,6 +182,60 @@ describe('useMobileStructuredSessionWrites', () => {
expect(api!.outbox).toEqual([])
})
it('does not install a saved outbox entry after switching sessions', async () => {
const saved = deferred<void>()
vi.mocked(saveMobileStructuredOutbox).mockReturnValueOnce(saved.promise)
let sending!: Promise<boolean>
act(() => {
sending = api!.send('old session message')
})
await vi.waitFor(() => expect(saveMobileStructuredOutbox).toHaveBeenCalledTimes(1))
sessionId = 'mobile_2'
await act(async () => {
renderer!.update(createElement(Probe))
await Promise.resolve()
})
expect(api!.outbox).toEqual([])
await act(async () => {
saved.resolve()
await expect(sending).resolves.toBe(false)
await Promise.resolve()
})
expect(api!.outbox).toEqual([])
expect(sendRequest).not.toHaveBeenCalled()
})
it('serializes concurrent sends against the latest durable outbox', async () => {
connected = false
act(() => renderer!.update(createElement(Probe)))
const firstSave = deferred<void>()
vi.mocked(saveMobileStructuredOutbox).mockReturnValueOnce(firstSave.promise)
let first!: Promise<boolean>
let second!: Promise<boolean>
act(() => {
first = api!.send('one')
second = api!.send('two')
})
await vi.waitFor(() => expect(saveMobileStructuredOutbox).toHaveBeenCalledTimes(1))
await act(async () => {
firstSave.resolve()
await expect(Promise.all([first, second])).resolves.toEqual([true, true])
})
expect(api!.outbox.map((entry) => entry.body.blocks[0])).toEqual([
{ type: 'text', text: 'one' },
{ type: 'text', text: 'two' }
])
expect(saveMobileStructuredOutbox).toHaveBeenLastCalledWith('mobile_1', api!.outbox)
expect(sendRequest).not.toHaveBeenCalled()
})
it.each(['agent_session_operation_conflict', 'agent_session_operation_expired'] as const)(
'rotates a send operation after %s',
async (code) => {
@@ -4,11 +4,7 @@ import type {
AgentSessionMutationResult,
AgentSessionSendResult
} from '../../../src/shared/agent-session-wire'
import {
reconcileStructuredAgentSessionOutbox,
structuredAgentSessionSendRequest,
updateStructuredAgentSessionOutboxEntry
} from '../../../src/shared/structured-agent-session-outbox'
import { structuredAgentSessionSendRequest } from '../../../src/shared/structured-agent-session-outbox'
import type { RpcClient } from '../transport/rpc-client'
import type { PendingNativeChatImage } from './mobile-native-chat-image-attachment'
import {
@@ -23,6 +19,7 @@ import {
} from './mobile-structured-outbox-entry'
import type { MobileStructuredOutboxEntry } from './mobile-structured-outbox-store'
import { useMobileStructuredOutboxPersistence } from './use-mobile-structured-outbox-persistence'
import { useMobileStructuredOutboxMutations } from './use-mobile-structured-outbox-mutations'
import { useMobileStructuredSessionMutations } from './use-mobile-structured-session-mutations'
import type { MobileStructuredSessionWrites } from './mobile-structured-session-write-types'
@@ -56,8 +53,16 @@ export function useMobileStructuredSessionWrites(args: {
handoffOperationId: args.handoffOperationId,
onRefusal: setError
})
const outboxMutations = useMobileStructuredOutboxMutations({
activeSessionRef,
hydrationRef,
outboxRef,
setOutbox,
persist
})
useEffect(() => {
outboxMutations.beginSession()
activeSessionRef.current = sessionId
dispatchGenerationRef.current += 1
dispatchingRef.current = false
@@ -77,22 +82,14 @@ export function useMobileStructuredSessionWrites(args: {
})
hydrationRef.current = hydration
return hydration.cancel
}, [sessionId])
}, [outboxMutations, sessionId])
useEffect(() => {
if (!sessionId || submissions.length === 0 || outboxRef.current.length === 0) {
return
}
const next = reconcileStructuredAgentSessionOutbox(outboxRef.current, submissions)
if (
next.length !== outboxRef.current.length ||
next.some((entry, index) => entry !== outboxRef.current[index])
) {
outboxRef.current = next
setOutbox(next)
void persist(sessionId, next)
}
}, [persist, sessionId, submissions])
void outboxMutations.reconcile(sessionId, submissions)
}, [outboxMutations, sessionId, submissions])
useEffect(() => {
activeSessionRef.current = sessionId
@@ -101,20 +98,12 @@ export function useMobileStructuredSessionWrites(args: {
dispatchingRef.current = false
blockedIdRef.current = null
if (sessionId) {
const current = outboxRef.current
const next = current.map((entry) =>
entry.state === 'dispatching' ? { ...entry, state: 'queued' as const } : entry
)
if (next.some((entry, index) => entry !== current[index])) {
outboxRef.current = next
setOutbox(next)
void persist(sessionId, next)
}
void outboxMutations.normalizeDispatching(sessionId)
}
if (connected) {
setDispatchVersion((current) => current + 1)
}
}, [connected, fence, persist, sessionId])
}, [connected, fence, outboxMutations, sessionId])
const replaceEntry = useCallback(
async (
@@ -124,12 +113,9 @@ export function useMobileStructuredSessionWrites(args: {
if (!sessionId) {
return
}
const next = updateStructuredAgentSessionOutboxEntry(outboxRef.current, id, update)
outboxRef.current = next
setOutbox(next)
await persist(sessionId, next)
await outboxMutations.replaceEntry(sessionId, id, update)
},
[persist, sessionId]
[outboxMutations, sessionId]
)
useEffect(() => {
@@ -239,7 +225,7 @@ export function useMobileStructuredSessionWrites(args: {
return false
}
const hydration = hydrationRef.current
if (!(await waitForMobileStructuredOutboxHydration(hydration, sessionId))) {
if (!hydration || !(await waitForMobileStructuredOutboxHydration(hydration, sessionId))) {
setError('Chat session is still loading')
return false
}
@@ -247,19 +233,18 @@ export function useMobileStructuredSessionWrites(args: {
return false
}
const entry = createQueuedMobileStructuredOutboxEntry({ sessionId, text, attachments })
const next = [...outboxRef.current, entry]
try {
await persist(sessionId, next)
const saved = await outboxMutations.enqueue(sessionId, hydration, entry)
if (saved) {
setError(null)
}
return saved
} catch {
setError('Message could not be saved to the outbox')
return false
}
outboxRef.current = next
setOutbox(next)
setError(null)
return true
},
[persist, sessionId]
[outboxMutations, sessionId]
)
const takeQueuedForEdit = useCallback(
@@ -53,7 +53,13 @@ describe('local structured session tab projection', () => {
} satisfies RuntimeMobileSessionTabsResult
expect(projectLocalStructuredSessionTabs(snapshot)).toMatchObject({
tabGroups: undefined,
tabGroups: [
{
id: 'structured-group',
activeTabId: 'agent-session:codex-1',
tabOrder: ['agent-session:codex-1']
}
],
tabGroupLayout: undefined,
tabs: [expect.objectContaining({ type: 'agent-session', agent: 'codex' })]
})
@@ -20,8 +20,16 @@ export function projectLocalStructuredSessionTabs(
return {
...snapshot,
tabs: snapshot.tabs.filter((tab) => structuredIds.has(tab.id)),
// Why: replaying local terminal topology under the structured owner clones the workspace layout.
tabGroups: undefined,
tabGroups: snapshot.tabGroups
?.map((group) => ({
...group,
tabOrder: group.tabOrder.filter((id) => structuredIds.has(id)),
activeTabId:
group.activeTabId && structuredIds.has(group.activeTabId) ? group.activeTabId : null,
recentTabIds: group.recentTabIds?.filter((id) => structuredIds.has(id))
}))
.filter((group) => group.tabOrder.length > 0),
// Why: group membership locates chats; the renderer's split tree remains locally authoritative.
tabGroupLayout: undefined
}
}
@@ -33,7 +41,9 @@ function applySnapshots(snapshots: readonly RuntimeMobileSessionTabsResult[]): v
const patch = applyWebSessionTabsSnapshot(
next,
projectLocalStructuredSessionTabs(snapshot),
LOCAL_STRUCTURED_SESSION_OWNER
LOCAL_STRUCTURED_SESSION_OWNER,
Date.now(),
{ preserveLocalLayout: true }
)
next = patch === next ? next : ({ ...next, ...patch } as typeof state)
}
@@ -122,6 +122,10 @@ type SessionTabsRemovalFence = {
pendingCount: number
}
export type WebSessionTabsSnapshotApplyOptions = {
preserveLocalLayout?: boolean
}
type TrackedWebSessionTabsWorktree = {
worktree: string
freshness: SnapshotFreshness
@@ -2342,7 +2346,8 @@ function applyWebSessionTabsSnapshotWithContext(
rawSnapshot: RuntimeMobileSessionTabsResult,
environmentId: string,
now = Date.now(),
batchContext?: WebSessionTabsBatchContext
batchContext?: WebSessionTabsBatchContext,
options?: WebSessionTabsSnapshotApplyOptions
): WebSessionTabsSyncState | Partial<WebSessionTabsSyncState> {
const worktreeId = rawSnapshot.worktree
if (worktreeId === FLOATING_TERMINAL_WORKTREE_ID) {
@@ -3110,7 +3115,7 @@ function applyWebSessionTabsSnapshotWithContext(
)
: state.activeGroupIdByWorktree
const nextLayoutByWorktree = (() => {
if (!nextGroups) {
if (!nextGroups || options?.preserveLocalLayout) {
return state.layoutByWorktree
}
const validGroupIds = new Set(nextGroups.map((group) => group.id))
@@ -3370,9 +3375,17 @@ export function applyWebSessionTabsSnapshot(
state: WebSessionTabsSyncState,
rawSnapshot: RuntimeMobileSessionTabsResult,
environmentId: string,
now = Date.now()
now = Date.now(),
options?: WebSessionTabsSnapshotApplyOptions
): WebSessionTabsSyncState | Partial<WebSessionTabsSyncState> {
return applyWebSessionTabsSnapshotWithContext(state, rawSnapshot, environmentId, now)
return applyWebSessionTabsSnapshotWithContext(
state,
rawSnapshot,
environmentId,
now,
undefined,
options
)
}
export function applyWebSessionTabsSnapshots(