feat: meter Gemini thinking tokens and grounding requests (#3178)

* feat: meter Gemini thinking tokens and grounding requests

- Thinking tokens: Extracted from standard completion tokens to ensure they are billed accurately at the correct model-specific rate.
- Grounding requests: Added flat-fee metering for Google Search by tracking grounding_metadata across both streaming and non-streaming responses.
- Pricing updates: Corrected stale rates for Gemini 2.5 Flash output, cached tokens, thinking tokens, and grounding requests.

* fix: correct Gemini 2.x metering rates and harden grounding capture

Pricing corrections (verified against ai.google.dev/gemini-api/docs/pricing):
- gemini-2.5-flash output is $2.50/M, not $1.00/M: restore
  completion_tokens to 250 and bill thinking_tokens at the same output
  rate (250). The previous 100 under-billed output ~60%, and this is the
  provider's default model.
- gemini-2.5-flash cache read is $0.03/M: restore cached_tokens to 3
  (the 7.5 value over-billed).
- Grounding with Google Search is $35 / 1,000 requests for Gemini 2.x
  models and $14 / 1,000 for 3.x. Set grounding_requests to 3_500_000
  for gemini-2.0-flash, gemini-2.5-flash, gemini-2.5-flash-lite and
  gemini-2.5-pro; 3.x models keep 1_400_000.

Streaming robustness:
- In create_chat_stream_handler, don't let a later extra_content chunk
  without grounding_metadata overwrite an earlier one that carried it,
  so grounding requests are still metered.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: Daniel Salazar <daniel.salazar@puter.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

src/backend/drivers/ai-chat/providers/gemini/GeminiChatProvider.test.ts

@@ -355,6 +355,8 @@ describe('GeminiChatProvider.complete non-stream output', () => {
             prompt_tokens: 90,
             completion_tokens: 50,
             cached_tokens: 10,
+            thinking_tokens: 0,
+            grounding_requests: 0,
         });
 
         // gemini-2.5-flash costs: prompt=30, completion=250, cached=3.
@@ -367,11 +369,15 @@ describe('GeminiChatProvider.complete non-stream output', () => {
             prompt_tokens: 90,
             completion_tokens: 50,
             cached_tokens: 10,
+            thinking_tokens: 0,
+            grounding_requests: 0,
         });
         expect(overrides).toEqual({
             prompt_tokens: 90 * Number(flash.costs.prompt_tokens),
             completion_tokens: 50 * Number(flash.costs.completion_tokens),
             cached_tokens: 10 * Number(flash.costs.cached_tokens ?? 0),
+            thinking_tokens: 0,
+            grounding_requests: 0,
         });
     });
 
@@ -446,6 +452,8 @@ describe('GeminiChatProvider.complete streaming', () => {
             prompt_tokens: 3, // 4 - 1 cached
             completion_tokens: 2,
             cached_tokens: 1,
+            thinking_tokens: 0,
+            grounding_requests: 0,
         });
 
         const flash = GEMINI_MODELS.find((m) => m.id === 'gemini-2.5-flash')!;
@@ -456,10 +464,192 @@ describe('GeminiChatProvider.complete streaming', () => {
             prompt_tokens: 3 * Number(flash.costs.prompt_tokens),
             completion_tokens: 2 * Number(flash.costs.completion_tokens),
             cached_tokens: 1 * Number(flash.costs.cached_tokens ?? 0),
+            thinking_tokens: 0,
+            grounding_requests: 0,
         });
     });
 });
 
