orca/src/main/codex/codex-structured-session-state.ts

import type { StructuredAgentSessionLogger } from '../native-chat/agent-session-wire/structured-agent-session-logger'
import type {
  AgentJournalItemIdentity,
  AgentSessionJournalIdentity
} from '../../shared/agent-session-journal-types'
import { randomUUID } from 'node:crypto'
import type { AgentJournalDispatchRejection } from '../../shared/agent-session-failure-words'
import { cancelProcessAcquisition } from '../../shared/child-process/cancel-process-acquisition'
import type {
  CodexAppServerConnection,
  openCodexAppServerConnection
} from './codex-app-server-connection'
import { CodexAcquisitionWindow } from './codex-structured-acquisition-window'
import {
  createCodexTurnOpenWaits,
  type CodexTurnOpenWaits
} from './codex-structured-turn-open-wait'
import type { CodexDispatchEchoes } from './codex-structured-dispatch-echo'
import type { AgentChildWorkEvidence } from '../../shared/agent-status-child-work-evidence'
import type { CodexBackgroundTaskTracker } from './codex-background-task-tracker'
import type { CodexJournalTranslator } from './codex-structured-journal-translation'
import type { StructuredAgentSessionEndedEvent } from '../native-chat/agent-session-wire/structured-agent-session-adapter'
import type { CodexStructuredPermissionPolicy } from './codex-structured-permission-policy'
import type {
  AgentModelCatalogSessionAccess,
  AgentModelCatalogStore
} from '../native-chat/agent-model-catalog/agent-model-catalog-store'

export type CodexSessionCatalogAccess = AgentModelCatalogSessionAccess

export type CodexStructuredLaunch = {
  command: string
  args: string[]
  cwd: string
  codexHome: string | null
  resumeThreadId: string | null
  resumePath?: string | null
  /** The resumed thread is this session's own creation: when Codex answers that it holds no
   *  rollout for it, start a new thread in its place. Never set for a thread a resume proved. */
  supersedeIfUnsaved?: boolean
  permissionPolicy?: CodexStructuredPermissionPolicy
  /** The model the session chose; the thread opens on it so its first turn is not a switch. */
  model?: string
  env?: Record<string, string>
}

export type CodexStructuredSessionEvent =
  | {
      type: 'notification'
      sessionId: string
      threadId: string
      method: string
      params: unknown
      /** Host receipt time of a turn boundary; survives retry and deferral so a replay is not re-stamped. */
      observedAt?: number
      /** Highest dispatch sequence armed when this turn-start was first received. */
      dispatchSequenceAtReceipt?: number
    }
  | { type: 'server-request'; sessionId: string; threadId: string; method: string; params: unknown }
  | { type: 'provider-frame'; sessionId: string; threadId: string; kind: string; payload: unknown }
  | {
      type: 'prompt'
      sessionId: string
      threadId: string
      method: string
      params: unknown
      codexItemId: string
      promptKey: string
    }
  | StructuredAgentSessionEndedEvent
  /** Translator-only compatibility for callers that do not participate in host recovery. */
  | { type: 'ended'; sessionId: string; reason: string; observedAt?: number }

export type CodexStructuredSessionAdapterDeps = {
  resolveLaunch: (input: {
    identity: AgentSessionJournalIdentity
  }) => Promise<CodexStructuredLaunch>
  /** Host capability seam; production uses the native Windows process table. */
  isWindowsProcessStartTimeAvailable?: () => boolean
  onEvent?: (event: CodexStructuredSessionEvent) => void
  /** Where bookkeeping a close or exit does after the child is gone reports a failure. */
  logger?: StructuredAgentSessionLogger
  /** What the session's child work did, delivered after the journal handled the frame. */
  onChildWorkEvidence?: (sessionId: string, evidence: AgentChildWorkEvidence[]) => void
  /** A send admitted earlier: its identity once Codex echoes it, or its rejection when the turn
   *  Codex answered it into ended without taking it. */
  onDispatchSettledLate?: (
    input: { sessionId: string; clientMessageId: string } & (
      | { providerIdentity: AgentJournalItemIdentity }
      | ({ state: 'rejected' } & AgentJournalDispatchRejection)
    )
  ) => void
  /** Codex reported its thread not running with no turn open: a send whose
   *  dispatch was never answered is owed nothing after this. */
  onPrimaryThreadStoppedRunning?: (input: { sessionId: string }) => void
  openConnection?: typeof openCodexAppServerConnection
  readProcessStartTime?: (pid: number) => Promise<number | null>
  mintLinkId?: () => string
  mintAcquisitionGeneration?: () => string
  now?: () => number
  requestTimeoutMs?: number
  /** Host model catalog; sessions write their listings through and read back. */
  modelCatalog?: AgentModelCatalogStore
}

