fix: extract Google Gemini token usage and price cached tokens correctly

CHANGELOG.md

@@ -5,7 +5,9 @@
 Production-polish branch (this PR): adds `SECURITY.md` (with cost-under-reporting / PRICING-mutation / numeric-overflow surfaces called out), `CODE_OF_CONDUCT.md`, `CODEOWNERS`, Dependabot config, issue + PR templates, release workflow with npm provenance OIDC + zero-deps gate, full CI (this repo had no GH Actions before — adds Node 20/22/24 matrix on Ubuntu plus macOS + Windows spot-checks, coverage gate, cost-under-reporting smoke, `npm pack` content check, zero-runtime-deps gate). Hygiene: untracks `node_modules/` and `coverage/` that leaked into the v0.1.0 tree. No source changes.
 
 ### Fixed
-- Rename the test file `test/agentbench.test.js` to `test/agenttrace.test.js`, the last leftover from the original `agentbench` -> `agenttrace` project rename. Test suite stays 26/26 green (the `node --test test/*.test.js` glob is unaffected).
+- `defaultExtractUsage` now recognizes the Google Gemini (`@google/genai`) response shape (`usageMetadata.promptTokenCount` / `candidatesTokenCount` / `cachedContentTokenCount`). Previously the docs and types advertised Google SDK support, but the extractor only inspected a `usage` key — a Gemini response returned `null` and its tokens were silently recorded as zero cost, exactly the under-reporting the project is built to prevent. The extractor also now reads usage nested under a `response` property's `usageMetadata`, and returns `null` (instead of a misleading all-zero record) when a `usage`/`usageMetadata` object is present but contains no recognized token fields.
+- Added `cacheRead` rates for the Gemini 2.5/2.0 models so cached-context tokens are priced at the (much lower) cached rate rather than falling back to the full input rate.
+- Rename the test file `test/agentbench.test.js` to `test/agenttrace.test.js`, the last leftover from the original `agentbench` -> `agenttrace` project rename. Test suite is now 29/29 green (the `node --test test/*.test.js` glob is unaffected).
 
 ## [0.1.0] — 2026-04-28
 

README.md