+// ── Thinking-token metering ─────────────────────────────────────────
+
+describe('GeminiChatProvider.complete thinking token metering', () => {
+    it('splits thinking tokens out of completion_tokens and bills each at its own rate', async () => {
+        const { provider } = makeProvider();
+        createMock.mockResolvedValueOnce({
+            choices: [
+                {
+                    message: { content: 'result', role: 'assistant' },
+                    finish_reason: 'stop',
+                },
+            ],
+            usage: {
+                prompt_tokens: 100,
+                completion_tokens: 60,
+                completion_tokens_details: { reasoning_tokens: 20 },
+                prompt_tokens_details: { cached_tokens: 10 },
+            },
+        });
+
+        const result = await withTestActor(() =>
+            provider.complete({
+                model: 'gemini-2.5-flash',
+                messages: [{ role: 'user', content: 'think hard' }],
+            }),
+        );
+
+        expect((result as { usage: unknown }).usage).toEqual({
+            prompt_tokens: 90,       // 100 - 10 cached
+            completion_tokens: 40,   // 60 - 20 thinking
+            cached_tokens: 10,
+            thinking_tokens: 20,
+            grounding_requests: 0,
+        });
+
+        const flash = GEMINI_MODELS.find((m) => m.id === 'gemini-2.5-flash')!;
+        expect(recordSpy).toHaveBeenCalledTimes(1);
+        const [usage, , , overrides] = recordSpy.mock.calls[0]!;
+        expect(usage.thinking_tokens).toBe(20);
+        expect(usage.completion_tokens).toBe(40);
+        expect(overrides).toMatchObject({
+            thinking_tokens: 20 * Number(flash.costs.thinking_tokens),
+            completion_tokens: 40 * Number(flash.costs.completion_tokens),
+        });
+    });
+
+    it('sets thinking_tokens to 0 when completion_tokens_details is absent', async () => {
+        const { provider } = makeProvider();
+        createMock.mockResolvedValueOnce({
+            choices: [
+                {
+                    message: { content: 'ok', role: 'assistant' },
+                    finish_reason: 'stop',
+                },
+            ],
+            usage: { prompt_tokens: 5, completion_tokens: 3 },
+        });
+
+        await withTestActor(() =>
+            provider.complete({
+                model: 'gemini-2.5-flash',
+                messages: [{ role: 'user', content: 'hi' }],
+            }),
+        );
+
+        const [usage] = recordSpy.mock.calls[0]!;
+        expect(usage.thinking_tokens).toBe(0);
+        expect(usage.completion_tokens).toBe(3);
+    });
+});
+
+// ── Grounding-request metering ───────────────────────────────────────
+
+describe('GeminiChatProvider.complete grounding request metering', () => {
+    it('charges one grounding request when non-stream response contains grounding_metadata', async () => {
+        const { provider } = makeProvider();
+        createMock.mockResolvedValueOnce({
+            choices: [
+                {
+                    message: {
+                        content: 'result',
+                        role: 'assistant',
+                        extra_content: {
+                            grounding_metadata: { web_search_queries: ['foo'] },
+                        },
+                    },
+                    finish_reason: 'stop',
+                },
+            ],
+            usage: { prompt_tokens: 10, completion_tokens: 5 },
+        });
+
+        await withTestActor(() =>
+            provider.complete({
+                model: 'gemini-2.5-flash',
+                messages: [{ role: 'user', content: 'search for foo' }],
+            }),
+        );
+
+        const flash = GEMINI_MODELS.find((m) => m.id === 'gemini-2.5-flash')!;
+        const [usage, , , overrides] = recordSpy.mock.calls[0]!;
+        expect(usage.grounding_requests).toBe(1);
+        expect(overrides!.grounding_requests).toBe(
+            1 * Number(flash.costs.grounding_requests),
+        );
+    });
+
+    it('does not charge a grounding request when no grounding_metadata is present', async () => {
+        const { provider } = makeProvider();
+        createMock.mockResolvedValueOnce({
+            choices: [
+                {
+                    message: { content: 'ok', role: 'assistant' },
+                    finish_reason: 'stop',
+                },
+            ],
+            usage: { prompt_tokens: 5, completion_tokens: 3 },
+        });
+
+        await withTestActor(() =>
+            provider.complete({
+                model: 'gemini-2.5-flash',
+                messages: [{ role: 'user', content: 'hi' }],
+            }),
+        );
+
+        const [usage] = recordSpy.mock.calls[0]!;
+        expect(usage.grounding_requests).toBe(0);
+    });
+
+    it('charges one grounding request when streaming response contains grounding extra_content', async () => {
+        const { provider } = makeProvider();
+        createMock.mockReturnValueOnce(
+            asAsyncIterable([
+                { choices: [{ delta: { content: 'ans' } }] },
+                {
+                    choices: [
+                        {
+                            delta: {
+                                extra_content: {
+                                    grounding_metadata: {
+                                        web_search_queries: ['bar'],
+                                    },
+                                },
+                            },
+                        },
+                    ],
+                },
+                {
+                    choices: [{ delta: {} }],
+                    usage: { prompt_tokens: 8, completion_tokens: 4 },
+                },
+            ]),
+        );
+
+        const result = await withTestActor(() =>
+            provider.complete({
+                model: 'gemini-2.5-flash',
+                messages: [{ role: 'user', content: 'search bar' }],
+                stream: true,
+            }),
+        );
+
+        const harness = makeCapturingChatStream();
+        await (
+            result as {
+                init_chat_stream: (p: { chatStream: unknown }) => Promise<void>;
+            }
+        ).init_chat_stream({ chatStream: harness.chatStream });
+
+        const flash = GEMINI_MODELS.find((m) => m.id === 'gemini-2.5-flash')!;
+        expect(recordSpy).toHaveBeenCalledTimes(1);
+        const [usage, , , overrides] = recordSpy.mock.calls[0]!;
+        expect(usage.grounding_requests).toBe(1);
+        expect(overrides!.grounding_requests).toBe(
+            1 * Number(flash.costs.grounding_requests),
+        );
+    });
+});
+
 // ── Error mapping ───────────────────────────────────────────────────
 
 describe('GeminiChatProvider.complete error mapping', () => {

src/backend/drivers/ai-chat/providers/gemini/GeminiChatProvider.ts

@@ -102,24 +102,52 @@ export class GeminiChatProvider implements IChatProvider {
         }
 
         return handle_completion_output({
-            usage_calculator: ({ usage }) => {
+            usage_calculator: (args) => {
+                // Cast to access Gemini-specific extras passed alongside usage:
+                // - choices: non-stream grounding metadata lives in choices[0].message.extra_content
+                // - extra_content: streaming grounding metadata accumulated by the stream handler
+                const { usage, choices, extra_content } = args as {
+                    usage: typeof args.usage;
+                    choices?: Array<{
+                        message?: {
+                            extra_content?: { grounding_metadata?: unknown };
+                        };
+                    }>;
+                    extra_content?: { grounding_metadata?: unknown };
+                };
+
+                const cached_tokens =
+                    usage?.prompt_tokens_details?.cached_tokens ?? 0;
+
+                // Thinking tokens are a subset of completion_tokens billed at a different rate
+                const thinking_tokens =
+                    usage?.completion_tokens_details?.reasoning_tokens ?? 0;
+
                 const trackedUsage = {
-                    prompt_tokens:
-                        (usage.prompt_tokens ?? 0) -
-                        (usage.prompt_tokens_details?.cached_tokens ?? 0),
-                    completion_tokens: usage.completion_tokens ?? 0,
-                    cached_tokens:
-                        usage.prompt_tokens_details?.cached_tokens ?? 0,
+                    prompt_tokens: (usage?.prompt_tokens ?? 0) - cached_tokens,
+                    completion_tokens: Math.max(
+                        0,
+                        (usage?.completion_tokens ?? 0) - thinking_tokens,
+                    ),
+                    cached_tokens,
+                    thinking_tokens,
+                    // Grounding search is a per-request fee not reflected in token counts
+                    grounding_requests:
+                        (choices?.[0]?.message?.extra_content
+                            ?.grounding_metadata ??
+                        extra_content?.grounding_metadata)
+                            ? 1
+                            : 0,
                 };
 
                 const costsOverrideFromModel = Object.fromEntries(
                     Object.entries(trackedUsage).map(([k, v]) => {
-                        return [k, v * modelUsed.costs[k]];
+                        return [k, v * (modelUsed.costs[k] ?? 0)];
                     }),
                 );
                 this.meteringService.utilRecordUsageObject(
                     trackedUsage,
-                    actor,
+                    actor!,
                     `gemini:${modelUsed?.id}`,
                     costsOverrideFromModel,
                 );

src/backend/drivers/ai-chat/providers/gemini/models.ts

@@ -43,7 +43,10 @@ export const GEMINI_MODELS: IChatModel[] = [
             tokens: 1_000_000,
             prompt_tokens: 150,
             completion_tokens: 900,
+            thinking_tokens: 900,
             cached_tokens: 15,
+            // Gemini 3.x grounding is $14 / 1,000 requests
+            grounding_requests: 1_400_000,
         },
     },
     {
@@ -68,6 +71,8 @@ export const GEMINI_MODELS: IChatModel[] = [
             prompt_tokens: 10,
             completion_tokens: 40,
             cached_tokens: 3,
+            // Gemini 2.x grounding is $35 / 1,000 requests
+            grounding_requests: 3_500_000,
         },
         max_tokens: 8192,
     },
@@ -115,8 +120,13 @@ export const GEMINI_MODELS: IChatModel[] = [
         costs: {
             tokens: 1_000_000,
             prompt_tokens: 30,
+            // Output is $2.50/M; thinking tokens bill at the same output rate
             completion_tokens: 250,
+            thinking_tokens: 250,
+            // Cache read is $0.03/M (10% of input)
             cached_tokens: 3,
+            // Gemini 2.x grounding is $35 / 1,000 requests
+            grounding_requests: 3_500_000,
         },
         max_tokens: 65536,
     },
@@ -141,7 +151,10 @@ export const GEMINI_MODELS: IChatModel[] = [
             tokens: 1_000_000,
             prompt_tokens: 10,
             completion_tokens: 40,
+            thinking_tokens: 40,
             cached_tokens: 1,
+            // Gemini 2.x grounding is $35 / 1,000 requests
+            grounding_requests: 3_500_000,
         },
         max_tokens: 65536,
     },
@@ -166,7 +179,10 @@ export const GEMINI_MODELS: IChatModel[] = [
             tokens: 1_000_000,
             prompt_tokens: 125,
             completion_tokens: 1000,
-            cached_tokens: 13,
+            thinking_tokens: 1000,
+            cached_tokens: 31,
+            // Gemini 2.x grounding is $35 / 1,000 requests
+            grounding_requests: 3_500_000,
         },
         max_tokens: 200_000,
     },
@@ -191,7 +207,9 @@ export const GEMINI_MODELS: IChatModel[] = [
             tokens: 1_000_000,
             prompt_tokens: 200,
             completion_tokens: 1200,
+            thinking_tokens: 1200,
             cached_tokens: 20,
+            grounding_requests: 1_400_000,
         },
         max_tokens: 65536,
     },
@@ -216,7 +234,9 @@ export const GEMINI_MODELS: IChatModel[] = [
             tokens: 1_000_000,
             prompt_tokens: 50,
             completion_tokens: 300,
+            thinking_tokens: 300,
             cached_tokens: 5,
+            grounding_requests: 1_400_000,
         },
         max_tokens: 65536,
     },
@@ -245,7 +265,9 @@ export const GEMINI_MODELS: IChatModel[] = [
             tokens: 1_000_000,
             prompt_tokens: 25,
             completion_tokens: 150,
+            thinking_tokens: 150,
             cached_tokens: 2.5,
+            grounding_requests: 1_400_000,
         },
         max_tokens: 65536,
     },

