From 45faf4d31b10792582cdce5ca693fc5a548d1edc Mon Sep 17 00:00:00 2001 From: Jack Wotherspoon Date: Thu, 12 Mar 2026 14:38:09 +0100 Subject: [PATCH 1/7] fix: register themes on extension load not start (#22148) --- .../cli/src/config/extension-manager.test.ts | 65 ++++++++++++++++++- packages/cli/src/config/extension-manager.ts | 9 ++- packages/cli/src/ui/themes/theme-manager.ts | 11 ++++ 3 files changed, 83 insertions(+), 2 deletions(-) diff --git a/packages/cli/src/config/extension-manager.test.ts b/packages/cli/src/config/extension-manager.test.ts index 5b44c07194..13c1de15fa 100644 --- a/packages/cli/src/config/extension-manager.test.ts +++ b/packages/cli/src/config/extension-manager.test.ts @@ -12,12 +12,13 @@ import { ExtensionManager } from './extension-manager.js'; import { createTestMergedSettings, type MergedSettings } from './settings.js'; import { createExtension } from '../test-utils/createExtension.js'; import { EXTENSIONS_DIRECTORY_NAME } from './extensions/variables.js'; +import { themeManager } from '../ui/themes/theme-manager.js'; import { TrustLevel, loadTrustedFolders, isWorkspaceTrusted, } from './trustedFolders.js'; -import { getRealPath } from '@google/gemini-cli-core'; +import { getRealPath, type CustomTheme } from '@google/gemini-cli-core'; const mockHomedir = vi.hoisted(() => vi.fn(() => '/tmp/mock-home')); @@ -38,6 +39,26 @@ vi.mock('@google/gemini-cli-core', async (importOriginal) => { }; }); +const testTheme: CustomTheme = { + type: 'custom', + name: 'MyTheme', + background: { + primary: '#282828', + diff: { added: '#2b3312', removed: '#341212' }, + }, + text: { + primary: '#ebdbb2', + secondary: '#a89984', + link: '#83a598', + accent: '#d3869b', + }, + status: { + success: '#b8bb26', + warning: '#fabd2f', + error: '#fb4934', + }, +}; + describe('ExtensionManager', () => { let tempHomeDir: string; let tempWorkspaceDir: string; @@ -65,6 +86,7 @@ describe('ExtensionManager', () => { }); afterEach(() => { + themeManager.clearExtensionThemes(); try { fs.rmSync(tempHomeDir, { recursive: true, force: true }); } catch (_e) { @@ -484,4 +506,45 @@ describe('ExtensionManager', () => { ).rejects.toThrow(/already installed/); }); }); + + describe('early theme registration', () => { + it('should register themes with ThemeManager during loadExtensions for active extensions', async () => { + createExtension({ + extensionsDir: userExtensionsDir, + name: 'themed-ext', + version: '1.0.0', + themes: [testTheme], + }); + + await extensionManager.loadExtensions(); + + expect(themeManager.getCustomThemeNames()).toContain( + 'MyTheme (themed-ext)', + ); + }); + + it('should not register themes for inactive extensions', async () => { + createExtension({ + extensionsDir: userExtensionsDir, + name: 'disabled-ext', + version: '1.0.0', + themes: [testTheme], + }); + + // Disable the extension by creating an enablement override + const manager = new ExtensionManager({ + enabledExtensionOverrides: ['none'], + settings: createTestMergedSettings(), + workspaceDir: tempWorkspaceDir, + requestConsent: vi.fn().mockResolvedValue(true), + requestSetting: null, + }); + + await manager.loadExtensions(); + + expect(themeManager.getCustomThemeNames()).not.toContain( + 'MyTheme (disabled-ext)', + ); + }); + }); }); diff --git a/packages/cli/src/config/extension-manager.ts b/packages/cli/src/config/extension-manager.ts index 80c48193e2..68617bcbcd 100644 --- a/packages/cli/src/config/extension-manager.ts +++ b/packages/cli/src/config/extension-manager.ts @@ -564,7 +564,7 @@ Would you like to attempt to install via "git clone" instead?`, protected override async startExtension(extension: GeminiCLIExtension) { await super.startExtension(extension); - if (extension.themes) { + if (extension.themes && !themeManager.hasExtensionThemes(extension.name)) { themeManager.registerExtensionThemes(extension.name, extension.themes); } } @@ -624,6 +624,13 @@ Would you like to attempt to install via "git clone" instead?`, this.loadedExtensions = builtExtensions; + // Register extension themes early so they're available at startup. + for (const ext of this.loadedExtensions) { + if (ext.isActive && ext.themes) { + themeManager.registerExtensionThemes(ext.name, ext.themes); + } + } + await Promise.all( this.loadedExtensions.map((ext) => this.maybeStartExtension(ext)), ); diff --git a/packages/cli/src/ui/themes/theme-manager.ts b/packages/cli/src/ui/themes/theme-manager.ts index 00fed5ce20..66826bb87e 100644 --- a/packages/cli/src/ui/themes/theme-manager.ts +++ b/packages/cli/src/ui/themes/theme-manager.ts @@ -240,6 +240,17 @@ class ThemeManager { } } + /** + * Checks if themes for a given extension are already registered. + * @param extensionName The name of the extension. + * @returns True if any themes from the extension are registered. + */ + hasExtensionThemes(extensionName: string): boolean { + return Array.from(this.extensionThemes.keys()).some((name) => + name.endsWith(`(${extensionName})`), + ); + } + /** * Clears all registered extension themes. * This is primarily for testing purposes to reset state between tests. From 18e8dd768aa65fcd9bb227de0a3ab1ac8bdb61d8 Mon Sep 17 00:00:00 2001 From: Sehoon Shon Date: Thu, 12 Mar 2026 09:46:58 -0400 Subject: [PATCH 2/7] feat(ui): Do not show Ultra users /upgrade hint (#22154) (#22156) --- .../src/ui/commands/upgradeCommand.test.ts | 20 +++++++++++ .../cli/src/ui/commands/upgradeCommand.ts | 10 ++++++ .../cli/src/ui/components/DialogManager.tsx | 1 + .../src/ui/components/ProQuotaDialog.test.tsx | 34 +++++++++++++++++++ .../cli/src/ui/components/ProQuotaDialog.tsx | 7 +++- .../src/ui/components/UserIdentity.test.tsx | 19 +++++++++++ .../cli/src/ui/components/UserIdentity.tsx | 5 ++- packages/cli/src/utils/tierUtils.test.ts | 28 +++++++++++++++ packages/cli/src/utils/tierUtils.ts | 15 ++++++++ 9 files changed, 137 insertions(+), 2 deletions(-) create mode 100644 packages/cli/src/utils/tierUtils.test.ts create mode 100644 packages/cli/src/utils/tierUtils.ts diff --git a/packages/cli/src/ui/commands/upgradeCommand.test.ts b/packages/cli/src/ui/commands/upgradeCommand.test.ts index d511f69c3a..9c54eb0191 100644 --- a/packages/cli/src/ui/commands/upgradeCommand.test.ts +++ b/packages/cli/src/ui/commands/upgradeCommand.test.ts @@ -37,6 +37,7 @@ describe('upgradeCommand', () => { getContentGeneratorConfig: vi.fn().mockReturnValue({ authType: AuthType.LOGIN_WITH_GOOGLE, }), + getUserTierName: vi.fn().mockReturnValue(undefined), }, }, } as unknown as CommandContext); @@ -115,4 +116,23 @@ describe('upgradeCommand', () => { }); expect(openBrowserSecurely).not.toHaveBeenCalled(); }); + + it('should return info message for ultra tiers', async () => { + vi.mocked(mockContext.services.config!.getUserTierName).mockReturnValue( + 'Advanced Ultra', + ); + + if (!upgradeCommand.action) { + throw new Error('The upgrade command must have an action.'); + } + + const result = await upgradeCommand.action(mockContext, ''); + + expect(result).toEqual({ + type: 'message', + messageType: 'info', + content: 'You are already on the highest tier: Advanced Ultra.', + }); + expect(openBrowserSecurely).not.toHaveBeenCalled(); + }); }); diff --git a/packages/cli/src/ui/commands/upgradeCommand.ts b/packages/cli/src/ui/commands/upgradeCommand.ts index 4904509df1..9bbea156ce 100644 --- a/packages/cli/src/ui/commands/upgradeCommand.ts +++ b/packages/cli/src/ui/commands/upgradeCommand.ts @@ -10,6 +10,7 @@ import { shouldLaunchBrowser, UPGRADE_URL_PAGE, } from '@google/gemini-cli-core'; +import { isUltraTier } from '../../utils/tierUtils.js'; import { CommandKind, type SlashCommand } from './types.js'; /** @@ -35,6 +36,15 @@ export const upgradeCommand: SlashCommand = { }; } + const tierName = context.services.config?.getUserTierName(); + if (isUltraTier(tierName)) { + return { + type: 'message', + messageType: 'info', + content: `You are already on the highest tier: ${tierName}.`, + }; + } + if (!shouldLaunchBrowser()) { return { type: 'message', diff --git a/packages/cli/src/ui/components/DialogManager.tsx b/packages/cli/src/ui/components/DialogManager.tsx index de62401e1e..e7e23c834d 100644 --- a/packages/cli/src/ui/components/DialogManager.tsx +++ b/packages/cli/src/ui/components/DialogManager.tsx @@ -87,6 +87,7 @@ export const DialogManager = ({ !!uiState.quota.proQuotaRequest.isModelNotFoundError } authType={uiState.quota.proQuotaRequest.authType} + tierName={config?.getUserTierName()} onChoice={uiActions.handleProQuotaChoice} /> ); diff --git a/packages/cli/src/ui/components/ProQuotaDialog.test.tsx b/packages/cli/src/ui/components/ProQuotaDialog.test.tsx index d97d53314e..2b69770582 100644 --- a/packages/cli/src/ui/components/ProQuotaDialog.test.tsx +++ b/packages/cli/src/ui/components/ProQuotaDialog.test.tsx @@ -202,6 +202,40 @@ describe('ProQuotaDialog', () => { ); unmount(); }); + + it('should NOT render upgrade option for LOGIN_WITH_GOOGLE if tier is Ultra', () => { + const { unmount } = render( + , + ); + + expect(RadioButtonSelect).toHaveBeenCalledWith( + expect.objectContaining({ + items: [ + { + label: 'Switch to gemini-2.5-flash', + value: 'retry_always', + key: 'retry_always', + }, + { + label: 'Stop', + value: 'retry_later', + key: 'retry_later', + }, + ], + }), + undefined, + ); + unmount(); + }); }); describe('when it is a capacity error', () => { diff --git a/packages/cli/src/ui/components/ProQuotaDialog.tsx b/packages/cli/src/ui/components/ProQuotaDialog.tsx index 82a679db8c..e9e869edb0 100644 --- a/packages/cli/src/ui/components/ProQuotaDialog.tsx +++ b/packages/cli/src/ui/components/ProQuotaDialog.tsx @@ -9,6 +9,7 @@ import { Box, Text } from 'ink'; import { RadioButtonSelect } from './shared/RadioButtonSelect.js'; import { theme } from '../semantic-colors.js'; import { AuthType } from '@google/gemini-cli-core'; +import { isUltraTier } from '../../utils/tierUtils.js'; interface ProQuotaDialogProps { failedModel: string; @@ -17,6 +18,7 @@ interface ProQuotaDialogProps { isTerminalQuotaError: boolean; isModelNotFoundError?: boolean; authType?: AuthType; + tierName?: string; onChoice: ( choice: 'retry_later' | 'retry_once' | 'retry_always' | 'upgrade', ) => void; @@ -29,6 +31,7 @@ export function ProQuotaDialog({ isTerminalQuotaError, isModelNotFoundError, authType, + tierName, onChoice, }: ProQuotaDialogProps): React.JSX.Element { let items; @@ -47,6 +50,8 @@ export function ProQuotaDialog({ }, ]; } else if (isModelNotFoundError || isTerminalQuotaError) { + const isUltra = isUltraTier(tierName); + // free users and out of quota users on G1 pro and Cloud Console gets an option to upgrade items = [ { @@ -54,7 +59,7 @@ export function ProQuotaDialog({ value: 'retry_always' as const, key: 'retry_always', }, - ...(authType === AuthType.LOGIN_WITH_GOOGLE + ...(authType === AuthType.LOGIN_WITH_GOOGLE && !isUltra ? [ { label: 'Upgrade for higher limits', diff --git a/packages/cli/src/ui/components/UserIdentity.test.tsx b/packages/cli/src/ui/components/UserIdentity.test.tsx index 2aade5675b..8caa21b808 100644 --- a/packages/cli/src/ui/components/UserIdentity.test.tsx +++ b/packages/cli/src/ui/components/UserIdentity.test.tsx @@ -182,4 +182,23 @@ describe('', () => { expect(output).toContain('/upgrade'); unmount(); }); + + it('should not render /upgrade indicator for ultra tiers', async () => { + const mockConfig = makeFakeConfig(); + vi.spyOn(mockConfig, 'getContentGeneratorConfig').mockReturnValue({ + authType: AuthType.LOGIN_WITH_GOOGLE, + model: 'gemini-pro', + } as unknown as ContentGeneratorConfig); + vi.spyOn(mockConfig, 'getUserTierName').mockReturnValue('Advanced Ultra'); + + const { lastFrame, waitUntilReady, unmount } = renderWithProviders( + , + ); + await waitUntilReady(); + + const output = lastFrame(); + expect(output).toContain('Plan: Advanced Ultra'); + expect(output).not.toContain('/upgrade'); + unmount(); + }); }); diff --git a/packages/cli/src/ui/components/UserIdentity.tsx b/packages/cli/src/ui/components/UserIdentity.tsx index fa2f5c5afa..5ce4452aa4 100644 --- a/packages/cli/src/ui/components/UserIdentity.tsx +++ b/packages/cli/src/ui/components/UserIdentity.tsx @@ -13,6 +13,7 @@ import { UserAccountManager, AuthType, } from '@google/gemini-cli-core'; +import { isUltraTier } from '../../utils/tierUtils.js'; interface UserIdentityProps { config: Config; @@ -33,6 +34,8 @@ export const UserIdentity: React.FC = ({ config }) => { [config, authType], ); + const isUltra = useMemo(() => isUltraTier(tierName), [tierName]); + if (!authType) { return null; } @@ -60,7 +63,7 @@ export const UserIdentity: React.FC = ({ config }) => { Plan: {tierName} - /upgrade + {!isUltra && /upgrade} )} diff --git a/packages/cli/src/utils/tierUtils.test.ts b/packages/cli/src/utils/tierUtils.test.ts new file mode 100644 index 0000000000..05cdaa22bd --- /dev/null +++ b/packages/cli/src/utils/tierUtils.test.ts @@ -0,0 +1,28 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, expect, it } from 'vitest'; +import { isUltraTier } from './tierUtils.js'; + +describe('tierUtils', () => { + describe('isUltraTier', () => { + it('should return true if tier name contains "ultra" (case-insensitive)', () => { + expect(isUltraTier('Advanced Ultra')).toBe(true); + expect(isUltraTier('gemini ultra')).toBe(true); + expect(isUltraTier('ULTRA')).toBe(true); + }); + + it('should return false if tier name does not contain "ultra"', () => { + expect(isUltraTier('Free')).toBe(false); + expect(isUltraTier('Pro')).toBe(false); + expect(isUltraTier('Standard')).toBe(false); + }); + + it('should return false if tier name is undefined', () => { + expect(isUltraTier(undefined)).toBe(false); + }); + }); +}); diff --git a/packages/cli/src/utils/tierUtils.ts b/packages/cli/src/utils/tierUtils.ts new file mode 100644 index 0000000000..7722a9a411 --- /dev/null +++ b/packages/cli/src/utils/tierUtils.ts @@ -0,0 +1,15 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Checks if the given tier name corresponds to an "Ultra" tier. + * + * @param tierName The name of the user's tier. + * @returns True if the tier is an "Ultra" tier, false otherwise. + */ +export function isUltraTier(tierName?: string): boolean { + return !!tierName?.toLowerCase().includes('ultra'); +} From a38aaa47fbf1442339d617b480d37b3bc3221c99 Mon Sep 17 00:00:00 2001 From: Jack Wotherspoon Date: Thu, 12 Mar 2026 14:51:36 +0100 Subject: [PATCH 3/7] chore: remove unnecessary log for themes (#22165) --- packages/cli/src/ui/themes/theme-manager.ts | 5 ----- 1 file changed, 5 deletions(-) diff --git a/packages/cli/src/ui/themes/theme-manager.ts b/packages/cli/src/ui/themes/theme-manager.ts index 66826bb87e..96b4fea4e3 100644 --- a/packages/cli/src/ui/themes/theme-manager.ts +++ b/packages/cli/src/ui/themes/theme-manager.ts @@ -174,11 +174,6 @@ class ThemeManager { return; } - debugLogger.log( - `Registering extension themes for "${extensionName}":`, - customThemes, - ); - for (const customThemeConfig of customThemes) { const namespacedName = `${customThemeConfig.name} (${extensionName})`; From 8432bcee752bb36e2b55274a9426ffac604812ab Mon Sep 17 00:00:00 2001 From: Abhi <43648792+abhipatel12@users.noreply.github.com> Date: Thu, 12 Mar 2026 10:17:36 -0400 Subject: [PATCH 4/7] fix(core): resolve MCP tool FQN validation, schema export, and wildcards in subagents (#22069) --- packages/core/src/agents/agentLoader.ts | 8 ++- packages/core/src/agents/local-executor.ts | 66 ++++++++++++++----- packages/core/src/tools/mcp-tool.ts | 20 +----- packages/core/src/tools/tool-names.test.ts | 40 ++++++----- packages/core/src/tools/tool-names.ts | 51 ++++++++++---- packages/core/src/tools/tool-registry.test.ts | 15 ++--- packages/core/src/tools/tool-registry.ts | 35 +++++----- 7 files changed, 136 insertions(+), 99 deletions(-) diff --git a/packages/core/src/agents/agentLoader.ts b/packages/core/src/agents/agentLoader.ts index 12337c6248..e0ccba0782 100644 --- a/packages/core/src/agents/agentLoader.ts +++ b/packages/core/src/agents/agentLoader.ts @@ -107,9 +107,11 @@ const localAgentSchema = z display_name: z.string().optional(), tools: z .array( - z.string().refine((val) => isValidToolName(val), { - message: 'Invalid tool name', - }), + z + .string() + .refine((val) => isValidToolName(val, { allowWildcards: true }), { + message: 'Invalid tool name', + }), ) .optional(), model: z.string().optional(), diff --git a/packages/core/src/agents/local-executor.ts b/packages/core/src/agents/local-executor.ts index cbc6260304..6a9dfe0151 100644 --- a/packages/core/src/agents/local-executor.ts +++ b/packages/core/src/agents/local-executor.ts @@ -17,7 +17,13 @@ import { type Schema, } from '@google/genai'; import { ToolRegistry } from '../tools/tool-registry.js'; -import { DiscoveredMCPTool } from '../tools/mcp-tool.js'; +import { type AnyDeclarativeTool } from '../tools/tools.js'; +import { + DiscoveredMCPTool, + isMcpToolName, + parseMcpToolName, + MCP_TOOL_PREFIX, +} from '../tools/mcp-tool.js'; import { CompressionStatus } from '../core/turn.js'; import { type ToolCallRequestInfo } from '../scheduler/types.js'; import { type Message } from '../confirmation-bus/types.js'; @@ -146,28 +152,55 @@ export class LocalAgentExecutor { context.config.getAgentRegistry().getAllAgentNames(), ); - const registerToolByName = (toolName: string) => { + const registerToolInstance = (tool: AnyDeclarativeTool) => { // Check if the tool is a subagent to prevent recursion. // We do not allow agents to call other agents. - if (allAgentNames.has(toolName)) { + if (allAgentNames.has(tool.name)) { debugLogger.warn( - `[LocalAgentExecutor] Skipping subagent tool '${toolName}' for agent '${definition.name}' to prevent recursion.`, + `[LocalAgentExecutor] Skipping subagent tool '${tool.name}' for agent '${definition.name}' to prevent recursion.`, ); return; } + agentToolRegistry.registerTool(tool); + }; + + const registerToolByName = (toolName: string) => { + // Handle global wildcard + if (toolName === '*') { + for (const tool of parentToolRegistry.getAllTools()) { + registerToolInstance(tool); + } + return; + } + + // Handle MCP wildcards + if (isMcpToolName(toolName)) { + if (toolName === `${MCP_TOOL_PREFIX}*`) { + for (const tool of parentToolRegistry.getAllTools()) { + if (tool instanceof DiscoveredMCPTool) { + registerToolInstance(tool); + } + } + return; + } + + const parsed = parseMcpToolName(toolName); + if (parsed.serverName && parsed.toolName === '*') { + for (const tool of parentToolRegistry.getToolsByServer( + parsed.serverName, + )) { + registerToolInstance(tool); + } + return; + } + } + // If the tool is referenced by name, retrieve it from the parent // registry and register it with the agent's isolated registry. const tool = parentToolRegistry.getTool(toolName); if (tool) { - if (tool instanceof DiscoveredMCPTool) { - // Subagents MUST use fully qualified names for MCP tools to ensure - // unambiguous tool calls and to comply with policy requirements. - // We automatically "upgrade" any MCP tool to its qualified version. - agentToolRegistry.registerTool(tool.asFullyQualifiedTool()); - } else { - agentToolRegistry.registerTool(tool); - } + registerToolInstance(tool); } }; @@ -1175,10 +1208,9 @@ export class LocalAgentExecutor { const { toolConfig, outputConfig } = this.definition; if (toolConfig) { - const toolNamesToLoad: string[] = []; for (const toolRef of toolConfig.tools) { if (typeof toolRef === 'string') { - toolNamesToLoad.push(toolRef); + // The names were already expanded and loaded into the registry during creation. } else if (typeof toolRef === 'object' && 'schema' in toolRef) { // Tool instance with an explicit schema property. toolsList.push(toolRef.schema); @@ -1187,10 +1219,8 @@ export class LocalAgentExecutor { toolsList.push(toolRef); } } - // Add schemas from tools that were registered by name. - toolsList.push( - ...this.toolRegistry.getFunctionDeclarationsFiltered(toolNamesToLoad), - ); + // Add schemas from tools that were explicitly registered by name or wildcard. + toolsList.push(...this.toolRegistry.getFunctionDeclarations()); } // Always inject complete_task. diff --git a/packages/core/src/tools/mcp-tool.ts b/packages/core/src/tools/mcp-tool.ts index 523eac62ad..5702f88a52 100644 --- a/packages/core/src/tools/mcp-tool.ts +++ b/packages/core/src/tools/mcp-tool.ts @@ -58,6 +58,7 @@ export function parseMcpToolName(name: string): { // Remove the prefix const withoutPrefix = name.slice(MCP_TOOL_PREFIX.length); // The first segment is the server name, the rest is the tool name + // Must be strictly `server_tool` where neither are empty const match = withoutPrefix.match(/^([^_]+)_(.+)$/); if (match) { return { @@ -390,25 +391,6 @@ export class DiscoveredMCPTool extends BaseDeclarativeTool< `${this.serverName}${MCP_QUALIFIED_NAME_SEPARATOR}${this.serverToolName}`, ); } - - asFullyQualifiedTool(): DiscoveredMCPTool { - return new DiscoveredMCPTool( - this.mcpTool, - this.serverName, - this.serverToolName, - this.description, - this.parameterSchema, - this.messageBus, - this.trust, - this.isReadOnly, - this.getFullyQualifiedName(), - this.cliConfig, - this.extensionName, - this.extensionId, - this._toolAnnotations, - ); - } - protected createInvocation( params: ToolParams, messageBus: MessageBus, diff --git a/packages/core/src/tools/tool-names.test.ts b/packages/core/src/tools/tool-names.test.ts index 8ff871986f..c631541171 100644 --- a/packages/core/src/tools/tool-names.test.ts +++ b/packages/core/src/tools/tool-names.test.ts @@ -25,7 +25,8 @@ vi.mock('./tool-names.js', async (importOriginal) => { ...actual, TOOL_LEGACY_ALIASES: mockedAliases, isValidToolName: vi.fn().mockImplementation((name: string, options) => { - if (mockedAliases[name]) return true; + if (Object.prototype.hasOwnProperty.call(mockedAliases, name)) + return true; return actual.isValidToolName(name, options); }), getToolAliases: vi.fn().mockImplementation((name: string) => { @@ -55,11 +56,9 @@ describe('tool-names', () => { expect(isValidToolName(`${DISCOVERED_TOOL_PREFIX}my_tool`)).toBe(true); }); - it('should validate MCP tool names (server__tool)', () => { - expect(isValidToolName('server__tool')).toBe(true); - expect(isValidToolName('my-server__my-tool')).toBe(true); - expect(isValidToolName('my.server__my:tool')).toBe(true); - expect(isValidToolName('my-server...truncated__tool')).toBe(true); + it('should validate modern MCP FQNs (mcp_server_tool)', () => { + expect(isValidToolName('mcp_server_tool')).toBe(true); + expect(isValidToolName('mcp_my-server_my-tool')).toBe(true); }); it('should validate legacy tool aliases', async () => { @@ -69,28 +68,33 @@ describe('tool-names', () => { } }); - it('should reject invalid tool names', () => { - expect(isValidToolName('')).toBe(false); - expect(isValidToolName('invalid-name')).toBe(false); - expect(isValidToolName('server__')).toBe(false); - expect(isValidToolName('__tool')).toBe(false); - expect(isValidToolName('server__tool__extra')).toBe(false); + it('should return false for invalid tool names', () => { + expect(isValidToolName('invalid-tool-name')).toBe(false); + expect(isValidToolName('mcp_server')).toBe(false); + expect(isValidToolName('mcp__tool')).toBe(false); + expect(isValidToolName('mcp_invalid server_tool')).toBe(false); + expect(isValidToolName('mcp_server_invalid tool')).toBe(false); + expect(isValidToolName('mcp_server_')).toBe(false); }); it('should handle wildcards when allowed', () => { // Default: not allowed expect(isValidToolName('*')).toBe(false); - expect(isValidToolName('server__*')).toBe(false); + expect(isValidToolName('mcp_*')).toBe(false); + expect(isValidToolName('mcp_server_*')).toBe(false); // Explicitly allowed expect(isValidToolName('*', { allowWildcards: true })).toBe(true); - expect(isValidToolName('server__*', { allowWildcards: true })).toBe(true); + expect(isValidToolName('mcp_*', { allowWildcards: true })).toBe(true); + expect(isValidToolName('mcp_server_*', { allowWildcards: true })).toBe( + true, + ); // Invalid wildcards - expect(isValidToolName('__*', { allowWildcards: true })).toBe(false); - expect(isValidToolName('server__tool*', { allowWildcards: true })).toBe( - false, - ); + expect(isValidToolName('mcp__*', { allowWildcards: true })).toBe(false); + expect( + isValidToolName('mcp_server_tool*', { allowWildcards: true }), + ).toBe(false); }); }); diff --git a/packages/core/src/tools/tool-names.ts b/packages/core/src/tools/tool-names.ts index 38a868d665..91b0574d9e 100644 --- a/packages/core/src/tools/tool-names.ts +++ b/packages/core/src/tools/tool-names.ts @@ -221,6 +221,12 @@ export const DISCOVERED_TOOL_PREFIX = 'discovered_tool_'; /** * List of all built-in tool names. */ +import { + isMcpToolName, + parseMcpToolName, + MCP_TOOL_PREFIX, +} from './mcp-tool.js'; + export const ALL_BUILTIN_TOOL_NAMES = [ GLOB_TOOL_NAME, WRITE_TODOS_TOOL_NAME, @@ -290,25 +296,44 @@ export function isValidToolName( return true; } - // MCP tools (format: server__tool) - if (name.includes('__')) { - const parts = name.split('__'); - if (parts.length !== 2 || parts[0].length === 0 || parts[1].length === 0) { + // Handle standard MCP FQNs (mcp_server_tool or wildcards mcp_*, mcp_server_*) + if (isMcpToolName(name)) { + // Global wildcard: mcp_* + if (name === `${MCP_TOOL_PREFIX}*` && options.allowWildcards) { + return true; + } + + // Explicitly reject names with empty server component (e.g. mcp__tool) + if (name.startsWith(`${MCP_TOOL_PREFIX}_`)) { return false; } - const server = parts[0]; - const tool = parts[1]; + const parsed = parseMcpToolName(name); + // Ensure that both components are populated. parseMcpToolName splits at the second _, + // so `mcp__tool` has serverName="", toolName="tool" + if (parsed.serverName && parsed.toolName) { + // Basic slug validation for server and tool names. + // We allow dots (.) and colons (:) as they are valid in function names and + // used for truncation markers. + const slugRegex = /^[a-z0-9_.:-]+$/i; - if (tool === '*') { - return !!options.allowWildcards; + if (!slugRegex.test(parsed.serverName)) { + return false; + } + + if (parsed.toolName === '*') { + return options.allowWildcards === true; + } + + // A tool name consisting only of underscores is invalid. + if (/^_*$/.test(parsed.toolName)) { + return false; + } + + return slugRegex.test(parsed.toolName); } - // Basic slug validation for server and tool names. - // We allow dots (.) and colons (:) as they are valid in function names and - // used for truncation markers. - const slugRegex = /^[a-z0-9_.:-]+$/i; - return slugRegex.test(server) && slugRegex.test(tool); + return false; } return false; diff --git a/packages/core/src/tools/tool-registry.test.ts b/packages/core/src/tools/tool-registry.test.ts index ea560865e6..21bbb0cc71 100644 --- a/packages/core/src/tools/tool-registry.test.ts +++ b/packages/core/src/tools/tool-registry.test.ts @@ -310,13 +310,13 @@ describe('ToolRegistry', () => { excludedTools: ['tool-a'], }, { - name: 'should match simple MCP tool names, when qualified or unqualified', - tools: [mcpTool, mcpTool.asFullyQualifiedTool()], + name: 'should match simple MCP tool names', + tools: [mcpTool], excludedTools: [mcpTool.name], }, { - name: 'should match qualified MCP tool names when qualified or unqualified', - tools: [mcpTool, mcpTool.asFullyQualifiedTool()], + name: 'should match qualified MCP tool names', + tools: [mcpTool], excludedTools: [mcpTool.name], }, { @@ -414,9 +414,9 @@ describe('ToolRegistry', () => { const toolName = 'my-tool'; const mcpTool = createMCPTool(serverName, toolName, 'desc'); - // Register same MCP tool twice (one as alias, one as qualified) + // Register same MCP tool twice + toolRegistry.registerTool(mcpTool); toolRegistry.registerTool(mcpTool); - toolRegistry.registerTool(mcpTool.asFullyQualifiedTool()); const toolNames = toolRegistry.getAllToolNames(); expect(toolNames).toEqual([`mcp_${serverName}_${toolName}`]); @@ -698,9 +698,8 @@ describe('ToolRegistry', () => { const toolName = 'my-tool'; const mcpTool = createMCPTool(serverName, toolName, 'description'); - // Register both alias and qualified toolRegistry.registerTool(mcpTool); - toolRegistry.registerTool(mcpTool.asFullyQualifiedTool()); + toolRegistry.registerTool(mcpTool); const declarations = toolRegistry.getFunctionDeclarations(); expect(declarations).toHaveLength(1); diff --git a/packages/core/src/tools/tool-registry.ts b/packages/core/src/tools/tool-registry.ts index 69695877c2..f8542112bb 100644 --- a/packages/core/src/tools/tool-registry.ts +++ b/packages/core/src/tools/tool-registry.ts @@ -222,14 +222,10 @@ export class ToolRegistry { */ registerTool(tool: AnyDeclarativeTool): void { if (this.allKnownTools.has(tool.name)) { - if (tool instanceof DiscoveredMCPTool) { - tool = tool.asFullyQualifiedTool(); - } else { - // Decide on behavior: throw error, log warning, or allow overwrite - debugLogger.warn( - `Tool with name "${tool.name}" is already registered. Overwriting.`, - ); - } + // Decide on behavior: throw error, log warning, or allow overwrite + debugLogger.warn( + `Tool with name "${tool.name}" is already registered. Overwriting.`, + ); } this.allKnownTools.set(tool.name, tool); } @@ -594,7 +590,17 @@ export class ToolRegistry { for (const name of toolNames) { const tool = this.getTool(name); if (tool) { - declarations.push(tool.getSchema(modelId)); + let schema = tool.getSchema(modelId); + + // Ensure the schema name matches the qualified name for MCP tools + if (tool instanceof DiscoveredMCPTool) { + schema = { + ...schema, + name: tool.getFullyQualifiedName(), + }; + } + + declarations.push(schema); } } return declarations; @@ -670,17 +676,6 @@ export class ToolRegistry { } } - if (!tool && name.includes('__')) { - for (const t of this.allKnownTools.values()) { - if (t instanceof DiscoveredMCPTool) { - if (t.getFullyQualifiedName() === name) { - tool = t; - break; - } - } - } - } - if (tool && this.isActiveTool(tool)) { return tool; } From 34709dc62d4350f415df27822ef02fcb1b6e6ec5 Mon Sep 17 00:00:00 2001 From: Jaisal K Jain <105512018+JaisalJain@users.noreply.github.com> Date: Thu, 12 Mar 2026 20:08:54 +0530 Subject: [PATCH 5/7] fix(cli): validate --model argument at startup (#21393) Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- packages/cli/src/config/config.test.ts | 8 ++-- packages/cli/src/config/config.ts | 14 ++++++ packages/core/src/config/models.test.ts | 64 +++++++++++++++++++++++++ packages/core/src/config/models.ts | 43 +++++++++++++++++ 4 files changed, 125 insertions(+), 4 deletions(-) diff --git a/packages/cli/src/config/config.test.ts b/packages/cli/src/config/config.test.ts index 422f6cd2ac..995be3fc61 100644 --- a/packages/cli/src/config/config.test.ts +++ b/packages/cli/src/config/config.test.ts @@ -1773,7 +1773,7 @@ describe('loadCliConfig model selection', () => { }); it('always prefers model from argv', async () => { - process.argv = ['node', 'script.js', '--model', 'gemini-2.5-flash-preview']; + process.argv = ['node', 'script.js', '--model', 'gemini-2.5-flash']; const argv = await parseArguments(createTestMergedSettings()); const config = await loadCliConfig( createTestMergedSettings({ @@ -1785,11 +1785,11 @@ describe('loadCliConfig model selection', () => { argv, ); - expect(config.getModel()).toBe('gemini-2.5-flash-preview'); + expect(config.getModel()).toBe('gemini-2.5-flash'); }); it('selects the model from argv if provided', async () => { - process.argv = ['node', 'script.js', '--model', 'gemini-2.5-flash-preview']; + process.argv = ['node', 'script.js', '--model', 'gemini-2.5-flash']; const argv = await parseArguments(createTestMergedSettings()); const config = await loadCliConfig( createTestMergedSettings({ @@ -1799,7 +1799,7 @@ describe('loadCliConfig model selection', () => { argv, ); - expect(config.getModel()).toBe('gemini-2.5-flash-preview'); + expect(config.getModel()).toBe('gemini-2.5-flash'); }); it('selects the default auto model if provided via auto alias', async () => { diff --git a/packages/cli/src/config/config.ts b/packages/cli/src/config/config.ts index 010fb8e17f..e910d47546 100755 --- a/packages/cli/src/config/config.ts +++ b/packages/cli/src/config/config.ts @@ -31,6 +31,8 @@ import { type HierarchicalMemory, coreEvents, GEMINI_MODEL_ALIAS_AUTO, + isValidModelOrAlias, + getValidModelsAndAliases, getAdminErrorMessage, isHeadlessMode, Config, @@ -671,6 +673,18 @@ export async function loadCliConfig( const specifiedModel = argv.model || process.env['GEMINI_MODEL'] || settings.model?.name; + // Validate the model if one was explicitly specified + if (specifiedModel && specifiedModel !== GEMINI_MODEL_ALIAS_AUTO) { + if (!isValidModelOrAlias(specifiedModel)) { + const validModels = getValidModelsAndAliases(); + + throw new FatalConfigError( + `Invalid model: "${specifiedModel}"\n\n` + + `Valid models and aliases:\n${validModels.map((m) => ` - ${m}`).join('\n')}\n\n` + + `Use /model to switch models interactively.`, + ); + } + } const resolvedModel = specifiedModel === GEMINI_MODEL_ALIAS_AUTO ? defaultModel diff --git a/packages/core/src/config/models.test.ts b/packages/core/src/config/models.test.ts index d62827ed91..b3f5db9430 100644 --- a/packages/core/src/config/models.test.ts +++ b/packages/core/src/config/models.test.ts @@ -22,6 +22,7 @@ import { GEMINI_MODEL_ALIAS_PRO, GEMINI_MODEL_ALIAS_FLASH, GEMINI_MODEL_ALIAS_AUTO, + GEMINI_MODEL_ALIAS_FLASH_LITE, PREVIEW_GEMINI_FLASH_MODEL, PREVIEW_GEMINI_MODEL_AUTO, DEFAULT_GEMINI_MODEL_AUTO, @@ -30,6 +31,10 @@ import { PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL, isPreviewModel, isProModel, + isValidModelOrAlias, + getValidModelsAndAliases, + VALID_GEMINI_MODELS, + VALID_ALIASES, } from './models.js'; describe('isPreviewModel', () => { @@ -389,3 +394,62 @@ describe('isActiveModel', () => { ).toBe(false); }); }); + +describe('isValidModelOrAlias', () => { + it('should return true for valid model names', () => { + expect(isValidModelOrAlias(DEFAULT_GEMINI_MODEL)).toBe(true); + expect(isValidModelOrAlias(PREVIEW_GEMINI_MODEL)).toBe(true); + expect(isValidModelOrAlias(DEFAULT_GEMINI_FLASH_MODEL)).toBe(true); + expect(isValidModelOrAlias(DEFAULT_GEMINI_FLASH_LITE_MODEL)).toBe(true); + expect(isValidModelOrAlias(PREVIEW_GEMINI_FLASH_MODEL)).toBe(true); + expect(isValidModelOrAlias(PREVIEW_GEMINI_3_1_MODEL)).toBe(true); + expect(isValidModelOrAlias(PREVIEW_GEMINI_3_1_CUSTOM_TOOLS_MODEL)).toBe( + true, + ); + }); + + it('should return true for valid aliases', () => { + expect(isValidModelOrAlias(GEMINI_MODEL_ALIAS_AUTO)).toBe(true); + expect(isValidModelOrAlias(GEMINI_MODEL_ALIAS_PRO)).toBe(true); + expect(isValidModelOrAlias(GEMINI_MODEL_ALIAS_FLASH)).toBe(true); + expect(isValidModelOrAlias(GEMINI_MODEL_ALIAS_FLASH_LITE)).toBe(true); + expect(isValidModelOrAlias(PREVIEW_GEMINI_MODEL_AUTO)).toBe(true); + expect(isValidModelOrAlias(DEFAULT_GEMINI_MODEL_AUTO)).toBe(true); + }); + + it('should return true for custom (non-gemini) models', () => { + expect(isValidModelOrAlias('gpt-4')).toBe(true); + expect(isValidModelOrAlias('claude-3')).toBe(true); + expect(isValidModelOrAlias('my-custom-model')).toBe(true); + }); + + it('should return false for invalid gemini model names', () => { + expect(isValidModelOrAlias('gemini-4-pro')).toBe(false); + expect(isValidModelOrAlias('gemini-99-flash')).toBe(false); + expect(isValidModelOrAlias('gemini-invalid')).toBe(false); + }); +}); + +describe('getValidModelsAndAliases', () => { + it('should return a sorted array', () => { + const result = getValidModelsAndAliases(); + const sorted = [...result].sort(); + expect(result).toEqual(sorted); + }); + + it('should include all valid models and aliases', () => { + const result = getValidModelsAndAliases(); + for (const model of VALID_GEMINI_MODELS) { + expect(result).toContain(model); + } + for (const alias of VALID_ALIASES) { + expect(result).toContain(alias); + } + }); + + it('should not contain duplicates', () => { + const result = getValidModelsAndAliases(); + const unique = [...new Set(result)]; + expect(result).toEqual(unique); + }); +}); diff --git a/packages/core/src/config/models.ts b/packages/core/src/config/models.ts index ffbf597793..59e7e4b457 100644 --- a/packages/core/src/config/models.ts +++ b/packages/core/src/config/models.ts @@ -32,6 +32,15 @@ export const GEMINI_MODEL_ALIAS_PRO = 'pro'; export const GEMINI_MODEL_ALIAS_FLASH = 'flash'; export const GEMINI_MODEL_ALIAS_FLASH_LITE = 'flash-lite'; +export const VALID_ALIASES = new Set([ + GEMINI_MODEL_ALIAS_AUTO, + GEMINI_MODEL_ALIAS_PRO, + GEMINI_MODEL_ALIAS_FLASH, + GEMINI_MODEL_ALIAS_FLASH_LITE, + PREVIEW_GEMINI_MODEL_AUTO, + DEFAULT_GEMINI_MODEL_AUTO, +]); + export const DEFAULT_GEMINI_EMBEDDING_MODEL = 'gemini-embedding-001'; // Cap the thinking at 8192 to prevent run-away thinking loops. @@ -283,3 +292,37 @@ export function isActiveModel( ); } } + +/** + * Checks if the model name is valid (either a valid model or a valid alias). + * + * @param model The model name to check. + * @returns True if the model is valid. + */ +export function isValidModelOrAlias(model: string): boolean { + // Check if it's a valid alias + if (VALID_ALIASES.has(model)) { + return true; + } + + // Check if it's a valid model name + if (VALID_GEMINI_MODELS.has(model)) { + return true; + } + + // Allow custom models (non-gemini models) + if (!model.startsWith('gemini-')) { + return true; + } + + return false; +} + +/** + * Gets a list of all valid model names and aliases for error messages. + * + * @returns Array of valid model names and aliases. + */ +export function getValidModelsAndAliases(): string[] { + return [...new Set([...VALID_ALIASES, ...VALID_GEMINI_MODELS])].sort(); +} From 7506b00488def4a5305a909185637dc51ababeab Mon Sep 17 00:00:00 2001 From: Michael Ramos Date: Thu, 12 Mar 2026 07:43:40 -0700 Subject: [PATCH 6/7] fix(core): handle policy ALLOW for exit_plan_mode (#21802) --- .../core/src/tools/exit-plan-mode.test.ts | 20 +++++++++++++++++++ packages/core/src/tools/exit-plan-mode.ts | 12 +++++++++-- 2 files changed, 30 insertions(+), 2 deletions(-) diff --git a/packages/core/src/tools/exit-plan-mode.test.ts b/packages/core/src/tools/exit-plan-mode.test.ts index 22de81fc7f..4b6b537d00 100644 --- a/packages/core/src/tools/exit-plan-mode.test.ts +++ b/packages/core/src/tools/exit-plan-mode.test.ts @@ -339,6 +339,26 @@ Ask the user for specific feedback on how to improve the plan.`, }); }); + describe('execute when shouldConfirmExecute is never called', () => { + it('should approve with DEFAULT mode when approvalPayload is null (policy ALLOW skips confirmation)', async () => { + const planRelativePath = createPlanFile('test.md', '# Content'); + const invocation = tool.build({ plan_path: planRelativePath }); + + // Simulate the scheduler's policy ALLOW path: execute() is called + // directly without ever calling shouldConfirmExecute(), leaving + // approvalPayload null. + const result = await invocation.execute(new AbortController().signal); + const expectedPath = path.join(mockPlansDir, 'test.md'); + + expect(result.llmContent).toContain('Plan approved'); + expect(result.returnDisplay).toContain('Plan approved'); + expect(mockConfig.setApprovalMode).toHaveBeenCalledWith( + ApprovalMode.DEFAULT, + ); + expect(mockConfig.setApprovedPlanPath).toHaveBeenCalledWith(expectedPath); + }); + }); + describe('getApprovalModeDescription (internal)', () => { it('should handle all valid approval modes', async () => { const planRelativePath = createPlanFile('test.md', '# Content'); diff --git a/packages/core/src/tools/exit-plan-mode.ts b/packages/core/src/tools/exit-plan-mode.ts index 442b00e5cb..b1615b18b4 100644 --- a/packages/core/src/tools/exit-plan-mode.ts +++ b/packages/core/src/tools/exit-plan-mode.ts @@ -203,8 +203,16 @@ export class ExitPlanModeInvocation extends BaseToolInvocation< }; } - const payload = this.approvalPayload; - if (payload?.approved) { + // When a user policy grants `allow` for exit_plan_mode, the scheduler + // skips the confirmation phase entirely and shouldConfirmExecute is never + // called, leaving approvalPayload null. Treat that as an approval with + // the default mode — consistent with the ALLOW branch inside + // shouldConfirmExecute. + const payload = this.approvalPayload ?? { + approved: true, + approvalMode: ApprovalMode.DEFAULT, + }; + if (payload.approved) { const newMode = payload.approvalMode ?? ApprovalMode.DEFAULT; if (newMode === ApprovalMode.PLAN || newMode === ApprovalMode.YOLO) { From 867dc0fdda7c59fab627f6a1db65cf0f3fde84c2 Mon Sep 17 00:00:00 2001 From: Gaurav <39389231+gsquared94@users.noreply.github.com> Date: Thu, 12 Mar 2026 08:16:27 -0700 Subject: [PATCH 7/7] feat(telemetry): add Clearcut instrumentation for AI credits billing events (#22153) --- .../clearcut-logger/clearcut-logger.test.ts | 101 ++++++++++++++++++ .../clearcut-logger/clearcut-logger.ts | 88 +++++++++++++++ .../clearcut-logger/event-metadata-key.ts | 24 ++++- packages/core/src/telemetry/loggers.ts | 19 ++++ 4 files changed, 231 insertions(+), 1 deletion(-) diff --git a/packages/core/src/telemetry/clearcut-logger/clearcut-logger.test.ts b/packages/core/src/telemetry/clearcut-logger/clearcut-logger.test.ts index 93eebd651e..dd641e3955 100644 --- a/packages/core/src/telemetry/clearcut-logger/clearcut-logger.test.ts +++ b/packages/core/src/telemetry/clearcut-logger/clearcut-logger.test.ts @@ -51,6 +51,12 @@ import { InstallationManager } from '../../utils/installationManager.js'; import si, { type Systeminformation } from 'systeminformation'; import * as os from 'node:os'; +import { + CreditsUsedEvent, + OverageOptionSelectedEvent, + EmptyWalletMenuShownEvent, + CreditPurchaseClickEvent, +} from '../billingEvents.js'; interface CustomMatchers { toHaveMetadataValue: ([key, value]: [EventMetadataKey, string]) => R; @@ -1551,4 +1557,99 @@ describe('ClearcutLogger', () => { ]); }); }); + + describe('logCreditsUsedEvent', () => { + it('logs an event with model, consumed, and remaining credits', () => { + const { logger } = setup(); + const event = new CreditsUsedEvent('gemini-3-pro-preview', 10, 490); + + logger?.logCreditsUsedEvent(event); + + const events = getEvents(logger!); + expect(events.length).toBe(1); + expect(events[0]).toHaveEventName(EventNames.CREDITS_USED); + expect(events[0]).toHaveMetadataValue([ + EventMetadataKey.GEMINI_CLI_BILLING_MODEL, + '"gemini-3-pro-preview"', + ]); + expect(events[0]).toHaveMetadataValue([ + EventMetadataKey.GEMINI_CLI_BILLING_CREDITS_CONSUMED, + '10', + ]); + expect(events[0]).toHaveMetadataValue([ + EventMetadataKey.GEMINI_CLI_BILLING_CREDITS_REMAINING, + '490', + ]); + }); + }); + + describe('logOverageOptionSelectedEvent', () => { + it('logs an event with model, selected option, and credit balance', () => { + const { logger } = setup(); + const event = new OverageOptionSelectedEvent( + 'gemini-3-pro-preview', + 'use_credits', + 350, + ); + + logger?.logOverageOptionSelectedEvent(event); + + const events = getEvents(logger!); + expect(events.length).toBe(1); + expect(events[0]).toHaveEventName(EventNames.OVERAGE_OPTION_SELECTED); + expect(events[0]).toHaveMetadataValue([ + EventMetadataKey.GEMINI_CLI_BILLING_MODEL, + '"gemini-3-pro-preview"', + ]); + expect(events[0]).toHaveMetadataValue([ + EventMetadataKey.GEMINI_CLI_BILLING_SELECTED_OPTION, + '"use_credits"', + ]); + expect(events[0]).toHaveMetadataValue([ + EventMetadataKey.GEMINI_CLI_BILLING_CREDIT_BALANCE, + '350', + ]); + }); + }); + + describe('logEmptyWalletMenuShownEvent', () => { + it('logs an event with the model', () => { + const { logger } = setup(); + const event = new EmptyWalletMenuShownEvent('gemini-3-pro-preview'); + + logger?.logEmptyWalletMenuShownEvent(event); + + const events = getEvents(logger!); + expect(events.length).toBe(1); + expect(events[0]).toHaveEventName(EventNames.EMPTY_WALLET_MENU_SHOWN); + expect(events[0]).toHaveMetadataValue([ + EventMetadataKey.GEMINI_CLI_BILLING_MODEL, + '"gemini-3-pro-preview"', + ]); + }); + }); + + describe('logCreditPurchaseClickEvent', () => { + it('logs an event with model and source', () => { + const { logger } = setup(); + const event = new CreditPurchaseClickEvent( + 'empty_wallet_menu', + 'gemini-3-pro-preview', + ); + + logger?.logCreditPurchaseClickEvent(event); + + const events = getEvents(logger!); + expect(events.length).toBe(1); + expect(events[0]).toHaveEventName(EventNames.CREDIT_PURCHASE_CLICK); + expect(events[0]).toHaveMetadataValue([ + EventMetadataKey.GEMINI_CLI_BILLING_MODEL, + '"gemini-3-pro-preview"', + ]); + expect(events[0]).toHaveMetadataValue([ + EventMetadataKey.GEMINI_CLI_BILLING_PURCHASE_SOURCE, + '"empty_wallet_menu"', + ]); + }); + }); }); diff --git a/packages/core/src/telemetry/clearcut-logger/clearcut-logger.ts b/packages/core/src/telemetry/clearcut-logger/clearcut-logger.ts index 5e19d7f49b..5953578eae 100644 --- a/packages/core/src/telemetry/clearcut-logger/clearcut-logger.ts +++ b/packages/core/src/telemetry/clearcut-logger/clearcut-logger.ts @@ -52,6 +52,12 @@ import type { TokenStorageInitializationEvent, StartupStatsEvent, } from '../types.js'; +import type { + CreditsUsedEvent, + OverageOptionSelectedEvent, + EmptyWalletMenuShownEvent, + CreditPurchaseClickEvent, +} from '../billingEvents.js'; import { EventMetadataKey } from './event-metadata-key.js'; import type { Config } from '../../config/config.js'; import { InstallationManager } from '../../utils/installationManager.js'; @@ -121,6 +127,10 @@ export enum EventNames { CONSECA_POLICY_GENERATION = 'conseca_policy_generation', CONSECA_VERDICT = 'conseca_verdict', STARTUP_STATS = 'startup_stats', + CREDITS_USED = 'credits_used', + OVERAGE_OPTION_SELECTED = 'overage_option_selected', + EMPTY_WALLET_MENU_SHOWN = 'empty_wallet_menu_shown', + CREDIT_PURCHASE_CLICK = 'credit_purchase_click', } export interface LogResponse { @@ -1806,6 +1816,84 @@ export class ClearcutLogger { this.flushIfNeeded(); } + // ========================================================================== + // Billing / AI Credits Events + // ========================================================================== + + logCreditsUsedEvent(event: CreditsUsedEvent): void { + const data: EventValue[] = [ + { + gemini_cli_key: EventMetadataKey.GEMINI_CLI_BILLING_MODEL, + value: JSON.stringify(event.model), + }, + { + gemini_cli_key: EventMetadataKey.GEMINI_CLI_BILLING_CREDITS_CONSUMED, + value: JSON.stringify(event.credits_consumed), + }, + { + gemini_cli_key: EventMetadataKey.GEMINI_CLI_BILLING_CREDITS_REMAINING, + value: JSON.stringify(event.credits_remaining), + }, + ]; + + this.enqueueLogEvent(this.createLogEvent(EventNames.CREDITS_USED, data)); + this.flushIfNeeded(); + } + + logOverageOptionSelectedEvent(event: OverageOptionSelectedEvent): void { + const data: EventValue[] = [ + { + gemini_cli_key: EventMetadataKey.GEMINI_CLI_BILLING_MODEL, + value: JSON.stringify(event.model), + }, + { + gemini_cli_key: EventMetadataKey.GEMINI_CLI_BILLING_SELECTED_OPTION, + value: JSON.stringify(event.selected_option), + }, + { + gemini_cli_key: EventMetadataKey.GEMINI_CLI_BILLING_CREDIT_BALANCE, + value: JSON.stringify(event.credit_balance), + }, + ]; + + this.enqueueLogEvent( + this.createLogEvent(EventNames.OVERAGE_OPTION_SELECTED, data), + ); + this.flushIfNeeded(); + } + + logEmptyWalletMenuShownEvent(event: EmptyWalletMenuShownEvent): void { + const data: EventValue[] = [ + { + gemini_cli_key: EventMetadataKey.GEMINI_CLI_BILLING_MODEL, + value: JSON.stringify(event.model), + }, + ]; + + this.enqueueLogEvent( + this.createLogEvent(EventNames.EMPTY_WALLET_MENU_SHOWN, data), + ); + this.flushIfNeeded(); + } + + logCreditPurchaseClickEvent(event: CreditPurchaseClickEvent): void { + const data: EventValue[] = [ + { + gemini_cli_key: EventMetadataKey.GEMINI_CLI_BILLING_MODEL, + value: JSON.stringify(event.model), + }, + { + gemini_cli_key: EventMetadataKey.GEMINI_CLI_BILLING_PURCHASE_SOURCE, + value: JSON.stringify(event.source), + }, + ]; + + this.enqueueLogEvent( + this.createLogEvent(EventNames.CREDIT_PURCHASE_CLICK, data), + ); + this.flushIfNeeded(); + } + /** * Adds default fields to data, and returns a new data array. This fields * should exist on all log events. diff --git a/packages/core/src/telemetry/clearcut-logger/event-metadata-key.ts b/packages/core/src/telemetry/clearcut-logger/event-metadata-key.ts index 20c983aa63..632730aeeb 100644 --- a/packages/core/src/telemetry/clearcut-logger/event-metadata-key.ts +++ b/packages/core/src/telemetry/clearcut-logger/event-metadata-key.ts @@ -7,7 +7,7 @@ // Defines valid event metadata keys for Clearcut logging. export enum EventMetadataKey { // Deleted enums: 24 - // Next ID: 180 + // Next ID: 191 GEMINI_CLI_KEY_UNKNOWN = 0, @@ -687,4 +687,26 @@ export enum EventMetadataKey { // Logs the error type for a network retry. GEMINI_CLI_NETWORK_RETRY_ERROR_TYPE = 182, + + // ========================================================================== + // Billing / AI Credits Event Keys + // ========================================================================== + + // Logs the model associated with a billing event. + GEMINI_CLI_BILLING_MODEL = 185, + + // Logs the number of AI credits consumed in a request. + GEMINI_CLI_BILLING_CREDITS_CONSUMED = 186, + + // Logs the remaining AI credits after a request. + GEMINI_CLI_BILLING_CREDITS_REMAINING = 187, + + // Logs the overage option selected by the user (e.g. use_credits, use_fallback, manage, stop). + GEMINI_CLI_BILLING_SELECTED_OPTION = 188, + + // Logs the user's credit balance when the overage menu was shown. + GEMINI_CLI_BILLING_CREDIT_BALANCE = 189, + + // Logs the source of a credit purchase click (e.g. overage_menu, empty_wallet_menu, manage). + GEMINI_CLI_BILLING_PURCHASE_SOURCE = 190, } diff --git a/packages/core/src/telemetry/loggers.ts b/packages/core/src/telemetry/loggers.ts index 52e0fb35bb..d5cc605e65 100644 --- a/packages/core/src/telemetry/loggers.ts +++ b/packages/core/src/telemetry/loggers.ts @@ -85,6 +85,12 @@ import { uiTelemetryService, type UiEvent } from './uiTelemetry.js'; import { ClearcutLogger } from './clearcut-logger/clearcut-logger.js'; import { debugLogger } from '../utils/debugLogger.js'; import type { BillingTelemetryEvent } from './billingEvents.js'; +import { + CreditsUsedEvent, + OverageOptionSelectedEvent, + EmptyWalletMenuShownEvent, + CreditPurchaseClickEvent, +} from './billingEvents.js'; export function logCliConfiguration( config: Config, @@ -877,4 +883,17 @@ export function logBillingEvent( }; logger.emit(logRecord); }); + + const cc = ClearcutLogger.getInstance(config); + if (cc) { + if (event instanceof CreditsUsedEvent) { + cc.logCreditsUsedEvent(event); + } else if (event instanceof OverageOptionSelectedEvent) { + cc.logOverageOptionSelectedEvent(event); + } else if (event instanceof EmptyWalletMenuShownEvent) { + cc.logEmptyWalletMenuShownEvent(event); + } else if (event instanceof CreditPurchaseClickEvent) { + cc.logCreditPurchaseClickEvent(event); + } + } }