From 27133a7ac620f8c2ee52eb0c50b69ee65ff06ae3 Mon Sep 17 00:00:00 2001
From: Diego Rodrigues de Sa e Souza
 <8016841+diegosouzapw@users.noreply.github.com>
Date: Thu, 25 Jun 2026 22:58:55 -0300
Subject: [PATCH] fix(sse): read Antigravity usage from the
 response.usageMetadata envelope (#4785)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Integrated into release/v3.8.37 — Antigravity response.usageMetadata envelope. Cherry-picked onto release tip (resolved test-tail adjacency with #4754 Ollama block); usage-extractor 23/23 green.
---
 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 f1758b7030..7528960ace 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 e156a989af..84d330e757 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);
+});