src/backend/drivers/ai-chat/utils/OpenAIUtil.js

@@ -253,6 +253,7 @@ export const create_chat_stream_handler =
         const tool_call_blocks = [];
 
         let last_usage = null;
+        let last_extra_content = null;
         for await (let chunk of completion) {
             chunk = deviations.chunk_but_like_actually(chunk);
             const chunk_usage = deviations.index_usage_from_stream_chunk(chunk);
@@ -285,6 +286,14 @@ export const create_chat_stream_handler =
                 // Apps have to choose to handle extra_content themselves, it doesn't seem like theres a way we can do it in a backwards
                 // compatible fashion since most streaming apps will handle chat history by continuously updating content themselves
                 // This doesn't present us a chance to add in an extra object for gemini's chat continuing features
+                // Don't let a later extra_content chunk without grounding clobber an
+                // earlier one that carried grounding_metadata (used for metering).
+                if (
+                    choice.delta.extra_content.grounding_metadata ||
+                    !last_extra_content?.grounding_metadata
+                ) {
+                    last_extra_content = choice.delta.extra_content;
+                }
                 textblock.addExtraContent(choice.delta.extra_content);
             }
 
@@ -315,7 +324,10 @@ export const create_chat_stream_handler =
         }
 
         // TODO DS: this is a bit too abstracted... this is basically just doing the metering now
-        const usage = usage_calculator({ usage: last_usage });
+        const usage = usage_calculator({
+            usage: last_usage,
+            extra_content: last_extra_content,
+        });
 
         if (mode === 'text') textblock.end();
         if (mode === 'tool') toolblock.end();