LLM Token-Usage Capture — Design

Design for capturing per-call LLM token usage across the platform. Draft for review, July 2026.

Scope: blitz backend only — the frontend (blitz-ui) makes no direct LLM calls; Finni routes through blitz-agents over WebSocket, so all usage is server-side.

Status: design only. Emission target (logs vs metrics vs DB) is deferred — kept behind a pluggable sink so that choice touches no call site. Related: blitz#898 (agents deployment).

Goal

Capture per-call token usage (input / output / total, plus cache & reasoning breakdown where available) for every LLM API call, with enough context (tenant, agent, session, model) to later attribute cost and build usage dashboards or billing.

The two touchpoints

All ~30 call sites collapse to two instrumentation points:

Path	Reaches	Instrument once at
A — LangChain / Gemini (model.invoke)	all 29 agent call sites in systemagents, financeagents, systemwf, docai/gemini.ts	the model factory createModel() — src/libs/agentsbase/src/models/modelConfig.ts:142
B — Google native SDK (generateContent)	1 site, doc extraction	the call site — src/libs/docai/src/googlegenerativeai.ts:40

Because withStructuredOutput(...) discards usage_metadata on its parsed return value, per-call-site reads would silently miss the majority of Path A. The factory-level callback handler avoids this: LangChain fires handleLLMEnd on the underlying model regardless of structured-output wrapping.

Path A — LangChain callback handler (primary)

Mechanism

Define a BaseCallbackHandler that reads token usage in handleLLMEnd and hands a normalized record to the sink. Attach it in createModel() so every model instance carries it — zero edits to the 29 call sites.

// src/libs/agentsbase/src/observability/tokenUsageHandler.ts  (new)
import { BaseCallbackHandler } from "@langchain/core/callbacks/base";
import type { LLMResult } from "@langchain/core/outputs";

export class TokenUsageHandler extends BaseCallbackHandler {
  name = "TokenUsageHandler";
  constructor(private modelType: string, private modelName: string, private sink: UsageSink) {
    super();
  }

  async handleLLMEnd(output: LLMResult, runId: string, parentRunId?: string,
                     tags?: string[]): Promise<void> {
    // Chat models: usage is on the generation message; llmOutput is a fallback.
    const gen = output.generations?.[0]?.[0] as any;
    const u = gen?.message?.usage_metadata ?? output.llmOutput?.tokenUsage;
    if (!u) return;                      // no usage → nothing to record

    this.sink.record({
      inputTokens:  u.input_tokens  ?? u.promptTokens     ?? 0,
      outputTokens: u.output_tokens ?? u.completionTokens ?? 0,
      totalTokens:  u.total_tokens  ?? u.totalTokens       ?? 0,
      cachedInputTokens:  u.input_token_details?.cache_read ?? 0,
      reasoningTokens:    u.output_token_details?.reasoning ?? 0,   // Gemini 2.5 "thoughts"
      modelType: this.modelType,         // "gemini" | "ollama" | "openai"
      modelName: this.modelName,         // e.g. "gemini-flash-latest"
      runId, parentRunId, tags,
      // context (tenant/agent/session) is merged in from run metadata — see below
    });
  }
}

Wiring in the factory:

// modelConfig.ts — createModel()
case "gemini":
  return new ChatGoogleGenerativeAI({
    model: config.model, temperature: config.temperature,
    maxOutputTokens: config.maxTokens, apiKey: config.apiKey,
    callbacks: [new TokenUsageHandler("gemini", config.model, usageSink)],
  });

Add the same one-line callbacks to the ollama and openai branches so free/local and any future paid provider are covered and distinguishable by modelType.

Context propagation (tenant / agent / session)

The handler knows what model and how many tokens, but not whose request. LangChain propagates metadata from a top-level .invoke(input, config) down to every nested model call and exposes it on the callback run. So context is injected once at each request entry and read back in the handler — no threading through internal functions.

