diff --git a/changelog.d/fixes/13055-antigravity-thought-token-usage.md b/changelog.d/fixes/13055-antigravity-thought-token-usage.md new file mode 100644 index 0000000000..f19f31a0a5 --- /dev/null +++ b/changelog.d/fixes/13055-antigravity-thought-token-usage.md @@ -0,0 +1 @@ +- **fix(antigravity):** Preserve upstream thought-token usage in normalized completion and reasoning token counts ([#13055](https://github.com/diegosouzapw/OmniRoute/pull/13055)) — thanks @pacocartones diff --git a/open-sse/executors/antigravity/sseCollect.ts b/open-sse/executors/antigravity/sseCollect.ts index 5b7ef3ea85..ee4de7c11b 100644 --- a/open-sse/executors/antigravity/sseCollect.ts +++ b/open-sse/executors/antigravity/sseCollect.ts @@ -18,8 +18,7 @@ export type AntigravityCollectedStream = { // Both run once per SSE data line / per text part (processAntigravitySSEPayload), // so the literals are hoisted to module constants. -const TEXTUAL_TOOL_CALL_RE = - /^[\s\S]*?\[Tool call:\s*([^\]\n]+)\]\s*\nArguments:\s*([\s\S]+?)\s*$/; +const TEXTUAL_TOOL_CALL_RE = /^[\s\S]*?\[Tool call:\s*([^\]\n]+)\]\s*\nArguments:\s*([\s\S]+?)\s*$/; export function stripZeroWidth(value: unknown): unknown { if (typeof value === "string") { @@ -145,10 +144,14 @@ export function processAntigravitySSEPayload( } if (parsed?.response?.usageMetadata) { const um = parsed.response.usageMetadata; + const thoughtsTokens = typeof um.thoughtsTokenCount === "number" ? um.thoughtsTokenCount : 0; collected.usage = { prompt_tokens: um.promptTokenCount || 0, - completion_tokens: um.candidatesTokenCount || 0, + completion_tokens: (um.candidatesTokenCount || 0) + thoughtsTokens, total_tokens: um.totalTokenCount || 0, + ...(thoughtsTokens > 0 + ? { completion_tokens_details: { reasoning_tokens: thoughtsTokens } } + : {}), }; } if (Array.isArray(parsed?.remainingCredits)) { diff --git a/tests/unit/executor-antigravity.test.ts b/tests/unit/executor-antigravity.test.ts index facd6ba68b..befff78723 100644 --- a/tests/unit/executor-antigravity.test.ts +++ b/tests/unit/executor-antigravity.test.ts @@ -36,6 +36,7 @@ type ChatCompletionPayload = { prompt_tokens: number; completion_tokens: number; total_tokens: number; + completion_tokens_details?: { reasoning_tokens: number }; }; }; @@ -482,6 +483,33 @@ test("AntigravityExecutor.collectStreamToResponse turns SSE Gemini chunks into a }); }); +test("AntigravityExecutor.collectStreamToResponse preserves upstream thought token usage", async () => { + const executor = new AntigravityExecutor(); + const response = new Response( + 'data: {"response":{"candidates":[{"content":{"parts":[{"text":"Done"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":5,"candidatesTokenCount":3,"thoughtsTokenCount":7,"totalTokenCount":15}}}\n\n', + { + status: 200, + headers: { "Content-Type": "text/event-stream" }, + } + ); + + const result = await executor.collectStreamToResponse( + response, + "gemini-3.7-pro-high", + "https://example.com", + { Authorization: "Bearer ag-token" }, + { request: {} } + ); + const payload = (await result.response.json()) as ChatCompletionPayload; + + assert.deepEqual(payload.usage, { + prompt_tokens: 5, + completion_tokens: 10, + total_tokens: 15, + completion_tokens_details: { reasoning_tokens: 7 }, + }); +}); + test("AntigravityExecutor.collectStreamToResponse converts textual tool call SSE to structured tool_calls", async () => { const executor = new AntigravityExecutor(); const response = new Response(