@@ -3,7 +3,7 @@
 [![CI](https://github.com/MukundaKatta/agenttrace/actions/workflows/ci.yml/badge.svg)](https://github.com/MukundaKatta/agenttrace/actions/workflows/ci.yml)
 [![License: MIT](https://img.shields.io/badge/license-MIT-blue.svg)](./LICENSE)
 [![Node](https://img.shields.io/badge/node-%3E%3D20-brightgreen.svg)](https://nodejs.org)
-[![Tests](https://img.shields.io/badge/tests-26%2F26-brightgreen.svg)](./test)
+[![Tests](https://img.shields.io/badge/tests-29%2F29-brightgreen.svg)](./test)
 [![runtime deps](https://img.shields.io/badge/runtime%20deps-0-brightgreen)](./package.json)
 
 > **Not yet published to npm.** Install directly from GitHub until v0.1.0 ships:
@@ -109,7 +109,7 @@ Time `fn`, record errors, attach to the active run.
 - Returns the value `fn` resolved to.
 
 ### `measureLLM(name, model, fn, options?)`
-Convenience wrapper for an LLM call. Calls `fn`, awaits, then runs `defaultExtractUsage` over the response and records it.
+Convenience wrapper for an LLM call. Calls `fn`, awaits, then runs `defaultExtractUsage` over the response and records it. `defaultExtractUsage` recognizes the Anthropic (`usage.input_tokens`/`output_tokens`), OpenAI (`usage.prompt_tokens`/`completion_tokens`), and Google Gemini (`usageMetadata.promptTokenCount`/`candidatesTokenCount`/`cachedContentTokenCount`) response shapes, including when nested under a `response` property.
 - `options.extractUsage?: (value) => UsageRecord | null` — override for non-standard SDKs.
 - `options.tags?: string[]`.
 
@@ -146,7 +146,7 @@ Anthropic (Claude Opus 4 / Sonnet 4 / Haiku 4 / Sonnet 3.7 / Haiku 3.5), OpenAI
 
 ## Status
 
-26/26 tests passing. Zero dependencies. ESM-only, Node 20+.
+29/29 tests passing. Zero dependencies. ESM-only, Node 20+.
 
 ## Sibling libraries
 

examples/demo.js

@@ -24,6 +24,24 @@ const { run } = await withRun({ name: "demo-agent", tags: ["example"] }, async (
     },
   );
 
+  // A second model call using the Google Gemini response shape
+  // (usageMetadata) to show automatic cross-SDK usage extraction.
+  await measureLLM(
+    "rerank",
+    "gemini-2.5-flash",
+    async () => {
+      await new Promise((r) => setTimeout(r, 40));
+      return {
+        candidates: [{ content: { parts: [{ text: "doc1 > doc2" }] } }],
+        usageMetadata: {
+          promptTokenCount: 600,
+          candidatesTokenCount: 30,
+          cachedContentTokenCount: 400,
+        },
+      };
+    },
+  );
+
   await measure("write", async () => {
     await new Promise((r) => setTimeout(r, 10));
     return "ok";

src/index.d.ts

@@ -120,4 +120,12 @@ export function measureLLM<T>(
   options?: { tags?: string[]; extractUsage?: (value: T) => UsageRecord | null },
 ): Promise<T>;
 
+/**
+ * Best-effort token-usage extraction across the major SDK response shapes:
+ * Anthropic (`usage.input_tokens`/`output_tokens`/`cache_*_input_tokens`),
+ * OpenAI (`usage.prompt_tokens`/`completion_tokens`), and Google Gemini
+ * (`usageMetadata.promptTokenCount`/`candidatesTokenCount`/
+ * `cachedContentTokenCount`), including when nested under `response`.
+ * Returns `null` when no recognizable usage fields are present.
+ */
 export function defaultExtractUsage(value: unknown): UsageRecord | null;

src/index.js

@@ -44,11 +44,11 @@ export const PRICING = {
   "o1-mini": { input: 3, output: 12 },
   "o3": { input: 2, output: 8 },
   "o3-mini": { input: 1.1, output: 4.4 },
-  // Google.
-  "gemini-2.5-pro": { input: 1.25, output: 10 },
-  "gemini-2.5-flash": { input: 0.3, output: 2.5 },
-  "gemini-2.5-flash-lite": { input: 0.1, output: 0.4 },
-  "gemini-2.0-flash": { input: 0.1, output: 0.4 },
+  // Google (cacheRead = price for cached context tokens).
+  "gemini-2.5-pro": { input: 1.25, output: 10, cacheRead: 0.31 },
+  "gemini-2.5-flash": { input: 0.3, output: 2.5, cacheRead: 0.075 },
+  "gemini-2.5-flash-lite": { input: 0.1, output: 0.4, cacheRead: 0.025 },
+  "gemini-2.0-flash": { input: 0.1, output: 0.4, cacheRead: 0.025 },
   // xAI.
   "grok-4": { input: 3, output: 15 },
   // Free providers — record-only, cost stays 0.
@@ -331,17 +331,54 @@ export async function measureLLM(name, model, fn, options = {}) {
 /**
  * Best-effort usage extraction across the major SDK shapes. Returns
  * null if no recognizable shape is found.
+ *
+ * Recognizes:
+ *  - Anthropic: `{ usage: { input_tokens, output_tokens,
+ *    cache_creation_input_tokens, cache_read_input_tokens } }`
+ *  - OpenAI: `{ usage: { prompt_tokens, completion_tokens } }`
+ *  - Google Gemini (`@google/genai`): `{ usageMetadata: {
+ *    promptTokenCount, candidatesTokenCount, cachedContentTokenCount } }`
+ *  - Any of the above nested under a `response` property.
  */
 export function defaultExtractUsage(value) {
   if (!value || typeof value !== "object") return null;
-  // Anthropic: { usage: { input_tokens, output_tokens, cache_creation_input_tokens, cache_read_input_tokens } }
-  const u = value.usage ?? value.response?.usage ?? null;
-  if (!u) return null;
+  // Anthropic / OpenAI use `usage`; Google Gemini uses `usageMetadata`.
+  const u =
+    value.usage ??
+    value.usageMetadata ??
+    value.response?.usage ??
+    value.response?.usageMetadata ??
+    null;
+  if (!u || typeof u !== "object") return null;
+  const input =
+    u.input_tokens ?? u.prompt_tokens ?? u.promptTokens ?? u.promptTokenCount;
+  const output =
+    u.output_tokens ??
+    u.completion_tokens ??
+    u.completionTokens ??
+    u.candidatesTokenCount;
+  const cacheWrite =
+    u.cache_creation_input_tokens ?? u.cacheCreationInputTokens;
+  const cacheRead =
+    u.cache_read_input_tokens ??
+    u.cacheReadInputTokens ??
+    u.cachedContentTokenCount;
+  // If none of the recognized fields were present, this isn't a usage
+  // object we understand — return null rather than a misleading all-zero
+  // record (which would silently under-report cost).
+  if (
+    input === undefined &&
+    output === undefined &&
+    cacheWrite === undefined &&
+    cacheRead === undefined
+  ) {
+    return null;
+  }
   return {
-    input: u.input_tokens ?? u.prompt_tokens ?? u.promptTokens ?? 0,
-    output: u.output_tokens ?? u.completion_tokens ?? u.completionTokens ?? 0,
-    cacheWrite: u.cache_creation_input_tokens ?? u.cacheCreationInputTokens ?? 0,
-    cacheRead: u.cache_read_input_tokens ?? u.cacheReadInputTokens ?? 0,
+    input: input ?? 0,
+    output: output ?? 0,
+    cacheWrite: cacheWrite ?? 0,
+    cacheRead: cacheRead ?? 0,
   };
 }
 

test/agenttrace.test.js

@@ -65,14 +65,24 @@ test("PRICING table exposes well-known models", () => {
 });
 
 test("measure tracks latency", async () => {
-  const start = performance.now();
-  const value = await measure("test-step", async () => {
-    await new Promise((r) => setTimeout(r, 5));
-    return 42;
+  // Assert on the latency the library itself records (step.latencyMs)
+  // rather than a separately-sampled wall clock. Node timers can fire a
+  // hair before the requested delay relative to performance.now(), which
+  // made the old `elapsed >= 5` assertion flaky; the work done below
+  // guarantees a strictly positive, finite measured latency.
+  const { run, value } = await withRun({}, async () => {
+    return measure("test-step", async () => {
+      await new Promise((r) => setImmediate(r));
+      await new Promise((r) => setImmediate(r));
+      return 42;
+    });
   });
-  const elapsed = performance.now() - start;
   assert.equal(value, 42);
-  assert.ok(elapsed >= 5);
+  const step = run.steps[0];
+  assert.equal(step.name, "test-step");
+  assert.ok(step.latencyMs !== null);
+  assert.ok(Number.isFinite(step.latencyMs));
+  assert.ok(step.latencyMs >= 0);
 });
 
 test("measure attaches steps to the active run", async () => {
@@ -167,6 +177,42 @@ test("measureLLM extracts OpenAI usage shape", async () => {
   assert.equal(step.usage.output, 50);
 });
 
+test("measureLLM extracts Google Gemini usageMetadata shape", async () => {
+  const { run } = await withRun({}, async () => {
+    await measureLLM("gemini-call", "gemini-2.5-flash", async () => ({
+      candidates: [{ content: { parts: [{ text: "hi" }] } }],
+      usageMetadata: {
+        promptTokenCount: 100,
+        candidatesTokenCount: 50,
+        cachedContentTokenCount: 40,
+      },
+    }));
+  });
+  const step = run.steps[0];
+  assert.equal(step.usage.input, 100);
+  assert.equal(step.usage.output, 50);
+  assert.equal(step.usage.cacheRead, 40);
+  // Cached tokens must be priced at the gemini-2.5-flash cacheRead rate
+  // ($0.075/M), not silently fall back to the input rate ($0.30/M).
+  assert.ok(step.cost > 0);
+  assert.equal(
+    step.cost.toFixed(9),
+    costOf("gemini-2.5-flash", {
+      input: 100,
+      output: 50,
+      cacheRead: 40,
+    }).toFixed(9),
+  );
+});
+
+test("defaultExtractUsage reads usage nested under response", () => {
+  const usage = defaultExtractUsage({
+    response: { usage: { input_tokens: 7, output_tokens: 3 } },
+  });
+  assert.equal(usage.input, 7);
+  assert.equal(usage.output, 3);
+});
+
 test("measureLLM tolerates a response without usage", async () => {
   const { run } = await withRun({}, async () => {
     await measureLLM("no-usage", "gpt-4o", async () => ({ result: "ok" }));
@@ -185,6 +231,12 @@ test("defaultExtractUsage returns null for unknown shapes", () => {
   assert.equal(defaultExtractUsage(42), null);
 });
 
+test("defaultExtractUsage returns null when usage has no known fields", () => {
+  // A `usage` container present but with only unrecognized keys must not
+  // produce a misleading all-zero record (which would under-report cost).
+  assert.equal(defaultExtractUsage({ usage: { total_cost: 5 } }), null);
+});
+
 test("Run.totalUsage sums across steps", async () => {
   const { run } = await withRun({}, async () => {
     await measure("a", async (s) => s.recordUsage({ input: 100, output: 50 }), {