diff --git a/open-sse/handlers/chatCore/requestDetail.js b/open-sse/handlers/chatCore/requestDetail.js index b3d24115..376d5f90 100644 --- a/open-sse/handlers/chatCore/requestDetail.js +++ b/open-sse/handlers/chatCore/requestDetail.js @@ -44,13 +44,14 @@ export function extractUsageFromResponse(responseBody) { }; } - // Gemini format - if (responseBody.usageMetadata) { + // Gemini format. Antigravity / gemini-cli wrap the payload in { response: {...} }. + const usageMetadata = responseBody.usageMetadata || responseBody.response?.usageMetadata; + if (usageMetadata) { return { - prompt_tokens: responseBody.usageMetadata.promptTokenCount || 0, - completion_tokens: responseBody.usageMetadata.candidatesTokenCount || 0, - cached_tokens: responseBody.usageMetadata.cachedContentTokenCount || 0, - reasoning_tokens: responseBody.usageMetadata.thoughtsTokenCount || 0 + prompt_tokens: usageMetadata.promptTokenCount || 0, + completion_tokens: usageMetadata.candidatesTokenCount || 0, + cached_tokens: usageMetadata.cachedContentTokenCount || 0, + reasoning_tokens: usageMetadata.thoughtsTokenCount || 0 }; } diff --git a/tests/unit/antigravity-nonstream-usage-3260.test.js b/tests/unit/antigravity-nonstream-usage-3260.test.js new file mode 100644 index 00000000..af7fc9a1 --- /dev/null +++ b/tests/unit/antigravity-nonstream-usage-3260.test.js @@ -0,0 +1,54 @@ +import { describe, expect, it, vi } from "vitest"; + +vi.mock("@/lib/usageDb.js", () => ({ + appendRequestLog: vi.fn(async () => {}), + saveRequestDetail: vi.fn(async () => {}), + saveRequestUsage: vi.fn(async () => {}) +})); + +const { extractUsageFromResponse } = await import("../../open-sse/handlers/chatCore/requestDetail.js"); + +const USAGE_METADATA = { + promptTokenCount: 1234, + candidatesTokenCount: 56, + cachedContentTokenCount: 78, + thoughtsTokenCount: 90, +}; + +const EXPECTED = { + prompt_tokens: 1234, + completion_tokens: 56, + cached_tokens: 78, + reasoning_tokens: 90, +}; + +describe("#3260 non-streaming usage extraction for enveloped Gemini responses", () => { + it("reads usageMetadata out of the antigravity { response } envelope", () => { + expect(extractUsageFromResponse({ response: { usageMetadata: USAGE_METADATA } })).toEqual(EXPECTED); + }); + + it("still reads a top-level usageMetadata", () => { + expect(extractUsageFromResponse({ usageMetadata: USAGE_METADATA })).toEqual(EXPECTED); + }); + + it("prefers the top-level metadata when both are present", () => { + const enveloped = { ...USAGE_METADATA, promptTokenCount: 1 }; + const out = extractUsageFromResponse({ + usageMetadata: USAGE_METADATA, + response: { usageMetadata: enveloped }, + }); + expect(out.prompt_tokens).toBe(1234); + }); + + it("leaves the OpenAI and Claude shapes alone", () => { + expect(extractUsageFromResponse({ usage: { prompt_tokens: 10, completion_tokens: 2 } })) + .toMatchObject({ prompt_tokens: 10, completion_tokens: 2 }); + expect(extractUsageFromResponse({ usage: { input_tokens: 10, output_tokens: 2 } })) + .toMatchObject({ prompt_tokens: 10, completion_tokens: 2 }); + }); + + it("returns null when there is no usage anywhere", () => { + expect(extractUsageFromResponse({ response: { candidates: [] } })).toBeNull(); + expect(extractUsageFromResponse(null)).toBeNull(); + }); +});