my-pi/pi-permission-system/src/auto-review/transcript.ts

import type { SessionEntry } from '@earendil-works/pi-coding-agent'

const MAX_MESSAGE_ENTRY_TOKENS = 2_000
const MAX_TOOL_ENTRY_TOKENS = 1_000
const TRUSTED_USER_INTERACTION_TOOLS = new Set(['ask_user_question', 'plan_mode_question'])

/** Transcript budgets, as configured on `autoReview.transcript`. */
export interface TranscriptBudgets {
  readonly messageTokens: number
  readonly toolTokens: number
  readonly maxUntrustedEntries: number
}

export const DEFAULT_TRANSCRIPT_BUDGETS: TranscriptBudgets = {
  messageTokens: 10_000,
  toolTokens: 10_000,
  maxUntrustedEntries: 40,
}

/**
 * One tool-result summary record, as returned by the injected lookup.
 *
 * Read structurally: the fields mirror pi-compaction's `ToolHistoryRecord`,
 * but the transcript must not import pi-compaction code — the record reaches
 * it through a lookup function the caller resolves (or leaves absent).
 */
export interface ToolResultSummaryRecord {
  readonly toolCallId: string
  readonly resultTimestamp?: unknown
  readonly semanticSummary?: unknown
  readonly deterministicSummary?: unknown
}

/**
 * Resolve summaries for a batch of tool-call ids, in one call per render.
 *
 * The caller binds the session id; the transcript asks for every toolResult
 * entry's id at once so the lookup performs a single service query.
 */
export type ToolResultSummaryQuery = (
  toolCallIds: readonly string[],
) => readonly ToolResultSummaryRecord[]

/** Options accepted by {@link renderTranscript}. */
export interface TranscriptRenderOptions {
  readonly budgets?: TranscriptBudgets
  /** Optional summary lookup; absent means no substitution. */
  readonly querySummaries?: ToolResultSummaryQuery
}

type TranscriptKind = 'user' | 'user_interaction' | 'assistant' | 'tool'

export interface TranscriptEntry {
  index: number
  /** The SessionEntry `id` this transcript entry was rendered from. */
  entryId: string
  kind: TranscriptKind
  label: string
  text: string
  /** Set on plain toolResult entries, so a summary lookup can match them. */
  toolCallId?: string
  /** The message's own timestamp, read structurally; may be absent. */
  timestamp?: unknown
  truncated?: boolean
}

export interface TranscriptStats {
  entriesRetained: number
  entriesOmitted: number
  entriesTruncated: number
  toolResultsSummarized: number
  latestTrustedEntryRetained: boolean
}

export interface RenderedTranscript {
  entries: string[]
  omittedCount: number
  stats: TranscriptStats
  /**
   * The SessionEntry `id` of the newest trusted (`user` or `user_interaction`)
   * entry in this render, or `null` when the branch has none. Lets a caller
   * key per-request caches to the conversation point that authorized the ask.
   */
  latestTrustedEntryId: string | null
}

interface ContentBlock {
  type?: unknown
  id?: unknown
  text?: unknown
  thinking?: unknown
  name?: unknown
  toolName?: unknown
  arguments?: unknown
}

interface MessageLike {
  role?: unknown
  content?: unknown
  command?: unknown
  output?: unknown
  summary?: unknown
  toolCallId?: unknown
  toolName?: unknown
  isError?: unknown
  details?: unknown
  timestamp?: unknown
}

interface UserInteractionDetails {
  cancelled?: unknown
  answers?: unknown
}

interface UserInteractionAnswer {
  question?: unknown
  answer?: unknown
  selected?: unknown
  notes?: unknown
}

/**
 * A conservative token estimate: ASCII characters count as a quarter token
 * each, every non-ASCII code point counts as a full token.
 *
 * The plain `length / 4` heuristic under-counts CJK by roughly 2–4x, which
 * would silently overflow every budget this module enforces.
 */
