orca/src/main/codex/codex-background-task-tracker.ts

import type {
  AgentSessionBackgroundTask,
  AgentSessionBackgroundTaskState
} from '../../shared/agent-session-wire'
import type { AgentChildWorkEvidence } from '../../shared/agent-status-child-work-evidence'
import {
  readCodexBackgroundTaskFrame,
  type CodexBackgroundTaskEvent
} from './codex-background-task-frames'
import { CodexSubagentExecutions } from './codex-subagent-executions'
import { CodexBackgroundCommandTracker } from './codex-background-command-tracker'
import { CodexChildWorkEvidence } from './codex-child-work-evidence'
import type { CodexAbandonedCommand } from './codex-prompt-registry'
import type { CodexStructuredSessionAdapterDeps } from './codex-structured-session-state'
import { boundSubagentField } from '../native-chat/agent-session-journal/journal-subagent-group-body'

/** Where a session's child-work evidence goes, and the host clock that stamps it. */
export type CodexChildWorkSink = {
  deliver: (evidence: AgentChildWorkEvidence[]) => void
  now: () => number
}

export function codexChildWorkSink(
  sessionId: string,
  deps: Pick<CodexStructuredSessionAdapterDeps, 'onChildWorkEvidence' | 'now'>
): CodexChildWorkSink {
  return {
    deliver: (evidence) => deps.onChildWorkEvidence?.(sessionId, evidence),
    now: () => deps.now?.() ?? Date.now()
  }
}

/** What a probe learned: `unknown` when the app-server gave no answer, so the next chance asks again. */
export type CodexTerminalStopSupport = 'supported' | 'unsupported' | 'unknown'

/** Projects the same child execution facts the durable roster consumes. */
export class CodexBackgroundTaskTracker {
  private publishedFingerprint = '[]'
  private publishedState: AgentSessionBackgroundTaskState | null = null
  private readonly commands: CodexBackgroundCommandTracker
  private readonly childWork: CodexChildWorkEvidence
  /** Whether this app-server can terminate a background command, settled by its first answer. */
  private terminalStops: CodexTerminalStopSupport | 'probing' = 'unknown'

  constructor(
    private readonly primaryThreadId: string,
    private readonly executions = new CodexSubagentExecutions(),
    private readonly childWorkSink?: CodexChildWorkSink
  ) {
    this.commands = new CodexBackgroundCommandTracker(primaryThreadId)
    this.childWork = new CodexChildWorkEvidence(
      primaryThreadId,
      executions,
      (threadId) => this.commands.threadCommands(threadId),
      () => this.stopsTerminals
    )
  }

  get stopsTerminals(): boolean {
    return this.terminalStops === 'supported'
  }

  /** True when a running command has a process a stop could name and the app-server has not
   *  answered yet: the caller probes then, and settles the answer here. */
  beginTerminalStopProbe(): boolean {
    if (this.terminalStops !== 'unknown' || !this.commands.holdsProcess) {
      return false
    }
    this.terminalStops = 'probing'
    return true
  }

  /** On yes, restates every running command so its record says it can be stopped; callers publish
   *  the child work after. */
  settleTerminalStopProbe(support: CodexTerminalStopSupport): void {
    if (this.terminalStops !== 'probing') {
      return
    }
    this.terminalStops = support
    if (support === 'supported') {
      this.childWork.restateCommands(this.commands.liveCommands())
    }
  }

  /** The processes behind the named tasks that a stop reaches. */
  backgroundProcesses(taskIds: readonly string[]): { threadId: string; processId: string }[] {
    return this.stopsTerminals ? this.commands.backgroundProcesses(taskIds) : []
  }

  /** What a Stop of the named sub-agents reaches: the run each one and each of its descendants is
   *  working on, and, when this app-server can terminate them, the processes their threads hold.
   *  A sub-agent row's providerId is its thread id. */
  subagentStopTargets(taskIds: readonly string[]): {
    turns: { threadId: string; turnId: string }[]
    terminals: { threadId: string; processId: string }[]
  } {
    const named = new Set(taskIds)
    const lineage = this.executions
      .workingChildren()
      .filter((child) => named.has(child.agentThreadId))
      .flatMap((child) => this.executions.lineage(child.agentThreadId))
    const threads = [...new Set(lineage.map((child) => child.agentThreadId))]
    return {
      turns: threads.flatMap((threadId) => {
        const execution = this.executions.find(threadId)?.execution
        return execution?.state === 'working' ? [{ threadId, turnId: execution.turnId }] : []
      }),
      terminals: this.stopsTerminals
        ? threads.flatMap((threadId) =>
            this.commands
              .threadCommands(threadId)
              .flatMap((command) =>
                command.type === 'started' && command.processId !== undefined
                  ? [{ threadId, processId: command.processId }]
                  : []
              )
          )
        : []
    }
  }

