path: root/src/core/metrics/place.ts

import type { Usage } from "@dispatch/wire";
import type { RenderGroup } from "../chunks";
import type { MetricsRow, TurnMetricsEntry } from "./types";

function groupStepId(g: RenderGroup): string | undefined {
  if (g.kind === "tool-batch") return g.stepId;
  const c = g.chunk.chunk;
  return c.type === "tool-call" || c.type === "tool-result" ? c.stepId : undefined;
}

/** Element-wise sum of two token usages (cache fields included only when nonzero). */
function addUsage(a: Usage, b: Usage): Usage {
  const out: Usage = {
    inputTokens: a.inputTokens + b.inputTokens,
    outputTokens: a.outputTokens + b.outputTokens,
  };
  const read = (a.cacheReadTokens ?? 0) + (b.cacheReadTokens ?? 0);
  const write = (a.cacheWriteTokens ?? 0) + (b.cacheWriteTokens ?? 0);
  if (read > 0) (out as { cacheReadTokens?: number }).cacheReadTokens = read;
  if (write > 0) (out as { cacheWriteTokens?: number }).cacheWriteTokens = write;
  return out;
}

/**
 * Interleave turn metrics into the rendered transcript.
 *
 * Splits groups into per-turn segments: a new segment begins at each `single`
 * group with `group.chunk.role === "user"`. Segments are matched to entries
 * by `stepId` presence when possible (robust against chat-limit trimming: when
 * a turn's user message is trimmed, positional alignment would be off, but
 * stepId matching still finds the right entry). Segments with no stepId-bearing
 * groups (text-only turns) fall back to POSITIONAL tail-alignment: since the
 * loaded transcript is always a SUFFIX of the full turn history (the chat limit
 * keeps the newest and unloads the oldest), segment `seg` ↔ entry `K - T + seg`.
 *
 * Within a segment that has a matched entry, each completed step's metrics
 * are placed INLINE right after the last group bearing that step's `stepId`.
 * Steps whose `stepId` does not appear in any group ("unanchored"):
 * - If the segment HAS stepId-bearing groups (tool chunks exist but this step's
 *   were trimmed): SKIPPED (no blank "step N · 0 tok" bubbles).
 * - If the segment has NO stepId-bearing groups (text-only turn): placed at the
 *   segment tail before the turn-metrics row (the original behavior).
 *
 * A `turn-metrics` row is emitted ONLY when `entry.total !== null` (i.e. the turn
 * is finalized via `done` or durable data). A still-generating turn emits no
 * turn-total row.
 *
 * Fully trimmed turns (entries whose content was unloaded by the chat limit and
 * which match no segment) are NOT rendered as standalone rows — that previously
 * piled a wall of stale cache badges at the top of a long, trimmed transcript.
 * Their usage still counts toward the per-turn "chat total" cumulative (computed
 * across ALL finalized turns in entry-array order), so the running cache rate
 * stays correct regardless of which turns were trimmed; paging earlier history
 * back in ("Show earlier messages") re-matches them and re-renders their rows.
 */