• Graph-based agents (most of apps/agents/src/controllers/agents/*): add metadata to the existing graph.invoke(state, config):

  await graph.invoke(state, {
    metadata: { tenantId, agentName: "finni", sessionId, userId },
  });

  LangGraph forwards this to every child model.invoke, so the handler sees it on all fan-out calls automatically.
• Direct model calls (e.g. src/packages/system/systemwf/src/activities/agentactivities.ts:31): pass the same metadata on model.invoke(messages, { metadata }).

The handler reads it via the run metadata exposed on handleLLMStart / handleLLMEnd (stash by runId in handleChatModelStart if needed).

Effort note: this is the only part that touches call sites, and only the entry sites (~a dozen controllers + the Temporal activity), not the 29 inner ones. If we accept "no per-tenant attribution in v1," this step can be skipped and added later.

Path B — Google native SDK (doc extraction)

One site, non-LangChain. Read usageMetadata off the response and call the same sink:

// googlegenerativeai.ts — after generateContent(...)
const result = await model.generateContent([...]);
const um = result.response.usageMetadata;   // { promptTokenCount, candidatesTokenCount,
                                            //   totalTokenCount, cachedContentTokenCount,
                                            //   thoughtsTokenCount }
usageSink.record({
  inputTokens: um?.promptTokenCount ?? 0,
  outputTokens: um?.candidatesTokenCount ?? 0,
  totalTokens: um?.totalTokenCount ?? 0,
  cachedInputTokens: um?.cachedContentTokenCount ?? 0,
  reasoningTokens: um?.thoughtsTokenCount ?? 0,
  modelType: "gemini",
  modelName: process.env.GOOGLE_GENAI_MODEL || "gemini-1.5-flash",
});

Normalized usage record (shared contract)

Both paths emit the same shape, so the sink and any downstream (dashboard/billing) see one schema:

interface TokenUsageEvent {
  timestamp: string;          // ISO, stamped by the sink
  modelType: "gemini" | "ollama" | "openai";
  modelName: string;          // resolved model, incl. floating aliases like gemini-flash-latest
  inputTokens: number;
  outputTokens: number;
  totalTokens: number;
  cachedInputTokens: number;  // billed at a lower rate; 0 if none
  reasoningTokens: number;    // Gemini 2.5 "thoughts"; subset of output, priced differently
  // context (nullable in v1 if propagation deferred)
  tenantId?: string;
  agentName?: string;
  sessionId?: string;
  userId?: string;
  runId?: string;
  parentRunId?: string;
}

Pluggable sink (defers the emission decision)

interface UsageSink { record(e: TokenUsageEvent): void; }  // fire-and-forget, never blocks the request

The sink is the only thing that changes when the emission target is chosen later:

Sink impl	Effort	Good for	Trade-off
Structured log (pino via existing Logger)	lowest	dashboards from log aggregation, quick start	not queryable as first-class data; retention tied to log retention
Metrics counter (labels: model, agent, tenant)	low	real-time totals, alerting	aggregate only — no per-call detail/audit
Postgres table (llm_usage in mgmt or finance)	medium	billing, chargeback, historical queries, audit	needs Prisma model + migration; write path must be async/batched

The sink must be non-blocking and failure-isolated — a usage-write error must never fail the LLM request (wrap in try/catch, drop on error, count drops).

Edge cases & correctness notes

• Ollama fallback skews cost. getModel() silently falls back to gemma4 (Ollama, local/free) on any Gemini failure (modelConfig.ts:180-190). modelType in the record separates billable Gemini from free local — always filter by it before costing.
• Floating alias drift. gemini-flash maps to gemini-flash-latest; the actual billed model can change under a stable config name. Record the resolved modelName; if a price table is added later, key it on resolved name and alert on unknown names.
• Errors → no usage. On invoke failure handleLLMError fires (not handleLLMEnd), so failed calls record nothing. Usually correct (no tokens billed), but a 429-after-partial can still incur input-token cost we won't see — accepted gap for v1.
• No streaming anywhere — all calls are single-response (agentactivities.ts:20-23 forces it), so no final-chunk aggregation is needed. If streaming is enabled later, handleLLMEnd still fires once at stream end with cumulative usage — design holds.
• No embeddings exist today — nothing to instrument, but embeddings would route through the same factory if added.
• Existing reader is now redundant. agentactivities.ts:36-40 already reads usage_metadata into AgentInvokeOutput. Harmless to keep; can be simplified to rely on the handler once landed.
• Retries (LangChain built-in) fire handleLLMEnd per successful attempt — usage counted per real API call, which is what we want for cost.

Rollout plan

1. Handler + sink interface + log sink in agentsbase/observability — covers all of Path A immediately with a log sink (no schema decisions). Smallest shippable slice — full capture, deferred emission.
2. Wire factory (createModel, 3 branches) — Path A live end-to-end.
3. Path B one-site edit — full coverage.
4. Context propagation at request-entry sites — adds tenant/agent/session.
5. Swap/duplicate sink to metrics or Postgres once the emission target is decided.

Steps 1–3 give complete token capture with no cross-cutting schema work; 4–5 add attribution and the chosen durable target.

Verification

• Unit: feed a mock LLMResult (Gemini-shaped usage_metadata, incl. input_token_details) to the handler; assert the emitted event.
• Integration: run one agent (e.g. Finni) against Gemini in staging, confirm one event per model call with non-zero tokens and correct modelName; run with the Gemini key unset to force Ollama fallback and confirm modelType: "ollama".
• Cross-check Path A totals against the numbers already surfaced at agentactivities.ts:47 for the wfw path.

Open decisions

#	Item	Status
1	Attribution granularity — per-tenant/session in v1, or per-model/agent to start?	Pending
2	Tokens only now vs computed cost (price table)	Pending — aliases + Ollama fallback argue for cost downstream
3	Durable table location if DB sink — mgmt (operator/billing) vs finance	Pending