@genesislcap/ai-assistant 15.4.1 → 15.6.0

package/src/utils/sum-usage.test.ts

@@ -0,0 +1,140 @@
+import type { ChatMessage } from '@genesislcap/foundation-ai';
+import { assert, createLogicSuite } from '@genesislcap/foundation-testing';
+import { addUsage, emptyUsage, sumUsage, totalTokens } from './sum-usage';
+
+const suite = createLogicSuite('sumUsage');
+
+const msg = (over: Partial<ChatMessage> = {}): ChatMessage =>
+  ({ role: 'assistant', content: '', ...over }) as ChatMessage;
+
+suite('returns zeroes when no message carries usage', () => {
+  const usage = sumUsage([msg({ role: 'user', content: 'hi' })]);
+  assert.equal(usage, emptyUsage());
+  assert.is(totalTokens(usage), 0);
+});
+
+suite('splits the prompt into three disjoint input buckets', () => {
+  // `inputTokens` is the WHOLE prompt; the cache fields break it down. So uncached is
+  // the remainder, and the three must re-add to the prompt exactly — that identity is
+  // what stops the buckets double-counting the prompt when summed.
+  const usage = sumUsage([
+    msg({ inputTokens: 1000, cacheReadTokens: 700, cacheWriteTokens: 200, outputTokens: 50 }),
+  ]);
+  assert.is(usage.uncachedInputTokens, 100);
+  assert.is(usage.cacheReadTokens, 700);
+  assert.is(usage.cacheWriteTokens, 200);
+  assert.is(usage.outputTokens, 50);
+  assert.is(
+    usage.uncachedInputTokens + usage.cacheReadTokens + usage.cacheWriteTokens,
+    1000,
+    'input buckets re-add to the prompt total',
+  );
+  assert.is(totalTokens(usage), 1050);
+});
+
+suite('treats the whole prompt as uncached when no cache split is reported', () => {
+  // Providers that report no cache breakdown (and messages persisted before the fields
+  // existed) must still contribute their full prompt, not drop it.
+  const usage = sumUsage([msg({ inputTokens: 400, outputTokens: 60 })]);
+  assert.is(usage.uncachedInputTokens, 400);
+  assert.is(usage.cacheReadTokens, 0);
+  assert.is(usage.cacheWriteTokens, 0);
+  assert.is(totalTokens(usage), 460);
+});
+
+suite('clamps uncached input at zero when buckets exceed the prompt total', () => {
+  // Malformed or hand-edited history must not be able to drive a cost display negative.
+  const usage = sumUsage([
+    msg({ inputTokens: 100, cacheReadTokens: 90, cacheWriteTokens: 90, outputTokens: 0 }),
+  ]);
+  assert.is(usage.uncachedInputTokens, 0);
+});
+
+suite('sums LLM cost and widget-reported external cost', () => {
+  const usage = sumUsage([msg({ cost: 0.25 }), msg({ externalCostUsd: 0.1, cost: 0.05 })]);
+  assert.is(Math.round(usage.costUsd * 100) / 100, 0.4);
+});
+
+suite('recurses into sub-agent traces', () => {
+  // A sub-agent's whole conversation is embedded in the parent's tool call, so this is
+  // the only place its spend is visible to a host total.
+  const usage = sumUsage([
+    msg({
+      toolCalls: [
+        {
+          id: 'tc-1',
+          name: 'delegate',
+          args: {},
+          subAgentTrace: [
+            msg({ cost: 0.2, inputTokens: 300, cacheReadTokens: 100, outputTokens: 40 }),
+          ],
+        },
+      ],
+    }),
+  ]);
+  assert.is(usage.costUsd, 0.2);
+  assert.is(usage.uncachedInputTokens, 200);
+  assert.is(usage.cacheReadTokens, 100);
+  assert.is(usage.outputTokens, 40);
+});
+
+suite('counts usage banked onto a compaction summary', () => {
+  // Compaction destroys the messages it summarises. Without the banked rollup, a total
+  // re-derived from the transcript would fall by whatever those turns had spent.
+  const before = sumUsage([
+    msg({ cost: 0.5, inputTokens: 1000, cacheReadTokens: 400, outputTokens: 100 }),
+    msg({ cost: 0.3, inputTokens: 500, outputTokens: 50 }),
+  ]);
+  const afterCompaction = sumUsage([
+    msg({
+      role: 'compacted-summary',
+      content: 'summary of the above',
+      compaction: {
+        compactedCount: 2,
+        createdAt: '2026-06-14T12:00:00.000Z',
+        rolledUpUsage: before,
+      },
+    }),
+  ]);
+  assert.equal(afterCompaction, before, 'compaction preserves the running total');
+});
+
+suite('ignores tokensBefore, which is a context reading rather than spend', () => {
+  // `tokensBefore` is one turn's context-window size. Counting it as cumulative spend
+  // would silently inflate the session total on every compaction.
+  const usage = sumUsage([
+    msg({
+      role: 'compacted-summary',
+      content: 'summary',
+      compaction: {
+        compactedCount: 3,
+        createdAt: '2026-06-14T12:00:00.000Z',
+        tokensBefore: 90_000,
+      },
+    }),
+  ]);
+  assert.is(totalTokens(usage), 0);
+});
+
+suite('addUsage adds every bucket and mutates neither operand', () => {
+  const a = { ...emptyUsage(), costUsd: 1.5, uncachedInputTokens: 100, outputTokens: 10 };
+  const b = { ...emptyUsage(), costUsd: 0.25, cacheReadTokens: 900, outputTokens: 5 };
+  const sum = addUsage(a, b);
+
+  assert.is(sum.costUsd, 1.75);
+  assert.is(sum.uncachedInputTokens, 100);
+  assert.is(sum.cacheReadTokens, 900);
+  assert.is(sum.outputTokens, 15);
+  // Both operands are live state (a row's banked figure and the session total), so a
+  // caller must be able to add without either being altered underneath it.
+  assert.is(a.costUsd, 1.5);
+  assert.is(b.cacheReadTokens, 900);
+});
+
+suite('addUsage with an empty operand is the identity', () => {
+  const a = { ...emptyUsage(), costUsd: 3, cacheWriteTokens: 42 };
+  assert.equal(addUsage(a, emptyUsage()), a);
+  assert.equal(addUsage(emptyUsage(), a), a);
+});
+
+suite.run();