export function estimateTokens(text: string): number {
  let ascii = 0
  let nonAscii = 0
  for (const character of text) {
    if ((character.codePointAt(0) ?? 0) > 127) {
      nonAscii += 1
    } else {
      ascii += 1
    }
  }
  return Math.ceil(ascii / 4) + nonAscii
}

function truncateToCharacters(text: string, maxCharacters: number): string {
  if (text.length <= maxCharacters) {
    return text
  }
  const tag = '\n...[truncated]...\n'
  const available = Math.max(0, maxCharacters - tag.length)
  const headLength = Math.floor(available * 0.7)
  const tailLength = available - headLength
  return `${text.slice(0, headLength)}${tag}${text.slice(-tailLength)}`
}

/**
 * Largest head/tail truncation whose estimate fits `maxTokens`.
 *
 * Character budgets and token budgets diverge on non-ASCII text, so the bound
 * is searched: binary search the character count while the estimate — not the
 * length — decides whether a candidate fits.
 */
export function truncateToApproximateTokens(text: string, maxTokens: number): string {
  if (estimateTokens(text) <= maxTokens) {
    return text
  }
  let lower = 0
  let upper = text.length
  let result = truncateToCharacters(text, 0)
  while (lower <= upper) {
    const middle = Math.floor((lower + upper) / 2)
    const candidate = truncateToCharacters(text, middle)
    if (estimateTokens(candidate) <= maxTokens) {
      result = candidate
      lower = middle + 1
    } else {
      upper = middle - 1
    }
  }
  return result
}

function serializeUnknown(value: unknown): string {
  if (typeof value === 'string') {
    return value
  }
  try {
    return JSON.stringify(value)
  } catch {
    return String(value)
  }
}

function normalizeAnswer(value: unknown): unknown {
  if (value === null || ['string', 'number', 'boolean'].includes(typeof value)) {
    return value
  }
  if (Array.isArray(value)) {
    return value.map(normalizeAnswer)
  }
  return serializeUnknown(value)
}

function textFromContent(content: unknown): string {
  if (typeof content === 'string') {
    return content
  }
  if (!Array.isArray(content)) {
    return serializeUnknown(content)
  }
  return content
    .map(rawBlock => {
      const block = rawBlock as ContentBlock
      if (block.type === 'text' && typeof block.text === 'string') {
        return block.text
      }
      if (block.type === 'image') {
        return '[image omitted]'
      }
      return ''
    })
    .filter(Boolean)
    .join('\n')
}

function normalizedAnswerEvidence(answer: UserInteractionAnswer): unknown | undefined {
  const primaryAnswer =
    answer.answer !== undefined && answer.answer !== null
      ? normalizeAnswer(answer.answer)
      : Array.isArray(answer.selected) && answer.selected.length > 0
        ? answer.selected.map(normalizeAnswer)
        : undefined
  const notes = typeof answer.notes === 'string' && answer.notes.length > 0 ? answer.notes : undefined
  if (primaryAnswer === undefined) {
    return notes
  }
  if (notes === undefined) {
    return primaryAnswer
  }
  return { selection: primaryAnswer, notes }
}

function normalizedUserInteraction(
  message: MessageLike,
  interactionToolCalls: ReadonlyMap<string, string>,
): TranscriptEntry['text'] | undefined {
  const name = typeof message.toolName === 'string' ? message.toolName : undefined
  const toolCallId = typeof message.toolCallId === 'string' ? message.toolCallId : undefined
  if (
    name === undefined ||
    toolCallId === undefined ||
    !TRUSTED_USER_INTERACTION_TOOLS.has(name) ||
    interactionToolCalls.get(toolCallId) !== name ||
    message.isError !== false
  ) {
    return undefined
  }
  if (message.details === null || typeof message.details !== 'object') {
    return undefined
  }

  const details = message.details as UserInteractionDetails
  if (details.cancelled !== false || !Array.isArray(details.answers) || details.answers.length === 0) {
    return undefined
  }

  const answers: Array<{ question: string; answer: unknown }> = []
  for (const rawAnswer of details.answers) {
    if (rawAnswer === null || typeof rawAnswer !== 'object') {
      return undefined
    }
    const answer = rawAnswer as UserInteractionAnswer
    const answerEvidence = normalizedAnswerEvidence(answer)
    if (typeof answer.question !== 'string' || answer.question.length === 0 || answerEvidence === undefined) {
      return undefined
    }
    answers.push({
      question: answer.question,
      answer: answerEvidence,
    })
  }
  return JSON.stringify(answers)
}

