fix(sse): read Antigravity usage from the response.usageMetadata envelope

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,
     });
   }
 

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);
+});