Extraction observability: usage capture, live status payload, true abort (Wave 4)

Extraction is real spend and background work; the dock needs to see
and stop it. User decisions: persist all cost data (chat record),
per-stage usage granularity, and TRUE cancellation for abort.

USAGE CAPTURE (live + persisted):
- extraction.ts: per-call onStepFinish captures each billed step's
  usage (SDK LanguageModelUsage -> CallUsage: input/output/cache
  read+write/reasoning) into the registry's live per-stage usage
  (extractor/auditor; the merger is deterministic — no model call).
- /api/graph/status runs now carry usage (model + steps per stage);
  cost/context folds are the consumer's (core/prices.ts — nothing
  stores dollars).
- ChatRecord.extractions: append-only per-run usage log (mirrors the
  turns log): startedAt key, at, durationMs, merge stats when the
  merger ran, stages. Persisted on success AND failure — a failed run
  still spent money. Idempotent per startedAt; the list payload
  (ChatSummary) never leaks it.
- EmissionStats moved to contracts/graph.ts (its contract home);
  core/graph.ts re-exports for existing importers.
- Verified end-to-end against the real pipeline (backfill one chat):
  per-step token counts, reasoning tokens, and the auditor's growing
  prompt-cache reads captured and persisted.

TRUE ABORT (the dock's abort button, user decision):
- RunningExtraction owns one AbortController per run, shared across
  phases; abortExtraction(chatId) fires it and returns whether a run
  was aborted (idempotent no-op otherwise). The pending model call
  receives the signal — the provider cancels the request, spend
  stops. throwIfAborted() at phase boundaries ends runs aborted in
  sync gaps.
- Abort is a deliberate stop, not a failure: same ledger discipline
  (nothing merged, nothing checkpointed, partial usage persisted) but
  the run reports zero emissions instead of throwing. Exported via
  src/server/observability.ts for master's /api/activity/abort route.
- Tests: mid-run abort (hanging mock that rejects on the signal),
  live-accumulation + persisted-record assertions, store-seam
  idempotency, and the summary-payload leak guard.
This commit is contained in:
amosblomqvist
2026-09-11 03:25:25 +03:00
parent 83a00800e9
commit 9174cde55e
10 changed files with 330 additions and 31 deletions
+18 -4
View File
@@ -11,7 +11,7 @@ import * as fs from "node:fs";
import * as path from "node:path";
import { PATHS } from "../../config.ts";
import type { ChatStore } from "../core/ports.ts";
import type { ChatRecord, ChatSummary, PlanTree, TurnUsage } from "../core/contracts.ts";
import type { ChatRecord, ChatSummary, ExtractionRunUsage, PlanTree, TurnUsage } from "../core/contracts.ts";
import type { UIMessage } from "ai";
import { findPendingToolCall } from "../core/pending.ts";
import { phaseOfTree } from "../core/phases.ts";
@@ -51,9 +51,10 @@ function listChats(): ChatSummary[] {
.filter((f) => f.endsWith(".json"))
.map((f) => getChat(f.slice(0, -5)))
.filter((r): r is ChatRecord => r !== null)
.map(({ messages, tree, turns, ...summary }) => ({
// turns stay out of the list payload (list = summaries; usage
// aggregates ride the full record, the picker never needs them).
.map(({ messages, tree, turns, extractions, ...summary }) => ({
// turns and extractions stay out of the list payload (list =
// summaries; usage aggregates ride the full record, the picker
// never needs them).
...summary,
messageCount: messages.length,
// Record-derived status for the session picker: phase from the
@@ -90,6 +91,18 @@ function appendTurnUsage(id: string, entry: TurnUsage): void {
fs.writeFileSync(chatFile(id), JSON.stringify(record, null, "\t"), "utf-8");
}
/** Append one extraction-run usage entry to the record's extractions
* log — idempotent per startedAt: a retry that reaches the append twice
* (abort replays) does not double-count spend. */
function appendExtractionUsage(id: string, entry: ExtractionRunUsage): void {
const record = getChat(id);
if (!record) throw new Error(`Chat not found: ${id}`);
record.extractions = record.extractions?.some((r) => r.startedAt === entry.startedAt)
? record.extractions
: [...(record.extractions ?? []), entry];
fs.writeFileSync(chatFile(id), JSON.stringify(record, null, "\t"), "utf-8");
}
/** Persist the session plan tree (plan tools call this in execute). */
function saveTree(id: string, tree: PlanTree): void {
const record = getChat(id);
@@ -107,4 +120,5 @@ export const chatStore: ChatStore = {
saveMessages,
saveTree,
appendTurnUsage,
appendExtractionUsage,
};