orca/src/shared/structured-agent-session-reducer.ts

import type {
  AgentJournalCursor,
  AgentJournalRenderItem,
  AgentJournalSubmission
} from './agent-session-journal-types'
import type {
  AgentSessionBackgroundTaskState,
  AgentSessionSlashCommand,
  AgentSessionHistoryPage,
  AgentSessionSubscribeEvent,
  AgentSessionTurnActivity
} from './agent-session-wire'
import { backgroundTaskStatesEqual } from './agent-session-background-task-state-equality'
import { agentJournalSubmissionKey } from './agent-session-journal-item-key'
import { readAgentJournalTurn } from './agent-session-turn-record'

/** The last host clock sample: `hostNow - receivedAt` is the client's skew from the host,
 *  which is what lets a client attaching mid-turn anchor its live counter on the real start. */
export type StructuredAgentHostClock = {
  hostNow: number
  receivedAt: number
}

export type StructuredAgentSessionState = {
  epoch: string | null
  cursor: AgentJournalCursor | null
  fence: number | null
  items: AgentJournalRenderItem[]
  submissions: AgentJournalSubmission[]
  /** Head-trim floor for `items`; paging back raises it so a live batch cannot undo the page. */
  retainedItemLimit: number
  hasOlder: boolean
  status: 'idle' | 'loading' | 'ready' | 'error'
  error?: string
  backgroundTasks?: AgentSessionBackgroundTaskState | null
  commands?: AgentSessionSlashCommand[] | null
  activity?: AgentSessionTurnActivity | null
  /** Absent until a frame from a host that stamps `hostNow` has been applied. */
  hostClock?: StructuredAgentHostClock
  /** Bumped per live batch that leaves a turn row's newest revision outside the window
   *  (dropped or trimmed), so a whole-journal answer derived from turn rows is asked for again. */
  unloadedTurnRevisions?: number
}

export type StructuredAgentSessionAction =
  | { type: 'loading' }
  | { type: 'error'; message: string }
  | { type: 'event'; event: AgentSessionSubscribeEvent }
  | { type: 'history-page'; page: AgentSessionHistoryPage }
  | { type: 'older-page'; requestedCursor: AgentJournalCursor; page: AgentSessionHistoryPage }

const MAX_RETAINED_SUBMISSIONS = 256
// Well above the renderer's initial read window (300) plus a page, so only genuinely
// long live sessions trim; anything trimmed is still reachable by paging older.
const MAX_RETAINED_ITEMS = 1024

export const EMPTY_STRUCTURED_AGENT_SESSION: StructuredAgentSessionState = {
  epoch: null,
  cursor: null,
  fence: null,
  items: [],
  submissions: [],
  retainedItemLimit: MAX_RETAINED_ITEMS,
  hasOlder: false,
  status: 'idle'
}

/** A frame without `hostNow` (older host) leaves the previous sample in place. */
function hostClockField(
  hostNow: number | undefined,
  receivedAt: number,
  previous: StructuredAgentHostClock | undefined
): { hostClock?: StructuredAgentHostClock } {
  const hostClock = hostNow !== undefined ? { hostNow, receivedAt } : previous
  return hostClock ? { hostClock } : {}
}

function replacePage(
  page: AgentSessionHistoryPage,
  fence: number | null,
  backgroundTasks?: AgentSessionBackgroundTaskState | null,
  activity?: AgentSessionTurnActivity | null
): StructuredAgentSessionState {
  return {
    epoch: page.epoch,
    cursor: page.liveCursor ?? page.window.nextCursor,
    fence,
    items: [...page.items].sort((left, right) => left.sequence - right.sequence),
    submissions: page.submissions,
    retainedItemLimit: Math.max(MAX_RETAINED_ITEMS, page.items.length),
    hasOlder: page.hasOlder,
    status: 'ready',
    activity: activity ?? null,
    ...(backgroundTasks !== undefined
      ? { backgroundTasks }
      : page.backgroundTasks !== undefined
        ? { backgroundTasks: page.backgroundTasks }
        : {})
  }
}

function mergeItems(
  current: readonly AgentJournalRenderItem[],
  incoming: readonly AgentJournalRenderItem[],
  removedIds: readonly string[]
): AgentJournalRenderItem[] {
  const removed = new Set(removedIds)
  const byId = new Map(
    current.filter((item) => !removed.has(item.itemId)).map((item) => [item.itemId, item])
  )
  for (const item of incoming) {
    const prior = byId.get(item.itemId)
    if (!prior || item.revision >= prior.revision) {
      byId.set(item.itemId, item)
    }
  }
  return [...byId.values()].sort((left, right) => left.sequence - right.sequence)
}

/**
 * Live rows the loaded window can take. The window is a contiguous suffix of the
 * journal, and its oldest row is the load-older anchor. A revision of a row older
 * than the window keeps that row's original sequence, so admitting it would move
 * the anchor below the window and paging `before` it would skip every row between.
 * The journal keeps the revision; the page reader serves it once the window
 * reaches the row. With nothing older on the host the window is the whole journal
 * and a row below the head (a revived tombstone) leaves no hole, so it is admitted.
 */
function liveItemsWithinWindow(
  state: StructuredAgentSessionState,
  incoming: readonly AgentJournalRenderItem[]
): readonly AgentJournalRenderItem[] {
  const head = state.items[0]
  if (!head || !state.hasOlder) {
    return incoming
  }
  return incoming.filter((item) => item.sequence >= head.sequence)
}

function trimRetainedItems(
  items: AgentJournalRenderItem[],
  limit: number
): AgentJournalRenderItem[] {
  return items.length <= limit ? items : items.slice(items.length - limit)
}

