From 6a7a8a9c27f79b79a0a8712e77dc3cf0f2572b50 Mon Sep 17 00:00:00 2001 From: Brennan Benson <79079362+brennanb2025@users.noreply.github.com> Date: Sun, 2 Aug 2026 16:19:37 -0700 Subject: [PATCH] fix(codex): harden account-switch goal resume --- .../codex/codex-account-switch-resume.test.ts | 48 ++++++++-- src/main/codex/codex-account-switch-resume.ts | 71 +++++++------- .../codex-app-server-capability-cache.test.ts | 57 ++++++++++++ .../codex-app-server-capability-cache.ts | 93 ++++++++++++++++++- .../components/terminal-pane/TerminalPane.tsx | 26 +++++- .../codex-account-restart-apply-state.test.ts | 34 +++++++ .../lib/codex-account-restart-apply-state.ts | 34 +++++++ 7 files changed, 316 insertions(+), 47 deletions(-) create mode 100644 src/renderer/src/lib/codex-account-restart-apply-state.test.ts create mode 100644 src/renderer/src/lib/codex-account-restart-apply-state.ts diff --git a/src/main/codex/codex-account-switch-resume.test.ts b/src/main/codex/codex-account-switch-resume.test.ts index 64a30e53b0e..93f6085046c 100644 --- a/src/main/codex/codex-account-switch-resume.test.ts +++ b/src/main/codex/codex-account-switch-resume.test.ts @@ -246,7 +246,7 @@ describe('prepareCodexAccountSwitchResume', () => { ]) }) - it('still resumes a goalless thread, validating it via thread/read alone', async () => { + it('clears a stale target goal when the source thread is goalless', async () => { const { deps, calls } = prepareDeps({ script: (call) => (call.method === 'thread/goal/get' ? { goal: null } : {}) }) @@ -254,7 +254,11 @@ describe('prepareCodexAccountSwitchResume', () => { const decision = await prepareCodexAccountSwitchResume({ threadId: THREAD_ID }, deps) expect(decision).toEqual({ outcome: 'resume', threadId: THREAD_ID }) - expect(calls.map((call) => call.method)).toEqual(['thread/goal/get', 'thread/read']) + expect(calls).toEqual([ + { home: OLD_HOME, method: 'thread/goal/get', params: { threadId: THREAD_ID } }, + { home: NEW_HOME, method: 'thread/read', params: { threadId: THREAD_ID } }, + { home: NEW_HOME, method: 'thread/goal/clear', params: { threadId: THREAD_ID } } + ]) }) it('normalizes a missing token budget to null for the write', async () => { @@ -268,10 +272,16 @@ describe('prepareCodexAccountSwitchResume', () => { expect(calls.at(-1)?.params).toMatchObject({ tokenBudget: null }) }) - it('skips the bridge entirely when the launch home is also the new home', async () => { + it('resumes immediately without app-server work when both homes are the same', async () => { + const cache = new CodexAppServerCapabilityCache() + cache.rememberUnsupported(CODEX_GOAL_RPC_HOST_KEY, 1_000) const { deps, calls, bridgedHomes } = prepareDeps({ - script: () => ({ goal: GOAL }), - homes: { oldCodexHomePath: OLD_HOME, newCodexHomePath: OLD_HOME } + script: () => { + throw new Error('same-home restart must not spawn an app-server') + }, + homes: { oldCodexHomePath: OLD_HOME, newCodexHomePath: OLD_HOME }, + cache, + nowMs: () => 1_001 }) const decision = await prepareCodexAccountSwitchResume({ threadId: THREAD_ID }, deps) @@ -279,7 +289,7 @@ describe('prepareCodexAccountSwitchResume', () => { // Why: one home means one goals DB and one rollout tree; resume carries both. expect(decision).toEqual({ outcome: 'resume', threadId: THREAD_ID }) expect(bridgedHomes).toEqual([]) - expect(calls.map((call) => call.method)).toEqual(['thread/goal/get']) + expect(calls).toEqual([]) }) it('declines a thread id that is not a bare rollout UUID without touching anything', async () => { @@ -333,6 +343,22 @@ describe('prepareCodexAccountSwitchResume', () => { }) }) + it('declines when clearing a stale target goal fails', async () => { + const { deps } = prepareDeps({ + script: (call) => { + if (call.method === 'thread/goal/clear') { + throw new Error('codex app-server thread/goal/clear failed: boom') + } + return call.method === 'thread/goal/get' ? { goal: null } : {} + } + }) + + expect(await prepareCodexAccountSwitchResume({ threadId: THREAD_ID }, deps)).toEqual({ + outcome: 'fresh', + reason: 'goal-write-failed' + }) + }) + it('declines on an unrecognizable goal payload instead of guessing', async () => { const { deps } = prepareDeps({ script: () => ({ goal: { status: 'usageLimited' } }) @@ -344,6 +370,16 @@ describe('prepareCodexAccountSwitchResume', () => { }) }) + it('declines a goal response that omits the goal property', async () => { + const { deps, calls } = prepareDeps({ script: () => ({}) }) + + expect(await prepareCodexAccountSwitchResume({ threadId: THREAD_ID }, deps)).toEqual({ + outcome: 'fresh', + reason: 'goal-shape-unexpected' + }) + expect(calls.map((call) => call.method)).toEqual(['thread/goal/get']) + }) + it('treats a transient read failure as fresh without poisoning the capability', async () => { const cache = new CodexAppServerCapabilityCache() const { deps } = prepareDeps({ diff --git a/src/main/codex/codex-account-switch-resume.ts b/src/main/codex/codex-account-switch-resume.ts index 1275a887479..30e37a051e4 100644 --- a/src/main/codex/codex-account-switch-resume.ts +++ b/src/main/codex/codex-account-switch-resume.ts @@ -19,7 +19,7 @@ import { getCodexPaneAccount, type CodexPaneAccountRecord } from './codex-pane-a // the launch CODEX_HOME's sqlite, so the goal became doubly unreachable (new // thread, other home). This module decides whether the restarted pane may // `codex resume` its live thread instead, and carries the goal row across homes -// through the sanctioned app-server RPCs (thread/goal/get → thread/goal/set). +// through the sanctioned app-server goal RPCs (get → set/clear). // Orca never reads or writes Codex's sqlite itself: the app-server owns the DB // and its WAL. Any doubt — unrecorded pane, non-host lane, custom home, missing // rollout, old CLI without the goal RPCs — degrades to today's fresh `codex` @@ -29,6 +29,7 @@ const CODEX_THREAD_ID_PATTERN = /^[0-9a-f]{8}(?:-[0-9a-f]{4}){3}-[0-9a-f]{12}$/i const GOAL_READ_TIMEOUT_MS = 10_000 const GOAL_WRITE_TIMEOUT_MS = 15_000 +const GOAL_RPC_UNSUPPORTED = Symbol('goal-rpc-unsupported') /** Goal RPCs are probed on the native host only; WSL/remote lanes never reach * the probe (their panes degrade to a fresh startup in resolveHomes). */ @@ -161,7 +162,7 @@ export type PrepareCodexAccountSwitchResumeDeps = { * Decides whether an account-switch restart may resume the pane's thread, and * performs the goal bridge as part of deciding: `resume` is only answered once * the rollout is present in the new home, `thread/read` proved it resumable - * there, and any goal row was written through `thread/goal/set`. + * there, and the target goal was set or cleared to match the source. */ export async function prepareCodexAccountSwitchResume( request: { threadId: string }, @@ -176,14 +177,15 @@ export async function prepareCodexAccountSwitchResume( return homesDecision } const { homes } = homesDecision + if ( + normalizeRuntimePathForComparison(homes.oldCodexHomePath) === + normalizeRuntimePathForComparison(homes.newCodexHomePath) + ) { + // Same home keeps the same goals DB and rollout tree; resume alone carries both. + return { outcome: 'resume', threadId } + } const cache = deps.capabilityCache ?? codexGoalRpcCapabilityCache const nowMs = deps.nowMs ?? ((): number => Date.now()) - if ( - !cache.isKnownSupported(CODEX_GOAL_RPC_HOST_KEY) && - !cache.shouldTry(CODEX_GOAL_RPC_HOST_KEY, nowMs()) - ) { - return { outcome: 'fresh', reason: 'goal-rpc-unsupported' } - } const runSession = deps.runAppServerSession ?? runCodexAppServerSession const buildInvocation = deps.buildInvocation ?? buildCodexGoalRpcInvocation @@ -191,32 +193,28 @@ export async function prepareCodexAccountSwitchResume( try { // Why: the get doubles as the capability probe — a method-not-found reply or // a missing app-server subcommand is the only signal that marks unsupported. - const result = await runSession( - buildInvocation(homes.oldCodexHomePath, GOAL_READ_TIMEOUT_MS), - (rpc) => rpc.request('thread/goal/get', { threadId }) + const result = await cache.runWithFallbackAsync( + CODEX_GOAL_RPC_HOST_KEY, + () => + runSession(buildInvocation(homes.oldCodexHomePath, GOAL_READ_TIMEOUT_MS), (rpc) => + rpc.request('thread/goal/get', { threadId }) + ), + () => GOAL_RPC_UNSUPPORTED, + isCodexAppServerUnsupportedError, + nowMs() ) + if (result === GOAL_RPC_UNSUPPORTED) { + return { outcome: 'fresh', reason: 'goal-rpc-unsupported' } + } const parsed = parseThreadGoalSnapshot(result) if (parsed.outcome === 'malformed') { return { outcome: 'fresh', reason: 'goal-shape-unexpected' } } goal = parsed.goal - cache.rememberSupported(CODEX_GOAL_RPC_HOST_KEY) - } catch (error) { - if (isCodexAppServerUnsupportedError(error)) { - cache.rememberUnsupported(CODEX_GOAL_RPC_HOST_KEY, nowMs()) - return { outcome: 'fresh', reason: 'goal-rpc-unsupported' } - } + } catch { return { outcome: 'fresh', reason: 'goal-read-failed' } } - if ( - normalizeRuntimePathForComparison(homes.oldCodexHomePath) === - normalizeRuntimePathForComparison(homes.newCodexHomePath) - ) { - // Same home keeps the same goals DB and rollout tree; resume alone carries both. - return { outcome: 'resume', threadId } - } - if (!(await deps.ensureRolloutBridged(homes))) { return { outcome: 'fresh', reason: 'rollout-not-bridged' } } @@ -229,14 +227,16 @@ export async function prepareCodexAccountSwitchResume( // bridged rollout and upserts the thread row, proving the thread is // resumable in this home before Orca commits the pane to `codex resume`. await rpc.request('thread/read', { threadId }) - if (goal) { - await rpc.request('thread/goal/set', { - threadId, - objective: goal.objective, - status: goal.status, - tokenBudget: goal.tokenBudget - }) - } + // Why: the target may retain an older row from a prior switch; mirroring + // null must remove it or the resumed thread revives stale state. + await (goal + ? rpc.request('thread/goal/set', { + threadId, + objective: goal.objective, + status: goal.status, + tokenBudget: goal.tokenBudget + }) + : rpc.request('thread/goal/clear', { threadId })) } ) } catch (error) { @@ -258,7 +258,10 @@ function parseThreadGoalSnapshot( return { outcome: 'malformed' } } const goal = (result as { goal?: unknown }).goal - if (goal === null || goal === undefined) { + if (goal === undefined) { + return { outcome: 'malformed' } + } + if (goal === null) { return { outcome: 'parsed', goal: null } } if (typeof goal !== 'object') { diff --git a/src/main/codex/codex-app-server-capability-cache.test.ts b/src/main/codex/codex-app-server-capability-cache.test.ts index c6dea6e0031..8734f3b4998 100644 --- a/src/main/codex/codex-app-server-capability-cache.test.ts +++ b/src/main/codex/codex-app-server-capability-cache.test.ts @@ -109,6 +109,63 @@ describe('CodexAppServerCapabilityCache', () => { expect(cache.shouldTry('native', 2)).toBe(true) }) + it('coalesces concurrent unsupported async probes per host', async () => { + const cache = new CodexAppServerCapabilityCache() + let rejectProbe!: (error: Error) => void + const preferred = vi.fn( + () => + new Promise((_resolve, reject) => { + rejectProbe = reject + }) + ) + + const first = cache.runWithFallbackAsync( + 'native', + preferred, + () => 'first-fallback', + isUnsupported, + 5 + ) + const second = cache.runWithFallbackAsync( + 'native', + preferred, + () => 'second-fallback', + isUnsupported, + 5 + ) + await Promise.resolve() + expect(preferred).toHaveBeenCalledTimes(1) + + rejectProbe(unsupportedError) + await expect(first).resolves.toBe('first-fallback') + await expect(second).resolves.toBe('second-fallback') + expect(preferred).toHaveBeenCalledTimes(1) + }) + + it('lets concurrent callers run their own work once the async probe succeeds', async () => { + const cache = new CodexAppServerCapabilityCache() + let resolveProbe!: (value: string) => void + const preferred = vi + .fn<() => Promise>() + .mockImplementationOnce( + () => + new Promise((resolve) => { + resolveProbe = resolve + }) + ) + .mockResolvedValueOnce('second-result') + + const first = cache.runWithFallbackAsync('native', preferred, () => 'fallback', isUnsupported) + const second = cache.runWithFallbackAsync('native', preferred, () => 'fallback', isUnsupported) + await Promise.resolve() + expect(preferred).toHaveBeenCalledTimes(1) + + resolveProbe('first-result') + await expect(first).resolves.toBe('first-result') + await expect(second).resolves.toBe('second-result') + expect(preferred).toHaveBeenCalledTimes(2) + }) + it('builds host keys that keep WSL distros apart', () => { expect(getCodexAppServerHostKey({ kind: 'native' })).toBe('native') expect(getCodexAppServerHostKey({ kind: 'wsl', distro: 'Ubuntu' })).toBe('wsl:Ubuntu') diff --git a/src/main/codex/codex-app-server-capability-cache.ts b/src/main/codex/codex-app-server-capability-cache.ts index 85816c9cf0c..784759f221f 100644 --- a/src/main/codex/codex-app-server-capability-cache.ts +++ b/src/main/codex/codex-app-server-capability-cache.ts @@ -7,6 +7,11 @@ export const CODEX_APP_SERVER_CAPABILITY_RETRY_INTERVAL_MS = 30 * 60_000 * the native host and from each other — each can carry a different codex. */ export type CodexAppServerHostKey = 'native' | `wsl:${string}` +type AsyncProbeState = + | { outcome: 'supported' } + | { outcome: 'unsupported' } + | { outcome: 'transient-error'; error: unknown } + export function getCodexAppServerHostKey( host: { kind: 'native' } | { kind: 'wsl'; distro: string } ): CodexAppServerHostKey { @@ -14,14 +19,14 @@ export function getCodexAppServerHostKey( } /** - * Capability cache for the codex app-server trust-grant RPC pair, modeled on - * GitCapabilityCache but with a synchronous runner: the grant client blocks - * the main thread by design (launch prep), so probes cannot overlap — the - * unsupported mark alone is what keeps later installs off the dead probe. + * Capability cache for codex app-server RPC surfaces, modeled on + * GitCapabilityCache. Synchronous launch probes use the direct cache; async + * restart probes coalesce the first unknown-capability call per host. */ export class CodexAppServerCapabilityCache { private readonly retryAfterByHost = new Map() private readonly supportedHosts = new Set() + private readonly inFlightProbeByHost = new Map>() shouldTry(hostKey: CodexAppServerHostKey, nowMs = Date.now()): boolean { const retryAfterMs = this.retryAfterByHost.get(hostKey) @@ -75,9 +80,89 @@ export class CodexAppServerCapabilityCache { } } + async runWithFallbackAsync( + hostKey: CodexAppServerHostKey, + runPreferred: () => Promise, + runFallback: () => T | Promise, + isUnsupportedError: (error: unknown) => boolean, + nowMs = Date.now() + ): Promise { + if (this.supportedHosts.has(hostKey)) { + return this.runKnownAsync(runPreferred, runFallback, hostKey, isUnsupportedError, nowMs) + } + if (!this.shouldTry(hostKey, nowMs)) { + return runFallback() + } + + const existingProbe = this.inFlightProbeByHost.get(hostKey) + if (existingProbe) { + const state = await existingProbe + if (state.outcome === 'unsupported') { + return runFallback() + } + if (state.outcome === 'transient-error') { + throw state.error + } + return this.runKnownAsync(runPreferred, runFallback, hostKey, isUnsupportedError, nowMs) + } + + let preferredResult!: T + let probe!: Promise + probe = (async (): Promise => { + try { + // Defer invocation until after the promise is registered so a + // synchronous throw still updates the shared capability state. + preferredResult = await Promise.resolve().then(runPreferred) + if (this.inFlightProbeByHost.get(hostKey) === probe) { + this.rememberSupported(hostKey) + } + return { outcome: 'supported' } + } catch (error) { + if (isUnsupportedError(error)) { + if (this.inFlightProbeByHost.get(hostKey) === probe) { + this.rememberUnsupported(hostKey, nowMs) + } + return { outcome: 'unsupported' } + } + return { outcome: 'transient-error', error } + } + })() + this.inFlightProbeByHost.set(hostKey, probe) + const state = await probe + if (this.inFlightProbeByHost.get(hostKey) === probe) { + this.inFlightProbeByHost.delete(hostKey) + } + if (state.outcome === 'unsupported') { + return runFallback() + } + if (state.outcome === 'transient-error') { + throw state.error + } + return preferredResult + } + + private async runKnownAsync( + runPreferred: () => Promise, + runFallback: () => T | Promise, + hostKey: CodexAppServerHostKey, + isUnsupportedError: (error: unknown) => boolean, + nowMs: number + ): Promise { + try { + return await runPreferred() + } catch (error) { + if (!isUnsupportedError(error)) { + throw error + } + this.rememberUnsupported(hostKey, nowMs) + return runFallback() + } + } + clear(): void { this.retryAfterByHost.clear() this.supportedHosts.clear() + this.inFlightProbeByHost.clear() } } diff --git a/src/renderer/src/components/terminal-pane/TerminalPane.tsx b/src/renderer/src/components/terminal-pane/TerminalPane.tsx index b3b81a18be7..8e85bb8a0bd 100644 --- a/src/renderer/src/components/terminal-pane/TerminalPane.tsx +++ b/src/renderer/src/components/terminal-pane/TerminalPane.tsx @@ -149,6 +149,10 @@ import { resolveCodexAccountRestartStartup, type CodexAccountRestartStartup } from '@/lib/codex-session-restart' +import { + getSelectedHostCodexAccountId, + resolveCodexAccountRestartApplyState +} from '@/lib/codex-account-restart-apply-state' import { WORKSPACE_FILE_PATH_MIME, WORKSPACE_FILE_PATHS_MIME } from '@/lib/workspace-file-drag' import { isTerminalSessionStateSaveFailure } from '../../../../shared/terminal-session-state-save-failure' import { isTerminalZeroDimensionsDiagnostic } from '../../../../shared/terminal-zero-dimensions-diagnostic' @@ -1631,6 +1635,7 @@ function TerminalPane( const transport = paneTransportsRef.current.get(paneId) const panePtyBinding = panePtyBindingsRef.current.get(paneId) const existingPtyId = transport?.getPtyId() + const selectedHostAccountId = getSelectedHostCodexAccountId(useAppStore.getState().settings) // Why before teardown: resume preparation reads the live session's /goal // through the old home's app-server and bridges it into the new home; a @@ -1648,9 +1653,24 @@ function TerminalPane( // second restart can interleave before this one finishes. restartingCodexPaneIdsRef.current.delete(paneId) } - if (paneTransportsRef.current.get(paneId) !== transport) { - // Why: the pane reconnected while preparation ran; a second teardown - // would kill the replacement PTY. + const applyState = resolveCodexAccountRestartApplyState({ + capturedTransport: transport, + currentTransport: paneTransportsRef.current.get(paneId), + capturedPtyId: existingPtyId, + capturedHostAccountId: selectedHostAccountId, + currentHostAccountId: getSelectedHostCodexAccountId(useAppStore.getState().settings) + }) + if (applyState === 'target-replaced') { + // Why: reconnect can replace the PTY on the same transport object; + // teardown authority belongs to the captured object + PTY generation. + return + } + if (applyState === 'selection-changed') { + // Why: preparation targeted the earlier selection. Retry against the + // current account instead of launching its shell with another home's bridge. + if (existingPtyId) { + useAppStore.getState().queueCodexPaneRestarts([existingPtyId]) + } return } diff --git a/src/renderer/src/lib/codex-account-restart-apply-state.test.ts b/src/renderer/src/lib/codex-account-restart-apply-state.test.ts new file mode 100644 index 00000000000..80043c195b8 --- /dev/null +++ b/src/renderer/src/lib/codex-account-restart-apply-state.test.ts @@ -0,0 +1,34 @@ +import { describe, expect, it } from 'vitest' +import { resolveCodexAccountRestartApplyState } from './codex-account-restart-apply-state' + +describe('resolveCodexAccountRestartApplyState', () => { + it('rejects a PTY generation replaced on the same transport during preparation', () => { + let ptyId = 'pty-old' + const transport = { getPtyId: (): string => ptyId } + ptyId = 'pty-new' + + expect( + resolveCodexAccountRestartApplyState({ + capturedTransport: transport, + currentTransport: transport, + capturedPtyId: 'pty-old', + capturedHostAccountId: 'account-b', + currentHostAccountId: 'account-b' + }) + ).toBe('target-replaced') + }) + + it('retries when the selected account changes during preparation', () => { + const transport = { getPtyId: (): string => 'pty-1' } + + expect( + resolveCodexAccountRestartApplyState({ + capturedTransport: transport, + currentTransport: transport, + capturedPtyId: 'pty-1', + capturedHostAccountId: 'account-b', + currentHostAccountId: 'account-c' + }) + ).toBe('selection-changed') + }) +}) diff --git a/src/renderer/src/lib/codex-account-restart-apply-state.ts b/src/renderer/src/lib/codex-account-restart-apply-state.ts new file mode 100644 index 00000000000..95d21ef1f36 --- /dev/null +++ b/src/renderer/src/lib/codex-account-restart-apply-state.ts @@ -0,0 +1,34 @@ +import type { GlobalSettings } from '../../../shared/types' + +type CodexAccountRestartTransport = { + getPtyId: () => string | null +} + +export function getSelectedHostCodexAccountId( + settings: + | Pick + | null + | undefined +): string | null { + return ( + settings?.activeCodexManagedAccountIdsByRuntime?.host ?? + settings?.activeCodexManagedAccountId ?? + null + ) +} + +export function resolveCodexAccountRestartApplyState(args: { + capturedTransport: CodexAccountRestartTransport | undefined + currentTransport: CodexAccountRestartTransport | undefined + capturedPtyId: string | null | undefined + capturedHostAccountId: string | null + currentHostAccountId: string | null +}): 'apply' | 'target-replaced' | 'selection-changed' { + if ( + args.currentTransport !== args.capturedTransport || + args.currentTransport?.getPtyId() !== args.capturedPtyId + ) { + return 'target-replaced' + } + return args.currentHostAccountId === args.capturedHostAccountId ? 'apply' : 'selection-changed' +}