1 files changed, 345 insertions, 0 deletions
diff --git a/src/core/metrics/reducer.ts b/src/core/metrics/reducer.ts
new file mode 100644
index 0000000..39fc5ee
--- /dev/null
+++ b/src/core/metrics/reducer.ts
@@ -0,0 +1,345 @@
+import type { AgentEvent, StepId, StepMetrics, TurnMetrics, Usage } from "@dispatch/wire";
+import type { BuildingStep, LiveTurn, MetricsState, TurnMetricsEntry } from "./types";
+
+function sumStepUsages(steps: readonly BuildingStep[]): Usage {
+  let inputTokens = 0;
+  let outputTokens = 0;
+  let hasCacheRead = false;
+  let hasCacheWrite = false;
+  let cacheReadTokens = 0;
+  let cacheWriteTokens = 0;
+
+  for (const step of steps) {
+    if (step.usage === undefined) continue;
+    inputTokens += step.usage.inputTokens;
+    outputTokens += step.usage.outputTokens;
+    if (step.usage.cacheReadTokens !== undefined && step.usage.cacheReadTokens > 0) {
+      hasCacheRead = true;
+      cacheReadTokens += step.usage.cacheReadTokens;
+    }
+    if (step.usage.cacheWriteTokens !== undefined && step.usage.cacheWriteTokens > 0) {
+      hasCacheWrite = true;
+      cacheWriteTokens += step.usage.cacheWriteTokens;
+    }
+  }
+
+  const base: Usage = { inputTokens, outputTokens };
+  if (hasCacheRead) {
+    (base as { cacheReadTokens?: number }).cacheReadTokens = cacheReadTokens;
+  }
+  if (hasCacheWrite) {
+    (base as { cacheWriteTokens?: number }).cacheWriteTokens = cacheWriteTokens;
+  }
+  return base;
+}
+
+function buildingStepToMetrics(bs: BuildingStep): StepMetrics {
+  const usage: Usage = bs.usage ?? { inputTokens: 0, outputTokens: 0 };
+  const base: StepMetrics = { stepId: bs.stepId as StepId, usage };
+  if (bs.ttftMs !== undefined) {
+    (base as { ttftMs?: number }).ttftMs = bs.ttftMs;
+  }
+  if (bs.decodeMs !== undefined) {
+    (base as { decodeMs?: number }).decodeMs = bs.decodeMs;
+  }
+  if (bs.genTotalMs !== undefined) {
+    (base as { genTotalMs?: number }).genTotalMs = bs.genTotalMs;
+  }
+  return base;
+}
+
+function getStep(lt: LiveTurn, id: string): BuildingStep {
+  const step = lt.stepMap.get(id);
+  if (step === undefined) throw new Error(`Missing step ${id} in live turn`);
+  return step;
+}
+
+function liveTurnToMetrics(lt: LiveTurn): TurnMetrics {
+  const buildingSteps = lt.stepOrder.map((id) => getStep(lt, id));
+  const steps = buildingSteps.map((bs) => buildingStepToMetrics(bs));
+  const usage = lt.doneUsage ?? sumStepUsages(buildingSteps);
+  const base: TurnMetrics = { turnId: lt.turnId, usage, steps };
+  if (lt.durationMs !== undefined) {
+    (base as { durationMs?: number }).durationMs = lt.durationMs;
+  }
+  if (lt.doneContextSize !== undefined) {
+    (base as { contextSize?: number }).contextSize = lt.doneContextSize;
+  }
+  return base;
+}
+
+/**
+ * A step's contribution to the live context size: `inputTokens + outputTokens`,
+ * or `undefined` when the step has no usage yet OR its counters are not safe to
+ * sum (non-finite / negative — defensive: a corrupt provider report must never
+ * reach the status bar as NaN/Infinity). Cache tokens are deliberately NOT
+ * included: `cacheReadTokens` / `cacheWriteTokens` are a SUBSET of
+ * `inputTokens`, so adding them would double-count.
+ */
+function stepContextSize(usage: Usage | undefined): number | undefined {
+  if (usage === undefined) return undefined;
+  const { inputTokens, outputTokens } = usage;
+  if (!Number.isFinite(inputTokens) || !Number.isFinite(outputTokens)) return undefined;
+  if (inputTokens < 0 || outputTokens < 0) return undefined;
+  return inputTokens + outputTokens;
+}
+
+/**
+ * The context size an IN-FLIGHT (not-done) turn occupies right now — for
+ * progressive display DURING a turn (before it seals), so the indicator updates
+ * after each step instead of waiting for `done`.
+ *
+ * CONTRACT: only call this on a turn whose `done` event has NOT arrived (the
+ * caller, `selectCurrentContextSize`, reaches it solely for entries with
+ * `total === null`, i.e. `lt.done === false`). Finalized turns use their
+ * authoritative `contextSize` instead; `doneContextSize` is read on the
+ * `total` path, never here.
+ *
+ * Returns the most recent step WITH USABLE USAGE's `inputTokens + outputTokens`
+ * (scanning newest → oldest by first-seen step order): each step's input
+ * already includes all prior context (the prompt is re-prefilled every step), so
+ * the last step's input+output is the true occupancy — the same definition
+ * `TurnDoneEvent.contextSize` stamps at turn end. A just-reported step's usage
+ * wins immediately, even mid-stream. Steps with no usage or unsafe usage are
+ * skipped, falling back to the next older usable step. `undefined` when no step
+ * has reported usable usage yet.
+ */
+function liveTurnContextSize(lt: LiveTurn): number | undefined {
+  for (let i = lt.stepOrder.length - 1; i >= 0; i--) {
+    const step = lt.stepMap.get(lt.stepOrder[i] ?? "");
+    const ctx = stepContextSize(step?.usage);
+    if (ctx !== undefined) return ctx;
+  }
+  return undefined;
+}
+
+function ensureLiveTurn(state: MetricsState, turnId: string): [MetricsState, LiveTurn] {
+  const existing = state.live.get(turnId);
+  if (existing !== undefined) return [state, existing];
+
+  const newTurn: LiveTurn = {
+    turnId,
+    done: false,
+    durationMs: undefined,
+    doneUsage: undefined,
+    doneContextSize: undefined,
+    stepMap: new Map(),
+    stepOrder: [],
+  };
+  const newLive = new Map(state.live);
+  newLive.set(turnId, newTurn);
+  return [{ ...state, live: newLive, liveOrder: [...state.liveOrder, turnId] }, newTurn];
+}
+
+function upsertStep(lt: LiveTurn, stepId: string, update: Partial<BuildingStep>): LiveTurn {
+  const existing = lt.stepMap.get(stepId);
+  if (existing !== undefined) {
+    const merged: BuildingStep = {
+      stepId,
+      usage: update.usage ?? existing.usage,
+      ttftMs: update.ttftMs ?? existing.ttftMs,
+      decodeMs: update.decodeMs ?? existing.decodeMs,
+      genTotalMs: update.genTotalMs ?? existing.genTotalMs,
+      complete: update.complete ?? existing.complete,
+    };
+    const newMap = new Map(lt.stepMap);
+    newMap.set(stepId, merged);
+    return { ...lt, stepMap: newMap };
+  }
+
+  const fresh: BuildingStep = {
+    stepId,
+    usage: update.usage,
+    ttftMs: update.ttftMs,
+    decodeMs: update.decodeMs,
+    genTotalMs: update.genTotalMs,
+    complete: update.complete ?? false,
+  };
+  const newMap = new Map(lt.stepMap);
+  newMap.set(stepId, fresh);
+  return { ...lt, stepMap: newMap, stepOrder: [...lt.stepOrder, stepId] };
+}
+
+/** The initial empty metrics state. */
+export function initialMetricsState(): MetricsState {
+  return {
+    live: new Map(),
+    liveOrder: [],
+    durable: new Map(),
+    durableOrder: [],
+  };
+}
+
+/**
+ * Fold one live AgentEvent into the metrics state.
+ *
+ * - `usage` with `stepId`: upsert that step's usage.
+ * - `usage` without `stepId`: ignored.
+ * - `step-complete`: upsert that step's timing; default usage to zeros if absent.
+ * - `done`: set turn's `durationMs`, optional aggregate `usage`, and optional `contextSize`.
+ * - All other event types: return state unchanged.
+ */
+export function foldMetricsEvent(state: MetricsState, event: AgentEvent): MetricsState {
+  switch (event.type) {
+    case "usage": {
+      if (event.stepId === undefined) return state;
+      const [s1, lt] = ensureLiveTurn(state, event.turnId);
+      const updated = upsertStep(lt, event.stepId, { usage: event.usage });
+      const newLive = new Map(s1.live);
+      newLive.set(event.turnId, updated);
+      return { ...s1, live: newLive };
+    }
+
+    case "step-complete": {
+      const [s1, lt] = ensureLiveTurn(state, event.turnId);
+      const updated = upsertStep(lt, event.stepId, {
+        ttftMs: event.ttftMs,
+        decodeMs: event.decodeMs,
+        genTotalMs: event.genTotalMs,
+        complete: true,
+      });
+      const newLive = new Map(s1.live);
+      newLive.set(event.turnId, updated);
+      return { ...s1, live: newLive };
+    }
+
+    case "done": {
+      const [s1, lt] = ensureLiveTurn(state, event.turnId);
+      const updated: LiveTurn = {
+        ...lt,
+        done: true,
+        durationMs: event.durationMs ?? lt.durationMs,
+        doneUsage: event.usage ?? lt.doneUsage,
+        doneContextSize: event.contextSize ?? lt.doneContextSize,
+      };
+      const newLive = new Map(s1.live);
+      newLive.set(event.turnId, updated);
+      return { ...s1, live: newLive };
+    }
+
+    default:
+      return state;
+  }
+}
+
+/**
+ * Store durable (sealed) metrics from the backend. These win over live data
+ * for any shared `turnId`.
+ *
+ * Once durable (authoritative) data covers a turn, its live (in-memory) copy
+ * is REDUNDANT and is pruned from `state.live` / `liveOrder` so the live map
+ * doesn't grow unbounded over a long conversation. There is no display gap:
+ * the durable entry replaces the live one atomically in the same fold, and
+ * `selectOrderedTurnMetrics` / `selectCurrentContextSize` read durable for it.
+ */
+export function applyDurableMetrics(
+  state: MetricsState,
+  turns: readonly TurnMetrics[],
+): MetricsState {
+  const newDurable = new Map(state.durable);
+  const newDurableOrder = [...state.durableOrder];
+  const prunedIds = new Set<string>();
+  for (const turn of turns) {
+    if (!newDurable.has(turn.turnId)) {
+      newDurableOrder.push(turn.turnId);
+    }
+    newDurable.set(turn.turnId, turn);
+    if (state.live.has(turn.turnId)) prunedIds.add(turn.turnId);
+  }
+
+  if (prunedIds.size === 0) {
+    return { ...state, durable: newDurable, durableOrder: newDurableOrder };
+  }
+
+  const newLive = new Map(state.live);
+  for (const id of prunedIds) newLive.delete(id);
+  const newLiveOrder = state.liveOrder.filter((id) => !prunedIds.has(id));
+
+  return {
+    ...state,
+    live: newLive,
+    liveOrder: newLiveOrder,
+    durable: newDurable,
+    durableOrder: newDurableOrder,
+  };
+}
+
+/**
+ * Select the merged ordered list of turn metrics entries.
+ * Durable turns come first (in their order), then any live turns whose
+ * `turnId` is not in durable (in live first-seen order).
+ *
+ * Each entry contains the completed steps so far and an optional total
+ * (null until the turn is finalized via `done` or durable data).
+ * Live turns with no completed steps and not done are omitted.
+ */
+export function selectOrderedTurnMetrics(state: MetricsState): readonly TurnMetricsEntry[] {
+  const result: TurnMetricsEntry[] = [];
+  const seen = new Set<string>();
+
+  for (const turnId of state.durableOrder) {
+    const tm = state.durable.get(turnId);
+    if (tm !== undefined) {
+      result.push({ turnId, steps: tm.steps, total: tm });
+      seen.add(turnId);
+    }
+  }
+
+  for (const turnId of state.liveOrder) {
+    if (seen.has(turnId)) continue;
+    const lt = state.live.get(turnId);
+    if (lt === undefined) continue;
+
+    const completeSteps = lt.stepOrder
+      .map((id) => lt.stepMap.get(id))
+      .filter((s): s is BuildingStep => s?.complete === true)
+      .map((s) => buildingStepToMetrics(s));
+
+    if (completeSteps.length === 0 && !lt.done) continue;
+
+    result.push({
+      turnId,
+      steps: completeSteps,
+      total: lt.done ? liveTurnToMetrics(lt) : null,
+    });
+  }
+
+  return result;
+}
+
+/**
+ * Select the conversation's CURRENT context size — the tokens it occupies right
+ * now. Per the wire contract a client reads the LATEST turn's `contextSize`; we
+ * scan the merged ordered turns NEWEST → OLDEST and return the first DEFINED
+ * value.
+ *
+ * For a FINALIZED turn (`done` event or durable data) we use its authoritative
+ * `contextSize`. For an IN-FLIGHT (not-done) turn we compute it PROGRESSIVELY
+ * from the most recent step WITH USAGE — its `inputTokens + outputTokens` is the
+ * current occupancy (mirroring `TurnDoneEvent.contextSize`'s definition) — so
+ * the indicator updates after each step completes instead of waiting for the
+ * turn to seal. An in-flight turn with no step usage yet is skipped, falling
+ * back to the next older finalized turn.
+ *
+ * Returns `undefined` ("unknown") when no turn carries a context size — the
+ * caller renders a placeholder, NEVER `0`. Durable (sealed) data wins over
+ * live for a shared `turnId` (it is the persisted, authoritative value).
+ */
+export function selectCurrentContextSize(state: MetricsState): number | undefined {
+  const ordered = selectOrderedTurnMetrics(state);
+  for (let i = ordered.length - 1; i >= 0; i--) {
+    const entry = ordered[i];
+    if (entry === undefined) continue;
+    if (entry.total !== null) {
+      if (entry.total.contextSize !== undefined) return entry.total.contextSize;
+      continue;
+    }
+    // In-flight turn: progressive context size from the latest step with usage.
+    const lt = state.live.get(entry.turnId);
+    if (lt !== undefined) {
+      const live = liveTurnContextSize(lt);
+      if (live !== undefined) return live;
+    }
+  }
+  return undefined;
+}