MediaMetz/gemini-cli
1
0
Fork 0
mirror of https://github.com/google-gemini/gemini-cli.git synced 2026-03-15 16:41:11 -07:00
Code Issues Packages Projects Releases Wiki Activity

Feature/quota visibility 16795 (#18203)

Browse Source
This commit is contained in:
Spencer
2026-02-09 21:53:10 -05:00
committed by GitHub
parent 0a3ecf3a75
commit 6dae3a5402
43 changed files with 1315 additions and 317 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
packages/core/src/code_assist/codeAssist.test.ts
View File

@@ -1,6 +1,6 @@
/**
* @license
* Copyright 2025 Google LLC
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
@@ -43,6 +43,7 @@ describe('codeAssist', () => {
const mockUserData = {
projectId: 'test-project',
userTier: UserTierId.FREE,
userTierName: 'free-tier-name',
};
it('should create a server for LOGIN_WITH_GOOGLE', async () => {
@@ -70,7 +71,7 @@ describe('codeAssist', () => {
httpOptions,
'session-123',
'free-tier',
undefined,
'free-tier-name',
);
expect(generator).toBeInstanceOf(MockedCodeAssistServer);
});
@@ -99,7 +100,7 @@ describe('codeAssist', () => {
httpOptions,
undefined, // No session ID
'free-tier',
undefined,
'free-tier-name',
);
expect(generator).toBeInstanceOf(MockedCodeAssistServer);
});

272
packages/core/src/config/config.test.ts
View File

@@ -1,6 +1,6 @@
/**
* @license
* Copyright 2025 Google LLC
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
@@ -38,8 +38,9 @@ import { RipgrepFallbackEvent } from '../telemetry/types.js';
import { ToolRegistry } from '../tools/tool-registry.js';
import { ACTIVATE_SKILL_TOOL_NAME } from '../tools/tool-names.js';
import type { SkillDefinition } from '../skills/skillLoader.js';
import type { McpClientManager } from '../tools/mcp-client-manager.js';
import { DEFAULT_MODEL_CONFIGS } from './defaultModelConfigs.js';
import { DEFAULT_GEMINI_MODEL, PREVIEW_GEMINI_MODEL } from './models.js';
import { DEFAULT_GEMINI_MODEL } from './models.js';
vi.mock('fs', async (importOriginal) => {
const actual = await importOriginal<typeof import('fs')>();
@@ -169,6 +170,7 @@ const mockCoreEvents = vi.hoisted(() => ({
emitFeedback: vi.fn(),
emitModelChanged: vi.fn(),
emitConsoleLog: vi.fn(),
emitQuotaChanged: vi.fn(),
on: vi.fn(),
}));
@@ -203,7 +205,9 @@ import { getCodeAssistServer } from '../code_assist/codeAssist.js';
import { getExperiments } from '../code_assist/experiments/experiments.js';
import type { CodeAssistServer } from '../code_assist/server.js';
import { ContextManager } from '../services/contextManager.js';
import { UserTierId } from 'src/code_assist/types.js';
import { UserTierId } from '../code_assist/types.js';
import type { ModelConfigService } from '../services/modelConfigService.js';
import type { ModelConfigServiceConfig } from '../services/modelConfigService.js';
import { ExitPlanModeTool } from '../tools/exit-plan-mode.js';
import { EnterPlanModeTool } from '../tools/enter-plan-mode.js';
@@ -253,7 +257,7 @@ describe('Server Config (config.ts)', () => {
describe('initialize', () => {
it('should throw an error if checkpointing is enabled and GitService fails', async () => {
const gitError = new Error('Git is not installed');
(GitService.prototype.initialize as Mock).mockRejectedValue(gitError);
vi.mocked(GitService.prototype.initialize).mockRejectedValue(gitError);
const config = new Config({
...baseParams,
@@ -265,7 +269,7 @@ describe('Server Config (config.ts)', () => {
it('should not throw an error if checkpointing is disabled and GitService fails', async () => {
const gitError = new Error('Git is not installed');
(GitService.prototype.initialize as Mock).mockRejectedValue(gitError);
vi.mocked(GitService.prototype.initialize).mockRejectedValue(gitError);
const config = new Config({
...baseParams,
@@ -299,13 +303,16 @@ describe('Server Config (config.ts)', () => {
);
let mcpStarted = false;
(McpClientManager as unknown as Mock).mockImplementation(() => ({
startConfiguredMcpServers: vi.fn().mockImplementation(async () => {
await new Promise((resolve) => setTimeout(resolve, 50));
mcpStarted = true;
}),
getMcpInstructions: vi.fn(),
}));
vi.mocked(McpClientManager).mockImplementation(
() =>
({
startConfiguredMcpServers: vi.fn().mockImplementation(async () => {
await new Promise((resolve) => setTimeout(resolve, 50));
mcpStarted = true;
}),
getMcpInstructions: vi.fn(),
}) as Partial<McpClientManager> as McpClientManager,
);
await config.initialize();
@@ -329,13 +336,16 @@ describe('Server Config (config.ts)', () => {
resolveMcp = resolve;
});
(McpClientManager as unknown as Mock).mockImplementation(() => ({
startConfiguredMcpServers: vi.fn().mockImplementation(async () => {
await mcpPromise;
mcpStarted = true;
}),
getMcpInstructions: vi.fn(),
}));
(McpClientManager as unknown as Mock).mockImplementation(
() =>
({
startConfiguredMcpServers: vi.fn().mockImplementation(async () => {
await mcpPromise;
mcpStarted = true;
}),
getMcpInstructions: vi.fn(),
}) as Partial<McpClientManager> as McpClientManager,
);
await config.initialize();
@@ -459,7 +469,9 @@ describe('Server Config (config.ts)', () => {
vi.mocked(createContentGeneratorConfig).mockImplementation(
async (_: Config, authType: AuthType | undefined) =>
({ authType }) as unknown as ContentGeneratorConfig,
({
authType,
}) as Partial<ContentGeneratorConfig> as ContentGeneratorConfig,
);
await config.refreshAuth(AuthType.USE_GEMINI);
@@ -472,7 +484,9 @@ describe('Server Config (config.ts)', () => {
vi.mocked(createContentGeneratorConfig).mockImplementation(
async (_: Config, authType: AuthType | undefined) =>
({ authType }) as unknown as ContentGeneratorConfig,
({
authType,
}) as Partial<ContentGeneratorConfig> as ContentGeneratorConfig,
);
await config.refreshAuth(AuthType.USE_GEMINI);
@@ -489,7 +503,9 @@ describe('Server Config (config.ts)', () => {
vi.mocked(createContentGeneratorConfig).mockImplementation(
async (_: Config, authType: AuthType | undefined) =>
({ authType }) as unknown as ContentGeneratorConfig,
({
authType,
}) as Partial<ContentGeneratorConfig> as ContentGeneratorConfig,
);
await config.refreshAuth(AuthType.USE_GEMINI);
@@ -506,7 +522,9 @@ describe('Server Config (config.ts)', () => {
vi.mocked(createContentGeneratorConfig).mockImplementation(
async (_: Config, authType: AuthType | undefined) =>
({ authType }) as unknown as ContentGeneratorConfig,
({
authType,
}) as Partial<ContentGeneratorConfig> as ContentGeneratorConfig,
);
await config.refreshAuth(AuthType.USE_VERTEX_AI);
@@ -1268,7 +1286,7 @@ describe('setApprovalMode with folder trust', () => {
getTool: vi.fn().mockReturnValue(undefined),
unregisterTool: vi.fn(),
registerTool: vi.fn(),
} as unknown as ReturnType<Config['getToolRegistry']>);
} as Partial<ToolRegistry> as ToolRegistry);
const updateSpy = vi.spyOn(config, 'updateSystemInstructionIfInitialized');
config.setApprovalMode(ApprovalMode.PLAN);
@@ -1286,7 +1304,7 @@ describe('setApprovalMode with folder trust', () => {
getTool: vi.fn().mockReturnValue(undefined),
unregisterTool: vi.fn(),
registerTool: vi.fn(),
} as unknown as ReturnType<Config['getToolRegistry']>);
} as Partial<ToolRegistry> as ToolRegistry);
const updateSpy = vi.spyOn(config, 'updateSystemInstructionIfInitialized');
config.setApprovalMode(ApprovalMode.DEFAULT);
@@ -1310,11 +1328,11 @@ describe('setApprovalMode with folder trust', () => {
});
it('should register RipGrepTool when useRipgrep is true and it is available', async () => {
(canUseRipgrep as Mock).mockResolvedValue(true);
vi.mocked(canUseRipgrep).mockResolvedValue(true);
const config = new Config({ ...baseParams, useRipgrep: true });
await config.initialize();
const calls = (ToolRegistry.prototype.registerTool as Mock).mock.calls;
const calls = vi.mocked(ToolRegistry.prototype.registerTool).mock.calls;
const wasRipGrepRegistered = calls.some(
(call) => call[0] instanceof vi.mocked(RipGrepTool),
);
@@ -1328,11 +1346,11 @@ describe('setApprovalMode with folder trust', () => {
});
it('should register GrepTool as a fallback when useRipgrep is true but it is not available', async () => {
(canUseRipgrep as Mock).mockResolvedValue(false);
vi.mocked(canUseRipgrep).mockResolvedValue(false);
const config = new Config({ ...baseParams, useRipgrep: true });
await config.initialize();
const calls = (ToolRegistry.prototype.registerTool as Mock).mock.calls;
const calls = vi.mocked(ToolRegistry.prototype.registerTool).mock.calls;
const wasRipGrepRegistered = calls.some(
(call) => call[0] instanceof vi.mocked(RipGrepTool),
);
@@ -1346,17 +1364,17 @@ describe('setApprovalMode with folder trust', () => {
config,
expect.any(RipgrepFallbackEvent),
);
const event = (logRipgrepFallback as Mock).mock.calls[0][1];
const event = vi.mocked(logRipgrepFallback).mock.calls[0][1];
expect(event.error).toBeUndefined();
});
it('should register GrepTool as a fallback when canUseRipgrep throws an error', async () => {
const error = new Error('ripGrep check failed');
(canUseRipgrep as Mock).mockRejectedValue(error);
vi.mocked(canUseRipgrep).mockRejectedValue(error);
const config = new Config({ ...baseParams, useRipgrep: true });
await config.initialize();
const calls = (ToolRegistry.prototype.registerTool as Mock).mock.calls;
const calls = vi.mocked(ToolRegistry.prototype.registerTool).mock.calls;
const wasRipGrepRegistered = calls.some(
(call) => call[0] instanceof vi.mocked(RipGrepTool),
);
@@ -1370,7 +1388,7 @@ describe('setApprovalMode with folder trust', () => {
config,
expect.any(RipgrepFallbackEvent),
);
const event = (logRipgrepFallback as Mock).mock.calls[0][1];
const event = vi.mocked(logRipgrepFallback).mock.calls[0][1];
expect(event.error).toBe(String(error));
});
@@ -1378,7 +1396,7 @@ describe('setApprovalMode with folder trust', () => {
const config = new Config({ ...baseParams, useRipgrep: false });
await config.initialize();
const calls = (ToolRegistry.prototype.registerTool as Mock).mock.calls;
const calls = vi.mocked(ToolRegistry.prototype.registerTool).mock.calls;
const wasRipGrepRegistered = calls.some(
(call) => call[0] instanceof vi.mocked(RipGrepTool),
);
@@ -1526,8 +1544,11 @@ describe('Generation Config Merging (HACK)', () => {
};
const config = new Config(params);
// eslint-disable-next-line @typescript-eslint/no-explicit-any
const serviceConfig = (config.modelConfigService as any).config;
const serviceConfig = (
config.modelConfigService as Partial<ModelConfigService> as {
config: ModelConfigServiceConfig;
}
).config;
// Assert that the default aliases are present
expect(serviceConfig.aliases).toEqual(DEFAULT_MODEL_CONFIGS.aliases);
@@ -1550,8 +1571,11 @@ describe('Generation Config Merging (HACK)', () => {
};
const config = new Config(params);
// eslint-disable-next-line @typescript-eslint/no-explicit-any
const serviceConfig = (config.modelConfigService as any).config;
const serviceConfig = (
config.modelConfigService as Partial<ModelConfigService> as {
config: ModelConfigServiceConfig;
}
).config;
// Assert that the user's aliases are present
expect(serviceConfig.aliases).toEqual(userAliases);
@@ -1574,8 +1598,11 @@ describe('Generation Config Merging (HACK)', () => {
};
const config = new Config(params);
// eslint-disable-next-line @typescript-eslint/no-explicit-any
const serviceConfig = (config.modelConfigService as any).config;
const serviceConfig = (
config.modelConfigService as Partial<ModelConfigService> as {
config: ModelConfigServiceConfig;
}
).config;
// Assert that the user's aliases are used, not the defaults
expect(serviceConfig.aliases).toEqual(userAliases);
@@ -1585,8 +1612,11 @@ describe('Generation Config Merging (HACK)', () => {
const params: ConfigParameters = { ...baseParams };
const config = new Config(params);
// eslint-disable-next-line @typescript-eslint/no-explicit-any
const serviceConfig = (config.modelConfigService as any).config;
const serviceConfig = (
config.modelConfigService as Partial<ModelConfigService> as {
config: ModelConfigServiceConfig;
}
).config;
// Assert that the full default config is used
expect(serviceConfig).toEqual(DEFAULT_MODEL_CONFIGS);
@@ -1942,8 +1972,10 @@ describe('Hooks configuration', () => {
describe('Config Quota & Preview Model Access', () => {
let config: Config;
// eslint-disable-next-line @typescript-eslint/no-explicit-any
let mockCodeAssistServer: any;
let mockCodeAssistServer: {
projectId: string;
retrieveUserQuota: Mock;
};
const baseParams: ConfigParameters = {
cwd: '/tmp',
@@ -1965,14 +1997,22 @@ describe('Config Quota & Preview Model Access', () => {
projectId: 'test-project',
retrieveUserQuota: vi.fn(),
};
vi.mocked(getCodeAssistServer).mockReturnValue(mockCodeAssistServer);
vi.mocked(getCodeAssistServer).mockReturnValue(
mockCodeAssistServer as Partial<CodeAssistServer> as CodeAssistServer,
);
config = new Config(baseParams);
});
describe('refreshUserQuota', () => {
it('should update hasAccessToPreviewModel to true if quota includes preview model', async () => {
mockCodeAssistServer.retrieveUserQuota.mockResolvedValue({
buckets: [{ modelId: PREVIEW_GEMINI_MODEL }],
buckets: [
{
modelId: 'gemini-3-pro-preview',
remainingAmount: '100',
remainingFraction: 1.0,
},
],
});
await config.refreshUserQuota();
@@ -1981,13 +2021,82 @@ describe('Config Quota & Preview Model Access', () => {
it('should update hasAccessToPreviewModel to false if quota does not include preview model', async () => {
mockCodeAssistServer.retrieveUserQuota.mockResolvedValue({
buckets: [{ modelId: 'some-other-model' }],
buckets: [
{
modelId: 'some-other-model',
remainingAmount: '10',
remainingFraction: 0.1,
},
],
});
await config.refreshUserQuota();
expect(config.getHasAccessToPreviewModel()).toBe(false);
});
it('should calculate pooled quota correctly for auto models', async () => {
mockCodeAssistServer.retrieveUserQuota.mockResolvedValue({
buckets: [
{
modelId: 'gemini-2.5-pro',
remainingAmount: '10',
remainingFraction: 0.2,
},
{
modelId: 'gemini-2.5-flash',
remainingAmount: '80',
remainingFraction: 0.8,
},
],
});
config.setModel('auto-gemini-2.5');
await config.refreshUserQuota();
const pooled = (
config as Partial<Config> as {
getPooledQuota: () => {
remaining?: number;
limit?: number;
resetTime?: string;
};
}
).getPooledQuota();
// Pro: 10 / 0.2 = 50 total.
// Flash: 80 / 0.8 = 100 total.
// Pooled: (10 + 80) / (50 + 100) = 90 / 150 = 0.6
expect(pooled?.remaining).toBe(90);
expect(pooled?.limit).toBe(150);
expect((pooled?.remaining ?? 0) / (pooled?.limit ?? 1)).toBeCloseTo(0.6);
});
it('should return undefined pooled quota for non-auto models', async () => {
mockCodeAssistServer.retrieveUserQuota.mockResolvedValue({
buckets: [
{
modelId: 'gemini-2.5-pro',
remainingAmount: '10',
remainingFraction: 0.2,
},
],
});
config.setModel('gemini-2.5-pro');
await config.refreshUserQuota();
expect(
(
config as Partial<Config> as {
getPooledQuota: () => {
remaining?: number;
limit?: number;
resetTime?: string;
};
}
).getPooledQuota(),
).toEqual({});
});
it('should update hasAccessToPreviewModel to false if buckets are undefined', async () => {
mockCodeAssistServer.retrieveUserQuota.mockResolvedValue({});
@@ -2013,6 +2122,73 @@ describe('Config Quota & Preview Model Access', () => {
});
});
describe('refreshUserQuotaIfStale', () => {
beforeEach(() => {
vi.useFakeTimers();
vi.setSystemTime(new Date('2025-01-01T00:00:00Z'));
});
afterEach(() => {
vi.useRealTimers();
});
it('should refresh quota if stale', async () => {
mockCodeAssistServer.retrieveUserQuota.mockResolvedValue({
buckets: [],
});
// First call to initialize lastQuotaFetchTime
await config.refreshUserQuota();
expect(mockCodeAssistServer.retrieveUserQuota).toHaveBeenCalledTimes(1);
// Advance time by 31 seconds (default TTL is 30s)
vi.setSystemTime(Date.now() + 31_000);
await config.refreshUserQuotaIfStale();
expect(mockCodeAssistServer.retrieveUserQuota).toHaveBeenCalledTimes(2);
});
it('should not refresh quota if fresh', async () => {
mockCodeAssistServer.retrieveUserQuota.mockResolvedValue({
buckets: [],
});
// First call
await config.refreshUserQuota();
expect(mockCodeAssistServer.retrieveUserQuota).toHaveBeenCalledTimes(1);
// Advance time by only 10 seconds
vi.setSystemTime(Date.now() + 10_000);
await config.refreshUserQuotaIfStale();
expect(mockCodeAssistServer.retrieveUserQuota).toHaveBeenCalledTimes(1);
});
it('should respect custom staleMs', async () => {
mockCodeAssistServer.retrieveUserQuota.mockResolvedValue({
buckets: [],
});
// First call
await config.refreshUserQuota();
expect(mockCodeAssistServer.retrieveUserQuota).toHaveBeenCalledTimes(1);
// Advance time by 5 seconds
vi.setSystemTime(Date.now() + 5_000);
// Refresh with 2s staleMs -> should refresh
await config.refreshUserQuotaIfStale(2_000);
expect(mockCodeAssistServer.retrieveUserQuota).toHaveBeenCalledTimes(2);
// Advance by another 5 seconds
vi.setSystemTime(Date.now() + 5_000);
// Refresh with 10s staleMs -> should NOT refresh
await config.refreshUserQuotaIfStale(10_000);
expect(mockCodeAssistServer.retrieveUserQuota).toHaveBeenCalledTimes(2);
});
});
describe('getUserTier and getUserTierName', () => {
it('should return undefined if contentGenerator is not initialized', () => {
const config = new Config(baseParams);
@@ -2032,7 +2208,7 @@ describe('Config Quota & Preview Model Access', () => {
vi.mocked(createContentGenerator).mockResolvedValue({
userTier: mockTier,
userTierName: mockTierName,
} as unknown as CodeAssistServer);
} as Partial<CodeAssistServer> as CodeAssistServer);
await config.refreshAuth(AuthType.USE_GEMINI);

269
packages/core/src/config/config.ts
View File

@@ -1,6 +1,6 @@
/**
* @license
* Copyright 2025 Google LLC
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
@@ -53,9 +53,14 @@ import { tokenLimit } from '../core/tokenLimits.js';
import {
DEFAULT_GEMINI_EMBEDDING_MODEL,
DEFAULT_GEMINI_FLASH_MODEL,
DEFAULT_GEMINI_MODEL,
DEFAULT_GEMINI_MODEL_AUTO,
isAutoModel,
isPreviewModel,
PREVIEW_GEMINI_FLASH_MODEL,
PREVIEW_GEMINI_MODEL,
PREVIEW_GEMINI_MODEL_AUTO,
resolveModel,
} from './models.js';
import { shouldAttemptBrowserLaunch } from '../utils/browser.js';
import type { MCPOAuthConfig } from '../mcp/oauth-provider.js';
@@ -87,14 +92,14 @@ import { ContextManager } from '../services/contextManager.js';
import type { GenerateContentParameters } from '@google/genai';
// Re-export OAuth config type
export type { MCPOAuthConfig, AnyToolInvocation };
import type { AnyToolInvocation } from '../tools/tools.js';
export type { MCPOAuthConfig, AnyToolInvocation, AnyDeclarativeTool };
import type { AnyToolInvocation, AnyDeclarativeTool } from '../tools/tools.js';
import { WorkspaceContext } from '../utils/workspaceContext.js';
import { Storage } from './storage.js';
import type { ShellExecutionConfig } from '../services/shellExecutionService.js';
import { FileExclusions } from '../utils/ignorePatterns.js';
import type { EventEmitter } from 'node:events';
import { MessageBus } from '../confirmation-bus/message-bus.js';
import type { EventEmitter } from 'node:events';
import { PolicyEngine } from '../policy/policy-engine.js';
import { ApprovalMode, type PolicyEngineConfig } from '../policy/types.js';
import { HookSystem } from '../hooks/index.js';
@@ -564,6 +569,31 @@ export class Config {
fallbackModelHandler?: FallbackModelHandler;
validationHandler?: ValidationHandler;
private quotaErrorOccurred: boolean = false;
private modelQuotas: Map<
string,
{ remaining: number; limit: number; resetTime?: string }
> = new Map();
private lastRetrievedQuota?: RetrieveUserQuotaResponse;
private lastQuotaFetchTime = 0;
private lastEmittedQuotaRemaining: number | undefined;
private lastEmittedQuotaLimit: number | undefined;
private emitQuotaChangedEvent(): void {
const pooled = this.getPooledQuota();
if (
this.lastEmittedQuotaRemaining !== pooled.remaining ||
this.lastEmittedQuotaLimit !== pooled.limit
) {
this.lastEmittedQuotaRemaining = pooled.remaining;
this.lastEmittedQuotaLimit = pooled.limit;
coreEvents.emitQuotaChanged(
pooled.remaining,
pooled.limit,
pooled.resetTime,
);
}
}
private readonly summarizeToolOutput:
| Record<string, SummarizeToolOutputSettings>
| undefined;
@@ -1206,6 +1236,90 @@ export class Config {
return this.quotaErrorOccurred;
}
setQuota(
remaining: number | undefined,
limit: number | undefined,
modelId?: string,
): void {
const activeModel = modelId ?? this.getActiveModel();
if (remaining !== undefined && limit !== undefined) {
const current = this.modelQuotas.get(activeModel);
if (
!current ||
current.remaining !== remaining ||
current.limit !== limit
) {
this.modelQuotas.set(activeModel, { remaining, limit });
this.emitQuotaChangedEvent();
}
}
}
private getPooledQuota(): {
remaining?: number;
limit?: number;
resetTime?: string;
} {
const model = this.getModel();
if (!isAutoModel(model)) {
return {};
}
const isPreview =
model === PREVIEW_GEMINI_MODEL_AUTO ||
isPreviewModel(this.getActiveModel());
const proModel = isPreview ? PREVIEW_GEMINI_MODEL : DEFAULT_GEMINI_MODEL;
const flashModel = isPreview
? PREVIEW_GEMINI_FLASH_MODEL
: DEFAULT_GEMINI_FLASH_MODEL;
const proQuota = this.modelQuotas.get(proModel);
const flashQuota = this.modelQuotas.get(flashModel);
if (proQuota || flashQuota) {
// For reset time, take the one that is furthest in the future (most conservative)
const resetTime = [proQuota?.resetTime, flashQuota?.resetTime]
.filter((t): t is string => !!t)
.sort()
.reverse()[0];
return {
remaining: (proQuota?.remaining ?? 0) + (flashQuota?.remaining ?? 0),
limit: (proQuota?.limit ?? 0) + (flashQuota?.limit ?? 0),
resetTime,
};
}
return {};
}
getQuotaRemaining(): number | undefined {
const pooled = this.getPooledQuota();
if (pooled.remaining !== undefined) {
return pooled.remaining;
}
const primaryModel = resolveModel(this.getModel());
return this.modelQuotas.get(primaryModel)?.remaining;
}
getQuotaLimit(): number | undefined {
const pooled = this.getPooledQuota();
if (pooled.limit !== undefined) {
return pooled.limit;
}
const primaryModel = resolveModel(this.getModel());
return this.modelQuotas.get(primaryModel)?.limit;
}
getQuotaResetTime(): string | undefined {
const pooled = this.getPooledQuota();
if (pooled.resetTime !== undefined) {
return pooled.resetTime;
}
const primaryModel = resolveModel(this.getModel());
return this.modelQuotas.get(primaryModel)?.resetTime;
}
getEmbeddingModel(): string {
return this.embeddingModel;
}
@@ -1285,6 +1399,35 @@ export class Config {
const quota = await codeAssistServer.retrieveUserQuota({
project: codeAssistServer.projectId,
});
if (quota.buckets) {
this.lastRetrievedQuota = quota;
this.lastQuotaFetchTime = Date.now();
for (const bucket of quota.buckets) {
if (
bucket.modelId &&
bucket.remainingAmount &&
bucket.remainingFraction != null
) {
const remaining = parseInt(bucket.remainingAmount, 10);
const limit =
bucket.remainingFraction > 0
? Math.round(remaining / bucket.remainingFraction)
: (this.modelQuotas.get(bucket.modelId)?.limit ?? 0);
if (!isNaN(remaining) && Number.isFinite(limit) && limit > 0) {
this.modelQuotas.set(bucket.modelId, {
remaining,
limit,
resetTime: bucket.resetTime,
});
}
}
}
this.emitQuotaChangedEvent();
}
const hasAccess =
quota.buckets?.some((b) => b.modelId === PREVIEW_GEMINI_MODEL) ?? false;
this.setHasAccessToPreviewModel(hasAccess);
@@ -1295,6 +1438,41 @@ export class Config {
}
}
async refreshUserQuotaIfStale(
staleMs = 30_000,
): Promise<RetrieveUserQuotaResponse | undefined> {
const now = Date.now();
if (now - this.lastQuotaFetchTime > staleMs) {
return this.refreshUserQuota();
}
return this.lastRetrievedQuota;
}
getLastRetrievedQuota(): RetrieveUserQuotaResponse | undefined {
return this.lastRetrievedQuota;
}
getRemainingQuotaForModel(modelId: string):
| {
remainingAmount?: number;
remainingFraction?: number;
resetTime?: string;
}
| undefined {
const bucket = this.lastRetrievedQuota?.buckets?.find(
(b) => b.modelId === modelId,
);
if (!bucket) return undefined;
return {
remainingAmount: bucket.remainingAmount
? parseInt(bucket.remainingAmount, 10)
: undefined,
remainingFraction: bucket.remainingFraction,
resetTime: bucket.resetTime,
};
}
getCoreTools(): string[] | undefined {
return this.coreTools;
}
@@ -2160,10 +2338,12 @@ export class Config {
const registry = new ToolRegistry(this, this.messageBus);
// helper to create & register core tools that are enabled
// eslint-disable-next-line @typescript-eslint/no-explicit-any
const registerCoreTool = (ToolClass: any, ...args: unknown[]) => {
const className = ToolClass.name;
const toolName = ToolClass.Name || className;
const maybeRegister = (
toolClass: { name: string; Name?: string },
registerFn: () => void,
) => {
const className = toolClass.name;
const toolName = toolClass.Name || className;
const coreTools = this.getCoreTools();
// On some platforms, the className can be minified to _ClassName.
const normalizedClassName = className.replace(/^_+/, '');
@@ -2180,15 +2360,16 @@ export class Config {
}
if (isEnabled) {
// Pass message bus to tools (required for policy engine integration)
const toolArgs = [...args, this.getMessageBus()];
registry.registerTool(new ToolClass(...toolArgs));
registerFn();
}
};
registerCoreTool(LSTool, this);
registerCoreTool(ReadFileTool, this);
maybeRegister(LSTool, () =>
registry.registerTool(new LSTool(this, this.messageBus)),
);
maybeRegister(ReadFileTool, () =>
registry.registerTool(new ReadFileTool(this, this.messageBus)),
);
if (this.getUseRipgrep()) {
let useRipgrep = false;
@@ -2199,30 +2380,60 @@ export class Config {
errorString = String(error);
}
if (useRipgrep) {
registerCoreTool(RipGrepTool, this);
maybeRegister(RipGrepTool, () =>
registry.registerTool(new RipGrepTool(this, this.messageBus)),
);
} else {
logRipgrepFallback(this, new RipgrepFallbackEvent(errorString));
registerCoreTool(GrepTool, this);
maybeRegister(GrepTool, () =>
registry.registerTool(new GrepTool(this, this.messageBus)),
);
}
} else {
registerCoreTool(GrepTool, this);
maybeRegister(GrepTool, () =>
registry.registerTool(new GrepTool(this, this.messageBus)),
);
}
registerCoreTool(GlobTool, this);
registerCoreTool(ActivateSkillTool, this);
registerCoreTool(EditTool, this);
registerCoreTool(WriteFileTool, this);
registerCoreTool(WebFetchTool, this);
registerCoreTool(ShellTool, this);
registerCoreTool(MemoryTool);
registerCoreTool(WebSearchTool, this);
registerCoreTool(AskUserTool);
maybeRegister(GlobTool, () =>
registry.registerTool(new GlobTool(this, this.messageBus)),
);
maybeRegister(ActivateSkillTool, () =>
registry.registerTool(new ActivateSkillTool(this, this.messageBus)),
);
maybeRegister(EditTool, () =>
registry.registerTool(new EditTool(this, this.messageBus)),
);
maybeRegister(WriteFileTool, () =>
registry.registerTool(new WriteFileTool(this, this.messageBus)),
);
maybeRegister(WebFetchTool, () =>
registry.registerTool(new WebFetchTool(this, this.messageBus)),
);
maybeRegister(ShellTool, () =>
registry.registerTool(new ShellTool(this, this.messageBus)),
);
maybeRegister(MemoryTool, () =>
registry.registerTool(new MemoryTool(this.messageBus)),
);
maybeRegister(WebSearchTool, () =>
registry.registerTool(new WebSearchTool(this, this.messageBus)),
);
maybeRegister(AskUserTool, () =>
registry.registerTool(new AskUserTool(this.messageBus)),
);
if (this.getUseWriteTodos()) {
registerCoreTool(WriteTodosTool);
maybeRegister(WriteTodosTool, () =>
registry.registerTool(new WriteTodosTool(this.messageBus)),
);
}
if (this.isPlanEnabled()) {
registerCoreTool(ExitPlanModeTool, this);
registerCoreTool(EnterPlanModeTool, this);
maybeRegister(ExitPlanModeTool, () =>
registry.registerTool(new ExitPlanModeTool(this, this.messageBus)),
);
maybeRegister(EnterPlanModeTool, () =>
registry.registerTool(new EnterPlanModeTool(this, this.messageBus)),
);
}
// Register Subagents as Tools

51
packages/core/src/core/contentGenerator.test.ts
View File

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

3
packages/core/src/core/loggingContentGenerator.test.ts
View File

@@ -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();

9
packages/core/src/core/loggingContentGenerator.ts
View File

@@ -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,

25
packages/core/src/utils/events.ts
View File

@@ -1,6 +1,6 @@
/**
* @license
* Copyright 2025 Google LLC
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
@@ -127,6 +127,15 @@ export interface AgentsDiscoveredPayload {
agents: AgentDefinition[];
}
/**
* Payload for the 'quota-changed' event.
*/
export interface QuotaChangedPayload {
remaining: number | undefined;
limit: number | undefined;
resetTime?: string;
}
export enum CoreEvent {
UserFeedback = 'user-feedback',
ModelChanged = 'model-changed',
@@ -146,6 +155,7 @@ export enum CoreEvent {
AgentsDiscovered = 'agents-discovered',
RequestEditorSelection = 'request-editor-selection',
EditorSelected = 'editor-selected',
QuotaChanged = 'quota-changed',
}
/**
@@ -161,6 +171,7 @@ export interface CoreEvents extends ExtensionEvents {
[CoreEvent.ConsoleLog]: [ConsoleLogPayload];
[CoreEvent.Output]: [OutputPayload];
[CoreEvent.MemoryChanged]: [MemoryChangedPayload];
[CoreEvent.QuotaChanged]: [QuotaChangedPayload];
[CoreEvent.ExternalEditorClosed]: never[];
[CoreEvent.McpClientUpdate]: Array<Map<string, McpClient> | never>;
[CoreEvent.OauthDisplayMessage]: string[];
@@ -311,6 +322,18 @@ export class CoreEventEmitter extends EventEmitter<CoreEvents> {
this._emitOrQueue(CoreEvent.AgentsDiscovered, payload);
}
/**
* Notifies subscribers that the quota has changed.
*/
emitQuotaChanged(
remaining: number | undefined,
limit: number | undefined,
resetTime?: string,
): void {
const payload: QuotaChangedPayload = { remaining, limit, resetTime };
this.emit(CoreEvent.QuotaChanged, payload);
}
/**
* Flushes buffered messages. Call this immediately after primary UI listener
* subscribes.