export function interleaveTurnMetrics(
  groups: readonly RenderGroup[],
  entries: readonly TurnMetricsEntry[],
): readonly MetricsRow[] {
  if (entries.length === 0) {
    return groups.map((g) => ({ kind: "group" as const, group: g }));
  }

  const segmentStarts: number[] = [];
  for (let i = 0; i < groups.length; i++) {
    const g = groups[i];
    if (g !== undefined && g.kind === "single" && g.chunk.role === "user") {
      segmentStarts.push(i);
    }
  }

  let T = segmentStarts.length;

  // No user messages — e.g. a compacted conversation whose history starts
  // with a system summary. Treat the entire transcript as one segment so
  // turn/step metrics can still be placed.
  if (T === 0 && entries.length > 0) {
    segmentStarts.push(0);
    T = 1;
  }

  if (T === 0) {
    return groups.map((g) => ({ kind: "group" as const, group: g }));
  }

  const K = entries.length;

  // Build stepId → entry-index lookup for matching.
  const entryStepIds: Set<string>[] = entries.map((e) => new Set(e.steps.map((s) => s.stepId)));

  // Match segments to entries. Pass 1: match by stepId overlap (handles
  // trimming where positional alignment alone could be ambiguous). Pass 2:
  // positional tail-alignment fallback for unmatched segments (text-only turns
  // with no stepId-bearing groups).
  const usedEntries = new Set<number>();
  const segmentEntry = new Map<number, TurnMetricsEntry>();
  const segmentEntryIndex = new Map<number, number>();

  // Pass 1: stepId matching.
  for (let seg = 0; seg < T; seg++) {
    const start = segmentStarts[seg] ?? 0;
    const end = seg + 1 < T ? (segmentStarts[seg + 1] ?? groups.length) : groups.length;

    const segStepIds = new Set<string>();
    for (let i = start; i < end; i++) {
      const g = groups[i];
      if (g === undefined) continue;
      const sid = groupStepId(g);
      if (sid !== undefined) segStepIds.add(sid);
    }
    if (segStepIds.size === 0) continue; // text-only — defer to pass 2

    let bestEntry = -1;
    let bestMatch = 0;
    for (let i = 0; i < K; i++) {
      if (usedEntries.has(i)) continue;
      let match = 0;
      for (const sid of segStepIds) {
        if (entryStepIds[i]?.has(sid)) match++;
      }
      if (match > bestMatch) {
        bestMatch = match;
        bestEntry = i;
      }
    }
    if (bestEntry >= 0) {
      usedEntries.add(bestEntry);
      const e = entries[bestEntry];
      if (e !== undefined) {
        segmentEntry.set(seg, e);
        segmentEntryIndex.set(seg, bestEntry);
      }
    }
  }

  // Pass 2: positional fallback for segments pass 1 left unmatched
  // (text-only turns with no stepId-bearing groups to anchor on).
  //
  // The loaded transcript is always a SUFFIX of the full turn history —
  // chat-limit/windowing keeps the NEWEST chunks and unloads the OLDEST — so
  // the T loaded segments correspond to the LAST T entries. TAIL-ALIGNMENT
  // (segment `seg` ↔ entry `K - T + seg`) is therefore correct whenever the
  // metrics hold at least as many turns as there are loaded segments
  // (`K >= T`): the leading `K - T` entries are TRIMMED turns (their content
  // was unloaded) and must be skipped, never matched to a newer segment.
  //
  // This MUST run even when pass 1 matched SOME segments (tool turns). The
  // earlier code only tail-aligned when pass 1 matched NONE, falling back to
  // HEAD-alignment otherwise — which, with leading trimmed entries, matched a
  // brand-new text-only turn to an old (trimmed) entry's STALE metrics (the
  // "new steps show no / wrong cache" failure). Tail-aligning by position is
  // safe alongside pass 1: stepIds are unique per turn, so pass 1 already
  // grabbed each tool turn's positionally-correct entry, leaving the right
  // entry free for each text-only turn.
  //
  // Only when `K < T` (fewer entries than segments — some loaded turns have no
  // metrics yet, e.g. a metrics sync still pending or a freshly loaded
  // transcript) do we head-align, assigning the first K entries to the first K
  // unmatched segments (the turns that DO have metrics sit at the front).
  if (K >= T) {
    // Tail-align: skip the first K-T entries (trimmed turns).
    for (let seg = 0; seg < T; seg++) {
      if (segmentEntry.has(seg)) continue;
      const entryIdx = K - T + seg;
      if (entryIdx >= 0 && entryIdx < K && !usedEntries.has(entryIdx)) {
        usedEntries.add(entryIdx);
        const e = entries[entryIdx];
        if (e !== undefined) {
          segmentEntry.set(seg, e);
          segmentEntryIndex.set(seg, entryIdx);
        }
      }
    }
  } else {
    // Head-align fallback (K < T): first K entries to first K unmatched segments.
    let nextUnused = 0;
    for (let seg = 0; seg < T; seg++) {
      if (segmentEntry.has(seg)) continue;
      while (nextUnused < K && usedEntries.has(nextUnused)) nextUnused++;
      if (nextUnused < K) {
        usedEntries.add(nextUnused);
        const e = entries[nextUnused];
        if (e !== undefined) {
          segmentEntry.set(seg, e);
          segmentEntryIndex.set(seg, nextUnused);
        }
        nextUnused++;
      }
    }
  }

  // Running cumulative usage across ALL finalized turns (in entry order), for
  // the per-turn "chat total" cache rate. Alongside it, the previous finalized
  // turn's usage at each index — the baseline for cross-turn retention.
  const cumulativeByEntry: Usage[] = [];
  const prevUsageByEntry: (Usage | null)[] = [];
  let runningUsage: Usage = { inputTokens: 0, outputTokens: 0 };
  let lastFinalizedUsage: Usage | null = null;
  for (const e of entries) {
    prevUsageByEntry.push(lastFinalizedUsage);
    if (e.total !== null) {
      runningUsage = addUsage(runningUsage, e.total.usage);
      lastFinalizedUsage = e.total.usage;
    }
    cumulativeByEntry.push(runningUsage);
  }

  const rows: MetricsRow[] = [];

  const firstUserIdx = segmentStarts[0] ?? 0;

  for (let i = 0; i < firstUserIdx; i++) {
    const g = groups[i];
    if (g !== undefined) {
      rows.push({ kind: "group", group: g });
    }
  }

  for (let seg = 0; seg < T; seg++) {
    const start = segmentStarts[seg] ?? 0;
    const end = seg + 1 < T ? (segmentStarts[seg + 1] ?? groups.length) : groups.length;

    const entry = segmentEntry.get(seg);

    if (entry === undefined) {
      for (let i = start; i < end; i++) {
        const g = groups[i];
        if (g !== undefined) {
          rows.push({ kind: "group", group: g });
        }
      }
      continue;
    }

    const entryIdx = segmentEntryIndex.get(seg) ?? 0;

    // Build anchor map: for each stepId, the LAST group index in this segment.
    const anchorByStepId = new Map<string, number>();
    for (let i = start; i < end; i++) {
      const g = groups[i];
      if (g === undefined) continue;
      const sid = groupStepId(g);
      if (sid !== undefined) {
        anchorByStepId.set(sid, i);
      }
    }

    // Classify each step as anchored or unanchored. Unanchored steps
    // (content trimmed, or text-only steps with no tool chunks) are SKIPPED —
    // step-metrics are only shown inline next to the content they describe.
    const anchored: Map<number, { stepIndex: number; step: (typeof entry.steps)[number] }[]> =
      new Map();

    for (let i = 0; i < entry.steps.length; i++) {
      const step = entry.steps[i];
      if (step === undefined) continue;
      const anchorGroupIdx = anchorByStepId.get(step.stepId);
      if (anchorGroupIdx !== undefined) {
        let arr = anchored.get(anchorGroupIdx);
        if (arr === undefined) {
          arr = [];
          anchored.set(anchorGroupIdx, arr);
        }
        arr.push({ stepIndex: i, step });
      }
      // Unanchored steps (no matching group) are skipped — no tail bubbles.
    }

    // Emit groups; after each anchored group, emit its step-metrics rows.
    for (let i = start; i < end; i++) {
      const g = groups[i];
      if (g !== undefined) {
        rows.push({ kind: "group", group: g });
      }
      const stepsHere = anchored.get(i);
      if (stepsHere !== undefined) {
        stepsHere.sort((a, b) => a.stepIndex - b.stepIndex);
        for (const { step, stepIndex } of stepsHere) {
          rows.push({ kind: "step-metrics", step, index: stepIndex });
        }
      }
    }

    // Turn-metrics row (only when the turn is finalized). Unanchored steps
    // are skipped — no tail bubbles.
    if (entry.total !== null) {
      rows.push({
        kind: "turn-metrics",
        turn: entry.total,
        turnNumber: entryIdx + 1,
        cumulativeUsage: cumulativeByEntry[entryIdx] ?? entry.total.usage,
        prevTurnUsage: prevUsageByEntry[entryIdx] ?? null,
      });
    }
  }

  return rows;
}