package/src/utils/sum-usage.ts

@@ -0,0 +1,130 @@
+import type { AggregateUsage, ChatMessage } from '@genesislcap/foundation-ai';
+
+/**
+ * Zeroed {@link AggregateUsage}. A factory, not a shared const, because callers
+ * accumulate into the returned object.
+ *
+ * @internal
+ */
+export function emptyUsage(): AggregateUsage {
+  return {
+    costUsd: 0,
+    uncachedInputTokens: 0,
+    cacheReadTokens: 0,
+    cacheWriteTokens: 0,
+    outputTokens: 0,
+  };
+}
+
+/**
+ * Total tokens processed — the four buckets added up. Safe to sum because the
+ * buckets are disjoint (see {@link AggregateUsage}); the equivalent sum over
+ * per-message fields would double-count the prompt.
+ *
+ * Bear in mind what this number is: because every turn resends the conversation,
+ * a turn's prompt is counted again in the next turn's. It is billed throughput,
+ * NOT the count of distinct tokens in the conversation, so it grows
+ * super-linearly with turn count and is expected to dwarf the context size.
+ *
+ * @internal
+ */
+export function totalTokens(usage: AggregateUsage): number {
+  return (
+    usage.uncachedInputTokens + usage.cacheReadTokens + usage.cacheWriteTokens + usage.outputTokens
+  );
+}
+
+/**
+ * `a + b`, per bucket, as a new object.
+ *
+ * Used to lay a row's banked (pre-transcript) spend over what the live transcript
+ * proves. Pure rather than in-place because both operands are live state a caller
+ * must not mutate.
+ *
+ * @internal
+ */
+export function addUsage(a: AggregateUsage, b: AggregateUsage): AggregateUsage {
+  return {
+    costUsd: a.costUsd + b.costUsd,
+    uncachedInputTokens: a.uncachedInputTokens + b.uncachedInputTokens,
+    cacheReadTokens: a.cacheReadTokens + b.cacheReadTokens,
+    cacheWriteTokens: a.cacheWriteTokens + b.cacheWriteTokens,
+    outputTokens: a.outputTokens + b.outputTokens,
+  };
+}
+
+/**
+ * Sum cost and per-bucket token usage across a message list.
+ *
+ * One walk produces everything a usage display needs, replacing separate cost and
+ * token passes over the same (potentially deep) transcript. Recurses into each
+ * tool call's `subAgentTrace`, so LLM work a sub-agent did — possibly on another
+ * provider at other rates — lands in the total; the driver embeds a whole
+ * sub-agent conversation in the parent's tool result, so this is the only place
+ * that spend is visible.
+ *
+ * Splitting the prompt into buckets is the point: `cost` already has the cache
+ * discount applied, but a single token total cannot show that a session was cheap
+ * because most of its prompt was a cache hit. The three input buckets are derived
+ * from the per-message fields, where `inputTokens` is the whole prompt and the
+ * cache fields break it down — so uncached input is the remainder, never an
+ * addition.
+ *
+ * Also counts `compaction.rolledUpUsage`: compaction destroys the messages it
+ * summarises, and without their banked usage every total re-derived from the
+ * transcript would drop by whatever those turns had already spent.
+ *
+ * Returns a zeroed total when no message carries usage — providers that report
+ * none (e.g. Chrome built-in) are skipped silently rather than guessed at.
+ *
+ * @internal
+ */
+export function sumUsage(messages: readonly ChatMessage[]): AggregateUsage {
+  const total = emptyUsage();
+  accumulate(messages, total);
+  return total;
+}
+
+/** Add `from` into `into`, in place. */
+function addInto(into: AggregateUsage, from: AggregateUsage): void {
+  into.costUsd += from.costUsd;
+  into.uncachedInputTokens += from.uncachedInputTokens;
+  into.cacheReadTokens += from.cacheReadTokens;
+  into.cacheWriteTokens += from.cacheWriteTokens;
+  into.outputTokens += from.outputTokens;
+}
+
+/**
+ * Walk `messages`, accumulating into `into`. Mutates a single accumulator rather
+ * than merging returned objects — this runs on every transcript change, so the
+ * per-message allocation is worth avoiding.
+ */
+function accumulate(messages: readonly ChatMessage[], into: AggregateUsage): void {
+  for (const m of messages) {
+    // Provider/LLM cost plus any non-LLM cost a widget reported for its own
+    // external service calls. Both are real spend on this session.
+    if (m.cost != null) into.costUsd += m.cost;
+    if (m.externalCostUsd != null) into.costUsd += m.externalCostUsd;
+
+    const cacheRead = m.cacheReadTokens ?? 0;
+    const cacheWrite = m.cacheWriteTokens ?? 0;
+    into.cacheReadTokens += cacheRead;
+    into.cacheWriteTokens += cacheWrite;
+    // The uncached remainder of the prompt. Clamped at zero so malformed or
+    // hand-edited history whose buckets exceed the prompt total can't drive the
+    // session total negative — a wrong-but-plausible number is worse than a
+    // conservative one here, since this feeds a cost display.
+    into.uncachedInputTokens += Math.max(0, (m.inputTokens ?? 0) - cacheRead - cacheWrite);
+    into.outputTokens += m.outputTokens ?? 0;
+
+    // Spend banked by a compaction on this summary's behalf. NOTE: deliberately
+    // not `compaction.tokensBefore` — that is a one-turn context-window reading,
+    // not cumulative spend, and counting it here would inflate the total.
+    const rolled = m.compaction?.rolledUpUsage;
+    if (rolled) addInto(into, rolled);
+
+    for (const tc of m.toolCalls ?? []) {
+      if (tc.subAgentTrace) accumulate(tc.subAgentTrace, into);
+    }
+  }
+}