diff --git a/src/main/db/connection.ts b/src/main/db/connection.ts index 7aed52588..49486d4dc 100644 --- a/src/main/db/connection.ts +++ b/src/main/db/connection.ts @@ -179,6 +179,17 @@ export function initDatabase(dbPath: string) { value INTEGER NOT NULL DEFAULT 1 ); CREATE INDEX IF NOT EXISTS idx_usage_events_kind ON usage_events (kind); + CREATE TABLE IF NOT EXISTS usage_token_ledger ( + provider TEXT NOT NULL, + scope_id TEXT NOT NULL, + epoch INTEGER NOT NULL, + last_counter INTEGER NOT NULL, + PRIMARY KEY (provider, scope_id, epoch) + ); + CREATE TABLE IF NOT EXISTS usage_token_samples ( + sample_id TEXT PRIMARY KEY, + ts INTEGER NOT NULL + ); CREATE TABLE IF NOT EXISTS scheduled_tasks ( id TEXT PRIMARY KEY, name TEXT NOT NULL, @@ -225,7 +236,7 @@ export function initDatabase(dbPath: string) { // Baseline schema version for future DB migrations. // New upgrade steps should live behind this gate when we need them. - const SCHEMA_VERSION = 26; + const SCHEMA_VERSION = 27; const storedVersion = Number( ( @@ -496,6 +507,25 @@ export function initDatabase(dbPath: string) { } } + if (storedVersion < 27) { + // Main-process token ledger: cumulative-counter state keyed by provider + // scope+epoch, and exact-once dedup of per-call usage samples. Counted + // amounts land in usage_events as kind="tokens_v2" rows. + sqlite.exec(` + CREATE TABLE IF NOT EXISTS usage_token_ledger ( + provider TEXT NOT NULL, + scope_id TEXT NOT NULL, + epoch INTEGER NOT NULL, + last_counter INTEGER NOT NULL, + PRIMARY KEY (provider, scope_id, epoch) + ); + CREATE TABLE IF NOT EXISTS usage_token_samples ( + sample_id TEXT PRIMARY KEY, + ts INTEGER NOT NULL + ); + `); + } + sqlite .prepare( "INSERT INTO app_state (key, value) VALUES ('schema_version', ?) ON CONFLICT(key) DO UPDATE SET value = excluded.value", diff --git a/src/main/db/runtimeItems.ts b/src/main/db/runtimeItems.ts index 100aa2758..952822e75 100644 --- a/src/main/db/runtimeItems.ts +++ b/src/main/db/runtimeItems.ts @@ -498,6 +498,11 @@ export function dbApplyThreadRuntimeEvents( } break; + case "usage.spent": + // Token consumption is not a chat item; the usage ledger persists it + // (recordUsageSpentFromRuntimeEvents) alongside this function. + break; + default: break; } diff --git a/src/main/profile/tokenStats.test.ts b/src/main/profile/tokenStats.test.ts new file mode 100644 index 000000000..9ab1d39b2 --- /dev/null +++ b/src/main/profile/tokenStats.test.ts @@ -0,0 +1,119 @@ +import { execFileSync } from "node:child_process"; +import { existsSync, mkdtempSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import Database from "better-sqlite3"; +import { afterEach, beforeEach, describe, expect, it } from "vitest"; +import { closeDatabase, initDatabase } from "../db/connection"; +import { dbAppendUsageEvents } from "../db/usageEvents"; +import { computeProfileTokenStats } from "./tokenStats"; + +// node_modules/better-sqlite3 may be compiled for Electron's ABI. Fall back to +// the Node-ABI binding used by the headless server (mirrors src/main/db tests). +const serverNativeBinding = join(process.cwd(), "dist", "server-native", "better_sqlite3.node"); +let nativeBindingEnv: string | undefined; + +function databaseOpens(nativeBinding?: string): boolean { + if (nativeBinding && !existsSync(nativeBinding)) return false; + try { + const database = nativeBinding + ? new Database(":memory:", { nativeBinding }) + : new Database(":memory:"); + database.close(); + return true; + } catch { + return false; + } +} + +if (!databaseOpens()) { + if (!databaseOpens(serverNativeBinding)) { + execFileSync(process.execPath, [join(process.cwd(), "scripts", "prepare-server-native.mjs")], { + stdio: "inherit", + }); + } + if (!databaseOpens(serverNativeBinding)) { + throw new Error("Unable to prepare a Node-compatible better-sqlite3 binding for tests."); + } + nativeBindingEnv = serverNativeBinding; +} + +describe("computeProfileTokenStats (real sqlite)", () => { + let dir: string; + + beforeEach(() => { + if (nativeBindingEnv) { + process.env.PORACODE_BETTER_SQLITE3_NATIVE_BINDING = nativeBindingEnv; + } + dir = mkdtempSync(join(tmpdir(), "lc-tokenstats-test-")); + initDatabase(join(dir, "state.sqlite")); + }); + + afterEach(() => { + closeDatabase(); + rmSync(dir, { recursive: true, force: true }); + }); + + it("sums legacy tokens and exact tokens_v2 rows together", () => { + const now = Date.now(); + dbAppendUsageEvents([ + { ts: now, kind: "tokens", provider: "claude", value: 100 }, + { ts: now, kind: "tokens_v2", provider: "claude", value: 25 }, + { ts: now, kind: "tokens_v2", provider: "codex", value: 50 }, + // Zero/negative deltas and unrelated kinds never reach the sums. + { ts: now, kind: "tokens_v2", provider: "codex", value: 0 }, + { ts: now, kind: "turn", provider: "codex", value: 4000 }, + ]); + + const stats = computeProfileTokenStats({ utcOffsetMinutes: 0 }); + expect(stats.available).toBe(true); + expect(stats.lifetimeTokens).toBe(175); + expect(stats.providers.find((p) => p.provider === "claude")?.tokens).toBe(125); + expect(stats.providers.find((p) => p.provider === "codex")?.tokens).toBe(50); + }); + + it("lists base providers with activity but no tokens_v2 rows as unavailable", () => { + const now = Date.now(); + dbAppendUsageEvents([ + { ts: now, kind: "tokens_v2", provider: "claude", value: 10 }, + { ts: now, kind: "thread_started", provider: "claude" }, + { ts: now, kind: "thread_started", provider: "kimi" }, + { ts: now, kind: "turn", provider: "qwen", value: 5000 }, + // Account-scoped kinds fold to their base provider for the honesty list. + { ts: now, kind: "turn", provider: "kimi:for-coding", value: 3000 }, + // A zero-value tokens_v2 row still proves exact coverage. + { ts: now, kind: "thread_started", provider: "codex" }, + { ts: now, kind: "tokens_v2", provider: "codex", value: 0 }, + ]); + + const stats = computeProfileTokenStats({ utcOffsetMinutes: 0 }); + expect(stats.unavailableProviders).toEqual(["kimi", "qwen"]); + }); + + it("reports no unavailable providers when the log has no token rows", () => { + dbAppendUsageEvents([{ ts: Date.now(), kind: "message", provider: "claude" }]); + + const stats = computeProfileTokenStats({ utcOffsetMinutes: 0 }); + expect(stats.available).toBe(false); + expect(stats.lifetimeTokens).toBe(0); + expect(stats.unavailableProviders).toEqual([]); + }); + + it("ignores activity from before the first exact (tokens_v2) row", () => { + const now = Date.now(); + const day = 86_400_000; + dbAppendUsageEvents([ + // Pre-ledger-era activity: qwen ran before any exact telemetry existed on + // this device — not evidence of missing coverage. + { ts: now - 10 * day, kind: "thread_started", provider: "qwen" }, + { ts: now - 10 * day, kind: "turn", provider: "qwen", value: 5000 }, + // Ledger era starts here. + { ts: now - day, kind: "tokens_v2", provider: "claude", value: 10 }, + // kimi ran in the exact era but produced no exact rows. + { ts: now, kind: "thread_started", provider: "kimi" }, + ]); + + const stats = computeProfileTokenStats({ utcOffsetMinutes: 0 }); + expect(stats.unavailableProviders).toEqual(["kimi"]); + }); +}); diff --git a/src/main/profile/tokenStats.ts b/src/main/profile/tokenStats.ts index 64eb82e1c..53def30d0 100644 --- a/src/main/profile/tokenStats.ts +++ b/src/main/profile/tokenStats.ts @@ -19,10 +19,13 @@ import { accountLabel, modelKey, modelLabel, providerLabel } from "./labels"; /** * Token usage from Poracode's own activity, read from the durable `usage_events` - * log (kind="tokens" - per-turn deltas captured at the canonical-event layer for - * every provider, incl. all ACP agents). No external transcript scanning, and no - * dependency on threads (survives delete/archive). Reported both globally (folded - * to the base provider) and per account (each profile separately). + * log. Two row kinds feed the sums: kind="tokens_v2" (exact spend counted by the + * main-process usage ledger from provider-reported counters) and kind="tokens" + * (legacy per-turn deltas captured at the canonical-event layer — approximate + * occupancy growth, kept for historical continuity). No external transcript + * scanning, and no dependency on threads (survives delete/archive). Reported + * both globally (folded to the base provider) and per account (each profile + * separately). */ function round1(value: number): number { @@ -50,6 +53,7 @@ function emptyTokenStats( accounts: [], models: [], tokenHeatmap: heatmap, + unavailableProviders: [], }; } @@ -80,12 +84,32 @@ export function computeProfileTokenStats(req: ProfileStatsRequest): ProfileToken const perProvider = new Map(); const perAccount = new Map(); const perModel = new Map(); + // Coverage sets for the honesty list: providers that ran threads or turns vs + // providers with at least one exact (tokens_v2) token row. A zero-value + // tokens_v2 row still proves the provider reports exact usage. Computed over + // the whole device log (ignoring the active account/window filters), but + // activity only counts once the exact-telemetry era started on this device + // (first tokens_v2 row) — otherwise every provider would flash "unavailable" + // right after upgrade until it next records exact rows. + const activityFirstTs = new Map(); + const exactTokenProviders = new Set(); + let firstExactTs: number | undefined; let lifetimeTokens = 0; const windowDays = statsWindowDays(req.window); const startDayIndex = windowDays ? windowStartIndex(todayIndex, windowDays) : undefined; for (const row of dbGetAllUsageEvents()) { - if (row.kind !== "tokens" || row.value <= 0) continue; + if (row.provider) { + const base = baseAgentKind(row.provider); + if (row.kind === "tokens_v2") { + exactTokenProviders.add(base); + if (firstExactTs === undefined || row.ts < firstExactTs) firstExactTs = row.ts; + } else if (row.kind === "thread_started" || row.kind === "turn") { + const prev = activityFirstTs.get(base); + if (prev === undefined || row.ts < prev) activityFirstTs.set(base, row.ts); + } + } + if ((row.kind !== "tokens" && row.kind !== "tokens_v2") || row.value <= 0) continue; // Scope to the selected account (exact account-kind match) when filtering. if (req.provider && row.provider !== req.provider) continue; const dayIndex = localDayIndex(row.ts, offset); @@ -146,6 +170,17 @@ export function computeProfileTokenStats(req: ProfileStatsRequest): ProfileToken })) .sort((a, b) => b.count - a.count); + const unavailableProviders = + firstExactTs === undefined + ? [] + : [...activityFirstTs] + .filter( + ([provider, firstActivityTs]) => + firstActivityTs >= firstExactTs && !exactTokenProviders.has(provider), + ) + .map(([provider]) => provider) + .sort(); + const result: ProfileTokenStats = { available: lifetimeTokens > 0, scope: req.scope ?? "device", @@ -160,6 +195,7 @@ export function computeProfileTokenStats(req: ProfileStatsRequest): ProfileToken accounts, models, tokenHeatmap, + unavailableProviders, }; tokenCache.set(cacheKey, { generation, result }); return result; diff --git a/src/main/profile/usageLedger.test.ts b/src/main/profile/usageLedger.test.ts new file mode 100644 index 000000000..3ca460368 --- /dev/null +++ b/src/main/profile/usageLedger.test.ts @@ -0,0 +1,255 @@ +import { existsSync, mkdtempSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import Database from "better-sqlite3"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import type { RuntimeEvent, Thread, UsageSpent } from "@/shared/contracts"; +import { closeDatabase, getSqlite, initDatabase } from "../db/connection"; +import { dbUpsertProject, dbUpsertThread } from "../db/projectsThreads"; +import { recordUsageSpentFromRuntimeEvents } from "./usageLedger"; + +const serverNativeBinding = join(process.cwd(), "dist", "server-native", "better_sqlite3.node"); +let nativeBindingEnv: string | undefined; +let sqliteAvailable = true; +try { + new Database(":memory:").close(); +} catch { + if (existsSync(serverNativeBinding)) { + nativeBindingEnv = serverNativeBinding; + } else { + sqliteAvailable = false; + } +} + +const THREAD_ID = "thread-1"; + +function testThread(): Thread { + return { + id: THREAD_ID, + projectId: "project-1", + title: "Usage ledger", + // Account-scoped kind on purpose: the ledger must keep the full kind so + // different profiles of one provider are counted separately. + agentKind: "claude:work", + config: { model: "claude-sonnet-4" }, + status: "working", + attention: "working", + canResumeWithConfig: false, + archived: false, + done: false, + starred: false, + presentationMode: "gui", + createdAt: "2026-01-01T00:00:00.000Z", + updatedAt: "2026-01-01T00:00:00.000Z", + }; +} + +let seq = 0; +function spent(usage: Partial & { counter: number }): RuntimeEvent { + seq += 1; + return { + type: "usage.spent", + threadId: THREAD_ID, + usage: { + counterKind: "cumulative", + scopeId: "scope-1", + epoch: 0, + sampleId: `sample-${seq}`, + ...usage, + }, + }; +} + +interface TokenRow { + ts: number; + kind: string; + provider: string | null; + model: string | null; + value: number; +} + +// Read rows straight from sqlite: dbGetAllUsageEvents caches on the (process- +// global) profile generation, which is not reset between this file's tests. +function tokenRows(): TokenRow[] { + return getSqlite() + .prepare("SELECT ts, kind, provider, model, value FROM usage_events ORDER BY id") + .all() as TokenRow[]; +} + +function ledgerCounter(scopeId: string, epoch: number): number | undefined { + const row = getSqlite() + .prepare( + "SELECT last_counter FROM usage_token_ledger WHERE provider = 'claude:work' AND scope_id = ? AND epoch = ?", + ) + .get(scopeId, epoch) as { last_counter: number } | undefined; + return row?.last_counter; +} + +function sampleCount(): number { + const row = getSqlite().prepare("SELECT COUNT(*) AS n FROM usage_token_samples").get() as { + n: number; + }; + return row.n; +} + +describe.skipIf(!sqliteAvailable)("usageLedger (real sqlite round-trip)", () => { + let dir: string; + + beforeEach(() => { + if (nativeBindingEnv) { + process.env.PORACODE_BETTER_SQLITE3_NATIVE_BINDING = nativeBindingEnv; + } + seq = 0; + dir = mkdtempSync(join(tmpdir(), "poracode-usage-ledger-test-")); + initDatabase(join(dir, "state.sqlite")); + dbUpsertProject( + { + id: "project-1", + name: "Test project", + location: { kind: "posix", path: "/tmp/project" }, + createdAt: "2026-01-01T00:00:00.000Z", + }, + 0, + ); + dbUpsertThread(testThread(), 0); + }); + + afterEach(() => { + closeDatabase(); + rmSync(dir, { recursive: true, force: true }); + delete process.env.PORACODE_BETTER_SQLITE3_NATIVE_BINDING; + vi.restoreAllMocks(); + }); + + it("counts the full first value for a fresh cumulative scope", () => { + recordUsageSpentFromRuntimeEvents(THREAD_ID, [ + spent({ counter: 1200, fresh: true, occurredAt: 1000 }), + ]); + + expect(tokenRows()).toEqual([ + { + ts: 1000, + kind: "tokens_v2", + provider: "claude:work", + model: "claude-sonnet-4", + value: 1200, + }, + ]); + expect(ledgerCounter("scope-1", 0)).toBe(1200); + }); + + it("establishes a zero baseline for a non-fresh first sample (resume-safe)", () => { + // A resumed scope's first counter reflects already-counted history, so it + // only establishes the baseline; genuine growth after it counts the delta. + recordUsageSpentFromRuntimeEvents(THREAD_ID, [spent({ counter: 5000 })]); + expect(tokenRows()).toEqual([]); + expect(ledgerCounter("scope-1", 0)).toBe(5000); + + recordUsageSpentFromRuntimeEvents(THREAD_ID, [spent({ counter: 5200 })]); + expect(tokenRows().map((row) => row.value)).toEqual([200]); + }); + + it("counts increases, and counts nothing for equal or decreased counters", () => { + recordUsageSpentFromRuntimeEvents(THREAD_ID, [spent({ counter: 1000, fresh: true })]); + recordUsageSpentFromRuntimeEvents(THREAD_ID, [spent({ counter: 1600 })]); // +600 + recordUsageSpentFromRuntimeEvents(THREAD_ID, [spent({ counter: 1600 })]); // equal: 0 + // Decrease: no reset heuristic — out-of-order/replayed samples count 0 and + // the high-water mark stays, so the next increase is measured from 1600. + recordUsageSpentFromRuntimeEvents(THREAD_ID, [spent({ counter: 900 })]); + recordUsageSpentFromRuntimeEvents(THREAD_ID, [spent({ counter: 1700 })]); // +100 + + expect(tokenRows().map((row) => row.value)).toEqual([1000, 600, 100]); + expect(ledgerCounter("scope-1", 0)).toBe(1700); + }); + + it("starts a fresh baseline when the epoch bumps for the same scope", () => { + recordUsageSpentFromRuntimeEvents(THREAD_ID, [spent({ counter: 1000, fresh: true })]); + // New epoch, non-fresh: first sample in the new epoch is baseline-only + // even though it is lower than the previous epoch's counter. + recordUsageSpentFromRuntimeEvents(THREAD_ID, [spent({ counter: 300, epoch: 1 })]); + expect(tokenRows().map((row) => row.value)).toEqual([1000]); + expect(ledgerCounter("scope-1", 1)).toBe(300); + + recordUsageSpentFromRuntimeEvents(THREAD_ID, [spent({ counter: 450, epoch: 1 })]); + expect(tokenRows().map((row) => row.value)).toEqual([1000, 150]); + }); + + it("sums per-call samples and ignores replayed sample ids", () => { + const batch = [ + spent({ counterKind: "per-call", counter: 100, sampleId: "call-1" }), + spent({ counterKind: "per-call", counter: 200, sampleId: "call-2" }), + ]; + recordUsageSpentFromRuntimeEvents(THREAD_ID, batch); + // A replay of the same batch (e.g. after a resume) must count nothing. + recordUsageSpentFromRuntimeEvents(THREAD_ID, batch); + recordUsageSpentFromRuntimeEvents(THREAD_ID, [ + spent({ counterKind: "per-call", counter: 50, sampleId: "call-3" }), + ]); + + expect(tokenRows().map((row) => row.value)).toEqual([100, 200, 50]); + expect(sampleCount()).toBe(3); + }); + + it("commits a mixed batch together and ignores non-usage events", () => { + recordUsageSpentFromRuntimeEvents(THREAD_ID, [ + spent({ + counter: 800, + fresh: true, + scopeId: "scope-a", + occurredAt: 5000, + model: "claude-opus-4", + }), + spent({ + counterKind: "per-call", + counter: 120, + scopeId: "scope-b", + sampleId: "call-x", + occurredAt: 5001, + }), + { + type: "item.started", + threadId: THREAD_ID, + itemId: "item-1", + itemType: "assistant_message", + }, + { type: "context.updated", threadId: THREAD_ID, usage: { usedTokens: 9999 } }, + ]); + + // Event-level model wins; without it the row falls back to the thread model. + expect(tokenRows()).toEqual([ + { + ts: 5000, + kind: "tokens_v2", + provider: "claude:work", + model: "claude-opus-4", + value: 800, + }, + { + ts: 5001, + kind: "tokens_v2", + provider: "claude:work", + model: "claude-sonnet-4", + value: 120, + }, + ]); + // Ledger state, dedup samples, and rows all committed from the one batch. + expect(ledgerCounter("scope-a", 0)).toBe(800); + expect(sampleCount()).toBe(1); + }); + + it("falls back to Date.now when occurredAt is absent", () => { + vi.spyOn(Date, "now").mockReturnValue(42000); + recordUsageSpentFromRuntimeEvents(THREAD_ID, [spent({ counter: 64, fresh: true })]); + + expect(tokenRows().map((row) => row.ts)).toEqual([42000]); + }); + + it("skips events for an unknown thread without throwing", () => { + expect(() => + recordUsageSpentFromRuntimeEvents("missing-thread", [spent({ counter: 100, fresh: true })]), + ).not.toThrow(); + + expect(tokenRows()).toEqual([]); + expect(sampleCount()).toBe(0); + }); +}); diff --git a/src/main/profile/usageLedger.ts b/src/main/profile/usageLedger.ts new file mode 100644 index 000000000..497524e7a --- /dev/null +++ b/src/main/profile/usageLedger.ts @@ -0,0 +1,105 @@ +import type { RuntimeEvent } from "@/shared/contracts"; +import { dbAppendUsageEvents, dbGetThread } from "../db"; +import { getSqlite } from "../db/connection"; +import type { UsageEventInput } from "../db/usageEvents"; + +/** + * Main-process token ledger: turns canonical `usage.spent` events into durable + * `usage_events` rows (kind="tokens_v2"). Lives in main (not the renderer) so + * recording works with the window closed and on headless hosts — every + * supervisor event already routes through runtimePersistence on both. + * + * Semantics per `usageSpentSchema`: + * - cumulative: counter state keyed by (provider, scopeId, epoch) in + * `usage_token_ledger`. First sample establishes the baseline (counts 0, + * resume-safe) unless the adapter marks the scope `fresh` (baseline 0, first + * counter counts in full). Increases count the delta; equal/lower samples + * count 0 (out-of-order safe). Resets are signalled by an epoch bump, never + * inferred from a decrease. + * - per-call: each sample sums directly, deduped exact-once by `sampleId` in + * `usage_token_samples` (replay-safe). + * + * The whole event batch — ledger-state updates, sample inserts, and the + * usage_events append — commits in ONE sqlite transaction (dbAppendUsageEvents + * nests via savepoint), so a crash can never count without persisting state or + * vice versa. + */ + +type UsageSpentEvent = Extract; + +// Threads legitimately disappear (delete) while a supervisor stream is still +// flushing. Log a missing thread once per process, not per event. +const loggedMissingThreads = new Set(); + +export function recordUsageSpentFromRuntimeEvents( + threadId: string, + events: readonly RuntimeEvent[], +): void { + const spentEvents: UsageSpentEvent[] = []; + for (const event of events) { + if (event.type === "usage.spent") spentEvents.push(event); + } + if (spentEvents.length === 0) return; + + const thread = dbGetThread(threadId); + if (!thread) { + if (!loggedMissingThreads.has(threadId)) { + loggedMissingThreads.add(threadId); + console.warn(`[usage-ledger] skipping usage.spent for unknown thread ${threadId}`); + } + return; + } + // Full account-scoped kind (e.g. "claude:work"), matching existing + // usage_events conventions; the global rollup folds to the base provider. + const provider = thread.agentKind; + + const sqlite = getSqlite(); + const readLedger = sqlite.prepare( + "SELECT last_counter FROM usage_token_ledger WHERE provider = ? AND scope_id = ? AND epoch = ?", + ); + const writeLedger = sqlite.prepare( + `INSERT INTO usage_token_ledger (provider, scope_id, epoch, last_counter) VALUES (?, ?, ?, ?) + ON CONFLICT(provider, scope_id, epoch) DO UPDATE SET last_counter = excluded.last_counter`, + ); + const insertSample = sqlite.prepare( + "INSERT OR IGNORE INTO usage_token_samples (sample_id, ts) VALUES (?, ?)", + ); + + sqlite.transaction(() => { + const rows: UsageEventInput[] = []; + for (const event of spentEvents) { + const usage = event.usage; + const ts = usage.occurredAt ?? Date.now(); + let amount = 0; + if (usage.counterKind === "cumulative") { + const existing = readLedger.get(provider, usage.scopeId, usage.epoch) as + | { last_counter: number } + | undefined; + if (!existing) { + writeLedger.run(provider, usage.scopeId, usage.epoch, usage.counter); + amount = usage.fresh === true ? usage.counter : 0; + } else { + const delta = usage.counter - existing.last_counter; + if (delta > 0) { + writeLedger.run(provider, usage.scopeId, usage.epoch, usage.counter); + amount = delta; + } + // Equal/out-of-order samples count nothing and leave state untouched. + } + } else { + // per-call: counted only when the dedup insert actually happened. + const inserted = insertSample.run(usage.sampleId, ts); + amount = inserted.changes === 1 ? usage.counter : 0; + } + if (amount <= 0) continue; + rows.push({ + ts, + kind: "tokens_v2", + provider, + model: usage.model ?? thread.config.model ?? null, + value: amount, + }); + } + dbAppendUsageEvents(rows); + })(); +} diff --git a/src/main/remote/server/runtimePersistence.ts b/src/main/remote/server/runtimePersistence.ts index 94514da1b..14ff85e2e 100644 --- a/src/main/remote/server/runtimePersistence.ts +++ b/src/main/remote/server/runtimePersistence.ts @@ -1,4 +1,5 @@ import { dbApplyThreadRuntimeEvents, dbReplaceThreadRuntimeSnapshot } from "../../db"; +import { recordUsageSpentFromRuntimeEvents } from "../../profile/usageLedger"; import type { RemoteBroadcastEvent } from "./context"; import { persistThreadStateEvent } from "./threadStatePersistence"; @@ -16,21 +17,25 @@ export function persistSupervisorEvent(event: RemoteBroadcastEvent): void { } /** - * Applies canonical supervisor runtime events directly to SQLite. The same - * writer is used by desktop main and the headless remote host, so durability - * never depends on a mounted renderer or a full transcript rewrite. + * Applies canonical supervisor runtime events directly to SQLite — chat items + * and, for `usage.spent`, the token usage ledger. The same writer is used by + * desktop main and the headless remote host, so durability never depends on a + * mounted renderer or a full transcript rewrite. */ function persistRuntimeEvent(event: RemoteBroadcastEvent): void { switch (event.type) { case "thread-runtime-event": dbApplyThreadRuntimeEvents(event.threadId, [event.event]); + recordUsageSpentFromRuntimeEvents(event.threadId, [event.event]); return; case "thread-runtime-events": dbApplyThreadRuntimeEvents(event.threadId, event.events); + recordUsageSpentFromRuntimeEvents(event.threadId, event.events); return; case "thread-runtime-events-multi": for (const batch of event.batches) { dbApplyThreadRuntimeEvents(batch.threadId, batch.events); + recordUsageSpentFromRuntimeEvents(batch.threadId, batch.events); } return; case "thread-reset": diff --git a/src/renderer/locales/de/messages.po b/src/renderer/locales/de/messages.po index 5bd6efc09..2366ef9bc 100644 --- a/src/renderer/locales/de/messages.po +++ b/src/renderer/locales/de/messages.po @@ -10150,6 +10150,10 @@ msgstr "Das Dateibaum-Panel umschalten" msgid "Toggle Work or Plan" msgstr "Zwischen Arbeit und Plan umschalten" +#: src/renderer/views/ProfileOverlay/parts/StatStrip.tsx +msgid "Token usage unavailable for: {unavailableLabels}" +msgstr "Token-Nutzung nicht verfügbar für: {unavailableLabels}" + #: src/renderer/views/ProfileOverlay/parts/ActivitySection.tsx msgid "Tokens" msgstr "Tokens" diff --git a/src/renderer/locales/en/messages.po b/src/renderer/locales/en/messages.po index 3737001f3..f0f5bbae6 100644 --- a/src/renderer/locales/en/messages.po +++ b/src/renderer/locales/en/messages.po @@ -10150,6 +10150,10 @@ msgstr "Toggle the file tree panel" msgid "Toggle Work or Plan" msgstr "Toggle Work or Plan" +#: src/renderer/views/ProfileOverlay/parts/StatStrip.tsx +msgid "Token usage unavailable for: {unavailableLabels}" +msgstr "Token usage unavailable for: {unavailableLabels}" + #: src/renderer/views/ProfileOverlay/parts/ActivitySection.tsx msgid "Tokens" msgstr "Tokens" diff --git a/src/renderer/locales/es/messages.po b/src/renderer/locales/es/messages.po index bd3b03a71..faaeaf18d 100644 --- a/src/renderer/locales/es/messages.po +++ b/src/renderer/locales/es/messages.po @@ -10150,6 +10150,10 @@ msgstr "Alternar el panel del árbol de archivos" msgid "Toggle Work or Plan" msgstr "Alternar Trabajo o Plan" +#: src/renderer/views/ProfileOverlay/parts/StatStrip.tsx +msgid "Token usage unavailable for: {unavailableLabels}" +msgstr "Uso de tokens no disponible para: {unavailableLabels}" + #: src/renderer/views/ProfileOverlay/parts/ActivitySection.tsx msgid "Tokens" msgstr "Tokens" diff --git a/src/renderer/locales/fr/messages.po b/src/renderer/locales/fr/messages.po index 32a64668d..33bd4d9c3 100644 --- a/src/renderer/locales/fr/messages.po +++ b/src/renderer/locales/fr/messages.po @@ -10149,6 +10149,10 @@ msgstr "Basculer le panneau d'arborescence de fichiers" msgid "Toggle Work or Plan" msgstr "Basculer entre Travail et Plan" +#: src/renderer/views/ProfileOverlay/parts/StatStrip.tsx +msgid "Token usage unavailable for: {unavailableLabels}" +msgstr "Utilisation des tokens indisponible pour : {unavailableLabels}" + #: src/renderer/views/ProfileOverlay/parts/ActivitySection.tsx msgid "Tokens" msgstr "Tokens" diff --git a/src/renderer/locales/ja/messages.po b/src/renderer/locales/ja/messages.po index adafc66a5..2767f7933 100644 --- a/src/renderer/locales/ja/messages.po +++ b/src/renderer/locales/ja/messages.po @@ -10148,6 +10148,10 @@ msgstr "ファイルツリーパネルを切り替えます" msgid "Toggle Work or Plan" msgstr "作業と計画を切り替え" +#: src/renderer/views/ProfileOverlay/parts/StatStrip.tsx +msgid "Token usage unavailable for: {unavailableLabels}" +msgstr "トークン使用量を取得できません: {unavailableLabels}" + #: src/renderer/views/ProfileOverlay/parts/ActivitySection.tsx msgid "Tokens" msgstr "トークン" diff --git a/src/renderer/locales/ko/messages.po b/src/renderer/locales/ko/messages.po index f66a4620f..03816989a 100644 --- a/src/renderer/locales/ko/messages.po +++ b/src/renderer/locales/ko/messages.po @@ -10150,6 +10150,10 @@ msgstr "파일 트리 패널 표시 전환" msgid "Toggle Work or Plan" msgstr "작업 또는 계획 전환" +#: src/renderer/views/ProfileOverlay/parts/StatStrip.tsx +msgid "Token usage unavailable for: {unavailableLabels}" +msgstr "토큰 사용량을 사용할 수 없음: {unavailableLabels}" + #: src/renderer/views/ProfileOverlay/parts/ActivitySection.tsx msgid "Tokens" msgstr "토큰" diff --git a/src/renderer/locales/pl/messages.po b/src/renderer/locales/pl/messages.po index 71d871343..cda0620b9 100644 --- a/src/renderer/locales/pl/messages.po +++ b/src/renderer/locales/pl/messages.po @@ -10150,6 +10150,10 @@ msgstr "Przełącz panel drzewa plików" msgid "Toggle Work or Plan" msgstr "Przełącz Pracę lub Plan" +#: src/renderer/views/ProfileOverlay/parts/StatStrip.tsx +msgid "Token usage unavailable for: {unavailableLabels}" +msgstr "Zużycie tokenów niedostępne dla: {unavailableLabels}" + #: src/renderer/views/ProfileOverlay/parts/ActivitySection.tsx msgid "Tokens" msgstr "Tokeny" diff --git a/src/renderer/locales/pt-BR/messages.po b/src/renderer/locales/pt-BR/messages.po index 8a16c50c0..302e4f532 100644 --- a/src/renderer/locales/pt-BR/messages.po +++ b/src/renderer/locales/pt-BR/messages.po @@ -10150,6 +10150,10 @@ msgstr "Alternar o painel da árvore de arquivos" msgid "Toggle Work or Plan" msgstr "Alternar entre Trabalho e Plano" +#: src/renderer/views/ProfileOverlay/parts/StatStrip.tsx +msgid "Token usage unavailable for: {unavailableLabels}" +msgstr "Uso de tokens indisponível para: {unavailableLabels}" + #: src/renderer/views/ProfileOverlay/parts/ActivitySection.tsx msgid "Tokens" msgstr "Tokens" diff --git a/src/renderer/locales/ru/messages.po b/src/renderer/locales/ru/messages.po index 2bd6f18d4..7362a50d4 100644 --- a/src/renderer/locales/ru/messages.po +++ b/src/renderer/locales/ru/messages.po @@ -10150,6 +10150,10 @@ msgstr "Переключить панель дерева файлов" msgid "Toggle Work or Plan" msgstr "Переключить Работа или План" +#: src/renderer/views/ProfileOverlay/parts/StatStrip.tsx +msgid "Token usage unavailable for: {unavailableLabels}" +msgstr "Данные об использовании токенов недоступны для: {unavailableLabels}" + #: src/renderer/views/ProfileOverlay/parts/ActivitySection.tsx msgid "Tokens" msgstr "Токены" diff --git a/src/renderer/locales/tr/messages.po b/src/renderer/locales/tr/messages.po index 8e2bd2b61..39cc6e92c 100644 --- a/src/renderer/locales/tr/messages.po +++ b/src/renderer/locales/tr/messages.po @@ -10150,6 +10150,10 @@ msgstr "Dosya ağacı panelini aç/kapat" msgid "Toggle Work or Plan" msgstr "İş veya Plan modunu aç/kapat" +#: src/renderer/views/ProfileOverlay/parts/StatStrip.tsx +msgid "Token usage unavailable for: {unavailableLabels}" +msgstr "Token kullanımı şunlar için kullanılamıyor: {unavailableLabels}" + #: src/renderer/views/ProfileOverlay/parts/ActivitySection.tsx msgid "Tokens" msgstr "Token" diff --git a/src/renderer/locales/uk/messages.po b/src/renderer/locales/uk/messages.po index ff0631778..5a94a5d55 100644 --- a/src/renderer/locales/uk/messages.po +++ b/src/renderer/locales/uk/messages.po @@ -10150,6 +10150,10 @@ msgstr "Перемкнути панель дерева файлів" msgid "Toggle Work or Plan" msgstr "Перемкнути Робота або План" +#: src/renderer/views/ProfileOverlay/parts/StatStrip.tsx +msgid "Token usage unavailable for: {unavailableLabels}" +msgstr "Дані про використання токенів недоступні для: {unavailableLabels}" + #: src/renderer/views/ProfileOverlay/parts/ActivitySection.tsx msgid "Tokens" msgstr "Токени" diff --git a/src/renderer/locales/vi/messages.po b/src/renderer/locales/vi/messages.po index 70b898637..c9021eb25 100644 --- a/src/renderer/locales/vi/messages.po +++ b/src/renderer/locales/vi/messages.po @@ -10150,6 +10150,10 @@ msgstr "Bật/tắt bảng điều khiển cây tập tin" msgid "Toggle Work or Plan" msgstr "Chuyển đổi công việc hoặc kế hoạch" +#: src/renderer/views/ProfileOverlay/parts/StatStrip.tsx +msgid "Token usage unavailable for: {unavailableLabels}" +msgstr "Không có dữ liệu sử dụng token cho: {unavailableLabels}" + #: src/renderer/views/ProfileOverlay/parts/ActivitySection.tsx msgid "Tokens" msgstr "Tokens" diff --git a/src/renderer/locales/zh-CN/messages.po b/src/renderer/locales/zh-CN/messages.po index 24554cb00..2770730f6 100644 --- a/src/renderer/locales/zh-CN/messages.po +++ b/src/renderer/locales/zh-CN/messages.po @@ -10149,6 +10149,10 @@ msgstr "切换文件树面板" msgid "Toggle Work or Plan" msgstr "切换工作或计划" +#: src/renderer/views/ProfileOverlay/parts/StatStrip.tsx +msgid "Token usage unavailable for: {unavailableLabels}" +msgstr "以下提供方的 token 用量不可用:{unavailableLabels}" + #: src/renderer/views/ProfileOverlay/parts/ActivitySection.tsx msgid "Tokens" msgstr "Token" diff --git a/src/renderer/state/slices/runtimeEventReducer.ts b/src/renderer/state/slices/runtimeEventReducer.ts index 917eda675..142cdef41 100644 --- a/src/renderer/state/slices/runtimeEventReducer.ts +++ b/src/renderer/state/slices/runtimeEventReducer.ts @@ -209,6 +209,11 @@ function applyRuntimeEventToRuntimeState( // No item state to mutate. Status flows through the existing thread-state channel. return {}; + case "usage.spent": + // Token consumption is persisted by the main-process usage ledger; the + // renderer keeps no token state (the dock reads context.updated only). + return {}; + case "turn.started": // Mark the runtime turn open so live activity may (re)open the GUI turn. // No item state to mutate; status flows through the thread-state channel. diff --git a/src/renderer/state/usageRecorder.test.ts b/src/renderer/state/usageRecorder.test.ts index 06d8d6818..e999397cb 100644 --- a/src/renderer/state/usageRecorder.test.ts +++ b/src/renderer/state/usageRecorder.test.ts @@ -1,5 +1,5 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; -import type { RuntimeEvent, Thread, UsageEventInputPayload } from "@/shared/contracts"; +import type { Thread, UsageEventInputPayload } from "@/shared/contracts"; const bridgeMock = vi.hoisted(() => ({ appendUsageEvents: vi.fn<(payload: { events: UsageEventInputPayload[] }) => Promise>(), @@ -9,7 +9,7 @@ vi.mock("@/renderer/bridge", () => ({ readBridge: () => bridgeMock, })); -import { recordRuntimeUsage, recordThreadStarted } from "./usageRecorder"; +import { recordRuntimeUsage } from "./usageRecorder"; function makeThread(id: string, agentKind: string): Thread { return { @@ -20,29 +20,18 @@ function makeThread(id: string, agentKind: string): Thread { } as unknown as Thread; } -function contextUpdated(threadId: string, usedTokens: number): RuntimeEvent { - return { type: "context.updated", threadId, usage: { usedTokens } }; -} - // The recorder flushes its buffer synchronously on `pagehide`; dispatch it to // drain without waiting on the idle/timeout debounce. function flushNow(): void { window.dispatchEvent(new Event("pagehide")); } -function emittedTokenValues(provider: string): number[] { - return bridgeMock.appendUsageEvents.mock.calls - .flatMap((call) => call[0].events) - .filter((event) => event.kind === "tokens" && event.provider === provider) - .map((event) => event.value ?? 0); -} - function emittedEvents(kind?: string): UsageEventInputPayload[] { const events = bridgeMock.appendUsageEvents.mock.calls.flatMap((call) => call[0].events); return kind === undefined ? events : events.filter((event) => event.kind === kind); } -describe("usageRecorder token baseline", () => { +describe("usageRecorder item hit refinement", () => { beforeEach(() => { flushNow(); // drain anything a prior test left buffered bridgeMock.appendUsageEvents.mockReset(); @@ -52,33 +41,6 @@ describe("usageRecorder token baseline", () => { flushNow(); }); - it("does not re-count a resumed thread's restored context, but counts later growth", () => { - const provider = "resumed-provider"; - const thread = makeThread("resumed-thread", provider); - // Resumed thread: recordThreadStarted was NOT called this session, so there - // is no seeded baseline. Its first context.updated reports the restored - // context (already counted in a prior session) and must emit nothing. - recordRuntimeUsage("resumed-thread", [contextUpdated("resumed-thread", 50_000)], [thread]); - flushNow(); - expect(emittedTokenValues(provider)).toEqual([]); - - // A later context.updated reflects genuine growth and IS counted. - recordRuntimeUsage("resumed-thread", [contextUpdated("resumed-thread", 50_500)], [thread]); - flushNow(); - expect(emittedTokenValues(provider)).toEqual([500]); - }); - - it("counts the full initial context for a thread started this session", () => { - const provider = "new-provider"; - const thread = makeThread("new-thread", provider); - // recordThreadStarted seeds the baseline to 0, so the first context.updated - // counts the whole new context as a delta from zero. - recordThreadStarted(thread); - recordRuntimeUsage("new-thread", [contextUpdated("new-thread", 1_200)], [thread]); - flushNow(); - expect(emittedTokenValues(provider)).toEqual([1_200]); - }); - it("records a skill name from the completed payload when the start was generic", () => { const thread = makeThread("skill-thread", "skill-provider"); recordRuntimeUsage( diff --git a/src/renderer/state/usageRecorder.ts b/src/renderer/state/usageRecorder.ts index 2efeecf53..48eeb1919 100644 --- a/src/renderer/state/usageRecorder.ts +++ b/src/renderer/state/usageRecorder.ts @@ -13,7 +13,9 @@ import { readBridge } from "@/renderer/bridge"; * `RuntimeEvent` stream before they reach here, so this is the single place we * derive usage facts - no per-provider splits. Each fact is written to the * durable `usage_events` log (no thread FK), with provider/model/mode embedded, - * so the stats survive thread delete/archive. + * so the stats survive thread delete/archive. Token consumption is the one + * exception: it is counted by the main-process usage ledger from `usage.spent` + * events (kind="tokens_v2"), never here. * * Writes are buffered and flushed on a debounce (fire-and-forget IPC) so the hot * event path is never blocked. This module intentionally does NOT import the app @@ -33,25 +35,6 @@ const DEDUP_CAP = 20000; let buffer: UsageEventInputPayload[] = []; let scheduled: { cancel: () => void } | null = null; const turnStartByThread = new Map(); -// Cumulative usedTokens per thread, used for the LAG delta. Intentionally NOT -// cleared per turn: usedTokens is cumulative across a thread's whole life, so -// resetting it would make the next delta count the full context again. One int -// per thread that streamed tokens this session - cleared on app restart. -// -// Presence (`.has`), not just value, carries meaning. A thread created THIS -// session is seeded to 0 by recordThreadStarted, so its first context.updated -// counts the whole new context (delta from 0). A thread RESUMED from a prior -// session has no entry: its first context.updated reports a context whose -// tokens were already counted last session, so we only establish the baseline -// and emit nothing - otherwise the restored context would be double-counted on -// every restart, inflating lifetime/peak. See the context.updated case below. -const lastUsedByThread = new Map(); -// Token deltas coalesced per provider|model between flushes, so a turn that -// streams many context.updated events produces ONE row, not dozens. -const pendingTokens = new Map< - string, - { provider: string | null; model: string | null; value: number } ->(); const seenItems = new Set(); const seenTurns = new Set(); const pendingItemHits = new Map(); @@ -62,13 +45,6 @@ function flush(): void { scheduled.cancel(); scheduled = null; } - if (pendingTokens.size > 0) { - const ts = Date.now(); - for (const t of pendingTokens.values()) { - buffer.push({ ts, kind: "tokens", provider: t.provider, model: t.model, value: t.value }); - } - pendingTokens.clear(); - } if (buffer.length === 0) return; const events = buffer; buffer = []; @@ -301,10 +277,6 @@ export function recordAiAction(type: AiActionType, provider: string, model: stri /** Record that a thread was started (provider/model/mode/fast/effort). */ export function recordThreadStarted(thread: Thread): void { const m = metaOf(thread); - // Seed the token baseline so this freshly-started thread's first - // context.updated counts its initial context as new (delta from 0). Resumed - // threads get no seed and are handled in the context.updated case below. - if (!lastUsedByThread.has(thread.id)) lastUsedByThread.set(thread.id, 0); push({ ts: Date.now(), kind: "thread_started", @@ -362,7 +334,6 @@ export function recordRuntimeUsage( }; const now = Date.now(); - let tokensTouched = false; for (const event of events) { switch (event.type) { @@ -392,32 +363,6 @@ export function recordRuntimeUsage( turnStartByThread.delete(threadId); break; } - case "context.updated": { - const used = event.usage.usedTokens; - if (typeof used === "number" && used > 0) { - // No baseline means this thread was resumed from a prior session: its - // usedTokens already reflects context counted then, so only establish - // the baseline and emit nothing. Counting delta-from-zero here would - // re-add the whole restored context. Threads started this session are - // seeded to 0 by recordThreadStarted, so they DO count their first - // context. - const hasBaseline = lastUsedByThread.has(threadId); - const prev = lastUsedByThread.get(threadId) ?? 0; - lastUsedByThread.set(threadId, used); - const delta = hasBaseline ? Math.max(0, used - prev) : 0; - if (delta > 0) { - const m = getMeta(); - if (m) { - const key = `${m.provider}|${m.model ?? ""}`; - const entry = pendingTokens.get(key); - if (entry) entry.value += delta; - else pendingTokens.set(key, { provider: m.provider, model: m.model, value: delta }); - tokensTouched = true; - } - } - } - break; - } case "item.started": { rememberInMap(itemTypesById, event.itemId, event.itemType); const hit = classifyItem(event.itemType, event.payload); @@ -458,6 +403,4 @@ export function recordRuntimeUsage( break; } } - - if (tokensTouched) scheduleFlush(); } diff --git a/src/renderer/views/ProfileOverlay/parts/StatStrip.tsx b/src/renderer/views/ProfileOverlay/parts/StatStrip.tsx index 1ca75d041..0b85ef6ca 100644 --- a/src/renderer/views/ProfileOverlay/parts/StatStrip.tsx +++ b/src/renderer/views/ProfileOverlay/parts/StatStrip.tsx @@ -1,6 +1,6 @@ import type { ReactNode } from "react"; import { msg } from "@lingui/core/macro"; -import { useLingui } from "@lingui/react/macro"; +import { Trans, useLingui } from "@lingui/react/macro"; import type { ProfileCoreStats, ProfileStatsWindow, ProfileTokenStats } from "@/shared/contracts"; import type { TranslateFn } from "@/renderer/i18n/i18n"; import { formatCompact, formatDayLabel, formatDuration } from "../format"; @@ -30,6 +30,17 @@ function formatDaysLabel(days: number, t: TranslateFn): string { return days === 1 ? t(msg`${days} day`) : t(msg`${days} days`); } +/** Display name for a provider key — the stats label when present, else title-case. */ +function unavailableProviderLabel(key: string, tokens: ProfileTokenStats): string { + const known = tokens.providers.find((p) => p.provider === key)?.label; + if (known) return known; + return key + .split(/[-_:]+/) + .filter(Boolean) + .map((part) => part.charAt(0).toUpperCase() + part.slice(1)) + .join(" "); +} + export function StatStrip(props: { core: ProfileCoreStats; tokens: ProfileTokenStats | null; @@ -56,17 +67,28 @@ export function StatStrip(props: { "-" ); + const unavailableLabels = + tokens?.unavailableProviders.map((key) => unavailableProviderLabel(key, tokens)).join(", ") ?? + ""; + return ( -
- - - - - +
+
+ + + + + +
+ {tokens && tokens.unavailableProviders.length > 0 ? ( +

+ Token usage unavailable for: {unavailableLabels} +

+ ) : null}
); } diff --git a/src/shared/contracts/profile.ts b/src/shared/contracts/profile.ts index 13e201c7b..b8f540143 100644 --- a/src/shared/contracts/profile.ts +++ b/src/shared/contracts/profile.ts @@ -245,6 +245,13 @@ export interface ProfileTokenStats { /** Token-weighted model mix. */ models: ProfileBreakdownEntry[]; tokenHeatmap: ProfileHeatmap; + /** + * Base-provider keys (e.g. "kimi") with recorded activity (thread_started or + * turn rows) but zero exact token rows (kind="tokens_v2") — their token + * spend is not measured. Sorted; empty when every active provider reports + * exact usage. + */ + unavailableProviders: string[]; } // -- IPC payloads ----------------------------------------------------- diff --git a/src/shared/contracts/runtimeEvent.ts b/src/shared/contracts/runtimeEvent.ts index a5fb6b686..1cee0f9dc 100644 --- a/src/shared/contracts/runtimeEvent.ts +++ b/src/shared/contracts/runtimeEvent.ts @@ -295,6 +295,38 @@ export const threadContextUsageSchema = z.object({ }); export type ThreadContextUsage = z.infer; +/** + * Provider-reported token CONSUMPTION, kept strictly separate from + * `context.updated` (which carries context-window occupancy for the dock). + * Adapters normalize their native usage payloads into this shape; the + * main-process usage ledger is the only consumer that persists it. + * + * - `counterKind: "cumulative"` — `counter` is an absolute, monotonically + * increasing total for the scope `(provider, scopeId, epoch)` (e.g. Codex + * `total_token_usage.total_tokens`). The ledger counts increases, ignores + * equal/lower values (out-of-order safe); resets are signalled by bumping + * `epoch`, never inferred from a counter decrease. + * - `counterKind: "per-call"` — `counter` is one API call's total + * (e.g. Claude assistant-message usage). The ledger sums these; `sampleId` + * gives exact-once dedup across replays. + * `scopeId` is the PROVIDER-side scope (session/thread id), not the Poracode + * thread id, so resume/fork semantics stay explicit. `fresh: true` marks a + * scope the adapter knows was just created (baseline 0); otherwise the first + * sample in a scope+epoch establishes the baseline and counts nothing. + */ +export const usageSpentSchema = z.object({ + counterKind: z.enum(["cumulative", "per-call"]), + counter: z.number().int().nonnegative(), + scopeId: z.string().min(1), + epoch: z.number().int().nonnegative(), + fresh: z.boolean().optional(), + sampleId: z.string().min(1), + turnId: z.string().optional(), + occurredAt: z.number().int().nonnegative().optional(), + model: z.string().optional(), +}); +export type UsageSpent = z.infer; + // ── Request payloads ───────────────────────────────────────────────── export const userInputOptionSchema = z.object({ @@ -465,6 +497,11 @@ export const runtimeEventSchema = z.discriminatedUnion("type", [ threadId: z.string(), usage: threadContextUsageSchema, }), + z.object({ + type: z.literal("usage.spent"), + threadId: z.string(), + usage: usageSpentSchema, + }), z.object({ type: z.literal("request.opened"), threadId: z.string(), diff --git a/src/supervisor/agents/acp/session.test.ts b/src/supervisor/agents/acp/session.test.ts index 2d31fdc3d..68dc53164 100644 --- a/src/supervisor/agents/acp/session.test.ts +++ b/src/supervisor/agents/acp/session.test.ts @@ -168,6 +168,9 @@ function makeConfigSyncSession( session["agentSessionCapabilities"] = undefined; session["cwd"] = "C:\\repo"; session["stableSessionRef"] = undefined; + session["usageScopeId"] = undefined; + session["usageEpoch"] = 0; + session["usageScopeFresh"] = false; session["launchOptions"] = {}; session["mcpServers"] = overrides.mcpServers ?? []; session["loadSessionErrorRewriter"] = rewriteLoadSessionError; @@ -309,6 +312,98 @@ describe("ACP empty-response provider guard", () => { }); }); +describe("ACP prompt-response usage → usage.spent", () => { + it("emits cumulative usage.spent from a new session's prompt usage, fresh once", async () => { + const { connection, listener, session } = makeConfigSyncSession(); + await session.openThread({ model: "model-a" }); + + connection.prompt.mockResolvedValueOnce({ + stopReason: "end_turn", + usage: { totalTokens: 1200, inputTokens: 1000, outputTokens: 200 }, + } as { stopReason: string }); + await session.startTurn("hello", { model: "model-a" }); + + expect(listener.onRuntimeEvent).toHaveBeenCalledWith({ + type: "usage.spent", + threadId: "thread-1", + usage: { + counterKind: "cumulative", + counter: 1200, + scopeId: "session-1", + epoch: 0, + fresh: true, + sampleId: "session-1:0:1200", + }, + }); + // The dock's context.updated is still emitted from the same payload. + expect(listener.onRuntimeEvent).toHaveBeenCalledWith( + expect.objectContaining({ type: "context.updated" }), + ); + + // Second turn: same scope, `fresh` consumed by the first sample. + connection.prompt.mockResolvedValueOnce({ + stopReason: "end_turn", + usage: { totalTokens: 1500 }, + } as { stopReason: string }); + await session.startTurn("again", { model: "model-a" }); + + expect(listener.onRuntimeEvent).toHaveBeenCalledWith({ + type: "usage.spent", + threadId: "thread-1", + usage: { + counterKind: "cumulative", + counter: 1500, + scopeId: "session-1", + epoch: 0, + sampleId: "session-1:0:1500", + }, + }); + }); + + it("marks resumed sessions non-fresh and emits nothing without prompt usage", async () => { + const { connection, listener, session } = makeConfigSyncSession(); + await session.openThread( + { model: "model-a" }, + { providerSessionId: "session-resume", discoveredAt: new Date().toISOString() }, + ); + + connection.prompt.mockResolvedValueOnce({ + stopReason: "end_turn", + usage: { totalTokens: 9000 }, + } as { stopReason: string }); + await session.startTurn("hello", { model: "model-a" }); + + const spentEvents = () => + listener.onRuntimeEvent.mock.calls + .map(([event]) => event) + .filter((event): event is Record => { + return ( + typeof event === "object" && + event !== null && + (event as { type?: string }).type === "usage.spent" + ); + }); + expect(spentEvents()).toHaveLength(1); + expect(spentEvents()[0]).toMatchObject({ + usage: { + counterKind: "cumulative", + counter: 9000, + scopeId: "session-resume", + epoch: 0, + sampleId: "session-resume:0:9000", + }, + }); + // Resumed scope: baseline-only first sample, no fresh flag. + expect(spentEvents()[0]).not.toMatchObject({ usage: { fresh: true } }); + + // Bridges without prompt-response usage emit no spend event at all. + listener.onRuntimeEvent.mockClear(); + connection.prompt.mockResolvedValueOnce({ stopReason: "end_turn" }); + await session.startTurn("hello again", { model: "model-a" }); + expect(spentEvents()).toHaveLength(0); + }); +}); + describe("rewriteLoadSessionError — user-facing copy for session/load failures", () => { it("rewrites a 'Session not found' invalidParams into resume-specific guidance", () => { const raw = RequestError.invalidParams({ message: 'Session "abc-123" not found' }); diff --git a/src/supervisor/agents/acp/session.ts b/src/supervisor/agents/acp/session.ts index 2eb0fb964..39ce96a95 100644 --- a/src/supervisor/agents/acp/session.ts +++ b/src/supervisor/agents/acp/session.ts @@ -104,6 +104,7 @@ import { AcpTerminalManager } from "./terminalManager"; import { appendInterruptAckTextTail, createAcpPromptUsageEvent, + createAcpPromptUsageSpentEvent, isAcpPromptCancellationError, normalizeAcpStopReason, resolveAcpPromptFailureMessage, @@ -188,6 +189,14 @@ export class AcpStructuredSession implements StructuredSessionHandle { private detachedTurnId: string | undefined; private readonly detachedTurnParentToolCallIds = new Set(); private stableSessionRef: SessionRef | undefined; + /** + * usage.spent ledger scope: the ACP session id plus an epoch that bumps if + * the id ever changes, and a `fresh` flag consumed by the first emitted + * sample (true only for sessions this handle created via `session/new`). + */ + private usageScopeId: string | undefined; + private usageEpoch = 0; + private usageScopeFresh = false; /** * True while a `connection.prompt()` call is in flight (between issue and * resolution). Used together with `pendingPromptInterrupt` to close the @@ -594,6 +603,19 @@ export class AcpStructuredSession implements StructuredSessionHandle { return kept; } + /** + * Track the usage.spent ledger scope. A changed session id ends the old + * counter lineage — bump the epoch rather than inferring a reset from the + * cumulative counter. `fresh` marks sessions created via `session/new` + * (baseline 0); resumed/loaded sessions get a baseline-only first sample. + */ + private trackUsageScope(sessionId: string, fresh: boolean): void { + if (this.usageScopeId === sessionId) return; + if (this.usageScopeId !== undefined) this.usageEpoch += 1; + this.usageScopeId = sessionId; + this.usageScopeFresh = fresh; + } + async openThread(config: ThreadConfig, sessionRef?: SessionRef): Promise { let availableModeIds: string[] = []; let configOptions: unknown[] | null | undefined; @@ -614,6 +636,7 @@ export class AcpStructuredSession implements StructuredSessionHandle { mcpServers, }); this.adoptSessionRef(sessionRef); + this.trackUsageScope(sessionRef.providerSessionId, false); availableModeIds = result.modes?.availableModes?.map((m) => m.id) ?? []; configOptions = result.configOptions; } catch (error) { @@ -633,6 +656,7 @@ export class AcpStructuredSession implements StructuredSessionHandle { mcpServers, }); this.adoptSessionRef(sessionRef); + this.trackUsageScope(sessionRef.providerSessionId, false); availableModeIds = result.modes?.availableModes?.map((m) => m.id) ?? []; configOptions = result.configOptions; } catch (error) { @@ -650,6 +674,7 @@ export class AcpStructuredSession implements StructuredSessionHandle { }); this.sessionId = result.sessionId; this.stableSessionRef = createKnownSessionRef(result.sessionId); + this.trackUsageScope(result.sessionId, true); availableModeIds = result.modes?.availableModes?.map((m) => m.id) ?? []; configOptions = result.configOptions; console.log("[acp] session created:", this.sessionId, "modes:", availableModeIds); @@ -756,6 +781,20 @@ export class AcpStructuredSession implements StructuredSessionHandle { }); const usageEvent = createAcpPromptUsageEvent(this.threadId, result.usage); if (usageEvent) this.emitRuntimeEvents([usageEvent]); + // The same prompt response also carries the session-cumulative counter + // for the token ledger (absent on most bridges — then nothing is emitted + // and the provider lands on the profile's unavailable list). + if (this.usageScopeId) { + const spentEvent = createAcpPromptUsageSpentEvent(this.threadId, result.usage, { + scopeId: this.usageScopeId, + epoch: this.usageEpoch, + ...(this.usageScopeFresh ? { fresh: true } : {}), + }); + if (spentEvent) { + this.emitRuntimeEvents([spentEvent]); + this.usageScopeFresh = false; + } + } // Map stopReason to Poracode status const normalizedStopReason = normalizeAcpStopReason(result.stopReason, { diff --git a/src/supervisor/agents/acp/sessionErrors.ts b/src/supervisor/agents/acp/sessionErrors.ts index 5dd3f341a..1cc471264 100644 --- a/src/supervisor/agents/acp/sessionErrors.ts +++ b/src/supervisor/agents/acp/sessionErrors.ts @@ -25,6 +25,37 @@ export function createAcpPromptUsageEvent( ); } +/** + * Cumulative `usage.spent` from the same `session/prompt` response usage the + * context event above parses: per the ACP schema, `usage.totalTokens` is the + * session-cumulative counter, so the ledger counts increases per + * (provider, scopeId, epoch). `sampleId` folds the counter value in, which + * makes replays of the same prompt response exact-once. When the agent + * returns no prompt-response usage (most bridges today), emit nothing — the + * profile honesty list covers those providers. + */ +export function createAcpPromptUsageSpentEvent( + threadId: string, + usage: unknown, + scope: { scopeId: string; epoch: number; fresh?: boolean }, +): RuntimeEvent | undefined { + if (!usage || typeof usage !== "object") return undefined; + const totalTokens = readNonNegativeInteger((usage as Record).totalTokens); + if (totalTokens === undefined) return undefined; + return { + type: "usage.spent", + threadId, + usage: { + counterKind: "cumulative", + counter: totalTokens, + scopeId: scope.scopeId, + epoch: scope.epoch, + ...(scope.fresh ? { fresh: true } : {}), + sampleId: `${scope.scopeId}:${scope.epoch}:${totalTokens}`, + }, + }; +} + /** * Replace the raw JSON-RPC error from `session/load` with a message the * renderer can show verbatim. Provider-agnostic on purpose: the same code diff --git a/src/supervisor/agents/claude/canonicalMapping/dispatch.ts b/src/supervisor/agents/claude/canonicalMapping/dispatch.ts index 49309fd65..c99454e86 100644 --- a/src/supervisor/agents/claude/canonicalMapping/dispatch.ts +++ b/src/supervisor/agents/claude/canonicalMapping/dispatch.ts @@ -10,6 +10,7 @@ import { fileChangeMetadataFromToolResult, inputFingerprint, newItemId, + readClaudeAssistantMessageId, tryParseJsonRecord, } from "./helpers"; import { applyPlanAggregatorInput, bindTaskCreateResult } from "./planMapping"; @@ -30,6 +31,7 @@ import { completeTextItem, ensureTextItem, closeClaudeOpenItems } from "./textIt import { createToolItemState, hasToolCallPayload, isSubAgentToolName } from "./toolClassification"; import { startToolItem, syncSubAgentModelProgress } from "./toolItems"; import { toolPayload } from "./toolPayload"; +import { createClaudeUsageSpentEvent, readClaudeAssistantSpendTokens } from "./usageSpent"; import { workflowFromToolUseResult } from "./workflowOutput"; export function readParentToolUseId(message: SDKMessage): string | undefined { @@ -37,12 +39,6 @@ export function readParentToolUseId(message: SDKMessage): string | undefined { return typeof value === "string" && value.length > 0 ? value : undefined; } -function readClaudeAssistantMessageId(message: unknown): string | undefined { - if (!message || typeof message !== "object") return undefined; - const value = (message as { id?: unknown }).id; - return typeof value === "string" && value.length > 0 ? value : undefined; -} - function tagParent( events: RuntimeEvent[], parentItemId: string | undefined, @@ -372,13 +368,23 @@ function mapClaudeSdkMessageInner( } if (message.type === "assistant") { + // Per-call spend for EVERY assistant API message — main-thread and + // sub-agent (parent-attributed) alike — emitted here, exactly once, before + // either rendering path. Never emitted from task_progress/task_notification + // or result.usage: those would double-count sidechain calls. + const usageSpentEvent = + state.usageScope && readClaudeAssistantSpendTokens(message) !== undefined + ? createClaudeUsageSpentEvent(state.threadId, message, state.usageScope.sample()) + : undefined; // Sub-agent (parent-attributed) whole messages must not touch the shared // main-lane per-index maps — index 0 of a sub-agent message would collide // with the main thread's streaming block at index 0. Emit self-contained, // already-complete child items instead (tagParent attaches parentItemId). if (readParentToolUseId(message)) { - return flushSubAgentAssistantMessage(message, state); + const subAgentEvents = flushSubAgentAssistantMessage(message, state); + return usageSpentEvent ? [usageSpentEvent, ...subAgentEvents] : subAgentEvents; } + if (usageSpentEvent) events.push(usageSpentEvent); const messageId = readClaudeAssistantMessageId(message.message); const skipTextSnapshot = messageId ? state.streamedAssistantMessageIds.has(messageId) : false; const content = (message.message as { content?: unknown }).content; diff --git a/src/supervisor/agents/claude/canonicalMapping/helpers.ts b/src/supervisor/agents/claude/canonicalMapping/helpers.ts index 5e38fe64e..df03ff92e 100644 --- a/src/supervisor/agents/claude/canonicalMapping/helpers.ts +++ b/src/supervisor/agents/claude/canonicalMapping/helpers.ts @@ -7,6 +7,13 @@ import type { FileChangeMetadata } from "../sdkCanonicalMappingState"; export { newItemId } from "../../contextUsage"; +/** The API message id (`msg_…`) of a BetaMessage-ish payload, when present. */ +export function readClaudeAssistantMessageId(message: unknown): string | undefined { + if (!message || typeof message !== "object") return undefined; + const value = (message as { id?: unknown }).id; + return typeof value === "string" && value.length > 0 ? value : undefined; +} + /** * Collect inline images out of a Claude `tool_result` content (Anthropic image * blocks: `{ type: "image", source: { type: "base64", media_type, data } }`) as diff --git a/src/supervisor/agents/claude/canonicalMapping/usageSpent.ts b/src/supervisor/agents/claude/canonicalMapping/usageSpent.ts new file mode 100644 index 000000000..dd5d599bb --- /dev/null +++ b/src/supervisor/agents/claude/canonicalMapping/usageSpent.ts @@ -0,0 +1,90 @@ +/** + * Per-call token spend (`usage.spent`) for Claude SDK sessions. + * + * Every SDK `assistant` message carries the API call's `message.usage`; the + * main-process usage ledger sums these per-call counters with exact-once + * dedup on `sampleId`. Sub-agent (parent_tool_use_id) messages are included — + * they are the only complete record of sidechain spend, so `task_progress` + * (cumulative per task) and `result.usage` must NOT also be counted. + */ + +import type { SDKAssistantMessage } from "@anthropic-ai/claude-agent-sdk"; +import type { RuntimeEvent } from "@/shared/contracts"; +import { readClaudeApiUsageSpendTokens } from "./result"; +import { readClaudeAssistantMessageId } from "./helpers"; + +/** + * Usage scope (SDK session id + epoch) for one Claude session. The SDK can + * adopt a different session id mid-thread (see `shouldAdoptSessionId` in + * sdkSession.ts); the session layer calls {@link adoptScope} so the ledger + * keys that span as a new `(scopeId, epoch)`. `fresh: true` marks a session + * started new (not `--resume`). + */ +export class ClaudeUsageScopeTracker { + private epoch = 0; + private freshPending: boolean; + + constructor( + private scopeId: string, + fresh: boolean, + ) { + this.freshPending = fresh; + } + + /** The SDK assigned/adopted a new session id: same conversation, new scope epoch. */ + adoptScope(scopeId: string): void { + this.scopeId = scopeId; + this.epoch += 1; + } + + sample(): { scopeId: string; epoch: number; fresh?: boolean } { + const fresh = this.freshPending; + this.freshPending = false; + return { + scopeId: this.scopeId, + epoch: this.epoch, + ...(fresh ? { fresh: true } : {}), + }; + } +} + +/** One assistant API call's total spend: input + output + cache creation + cache read. */ +export function readClaudeAssistantSpendTokens(message: SDKAssistantMessage): number | undefined { + return readClaudeApiUsageSpendTokens(message.message?.usage); +} + +export function createClaudeUsageSpentEvent( + threadId: string, + message: SDKAssistantMessage, + meta: { scopeId: string; epoch: number; fresh?: boolean }, +): RuntimeEvent | undefined { + const counter = readClaudeAssistantSpendTokens(message); + if (counter === undefined) return undefined; + const apiMessageId = readClaudeAssistantMessageId(message.message); + const requestId = + typeof message.request_id === "string" && message.request_id.length > 0 + ? message.request_id + : undefined; + // Stable per API message (`msg_…:req_…`) so replays dedup exactly once in + // the ledger; fall back to the envelope uuid when the payload has no id. + const sampleId = apiMessageId + ? requestId + ? `${apiMessageId}:${requestId}` + : apiMessageId + : `uuid:${message.uuid}`; + const model = typeof message.message?.model === "string" ? message.message.model : undefined; + return { + type: "usage.spent", + threadId, + usage: { + counterKind: "per-call", + counter, + scopeId: meta.scopeId, + epoch: meta.epoch, + ...(meta.fresh ? { fresh: true } : {}), + sampleId, + occurredAt: Date.now(), + ...(model ? { model } : {}), + }, + }; +} diff --git a/src/supervisor/agents/claude/sdkCanonicalMapping.test.ts b/src/supervisor/agents/claude/sdkCanonicalMapping.test.ts index 2f0bab837..ba570489d 100644 --- a/src/supervisor/agents/claude/sdkCanonicalMapping.test.ts +++ b/src/supervisor/agents/claude/sdkCanonicalMapping.test.ts @@ -2,6 +2,7 @@ import type { SDKControlGetContextUsageResponse, SDKMessage } from "@anthropic-a import { describe, expect, it, vi } from "vitest"; import { buildClaudeQuestionAnswerEvents, + ClaudeUsageScopeTracker, createClaudeMapperState, emitActiveGoalTokenUpdate, mapClaudeContextUsageResponse, @@ -3656,3 +3657,216 @@ describe("sdkCanonicalMapping — emitActiveGoalTokenUpdate", () => { expect(emitActiveGoalTokenUpdate(state, 1_000)).toBeUndefined(); }); }); + +describe("sdkCanonicalMapping — usage.spent", () => { + function assistantMessage(options: { + id: string; + uuid?: string; + parentToolUseId?: string; + requestId?: string; + model?: string; + usage?: { + input_tokens: number; + output_tokens: number; + cache_creation_input_tokens?: number; + cache_read_input_tokens?: number; + }; + }): SDKMessage { + return { + type: "assistant", + session_id: "claude-session", + uuid: options.uuid ?? `uuid-${options.id}`, + parent_tool_use_id: options.parentToolUseId ?? null, + ...(options.requestId ? { request_id: options.requestId } : {}), + message: { + id: options.id, + role: "assistant", + model: options.model ?? "claude-opus-4-8", + content: [{ type: "text", text: "hello" }], + usage: options.usage ?? { input_tokens: 10, output_tokens: 5 }, + }, + } as unknown as SDKMessage; + } + + function spentOf(state: ReturnType, message: SDKMessage) { + return mapClaudeSdkMessage(message, state).find((event) => event.type === "usage.spent"); + } + + it("emits a per-call usage.spent event with the four-field usage sum", () => { + const state = createClaudeMapperState("thread-1"); + state.usageScope = new ClaudeUsageScopeTracker("claude-session", true); + + const spent = spentOf( + state, + assistantMessage({ + id: "msg-1", + requestId: "req-1", + usage: { + input_tokens: 10, + output_tokens: 5, + cache_creation_input_tokens: 3, + cache_read_input_tokens: 2, + }, + }), + ); + + expect(spent).toEqual({ + type: "usage.spent", + threadId: "thread-1", + usage: { + counterKind: "per-call", + counter: 20, + scopeId: "claude-session", + epoch: 0, + fresh: true, + sampleId: "msg-1:req-1", + occurredAt: expect.any(Number), + model: "claude-opus-4-8", + }, + }); + }); + + it("emits usage.spent for sidechain (parent-attributed) assistant messages", () => { + const state = createClaudeMapperState("thread-1"); + state.usageScope = new ClaudeUsageScopeTracker("claude-session", false); + + const events = mapClaudeSdkMessage( + assistantMessage({ + id: "msg-sub-1", + parentToolUseId: "toolu_parent", + requestId: "req-sub-1", + model: "claude-haiku-4-5-20251001", + usage: { + input_tokens: 100, + output_tokens: 20, + cache_creation_input_tokens: 0, + cache_read_input_tokens: 80, + }, + }), + state, + ); + + expect(events[0]).toEqual({ + type: "usage.spent", + threadId: "thread-1", + usage: { + counterKind: "per-call", + counter: 200, + scopeId: "claude-session", + epoch: 0, + sampleId: "msg-sub-1:req-sub-1", + occurredAt: expect.any(Number), + model: "claude-haiku-4-5-20251001", + }, + }); + // The child items still render alongside the spend event. + expect( + events.some( + (event) => event.type === "item.started" && event.itemType === "assistant_message", + ), + ).toBe(true); + }); + + it("does not emit usage.spent from task_progress, task_notification, or result usage", () => { + const state = createClaudeMapperState("thread-1"); + state.usageScope = new ClaudeUsageScopeTracker("claude-session", false); + + const progress = mapClaudeSdkMessage( + { + type: "system", + subtype: "task_progress", + session_id: "claude-session", + task_id: "task-1", + tool_use_id: "toolu_T1", + description: "Searching", + usage: { total_tokens: 4_200, tool_uses: 3, duration_ms: 1_500 }, + } as unknown as SDKMessage, + state, + ); + expect(progress.some((event) => event.type === "usage.spent")).toBe(false); + + const notification = mapClaudeSdkMessage( + { + type: "system", + subtype: "task_notification", + session_id: "claude-session", + task_id: "task-1", + status: "completed", + summary: "Done", + usage: { total_tokens: 98_765, tool_uses: 8, duration_ms: 12_000 }, + } as unknown as SDKMessage, + state, + ); + expect(notification.some((event) => event.type === "usage.spent")).toBe(false); + + const result = mapClaudeSdkMessage( + { + type: "result", + subtype: "success", + session_id: "claude-session", + usage: { input_tokens: 10, output_tokens: 5 }, + } as unknown as SDKMessage, + state, + ); + expect(result.some((event) => event.type === "usage.spent")).toBe(false); + }); + + it("keeps the sampleId stable per API message (message.id + request_id)", () => { + const state = createClaudeMapperState("thread-1"); + state.usageScope = new ClaudeUsageScopeTracker("claude-session", false); + + const first = spentOf(state, assistantMessage({ id: "msg-1", requestId: "req-1" })); + const replay = spentOf( + state, + assistantMessage({ id: "msg-1", uuid: "uuid-replay", requestId: "req-1" }), + ); + expect(first?.type === "usage.spent" && first.usage.sampleId).toBe("msg-1:req-1"); + expect(replay?.type === "usage.spent" && replay.usage.sampleId).toBe("msg-1:req-1"); + + // Without an SDK request_id the API message id alone is the dedup key. + const bare = spentOf(state, assistantMessage({ id: "msg-2" })); + expect(bare?.type === "usage.spent" && bare.usage.sampleId).toBe("msg-2"); + }); + + it("emits nothing for assistant messages without usage, and marks only the first sample fresh", () => { + const state = createClaudeMapperState("thread-1"); + state.usageScope = new ClaudeUsageScopeTracker("claude-session", true); + + expect( + spentOf( + state, + assistantMessage({ id: "msg-0", usage: { input_tokens: 0, output_tokens: 0 } }), + ), + ).toBeUndefined(); + // The zero-usage message did not consume the fresh baseline. + const first = spentOf(state, assistantMessage({ id: "msg-1" })); + expect(first).toMatchObject({ usage: { fresh: true, epoch: 0 } }); + const second = spentOf(state, assistantMessage({ id: "msg-2" })); + expect(second).toMatchObject({ usage: { epoch: 0 } }); + expect(second?.type === "usage.spent" && second.usage).not.toHaveProperty("fresh"); + }); + + it("bumps the epoch when the SDK session id is adopted mid-thread", () => { + const state = createClaudeMapperState("thread-1"); + state.usageScope = new ClaudeUsageScopeTracker("claude-session", false); + + expect(spentOf(state, assistantMessage({ id: "msg-1" }))).toMatchObject({ + usage: { scopeId: "claude-session", epoch: 0, sampleId: "msg-1" }, + }); + + state.usageScope.adoptScope("claude-session-2"); + + expect(spentOf(state, assistantMessage({ id: "msg-2" }))).toMatchObject({ + usage: { scopeId: "claude-session-2", epoch: 1, sampleId: "msg-2" }, + }); + }); + + it("emits no usage.spent without a usage scope", () => { + const state = createClaudeMapperState("thread-1"); + expect( + mapClaudeSdkMessage(assistantMessage({ id: "msg-1" }), state).some( + (event) => event.type === "usage.spent", + ), + ).toBe(false); + }); +}); diff --git a/src/supervisor/agents/claude/sdkCanonicalMapping.ts b/src/supervisor/agents/claude/sdkCanonicalMapping.ts index 8fc1f0f01..f8d579419 100644 --- a/src/supervisor/agents/claude/sdkCanonicalMapping.ts +++ b/src/supervisor/agents/claude/sdkCanonicalMapping.ts @@ -26,3 +26,7 @@ export { readClaudeApiUsageSpendTokens, } from "./canonicalMapping/result"; export { mapClaudeSdkMessage, readParentToolUseId } from "./canonicalMapping/dispatch"; +export { + ClaudeUsageScopeTracker, + createClaudeUsageSpentEvent, +} from "./canonicalMapping/usageSpent"; diff --git a/src/supervisor/agents/claude/sdkCanonicalMappingState.ts b/src/supervisor/agents/claude/sdkCanonicalMappingState.ts index b4c0c47f0..d9e0a6c51 100644 --- a/src/supervisor/agents/claude/sdkCanonicalMappingState.ts +++ b/src/supervisor/agents/claude/sdkCanonicalMappingState.ts @@ -1,5 +1,6 @@ import type { CanonicalItemType, ToolCallProgress, ToolCallWorkflow } from "@/shared/contracts"; import type { PlanAggregatorState } from "../planAggregator"; +import type { ClaudeUsageScopeTracker } from "./canonicalMapping/usageSpent"; export interface TextItemState { itemId: string; @@ -58,6 +59,12 @@ export interface ToolItemState { export interface ClaudeMapperState { threadId: string; + /** + * Per-call usage scope (SDK session id + epoch) for `usage.spent` emission. + * Owned by the session layer (sdkSession.ts); undefined in tests/terminal + * mode, where no spend events are emitted. + */ + usageScope?: ClaudeUsageScopeTracker; currentTurnId?: string; assistantTextItems: Map; reasoningItems: Map; diff --git a/src/supervisor/agents/claude/sdkSession.ts b/src/supervisor/agents/claude/sdkSession.ts index f3954b496..0b6c0d654 100644 --- a/src/supervisor/agents/claude/sdkSession.ts +++ b/src/supervisor/agents/claude/sdkSession.ts @@ -45,6 +45,7 @@ import { DeferredTurnCompletion } from "./deferredTurnCompletion"; import { applyClaudeContextSuffix } from "./argv"; import { buildClaudeQuestionAnswerEvents, + ClaudeUsageScopeTracker, closeClaudeOpenItems, completeActiveGoalOnTaskDrainEvents, createClaudeMapperState, @@ -232,6 +233,12 @@ export class ClaudeSdkSession implements StructuredSessionHandle { this.currentConfig = config; this.sessionId = sessionRef?.providerSessionId ?? randomUUID(); this.openedResumeSessionId = sessionRef?.providerSessionId; + // Usage scope for per-call spend events: fresh only when the SDK session + // starts brand-new (a resumed session's history is not new spend). + this.mapperState.usageScope = new ClaudeUsageScopeTracker( + this.sessionId, + this.openedResumeSessionId === undefined, + ); this.startQuery(sessionRef?.providerSessionId); await this.requireQuery(); return this.sessionId ?? ""; @@ -857,6 +864,8 @@ export class ClaudeSdkSession implements StructuredSessionHandle { : undefined; if (sessionId && sessionId !== this.sessionId && this.shouldAdoptSessionId(message)) { this.sessionId = sessionId; + // Same conversation under a new provider session id: new usage scope epoch. + this.mapperState.usageScope?.adoptScope(sessionId); this.emitUpdate({ status: this.currentStatus, attention: this.currentAttention, diff --git a/src/supervisor/agents/codex/acp.ts b/src/supervisor/agents/codex/acp.ts index ab8572c21..ce593480b 100644 --- a/src/supervisor/agents/codex/acp.ts +++ b/src/supervisor/agents/codex/acp.ts @@ -29,6 +29,7 @@ import { import { buildCodexAppServerCommand } from "./argv"; import { createCodexMapperState, + CodexUsageScopeTracker, mapCodexNotification, type CodexMapperState, } from "./canonicalMapping"; @@ -489,6 +490,11 @@ export class CodexStructuredSession implements StructuredSessionHandle { const threadOverrides = buildCodexThreadOverrides(config); let threadId: string; + // `fresh` marks a brand-new provider thread (usage ledger baseline 0). A + // resumed thread carries inherited history — its first cumulative sample + // is a baseline that counts nothing. A failed resume falling back to + // `thread/start` DID create a new thread, so it is fresh again. + let createdNewThread = false; if (sessionRef) { this.beginResumeActiveStatusSuppression(sessionRef.providerSessionId); @@ -513,6 +519,7 @@ export class CodexStructuredSession implements StructuredSessionHandle { cause: error, }); } + createdNewThread = true; this.extractRolloutMeta(result); } } else { @@ -521,10 +528,12 @@ export class CodexStructuredSession implements StructuredSessionHandle { if (!threadId) { throw new Error("thread/start response did not contain a thread id."); } + createdNewThread = true; this.extractRolloutMeta(result); } this.remoteThreadId = threadId; + this.ensureMapperState().usageScope = new CodexUsageScopeTracker(threadId, createdNewThread); this.launchOptions = { ...this.launchOptions, resumeThreadId: threadId }; if (sessionRef) { void this.syncRemoteThreadState(threadId, toSessionRef(threadId)); @@ -820,6 +829,11 @@ export class CodexStructuredSession implements StructuredSessionHandle { this.remoteThreadId = newThreadId; this.launchOptions = { ...this.launchOptions, resumeThreadId: newThreadId }; + // The fork created a NEW provider thread carrying inherited history: bump + // the usage scope epoch so the first sample on it is a baseline, and do it + // BEFORE replaying buffered notifications so their samples land in the new + // scope. + this.mapperState?.usageScope?.replaceScope(newThreadId); this.replayForkNotifications(newThreadId); void this.rpc.request("thread/unsubscribe", { threadId }).catch((error) => { console.log("[codex] failed to unsubscribe old thread after fork:", error); @@ -1065,7 +1079,14 @@ export class CodexStructuredSession implements StructuredSessionHandle { const notifications = this.forkNotificationBuffer ?? []; this.forkNotificationBuffer = undefined; for (const notification of notifications) { - if (notification.threadId === threadId) { + // tokenUsage samples are scope-keyed (the mapper resolves scopeId from + // the notification's own scope tracker), so replay them even when they + // were captured mid-fork for another thread (e.g. a child subagent + // thread) — dropping them would punch a gap in that scope's counter. + if ( + notification.threadId === threadId || + notification.method === "thread/tokenUsage/updated" + ) { this.handleNotification(notification.method, notification.params); } } diff --git a/src/supervisor/agents/codex/canonicalMapping.test.ts b/src/supervisor/agents/codex/canonicalMapping.test.ts index 9b0582d74..651ab7146 100644 --- a/src/supervisor/agents/codex/canonicalMapping.test.ts +++ b/src/supervisor/agents/codex/canonicalMapping.test.ts @@ -3,6 +3,7 @@ import { tmpdir } from "node:os"; import { join } from "node:path"; import { describe, expect, it } from "vitest"; import { + CodexUsageScopeTracker, createCodexMapperState, mapCodexNotification, mapCodexServerRequest, @@ -205,6 +206,177 @@ describe("mapCodexNotification — turn lifecycle", () => { }); }); +describe("mapCodexNotification — usage.spent", () => { + function tokenUsageParams(totalTokens: number, lastTotalTokens = 130): Record { + return { + threadId: "codex-thread-1", + turnId: "turn-1", + tokenUsage: { + total: { + inputTokens: totalTokens - 6, + cachedInputTokens: 0, + outputTokens: 6, + reasoningOutputTokens: 0, + totalTokens, + }, + last: { + inputTokens: lastTotalTokens - 6, + cachedInputTokens: 0, + outputTokens: 6, + reasoningOutputTokens: 0, + totalTokens: lastTotalTokens, + }, + modelContextWindow: 258_400, + }, + }; + } + + function spentEventOf( + state: ReturnType, + params: Record, + ) { + return mapCodexNotification("thread/tokenUsage/updated", params, state).find( + (event) => event.type === "usage.spent", + ); + } + + it("emits a cumulative usage.spent sample from the total (never last) alongside context.updated", () => { + const state = createCodexMapperState("t-codex"); + state.usageScope = new CodexUsageScopeTracker("codex-thread-1", true); + + const events = mapCodexNotification( + "thread/tokenUsage/updated", + tokenUsageParams(11_839), + state, + ); + + expect(events).toEqual([ + { + type: "context.updated", + threadId: "t-codex", + usage: { + usedTokens: 130, + maxTokens: 258_400, + breakdown: [ + { id: "input", label: "Input", tokens: 124 }, + { id: "output", label: "Output", tokens: 6 }, + ], + }, + }, + { + type: "usage.spent", + threadId: "t-codex", + usage: { + counterKind: "cumulative", + counter: 11_839, + scopeId: "codex-thread-1", + epoch: 0, + fresh: true, + sampleId: "codex-thread-1:0:11839", + occurredAt: expect.any(Number), + }, + }, + ]); + }); + + it("reads the cumulative total from legacy snake_case payloads", () => { + const state = createCodexMapperState("t-codex"); + state.usageScope = new CodexUsageScopeTracker("codex-thread-1", false); + + const spent = spentEventOf(state, { + threadId: "codex-thread-1", + token_usage: { + total_token_usage: { total_tokens: 5_000, input_tokens: 4_000, output_tokens: 1_000 }, + last_token_usage: { total_tokens: 42 }, + model_context_window: 200_000, + }, + }); + + expect(spent).toEqual({ + type: "usage.spent", + threadId: "t-codex", + usage: { + counterKind: "cumulative", + counter: 5_000, + scopeId: "codex-thread-1", + epoch: 0, + sampleId: "codex-thread-1:0:5000", + occurredAt: expect.any(Number), + }, + }); + }); + + it("marks only the first sample of a brand-new thread fresh", () => { + const state = createCodexMapperState("t-codex"); + state.usageScope = new CodexUsageScopeTracker("codex-thread-1", true); + + expect(spentEventOf(state, tokenUsageParams(1_000))).toMatchObject({ + usage: { fresh: true, epoch: 0 }, + }); + const second = spentEventOf(state, tokenUsageParams(1_500)); + expect(second).toMatchObject({ usage: { counter: 1_500, epoch: 0 } }); + expect(second?.type === "usage.spent" && second.usage).not.toHaveProperty("fresh"); + }); + + it("emits no usage.spent without a usage scope (terminal mode)", () => { + const state = createCodexMapperState("t-codex"); + + const events = mapCodexNotification( + "thread/tokenUsage/updated", + tokenUsageParams(11_839), + state, + ); + + expect(events.every((event) => event.type === "context.updated")).toBe(true); + }); + + it("bumps the epoch when the cumulative counter resets", () => { + const state = createCodexMapperState("t-codex"); + state.usageScope = new CodexUsageScopeTracker("codex-thread-1", false); + + expect(spentEventOf(state, tokenUsageParams(1_000))).toMatchObject({ + usage: { counter: 1_000, epoch: 0, sampleId: "codex-thread-1:0:1000" }, + }); + expect(spentEventOf(state, tokenUsageParams(1_500))).toMatchObject({ + usage: { counter: 1_500, epoch: 0, sampleId: "codex-thread-1:0:1500" }, + }); + // Upstream reset (e.g. ContextWindowExceeded): a lower total starts a new epoch. + expect(spentEventOf(state, tokenUsageParams(300))).toMatchObject({ + usage: { counter: 300, epoch: 1, sampleId: "codex-thread-1:1:300" }, + }); + expect(spentEventOf(state, tokenUsageParams(400))).toMatchObject({ + usage: { counter: 400, epoch: 1, sampleId: "codex-thread-1:1:400" }, + }); + }); + + it("starts a new epoch on fork (replaceScope) with a baseline, non-fresh sample", () => { + const state = createCodexMapperState("t-codex"); + state.usageScope = new CodexUsageScopeTracker("codex-thread-1", true); + + expect(spentEventOf(state, tokenUsageParams(1_000))).toMatchObject({ + usage: { scopeId: "codex-thread-1", epoch: 0, fresh: true }, + }); + + state.usageScope.replaceScope("codex-thread-2"); + + // The forked thread's first sample carries inherited history: new epoch, + // baseline-only (no fresh), and the lower counter is NOT read as a reset. + const forked = spentEventOf(state, { + ...tokenUsageParams(500), + threadId: "codex-thread-2", + }); + expect(forked).toMatchObject({ + usage: { + counter: 500, + scopeId: "codex-thread-2", + epoch: 1, + sampleId: "codex-thread-2:1:500", + }, + }); + expect(forked?.type === "usage.spent" && forked.usage).not.toHaveProperty("fresh"); + }); +}); + describe("mapCodexNotification — goals", () => { it("maps Codex thread goal updates to a shared goal chat item", () => { const state = createCodexMapperState("t-codex"); diff --git a/src/supervisor/agents/codex/canonicalMapping.ts b/src/supervisor/agents/codex/canonicalMapping.ts index 958ada646..1cf6c909a 100644 --- a/src/supervisor/agents/codex/canonicalMapping.ts +++ b/src/supervisor/agents/codex/canonicalMapping.ts @@ -27,3 +27,5 @@ export { mapCodexServerRequest, translateCodexCanonicalResponse, } from "./canonicalMapping/serverRequest"; +export { createCodexUsageSpentEvent } from "./canonicalMapping/usage"; +export { CodexUsageScopeTracker } from "./canonicalMapping/usageScope"; diff --git a/src/supervisor/agents/codex/canonicalMapping/dispatch.ts b/src/supervisor/agents/codex/canonicalMapping/dispatch.ts index f66382d1d..ec9bf94a5 100644 --- a/src/supervisor/agents/codex/canonicalMapping/dispatch.ts +++ b/src/supervisor/agents/codex/canonicalMapping/dispatch.ts @@ -32,7 +32,12 @@ import { } from "./readers"; import { readStringField } from "../../fileChangeSummary"; import { extractCodexFileChangePath, readCommandAggregatedOutput } from "./toolExtraction"; -import { createCodexContextUsageEvent, createCodexTokenUsageEvent } from "./usage"; +import { + createCodexContextUsageEvent, + createCodexTokenUsageEvent, + createCodexUsageSpentEvent, + readCodexCumulativeTotalTokens, +} from "./usage"; export function mapCodexNotification( method: string, @@ -44,7 +49,18 @@ export function mapCodexNotification( if (method === "thread/tokenUsage/updated") { const usageEvent = createCodexTokenUsageEvent(threadId, params); - return usageEvent ? [usageEvent] : []; + const events: RuntimeEvent[] = usageEvent ? [usageEvent] : []; + // Ledger spend sample alongside the dock's context.updated: the dock keeps + // the per-call `last` occupancy, the ledger consumes the cumulative total. + const scope = state.usageScope; + if (scope) { + const counter = readCodexCumulativeTotalTokens(params); + if (counter !== undefined) { + const spentEvent = createCodexUsageSpentEvent(threadId, params, scope.sample(counter)); + if (spentEvent) events.push(spentEvent); + } + } + return events; } if (method === "turn/started") { diff --git a/src/supervisor/agents/codex/canonicalMapping/usage.ts b/src/supervisor/agents/codex/canonicalMapping/usage.ts index b39f90ff5..ee420a9dd 100644 --- a/src/supervisor/agents/codex/canonicalMapping/usage.ts +++ b/src/supervisor/agents/codex/canonicalMapping/usage.ts @@ -23,6 +23,59 @@ export function createCodexContextUsageEvent( return createCodexUsageEvent(threadId, usage as Record); } +/** + * Session-cumulative spend sample from `thread/tokenUsage/updated`: the + * absolute `total.totalTokens` counter (NEVER `last` — that is per-call and + * gets rewritten by upstream compaction). The session layer supplies the + * scope meta (scopeId/epoch/fresh) via `CodexUsageScopeTracker.sample`. + */ +export function createCodexUsageSpentEvent( + threadId: string, + params: Record | undefined, + meta: { scopeId: string; epoch: number; fresh?: boolean }, +): RuntimeEvent | undefined { + const counter = readCodexCumulativeTotalTokens(params); + if (counter === undefined) return undefined; + return { + type: "usage.spent", + threadId, + usage: { + counterKind: "cumulative", + counter, + scopeId: meta.scopeId, + epoch: meta.epoch, + ...(meta.fresh ? { fresh: true } : {}), + // Idempotent replays of the same counter value dedup naturally. + sampleId: `${meta.scopeId}:${meta.epoch}:${counter}`, + occurredAt: Date.now(), + }, + }; +} + +/** Read the session-cumulative total across the current and legacy payload shapes. */ +export function readCodexCumulativeTotalTokens( + params: Record | undefined, +): number | undefined { + const currentTokenUsage = readRecord(params?.tokenUsage); + const fromCurrent = readTotalTokens(readRecord(currentTokenUsage?.total)); + if (fromCurrent !== undefined) return fromCurrent; + + const legacy = readRecord(params?.token_usage) ?? currentTokenUsage; + const fromLegacy = readTotalTokens( + readRecord(legacy?.total) ?? + readRecord(legacy?.totalTokenUsage) ?? + readRecord(legacy?.total_token_usage), + ); + if (fromLegacy !== undefined) return fromLegacy; + + const info = readRecord(params?.info); + return readTotalTokens(readRecord(info?.total_token_usage) ?? readRecord(info?.totalTokenUsage)); +} + +function readTotalTokens(total: Record | undefined): number | undefined { + return readNonNegativeInteger(total?.totalTokens) ?? readNonNegativeInteger(total?.total_tokens); +} + export function createCodexTokenUsageEvent( threadId: string, params: Record | undefined, diff --git a/src/supervisor/agents/codex/canonicalMapping/usageScope.ts b/src/supervisor/agents/codex/canonicalMapping/usageScope.ts new file mode 100644 index 000000000..917ce09f5 --- /dev/null +++ b/src/supervisor/agents/codex/canonicalMapping/usageScope.ts @@ -0,0 +1,56 @@ +/** + * Cumulative token-usage scope tracking for one Codex session. + * + * The app-server's `total.totalTokens` counter is an absolute, monotonically + * increasing total per provider (Codex) thread. The main-process usage ledger + * counts counter increases within a `(provider, scopeId, epoch)` key, so this + * tracker draws those boundaries for one mapper: + * + * - `scopeId` is the CODEX thread id (not the Poracode thread id). A fork or + * replacement provider thread switches the scope via {@link replaceScope}, + * starting a new epoch whose first sample is a baseline (forked threads + * carry inherited history — that is not new spend). + * - An incoming counter LOWER than the last emitted one means upstream reset + * the total (e.g. ContextWindowExceeded recovery); the adapter bumps the + * epoch here because the ledger itself stays heuristic-free. + * - `fresh: true` is emitted only on the first sample of a brand-new provider + * thread (created via `thread/start`, not resumed or forked), telling the + * ledger the baseline is 0 so the first sample counts in full. + */ +export class CodexUsageScopeTracker { + private epoch = 0; + private lastCounter: number | undefined; + private freshPending: boolean; + + constructor( + private scopeId: string, + fresh: boolean, + ) { + this.freshPending = fresh; + } + + /** Switch to a replacement provider thread (fork/new session): new epoch, baseline sample. */ + replaceScope(scopeId: string): void { + this.scopeId = scopeId; + this.epoch += 1; + this.lastCounter = undefined; + this.freshPending = false; + } + + /** Meta for one incoming cumulative sample; a counter decrease signals an upstream reset. */ + sample(counter: number): { scopeId: string; epoch: number; fresh?: boolean } { + if (this.lastCounter !== undefined && counter < this.lastCounter) { + this.epoch += 1; + this.lastCounter = undefined; + this.freshPending = false; + } + this.lastCounter = counter; + const fresh = this.freshPending; + this.freshPending = false; + return { + scopeId: this.scopeId, + epoch: this.epoch, + ...(fresh ? { fresh: true } : {}), + }; + } +} diff --git a/src/supervisor/agents/codex/canonicalMappingState.ts b/src/supervisor/agents/codex/canonicalMappingState.ts index d0d750316..0c7cd5c44 100644 --- a/src/supervisor/agents/codex/canonicalMappingState.ts +++ b/src/supervisor/agents/codex/canonicalMappingState.ts @@ -1,7 +1,14 @@ import type { CanonicalItemType } from "@/shared/contracts"; +import type { CodexUsageScopeTracker } from "./canonicalMapping/usageScope"; export interface CodexMapperState { threadId: string; + /** + * Cumulative usage scope (codex thread id + epoch) for `usage.spent` + * emission. Owned by the session layer (acp.ts / subAgentRouting.ts); + * undefined in terminal mode, where no spend events are emitted. + */ + usageScope?: CodexUsageScopeTracker; /** Most recent turn id reported via `turn.started`. */ currentTurnId?: string; /** Open assistant_message item id, if any (closed on `turn/completed`). */ diff --git a/src/supervisor/agents/codex/codex.test.ts b/src/supervisor/agents/codex/codex.test.ts index a038fa5b5..0618e48e9 100644 --- a/src/supervisor/agents/codex/codex.test.ts +++ b/src/supervisor/agents/codex/codex.test.ts @@ -15,6 +15,7 @@ import { } from "./detection"; import { CodexStructuredSession } from "./acp"; import { CodexRpcResponseError, type CodexAppServerRpcListener } from "./appServerRpc"; +import { createCodexMapperState, CodexUsageScopeTracker } from "./canonicalMapping"; import type { CodexThreadStatus } from "./acpProtocol"; import type { OscNotification, OscTitle } from "@/shared/osc"; import type { RuntimeEvent, ToolCallPayload } from "@/shared/contracts"; @@ -794,6 +795,81 @@ describe("CodexSubAgentRouter", () => { }), ); }); + + it("routes child tokenUsage as usage.spent in the child scope, dropping context.updated", () => { + const router = createRouterWithChild(); + + const events = router.routeChildNotification( + "thread/tokenUsage/updated", + { + threadId: "child-thread", + tokenUsage: { + total: { + inputTokens: 900, + cachedInputTokens: 0, + outputTokens: 100, + reasoningOutputTokens: 0, + totalTokens: 1_000, + }, + last: { + inputTokens: 40, + cachedInputTokens: 0, + outputTokens: 10, + reasoningOutputTokens: 0, + totalTokens: 50, + }, + modelContextWindow: 258_400, + }, + }, + "provider-thread", + ); + + expect(events?.filter((event) => event.type === "usage.spent")).toEqual([ + { + type: "usage.spent", + threadId: "local-thread", + usage: { + counterKind: "cumulative", + counter: 1_000, + scopeId: "child-thread", + epoch: 0, + fresh: true, + sampleId: "child-thread:0:1000", + occurredAt: expect.any(Number), + }, + }, + ]); + // Child context.updated stays dropped — the dock keeps the main thread's occupancy. + expect(events?.some((event) => event.type === "context.updated")).toBe(false); + + const next = router.routeChildNotification( + "thread/tokenUsage/updated", + { + threadId: "child-thread", + tokenUsage: { + total: { + inputTokens: 1_400, + cachedInputTokens: 0, + outputTokens: 100, + reasoningOutputTokens: 0, + totalTokens: 1_500, + }, + last: { + inputTokens: 40, + cachedInputTokens: 0, + outputTokens: 10, + reasoningOutputTokens: 0, + totalTokens: 50, + }, + modelContextWindow: 258_400, + }, + }, + "provider-thread", + ); + expect(next?.find((event) => event.type === "usage.spent")).toMatchObject({ + usage: { counter: 1_500, scopeId: "child-thread", epoch: 0 }, + }); + }); }); function createRouterWithChild(): CodexSubAgentRouter { @@ -1022,6 +1098,79 @@ describe("CodexStructuredSession", () => { }); }); + it("starts a new usage scope epoch on fork and replays buffered tokenUsage into it", async () => { + const requests: Array<{ method: string; params: Record }> = []; + const structuredSession = makeStructuredSession(requests); + const runtimeEvents: RuntimeEvent[] = []; + (structuredSession as unknown as Record)["listener"] = { + onRuntimeEvent: (event: RuntimeEvent) => runtimeEvents.push(event), + }; + const mapperState = createCodexMapperState("local-thread"); + mapperState.usageScope = new CodexUsageScopeTracker("provider-thread", false); + (structuredSession as unknown as Record)["mapperState"] = mapperState; + (structuredSession as unknown as Record)["rpc"] = { + request: (method: string, params: Record) => { + requests.push({ method, params }); + if (method === "thread/read" && params.includeTurns === true) { + return Promise.resolve({ + thread: { + turns: ["turn-1", "turn-2", "turn-3"].map((id) => ({ id })), + }, + }); + } + if (method === "thread/fork") { + const response = Promise.resolve({ thread: { id: "forked-thread" } }); + // Arrives mid-fork (buffered), carrying the forked thread's inherited history. + dispatchNotification(structuredSession, { + jsonrpc: "2.0", + method: "thread/tokenUsage/updated", + params: { + threadId: "forked-thread", + turnId: "turn-2", + tokenUsage: { + total: { + inputTokens: 4_800, + cachedInputTokens: 0, + outputTokens: 200, + reasoningOutputTokens: 0, + totalTokens: 5_000, + }, + last: { + inputTokens: 80, + cachedInputTokens: 0, + outputTokens: 15, + reasoningOutputTokens: 5, + totalTokens: 100, + }, + modelContextWindow: 258_400, + }, + }, + }); + return response; + } + return Promise.resolve({ thread: { status: { type: "idle" } } }); + }, + }; + + const history = await structuredSession.rollbackThread(1); + + expect(history).toEqual({ providerSessionId: "forked-thread", messages: [] }); + // The buffered tokenUsage replayed into the new scope: epoch 1, baseline + // sample (no fresh — forked threads carry inherited history). + expect(runtimeEvents).toContainEqual({ + type: "usage.spent", + threadId: "local-thread", + usage: { + counterKind: "cumulative", + counter: 5_000, + scopeId: "forked-thread", + epoch: 1, + sampleId: "forked-thread:1:5000", + occurredAt: expect.any(Number), + }, + }); + }); + it("falls back to thread/rollback when thread/fork is unavailable", async () => { const requests: Array<{ method: string; params: Record }> = []; const structuredSession = makeStructuredSession(requests); diff --git a/src/supervisor/agents/codex/subAgentRouting.ts b/src/supervisor/agents/codex/subAgentRouting.ts index 7c1c9e72e..c8311d3b8 100644 --- a/src/supervisor/agents/codex/subAgentRouting.ts +++ b/src/supervisor/agents/codex/subAgentRouting.ts @@ -1,6 +1,7 @@ import type { RuntimeEvent, ToolCallPayload, ToolCallProgress } from "@/shared/contracts"; import { createCodexMapperState, + CodexUsageScopeTracker, mapCodexNotification, type CodexMapperState, } from "./canonicalMapping"; @@ -157,7 +158,11 @@ export class CodexSubAgentRouter { event.type === "item.started" || event.type === "item.updated" || event.type === "item.completed" || - event.type === "content.delta", + event.type === "content.delta" || + // Child tokenUsage flows to the usage ledger as usage.spent (scoped to + // the child codex thread); child context.updated stays dropped so the + // dock keeps showing the main thread's occupancy. + event.type === "usage.spent", ); const routed: RuntimeEvent[] = []; let progressChanged = false; @@ -342,9 +347,13 @@ export class CodexSubAgentRouter { existing.active = active; return []; } + const mapperState = createCodexMapperState(this.localThreadId); + // Child threads are brand-new provider threads (counter starts at 0), so + // their usage scope is fresh with its own epoch tracking. + mapperState.usageScope = new CodexUsageScopeTracker(threadId, true); this.children.set(threadId, { parentItemId, - mapperState: createCodexMapperState(this.localThreadId), + mapperState, itemIds: new Set(), active, failed: false, diff --git a/src/supervisor/agents/opencode/canonicalMapping/dispatch.ts b/src/supervisor/agents/opencode/canonicalMapping/dispatch.ts index 8b3262dcf..5dd277daa 100644 --- a/src/supervisor/agents/opencode/canonicalMapping/dispatch.ts +++ b/src/supervisor/agents/opencode/canonicalMapping/dispatch.ts @@ -13,7 +13,11 @@ import type { EventSubscribeResponse, Part } from "@opencode-ai/sdk/v2"; import type { RuntimeEvent } from "@/shared/contracts"; import { newItemId } from "../../contextUsage"; -import type { OpenCodeMapperState } from "../sdkCanonicalMappingState"; +import { + markOpenCodeUsageScopeSampled, + openCodeUsageScopeForSession, + type OpenCodeMapperState, +} from "../sdkCanonicalMappingState"; import { normalizeToolName } from "./readers"; import { classifyPermissionRequestType, @@ -35,7 +39,7 @@ import { } from "./textItems"; import { classifyToolItemType } from "./toolClassification"; import { toolPayload } from "./toolPayload"; -import { createOpenCodeContextUsageEvent } from "./usage"; +import { createOpenCodeContextUsageEvent, createOpenCodeUsageSpentEvent } from "./usage"; function handlePart(state: OpenCodeMapperState, part: Part, events: RuntimeEvent[]): void { if (part.type === "text") { @@ -206,6 +210,19 @@ function mapCanonicalEvent( // the reasoning Part snapshot didn't carry `time.end` before the message // wrapped up. if (info.role === "assistant" && info.time?.completed) { + // Token spend is final only on the completed snapshot (earlier + // message.updated snapshots still evolve), and emitted exactly once + // per message id — the ledger dedups per-call samples by sampleId. + // Child (subagent) sessions scope to their own session id. + if (!state.usageSpentMessages.has(info.id)) { + const scope = openCodeUsageScopeForSession(state, info.sessionID); + const spentEvent = createOpenCodeUsageSpentEvent(state.threadId, info, scope); + if (spentEvent) { + state.usageSpentMessages.add(info.id); + markOpenCodeUsageScopeSampled(state, scope.scopeId); + events.push(spentEvent); + } + } const itemId = state.assistantItems.get(info.id); if (itemId) { events.push({ @@ -354,9 +371,10 @@ export function mapOpenCodeEvent( // child-session message/part IDs are independent UUIDs from OpenCode, so // they don't collide with parent items in the mapper's shared state maps. tagChildEventsWithParent(canonicalEvents, child.itemId); - // Suppress context.updated events from child sessions — token accounting - // on the thread tracks the main session only; child sessions have their - // own budgets that don't roll up into the parent's display. + // Suppress context.updated events from child sessions — the context dock + // tracks the main session only; child sessions have their own budgets + // that don't roll up into the parent's display. usage.spent still flows + // through: the ledger keys it by the child's own scope id. for (const ev of canonicalEvents) { if (ev.type === "context.updated") continue; events.push(ev); diff --git a/src/supervisor/agents/opencode/canonicalMapping/usage.ts b/src/supervisor/agents/opencode/canonicalMapping/usage.ts index 2be4ae634..fb2244d8e 100644 --- a/src/supervisor/agents/opencode/canonicalMapping/usage.ts +++ b/src/supervisor/agents/opencode/canonicalMapping/usage.ts @@ -30,3 +30,63 @@ export function createOpenCodeContextUsageEvent( }), ); } + +function sumTokenBuckets(buckets: ReadonlyArray): number | undefined { + let total = 0; + let seen = false; + for (const bucket of buckets) { + if (bucket === undefined) continue; + total += bucket; + seen = true; + } + return seen ? total : undefined; +} + +/** + * Per-call `usage.spent` from an assistant message's FINAL completed snapshot + * (message snapshots evolve while `message.updated` repeats — the caller gates + * on `time.completed` and dedups by message id). The SDK's `tokens.total` is + * authoritative; the bucket sum is the fallback for older servers. `sampleId` + * is the provider message id, giving the ledger exact-once dedup across + * replayed snapshots. + */ +export function createOpenCodeUsageSpentEvent( + threadId: string, + info: unknown, + scope: { scopeId: string; epoch: number; fresh?: boolean }, +): RuntimeEvent | undefined { + if (!info || typeof info !== "object") return undefined; + const obj = info as Record; + if (obj.role !== "assistant") return undefined; + const id = typeof obj.id === "string" && obj.id.length > 0 ? obj.id : undefined; + if (!id) return undefined; + const tokens = obj.tokens; + if (!tokens || typeof tokens !== "object") return undefined; + const tok = tokens as Record; + const cache = + tok.cache && typeof tok.cache === "object" ? (tok.cache as Record) : {}; + const total = + readNonNegativeInteger(tok.total) ?? + sumTokenBuckets([ + readNonNegativeInteger(tok.input), + readNonNegativeInteger(tok.output), + readNonNegativeInteger(tok.reasoning), + readNonNegativeInteger(cache.read), + readNonNegativeInteger(cache.write), + ]); + if (total === undefined || total <= 0) return undefined; + const model = typeof obj.modelID === "string" && obj.modelID.length > 0 ? obj.modelID : undefined; + return { + type: "usage.spent", + threadId, + usage: { + counterKind: "per-call", + counter: total, + scopeId: scope.scopeId, + epoch: scope.epoch, + ...(scope.fresh ? { fresh: true } : {}), + sampleId: id, + ...(model ? { model } : {}), + }, + }; +} diff --git a/src/supervisor/agents/opencode/sdkCanonicalMapping.test.ts b/src/supervisor/agents/opencode/sdkCanonicalMapping.test.ts index d7d25d38e..9790e3dae 100644 --- a/src/supervisor/agents/opencode/sdkCanonicalMapping.test.ts +++ b/src/supervisor/agents/opencode/sdkCanonicalMapping.test.ts @@ -1388,3 +1388,246 @@ describe("sdkCanonicalMapping — closeOpenItems", () => { expect(closed[0]).toMatchObject({ type: "item.completed", threadId: "thread-1" }); }); }); + +describe("sdkCanonicalMapping — usage.spent", () => { + it("emits usage.spent only on the final completed snapshot, preferring tokens.total", () => { + const state = createOpenCodeMapperState("thread-1"); + setOpenCodeMainSessionId(state, "ses_test", { fresh: true }); + + // Evolving (not yet completed) snapshot: context.updated for the dock, but + // no spend sample yet. + const evolving = mapOpenCodeEvent( + messageUpdatedEvent( + assistantMessage("msg_1", { + tokens: { total: 100, input: 90, output: 10, reasoning: 0, cache: { read: 0, write: 0 } }, + }), + ), + state, + ); + expect(evolving.find((e) => e.type === "usage.spent")).toBeUndefined(); + expect(evolving.find((e) => e.type === "context.updated")).toBeDefined(); + + // Final completed snapshot: exactly one per-call sample keyed by message id. + const completed = mapOpenCodeEvent( + messageUpdatedEvent( + assistantMessage("msg_1", { + time: { created: 0, completed: 5 }, + tokens: { total: 150, input: 90, output: 60, reasoning: 0, cache: { read: 0, write: 0 } }, + }), + ), + state, + ); + const spent = completed.find((e) => e.type === "usage.spent"); + expect(spent).toMatchObject({ + type: "usage.spent", + threadId: "thread-1", + usage: { + counterKind: "per-call", + counter: 150, + scopeId: "ses_test", + epoch: 0, + fresh: true, + sampleId: "msg_1", + model: "big-pickle", + }, + }); + + // A repeated completed snapshot of the same message does not re-emit. + const replay = mapOpenCodeEvent( + messageUpdatedEvent( + assistantMessage("msg_1", { + time: { created: 0, completed: 5 }, + tokens: { total: 150, input: 90, output: 60, reasoning: 0, cache: { read: 0, write: 0 } }, + }), + ), + state, + ); + expect(replay.find((e) => e.type === "usage.spent")).toBeUndefined(); + }); + + it("falls back to summing token buckets when tokens.total is missing", () => { + const state = createOpenCodeMapperState("thread-1"); + setOpenCodeMainSessionId(state, "ses_test"); + + const events = mapOpenCodeEvent( + messageUpdatedEvent( + assistantMessage("msg_1", { + time: { created: 0, completed: 5 }, + tokens: { input: 10, output: 5, reasoning: 2, cache: { read: 3, write: 1 } }, + }), + ), + state, + ); + const spent = events.find((e) => e.type === "usage.spent"); + expect(spent).toMatchObject({ + usage: { counterKind: "per-call", counter: 21, sampleId: "msg_1" }, + }); + // fresh was not requested for this scope. + expect(spent && "usage" in spent ? spent.usage : undefined).not.toMatchObject({ + fresh: true, + }); + }); + + it("reports fresh only on the scope's first sample", () => { + const state = createOpenCodeMapperState("thread-1"); + setOpenCodeMainSessionId(state, "ses_test", { fresh: true }); + + const tokens = { + total: 10, + input: 5, + output: 5, + reasoning: 0, + cache: { read: 0, write: 0 }, + }; + const first = mapOpenCodeEvent( + messageUpdatedEvent( + assistantMessage("msg_1", { time: { created: 0, completed: 5 }, tokens }), + ), + state, + ); + const second = mapOpenCodeEvent( + messageUpdatedEvent( + assistantMessage("msg_2", { time: { created: 0, completed: 6 }, tokens }), + ), + state, + ); + expect(first.find((e) => e.type === "usage.spent")).toMatchObject({ + usage: { fresh: true }, + }); + const secondSpent = second.find((e) => e.type === "usage.spent"); + expect(secondSpent && "usage" in secondSpent ? secondSpent.usage : undefined).toMatchObject({ + scopeId: "ses_test", + epoch: 0, + }); + expect(secondSpent && "usage" in secondSpent ? secondSpent.usage : undefined).not.toMatchObject( + { fresh: true }, + ); + }); + + it("bumps the epoch when the main session id changes", () => { + const state = createOpenCodeMapperState("thread-1"); + setOpenCodeMainSessionId(state, "ses_a", { fresh: true }); + const tokens = { + total: 10, + input: 5, + output: 5, + reasoning: 0, + cache: { read: 0, write: 0 }, + }; + mapOpenCodeEvent( + messageUpdatedEvent( + assistantMessage("msg_1", { + sessionID: "ses_a", + time: { created: 0, completed: 5 }, + tokens, + }), + ), + state, + ); + + // Re-opened against a different provider session: new scope lineage, no + // fresh flag (not created new by this handle). + setOpenCodeMainSessionId(state, "ses_b"); + const events = mapOpenCodeEvent( + messageUpdatedEvent( + assistantMessage("msg_2", { + sessionID: "ses_b", + time: { created: 0, completed: 6 }, + tokens, + }), + ), + state, + ); + const spent = events.find((e) => e.type === "usage.spent"); + expect(spent).toMatchObject({ + usage: { scopeId: "ses_b", epoch: 1, sampleId: "msg_2" }, + }); + expect(spent && "usage" in spent ? spent.usage : undefined).not.toMatchObject({ fresh: true }); + }); + + it("emits usage.spent for child-session messages while still suppressing child context.updated", () => { + const state = createOpenCodeMapperState("thread-1"); + setOpenCodeMainSessionId(state, "ses_main", { fresh: true }); + + // Parent task tool starts, then its child session is announced. + mapOpenCodeEvent( + { + id: "evt-1", + type: "message.part.updated", + properties: { + sessionID: "ses_main", + time: 0, + part: { + id: "prt_task", + sessionID: "ses_main", + messageID: "msg_assistant", + type: "tool", + tool: "task", + callID: "call_task", + state: { + status: "running", + input: { description: "Explore code" }, + time: { start: 0 }, + }, + } as ToolPart, + }, + }, + state, + ); + const childSession: Session = { + id: "ses_child", + slug: "child", + projectID: "proj", + directory: "/", + parentID: "ses_main", + title: "subagent", + version: "1.0.0", + time: { created: 1, updated: 1 }, + }; + mapOpenCodeEvent( + { + id: "evt-2", + type: "session.created", + properties: { sessionID: "ses_child", info: childSession }, + }, + state, + ); + + // A completed assistant message in the child session. + const events = mapOpenCodeEvent( + { + id: "evt-3", + type: "message.updated", + properties: { + sessionID: "ses_child", + info: assistantMessage("msg_child_1", { + sessionID: "ses_child", + time: { created: 0, completed: 5 }, + tokens: { + total: 42, + input: 40, + output: 2, + reasoning: 0, + cache: { read: 0, write: 0 }, + }, + }), + }, + }, + state, + ); + + expect(events.find((e) => e.type === "context.updated")).toBeUndefined(); + expect(events.find((e) => e.type === "usage.spent")).toMatchObject({ + type: "usage.spent", + threadId: "thread-1", + usage: { + counterKind: "per-call", + counter: 42, + scopeId: "ses_child", + epoch: 0, + fresh: true, + sampleId: "msg_child_1", + }, + }); + }); +}); diff --git a/src/supervisor/agents/opencode/sdkCanonicalMapping.ts b/src/supervisor/agents/opencode/sdkCanonicalMapping.ts index e95bef72f..bb37ae3cc 100644 --- a/src/supervisor/agents/opencode/sdkCanonicalMapping.ts +++ b/src/supervisor/agents/opencode/sdkCanonicalMapping.ts @@ -17,6 +17,8 @@ export { createOpenCodeMapperState, isOpenCodeChildSession, + markOpenCodeUsageScopeSampled, + openCodeUsageScopeForSession, setOpenCodeMainSessionId, type OpenCodeMapperState, type OpenCodeSubAgentSessionState, diff --git a/src/supervisor/agents/opencode/sdkCanonicalMappingState.ts b/src/supervisor/agents/opencode/sdkCanonicalMappingState.ts index 1d9ccd83e..139f3fc03 100644 --- a/src/supervisor/agents/opencode/sdkCanonicalMappingState.ts +++ b/src/supervisor/agents/opencode/sdkCanonicalMappingState.ts @@ -77,6 +77,19 @@ export interface OpenCodeMapperState { unclaimedChildSessions: string[]; /** Map child session id → live progress state. */ subAgentSessions: Map; + /** + * Provider session id the `usage.spent` ledger scope currently maps to. + * Mirrors `mainSessionId` but tracks epoch/fresh for the token ledger. + */ + usageScopeId: string | null; + /** Epoch of the main usage scope — bumped when the provider session id changes. */ + usageEpoch: number; + /** True when the current main scope's session was created new (baseline 0). */ + usageScopeFresh: boolean; + /** Scope ids that have already emitted a `usage.spent` sample. */ + usageSampledScopes: Set; + /** Assistant message ids that already emitted `usage.spent` (exact-once). */ + usageSpentMessages: Set; } export function createOpenCodeMapperState(threadId: string): OpenCodeMapperState { @@ -95,16 +108,58 @@ export function createOpenCodeMapperState(threadId: string): OpenCodeMapperState taskToolsAwaitingChild: [], unclaimedChildSessions: [], subAgentSessions: new Map(), + usageScopeId: null, + usageEpoch: 0, + usageScopeFresh: false, + usageSampledScopes: new Set(), + usageSpentMessages: new Set(), }; } /** * Record the main session id once `openThread` has resolved. The mapper uses * this to recognise sub-sessions (`Session.parentID === mainSessionId`) when - * the `task` tool spawns a subagent. + * the `task` tool spawns a subagent. Also establishes the `usage.spent` ledger + * scope: a changed session id ends the old counter lineage, so the epoch bumps + * rather than inferring a reset from counter values. `fresh` marks a session + * the runtime just created (vs resumed) — its first sample counts in full. */ -export function setOpenCodeMainSessionId(state: OpenCodeMapperState, sessionId: string): void { +export function setOpenCodeMainSessionId( + state: OpenCodeMapperState, + sessionId: string, + options?: { fresh?: boolean }, +): void { state.mainSessionId = sessionId; + if (state.usageScopeId !== sessionId) { + if (state.usageScopeId !== null) state.usageEpoch += 1; + state.usageScopeId = sessionId; + state.usageScopeFresh = options?.fresh === true; + } +} + +/** + * Resolve the `usage.spent` scope for a message's provider session. The main + * session carries the tracked epoch/fresh; child (subagent) sessions are + * independent scopes — always created fresh per `task` run, epoch 0. + * `fresh` is reported only for a scope's first emitted sample. + */ +export function openCodeUsageScopeForSession( + state: OpenCodeMapperState, + sessionId: string, +): { scopeId: string; epoch: number; fresh: boolean } { + if (sessionId === state.usageScopeId) { + return { + scopeId: sessionId, + epoch: state.usageEpoch, + fresh: state.usageScopeFresh && !state.usageSampledScopes.has(sessionId), + }; + } + return { scopeId: sessionId, epoch: 0, fresh: !state.usageSampledScopes.has(sessionId) }; +} + +/** Mark that a `usage.spent` sample was emitted for the given scope. */ +export function markOpenCodeUsageScopeSampled(state: OpenCodeMapperState, scopeId: string): void { + state.usageSampledScopes.add(scopeId); } /** diff --git a/src/supervisor/agents/opencode/sdkSession.ts b/src/supervisor/agents/opencode/sdkSession.ts index 8fcfc9096..262f9a77f 100644 --- a/src/supervisor/agents/opencode/sdkSession.ts +++ b/src/supervisor/agents/opencode/sdkSession.ts @@ -279,7 +279,7 @@ export class OpencodeSdkSession implements StructuredSessionHandle { this.crossagentDeniedPermissionRules().length === 0 ? this.permissionSyncKey(config) : undefined; - if (this.mapperState) setOpenCodeMainSessionId(this.mapperState, id); + if (this.mapperState) setOpenCodeMainSessionId(this.mapperState, id, { fresh: true }); await this.refreshSlashCommands(); return id; } diff --git a/src/supervisor/agents/pi/rpcSession.test.ts b/src/supervisor/agents/pi/rpcSession.test.ts index fcf4adece..aee58be37 100644 --- a/src/supervisor/agents/pi/rpcSession.test.ts +++ b/src/supervisor/agents/pi/rpcSession.test.ts @@ -78,7 +78,7 @@ rl.on("line", (line) => { return; } if (type === "get_session_stats") { - send({ type: "response", id, command: "get_session_stats", success: true, data: { sessionId: "mock-session-1", contextUsage: { tokens: 100, contextWindow: 1000, percent: 10 } } }); + send({ type: "response", id, command: "get_session_stats", success: true, data: { sessionId: "mock-session-1", contextUsage: { tokens: 100, contextWindow: 1000, percent: 10 }, tokens: { input: 60, output: 25, cacheRead: 10, cacheWrite: 5, total: 100 }, cost: 0.01 } }); return; } if (type === "get_commands") { @@ -186,6 +186,29 @@ describe("PiRpcSession (mock pi --mode rpc)", () => { await session.dispose(); }); + it("publishes cumulative usage.spent from get_session_stats tokens.total", async () => { + const { session, events } = await createSession(); + await session.startTurn?.("hello", { model: "mock/model", effort: "off" }); + + // Spend publishes after the turn settles, off the same stats response as + // context.updated — tokens.total (billed cumulative), not contextUsage. + const spent = await waitFor(events, (e) => e.type === "usage.spent"); + expect(spent).toMatchObject({ + type: "usage.spent", + threadId: "thread-mock", + usage: { + counterKind: "cumulative", + counter: 100, + scopeId: "mock-session-1", + epoch: 0, + fresh: true, + sampleId: "mock-session-1:0:100", + model: "mock/model", + }, + }); + await session.dispose(); + }); + it("surfaces a provider error as a failed turn", async () => { const { session, events } = await createSession(); await session.startTurn?.("FAIL please", { model: "mock/model", effort: "off" }); diff --git a/src/supervisor/agents/pi/rpcSession.ts b/src/supervisor/agents/pi/rpcSession.ts index 8667dc87f..619c7892c 100644 --- a/src/supervisor/agents/pi/rpcSession.ts +++ b/src/supervisor/agents/pi/rpcSession.ts @@ -12,6 +12,7 @@ import { type StructuredSessionHandle, type StructuredSessionListener, } from "../base"; +import { readNonNegativeInteger } from "../contextUsage"; import { goalPayloadFromProviderState, startGoalItemEvents, @@ -133,6 +134,12 @@ export class PiRpcSession implements StructuredSessionHandle { private sessionRef: ReturnType | undefined; private disposed = false; private interruptRequested = false; + /** True when the CLI was launched with `--session ` (resumed, not fresh). */ + private readonly launchedWithResume: boolean; + /** usage.spent ledger scope: pi session id + epoch (bumped if the id changes). */ + private usageScopeId: string | undefined; + private usageEpoch = 0; + private usageScopeFresh = false; private constructor( private readonly input: CreateStructuredSessionInput, @@ -146,6 +153,7 @@ export class PiRpcSession implements StructuredSessionHandle { this.client = client; this.mcpExtensionPath = mcpExtensionPath; this.currentConfig = input.config; + this.launchedWithResume = input.sessionRef?.providerSessionId !== undefined; this.unsubscribeEvents = client.onEvent((event) => this.handleEvent(event)); this.unsubscribeExit = client.onExit(() => this.handleExit()); } @@ -864,6 +872,7 @@ export class PiRpcSession implements StructuredSessionHandle { const sessionId = typeof stats?.sessionId === "string" ? stats.sessionId : ""; if (sessionId) { this.sessionRef = createKnownSessionRef(sessionId); + this.publishUsageSpent(stats, sessionId); } if (!usage) return; const tokens = typeof usage.tokens === "number" ? usage.tokens : null; @@ -878,6 +887,40 @@ export class PiRpcSession implements StructuredSessionHandle { }); } + /** + * Cumulative `usage.spent` from `get_session_stats` — `stats.tokens.total` + * is pi's billed total for the session and keeps growing across compaction, + * unlike `contextUsage.tokens` (context-window occupancy for the dock). The + * ledger counts counter increases per (provider, scopeId, epoch); the epoch + * bumps if pi ever reports a different session id (session switch/fork), and + * `fresh` marks a session this process started new (resumed sessions get a + * baseline-only first sample — inherited history is not new spend). + */ + private publishUsageSpent(stats: Record | undefined, sessionId: string): void { + const totalTokens = readNonNegativeInteger(recordOf(stats?.tokens)?.total); + if (totalTokens === undefined) return; + if (this.usageScopeId !== sessionId) { + const firstScope = this.usageScopeId === undefined; + if (!firstScope) this.usageEpoch += 1; + this.usageScopeId = sessionId; + this.usageScopeFresh = firstScope && !this.launchedWithResume; + } + this.emit({ + type: "usage.spent", + threadId: this.input.threadId, + usage: { + counterKind: "cumulative", + counter: totalTokens, + scopeId: sessionId, + epoch: this.usageEpoch, + ...(this.usageScopeFresh ? { fresh: true } : {}), + sampleId: `${sessionId}:${this.usageEpoch}:${totalTokens}`, + ...(this.currentConfig.model ? { model: this.currentConfig.model } : {}), + }, + }); + this.usageScopeFresh = false; + } + private async publishSlashCommands(): Promise { const response = await this.client.request("get_commands").catch(() => undefined); const list = recordOf(response?.data)?.commands;