diff --git a/mobile/src/session/mobile-structured-agent-session-coalescer.test.ts b/mobile/src/session/mobile-structured-agent-session-coalescer.test.ts index da03cdb5202..dffc0bb45d7 100644 --- a/mobile/src/session/mobile-structured-agent-session-coalescer.test.ts +++ b/mobile/src/session/mobile-structured-agent-session-coalescer.test.ts @@ -38,6 +38,42 @@ describe('mobile structured event coalescer', () => { vi.useRealTimers() }) + it('retains ownership projection while merging a reconnect batch', () => { + vi.useFakeTimers() + const seen: AgentSessionSubscribeEvent[] = [] + const coalescer = createMobileStructuredEventCoalescer((event) => seen.push(event)) + coalescer.push({ + type: 'batch', + sessionId: 'session-a', + batch: { + cursor: { epoch: 'epoch-a', sequence: 1 }, + items: [], + removedItemIds: [], + submissions: [] + }, + fence: 3, + handoff: { + owner: 'tui', + direction: null, + phase: 'idle', + stage: null, + operationId: null, + terminal: { handle: 'term-1', tabId: 'tab-1', paneKey: 'tab-1:leaf-1' } + } + }) + coalescer.push(batch(2, 'continued')) + vi.advanceTimersByTime(48) + + expect(seen).toEqual([ + expect.objectContaining({ + type: 'batch', + fence: 3, + handoff: expect.objectContaining({ owner: 'tui', phase: 'idle' }) + }) + ]) + vi.useRealTimers() + }) + it('flushes pending text before a lifecycle row', () => { vi.useFakeTimers() const seen: AgentSessionSubscribeEvent[] = [] diff --git a/mobile/src/session/mobile-structured-agent-session-reducer.test.ts b/mobile/src/session/mobile-structured-agent-session-reducer.test.ts index a647a5013fb..31a1a589470 100644 --- a/mobile/src/session/mobile-structured-agent-session-reducer.test.ts +++ b/mobile/src/session/mobile-structured-agent-session-reducer.test.ts @@ -1,5 +1,6 @@ import { describe, expect, it } from 'vitest' import type { AgentJournalRenderItem } from '../../../src/shared/agent-session-journal-types' +import type { AgentSessionHandoffStatus } from '../../../src/shared/agent-session-wire' import { EMPTY_MOBILE_STRUCTURED_AGENT_SESSION, reduceMobileStructuredAgentSession @@ -25,6 +26,90 @@ function snapshot(epoch: string, items: AgentJournalRenderItem[]) { } describe('mobile structured session reducer', () => { + it('projects every stage across transfer, recovery, and return', () => { + const handoffs: AgentSessionHandoffStatus[] = [ + { owner: 'native', direction: null, phase: 'idle', stage: null, operationId: null }, + { + owner: 'native', + direction: 'to-tui', + phase: 'switching', + stage: 'preparing', + operationId: 'forward' + }, + { + owner: 'none', + direction: 'to-tui', + phase: 'switching', + stage: 'old-owner-stopped', + operationId: 'forward' + }, + { + owner: 'none', + direction: 'to-tui', + phase: 'switching', + stage: 'new-owner-proving', + operationId: 'forward' + }, + { + owner: 'tui', + direction: null, + phase: 'idle', + stage: null, + operationId: null, + terminal: { handle: 'term-1', tabId: 'tab-1', paneKey: 'tab-1:leaf-1' } + }, + { + owner: 'none', + direction: null, + phase: 'failed', + stage: 'manual-recovery', + operationId: null, + error: { message: 'Recovery required', recoverableOwner: 'none' } + }, + { + owner: 'tui', + direction: null, + phase: 'idle', + stage: null, + operationId: null, + terminal: { handle: 'term-2', tabId: 'tab-1', paneKey: 'tab-1:leaf-1' } + }, + { + owner: 'tui', + direction: 'to-native', + phase: 'switching', + stage: 'preparing', + operationId: 'reverse' + }, + { owner: 'native', direction: null, phase: 'idle', stage: null, operationId: null } + ] + let state = EMPTY_MOBILE_STRUCTURED_AGENT_SESSION + const projected: Array> = [] + + for (const [index, handoff] of handoffs.entries()) { + state = reduceMobileStructuredAgentSession(state, { + type: 'event', + event: { + ...snapshot('epoch-a', []), + fence: index + 1, + snapshot: { + ...snapshot('epoch-a', []).snapshot, + cursor: { epoch: 'epoch-a', sequence: index } + }, + handoff + } + }) + projected.push({ + owner: state.handoff!.owner, + phase: state.handoff!.phase, + stage: state.handoff!.stage + }) + } + + expect(projected).toEqual(handoffs.map(({ owner, phase, stage }) => ({ owner, phase, stage }))) + expect(state.handoff).toMatchObject({ owner: 'native', phase: 'idle', stage: null }) + }) + it('boots from a bounded tail page and retains its live resume cursor', () => { const state = reduceMobileStructuredAgentSession(EMPTY_MOBILE_STRUCTURED_AGENT_SESSION, { type: 'tail-page', diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-handoff-restart.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-handoff-restart.ts index 85dde4a990b..79038dc058b 100644 --- a/src/main/native-chat/agent-session-wire/structured-agent-session-handoff-restart.ts +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-handoff-restart.ts @@ -32,7 +32,10 @@ export async function restoreStructuredAgentSessionHandoff( const initial = input.requireRecord(sessionId) const operationId = initial.lease.handoffOperationId const initialStage = initial.lease.handoffStage - if (initialStage === 'recovering' || initialStage === 'manual-recovery') { + if ( + (initialStage === 'recovering' || initialStage === 'manual-recovery') && + !canRestoreLiveTuiOwner(initial) + ) { if (operationId) { await input.deps.store.recordOperationOutcome({ operationId, @@ -93,6 +96,10 @@ export async function restoreStructuredAgentSessionHandoff( } async function restoreOnce(input: RestartAccess, record: AgentSessionRecord): Promise { + if (canRestoreLiveTuiOwner(record)) { + await restoreRecoverableLiveTui(input, record) + return + } if (!input.deps.transport) { if (record.lease.handoffStage !== null || record.lease.runtimeKind === 'tui') { throw new Error('Agent TUI handoff recovery is unavailable on this host.') @@ -119,6 +126,45 @@ async function restoreOnce(input: RestartAccess, record: AgentSessionRecord): Pr } } +export function canRestoreLiveTuiOwner(record: AgentSessionRecord): boolean { + return ( + (record.lease.handoffStage === 'recovering' || + record.lease.handoffStage === 'manual-recovery') && + record.lease.runtimeKind === 'tui' && + record.lease.claimStatus === 'live' && + record.lease.ownerProcess !== null && + record.lease.handoffOperationId === null + ) +} + +async function restoreRecoverableLiveTui( + input: RestartAccess, + record: AgentSessionRecord +): Promise { + if (!input.deps.transport) { + throw new Error('Agent TUI handoff recovery is unavailable on this host.') + } + const recovered = await input.deps.transport.recoverTuiOwner(record) + const owner = await input.deps.transport.reproveTuiOwner({ record, owner: recovered }) + await setStoredAgentSessionHandoffStage(input.deps.store, { + sessionId: record.sessionId, + fence: record.lease.runtimeFence, + stage: null, + handoffOperationId: null, + now: input.deps.now() + }) + input.retainOwner(record.sessionId, owner) + input.setStatus(record.sessionId, { + owner: 'tui', + direction: null, + phase: 'idle', + stage: null, + operationId: null, + terminal: owner.terminal, + hostLabel: input.deps.transport.hostLabel + }) +} + async function restoreLiveTui(input: RestartAccess, record: AgentSessionRecord): Promise { const owner = await input.deps.transport!.recoverTuiOwner(record) input.retainOwner(record.sessionId, owner) diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-handoff.test.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-handoff.test.ts index 1e25b0f80e9..53b4c03d1a2 100644 --- a/src/main/native-chat/agent-session-wire/structured-agent-session-handoff.test.ts +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-handoff.test.ts @@ -480,6 +480,34 @@ describe('structured session ownership handoff', () => { expect(coordinator.status(SESSION)).toMatchObject({ owner: 'tui', phase: 'idle' }) }) + it('clears stale recovery after re-proving the live TUI owner', async () => { + await moveToTui() + const record = store.getRecord(SESSION)! + await setStoredAgentSessionHandoffStage(store, { + sessionId: SESSION, + fence: record.lease.runtimeFence, + stage: 'manual-recovery', + handoffOperationId: null, + now: NOW + }) + coordinator = createCoordinator() + + await coordinator.restore(SESSION) + expect(coordinator.status(SESSION)).toMatchObject({ + owner: 'tui', + phase: 'idle', + stage: null, + terminal: { handle: 'term-tui' } + }) + expect(reproveTuiOwner).toHaveBeenCalledOnce() + expect(store.getRecord(SESSION)?.lease).toMatchObject({ + runtimeKind: 'tui', + claimStatus: 'live', + handoffStage: null, + handoffOperationId: null + }) + }) + it('rejects cancellation once a queued handoff is no longer pending', async () => { expect( await submit(request('to-tui', 'after-turn', { action: 'cancel-queued' })) @@ -649,6 +677,11 @@ describe('structured session ownership handoff', () => { error: { recoverableOwner: 'none' } }) expect(store.getRecord(SESSION)?.lease.handoffStage).toBe('manual-recovery') + await coordinator.restore(SESSION) + expect(coordinator.status(SESSION)).toMatchObject({ + owner: 'none', + stage: 'manual-recovery' + }) expect( await submit(request('to-tui', 'now', { action: 'retry', operationId: failedOperation })) ).toMatchObject({ ok: false, refusal: { code: 'agent_session_operation_conflict' } }) 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 9131b02a306..de77ba98b48 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 @@ -1,11 +1,10 @@ import type { AgentSessionRecord } from '../../../shared/agent-session-record' import { importLegacyTranscriptIntoJournal } from '../agent-session-journal/journal-legacy-import' import { journalIdentityFor } from './structured-agent-session-attach' +import { canRestoreLiveTuiOwner } from './structured-agent-session-handoff-restart' import type { DeferredStructuredAgentSessionEventSink } from './structured-agent-session-event-sink' -import type { - StructuredAgentSessionHostDeps, - StructuredAgentSessionHostSession -} from './structured-agent-session-host' +import type { StructuredAgentSessionHostDeps } from './structured-agent-session-host' +import type { StructuredAgentSessionHostSession } from './structured-agent-session-host-types' import { StructuredAgentSessionHandoffCoordinator } from './structured-agent-session-handoff' import type { AgentSessionSubscribers } from './structured-agent-session-subscribers' @@ -20,6 +19,18 @@ type HostHandoffAccess = { export type StructuredAgentSessionHostHandoff = StructuredAgentSessionHandoffCoordinator +export async function refreshRecoverableStructuredHandoffStatus( + handoff: StructuredAgentSessionHostHandoff, + store: StructuredAgentSessionHostDeps['store'], + sessionId: string +) { + const record = store.getRecord(sessionId) + if (record && canRestoreLiveTuiOwner(record)) { + await handoff.restore(sessionId) + } + return handoff.status(sessionId) +} + export function createStructuredAgentSessionHostHandoff( deps: StructuredAgentSessionHostDeps, host: HostHandoffAccess diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-host-runtime-state.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-host-runtime-state.ts index c1df075905f..e2f74f719f3 100644 --- a/src/main/native-chat/agent-session-wire/structured-agent-session-host-runtime-state.ts +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-host-runtime-state.ts @@ -11,11 +11,15 @@ export class StructuredAgentSessionHostRuntimeState { private readonly eventSinks = new Map() private readonly leaseRenewer: StructuredAgentSessionLeaseRenewer - constructor(private readonly deps: StructuredAgentSessionHostDeps) { + constructor( + private readonly deps: StructuredAgentSessionHostDeps, + onLeaseRenewed?: (record: AgentSessionRecord) => Promise + ) { this.leaseRenewer = new StructuredAgentSessionLeaseRenewer({ store: deps.store, probe: (record) => this.probeRecord(record), now: () => deps.now?.() ?? Date.now(), + ...(onLeaseRenewed ? { onRenewed: onLeaseRenewed } : {}), onError: ({ sessionId, error }) => deps.onEventSinkError?.({ sessionId, error }) }) } diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-host-tabs.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-host-tabs.ts new file mode 100644 index 00000000000..c2947891c43 --- /dev/null +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-host-tabs.ts @@ -0,0 +1,15 @@ +export type StructuredAgentSessionTab = { + sessionId: string + workspaceId: string + agent: 'codex' +} + +export function listStructuredAgentSessionTabs( + sessions: ReadonlyMap +): StructuredAgentSessionTab[] { + return [...sessions.entries()].map(([sessionId, session]) => ({ + sessionId, + workspaceId: session.params.location.workspaceId, + agent: 'codex' + })) +} diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-host-test-data.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-host-test-data.ts new file mode 100644 index 00000000000..69e6d029d60 --- /dev/null +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-host-test-data.ts @@ -0,0 +1,63 @@ +import type { AgentJournalMessageItem } from '../../../shared/agent-session-journal-types' +import { computeAgentSessionPayloadFingerprint } from '../../../shared/agent-session-mutation-envelope' +import type { AgentSessionExecutionLocation } from '../../../shared/agent-session-record' +import { attachFingerprintFields } from './structured-agent-session-attach' +import type { AgentSessionAttachParams } from './structured-agent-session-attach' + +export const HOST_TEST_NOW = 1_800_000_000_000 +export const HOST_TEST_SESSION = 'session-alpha' +export const HOST_TEST_THREAD = '019fd532-7c11-7a90-b6de-4e1a2c3d5f60' + +export const HOST_TEST_LOCATION: AgentSessionExecutionLocation = { + executionHostId: 'local', + wslDistro: null, + workspaceId: 'workspace-1', + workspaceKind: 'git-worktree' +} + +let operations = 0 + +export function resetHostTestOperationIds(): void { + operations = 0 +} + +export function hostTestOperationId(): string { + operations += 1 + return `${HOST_TEST_NOW}-${operations.toString(16).padStart(32, '0')}` +} + +export function hostTestMessage(text: string): AgentJournalMessageItem { + return { kind: 'message', role: 'user', blocks: [{ type: 'text', text }] } +} + +export function hostTestAttachParams( + expectedRuntimeFence: number | null, + overrides: Partial = {} +): AgentSessionAttachParams { + const params: AgentSessionAttachParams = { + envelope: { + sessionId: HOST_TEST_SESSION, + clientOperationId: hostTestOperationId(), + expectedRuntimeFence, + payloadFingerprint: '0'.repeat(64) + }, + location: HOST_TEST_LOCATION, + provider: 'codex', + agent: 'codex', + accountHome: { variable: 'CODEX_HOME', path: '/home/dev/.codex' }, + runtimeKind: 'native', + providerHandle: { kind: 'codex', threadId: HOST_TEST_THREAD }, + ...overrides + } + return { + ...params, + envelope: { + ...params.envelope, + payloadFingerprint: computeAgentSessionPayloadFingerprint({ + method: 'agentSession.attach', + sessionId: params.envelope.sessionId, + fields: attachFingerprintFields(params) + }) + } + } +} diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-host-types.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-host-types.ts new file mode 100644 index 00000000000..e0aa61067c7 --- /dev/null +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-host-types.ts @@ -0,0 +1,10 @@ +import type { AgentSessionJournal } from '../agent-session-journal/journal-store' +import type { AgentSessionAttachParams } from './structured-agent-session-attach' + +export type StructuredAgentSessionCaller = { callerKey: string } + +export type StructuredAgentSessionHostSession = { + journal: AgentSessionJournal + params: AgentSessionAttachParams + fence: number +} diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-host.test.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-host.test.ts index 59321c2ecf5..39206eeace0 100644 --- a/src/main/native-chat/agent-session-wire/structured-agent-session-host.test.ts +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-host.test.ts @@ -1,19 +1,11 @@ -// The host against a real record store and a real journal, with only the -// provider adapter stubbed — admission, idempotency, and the journal writes are -// exactly the ones that ship. - import { mkdtemp, rm } from 'node:fs/promises' import { tmpdir } from 'node:os' import { join } from 'node:path' import { afterEach, beforeEach, describe, expect, it, vi, type Mock } from 'vitest' import { agentJournalItemKey } from '../../../shared/agent-session-journal-item-key' -import type { AgentJournalMessageItem } from '../../../shared/agent-session-journal-types' import type { AgentSessionOwnerProbe } from '../../../shared/agent-session-lease-adjudication' import { computeAgentSessionPayloadFingerprint } from '../../../shared/agent-session-mutation-envelope' -import type { - AgentSessionExecutionLocation, - AgentSessionRecord -} from '../../../shared/agent-session-record' +import type { AgentSessionRecord } from '../../../shared/agent-session-record' import type { AgentSessionMutationEnvelope, AgentSessionSubscribeEvent @@ -28,35 +20,20 @@ import type { AgentSessionDispatchOutcome, StructuredAgentSessionAdapter } from './structured-agent-session-adapter' -import { attachFingerprintFields } from './structured-agent-session-attach' import type { AgentSessionAttachParams } from './structured-agent-session-attach' import { StructuredAgentSessionHost } from './structured-agent-session-host' - -const NOW = 1_800_000_000_000 -const SESSION = 'session-alpha' -const THREAD = '019fd532-7c11-7a90-b6de-4e1a2c3d5f60' - -const LOCATION: AgentSessionExecutionLocation = { - executionHostId: 'local', - wslDistro: null, - workspaceId: 'workspace-1', - workspaceKind: 'git-worktree' -} +import { + HOST_TEST_NOW as NOW, + HOST_TEST_SESSION as SESSION, + HOST_TEST_THREAD as THREAD, + hostTestAttachParams, + hostTestMessage, + hostTestOperationId, + resetHostTestOperationIds +} from './structured-agent-session-host-test-data' const CALLER = { callerKey: 'client-1' } -let operations = 0 - -/** `<13-digit ms>-<32 hex>`, the only shape the durable ledger accepts. */ -function operationId(): string { - operations += 1 - return `${NOW}-${operations.toString(16).padStart(32, '0')}` -} - -function message(text: string): AgentJournalMessageItem { - return { kind: 'message', role: 'user', blocks: [{ type: 'text', text }] } -} - function envelope( method: string, fields: Record, @@ -64,9 +41,7 @@ function envelope( ): AgentSessionMutationEnvelope { return { sessionId: SESSION, - clientOperationId: operationId(), - // A mutation names the fence it believes it is writing under; only attach - // may leave it null. + clientOperationId: hostTestOperationId(), expectedRuntimeFence: store.getRecord(SESSION)?.lease.runtimeFence ?? 1, payloadFingerprint: computeAgentSessionPayloadFingerprint({ method, @@ -78,44 +53,11 @@ function envelope( } function attachParams(overrides: Partial = {}): AgentSessionAttachParams { - const params: AgentSessionAttachParams = { - envelope: { - sessionId: SESSION, - clientOperationId: operationId(), - expectedRuntimeFence: null, - payloadFingerprint: '0'.repeat(64) - }, - location: LOCATION, - provider: 'codex', - agent: 'codex', - accountHome: { variable: 'CODEX_HOME', path: '/home/dev/.codex' }, - runtimeKind: 'native', - providerHandle: { kind: 'codex', threadId: THREAD }, - ...overrides - } - return { - ...params, - envelope: { - ...params.envelope, - payloadFingerprint: computeAgentSessionPayloadFingerprint({ - method: 'agentSession.attach', - sessionId: params.envelope.sessionId, - fields: attachFingerprintFields(params) - }) - } - } + return hostTestAttachParams(null, overrides) } -/** `ensure`: the same transition as create, from a fence the client already holds. */ function ensureParams(fence: number): AgentSessionAttachParams { - return attachParams({ - envelope: { - sessionId: SESSION, - clientOperationId: operationId(), - expectedRuntimeFence: fence, - payloadFingerprint: '0'.repeat(64) - } - }) + return hostTestAttachParams(fence) } let root: string @@ -185,7 +127,7 @@ async function seedApproval(optionId = 'allow'): Promise<{ itemId: string; revis beforeEach(async () => { root = await mkdtemp(join(tmpdir(), 'orca-wire-host-')) - operations = 0 + resetHostTestOperationIds() ordinal = 0 acquire = vi.fn(async ({ fence }) => ({ process: { @@ -197,8 +139,6 @@ beforeEach(async () => { link: { linkId: `link-${fence}`, handle: { provider: 'codex', threadId: THREAD }, - // A restarted host re-proves the thread it inherited; only the first - // owner of a session may claim to have created it. origin: store.getRecord(SESSION)?.providerHandleChain.length ? 'resumed' : 'created', mintedAtFence: fence, observedAt: NOW @@ -409,7 +349,7 @@ describe('attach', () => { describe('send', () => { it('writes the submission before dispatching and resolves it accepted', async () => { await attach() - const body = message('add a retry') + const body = hostTestMessage('add a retry') const result = await host.send(CALLER, { envelope: envelope('agentSession.send', { body }), body @@ -427,7 +367,7 @@ describe('send', () => { it('settles a thrown dispatch as unknown, never as a rejection', async () => { await attach() dispatch.mockRejectedValueOnce(new Error('socket closed')) - const body = message('add a retry') + const body = hostTestMessage('add a retry') const result = await host.send(CALLER, { envelope: envelope('agentSession.send', { body }), body @@ -437,7 +377,7 @@ describe('send', () => { it('replays a retried send from the journal without dispatching twice', async () => { await attach() - const body = message('add a retry') + const body = hostTestMessage('add a retry') const params = { envelope: envelope('agentSession.send', { body }), body } await host.send(CALLER, params) const retry = await host.send(CALLER, params) @@ -450,7 +390,7 @@ describe('send', () => { dispatch .mockRejectedValueOnce(new Error('socket closed')) .mockImplementationOnce(async () => accepted()) - const body = message('possibly delivered') + const body = hostTestMessage('possibly delivered') const params = { envelope: envelope('agentSession.send', { body }), body } const first = await host.send(CALLER, params) @@ -472,7 +412,7 @@ describe('send', () => { it('refuses a stale fence and hands back the current one', async () => { const record = await attach() - const body = message('add a retry') + const body = hostTestMessage('add a retry') const result = await host.send(CALLER, { envelope: envelope( 'agentSession.send', @@ -489,7 +429,7 @@ describe('send', () => { it('does not let a refused call leave a ledger row that replays past the fence', async () => { const record = await attach() - const body = message('add a retry') + const body = hostTestMessage('add a retry') const params = { envelope: envelope( 'agentSession.send', @@ -507,7 +447,7 @@ describe('send', () => { }) it('refuses any mutation against a session this host has not attached', async () => { - const body = message('add a retry') + const body = hostTestMessage('add a retry') expect( await host.send(CALLER, { envelope: envelope('agentSession.send', { body }), body }) ).toMatchObject({ ok: false, refusal: { code: 'agent_session_ownership_unknown' } }) @@ -738,7 +678,7 @@ describe('restart', () => { it('restores durable journals for read-only history without acquiring a provider', async () => { await attach() - const body = message('persisted conversation') + const body = hostTestMessage('persisted conversation') await host.send(CALLER, { envelope: envelope('agentSession.send', { body }), body }) await reboot(async () => ({ outcome: 'indeterminate', reason: 'read does not need ownership' })) acquire.mockClear() @@ -753,13 +693,40 @@ describe('restart', () => { expect(acquire).not.toHaveBeenCalled() }) + it('clears stale TUI recovery when restart successfully reacquires the native owner', async () => { + await attach() + await store.transitionHandoff(SESSION, (record) => ({ + ...record, + lease: { + ...record.lease, + runtimeKind: 'tui', + handoffStage: 'manual-recovery' + } + })) + await reboot(async () => ({ outcome: 'pid-absent' })) + acquire.mockClear() + + await host.restoreReadableSessions() + + expect(acquire).toHaveBeenCalledOnce() + expect(store.getRecord(SESSION)?.lease).toMatchObject({ + runtimeKind: 'native', + claimStatus: 'live', + handoffStage: null, + handoffOperationId: null + }) + await expect(host.handoffStatus(SESSION)).resolves.toMatchObject({ + owner: 'native', + phase: 'idle', + stage: null + }) + }) + it("keeps a session whose owner cannot be probed out of a live writer's hands", async () => { await attach() const held = store.getRecord(SESSION)?.lease.runtimeFence ?? 0 await reboot(async () => ({ outcome: 'indeterminate', reason: 'no probe on this host' })) - // Nothing proved the previous process dead, so the lease routes to manual - // recovery rather than letting a second writer in behind it. expect(await host.attach(CALLER, ensureParams(held))).toMatchObject({ ok: false, refusal: { code: 'agent_session_ownership_unknown' } @@ -776,7 +743,6 @@ describe('restart', () => { await reboot(probe) await expect(host.attach(CALLER, ensureParams(held))).rejects.toThrow('probe exploded') - // One unlucky startup must not strand the session for this host's lifetime. const reattached = await host.attach(CALLER, ensureParams(await staleFenceFrom(held))) expect(reattached).toMatchObject({ ok: true }) expect(probe).toHaveBeenCalledTimes(2) @@ -792,7 +758,7 @@ describe('subscribe', () => { sessionId: SESSION, emit: (event) => events.push(event) }) - const body = message('add a retry') + const body = hostTestMessage('add a retry') await host.send(CALLER, { envelope: envelope('agentSession.send', { body }), body }) expect(events[0]?.type).toBe('snapshot') @@ -807,7 +773,7 @@ describe('subscribe', () => { it('resumes from a client cursor with only the rows it missed', async () => { await attach() - const body = message('add a retry') + const body = hostTestMessage('add a retry') const first = await host.send(CALLER, { envelope: envelope('agentSession.send', { body }), body @@ -823,10 +789,9 @@ describe('subscribe', () => { emit: (event) => events.push(event), cursor: first.cursor }) - // Nothing new since that cursor: a resume is silent, not a redundant replay. - expect(events).toHaveLength(0) + expect(events[0]).toMatchObject({ type: 'batch', handoff: { owner: 'native', phase: 'idle' } }) - const second = message('and a timeout') + const second = hostTestMessage('and a timeout') await host.send(CALLER, { envelope: envelope('agentSession.send', { body: second }), body: second @@ -850,7 +815,7 @@ describe('subscribe', () => { sessionId: SESSION, emit: (event) => events.push(event) }) - const body = message('survive subscriber failure') + const body = hostTestMessage('survive subscriber failure') const result = await host.send(CALLER, { envelope: envelope('agentSession.send', { body }), 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 bd8223bf8f5..cbf8408f33e 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 @@ -1,11 +1,5 @@ -// The structured agent-session host: one place where the lease, the journal, -// and the provider adapter meet. -// -// Every mutating call takes the same route — recompute the fingerprint, admit -// through the durable operation ledger, check the lease, then act — so no -// method can grow its own admission rules. Calls against one session are -// serialized; the journal's own queue orders writes, but admission has to be -// single-file too or two clients could both read one prompt as pending. +// Structured agent-session host: where the lease, journal, and provider adapter meet. +// Mutations share one durable admission path and serialize per session. import { randomUUID } from 'node:crypto' import type { AgentJournalMessageItem } from '../../../shared/agent-session-journal-types' @@ -31,7 +25,6 @@ import type { AgentSessionWireRefusal } from '../../../shared/agent-session-wire' import type { AgentSessionRecordStore } from '../../runtime/agent-session-record-store' -import type { AgentSessionJournal } from '../agent-session-journal/journal-store' import { readAgentSessionHistory } from './agent-session-history-page' import type { StructuredAgentSessionAdapter } from './structured-agent-session-adapter' import type { AgentSessionAttachParams } from './structured-agent-session-attach' @@ -58,38 +51,28 @@ import { restoreStructuredAgentSessionsOnRestart } from './structured-agent-sess import type { StructuredAgentSessionHandoffTransport } from './structured-agent-session-handoff-types' import { createStructuredAgentSessionHostHandoff, + refreshRecoverableStructuredHandoffStatus, type StructuredAgentSessionHostHandoff } from './structured-agent-session-host-handoff' import { StructuredAgentSessionHostRuntimeState } from './structured-agent-session-host-runtime-state' - -export type StructuredAgentSessionCaller = { - /** Stable per-client identity; scopes the operation ledger and records who - * answered a prompt. */ - callerKey: string -} +import { listStructuredAgentSessionTabs } from './structured-agent-session-host-tabs' +import type { + StructuredAgentSessionCaller, + StructuredAgentSessionHostSession +} from './structured-agent-session-host-types' export type StructuredAgentSessionHostDeps = { store: AgentSessionRecordStore adapter: StructuredAgentSessionAdapter journalRoot: string - /** Key id the host's claims are minted under. */ claimKeyId: string probeOwner?: (record: AgentSessionRecord) => Promise mintSpawnToken?: () => string now?: () => number - /** A journal append the provider streamed that could not be written. Unset - * drops it: the alternative is throwing inside the provider's notification - * callback, which would take the connection down over one lost row. */ onEventSinkError?: (input: { sessionId: string; error: unknown }) => void handoffTransport?: StructuredAgentSessionHandoffTransport } -export type StructuredAgentSessionHostSession = { - journal: AgentSessionJournal - params: AgentSessionAttachParams - fence: number -} - export class StructuredAgentSessionHost { private readonly sessions = new Map() private readonly subscribers = new AgentSessionSubscribers() @@ -99,7 +82,9 @@ export class StructuredAgentSessionHost { private readonly handoffs: StructuredAgentSessionHostHandoff constructor(readonly deps: StructuredAgentSessionHostDeps) { - this.runtimeState = new StructuredAgentSessionHostRuntimeState(deps) + this.runtimeState = new StructuredAgentSessionHostRuntimeState(deps, (record) => + this.restoreRenewedHandoff(record.sessionId) + ) this.reconcileLeases = createRestartReconciler({ store: deps.store, probe: (record) => this.runtimeState.probeRecord(record), @@ -118,24 +103,14 @@ export class StructuredAgentSessionHost { private now = (): number => this.deps.now?.() ?? Date.now() - hasSession(sessionId: string): boolean { - return this.sessions.has(sessionId) - } + hasSession = (sessionId: string): boolean => this.sessions.has(sessionId) supportsCreate(location: AgentSessionExecutionLocation, agent: string): boolean { return agent === 'codex' && (this.deps.adapter.supportsLocation?.(location) ?? false) } - listSessionTabs(): { - sessionId: string - workspaceId: string - agent: 'codex' - }[] { - return [...this.sessions.entries()].map(([sessionId, session]) => ({ - sessionId, - workspaceId: session.params.location.workspaceId, - agent: 'codex' as const - })) + listSessionTabs() { + return listStructuredAgentSessionTabs(this.sessions) } async restoreReadableSessions(): Promise { @@ -161,10 +136,14 @@ export class StructuredAgentSessionHost { return this.tasks.serialize(sessionId, task) } - /** - * `create` and `ensure` in one: a null expected fence means the session must - * not exist yet, any other value is a compare-and-swap against the lease. - */ + private restoreRenewedHandoff(sessionId: string): Promise { + return this.serialize(sessionId, async () => { + if (this.sessions.has(sessionId)) { + await refreshRecoverableStructuredHandoffStatus(this.handoffs, this.deps.store, sessionId) + } + }) + } + attach( caller: StructuredAgentSessionCaller, params: AgentSessionAttachParams @@ -206,9 +185,6 @@ export class StructuredAgentSessionHost { } else { this.subscribers.publish(sessionId, attached.journal) } - // Bound last: everything the provider streamed while the journal was - // opening lands after the attach result the client is about to read, - // never interleaved with it. eventSink.bind({ journal: attached.journal, fence, @@ -216,8 +192,6 @@ export class StructuredAgentSessionHost { }) } }) - // A refused attach that never reached a journal leaves the sink holding - // writes nothing will ever drain. if (!attached.ok && !this.sessions.has(sessionId)) { eventSink.close() this.runtimeState.discardEventSink(sessionId) @@ -227,16 +201,12 @@ export class StructuredAgentSessionHost { return this.tasks.trackAttach(attaching) } - /** Settles every row the provider streamed. Callers that must read the - * journal immediately after provider activity await this first. */ flushStreamedEvents(sessionId: string): Promise { return this.runtimeState.flushEventSink(sessionId) } - /** Settles final provider rows after every child has stopped producing them. */ async flushAllStreamedEvents(): Promise { this.runtimeState.stopLeaseRenewal() - // An acquired provider can still be opening its journal; bind before draining. await this.tasks.drainAttaches() await this.runtimeState.flushAllEventSinks() } @@ -317,16 +287,13 @@ export class StructuredAgentSessionHost { return this.handoffs.request(caller.callerKey, params) } - handoffStatus(sessionId: string): AgentSessionHandoffStatus { + async handoffStatus(sessionId: string): Promise { this.requireSession(sessionId) - return this.handoffs.status(sessionId) + return this.serialize(sessionId, () => + refreshRecoverableStructuredHandoffStatus(this.handoffs, this.deps.store, sessionId) + ) } - /** - * Paged read. Throws `agent_session_ownership_unknown` for a session this host - * has not attached — a reset would tell the client to reload something that - * does not exist here. - */ history(request: AgentSessionHistoryRequest): AgentSessionHistoryResult { const result = readAgentSessionHistory(this.requireSession(request.sessionId).journal, request) const fence = this.deps.store.getRecord(request.sessionId)?.lease.runtimeFence diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-lease-renewer.test.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-lease-renewer.test.ts index e0dd07c344c..17bfff24c8b 100644 --- a/src/main/native-chat/agent-session-wire/structured-agent-session-lease-renewer.test.ts +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-lease-renewer.test.ts @@ -97,16 +97,21 @@ describe('structured agent-session lease renewal', () => { outcome: 'identity-matched' as const, matchedOn: ['process-start-time' as const] })) + const onRenewed = vi.fn() const renewer = new StructuredAgentSessionLeaseRenewer({ store, probe, - now: () => NOW + 10_000 + now: () => NOW + 10_000, + onRenewed }) await renewer.renewNow() expect(probe).toHaveBeenCalledOnce() expect(store.getRecord('session-renewal')?.lease.lastRenewedAt).toBe(NOW + 10_000) + expect(onRenewed).toHaveBeenCalledWith( + expect.objectContaining({ sessionId: 'session-renewal' }) + ) }) it('stops extending the lease when child proof is no longer sufficient', async () => { diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-lease-renewer.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-lease-renewer.ts index f9f9b7b6773..dcec74d1d1a 100644 --- a/src/main/native-chat/agent-session-wire/structured-agent-session-lease-renewer.ts +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-lease-renewer.ts @@ -16,6 +16,7 @@ export class StructuredAgentSessionLeaseRenewer { store: AgentSessionRecordStore probe: (record: AgentSessionRecord) => Promise now: () => number + onRenewed?: (record: AgentSessionRecord) => Promise onError?: (input: { sessionId: string; error: unknown }) => void intervalMs?: number } @@ -60,12 +61,13 @@ export class StructuredAgentSessionLeaseRenewer { private async renewRecord(record: AgentSessionRecord): Promise { try { - await this.input.store.renewLease({ + const renewed = await this.input.store.renewLease({ sessionId: record.sessionId, fence: record.lease.runtimeFence, childProbe: await this.input.probe(record), now: this.input.now() }) + await this.input.onRenewed?.(renewed) } catch (error) { this.input.onError?.({ sessionId: record.sessionId, error }) } diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-read-restore.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-read-restore.ts index 82d3116c9ad..01a2578bcbf 100644 --- a/src/main/native-chat/agent-session-wire/structured-agent-session-read-restore.ts +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-read-restore.ts @@ -1,4 +1,7 @@ -import type { AgentSessionRecord } from '../../../shared/agent-session-record' +import type { + AgentSessionOwnerRuntimeKind, + AgentSessionRecord +} from '../../../shared/agent-session-record' import type { AgentSessionRecordStore } from '../../runtime/agent-session-record-store' import { loadJournal } from '../agent-session-journal/journal-open' import { journalDirectoryFor } from '../agent-session-journal/journal-paths' @@ -49,7 +52,11 @@ export async function restoreStructuredAgentSessionRead( export function attachParamsForRecord( record: AgentSessionRecord, - input: { clientOperationId: string; expectedRuntimeFence: number } + input: { + clientOperationId: string + expectedRuntimeFence: number + runtimeKind?: AgentSessionOwnerRuntimeKind + } ): AgentSessionAttachParams { const params: AgentSessionAttachParams = { envelope: { @@ -62,7 +69,7 @@ export function attachParamsForRecord( provider: 'codex', agent: 'codex', accountHome: record.accountHome, - runtimeKind: record.lease.runtimeKind + runtimeKind: input.runtimeKind ?? record.lease.runtimeKind } return { ...params, 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 082f79ca194..9a45c53f5bf 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 @@ -31,7 +31,8 @@ export async function restoreStructuredAgentSessionsOnRestart(input: { (await input.resume( attachParamsForRecord(current, { clientOperationId: input.operationId(), - expectedRuntimeFence: current.lease.runtimeFence + expectedRuntimeFence: current.lease.runtimeFence, + runtimeKind: 'native' }) )) ) { 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 9785ac91a34..a5262de99b8 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 @@ -156,6 +156,20 @@ export class AgentSessionSubscribers { return } if (since.rows.length === 0) { + if (handoff) { + this.emit(subscriber, { + type: 'batch', + sessionId: subscriber.sessionId, + batch: { + cursor: snapshot.cursor, + items: [], + removedItemIds: [], + submissions: [] + }, + fence: subscriber.fence, + handoff + }) + } return } const projected = projectJournalBatch({ diff --git a/src/main/runtime/rpc/methods/structured-agent-session.ts b/src/main/runtime/rpc/methods/structured-agent-session.ts index 3501f9a0a65..1225eeb7ea2 100644 --- a/src/main/runtime/rpc/methods/structured-agent-session.ts +++ b/src/main/runtime/rpc/methods/structured-agent-session.ts @@ -12,10 +12,8 @@ import { computeAgentSessionPayloadFingerprint } from '../../../../shared/agent-session-mutation-envelope' import { getStructuredAgentSessionHost } from '../../../native-chat/agent-session-wire/structured-agent-session-registry' -import type { - StructuredAgentSessionCaller, - StructuredAgentSessionHost -} from '../../../native-chat/agent-session-wire/structured-agent-session-host' +import type { StructuredAgentSessionHost } from '../../../native-chat/agent-session-wire/structured-agent-session-host' +import type { StructuredAgentSessionCaller } from '../../../native-chat/agent-session-wire/structured-agent-session-host-types' import { defineMethod, defineStreamingMethod, type RpcAnyMethod, type RpcContext } from '../core' import { AttachParams, diff --git a/src/shared/structured-agent-session-coalescer.ts b/src/shared/structured-agent-session-coalescer.ts index e1af3549097..51bc7fa0537 100644 --- a/src/shared/structured-agent-session-coalescer.ts +++ b/src/shared/structured-agent-session-coalescer.ts @@ -31,7 +31,11 @@ function mergeBatch( items: [...items.values()], removedItemIds: [...new Set([...left.batch.removedItemIds, ...right.batch.removedItemIds])], submissions: [...submissions.values()] - } + }, + ...(right.fence !== undefined || left.fence !== undefined + ? { fence: right.fence ?? left.fence } + : {}), + ...(right.handoff || left.handoff ? { handoff: right.handoff ?? left.handoff } : {}) } }