function mergeSubmissions(
  current: readonly AgentJournalSubmission[],
  incoming: readonly AgentJournalSubmission[],
  items: readonly AgentJournalRenderItem[]
): AgentJournalSubmission[] {
  const byId = new Map(current.map((submission) => [submission.clientMessageId, submission]))
  for (const submission of incoming) {
    byId.set(submission.clientMessageId, submission)
  }
  const sorted = [...byId.values()].sort((left, right) => left.submittedAt - right.submittedAt)
  const itemIds = new Set(
    items
      .filter((item) => item.body.kind === 'message' && item.body.role === 'user')
      .map((item) => item.itemId)
  )
  // Loaded user messages need their provider alias for durable turn attribution.
  return sorted.filter(
    (submission, index) =>
      index >= sorted.length - MAX_RETAINED_SUBMISSIONS ||
      itemIds.has(agentJournalSubmissionKey(submission.clientMessageId))
  )
}

/** `receivedAt` is the client clock at apply time; callers pass it so the reducer stays pure. */
export function reduceStructuredAgentSession(
  state: StructuredAgentSessionState,
  action: StructuredAgentSessionAction,
  receivedAt: number = Date.now()
): StructuredAgentSessionState {
  if (action.type === 'loading') {
    // Keep the last transcript visible while a reconnect rehydrates the stream.
    return { ...state, status: 'loading', error: undefined }
  }
  if (action.type === 'error') {
    return { ...state, status: 'error', error: action.message }
  }
  if (action.type === 'history-page') {
    return {
      ...replacePage(action.page, action.page.fence ?? null, state.backgroundTasks, state.activity),
      commands: state.commands,
      ...hostClockField(action.page.hostNow, receivedAt, state.hostClock)
    }
  }
  if (action.type === 'older-page') {
    const requested = action.requestedCursor
    if (state.epoch !== requested.epoch || action.page.epoch !== requested.epoch) {
      return state
    }
    const head = state.items[0]
    // A live batch head-trimmed past the anchor while this read was in flight, so the
    // page no longer abuts the retained window; merging it would leave a silent hole.
    // The caller re-anchors on the new head and asks again.
    if (head && head.sequence > requested.sequence) {
      return state
    }
    const items = mergeItems(state.items, action.page.items, action.page.removedItemIds)
    return {
      ...state,
      items,
      retainedItemLimit: Math.max(state.retainedItemLimit, items.length),
      submissions: mergeSubmissions(state.submissions, action.page.submissions, items),
      hasOlder: action.page.hasOlder,
      ...hostClockField(action.page.hostNow, receivedAt, state.hostClock)
    }
  }
  const event = action.event
  if (event.type === 'end') {
    return state
  }
  if (event.type === 'snapshot' || event.type === 'reset') {
    return {
      ...replacePage(event.page, event.fence, event.backgroundTasks, event.activity),
      commands: event.commands,
      ...hostClockField(event.hostNow, receivedAt, state.hostClock)
    }
  }
  if (state.epoch !== event.batch.cursor.epoch) {
    return state
  }
  if (state.cursor && event.batch.cursor.sequence < state.cursor.sequence) {
    return state
  }
  const backgroundTasks =
    event.backgroundTasks !== undefined ? event.backgroundTasks : state.backgroundTasks
  const activity = event.activity !== undefined ? event.activity : state.activity
  const liveItems = liveItemsWithinWindow(state, event.batch.items)
  const journalUnchanged =
    liveItems.length === 0 &&
    event.batch.removedItemIds.length === 0 &&
    event.batch.submissions.length === 0
  if (
    event.batch.cursor.sequence === state.cursor?.sequence &&
    journalUnchanged &&
    (event.fence === undefined || event.fence === state.fence) &&
    (event.commands === undefined || event.commands === state.commands) &&
    backgroundTaskStatesEqual(backgroundTasks, state.backgroundTasks) &&
    activity?.turnId === state.activity?.turnId &&
    activity?.text === state.activity?.text &&
    state.status === 'ready' &&
    state.error === undefined
  ) {
    return state
  }
  const merged = journalUnchanged
    ? state.items
    : mergeItems(state.items, liveItems, event.batch.removedItemIds)
  const items = trimRetainedItems(merged, state.retainedItemLimit)
  const outsideWindow = [
    ...(liveItems.length < event.batch.items.length
      ? event.batch.items.filter((item) => !liveItems.includes(item))
      : []),
    ...merged.slice(0, merged.length - items.length)
  ]
  const lostTurnRow = outsideWindow.some((item) => readAgentJournalTurn(item.body) !== null)
  return {
    ...state,
    cursor: event.batch.cursor,
    fence: event.fence ?? state.fence,
    items,
    // A trim leaves older items behind the cursor, so paging must stay offered.
    hasOlder: items.length < merged.length ? true : state.hasOlder,
    submissions:
      event.batch.submissions.length === 0 && event.batch.removedItemIds.length === 0
        ? state.submissions
        : mergeSubmissions(state.submissions, event.batch.submissions, items),
    status: 'ready',
    error: undefined,
    commands: event.commands !== undefined ? event.commands : state.commands,
    ...(backgroundTasks !== undefined ? { backgroundTasks } : {}),
    ...(activity !== undefined ? { activity } : {}),
    ...(lostTurnRow ? { unloadedTurnRevisions: (state.unloadedTurnRevisions ?? 0) + 1 } : {}),
    ...hostClockField(event.hostNow, receivedAt, state.hostClock)
  }
}

export function oldestStructuredAgentSessionCursor(
  state: StructuredAgentSessionState
): AgentJournalCursor | null {
  const oldest = state.items[0]
  return state.epoch && oldest ? { epoch: state.epoch, sequence: oldest.sequence } : null
}