From dcd26ef6775508b8db869ae5ce7645cb29d1b7c3 Mon Sep 17 00:00:00 2001 From: Brennan Benson <79079362+brennanb2025@users.noreply.github.com> Date: Sun, 9 Aug 2026 21:35:24 -0700 Subject: [PATCH] Close remaining structured mobile RC blockers --- .../mobile-structured-mutation-envelope.ts | 1 + ...use-mobile-structured-session-mutations.ts | 45 +++++++--- ...e-mobile-structured-session-writes.test.ts | 42 +++++++++- .../mobile-relay-rpc-streams.test.ts | 20 +++++ .../src/transport/mobile-relay-rpc-streams.ts | 35 ++++++-- .../rpc-client-structured-session.test.ts | 18 ++++ .../rpc-client-terminal-subscription.ts | 14 +++- mobile/src/transport/rpc-client.ts | 2 +- .../codex-structured-session-adapter.test.ts | 40 ++++++++- .../codex/codex-structured-session-adapter.ts | 66 +++++++-------- .../codex/codex-structured-session-close.ts | 24 ++++++ .../codex/codex-structured-session-state.ts | 11 +++ .../structured-agent-session-attach-flow.ts | 7 +- .../structured-agent-session-host.test.ts | 63 +++++++++++++- .../structured-agent-session-host.ts | 10 ++- ...ctured-agent-session-mutation-admission.ts | 3 +- ...structured-agent-session-mutation-plans.ts | 17 +++- ...structured-agent-session-replay-outcome.ts | 4 + .../structured-agent-session-turns.ts | 16 +++- .../structured-agent-session-schemas.ts | 1 + .../runtime/rpc/mobile-socket-wiring.test.ts | 82 +++++++++++++++++++ src/main/runtime/rpc/mobile-socket-wiring.ts | 4 +- ...ructured-agent-session-integration.test.ts | 2 - 23 files changed, 452 insertions(+), 75 deletions(-) create mode 100644 src/main/codex/codex-structured-session-close.ts diff --git a/mobile/src/session/mobile-structured-mutation-envelope.ts b/mobile/src/session/mobile-structured-mutation-envelope.ts index 7dad042ef61..3ed66491c2c 100644 --- a/mobile/src/session/mobile-structured-mutation-envelope.ts +++ b/mobile/src/session/mobile-structured-mutation-envelope.ts @@ -51,6 +51,7 @@ export function mobileStructuredSendRequest( fields }) }, + ...(entry.retryAfterUnknownSubmittedAt !== null ? { retryUnknown: true } : {}), ...fields } } diff --git a/mobile/src/session/use-mobile-structured-session-mutations.ts b/mobile/src/session/use-mobile-structured-session-mutations.ts index 3c22f852d74..50fa3b03712 100644 --- a/mobile/src/session/use-mobile-structured-session-mutations.ts +++ b/mobile/src/session/use-mobile-structured-session-mutations.ts @@ -34,13 +34,19 @@ export function useMobileStructuredSessionMutations(args: { onRefusal: (message: string | null) => void }): MobileStructuredSessionMutations { const operationIdsRef = useRef(new Map()) + const activeContextRef = useRef({ + client: args.client, + fence: args.fence, + sessionId: args.sessionId + }) + activeContextRef.current = { client: args.client, fence: args.fence, sessionId: args.sessionId } const { client, fence, onRefusal, sessionId } = args const mutate = useCallback( async (method: string, fingerprintMethod: string, fields: Record) => { if (!client || !sessionId || fence === null || client.getState() !== 'connected') { return null } - const mutationKey = `${fingerprintMethod}:${JSON.stringify(fields)}` + const mutationKey = `${sessionId}:${fingerprintMethod}:${JSON.stringify(fields)}` const clientOperationId = operationIdsRef.current.get(mutationKey) ?? createMobileStructuredOperationId('mobile-mutation', () => ExpoCrypto.randomUUID()) @@ -61,26 +67,34 @@ export function useMobileStructuredSessionMutations(args: { ...fields }) } catch (error) { - onRefusal( - isRpcDeliveryUnknown(error) || isLogicalClientCutoverError(error) - ? 'Response delivery unconfirmed' - : error instanceof Error - ? error.message - : 'Request was not sent' - ) + if (matchesActiveContext(activeContextRef.current, client, sessionId, fence)) { + onRefusal( + isRpcDeliveryUnknown(error) || isLogicalClientCutoverError(error) + ? 'Response delivery unconfirmed' + : error instanceof Error + ? error.message + : 'Request was not sent' + ) + } return null } if (!response.ok) { - onRefusal(response.error.message) + if (matchesActiveContext(activeContextRef.current, client, sessionId, fence)) { + onRefusal(response.error.message) + } return null } const result = response.result as AgentSessionMutationResult if (!result.ok) { - onRefusal(result.refusal.message) + if (matchesActiveContext(activeContextRef.current, client, sessionId, fence)) { + onRefusal(result.refusal.message) + } return null } operationIdsRef.current.delete(mutationKey) - onRefusal(null) + if (matchesActiveContext(activeContextRef.current, client, sessionId, fence)) { + onRefusal(null) + } return result.value }, [client, fence, onRefusal, sessionId] @@ -106,3 +120,12 @@ export function useMobileStructuredSessionMutations(args: { [mutate] ) } + +function matchesActiveContext( + active: { client: RpcClient | null; sessionId: string | null; fence: number | null }, + client: RpcClient, + sessionId: string, + fence: number +): boolean { + return active.client === client && active.sessionId === sessionId && active.fence === fence +} diff --git a/mobile/src/session/use-mobile-structured-session-writes.test.ts b/mobile/src/session/use-mobile-structured-session-writes.test.ts index 99af37e6d5b..eef091fbb72 100644 --- a/mobile/src/session/use-mobile-structured-session-writes.test.ts +++ b/mobile/src/session/use-mobile-structured-session-writes.test.ts @@ -65,13 +65,14 @@ describe('useMobileStructuredSessionWrites', () => { } as RpcClient let connected = true let fence = 3 + let sessionId = 'mobile_1' let submissions: AgentJournalSubmission[] = [] function Probe(): null { api = useMobileStructuredSessionWrites({ client, connected, - sessionId: 'mobile_1', + sessionId, fence, submissions }) @@ -85,6 +86,7 @@ describe('useMobileStructuredSessionWrites', () => { sendRequest.mockReset() connected = true fence = 3 + sessionId = 'mobile_1' submissions = [] const crypto = await import('expo-crypto') vi.mocked(crypto.randomUUID) @@ -187,6 +189,7 @@ describe('useMobileStructuredSessionWrites', () => { const retryId = (sendRequest.mock.calls[1]![1] as { envelope: { clientOperationId: string } }) .envelope.clientOperationId expect(retryId).toBe(firstId) + expect(sendRequest.mock.calls[1]![1]).toMatchObject({ retryUnknown: true }) await act(async () => { await Promise.resolve() await Promise.resolve() @@ -372,4 +375,41 @@ describe('useMobileStructuredSessionWrites', () => { expect(applied).toBe(false) expect(api!.error).toBe('not allowed') }) + + it('isolates mutation ids and late errors across session changes', async () => { + const first = deferred() + sendRequest.mockImplementationOnce(() => first.promise).mockResolvedValueOnce(accepted('b')) + + let mutationA!: Promise + act(() => { + mutationA = api!.setOption('model', 'gpt-5.6-sol') + }) + const operationA = ( + sendRequest.mock.calls[0]![1] as { envelope: { clientOperationId: string } } + ).envelope.clientOperationId + + sessionId = 'mobile_2' + await act(async () => { + renderer!.update(createElement(Probe)) + await Promise.resolve() + }) + await act(async () => { + await api!.setOption('model', 'gpt-5.6-sol') + }) + const operationB = ( + sendRequest.mock.calls[1]![1] as { envelope: { clientOperationId: string } } + ).envelope.clientOperationId + expect(operationB).not.toBe(operationA) + + await act(async () => { + first.resolve({ + id: 'a', + ok: false, + _meta: { runtimeId: 'runtime-1' }, + error: { code: 'failed', message: 'late session A failure' } + }) + await mutationA + }) + expect(api!.error).toBeNull() + }) }) diff --git a/mobile/src/transport/mobile-relay-rpc-streams.test.ts b/mobile/src/transport/mobile-relay-rpc-streams.test.ts index b0e59c6bf5b..cd27b458976 100644 --- a/mobile/src/transport/mobile-relay-rpc-streams.test.ts +++ b/mobile/src/transport/mobile-relay-rpc-streams.test.ts @@ -12,6 +12,26 @@ function rpcFailure(id: string): RpcFailure { } describe('MobileRelayRpcStreams failure parity', () => { + it('unsubscribes an agent session without waiting for nonexistent ready metadata', async () => { + let nextId = 0 + const sendFrame = vi.fn(() => true) + const streams = new MobileRelayRpcStreams({ + nextId: () => `stream-${++nextId}`, + sendFrame, + waitForConnected: async () => {} + }) + const cancel = streams.subscribe('agentSession.subscribe', { sessionId: 'session-a' }, vi.fn()) + await Promise.resolve() + + cancel() + + expect(sendFrame).toHaveBeenLastCalledWith({ + id: 'stream-2', + method: 'agentSession.unsubscribe', + params: { sessionId: 'session-a', subscriptionId: 'stream-1' } + }) + }) + it('emits an RPC failure exactly once before removing the stream', async () => { const listener = vi.fn() const sendFrame = vi.fn(() => true) diff --git a/mobile/src/transport/mobile-relay-rpc-streams.ts b/mobile/src/transport/mobile-relay-rpc-streams.ts index 7c485b27e02..76bdd82d6aa 100644 --- a/mobile/src/transport/mobile-relay-rpc-streams.ts +++ b/mobile/src/transport/mobile-relay-rpc-streams.ts @@ -4,6 +4,7 @@ import { type TerminalSnapshotState } from './rpc-client-terminal-binary-frame' import { + buildStreamUnsubscribe, buildTerminalUnsubscribeParams, updateTerminalSubscriptionViewport } from './rpc-client-terminal-subscription' @@ -22,6 +23,7 @@ type StreamRecord = { streamIds: Set subscriptionId?: string cancelled: boolean + sent: boolean } type StreamManagerOptions = { @@ -51,15 +53,21 @@ export class MobileRelayRpcStreams { listener, onBinaryFrame: subscribeOptions?.onBinaryFrame, streamIds: new Set(), - cancelled: false + cancelled: false, + sent: false } this.streams.set(id, stream) void this.options .waitForConnected() .then(() => { - if (!stream.cancelled && !this.options.sendFrame({ id, method, params: stream.params })) { - this.fail(id, stream, 'Connection interrupted') + if (stream.cancelled) { + return } + if (!this.options.sendFrame({ id, method, params: stream.params })) { + this.fail(id, stream, 'Connection interrupted') + return + } + stream.sent = true }) .catch((error: unknown) => { const message = error instanceof Error ? error.message : 'Connection interrupted' @@ -135,6 +143,10 @@ export class MobileRelayRpcStreams { return } stream.cancelled = true + if (!stream.sent) { + this.remove(id) + return + } if (stream.method === 'terminal.subscribe') { const params = buildTerminalUnsubscribeParams(stream.params) if (params) { @@ -144,12 +156,17 @@ export class MobileRelayRpcStreams { params }) } - } else if (stream.subscriptionId) { - this.options.sendFrame({ - id: this.options.nextId(), - method: stream.method.replace(/\.subscribe$/, '.unsubscribe'), - params: { subscriptionId: stream.subscriptionId } - }) + } else { + const unsubscribe = buildStreamUnsubscribe(stream.method, stream.params, id) + if (unsubscribe) { + this.options.sendFrame({ id: this.options.nextId(), ...unsubscribe }) + } else if (stream.subscriptionId) { + this.options.sendFrame({ + id: this.options.nextId(), + method: stream.method.replace(/\.subscribe$/, '.unsubscribe'), + params: { subscriptionId: stream.subscriptionId } + }) + } } this.remove(id) } diff --git a/mobile/src/transport/rpc-client-structured-session.test.ts b/mobile/src/transport/rpc-client-structured-session.test.ts index 1fe9b084ff7..a517548ac1b 100644 --- a/mobile/src/transport/rpc-client-structured-session.test.ts +++ b/mobile/src/transport/rpc-client-structured-session.test.ts @@ -81,4 +81,22 @@ describe('structured session RPC transport', () => { }) client.close() }) + + it('unsubscribes a structured stream by its request-scoped subscription id', () => { + const client = connect('ws://desktop.invalid', 'token', 'server-key') + const socket = mockSockets[0]! + socket.open() + socket.receive(JSON.stringify({ type: 'e2ee_ready' })) + socket.receive('encrypted:{"type":"e2ee_authenticated"}') + const dispose = client.subscribe('agentSession.subscribe', { sessionId: 'session-a' }, () => {}) + const subscribe = sentRequest(socket, 'agentSession.subscribe') + + dispose() + + expect(sentRequest(socket, 'agentSession.unsubscribe').params).toEqual({ + sessionId: 'session-a', + subscriptionId: subscribe.id + }) + client.close() + }) }) diff --git a/mobile/src/transport/rpc-client-terminal-subscription.ts b/mobile/src/transport/rpc-client-terminal-subscription.ts index 471bc3c4a4e..ce4fcd50174 100644 --- a/mobile/src/transport/rpc-client-terminal-subscription.ts +++ b/mobile/src/transport/rpc-client-terminal-subscription.ts @@ -34,11 +34,12 @@ export function updateTerminalSubscriptionViewport( } /** Build the unsubscribe RPC for a streaming method that needs the host told to - * tear down (session tabs, native chat), or null when none is required. Keeps + * tear down (tabs, native chat, structured sessions), or null when none is required. Keeps * the per-method echo logic out of the rpc-client teardown closure. */ export function buildStreamUnsubscribe( method: string | undefined, - params: unknown + params: unknown, + requestId?: string ): { method: string; params: Record } | null { if (!params || typeof params !== 'object') { return null @@ -61,6 +62,15 @@ export function buildStreamUnsubscribe( ? buildNativeChatUnsubscribe(agent, sessionId) : null } + if (method === 'agentSession.subscribe') { + const sessionId = (params as { sessionId?: unknown }).sessionId + return typeof sessionId === 'string' + ? { + method: 'agentSession.unsubscribe', + params: { sessionId, ...(requestId ? { subscriptionId: requestId } : {}) } + } + : null + } return null } diff --git a/mobile/src/transport/rpc-client.ts b/mobile/src/transport/rpc-client.ts index 47d64621912..7dbeab29f66 100644 --- a/mobile/src/transport/rpc-client.ts +++ b/mobile/src/transport/rpc-client.ts @@ -1104,7 +1104,7 @@ export function connect( }) } } else { - const unsub = buildStreamUnsubscribe(stream?.method, stream?.params) + const unsub = buildStreamUnsubscribe(stream?.method, stream?.params, id) if (unsub) { sendEncrypted({ id: nextId(), deviceToken, method: unsub.method, params: unsub.params }) } diff --git a/src/main/codex/codex-structured-session-adapter.test.ts b/src/main/codex/codex-structured-session-adapter.test.ts index d2a4fb89adc..d15bb08be77 100644 --- a/src/main/codex/codex-structured-session-adapter.test.ts +++ b/src/main/codex/codex-structured-session-adapter.test.ts @@ -319,6 +319,34 @@ describe('CodexStructuredSessionAdapter.acquire', () => { await closing expect(codex.connections[0]?.closeCount).toBe(1) }) + + it('fences an acquisition while launch resolution is still pending', async () => { + const launch = Promise.withResolvers() + const codex = fakeCodex() + const adapter = new CodexStructuredSessionAdapter({ + resolveLaunch: () => launch.promise, + openConnection: codex.openConnection, + readProcessStartTime: async () => 1_700_000_000_000 + }) + const acquiring = adapter.acquire({ + identity: identityFor('session-1'), + fence: 7, + spawnToken: 'spawn-9' + }) + + const closing = adapter.closeAll() + launch.resolve({ + command: 'codex', + args: ['app-server'], + cwd: '/work/repo', + codexHome: null, + resumeThreadId: null + }) + + await expect(acquiring).rejects.toThrow('superseded while being acquired') + await closing + expect(codex.connections).toHaveLength(0) + }) }) describe('CodexStructuredSessionAdapter.dispatch', () => { @@ -768,10 +796,11 @@ describe('CodexStructuredSessionAdapter lifecycle', () => { const events: CodexStructuredSessionEvent[] = [] const adapter = await acquired(codex, {}, events) await adapter.acquire({ identity: identityFor('session-1'), fence: 8, spawnToken: 'spawn-10' }) + const endedBeforeStaleExit = events.filter((event) => event.type === 'ended').length codex.connections[0].handlers.onExit?.(new Error('the superseded child died')) - expect(events.some((event) => event.type === 'ended')).toBe(false) + expect(events.filter((event) => event.type === 'ended')).toHaveLength(endedBeforeStaleExit) expect(await adapter.historyFilePath({ identity: identityFor('session-1') })).toBe( '/rollouts/abc.jsonl' ) @@ -796,13 +825,14 @@ describe('CodexStructuredSessionAdapter lifecycle', () => { it('flushes the final coalesced text before a graceful close', async () => { const codex = fakeCodex() const bodies: AgentJournalMessageItem[] = [] + const tombstones: unknown[] = [] const sink: StructuredAgentSessionEventSink = { appendItem: (_identity, body) => { if (body.kind === 'message') { bodies.push(body) } }, - appendTombstone: () => {}, + appendTombstone: (identity) => tombstones.push(identity), publish: () => {} } const adapter = adapterFor(codex) @@ -827,5 +857,11 @@ describe('CodexStructuredSessionAdapter lifecycle', () => { await adapter.closeSession('session-1') expect(bodies.at(-1)?.blocks).toEqual([{ type: 'text', text: 'last words' }]) + expect(tombstones).toContainEqual({ + provider: 'legacy', + agent: 'codex', + sessionId: 'session-1', + recordId: 'turn-lifecycle:turn-1' + }) }) }) diff --git a/src/main/codex/codex-structured-session-adapter.ts b/src/main/codex/codex-structured-session-adapter.ts index 63344a93bda..3b74e5d73c3 100644 --- a/src/main/codex/codex-structured-session-adapter.ts +++ b/src/main/codex/codex-structured-session-adapter.ts @@ -25,7 +25,9 @@ import { readCodexThreadId, readCodexTurnId } from './codex-structured-thread-fa import { openCodexThread } from './codex-structured-thread-open' import { dispatchCodexTurn, isCodexTurnOptionKey } from './codex-structured-turn-start' import { supportsCodexStructuredLocation } from './codex-structured-location-support' +import { closeCodexPublishedSession } from './codex-structured-session-close' import { + cancelCodexAcquisitionAttempt, createCodexAcquisitionAttempt, type CodexAcquisitionAttempt, type CodexSession, @@ -39,11 +41,7 @@ export type { CodexStructuredSessionEvent } from './codex-structured-session-state' -// The Codex half of the structured agent-session wire: one long-lived -// `codex app-server` child per session, started or resumed under the lease the -// host already reserved. Everything durable — journal, fence, idempotency — -// belongs to the wire; this adapter only starts the process, names the turn the -// provider accepted, and answers Codex's blocking prompt requests. +// Owns one app-server child per structured Codex session; durable state stays on the wire. export class CodexStructuredSessionAdapter implements StructuredAgentSessionAdapter { private readonly sessions = new Map() @@ -53,26 +51,19 @@ export class CodexStructuredSessionAdapter implements StructuredAgentSessionAdap supportsLocation = supportsCodexStructuredLocation - private get requestTimeoutMs(): number | undefined { - return this.deps.requestTimeoutMs - } - async acquire(input: { identity: AgentSessionJournalIdentity fence: number spawnToken: string events?: StructuredAgentSessionEventSink }): Promise { - // A re-acquire at a new fence must not leave the old child writing. - await this.closeSession(input.identity.sessionId) const sessionId = input.identity.sessionId - const launch = await this.deps.resolveLaunch({ identity: input.identity }) + const previousAttempt = this.acquiring.get(sessionId) const attempt = createCodexAcquisitionAttempt() const acquisition = attempt.window + this.acquiring.set(sessionId, attempt) let primaryThreadId = input.identity.providerHandle.kind === 'codex' ? input.identity.providerHandle.threadId : null - // Registered before the spawn, because the handshake itself can emit. - this.acquiring.set(sessionId, attempt) const translator = input.events ? createCodexJournalTranslator({ sink: input.events, @@ -84,6 +75,15 @@ export class CodexStructuredSessionAdapter implements StructuredAgentSessionAdap const open = this.deps.openConnection ?? openCodexAppServerConnection try { + await cancelCodexAcquisitionAttempt(previousAttempt) + await this.closePublishedSession(sessionId) + if (attempt.cancelled || this.acquiring.get(sessionId) !== attempt) { + throw new Error(`codex session ${sessionId} was superseded while being acquired`) + } + const launch = await this.deps.resolveLaunch({ identity: input.identity }) + if (attempt.cancelled || this.acquiring.get(sessionId) !== attempt) { + throw new Error(`codex session ${sessionId} was superseded while being acquired`) + } const connection = await open( { command: launch.command, @@ -106,7 +106,7 @@ export class CodexStructuredSessionAdapter implements StructuredAgentSessionAdap } ) acquisition.connection = connection - const opened = await openCodexThread(connection, launch, this.requestTimeoutMs) + const opened = await openCodexThread(connection, launch, this.deps.requestTimeoutMs) primaryThreadId = opened.threadId const acquired: AgentSessionAcquisition = { process: await codexProcessIdentity( @@ -121,10 +121,7 @@ export class CodexStructuredSessionAdapter implements StructuredAgentSessionAdap observedAt: this.deps.now?.() ?? Date.now() }) } - // Publication is one step at the end: everything the caller is promised is - // in hand, the child is still alive, and this attempt still owns the - // session. A half-published session would be a live map entry in front of - // a child nobody can reach. + // Publish only after every promised identity is proven and this attempt still owns the child. if (connection.closed) { throw new Error(`codex app-server for session ${sessionId} exited while being acquired`) } @@ -161,8 +158,7 @@ export class CodexStructuredSessionAdapter implements StructuredAgentSessionAdap } } - /** Buffers an event that arrived before the session was published, and drops - * one from a connection this session has already replaced. */ + /** Buffers pre-publication events and drops events from superseded children. */ private deliver( acquisition: CodexAcquisitionAttempt['window'], sessionId: string, @@ -176,8 +172,7 @@ export class CodexStructuredSessionAdapter implements StructuredAgentSessionAdap } } - /** Only the connection that currently owns the session may retire it, or a - * superseded child's death would evict its live replacement. */ + /** Only the current connection may retire a session. */ private handleExit( sessionId: string, acquisition: CodexAcquisitionAttempt['window'], @@ -210,8 +205,7 @@ export class CodexStructuredSessionAdapter implements StructuredAgentSessionAdap this.emit(session, { type: 'notification', sessionId, threadId, method, params }) } - /** Journal first, observers second: a test tap must never see a row the - * journal has not been told about. */ + /** Journal first so observers never see an event ahead of its durable row. */ private emit(session: CodexSession, event: CodexStructuredSessionEvent): void { session.translator?.handle(event) this.deps.onEvent?.(event) @@ -250,7 +244,7 @@ export class CodexStructuredSessionAdapter implements StructuredAgentSessionAdap body: AgentJournalMessageItem fence: number }): Promise { - return dispatchCodexTurn(this.session(input.sessionId), input, this.requestTimeoutMs) + return dispatchCodexTurn(this.session(input.sessionId), input, this.deps.requestTimeoutMs) } async cancelTurn(input: { @@ -263,7 +257,7 @@ export class CodexStructuredSessionAdapter implements StructuredAgentSessionAdap await session.connection.request( 'turn/interrupt', { threadId: session.threadId, turnId: input.turnId }, - { timeoutMs: this.requestTimeoutMs } + { timeoutMs: this.deps.requestTimeoutMs } ) return { cancelled: true } } catch (error) { @@ -315,15 +309,15 @@ export class CodexStructuredSessionAdapter implements StructuredAgentSessionAdap await attempt.window.connection?.close() await attempt.finished } - const session = this.sessions.get(sessionId) - if (!session) { - return - } - this.sessions.delete(sessionId) - session.prompts.clear() - session.translator?.flush() - session.translator?.dispose() - await session.connection.close() + await this.closePublishedSession(sessionId) + } + + private async closePublishedSession(sessionId: string): Promise { + await closeCodexPublishedSession({ + sessions: this.sessions, + sessionId, + onEvent: this.deps.onEvent + }) } async closeAll(): Promise { diff --git a/src/main/codex/codex-structured-session-close.ts b/src/main/codex/codex-structured-session-close.ts new file mode 100644 index 00000000000..dc3e401bf37 --- /dev/null +++ b/src/main/codex/codex-structured-session-close.ts @@ -0,0 +1,24 @@ +import type { CodexSession, CodexStructuredSessionEvent } from './codex-structured-session-state' + +export async function closeCodexPublishedSession(input: { + sessions: Map + sessionId: string + onEvent?: (event: CodexStructuredSessionEvent) => void +}): Promise { + const session = input.sessions.get(input.sessionId) + if (!session) { + return + } + input.sessions.delete(input.sessionId) + session.prompts.clear() + const event: CodexStructuredSessionEvent = { + type: 'ended', + sessionId: input.sessionId, + reason: 'codex session closed' + } + session.translator?.handle(event) + input.onEvent?.(event) + session.translator?.flush() + session.translator?.dispose() + await session.connection.close() +} diff --git a/src/main/codex/codex-structured-session-state.ts b/src/main/codex/codex-structured-session-state.ts index a6d94b2c175..9e4ca13a571 100644 --- a/src/main/codex/codex-structured-session-state.ts +++ b/src/main/codex/codex-structured-session-state.ts @@ -63,3 +63,14 @@ export function createCodexAcquisitionAttempt(): CodexAcquisitionAttempt { }) return { window: new CodexAcquisitionWindow(), cancelled: false, finished, finish } } + +export async function cancelCodexAcquisitionAttempt( + attempt: CodexAcquisitionAttempt | undefined +): Promise { + if (!attempt) { + return + } + attempt.cancelled = true + await attempt.window.connection?.close() + await attempt.finished +} diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-attach-flow.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-attach-flow.ts index 010c14fce7a..b75990348a8 100644 --- a/src/main/native-chat/agent-session-wire/structured-agent-session-attach-flow.ts +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-attach-flow.ts @@ -42,7 +42,9 @@ export type AttachFlowInput = { * host owns it and binds it to the journal inside `onAttached`. */ eventSink?: StructuredAgentSessionEventSink /** Stops acquisition-window events targeting the superseded journal. */ - onAcquiring?: () => void + onAcquiring?: () => Promise | void + /** Settles writes already captured by the superseded journal before opening another. */ + beforeJournalOpen?: () => Promise | void } export async function performAttach( @@ -80,6 +82,7 @@ export async function performAttach( } } + await input.beforeJournalOpen?.() const attached = await attachJournal({ record, params, @@ -119,7 +122,7 @@ async function acquireOwner( if (!spawnToken) { throw new Error('agent_session_ownership_unknown') } - input.onAcquiring?.() + await input.onAcquiring?.() const acquired = await input.adapter.acquire({ identity: journalIdentityFor(record, input.params), fence, 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 1dcb8837ec2..edbc0114db3 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 @@ -20,7 +20,10 @@ import type { } from '../../../shared/agent-session-wire' import { AgentSessionRecordStore } from '../../runtime/agent-session-record-store' import { journalDirectoryFor } from '../agent-session-journal/journal-paths' -import { openAgentSessionJournal } from '../agent-session-journal/journal-store' +import { + openAgentSessionJournal, + type AgentSessionJournal +} from '../agent-session-journal/journal-store' import type { AgentSessionDispatchOutcome, StructuredAgentSessionAdapter @@ -368,6 +371,39 @@ describe('attach', () => { expect(releaseAcquisition).toHaveBeenCalledWith({ sessionId: SESSION }) }) + + it('drains writes captured by the old journal before acquiring its replacement', async () => { + const record = await attach() + const events = acquire.mock.calls[0]?.[0].events + const oldJournal = ( + host as unknown as { sessions: Map } + ).sessions.get(SESSION)!.journal + const appendGate = Promise.withResolvers() + const originalAppend = oldJournal.appendItem.bind(oldJournal) + const append = vi.spyOn(oldJournal, 'appendItem').mockImplementationOnce(async (...args) => { + await appendGate.promise + return originalAppend(...args) + }) + events?.appendItem( + { provider: 'orca', clientMessageId: 'old-journal-write' }, + { kind: 'status', text: 'old journal write' } + ) + await vi.waitFor(() => expect(append).toHaveBeenCalledOnce()) + const released = await store.evictProvenDeadOwner({ + sessionId: SESSION, + expectedFence: record?.lease.runtimeFence ?? 1, + probe: { outcome: 'pid-absent' }, + now: NOW + }) + + const replacement = host.attach(CALLER, ensureParams(released.lease.runtimeFence)) + await new Promise((resolve) => setImmediate(resolve)) + expect(acquire).toHaveBeenCalledTimes(1) + + appendGate.resolve() + await expect(replacement).resolves.toMatchObject({ ok: true }) + expect(acquire).toHaveBeenCalledTimes(2) + }) }) describe('send', () => { @@ -409,6 +445,31 @@ describe('send', () => { expect(dispatch).toHaveBeenCalledTimes(1) }) + it('redispatches an explicitly retried durable unknown without appending a second submission', async () => { + await attach() + dispatch + .mockRejectedValueOnce(new Error('socket closed')) + .mockImplementationOnce(async () => accepted()) + const body = message('possibly delivered') + const params = { envelope: envelope('agentSession.send', { body }), body } + + const first = await host.send(CALLER, params) + expect(first).toMatchObject({ + ok: true, + value: { submission: { dispatchState: 'unknown' } } + }) + const retried = await host.send(CALLER, { ...params, retryUnknown: true }) + + expect(retried).toMatchObject({ + ok: true, + replayed: false, + value: { submission: { dispatchState: 'accepted' } } + }) + expect(dispatch).toHaveBeenCalledTimes(2) + const state = host.history({ sessionId: SESSION, direction: 'tail' }) + expect(state.ok && state.page.submissions).toHaveLength(1) + }) + it('refuses a stale fence and hands back the current one', async () => { const record = await attach() const body = message('add a retry') 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 dc845bfb0eb..4bbb85362ce 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 @@ -177,6 +177,10 @@ export class StructuredAgentSessionHost { journalRoot: this.deps.journalRoot, eventSink: eventSink.sink, onAcquiring: () => eventSink.unbind(), + beforeJournalOpen: async () => { + eventSink.unbind() + await eventSink.drained() + }, authority: { spawnToken: this.deps.mintSpawnToken?.() ?? randomUUID(), claimKeyId: this.deps.claimKeyId, @@ -263,7 +267,11 @@ export class StructuredAgentSessionHost { send( caller: StructuredAgentSessionCaller, - params: { envelope: AgentSessionMutationEnvelope; body: AgentJournalMessageItem } + params: { + envelope: AgentSessionMutationEnvelope + body: AgentJournalMessageItem + retryUnknown?: true + } ): Promise> { return this.mutate(caller, params.envelope, sendPlan(params)) } diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-mutation-admission.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-mutation-admission.ts index 4b6ac6d5f05..b43fa331270 100644 --- a/src/main/native-chat/agent-session-wire/structured-agent-session-mutation-admission.ts +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-mutation-admission.ts @@ -83,7 +83,8 @@ export async function admitAndRunAgentSessionMutation( const replay = resolveAgentSessionReplayOutcome({ operationId: envelope.clientOperationId, outcome: admission.row.outcome, - reconstruct: () => plan.replay(context, admission.row.outcome) + reconstruct: () => plan.replay(context, admission.row.outcome), + rerunWhenReplayMissing: plan.rerunWhenReplayMissing?.(context) }) if (replay.decision === 'refuse') { return refuseAgentSessionMutation(replay.refusal) diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-mutation-plans.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-mutation-plans.ts index 0d8bad1d19e..b87bba19a0d 100644 --- a/src/main/native-chat/agent-session-wire/structured-agent-session-mutation-plans.ts +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-mutation-plans.ts @@ -29,29 +29,42 @@ export type MutationPlan = { fields: Record run: (ctx: AgentSessionTurnContext) => Promise> replay: (ctx: AgentSessionTurnContext, outcome: AgentSessionOperationOutcome) => TValue | null + rerunWhenReplayMissing?: (ctx: AgentSessionTurnContext) => boolean } export function sendPlan(params: { envelope: AgentSessionMutationEnvelope body: AgentJournalMessageItem + retryUnknown?: true }): MutationPlan { // The operation id IS the client message id: one send, one durable row, one // key the client reconciles its optimistic bubble against. const clientMessageId = params.envelope.clientOperationId return { method: 'agentSession.send', + // A control signal is not payload; only the matching durable unknown unlocks redispatch. fields: { body: params.body }, + rerunWhenReplayMissing: (ctx) => + params.retryUnknown === true && + ctx.journal + .submissions() + .some( + (entry) => entry.clientMessageId === clientMessageId && entry.dispatchState === 'unknown' + ), run: (ctx) => performSend(ctx, { clientMessageId, payloadFingerprint: params.envelope.payloadFingerprint, - body: params.body + body: params.body, + retryUnknown: params.retryUnknown }), replay: (ctx) => { const submission = ctx.journal .submissions() .find((entry) => entry.clientMessageId === clientMessageId) - return submission ? { clientMessageId, submission } : null + return submission && !(params.retryUnknown && submission.dispatchState === 'unknown') + ? { clientMessageId, submission } + : null } } } diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-replay-outcome.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-replay-outcome.ts index 1121dbc8124..9417b5d7fd1 100644 --- a/src/main/native-chat/agent-session-wire/structured-agent-session-replay-outcome.ts +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-replay-outcome.ts @@ -14,6 +14,7 @@ export function resolveAgentSessionReplayOutcome(input: { operationId: string outcome: AgentSessionOperationOutcome reconstruct: () => TValue | null + rerunWhenReplayMissing?: boolean }): AgentSessionReplayOutcomeDecision { const { operationId, outcome } = input if (outcome.status === 'failed') { @@ -41,6 +42,9 @@ export function resolveAgentSessionReplayOutcome(input: { if (recorded) { return { decision: 'replay', value: recorded } } + if (input.rerunWhenReplayMissing) { + return { decision: 'rerun' } + } return outcome.status === 'succeeded' ? { decision: 'refuse', diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-turns.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-turns.ts index 33a022aa61f..61ee362c993 100644 --- a/src/main/native-chat/agent-session-wire/structured-agent-session-turns.ts +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-turns.ts @@ -79,10 +79,20 @@ async function appendStatus( export async function performSend( ctx: AgentSessionTurnContext, - input: { clientMessageId: string; payloadFingerprint: string; body: AgentJournalMessageItem } + input: { + clientMessageId: string + payloadFingerprint: string + body: AgentJournalMessageItem + retryUnknown?: true + } ): Promise> { - await ctx.journal.appendSubmission({ ...input, fence: ctx.fence }) - ctx.publish() + const existing = ctx.journal + .submissions() + .find((entry) => entry.clientMessageId === input.clientMessageId) + if (!(input.retryUnknown && existing?.dispatchState === 'unknown')) { + await ctx.journal.appendSubmission({ ...input, fence: ctx.fence }) + ctx.publish() + } const outcome = await dispatchSafely(ctx, input.clientMessageId, input.body) await ctx.journal.resolveDispatch( diff --git a/src/main/runtime/rpc/methods/structured-agent-session-schemas.ts b/src/main/runtime/rpc/methods/structured-agent-session-schemas.ts index 2e2c8db64c0..d6ec84051d7 100644 --- a/src/main/runtime/rpc/methods/structured-agent-session-schemas.ts +++ b/src/main/runtime/rpc/methods/structured-agent-session-schemas.ts @@ -133,6 +133,7 @@ const SendBlock = z.discriminatedUnion('type', [ export const SendParams = z .object({ envelope: MutationEnvelope, + retryUnknown: z.literal(true).optional(), body: z .object({ kind: z.literal('message'), diff --git a/src/main/runtime/rpc/mobile-socket-wiring.test.ts b/src/main/runtime/rpc/mobile-socket-wiring.test.ts index 505cf63feb6..f747fcb1fa0 100644 --- a/src/main/runtime/rpc/mobile-socket-wiring.test.ts +++ b/src/main/runtime/rpc/mobile-socket-wiring.test.ts @@ -348,4 +348,86 @@ describe('MobileSocketWiring', () => { expect(transport.setClientId).not.toHaveBeenCalled() expect(ws.close).toHaveBeenCalledWith(4001, 'Unauthorized') }) + + it('projects post-auth v2 capabilities into later RPC dispatches', () => { + const desktop = nacl.box.keyPair.fromSecretKey(new Uint8Array(32).fill(1)) + const phone = nacl.box.keyPair.fromSecretKey(new Uint8Array(32).fill(2)) + const ws = new FakeSocket() + const transport = new FakeTransport() + const onText = vi.fn() + const metadata: MobileSocketTransportMetadata = { + transport: 'relay', + relayHostId: 'AbCdEf0123_-xyZ9', + relayDeviceId: 'device-1', + basisConnId: 'connection-1', + credentialKind: 'resume' + } + const wiring = new MobileSocketWiring({ + deviceRegistry: registryFor('device-1', 'valid-token'), + e2eeKeypair: { + publicKey: desktop.publicKey, + secretKey: desktop.secretKey, + publicKeyB64: Buffer.from(desktop.publicKey).toString('base64') + }, + onText, + onBinary: vi.fn(), + onClose: vi.fn() + }) + wiring.attachTransport(transport, () => metadata) + const hello: MobileE2EEV2Hello = { + type: 'e2ee_hello', + v: 2, + clientPublicKeyB64: Buffer.from(phone.publicKey).toString('base64'), + clientNonceB64: Buffer.from(new Uint8Array(32).fill(3)).toString('base64'), + capabilities: { framing: [2], payloadKinds: ['text', 'binary'] }, + context: { + protocol: 'orca-mobile-e2ee', + initiator: 'mobile', + responder: 'desktop', + transport: 'relay', + relayHostId: metadata.relayHostId + } + } + transport.receive(ws, JSON.stringify(hello)) + const ready = JSON.parse(ws.sent[0]!.toString()) as MobileE2EEV2Ready + const handshake = validateMobileE2EEV2Handshake(hello, ready)! + const schedule = deriveMobileE2EEV2KeySchedule({ + sharedSecret: deriveSharedKey(phone.secretKey, desktop.publicKey), + transcript: encodeMobileE2EEV2Transcript(handshake), + clientNonce: handshake.clientNonce, + desktopNonce: handshake.desktopNonce + }) + const send = (value: unknown, counter: bigint): void => { + const frame = sealMobileE2EEV2Frame({ + payload: new TextEncoder().encode(JSON.stringify(value)), + key: schedule.mobileToDesktopKey, + sessionId: schedule.sessionId, + direction: 'mobile-to-desktop', + payloadKind: 'text', + counter + }) + transport.receive(ws, Buffer.from(frame).toString('base64')) + } + send( + { + type: 'e2ee_auth', + v: 2, + transcriptHashB64: Buffer.from(schedule.transcriptHash).toString('base64'), + deviceToken: 'valid-token' + }, + 0n + ) + send( + { + type: 'e2ee_client_capabilities', + v: 1, + clientCapabilities: ['agent-session.structured.v1'] + }, + 1n + ) + send({ id: 'rpc-1', method: 'agentSession.history', params: {} }, 2n) + + expect(onText).toHaveBeenCalledOnce() + expect(onText.mock.calls[0]?.[0].clientCapabilities).toEqual(['agent-session.structured.v1']) + }) }) diff --git a/src/main/runtime/rpc/mobile-socket-wiring.ts b/src/main/runtime/rpc/mobile-socket-wiring.ts index 43004be4582..384f403e9de 100644 --- a/src/main/runtime/rpc/mobile-socket-wiring.ts +++ b/src/main/runtime/rpc/mobile-socket-wiring.ts @@ -165,7 +165,9 @@ export class MobileSocketWiring { ws, connectionId, device, - clientCapabilities: channel.clientCapabilities, + get clientCapabilities() { + return channel.clientCapabilities + }, transport: metadata } this.authenticatedSockets.set(ws, socket) diff --git a/src/main/runtime/structured-agent-session-integration.test.ts b/src/main/runtime/structured-agent-session-integration.test.ts index c351a907a04..51b6c0689be 100644 --- a/src/main/runtime/structured-agent-session-integration.test.ts +++ b/src/main/runtime/structured-agent-session-integration.test.ts @@ -464,7 +464,6 @@ describe('a structured codex session over agentSession.*', () => { // no page overlaps another, and nothing the live stream showed is missing. expect([...older.page.items, ...tail.page.items].map((item) => item.body?.kind)).toEqual([ 'message', - 'status', 'message', 'tool-call', 'approval', @@ -473,7 +472,6 @@ describe('a structured codex session over agentSession.*', () => { ]) expect([...older.page.items, ...tail.page.items].map(textOf)).toEqual([ 'list files', - '', 'Two files.', '', '',