diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-host-handoff.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-host-handoff.ts index 7aa79f67e9a..5c84e97c81f 100644 --- a/src/main/native-chat/agent-session-wire/structured-agent-session-host-handoff.ts +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-host-handoff.ts @@ -54,7 +54,7 @@ export function createStructuredAgentSessionHostHandoff( publish: (sessionId, status) => { const session = host.session(sessionId) const fence = deps.store.getRecord(sessionId)?.lease.runtimeFence ?? session.fence - host.subscribers.handoff(sessionId, session.journal, fence, status) + host.subscribers.handoff(sessionId, fence, status) }, schedule: host.serialize, now: host.now diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-host.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-host.ts index 8722cf89d32..52bc00099be 100644 --- a/src/main/native-chat/agent-session-wire/structured-agent-session-host.ts +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-host.ts @@ -42,6 +42,7 @@ import { } from './structured-agent-session-subscribers' import { StructuredAgentSessionTaskQueue } from './structured-agent-session-task-queue' import { restoreStructuredAgentSessionsOnRestart } from './structured-agent-session-restart-restore' +import { StructuredAgentSessionRestartRestoreGate } from './structured-agent-session-restart-restore-gate' import * as providerRouting from './structured-agent-session-provider-routing' import { createStructuredAgentSessionHostHandoff, @@ -67,6 +68,7 @@ export class StructuredAgentSessionHost { private readonly runtimeState: StructuredAgentSessionHostRuntimeState private readonly reconcileLeases: (sessionId: string) => Promise private readonly handoffs: StructuredAgentSessionHostHandoff + private readonly restartRestore = new StructuredAgentSessionRestartRestoreGate() constructor(readonly deps: StructuredAgentSessionHostDeps) { this.runtimeState = new StructuredAgentSessionHostRuntimeState(deps, (record) => @@ -104,27 +106,29 @@ export class StructuredAgentSessionHost { })) } - async restoreReadableSessions(): Promise { - await restoreStructuredAgentSessionsOnRestart({ - store: this.deps.store, - journalRoot: this.deps.journalRoot, - records: this.deps.store - .listRecords() - .filter((record) => - providerRouting.adapterSupportsAgentSessionRecord(this.deps.adapter, record) - ), - reconcile: this.reconcileLeases, - operationId: () => - `${Math.trunc(this.now()).toString().padStart(13, '0')}-${randomUUID().replaceAll('-', '')}`, - resume: (params) => - this.attach({ callerKey: 'trusted-local:host-restart' }, params).then( - (result) => result.ok - ), - serialize: (sessionId, task) => this.serialize(sessionId, task), - hasSession: (sessionId) => this.sessions.has(sessionId), - onReadable: (sessionId, restored) => this.sessions.set(sessionId, restored), - restoreHandoff: (sessionId) => this.handoffs.restore(sessionId) - }) + restoreReadableSessions(): Promise { + return this.restartRestore.run(() => + restoreStructuredAgentSessionsOnRestart({ + store: this.deps.store, + journalRoot: this.deps.journalRoot, + records: this.deps.store + .listRecords() + .filter((record) => + providerRouting.adapterSupportsAgentSessionRecord(this.deps.adapter, record) + ), + reconcile: this.reconcileLeases, + operationId: () => + `${Math.trunc(this.now()).toString().padStart(13, '0')}-${randomUUID().replaceAll('-', '')}`, + resume: (params) => + this.attach({ callerKey: 'trusted-local:host-restart' }, params).then( + (result) => result.ok + ), + serialize: (sessionId, task) => this.serialize(sessionId, task), + hasSession: (sessionId) => this.sessions.has(sessionId), + onReadable: (sessionId, restored) => this.sessions.set(sessionId, restored), + restoreHandoff: (sessionId) => this.handoffs.restore(sessionId) + }) + ) } private serialize(sessionId: string, task: () => Promise): Promise { diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-restart-restore-gate.test.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-restart-restore-gate.test.ts new file mode 100644 index 00000000000..4567d84d5c0 --- /dev/null +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-restart-restore-gate.test.ts @@ -0,0 +1,27 @@ +import { describe, expect, it, vi } from 'vitest' +import { StructuredAgentSessionRestartRestoreGate } from './structured-agent-session-restart-restore-gate' + +describe('StructuredAgentSessionRestartRestoreGate', () => { + it('runs one successful restore across concurrent and later discovery', async () => { + const gate = new StructuredAgentSessionRestartRestoreGate() + const restore = vi.fn(async () => undefined) + + await Promise.all([gate.run(restore), gate.run(restore), gate.run(restore)]) + await gate.run(restore) + + expect(restore).toHaveBeenCalledOnce() + }) + + it('allows a failed restore to be retried', async () => { + const gate = new StructuredAgentSessionRestartRestoreGate() + const restore = vi + .fn<() => Promise>() + .mockRejectedValueOnce(new Error('restore failed')) + .mockResolvedValue(undefined) + + await expect(gate.run(restore)).rejects.toThrow('restore failed') + await expect(gate.run(restore)).resolves.toBeUndefined() + + expect(restore).toHaveBeenCalledTimes(2) + }) +}) diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-restart-restore-gate.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-restart-restore-gate.ts new file mode 100644 index 00000000000..dd4237dcb28 --- /dev/null +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-restart-restore-gate.ts @@ -0,0 +1,17 @@ +export class StructuredAgentSessionRestartRestoreGate { + private current: Promise | null = null + + run(restore: () => Promise): Promise { + if (this.current) { + return this.current + } + const tracked = restore().catch((error: unknown) => { + if (this.current === tracked) { + this.current = null + } + throw error + }) + this.current = tracked + return tracked + } +} diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-restart-restore.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-restart-restore.ts index 9a45c53f5bf..95fdaf11d1c 100644 --- a/src/main/native-chat/agent-session-wire/structured-agent-session-restart-restore.ts +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-restart-restore.ts @@ -40,6 +40,7 @@ export async function restoreStructuredAgentSessionsOnRestart(input: { } await input.serialize(sessionId, async () => { if (input.hasSession(sessionId)) { + await input.restoreHandoff(sessionId) return } const restored = await restoreStructuredAgentSessionRead( @@ -47,13 +48,12 @@ export async function restoreStructuredAgentSessionsOnRestart(input: { input.journalRoot, sessionId ) - if (restored) { - input.onReadable(sessionId, restored) + if (!restored) { + return } - }) - if (input.hasSession(sessionId)) { + input.onReadable(sessionId, restored) await input.restoreHandoff(sessionId) - } + }) }) ) } diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-subscribers.test.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-subscribers.test.ts new file mode 100644 index 00000000000..e9c8d428daa --- /dev/null +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-subscribers.test.ts @@ -0,0 +1,68 @@ +import { mkdtemp, rm } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { afterEach, beforeEach, describe, expect, it } from 'vitest' +import type { + AgentSessionHandoffStatus, + AgentSessionSubscribeEvent +} from '../../../shared/agent-session-wire' +import { openAgentSessionJournal } from '../agent-session-journal/journal-store' +import { AgentSessionSubscribers } from './structured-agent-session-subscribers' + +const SESSION = 'subscriber-session' + +let root: string + +beforeEach(async () => { + root = await mkdtemp(join(tmpdir(), 'orca-agent-subscribers-')) +}) + +afterEach(async () => { + await rm(root, { recursive: true, force: true }) +}) + +describe('AgentSessionSubscribers', () => { + it('publishes handoff-only changes without serializing a transcript snapshot', async () => { + const journal = await openAgentSessionJournal({ + identity: { + sessionId: SESSION, + workspaceId: 'workspace-1', + hostId: 'local', + agent: 'codex', + providerHandle: { kind: 'codex', threadId: 'thread-1' } + }, + journalDir: join(root, 'journal') + }) + const subscribers = new AgentSessionSubscribers() + const events: AgentSessionSubscribeEvent[] = [] + subscribers.open({ + id: 'subscriber-1', + sessionId: SESSION, + journal, + fence: 1, + emit: (event) => events.push(event) + }) + const handoff: AgentSessionHandoffStatus = { + owner: 'native', + direction: 'to-tui', + phase: 'switching', + stage: 'preparing', + operationId: 'handoff-1' + } + + subscribers.handoff(SESSION, 2, handoff) + + expect(events.at(-1)).toEqual({ + type: 'batch', + sessionId: SESSION, + batch: { + cursor: journal.cursor(), + items: [], + removedItemIds: [], + submissions: [] + }, + fence: 2, + handoff + }) + }) +}) diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-subscribers.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-subscribers.ts index a5262de99b8..59003786ede 100644 --- a/src/main/native-chat/agent-session-wire/structured-agent-session-subscribers.ts +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-subscribers.ts @@ -118,16 +118,20 @@ export class AgentSessionSubscribers { } } - handoff( - sessionId: string, - journal: AgentSessionJournal, - fence: number, - handoff: AgentSessionHandoffStatus - ): void { - const snapshot = journal.snapshot() + handoff(sessionId: string, fence: number, handoff: AgentSessionHandoffStatus): void { for (const subscriber of this.subscribers(sessionId)) { - this.emit(subscriber, { type: 'snapshot', sessionId, snapshot, fence, handoff }) - subscriber.cursor = snapshot.cursor + this.emit(subscriber, { + type: 'batch', + sessionId, + batch: { + cursor: subscriber.cursor, + items: [], + removedItemIds: [], + submissions: [] + }, + fence, + handoff + }) subscriber.fence = fence } }