Feature/quota visibility 16795 (#18203)

packages/core/src/core/contentGenerator.test.ts

@@ -1,6 +1,6 @@
 /**
  * @license
- * Copyright 2025 Google LLC
+ * Copyright 2026 Google LLC
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -139,12 +139,11 @@ describe('createContentGenerator', () => {
     expect(GoogleGenAI).toHaveBeenCalledWith({
       apiKey: 'test-api-key',
       vertexai: undefined,
-      httpOptions: {
-        headers: {
+      httpOptions: expect.objectContaining({
+        headers: expect.objectContaining({
           'User-Agent': expect.stringContaining('GeminiCLI/1.2.3/gemini-pro'),
-          'x-gemini-api-privileged-user-id': expect.any(String),
-        },
-      },
+        }),
+      }),
     });
     expect(generator).toEqual(
       new LoggingContentGenerator(mockGenerator.models, mockConfig),
@@ -209,21 +208,21 @@ describe('createContentGenerator', () => {
     expect(GoogleGenAI).toHaveBeenCalledWith({
       apiKey: 'test-api-key',
       vertexai: undefined,
-      httpOptions: {
+      httpOptions: expect.objectContaining({
         headers: expect.objectContaining({
           'User-Agent': expect.any(String),
           'X-Test-Header': 'test',
           Another: 'value',
         }),
-      },
+      }),
     });
     expect(GoogleGenAI).toHaveBeenCalledWith(
       expect.not.objectContaining({
-        httpOptions: {
+        httpOptions: expect.objectContaining({
           headers: expect.objectContaining({
             Authorization: expect.any(String),
           }),
-        },
+        }),
       }),
     );
   });
@@ -252,12 +251,12 @@ describe('createContentGenerator', () => {
     expect(GoogleGenAI).toHaveBeenCalledWith({
       apiKey: 'test-api-key',
       vertexai: undefined,
-      httpOptions: {
+      httpOptions: expect.objectContaining({
         headers: expect.objectContaining({
           'User-Agent': expect.any(String),
           Authorization: 'Bearer test-api-key',
         }),
-      },
+      }),
     });
   });
 
@@ -285,20 +284,20 @@ describe('createContentGenerator', () => {
     expect(GoogleGenAI).toHaveBeenCalledWith({
       apiKey: 'test-api-key',
       vertexai: undefined,
-      httpOptions: {
+      httpOptions: expect.objectContaining({
         headers: expect.objectContaining({
           'User-Agent': expect.any(String),
         }),
-      },
+      }),
     });
     // Explicitly assert that Authorization header is NOT present
     expect(GoogleGenAI).toHaveBeenCalledWith(
       expect.not.objectContaining({
-        httpOptions: {
+        httpOptions: expect.objectContaining({
           headers: expect.objectContaining({
             Authorization: expect.any(String),
           }),
-        },
+        }),
       }),
     );
   });
@@ -322,11 +321,11 @@ describe('createContentGenerator', () => {
     expect(GoogleGenAI).toHaveBeenCalledWith({
       apiKey: 'test-api-key',
       vertexai: undefined,
-      httpOptions: {
+      httpOptions: expect.objectContaining({
         headers: {
           'User-Agent': expect.any(String),
         },
-      },
+      }),
     });
     expect(generator).toEqual(
       new LoggingContentGenerator(mockGenerator.models, mockConfig),
@@ -357,11 +356,11 @@ describe('createContentGenerator', () => {
     expect(GoogleGenAI).toHaveBeenCalledWith({
       apiKey: 'test-api-key',
       vertexai: undefined,
-      httpOptions: {
+      httpOptions: expect.objectContaining({
         headers: expect.objectContaining({
           'User-Agent': expect.any(String),
         }),
-      },
+      }),
       apiVersion: 'v1',
     });
   });
@@ -389,11 +388,11 @@ describe('createContentGenerator', () => {
     expect(GoogleGenAI).toHaveBeenCalledWith({
       apiKey: 'test-api-key',
       vertexai: undefined,
-      httpOptions: {
+      httpOptions: expect.objectContaining({
         headers: expect.objectContaining({
           'User-Agent': expect.any(String),
         }),
-      },
+      }),
     });
 
     expect(GoogleGenAI).toHaveBeenCalledWith(
@@ -427,11 +426,11 @@ describe('createContentGenerator', () => {
     expect(GoogleGenAI).toHaveBeenCalledWith({
       apiKey: 'test-api-key',
       vertexai: undefined,
-      httpOptions: {
+      httpOptions: expect.objectContaining({
         headers: expect.objectContaining({
           'User-Agent': expect.any(String),
         }),
-      },
+      }),
     });
 
     expect(GoogleGenAI).toHaveBeenCalledWith(
@@ -466,11 +465,11 @@ describe('createContentGenerator', () => {
     expect(GoogleGenAI).toHaveBeenCalledWith({
       apiKey: 'test-api-key',
       vertexai: true,
-      httpOptions: {
+      httpOptions: expect.objectContaining({
         headers: expect.objectContaining({
           'User-Agent': expect.any(String),
         }),
-      },
+      }),
       apiVersion: 'v1alpha',
     });
   });

packages/core/src/core/loggingContentGenerator.test.ts

@@ -1,6 +1,6 @@
 /**
  * @license
- * Copyright 2025 Google LLC
+ * Copyright 2026 Google LLC
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -51,6 +51,7 @@ describe('LoggingContentGenerator', () => {
       getContentGeneratorConfig: vi.fn().mockReturnValue({
         authType: 'API_KEY',
       }),
+      refreshUserQuotaIfStale: vi.fn().mockResolvedValue(undefined),
     } as unknown as Config;
     loggingContentGenerator = new LoggingContentGenerator(wrapped, config);
     vi.useFakeTimers();

packages/core/src/core/loggingContentGenerator.ts

@@ -1,6 +1,6 @@
 /**
  * @license
- * Copyright 2025 Google LLC
+ * Copyright 2026 Google LLC
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -34,6 +34,7 @@ import { CodeAssistServer } from '../code_assist/server.js';
 import { toContents } from '../code_assist/converter.js';
 import { isStructuredError } from '../utils/quotaErrorDetection.js';
 import { runInDevTraceSpan, type SpanMetadata } from '../telemetry/trace.js';
+import { debugLogger } from '../utils/debugLogger.js';
 
 interface StructuredError {
   status: number;
@@ -234,6 +235,9 @@ export class LoggingContentGenerator implements ContentGenerator {
             req.config,
             serverDetails,
           );
+          this.config
+            .refreshUserQuotaIfStale()
+            .catch((e) => debugLogger.debug('quota refresh failed', e));
           return response;
         } catch (error) {
           const durationMs = Date.now() - startTime;
@@ -355,6 +359,9 @@ export class LoggingContentGenerator implements ContentGenerator {
         req.config,
         serverDetails,
       );
+      this.config
+        .refreshUserQuotaIfStale()
+        .catch((e) => debugLogger.debug('quota refresh failed', e));
       spanMetadata.output = {
         streamChunks: responses.map((r) => ({
           content: r.candidates?.[0]?.content ?? null,