  /** Whether the sub-agent is still working on this run. */
  runsSubagentTurn(threadId: string, turnId: string): boolean {
    const execution = this.executions.find(threadId)?.execution
    return execution?.state === 'working' && execution.turnId === turnId
  }

  get state(): AgentSessionBackgroundTaskState | null {
    // Journal admission precedes observe; readers must not see its pending facts.
    return this.publishedState
  }

  canObserve(event: CodexBackgroundTaskEvent): boolean {
    return this.commands.canObserve(event)
  }

  /** `unapproved`: commands whose approval the journal dropped with this frame's turn ending. */
  observe(
    event: CodexBackgroundTaskEvent,
    unapproved: readonly CodexAbandonedCommand[] = []
  ): boolean {
    const itemEvent = event.method === 'item/started' || event.method === 'item/completed'
    const command = itemEvent ? this.commands.observe(event) : null
    const frame = readCodexBackgroundTaskFrame(event, this.primaryThreadId)
    const commands = [
      ...unapproved.flatMap((abandoned) => this.commands.endUnapproved(abandoned) ?? []),
      ...(frame?.kind === 'turn' && frame.state !== 'working'
        ? this.commands.endTurn(frame.threadId, frame.turnId)
        : []),
      ...(event.method === 'thread/closed'
        ? this.commands.endThread(event.threadId)
        : command
          ? [command]
          : [])
    ]
    if (frame?.kind === 'subagents') {
      for (const child of frame.children) {
        this.executions.register(
          child.agentThreadId,
          child.label,
          child.parentTurnId,
          child.spawnerThreadId
        )
      }
    } else if (frame?.kind === 'thread-closed') {
      this.executions.closeThread(frame.threadId)
    } else if (frame && frame.threadId !== this.primaryThreadId) {
      this.executions.observeTurn(frame.threadId, frame.turnId, frame.state)
    }
    this.childWork.observe(event, frame, commands)
    if (!frame) {
      return itemEvent || commands.length > 0 ? this.refresh() : false
    }
    // A primary-turn frame only prompts a republish: turn end reveals children,
    // it never settles them. Codex `spawn_agent` children keep reporting well
    // past their parent turn, so nothing here may sweep the roster.
    return this.refresh()
  }

  clear(): boolean {
    this.executions.clear()
    this.childWork.clear(this.commands.clear())
    return this.refresh()
  }

  /** Everything the frames observed since the last drain said about the session's child work. */
  drainChildWorkEvidence(observedAt: number): AgentChildWorkEvidence[] {
    return this.childWork.drain(observedAt)
  }

  /** Hand the pending evidence to the host. Callers run this after the journal wrote the frame
   *  and the parent's own row republished, so a child record never lands ahead of either. */
  publishChildWork(): void {
    // Drained even with no sink, so undelivered evidence never accumulates.
    const evidence = this.drainChildWorkEvidence(this.childWorkSink?.now() ?? Date.now())
    if (evidence.length > 0) {
      this.childWorkSink?.deliver(evidence)
    }
  }

  private tasks(): AgentSessionBackgroundTask[] {
    const children = this.executions.workingChildren()
    const agents: AgentSessionBackgroundTask[] = children.map((child, index) => ({
      id: `codex-agent:${child.agentThreadId}`,
      kind: 'agent',
      ...(child.label ? { description: boundSubagentField(child.label, index) } : {})
    }))
    return [
      ...agents,
      ...this.commands.tasks(new Set(children.map((child) => child.agentThreadId)), (threadId) =>
        this.executions.label(threadId)
      )
    ]
  }

  private refresh(): boolean {
    const tasks = this.tasks()
    const fingerprint = JSON.stringify(tasks)
    if (fingerprint === this.publishedFingerprint) {
      return false
    }
    this.publishedFingerprint = fingerprint
    this.publishedState = tasks.length
      ? { state: 'monitoring', tasks, supportsStopAll: false }
      : null
    return true
  }
}