fix(codex): harden account-switch goal resume

2026-10-21 00:03:46 +00:00 · 2026-08-02 16:19:37 -07:00
parent b2f6b795e8
commit 6a7a8a9c27
7 changed files with 316 additions and 47 deletions
@@ -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({
@@ -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') {
@@ -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<string>((_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<string>>()
+      .mockImplementationOnce(
+        () =>
+          new Promise<string>((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')
@@ -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<CodexAppServerHostKey, number>()
  private readonly supportedHosts = new Set<CodexAppServerHostKey>()
+  private readonly inFlightProbeByHost = new Map<CodexAppServerHostKey, Promise<AsyncProbeState>>()

  shouldTry(hostKey: CodexAppServerHostKey, nowMs = Date.now()): boolean {
    const retryAfterMs = this.retryAfterByHost.get(hostKey)
@@ -75,9 +80,89 @@ export class CodexAppServerCapabilityCache {
    }
  }

+  async runWithFallbackAsync<T>(
+    hostKey: CodexAppServerHostKey,
+    runPreferred: () => Promise<T>,
+    runFallback: () => T | Promise<T>,
+    isUnsupportedError: (error: unknown) => boolean,
+    nowMs = Date.now()
+  ): Promise<T> {
+    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<AsyncProbeState>
+    probe = (async (): Promise<AsyncProbeState> => {
+      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<T>(
+    runPreferred: () => Promise<T>,
+    runFallback: () => T | Promise<T>,
+    hostKey: CodexAppServerHostKey,
+    isUnsupportedError: (error: unknown) => boolean,
+    nowMs: number
+  ): Promise<T> {
+    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()
  }
 }

@@ -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
      }

@@ -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')
+  })
+})
@@ -0,0 +1,34 @@
+import type { GlobalSettings } from '../../../shared/types'
+
+type CodexAccountRestartTransport = {
+  getPtyId: () => string | null
+}
+
+export function getSelectedHostCodexAccountId(
+  settings:
+    | Pick<GlobalSettings, 'activeCodexManagedAccountId' | 'activeCodexManagedAccountIdsByRuntime'>
+    | 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'
+}