export type CodexSession = {
  connection: CodexAppServerConnection
  ended: boolean
  /** First observed child exit survives rejected settlement admission. */
  exitObservedAt?: number
  /** The close Orca began for this child: asked for, or forced as a death, and why. Whatever ends
   *  the child after it (that close, or the exit the connection reports meanwhile) keeps this. */
  orcaClose?: { requested: boolean; reason: Error }
  fence: number
  acquisitionGeneration: string
  threadId: string
  historyMode?: 'legacy' | 'paginated'
  /** Primary-thread turns Codex reported started and not yet ended, as read off the wire: what
   *  rewind waits out and what a Stop naming no turn interrupts when the journal shows none. */
  activeTurnIds?: Set<string>
  /** Stops waiting for the turn Codex answered a send into to open. */
  turnOpenWaits: CodexTurnOpenWaits
  dispatchPending?: boolean
  prompts: CodexAcquisitionWindow['prompts']
  options: Map<string, string>
  reportedOptions: {
    model?: string
    effort?: string
    serviceTier?: string | null
    serviceTierKnown?: true
  }
  /** Exact provider-advertised Fast request value for each discovered model. */
  fastModeTierByModel: Map<string, string>
  /** Absent when the adapter runs without a host catalog store (tests). */
  catalogAccess?: CodexSessionCatalogAccess
  /** Sends whose identity is still to be settled by the provider echo. */
  dispatchEchoes: CodexDispatchEchoes
  translator: CodexJournalTranslator | null
  /** Ephemeral roster behind the background-tasks strip; never durable state. */
  backgroundTasks: CodexBackgroundTaskTracker
  unbindReadingControl?: () => void
  /** Terminates this exact child as an unexpected death and enters host recovery. */
  forceCloseUnexpected?: (reason: Error) => Promise<boolean>
}

export function mintCodexAcquisitionGeneration(deps: CodexStructuredSessionAdapterDeps): string {
  return deps.mintAcquisitionGeneration?.() ?? randomUUID()
}

export function codexSessionLifecycle(
  fence: number,
  acquisitionGeneration: string
): Pick<CodexSession, 'ended' | 'fence' | 'acquisitionGeneration' | 'turnOpenWaits'> {
  return {
    ended: false,
    fence,
    acquisitionGeneration,
    turnOpenWaits: createCodexTurnOpenWaits()
  }
}

/** A child that exited while being acquired never becomes the session's. */
export function assertCodexConnectionOpen(
  connection: Pick<CodexAppServerConnection, 'closed'>,
  sessionId: string
): void {
  if (connection.closed) {
    throw new Error(`codex app-server for session ${sessionId} exited while being acquired`)
  }
}

export function requireLiveCodexSession(
  sessions: Map<string, CodexSession>,
  sessionId: string
): CodexSession {
  const session = sessions.get(sessionId)
  if (!session || session.ended) {
    throw new Error(`no live codex app-server for session ${sessionId}`)
  }
  return session
}

export type CodexAcquisitionAttempt = {
  window: CodexAcquisitionWindow
  cancelled: boolean
  exitProven: boolean
  finished: Promise<void>
  finish: () => void
}

export function createCodexAcquisitionAttempt(): CodexAcquisitionAttempt {
  let finish = (): void => {}
  const finished = new Promise<void>((resolve) => {
    finish = resolve
  })
  return {
    window: new CodexAcquisitionWindow(),
    cancelled: false,
    exitProven: false,
    finished,
    finish
  }
}

export class CodexAcquisitionRegistry {
  private readonly attempts = new Map<string, CodexAcquisitionAttempt>()
  private closing = false

  get size(): number {
    return this.attempts.size
  }

  start(sessionId: string): {
    previousAttempt: CodexAcquisitionAttempt | undefined
    attempt: CodexAcquisitionAttempt
  } {
    if (this.closing) {
      throw new Error('codex structured session adapter is closing')
    }
    const previousAttempt = this.attempts.get(sessionId)
    const attempt = createCodexAcquisitionAttempt()
    this.attempts.set(sessionId, attempt)
    return { previousAttempt, attempt }
  }

  assertCurrent(sessionId: string, attempt: CodexAcquisitionAttempt): void {
    if (this.closing || attempt.cancelled || this.attempts.get(sessionId) !== attempt) {
      throw new Error(`codex session ${sessionId} was superseded while being acquired`)
    }
  }

  get(sessionId: string): CodexAcquisitionAttempt | undefined {
    return this.attempts.get(sessionId)
  }

  deleteIfCurrent(sessionId: string, attempt: CodexAcquisitionAttempt): void {
    if (this.attempts.get(sessionId) === attempt) {
      this.attempts.delete(sessionId)
    }
  }

  restoreIfCurrent(
    sessionId: string,
    replacement: CodexAcquisitionAttempt,
    previous: CodexAcquisitionAttempt
  ): void {
    if (this.attempts.get(sessionId) === replacement) {
      this.attempts.set(sessionId, previous)
    }
  }

  async closeFailedAttempt(sessionId: string, attempt: CodexAcquisitionAttempt): Promise<boolean> {
    const stopped = (await attempt.window.connection?.close()) ?? true
    if (stopped) {
      attempt.exitProven = true
      this.deleteIfCurrent(sessionId, attempt)
    }
    return stopped
  }

  sessionIds(): IterableIterator<string> {
    return this.attempts.keys()
  }

  close(): void {
    this.closing = true
  }
}

export async function cancelCodexAcquisitionAttempt(
  attempt: CodexAcquisitionAttempt | undefined
): Promise<boolean> {
  if (!attempt) {
    return true
  }
  return cancelProcessAcquisition({
    cancel: () => {
      attempt.cancelled = true
    },
    connection: () => attempt.window.connection,
    exitProven: () => attempt.exitProven,
    finished: attempt.finished
  })
}