function assistantEntries(
  message: MessageLike,
  index: number,
  entryId: string,
  interactionToolCalls: Map<string, string>,
): TranscriptEntry[] {
  const content = Array.isArray(message.content) ? message.content : []
  const text = textFromContent(message.content)
  const entries: TranscriptEntry[] = []
  if (text) {
    entries.push({ index, entryId, kind: 'assistant', label: 'assistant', text })
  }
  for (const rawBlock of content) {
    const block = rawBlock as ContentBlock
    if (block.type !== 'toolCall') {
      continue
    }
    const name =
      typeof block.name === 'string' ? block.name : typeof block.toolName === 'string' ? block.toolName : 'unknown'
    if (typeof block.id === 'string' && TRUSTED_USER_INTERACTION_TOOLS.has(name)) {
      interactionToolCalls.set(block.id, name)
    }
    entries.push({
      index,
      entryId,
      kind: 'tool',
      label: `tool:${name}`,
      text: serializeUnknown(block.arguments),
    })
  }
  return entries
}

function entriesFromMessage(
  message: MessageLike,
  index: number,
  entryId: string,
  interactionToolCalls: Map<string, string>,
): TranscriptEntry[] {
  switch (message.role) {
    case 'user': {
      const text = textFromContent(message.content)
      return text ? [{ index, entryId, kind: 'user', label: 'user', text }] : []
    }
    case 'assistant':
      return assistantEntries(message, index, entryId, interactionToolCalls)
    case 'toolResult': {
      const name = typeof message.toolName === 'string' ? message.toolName : 'unknown'
      const userInteraction = normalizedUserInteraction(message, interactionToolCalls)
      if (userInteraction !== undefined) {
        return [
          {
            index,
            entryId,
            kind: 'user_interaction',
            label: `user_interaction:${name}`,
            text: userInteraction,
          },
        ]
      }
      const suffix = message.isError === true ? ' (error)' : ''
      const text = textFromContent(message.content)
      return text
        ? [
            {
              index,
              entryId,
              kind: 'tool',
              label: `tool:${name}${suffix}`,
              text,
              ...(typeof message.toolCallId === 'string' ? { toolCallId: message.toolCallId } : {}),
              timestamp: message.timestamp,
            },
          ]
        : []
    }
    case 'bashExecution': {
      const command = serializeUnknown(message.command)
      const output = serializeUnknown(message.output)
      return [
        {
          index,
          entryId,
          kind: 'tool',
          label: 'tool:user-bash',
          text: `${command}\n${output}`,
        },
      ]
    }
    case 'branchSummary':
    case 'compactionSummary': {
      const text = serializeUnknown(message.summary)
      return text ? [{ index, entryId, kind: 'assistant', label: String(message.role), text }] : []
    }
    case 'custom': {
      const text = textFromContent(message.content)
      return text ? [{ index, entryId, kind: 'assistant', label: 'custom', text }] : []
    }
    default:
      return []
  }
}

export function collectTranscriptEntries(sessionEntries: SessionEntry[]): TranscriptEntry[] {
  const interactionToolCalls = new Map<string, string>()
  return sessionEntries.flatMap((entry, index) => {
    if (entry.type === 'message') {
      return entriesFromMessage(
        entry.message as MessageLike,
        index,
        entry.id,
        interactionToolCalls,
      )
    }
    if (entry.type === 'compaction' || entry.type === 'branch_summary') {
      return [
        {
          index,
          entryId: entry.id,
          kind: 'assistant' as const,
          label: entry.type,
          text: entry.summary,
        },
      ]
    }
    if (entry.type === 'custom_message') {
      const text = textFromContent(entry.content)
      return text
        ? [
            {
              index,
              entryId: entry.id,
              kind: 'assistant' as const,
              label: 'custom',
              text,
            },
          ]
        : []
    }
    return []
  })
}

