feat(token-meter): project durable token usage

packages/llm/token-meter/src/usage-projection.ts

@@ -0,0 +1,100 @@
+/**
+ * Pure fold for durable provider-reported token usage.
+ */
+
+import { z } from 'zod'
+import type { TokenUsage } from '@deepseek-ai/dsh-llm'
+import type { ProjectionDefinition } from '@deepseek-ai/dsh-session-projection'
+import type { TokenUsageProjection } from './projection.ts'
+
+interface UsageSample {
+  turn: number
+  step: number
+  buckets: TokenUsageProjection
+}
+
+interface TokenUsageState {
+  totals: TokenUsageProjection
+  last: UsageSample | null
+}
+
+const zeroBuckets = (): TokenUsageProjection => ({
+  uncachedInputTokens: 0,
+  outputTokens: 0,
+  cacheReadTokens: 0,
+  cacheWriteTokens: 0,
+})
+
+const bucketsFrom = (usage: TokenUsage): TokenUsageProjection => ({
+  uncachedInputTokens: usage.inputTokens,
+  outputTokens: usage.outputTokens,
+  cacheReadTokens: usage.cacheReadTokens ?? 0,
+  cacheWriteTokens: usage.cacheWriteTokens ?? 0,
+})
+
+const bucketsEqual = (left: TokenUsageProjection, right: TokenUsageProjection): boolean =>
+  left.uncachedInputTokens === right.uncachedInputTokens
+  && left.outputTokens === right.outputTokens
+  && left.cacheReadTokens === right.cacheReadTokens
+  && left.cacheWriteTokens === right.cacheWriteTokens
+
+const addReplacing = (
+  totals: TokenUsageProjection,
+  previous: TokenUsageProjection | undefined,
+  next: TokenUsageProjection,
+): TokenUsageProjection => ({
+  uncachedInputTokens: totals.uncachedInputTokens - (previous?.uncachedInputTokens ?? 0) + next.uncachedInputTokens,
+  outputTokens: totals.outputTokens - (previous?.outputTokens ?? 0) + next.outputTokens,
+  cacheReadTokens: totals.cacheReadTokens - (previous?.cacheReadTokens ?? 0) + next.cacheReadTokens,
+  cacheWriteTokens: totals.cacheWriteTokens - (previous?.cacheWriteTokens ?? 0) + next.cacheWriteTokens,
+})
+
+const projectionSchema = z.object({
+  uncachedInputTokens: z.number().int().nonnegative(),
+  outputTokens: z.number().int().nonnegative(),
+  cacheReadTokens: z.number().int().nonnegative(),
+  cacheWriteTokens: z.number().int().nonnegative(),
+}).strict()
+
+/**
+ * Token-meter's session projection unit.
+ *
+ * Usage chunks provide an early sample that survives a later request failure;
+ * an assistant message provides the final sample for the same turn/step. A
+ * repeated sample replaces that step's earlier value instead of double
+ * counting it.
+ */
+export const tokenUsageProjectionDefinition:
+ProjectionDefinition<'tokenUsage', TokenUsageState> = {
+  key: 'tokenUsage',
+  schema: projectionSchema,
+  init: () => ({ totals: zeroBuckets(), last: null }),
+  apply: (state, event) => {
+    let turn: number
+    let step: number
+    let usage: TokenUsage
+    if (event.type === 'assistant/chunk' && event.data.chunk.type === 'usage') {
+      ;({ turn, step } = event.data)
+      usage = event.data.chunk.usage
+    } else if (event.type === 'assistant/message' && event.data.usage !== undefined) {
+      ;({ turn, step, usage } = event.data)
+    } else {
+      return state
+    }
+
+    const buckets = bucketsFrom(usage)
+    const previous = state.last !== null
+      && state.last.turn === turn
+      && state.last.step === step
+      ? state.last.buckets
+      : undefined
+    if (previous !== undefined && bucketsEqual(previous, buckets)) return state
+
+    return {
+      totals: addReplacing(state.totals, previous, buckets),
+      last: { turn, step, buckets },
+    }
+  },
+  view: state => state.totals,
+  stateVersion: 1,
+}