From f193877a1cd1d76e77c7a2f9c1792b2831fefd88 Mon Sep 17 00:00:00 2001
From: Diego Rodrigues de Sa e Souza <souzamiriamrodrigues790@gmail.com>
Date: Tue, 23 Jun 2026 08:56:54 -0300
Subject: [PATCH] fix(sse): read Antigravity usage from the
 response.usageMetadata envelope

Antigravity (AG MITM) streaming chunks wrap the Gemini usageMetadata
inside a response envelope: { response: { usageMetadata: {...} } }.
extractUsage only inspected the top-level chunk.usageMetadata, so token
usage was silently dropped for AG-shaped chunks. Fall back to
chunk.response?.usageMetadata when the top-level field is absent.

Inspired-by: https://github.com/decolua/9router/commit/50990e84b40c76f2c06ca3d53b60171ca4b10d54
Co-authored-by: decolua <decoluadt@example.com>
---
 open-sse/utils/usageTracking.ts    | 16 ++++++-----
 tests/unit/usage-extractor.test.ts | 43 ++++++++++++++++++++++++++++++
 2 files changed, 53 insertions(+), 6 deletions(-)

diff --git a/open-sse/utils/usageTracking.ts b/open-sse/utils/usageTracking.ts
index f1758b70304..7528960ace3 100644
--- a/open-sse/utils/usageTracking.ts
+++ b/open-sse/utils/usageTracking.ts
@@ -392,13 +392,17 @@ export function extractUsage(chunk) {
   }
 
   // Gemini format (Antigravity)
-  if (chunk.usageMetadata && typeof chunk.usageMetadata === "object") {
+  // Antigravity wraps usageMetadata inside a `response` envelope:
+  // { response: { usageMetadata: {...} } } — fall back to it so AG-shaped
+  // chunks do not silently drop token usage.
+  const usageMeta = chunk.usageMetadata || chunk.response?.usageMetadata;
+  if (usageMeta && typeof usageMeta === "object") {
     return normalizeUsage({
-      prompt_tokens: chunk.usageMetadata?.promptTokenCount || 0,
-      completion_tokens: chunk.usageMetadata?.candidatesTokenCount || 0,
-      total_tokens: chunk.usageMetadata?.totalTokenCount,
-      cached_tokens: chunk.usageMetadata?.cachedContentTokenCount,
-      reasoning_tokens: chunk.usageMetadata?.thoughtsTokenCount,
+      prompt_tokens: usageMeta.promptTokenCount || 0,
+      completion_tokens: usageMeta.candidatesTokenCount || 0,
+      total_tokens: usageMeta.totalTokenCount,
+      cached_tokens: usageMeta.cachedContentTokenCount,
+      reasoning_tokens: usageMeta.thoughtsTokenCount,
     });
   }
 
diff --git a/tests/unit/usage-extractor.test.ts b/tests/unit/usage-extractor.test.ts
index e156a989afb..84d330e7579 100644
--- a/tests/unit/usage-extractor.test.ts
+++ b/tests/unit/usage-extractor.test.ts
@@ -374,3 +374,46 @@ test("extractUsage ignores non-final Ollama NDJSON chunks (done=false)", () => {
 
   assert.equal(usage, null);
 });
+
+// ── Antigravity (Gemini) streaming usageMetadata tests ──
+
+test("extractUsage reads top-level Gemini usageMetadata from a streaming chunk", () => {
+  const usage = extractUsage({
+    usageMetadata: {
+      promptTokenCount: 120,
+      candidatesTokenCount: 60,
+      totalTokenCount: 180,
+      cachedContentTokenCount: 30,
+      thoughtsTokenCount: 12,
+    },
+  });
+
+  assert.equal(usage.prompt_tokens, 120);
+  assert.equal(usage.completion_tokens, 60);
+  assert.equal(usage.total_tokens, 180);
+  assert.equal(usage.cached_tokens, 30);
+  assert.equal(usage.reasoning_tokens, 12);
+});
+
+test("extractUsage reads Antigravity usageMetadata wrapped inside a response envelope", () => {
+  // Antigravity (AG MITM) shapes usage as { response: { usageMetadata: {...} } }.
+  // Without the response.usageMetadata fallback, token usage is silently dropped.
+  const usage = extractUsage({
+    response: {
+      usageMetadata: {
+        promptTokenCount: 200,
+        candidatesTokenCount: 75,
+        totalTokenCount: 275,
+        cachedContentTokenCount: 40,
+        thoughtsTokenCount: 18,
+      },
+    },
+  });
+
+  assert.notEqual(usage, null);
+  assert.equal(usage.prompt_tokens, 200);
+  assert.equal(usage.completion_tokens, 75);
+  assert.equal(usage.total_tokens, 275);
+  assert.equal(usage.cached_tokens, 40);
+  assert.equal(usage.reasoning_tokens, 18);
+});