function renderTranscriptEntry(entry: TranscriptEntry): string {
  return JSON.stringify({
    index: entry.index,
    source: entry.kind,
    label: entry.label,
    content: entry.text,
  })
}

function transcriptEntryTokens(entry: TranscriptEntry): number {
  return estimateTokens(renderTranscriptEntry(entry))
}

function pretruncate(entry: TranscriptEntry): TranscriptEntry {
  const maxTokens = entry.kind === 'tool' ? MAX_TOOL_ENTRY_TOKENS : MAX_MESSAGE_ENTRY_TOKENS
  if (transcriptEntryTokens(entry) <= maxTokens) {
    return entry
  }

  let lower = 0
  let upper = entry.text.length
  let text = truncateToCharacters(entry.text, 0)
  while (lower <= upper) {
    const middle = Math.floor((lower + upper) / 2)
    const candidate = truncateToCharacters(entry.text, middle)
    if (estimateTokens(renderTranscriptEntry({ ...entry, text: candidate })) <= maxTokens) {
      text = candidate
      lower = middle + 1
    } else {
      upper = middle - 1
    }
  }
  return { ...entry, text, truncated: true }
}

function isTrusted(entry: TranscriptEntry): boolean {
  return entry.kind === 'user' || entry.kind === 'user_interaction'
}

function addWithinBudget(selected: Set<TranscriptEntry>, entries: TranscriptEntry[], budget: number): number {
  let used = 0
  for (const entry of entries) {
    const tokens = transcriptEntryTokens(entry)
    if (used + tokens > budget) {
      continue
    }
    selected.add(entry)
    used += tokens
  }
  return used
}

function summaryText(record: ToolResultSummaryRecord): string | undefined {
  const summary = record.semanticSummary ?? record.deterministicSummary
  return typeof summary === 'string' && summary.length > 0 ? summary : undefined
}

/**
 * Substitute pi-compaction summaries for plain tool results.
 *
 * One lookup call per render, over every toolResult's call id at once. A
 * record replaces the entry text only when it is strictly shorter — a summary
 * longer than its source compresses nothing — and the label gains
 * ` (summarized)` so the reviewer reads it as a digest, not the raw output.
 * Entries normalized to `user_interaction` and assistant toolCall argument
 * entries carry no `toolCallId`, so they are never touched.
 *
 * Returns the entries (substituted in place on copies) and how many were
 * summarized. Any lookup failure keeps the originals.
 */
function substituteSummaries(
  entries: TranscriptEntry[],
  query: ToolResultSummaryQuery | undefined,
): number {
  if (query === undefined) return 0
  const candidates = entries.filter(
    entry => entry.kind === 'tool' && entry.toolCallId !== undefined,
  )
  if (candidates.length === 0) return 0

  let records: readonly ToolResultSummaryRecord[]
  try {
    records = query(candidates.map(entry => entry.toolCallId as string))
  } catch {
    return 0
  }
  if (!Array.isArray(records) || records.length === 0) return 0

  const byCallId = new Map<string, ToolResultSummaryRecord[]>()
  for (const record of records) {
    if (record === null || typeof record !== 'object' || typeof record.toolCallId !== 'string') {
      continue
    }
    const list = byCallId.get(record.toolCallId)
    if (list === undefined) {
      byCallId.set(record.toolCallId, [record])
    } else {
      list.push(record)
    }
  }

  let summarized = 0
  for (const entry of candidates) {
    const list = byCallId.get(entry.toolCallId as string)
    if (list === undefined) continue
    const record =
      list.find(
        candidate =>
          entry.timestamp !== undefined &&
          candidate.resultTimestamp !== undefined &&
          candidate.resultTimestamp === entry.timestamp,
      ) ?? list[0]
    if (record === undefined) continue
    const summary = summaryText(record)
    if (summary === undefined) continue
    const replacement = `${summary}\n[summarized]`
    if (replacement.length >= entry.text.length) continue
    entry.text = replacement
    entry.label = `${entry.label} (summarized)`
    summarized += 1
  }
  return summarized
}

