orca/src/main/codex/codex-structured-item-streams.ts

import {
  AGENT_JOURNAL_THREAD_SCOPE,
  type AgentJournalRowAttribution
} from '../../shared/agent-session-journal-types'
import { createAgentSessionDeltaCoalescer } from '../native-chat/agent-session-wire/agent-session-delta-coalescer'
import { CodexItemStreamRetention } from './codex-item-stream-retention'
import { appendCodexItemAndPublish } from './codex-structured-journal-sink'
import { codexJournalItem, codexStreamingJournalItem } from './codex-structured-item-translation'
import { withJournalReasoningLifecycle } from '../native-chat/agent-session-journal/journal-reasoning-row'
import {
  codexStructuredItemKey,
  MAX_CODEX_ITEM_STREAM_PENDING_PATCHES,
  MAX_CODEX_ITEM_STREAM_PENDING_PATCH_BYTES,
  MAX_CODEX_ITEM_STREAM_RETAINED_BYTES,
  codexItemStreamState,
  pendingPatchBytes
} from './codex-structured-item-stream-bounds'
import {
  CODEX_ITEM_STREAM_TYPES,
  codexPatchChangeBytes,
  PATCH_UPDATED_METHOD,
  readCodexItemStreamRecord,
  readCodexItemStreamString,
  REASONING_PART_METHOD,
  TERMINAL_INTERACTION_METHOD
} from './codex-structured-item-stream-events'
import type {
  CodexItemStreamDeps,
  CodexItemStreamState,
  CodexPendingItemPatch,
  CodexStructuredItemStreamAdmission,
  CodexStructuredItemStreams
} from './codex-structured-item-stream-contracts'
export type {
  CodexStructuredItemStreamAdmission,
  CodexStructuredItemStreamHandleResult,
  CodexStructuredItemStreams
} from './codex-structured-item-stream-contracts'
export { codexStructuredItemKey } from './codex-structured-item-stream-bounds'
export {
  MAX_CODEX_ITEM_STREAM_PENDING_PATCH_BYTES,
  MAX_CODEX_ITEM_STREAM_RETAINED_BYTES,
  MAX_CODEX_ITEM_STREAM_PENDING_PATCHES,
  MAX_CODEX_ITEM_STREAM_STATES
} from './codex-structured-item-stream-bounds'
/** Delta-only item ids are provider input; retain only a deterministic recent window. */