export function renderTranscript(
  sessionEntries: SessionEntry[],
  options?: TranscriptRenderOptions,
): RenderedTranscript {
  const budgets = options?.budgets ?? DEFAULT_TRANSCRIPT_BUDGETS
  const collected = collectTranscriptEntries(sessionEntries)
  const toolResultsSummarized = substituteSummaries(collected, options?.querySummaries)
  const allEntries = collected.map(pretruncate)
  const selected = new Set<TranscriptEntry>()
  const trustedEntries = allEntries.filter(isTrusted)

  let messageTokens = 0
  if (trustedEntries.length > 0) {
    const first = trustedEntries[0]
    const latest = trustedEntries.at(-1)
    if (first !== undefined) {
      selected.add(first)
      messageTokens += transcriptEntryTokens(first)
    }
    if (latest !== undefined && latest !== first) {
      selected.add(latest)
      messageTokens += transcriptEntryTokens(latest)
    }
  }

  const remainingTrusted = trustedEntries.filter(entry => !selected.has(entry)).toReversed()
  messageTokens += addWithinBudget(selected, remainingTrusted, budgets.messageTokens - messageTokens)

  let toolTokens = 0
  let untrustedEntriesRetained = 0
  for (const entry of allEntries.toReversed()) {
    if (isTrusted(entry) || untrustedEntriesRetained >= budgets.maxUntrustedEntries) {
      continue
    }
    const tokens = transcriptEntryTokens(entry)
    if (entry.kind === 'tool') {
      if (toolTokens + tokens > budgets.toolTokens) {
        continue
      }
      toolTokens += tokens
    } else {
      if (messageTokens + tokens > budgets.messageTokens) {
        continue
      }
      messageTokens += tokens
    }
    selected.add(entry)
    untrustedEntriesRetained += 1
  }

  const retained = [...selected].sort((left, right) => left.index - right.index)
  const latestTrusted = trustedEntries.at(-1)
  const stats: TranscriptStats = {
    entriesRetained: retained.length,
    entriesOmitted: allEntries.length - retained.length,
    entriesTruncated: retained.filter(entry => entry.truncated === true).length,
    toolResultsSummarized,
    latestTrustedEntryRetained: latestTrusted !== undefined && selected.has(latestTrusted),
  }

  return {
    entries: retained.map(renderTranscriptEntry),
    omittedCount: stats.entriesOmitted,
    stats,
    latestTrustedEntryId: latestTrusted?.entryId ?? null,
  }
}

/**
 * A summary lookup bound to the session's pi-compaction service, or
 * `undefined` when no service is published.
 *
 * The service is a `Symbol.for()` global — read structurally so this module
 * never imports pi-compaction code, and a missing, throwing, or differently
 * shaped service simply yields no lookup rather than a render failure.
 */
export function compactionSummaryQuery(
  sessionId: string,
): ToolResultSummaryQuery | undefined {
  const service = (globalThis as Record<symbol, unknown>)[
    Symbol.for('my-pi.pi-compaction.service.v1')
  ]
  if (service === null || typeof service !== 'object') {
    return undefined
  }
  const query = (service as Record<string, unknown>).query
  if (typeof query !== 'function') {
    return undefined
  }
  const queryFn = query as (
    this: unknown,
    session: string,
    ids: readonly string[],
  ) => unknown
  return toolCallIds => {
    const records = queryFn.call(service, sessionId, toolCallIds)
    return Array.isArray(records) ? (records as ToolResultSummaryRecord[]) : []
  }
}