export function createCodexStructuredItemStreams(
  deps: CodexItemStreamDeps
): CodexStructuredItemStreams {
  const states = new CodexItemStreamRetention(deps.maxMetadataBytes)
  const checkpointLengths = new Map<string, number>()
  const pendingCheckpoints = new Set<string>()
  // Which thread and turn produced each stream, resolved to its attribution per append
  // so a parent learned after the first checkpoint still reaches the row.
  const producers = new Map<string, { threadId: string; turnId: string | null }>()
  const attributionOf = (key: string): AgentJournalRowAttribution => {
    const producer = producers.get(key)
    return producer
      ? deps.attributionFor(producer.threadId, producer.turnId)
      : { turnScope: AGENT_JOURNAL_THREAD_SCOPE }
  }
  // Patch updates are authoritative item snapshots. Keep the latest rejected
  // snapshot until the journal admits it; unlike streamed deltas, there is no
  // coalescer timer to retry these events for us.
  const pendingPatches = new Map<string, CodexPendingItemPatch>()
  let retainedPatchBytes = 0

  const forgetState = (key: string): void => {
    coalescer.forget(key)
    states.forget(key)
    checkpointLengths.delete(key)
    pendingCheckpoints.delete(key)
    producers.delete(key)
    const pending = pendingPatches.get(key)
    if (pending) {
      retainedPatchBytes = Math.max(0, retainedPatchBytes - pendingPatchBytes(pending))
      pendingPatches.delete(key)
    }
  }

  const trimStates = (): void => {
    while (states.overCapacity) {
      const oldest = states.oldestEvictable()
      if (typeof oldest !== 'string') {
        break
      }
      const pending = coalescer.snapshot(oldest)
      if (pending && pending.text.length > 0 && !persist(oldest, pending.text, true)) {
        // Keep the state (and its buffered text) until the sink recovers. A
        // bounded map is preferable to silently losing streamed output.
        break
      }
      forgetState(oldest)
    }
  }

  const trimPendingPatches = (): void => {
    while (pendingPatches.size > MAX_CODEX_ITEM_STREAM_PENDING_PATCHES) {
      const oldest = pendingPatches.keys().next().value
      if (typeof oldest !== 'string') {
        break
      }
      const pending = pendingPatches.get(oldest)
      if (pending) {
        retainedPatchBytes = Math.max(0, retainedPatchBytes - pendingPatchBytes(pending))
      }
      pendingPatches.delete(oldest)
    }
  }

  const append = (key: string, state: CodexItemStreamState, text: string): boolean => {
    const translated = codexStreamingJournalItem(state.item, text)
    if (!translated.body) {
      return true
    }
    // A stream only ever carries an item that has not completed yet.
    const body = withJournalReasoningLifecycle(translated.body, { state: 'running' })
    return appendCodexItemAndPublish(deps.sink, state.identity, body, {
      ...attributionOf(key),
      ...(state.startedAt === undefined ? {} : { observedAt: state.startedAt })
    }).accepted
  }

  const shouldPersist = (key: string, textLength: number): boolean => {
    const checkpointLength = checkpointLengths.get(key) ?? 0
    const nextLength = Math.max(checkpointLength + 32, Math.ceil(checkpointLength * 1.125))
    return !(checkpointLength > 0 && textLength < nextLength)
  }

  const persist = (key: string, text: string, force: boolean): boolean => {
    if (!force && !shouldPersist(key, text.length)) {
      return true
    }
    const state = states.get(key)
    if (state && append(key, state, text)) {
      checkpointLengths.set(key, text.length)
      pendingCheckpoints.delete(key)
      return true
    }
    return false
  }

  const coalescer = createAgentSessionDeltaCoalescer({
    windowMs: deps.coalesceMs,
    maxRetainedBytes: deps.maxRetainedBytes,
    maxTotalRetainedBytes: deps.maxTotalRetainedBytes,
    isProtected: (key) => states.isPersistent(key),
    schedule: deps.schedule,
    shouldEmit: shouldPersist,
    emit: (key, text) => {
      return persist(key, text, false)
    }
  })

  const ensureState = (
    threadId: string,
    itemId: string,
    type: string,
    params: unknown
  ): CodexItemStreamState | null => {
    const key = codexStructuredItemKey(threadId, itemId)
    const existing = states.get(key)
    if (existing) {
      return existing
    }
    const item = { type, id: itemId }
    const turnId = deps.turnIdFor(threadId, params)
    const state = { item, identity: deps.identityFor(threadId, turnId, item) }
    if (!states.retain(key, state)) {
      return null
    }
    producers.set(key, { threadId, turnId })
    trimStates()
    return state
  }

  const flush = (): boolean => {
    let flushed = coalescer.flushAll()
    for (const key of pendingCheckpoints) {
      const snapshot = coalescer.snapshot(key)
      if (snapshot && checkpointLengths.get(key) !== snapshot.text.length) {
        flushed = persist(key, snapshot.text, true) && flushed
      } else {
        pendingCheckpoints.delete(key)
      }
    }
    for (const key of pendingPatches.keys()) {
      flushed = flushPatch(key).accepted && flushed
    }
    return flushed
  }

  const flushPatch = (key: string): CodexStructuredItemStreamAdmission => {
    const pending = pendingPatches.get(key)
    if (!pending) {
      return { accepted: true }
    }
    const admission = appendCodexItemAndPublish(
      deps.sink,
      pending.identity,
      pending.body,
      attributionOf(key)
    )
    if (!admission.accepted) {
      return admission
    }
    retainedPatchBytes = Math.max(0, retainedPatchBytes - pendingPatchBytes(pending))
    pendingPatches.delete(key)
    return { accepted: true }
  }

  return {
    get persistentCount() {
      return states.persistentSize
    },
    canTrack: (threadId, item, identity) =>
      states.canRetain(
        codexStructuredItemKey(threadId, item.id),
        codexItemStreamState(item, identity)
      ),
    track: (threadId, turnId, item, identity, startedAt) => {
      const key = codexStructuredItemKey(threadId, item.id)
      if (!states.retain(key, codexItemStreamState(item, identity, startedAt))) {
        return false
      }
      producers.set(key, { threadId, turnId })
      trimStates()
      return true
    },
    handle: (threadId, method, params) => {
      const paramsRecord = readCodexItemStreamRecord(params)
      const itemId = readCodexItemStreamString(paramsRecord, 'itemId')
      if (method === PATCH_UPDATED_METHOD) {
        if (!itemId || !Array.isArray(paramsRecord.changes)) {
          return { handled: true, admission: { accepted: true } }
        }
        if (
          codexPatchChangeBytes(paramsRecord.changes) > MAX_CODEX_ITEM_STREAM_PENDING_PATCH_BYTES
        ) {
          return { handled: true, admission: { accepted: false, reason: 'backpressure' } }
        }
        const key = codexStructuredItemKey(threadId, itemId)
        const streamFlushed = coalescer.flush(key)
        const state = ensureState(threadId, itemId, 'fileChange', params)
        if (!state) {
          return { handled: true, admission: { accepted: false, reason: 'failed' } }
        }
        state.item = { ...state.item, changes: paramsRecord.changes }
        const translated = codexJournalItem(state.item)
        if (translated.body) {
          const nextPending: CodexPendingItemPatch = {
            identity: state.identity,
            body: translated.body
          }
          const previous = pendingPatches.get(key)
          const previousBytes = previous ? pendingPatchBytes(previous) : 0
          const nextBytes = pendingPatchBytes(nextPending)
          if (
            nextBytes > MAX_CODEX_ITEM_STREAM_PENDING_PATCH_BYTES ||
            retainedPatchBytes - previousBytes + nextBytes > MAX_CODEX_ITEM_STREAM_RETAINED_BYTES
          ) {
            return { handled: true, admission: { accepted: false, reason: 'backpressure' } }
          }
          retainedPatchBytes = Math.max(0, retainedPatchBytes - previousBytes) + nextBytes
          pendingPatches.set(key, nextPending)
          trimPendingPatches()
          if (streamFlushed) {
            const admission = flushPatch(key)
            if (!admission.accepted) {
              return { handled: true, admission }
            }
          }
        }
        return { handled: true, admission: { accepted: true } }
      }
      if (method === TERMINAL_INTERACTION_METHOD) {
        return { handled: true, admission: { accepted: true } }
      }
      const type = CODEX_ITEM_STREAM_TYPES[method as keyof typeof CODEX_ITEM_STREAM_TYPES]
      if (!type && method !== REASONING_PART_METHOD) {
        return { handled: false, admission: { accepted: true } }
      }
      if (!itemId) {
        return { handled: true, admission: { accepted: true } }
      }
      const state = ensureState(threadId, itemId, type ?? 'reasoning', params)
      if (!state) {
        return { handled: true, admission: { accepted: false, reason: 'failed' } }
      }
      const delta = method === REASONING_PART_METHOD ? '\n' : paramsRecord.delta
      if (typeof delta === 'string') {
        const key = codexStructuredItemKey(threadId, state.item.id)
        pendingCheckpoints.add(key)
        const accepted = coalescer.append(key, delta)
        if (!accepted) {
          return { handled: true, admission: { accepted: false, reason: 'backpressure' } }
        }
      }
      return { handled: true, admission: { accepted: true } }
    },
    forget: (threadId, itemId) => {
      const key = codexStructuredItemKey(threadId, itemId)
      forgetState(key)
    },
    flush,
    dispose: () => {
      coalescer.dispose()
      states.clear()
      checkpointLengths.clear()
      pendingCheckpoints.clear()
      producers.clear()
      pendingPatches.clear()
      retainedPatchBytes = 0
    },
    snapshot: (threadId, itemId) => {
      const key = codexStructuredItemKey(threadId, itemId)
      return coalescer.snapshot(key)
    }
  }
}