From 94362d67f0f74a4b3f3bab1441a3e040d20128a9 Mon Sep 17 00:00:00 2001 From: mkorwel Date: Fri, 20 Feb 2026 21:34:03 +0000 Subject: [PATCH] refactor: move experiment logic to ExperimentManager with Zod validation --- .../cli/src/ui/commands/experimentCommand.ts | 51 +-- packages/core/index.ts | 1 + .../src/code_assist/experiments/flagNames.ts | 168 +++++++++- packages/core/src/config/config.ts | 291 +++++++++--------- .../core/src/config/experimentManager.test.ts | 134 ++++++++ packages/core/src/config/experimentManager.ts | 224 ++++++++++++++ 6 files changed, 688 insertions(+), 181 deletions(-) create mode 100644 packages/core/src/config/experimentManager.test.ts create mode 100644 packages/core/src/config/experimentManager.ts diff --git a/packages/cli/src/ui/commands/experimentCommand.ts b/packages/cli/src/ui/commands/experimentCommand.ts index bdf59b3d05..10b2c0a18b 100644 --- a/packages/cli/src/ui/commands/experimentCommand.ts +++ b/packages/cli/src/ui/commands/experimentCommand.ts @@ -9,6 +9,7 @@ import { getExperimentFlagIdFromName, getExperimentFlagName, } from '@google/gemini-cli-core'; +import { z } from 'zod'; import { type CommandContext, CommandKind, @@ -26,7 +27,9 @@ const listExperimentsCommand: SlashCommand = { const { config } = context.services; if (!config) return; - const entries = Object.entries(ExperimentMetadata); + const entries = Object.entries(ExperimentMetadata).filter( + ([, metadata]) => !metadata.hidden, + ); if (entries.length === 0) { context.ui.addItem({ type: MessageType.INFO, @@ -40,7 +43,7 @@ const listExperimentsCommand: SlashCommand = { const id = parseInt(idStr, 10); const name = getExperimentFlagName(id) || `ID: ${id}`; const value = config.getExperimentValue(id); - output += `${name} (${metadata.type})\n`; + output += `${name}\n`; output += ` Value: ${value}\n`; output += ` Description: ${metadata.description}\n`; output += ` Default: ${metadata.defaultValue}\n\n`; @@ -82,31 +85,34 @@ const setExperimentCommand: SlashCommand = { } const metadata = ExperimentMetadata[id]; - let value: boolean | number | string; + let value: unknown; - if (metadata.type === 'boolean') { - if (rawValue === 'true' || rawValue === 'on') value = true; - else if (rawValue === 'false' || rawValue === 'off') value = false; - else { - context.ui.addItem({ - type: MessageType.ERROR, - text: `Invalid boolean value: ${rawValue}. Use true/false or on/off.`, - }); - return; + // Helper to parse strings based on Zod schema type + const parseValue = (raw: string, schema: z.ZodTypeAny): unknown => { + if (schema instanceof z.ZodBoolean) { + if (raw === 'true' || raw === 'on') return true; + if (raw === 'false' || raw === 'off') return false; + return raw; } - } else if (metadata.type === 'number') { - value = Number(rawValue); - if (isNaN(value)) { - context.ui.addItem({ - type: MessageType.ERROR, - text: `Invalid number value: ${rawValue}`, - }); - return; + if (schema instanceof z.ZodNumber) { + return Number(raw); } - } else { - value = rawValue; + return raw; + }; + + value = parseValue(rawValue, metadata.schema); + const result = metadata.schema.safeParse(value); + + if (!result.success) { + context.ui.addItem({ + type: MessageType.ERROR, + text: `Invalid value for ${name}: ${rawValue}. Error: ${result.error.errors[0].message}`, + }); + return; } + value = result.data; + const { settings, config } = context.services; if (!config) return; @@ -188,6 +194,7 @@ const unsetExperimentCommand: SlashCommand = { export const experimentCommand: SlashCommand = { name: 'experiment', description: 'Manage experimental features', + hidden: true, kind: CommandKind.BUILT_IN, autoExecute: false, subCommands: [ diff --git a/packages/core/index.ts b/packages/core/index.ts index f26941e4a1..cbfe0ebba0 100644 --- a/packages/core/index.ts +++ b/packages/core/index.ts @@ -49,6 +49,7 @@ export { getExperiments } from './src/code_assist/experiments/experiments.js'; export { ExperimentFlags, ExperimentMetadata, + type ExperimentMetadataEntry, getExperimentFlagName, getExperimentFlagIdFromName, } from './src/code_assist/experiments/flagNames.js'; diff --git a/packages/core/src/code_assist/experiments/flagNames.ts b/packages/core/src/code_assist/experiments/flagNames.ts index 411001aa39..5f692db3c5 100644 --- a/packages/core/src/code_assist/experiments/flagNames.ts +++ b/packages/core/src/code_assist/experiments/flagNames.ts @@ -4,6 +4,8 @@ * SPDX-License-Identifier: Apache-2.0 */ +import { z } from 'zod'; + export const ExperimentFlags = { CONTEXT_COMPRESSION_THRESHOLD: 45740197, USER_CACHING: 45740198, @@ -21,6 +23,20 @@ export const ExperimentFlags = { GEMINI_3_1_FLASH_LITE_LAUNCHED: 45771641, DEFAULT_REQUEST_TIMEOUT: 45773134, ENABLE_AWESOME: 45758820, + + // Migrated from hardcoded experimental settings + ENABLE_AGENTS: 45760001, + EXTENSION_MANAGEMENT: 45760002, + EXTENSION_CONFIG: 45760003, + EXTENSION_REGISTRY: 45760004, + EXTENSION_RELOADING: 45760005, + JIT_CONTEXT: 45760006, + USE_OSC52_PASTE: 45760007, + USE_OSC52_COPY: 45760008, + PLAN: 45760009, + MODEL_STEERING: 45760010, + DISABLE_LLM_CORRECTION: 45760011, + ENABLE_TOOL_OUTPUT_MASKING: 45760012, } as const; export type ExperimentFlagName = @@ -28,69 +44,181 @@ export type ExperimentFlagName = export interface ExperimentMetadataEntry { description: string; - type: 'boolean' | 'number' | 'string'; - defaultValue: boolean | number | string; + schema: z.ZodTypeAny; + defaultValue: unknown; + hidden?: boolean; + /** The key used in settings.json (defaults to kebab-case version of flag name) */ + settingKey?: string; } export const ExperimentMetadata: Record = { [ExperimentFlags.CONTEXT_COMPRESSION_THRESHOLD]: { description: 'Threshold at which context compression activates.', - type: 'number', + schema: z.number(), defaultValue: 0, }, [ExperimentFlags.USER_CACHING]: { description: 'Enables caching of user contexts.', - type: 'boolean', + schema: z.boolean(), defaultValue: false, }, [ExperimentFlags.BANNER_TEXT_NO_CAPACITY_ISSUES]: { description: 'Banner text displayed when there are no capacity issues.', - type: 'string', + schema: z.string(), defaultValue: '', }, [ExperimentFlags.BANNER_TEXT_CAPACITY_ISSUES]: { description: 'Banner text displayed during capacity issues.', - type: 'string', + schema: z.string(), defaultValue: '', }, [ExperimentFlags.ENABLE_PREVIEW]: { description: 'Enables preview features globally.', - type: 'boolean', + schema: z.boolean(), defaultValue: false, }, [ExperimentFlags.ENABLE_NUMERICAL_ROUTING]: { description: 'Enables numerical routing strategies for the model.', - type: 'boolean', + schema: z.boolean(), defaultValue: false, }, [ExperimentFlags.CLASSIFIER_THRESHOLD]: { description: 'Threshold for the intent classifier.', - type: 'number', + schema: z.number(), defaultValue: 0.5, }, [ExperimentFlags.ENABLE_ADMIN_CONTROLS]: { description: 'Enables admin control features in the CLI.', - type: 'boolean', + schema: z.boolean(), defaultValue: false, }, [ExperimentFlags.MASKING_PROTECTION_THRESHOLD]: { description: 'Threshold for masking protection logic.', - type: 'number', + schema: z.number(), defaultValue: 0, + settingKey: 'toolOutputMasking.toolProtectionThreshold', }, [ExperimentFlags.MASKING_PRUNABLE_THRESHOLD]: { description: 'Threshold for prunable masking.', - type: 'number', + schema: z.number(), defaultValue: 0, + settingKey: 'toolOutputMasking.minPrunableTokensThreshold', }, [ExperimentFlags.MASKING_PROTECT_LATEST_TURN]: { description: 'Protects the latest turn from being masked.', - type: 'boolean', + schema: z.boolean(), defaultValue: true, + settingKey: 'toolOutputMasking.protectLatestTurn', + }, + + // Migrated settings (marked hidden to keep /experiment list focused) + [ExperimentFlags.ENABLE_AGENTS]: { + description: 'Enable local and remote subagents.', + schema: z.boolean(), + defaultValue: false, + hidden: true, + settingKey: 'enableAgents', + }, + [ExperimentFlags.EXTENSION_MANAGEMENT]: { + description: 'Enable extension management features.', + schema: z.boolean(), + defaultValue: true, + hidden: true, + settingKey: 'extensionManagement', + }, + [ExperimentFlags.EXTENSION_CONFIG]: { + description: 'Enable requesting and fetching of extension settings.', + schema: z.boolean(), + defaultValue: true, + hidden: true, + settingKey: 'extensionConfig', + }, + [ExperimentFlags.EXTENSION_REGISTRY]: { + description: 'Enable extension registry explore UI.', + schema: z.boolean(), + defaultValue: false, + hidden: true, + settingKey: 'extensionRegistry', + }, + [ExperimentFlags.EXTENSION_RELOADING]: { + description: 'Enables extension loading/unloading within the CLI session.', + schema: z.boolean(), + defaultValue: false, + hidden: true, + settingKey: 'extensionReloading', + }, + [ExperimentFlags.JIT_CONTEXT]: { + description: 'Enable Just-In-Time (JIT) context loading.', + schema: z.boolean(), + defaultValue: false, + hidden: true, + settingKey: 'jitContext', + }, + [ExperimentFlags.USE_OSC52_PASTE]: { + description: 'Use OSC 52 for pasting.', + schema: z.boolean(), + defaultValue: false, + hidden: true, + settingKey: 'useOSC52Paste', + }, + [ExperimentFlags.USE_OSC52_COPY]: { + description: 'Use OSC 52 for copying.', + schema: z.boolean(), + defaultValue: false, + hidden: true, + settingKey: 'useOSC52Copy', + }, + [ExperimentFlags.PLAN]: { + description: 'Enable planning features (Plan Mode and tools).', + schema: z.boolean(), + defaultValue: false, + hidden: true, + settingKey: 'plan', + }, + [ExperimentFlags.MODEL_STEERING]: { + description: 'Enable model steering (user hints).', + schema: z.boolean(), + defaultValue: false, + hidden: true, + settingKey: 'modelSteering', + }, + [ExperimentFlags.DISABLE_LLM_CORRECTION]: { + description: 'Disable LLM-based error correction for edit tools.', + schema: z.boolean(), + defaultValue: true, + hidden: true, + settingKey: 'disableLLMCorrection', + }, + [ExperimentFlags.ENABLE_TOOL_OUTPUT_MASKING]: { + description: 'Enables tool output masking to save tokens.', + schema: z.boolean(), + defaultValue: true, + hidden: true, + settingKey: 'toolOutputMasking.enabled', + }, + [ExperimentFlags.GEMINI_3_1_PRO_LAUNCHED]: { + description: 'Indicates if Gemini 3.1 Pro has been launched.', + schema: z.boolean(), + defaultValue: false, + }, + [ExperimentFlags.PRO_MODEL_NO_ACCESS]: { + description: 'Indicates if the user has no access to Pro models.', + schema: z.boolean(), + defaultValue: false, + }, + [ExperimentFlags.GEMINI_3_1_FLASH_LITE_LAUNCHED]: { + description: 'Indicates if Gemini 3.1 Flash Lite has been launched.', + schema: z.boolean(), + defaultValue: false, + }, + [ExperimentFlags.DEFAULT_REQUEST_TIMEOUT]: { + description: 'The default request timeout in seconds.', + schema: z.number(), + defaultValue: 0, }, [ExperimentFlags.ENABLE_AWESOME]: { description: "When enabled, the ASCII art says 'matt'.", - type: 'boolean', + schema: z.boolean(), defaultValue: false, }, }; @@ -99,6 +227,11 @@ export const ExperimentMetadata: Record = { * Gets the name of an experiment flag from its ID. */ export function getExperimentFlagName(flagId: number): string | undefined { + const metadata = ExperimentMetadata[flagId]; + if (metadata?.settingKey) { + return metadata.settingKey; + } + for (const [name, id] of Object.entries(ExperimentFlags)) { if (id === flagId) { return name.toLowerCase().replace(/_/g, '-'); @@ -111,6 +244,13 @@ export function getExperimentFlagName(flagId: number): string | undefined { * Gets the ID of an experiment flag from its name (supports kebab-case or camelCase). */ export function getExperimentFlagIdFromName(name: string): number | undefined { + // Check metadata for explicit settingKey matches first + for (const [idStr, metadata] of Object.entries(ExperimentMetadata)) { + if (metadata.settingKey === name) { + return parseInt(idStr, 10); + } + } + // Convert enableNumericalRouting or enable-numerical-routing to ENABLE_NUMERICAL_ROUTING const constantName = name .replace(/([a-z])([A-Z])/g, '$1_$2') // camelCase to snake_case diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 753efc0172..578508b0c8 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -128,6 +128,7 @@ import { DEFAULT_MODEL_CONFIGS } from './defaultModelConfigs.js'; import { MemoryContextManager } from '../context/memoryContextManager.js'; import { TrackerService } from '../services/trackerService.js'; import type { GenerateContentParameters } from '@google/genai'; +import { ExperimentManager } from './experimentManager.js'; // Re-export OAuth config type export type { MCPOAuthConfig, AnyToolInvocation, AnyDeclarativeTool }; @@ -155,10 +156,6 @@ import type { } from '../code_assist/types.js'; import type { HierarchicalMemory } from './memory.js'; import { getCodeAssistServer } from '../code_assist/codeAssist.js'; -import { - getExperiments, - type Experiments, -} from '../code_assist/experiments/experiments.js'; import { AgentRegistry } from '../agents/registry.js'; import { AcknowledgedAgentsService } from '../agents/acknowledgedAgents.js'; import { setGlobalProxy, updateGlobalFetchTimeouts } from '../utils/fetch.js'; @@ -168,6 +165,10 @@ import { ExperimentMetadata, getExperimentFlagName, } from '../code_assist/experiments/flagNames.js'; +import { + getExperiments, + type Experiments, +} from '../code_assist/experiments/experiments.js'; import { debugLogger } from '../utils/debugLogger.js'; import { SkillManager, type SkillDefinition } from '../skills/skillManager.js'; import { startupProfiler } from '../telemetry/startupProfiler.js'; @@ -447,7 +448,6 @@ import { DEFAULT_MIN_PRUNABLE_TOKENS_THRESHOLD, DEFAULT_PROTECT_LATEST_TURN, } from '../context/toolOutputMaskingService.js'; - import { type ExtensionLoader, SimpleExtensionLoader, @@ -824,7 +824,6 @@ export class Config implements McpContext, AgentLoopContext { private readonly listExtensions: boolean; private readonly _extensionLoader: ExtensionLoader; private readonly _enabledExtensions: string[]; - private readonly enableExtensionReloading: boolean; fallbackModelHandler?: FallbackModelHandler; validationHandler?: ValidationHandler; private quotaErrorOccurred: boolean = false; @@ -882,6 +881,7 @@ export class Config implements McpContext, AgentLoopContext { private shellExecutionConfig: ShellExecutionConfig; private readonly extensionManagement: boolean = true; private readonly extensionRegistryURI: string | undefined; + private readonly enablePromptCompletion: boolean = false; private readonly truncateToolOutputThreshold: number; private compressionTruncationCounter = 0; private initialized = false; @@ -894,6 +894,7 @@ export class Config implements McpContext, AgentLoopContext { private readonly workspacePoliciesDir: string | undefined; private readonly _messageBus: MessageBus; private readonly policyEngine: PolicyEngine; + private readonly experimentManager: ExperimentManager; private policyUpdateConfirmationRequest: | PolicyUpdateConfirmationRequest | undefined; @@ -918,14 +919,7 @@ export class Config implements McpContext, AgentLoopContext { private pendingIncludeDirectories: string[]; private readonly enableHooksUI: boolean; private readonly enableHooks: boolean; - - private hooks: { [K in HookEventName]?: HookDefinition[] } | undefined; - private projectHooks: - | ({ [K in HookEventName]?: HookDefinition[] } & { disabled?: string[] }) - | undefined; - private disabledHooks: string[]; - private experiments: Experiments | undefined; - private experimentsPromise: Promise | undefined; + private experimentsPromise: Promise | undefined; private hookSystem?: HookSystem; private readonly onModelChange: ((model: string) => void) | undefined; private readonly onReload: @@ -942,8 +936,6 @@ export class Config implements McpContext, AgentLoopContext { private readonly enableAgents: boolean; private agents: AgentSettings; - private experimentalSettings: Record; - private readonly experimentalCliArgs: Record; private readonly enableEventDrivenScheduler: boolean; private readonly skillsSupport: boolean; private disabledSkills: string[]; @@ -959,6 +951,7 @@ export class Config implements McpContext, AgentLoopContext { private readonly modelSteering: boolean; private memoryContextManager?: MemoryContextManager; private readonly contextManagement: ContextManagementConfig; + private contextManager?: ContextManager; private terminalBackground: string | undefined = undefined; private remoteAdminSettings: AdminControlsSettings | undefined; private latestApiRequest: GenerateContentParameters | undefined; @@ -1107,11 +1100,45 @@ export class Config implements McpContext, AgentLoopContext { this.planEnabled = params.plan ?? true; this.trackerEnabled = params.tracker ?? false; this.planModeRoutingEnabled = params.planSettings?.modelRouting ?? true; + + // Merge legacy experimental parameters into a single object for the manager + const experimentalSettings = { + ...(params.experimentalSettings ?? {}), + }; + if (params.plan !== undefined) experimentalSettings['plan'] = params.plan; + if (params.experimentalJitContext !== undefined) + experimentalSettings['jitContext'] = params.experimentalJitContext; + if (params.enableAgents !== undefined) + experimentalSettings['enableAgents'] = params.enableAgents; + if (params.modelSteering !== undefined) + experimentalSettings['modelSteering'] = params.modelSteering; + if (params.enableExtensionReloading !== undefined) + experimentalSettings['extensionReloading'] = + params.enableExtensionReloading; + if (params.extensionManagement !== undefined) + experimentalSettings['extensionManagement'] = params.extensionManagement; + if (params.disableLLMCorrection !== undefined) + experimentalSettings['disableLLMCorrection'] = + params.disableLLMCorrection; + if (params.toolOutputMasking?.enabled !== undefined) { + experimentalSettings['toolOutputMasking'] = { + // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion + ...((experimentalSettings['toolOutputMasking'] as object) ?? {}), + enabled: params.toolOutputMasking.enabled, + }; + } + + this.experimentManager = new ExperimentManager({ + experimentalSettings, + experimentalCliArgs: params.experimentalCliArgs, + experiments: params.experiments, + }); this.enableEventDrivenScheduler = params.enableEventDrivenScheduler ?? true; this.skillsSupport = params.skillsSupport ?? true; this.disabledSkills = params.disabledSkills ?? []; this.adminSkillsEnabled = params.adminSkillsEnabled ?? true; this.modelAvailabilityService = new ModelAvailabilityService(); +<<<<<<< HEAD this.dynamicModelConfiguration = params.dynamicModelConfiguration ?? false; // HACK: The settings loading logic doesn't currently merge the default @@ -1355,7 +1382,6 @@ export class Config implements McpContext, AgentLoopContext { this.projectHooks = params.projectHooks; } - this.experiments = params.experiments; this.onModelChange = params.onModelChange; this.onReload = params.onReload; @@ -1420,7 +1446,7 @@ export class Config implements McpContext, AgentLoopContext { } // Add plans directory to workspace context for plan file storage - if (this.planEnabled) { + if (this.isPlanEnabled()) { const plansDir = this.storage.getPlansDir(); try { await fs.promises.access(plansDir); @@ -1502,9 +1528,15 @@ export class Config implements McpContext, AgentLoopContext { await this.hookSystem.initialize(); } +<<<<<<< HEAD if (this.experimentalJitContext) { this.memoryContextManager = new MemoryContextManager(this); await this.memoryContextManager.refresh(); +======= + if (this.isJitContextEnabled()) { + this.contextManager = new ContextManager(this); + await this.contextManager.refresh(); +>>>>>>> 49ab62d87 (refactor: move experiment logic to ExperimentManager with Zod validation) } await this._geminiClient.initialize(); @@ -1603,9 +1635,11 @@ export class Config implements McpContext, AgentLoopContext { this.setModel(DEFAULT_GEMINI_MODEL_AUTO); } - const adminControlsEnabled = - experiments?.flags[ExperimentFlags.ENABLE_ADMIN_CONTROLS]?.boolValue ?? - false; + // Fetch admin controls + await this.ensureExperimentsLoaded(); + const adminControlsEnabled = !!this.getExperimentValue( + ExperimentFlags.ENABLE_ADMIN_CONTROLS, + ); const adminControls = await fetchAdminControls( codeAssistServer, this.getRemoteAdminSettings(), @@ -1623,8 +1657,8 @@ export class Config implements McpContext, AgentLoopContext { } async getExperimentsAsync(): Promise { - if (this.experiments) { - return this.experiments; + if (this.experimentManager.getExperiments()) { + return this.experimentManager.getExperiments(); } const codeAssistServer = getCodeAssistServer(this); return getExperiments(codeAssistServer); @@ -2331,12 +2365,12 @@ export class Config implements McpContext, AgentLoopContext { } getUserMemory(): string | HierarchicalMemory { - if (this.experimentalJitContext && this.memoryContextManager) { + if (this.isJitContextEnabled() && this.contextManager) { return { - global: this.memoryContextManager.getGlobalMemory(), - extension: this.memoryContextManager.getExtensionMemory(), - project: this.memoryContextManager.getEnvironmentMemory(), - userProjectMemory: this.memoryContextManager.getUserProjectMemory(), + global: this.contextManager.getGlobalMemory(), + extension: this.contextManager.getExtensionMemory(), + project: this.contextManager.getEnvironmentMemory(), + userProjectMemory: this.contextManager.getUserProjectMemory(), }; } return this.userMemory; @@ -2346,8 +2380,8 @@ export class Config implements McpContext, AgentLoopContext { * Refreshes the MCP context, including memory, tools, and system instructions. */ async refreshMcpContext(): Promise { - if (this.experimentalJitContext && this.memoryContextManager) { - await this.memoryContextManager.refresh(); + if (this.isJitContextEnabled() && this.contextManager) { + await this.contextManager.refresh(); } else { const { refreshServerHierarchicalMemory } = await import( '../utils/memoryDiscovery.js' @@ -2422,7 +2456,7 @@ export class Config implements McpContext, AgentLoopContext { } isJitContextEnabled(): boolean { - return this.experimentalJitContext; + return this.experimentManager.isJitContextEnabled(); } isContextManagementEnabled(): boolean { @@ -2458,49 +2492,37 @@ export class Config implements McpContext, AgentLoopContext { } isModelSteeringEnabled(): boolean { - return this.modelSteering; + return this.experimentManager.isModelSteeringEnabled(); + } + + getToolOutputMaskingEnabled(): boolean { + return this.getExperimentValue( + ExperimentFlags.ENABLE_TOOL_OUTPUT_MASKING, + )!; } async getToolOutputMaskingConfig(): Promise { await this.ensureExperimentsLoaded(); - const remoteProtection = - this.experiments?.flags[ExperimentFlags.MASKING_PROTECTION_THRESHOLD] - ?.intValue; - const remotePrunable = - this.experiments?.flags[ExperimentFlags.MASKING_PRUNABLE_THRESHOLD] - ?.intValue; - const remoteProtectLatest = - this.experiments?.flags[ExperimentFlags.MASKING_PROTECT_LATEST_TURN] - ?.boolValue; - - const parsedProtection = remoteProtection - ? parseInt(remoteProtection, 10) - : undefined; - const parsedPrunable = remotePrunable - ? parseInt(remotePrunable, 10) - : undefined; - return { - protectionThresholdTokens: - parsedProtection !== undefined && !isNaN(parsedProtection) - ? parsedProtection - : this.contextManagement.tools.outputMasking - .protectionThresholdTokens, - minPrunableThresholdTokens: - parsedPrunable !== undefined && !isNaN(parsedPrunable) - ? parsedPrunable - : this.contextManagement.tools.outputMasking - .minPrunableThresholdTokens, - protectLatestTurn: - remoteProtectLatest ?? - this.contextManagement.tools.outputMasking.protectLatestTurn, + enabled: this.getExperimentValue( + ExperimentFlags.ENABLE_TOOL_OUTPUT_MASKING, + )!, + toolProtectionThreshold: this.getExperimentValue( + ExperimentFlags.MASKING_PROTECTION_THRESHOLD, + )!, + minPrunableTokensThreshold: this.getExperimentValue( + ExperimentFlags.MASKING_PRUNABLE_THRESHOLD, + )!, + protectLatestTurn: this.getExperimentValue( + ExperimentFlags.MASKING_PROTECT_LATEST_TURN, + )!, }; } getGeminiMdFileCount(): number { - if (this.experimentalJitContext && this.memoryContextManager) { - return this.memoryContextManager.getLoadedPaths().size; + if (this.isJitContextEnabled() && this.contextManager) { + return this.contextManager.getLoadedPaths().size; } return this.geminiMdFileCount; } @@ -2510,8 +2532,8 @@ export class Config implements McpContext, AgentLoopContext { } getGeminiMdFilePaths(): string[] { - if (this.experimentalJitContext && this.memoryContextManager) { - return Array.from(this.memoryContextManager.getLoadedPaths()); + if (this.isJitContextEnabled() && this.contextManager) { + return Array.from(this.contextManager.getLoadedPaths()); } return this.geminiMdFilePaths; } @@ -2609,6 +2631,45 @@ export class Config implements McpContext, AgentLoopContext { } /** + * Synchronizes enter/exit plan mode tools based on current mode. + */ + syncPlanModeTools(): void { + const registry = this.getToolRegistry(); + if (!registry) { + return; + } + const approvalMode = this.getApprovalMode(); + const isPlanMode = approvalMode === ApprovalMode.PLAN; + const isYoloMode = approvalMode === ApprovalMode.YOLO; + + if (isPlanMode) { + if (registry.getTool(ENTER_PLAN_MODE_TOOL_NAME)) { + registry.unregisterTool(ENTER_PLAN_MODE_TOOL_NAME); + } + if (!registry.getTool(EXIT_PLAN_MODE_TOOL_NAME)) { + registry.registerTool(new ExitPlanModeTool(this, this.messageBus)); + } + } else { + if (registry.getTool(EXIT_PLAN_MODE_TOOL_NAME)) { + registry.unregisterTool(EXIT_PLAN_MODE_TOOL_NAME); + } + if (this.isPlanEnabled() && !isYoloMode) { + if (!registry.getTool(ENTER_PLAN_MODE_TOOL_NAME)) { + registry.registerTool(new EnterPlanModeTool(this, this.messageBus)); + } + } else { + if (registry.getTool(ENTER_PLAN_MODE_TOOL_NAME)) { + registry.unregisterTool(ENTER_PLAN_MODE_TOOL_NAME); + } + } + } + + if (this.geminiClient?.isInitialized()) { + this.geminiClient.setTools().catch((err) => { + debugLogger.error('Failed to update tools', err); + }); + } + } * Logs the duration of the current approval mode. */ logCurrentModeDuration(mode: ApprovalMode): void { @@ -2840,7 +2901,9 @@ export class Config implements McpContext, AgentLoopContext { } getExtensionManagement(): boolean { - return this.extensionManagement; + return this.experimentManager.getExperimentValue( + ExperimentFlags.EXTENSION_MANAGEMENT, + ); } getExtensions(): GeminiCLIExtension[] { @@ -2858,15 +2921,15 @@ export class Config implements McpContext, AgentLoopContext { } getEnableExtensionReloading(): boolean { - return this.enableExtensionReloading; + return this.experimentManager.getEnableExtensionReloading(); } getDisableLLMCorrection(): boolean { - return this.disableLLMCorrection; + return this.experimentManager.getDisableLLMCorrection(); } isPlanEnabled(): boolean { - return this.planEnabled; + return this.experimentManager.isPlanEnabled(); } isTrackerEnabled(): boolean { @@ -2886,7 +2949,7 @@ export class Config implements McpContext, AgentLoopContext { } isAgentsEnabled(): boolean { - return this.enableAgents; + return this.experimentManager.isAgentsEnabled(); } isEventDrivenSchedulerEnabled(): boolean { @@ -3011,9 +3074,9 @@ export class Config implements McpContext, AgentLoopContext { await this.ensureExperimentsLoaded(); - const remoteThreshold = - this.experiments?.flags[ExperimentFlags.CONTEXT_COMPRESSION_THRESHOLD] - ?.floatValue; + const remoteThreshold = this.getExperimentValue( + ExperimentFlags.CONTEXT_COMPRESSION_THRESHOLD, + ); if (remoteThreshold === 0) { return undefined; } @@ -3021,7 +3084,7 @@ export class Config implements McpContext, AgentLoopContext { } async getUserCaching(): Promise { - return this.getExperimentValue(ExperimentFlags.USER_CACHING); + return this.experimentManager.getUserCaching(); } async getPlanModeRoutingEnabled(): Promise { @@ -3029,11 +3092,7 @@ export class Config implements McpContext, AgentLoopContext { } async getNumericalRoutingEnabled(): Promise { - return ( - this.getExperimentValue( - ExperimentFlags.ENABLE_NUMERICAL_ROUTING, - ) ?? false - ); + return this.experimentManager.isNumericalRoutingEnabled(); } /** @@ -3058,25 +3117,15 @@ export class Config implements McpContext, AgentLoopContext { } async getClassifierThreshold(): Promise { - return this.getExperimentValue( - ExperimentFlags.CLASSIFIER_THRESHOLD, - ); + return this.experimentManager.getClassifierThreshold(); } async getBannerTextNoCapacityIssues(): Promise { - return ( - this.getExperimentValue( - ExperimentFlags.BANNER_TEXT_NO_CAPACITY_ISSUES, - ) ?? '' - ); + return this.experimentManager.getBannerTextNoCapacityIssues(); } async getBannerTextCapacityIssues(): Promise { - return ( - this.getExperimentValue( - ExperimentFlags.BANNER_TEXT_CAPACITY_ISSUES, - ) ?? '' - ); + return this.experimentManager.getBannerTextCapacityIssues(); } /** @@ -3685,7 +3734,7 @@ export class Config implements McpContext, AgentLoopContext { * Get experiments configuration */ getExperiments(): Experiments | undefined { - return this.experiments; + return this.experimentManager.getExperiments(); } /** @@ -3698,69 +3747,21 @@ export class Config implements McpContext, AgentLoopContext { getExperimentValue( flagId: number, ): T | undefined { - const flagName = getExperimentFlagName(flagId); - if (!flagName) { - return undefined; - } - - // 1. Command-line argument - const cliValue = this.experimentalCliArgs[flagName]; - if (cliValue !== undefined) { - // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion - return cliValue as T; - } - - // 2. Local setting - const settingValue = this.experimentalSettings[flagName]; - if (settingValue !== undefined) { - // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion - return settingValue as T; - } - - // 3. Remote experiment - const remoteFlag = this.experiments?.flags[flagId]; - if (remoteFlag) { - const val = - remoteFlag.boolValue ?? - remoteFlag.floatValue ?? - remoteFlag.intValue ?? - remoteFlag.stringValue; - if (val !== undefined) { - // Handle string representation of numbers if necessary - if (typeof val === 'string' && !isNaN(Number(val))) { - const metadata = ExperimentMetadata[flagId]; - if (metadata?.type === 'number') { - // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion - return Number(val) as unknown as T; - } - } - // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion - return val as unknown as T; - } - } - - // 4. Default value from metadata - // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion - return ExperimentMetadata[flagId]?.defaultValue as unknown as T; + return this.experimentManager.getExperimentValue(flagId); } /** * Updates experimental settings. */ updateExperimentalSettings(settings: Record): void { - // Only update if settings have actually changed to avoid unnecessary re-initialization logic - // if we add any in the future. - this.experimentalSettings = { - ...this.experimentalSettings, - ...settings, - }; + this.experimentManager.updateExperimentalSettings(settings); } /** * Set experiments configuration */ setExperiments(experiments: Experiments): void { - this.experiments = experiments; + this.experimentManager.setExperiments(experiments); const flagSummaries = Object.entries(experiments.flags ?? {}) .sort(([a], [b]) => a.localeCompare(b)) .map(([flagId, flag]) => { diff --git a/packages/core/src/config/experimentManager.test.ts b/packages/core/src/config/experimentManager.test.ts new file mode 100644 index 0000000000..4e7c70eb2d --- /dev/null +++ b/packages/core/src/config/experimentManager.test.ts @@ -0,0 +1,134 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect } from 'vitest'; +import { ExperimentManager } from './experimentManager.js'; +import { ExperimentFlags } from '../code_assist/experiments/flagNames.js'; + +describe('ExperimentManager', () => { + const baseOptions = { + experimentalSettings: {}, + experimentalCliArgs: {}, + }; + + describe('getExperimentValue', () => { + it('should return default value when no overrides are present', () => { + const manager = new ExperimentManager(baseOptions); + // USER_CACHING default is false + expect(manager.getExperimentValue(ExperimentFlags.USER_CACHING)).toBe( + false, + ); + }); + + it('should prioritize CLI arguments over all else', () => { + const manager = new ExperimentManager({ + ...baseOptions, + experimentalCliArgs: { 'user-caching': true }, + experimentalSettings: { 'user-caching': false }, + experiments: { + flags: { + [ExperimentFlags.USER_CACHING]: { boolValue: false }, + }, + experimentIds: [], + }, + }); + expect(manager.getExperimentValue(ExperimentFlags.USER_CACHING)).toBe( + true, + ); + }); + + it('should prioritize local settings over remote experiments', () => { + const manager = new ExperimentManager({ + ...baseOptions, + experimentalSettings: { 'user-caching': true }, + experiments: { + flags: { + [ExperimentFlags.USER_CACHING]: { boolValue: false }, + }, + experimentIds: [], + }, + }); + expect(manager.getExperimentValue(ExperimentFlags.USER_CACHING)).toBe( + true, + ); + }); + + it('should use remote experiment if no local override', () => { + const manager = new ExperimentManager({ + ...baseOptions, + experiments: { + flags: { + [ExperimentFlags.USER_CACHING]: { boolValue: true }, + }, + experimentIds: [], + }, + }); + expect(manager.getExperimentValue(ExperimentFlags.USER_CACHING)).toBe( + true, + ); + }); + + it('should handle nested settings correctly', () => { + const manager = new ExperimentManager({ + ...baseOptions, + experimentalSettings: { + toolOutputMasking: { + enabled: false, + }, + }, + }); + expect( + manager.getExperimentValue(ExperimentFlags.ENABLE_TOOL_OUTPUT_MASKING), + ).toBe(false); + }); + + it('should validate values using Zod schema', () => { + const manager = new ExperimentManager({ + ...baseOptions, + experimentalSettings: { + 'classifier-threshold': 'not-a-number', + }, + }); + // CLASSIFIER_THRESHOLD default is 0.5 + expect( + manager.getExperimentValue(ExperimentFlags.CLASSIFIER_THRESHOLD), + ).toBe(0.5); + }); + + it('should parse numeric strings from remote flags if metadata type is number', () => { + const manager = new ExperimentManager({ + ...baseOptions, + experiments: { + flags: { + [ExperimentFlags.CLASSIFIER_THRESHOLD]: { stringValue: '0.8' }, + }, + experimentIds: [], + }, + }); + expect( + manager.getExperimentValue(ExperimentFlags.CLASSIFIER_THRESHOLD), + ).toBe(0.8); + }); + }); + + describe('convenience getters', () => { + it('isPlanEnabled should resolve correctly', () => { + const manager = new ExperimentManager({ + ...baseOptions, + experimentalSettings: { plan: true }, + }); + expect(manager.isPlanEnabled()).toBe(true); + }); + + it('isAgentsEnabled should resolve correctly', () => { + const manager = new ExperimentManager({ + ...baseOptions, + experimentalSettings: { enableAgents: true }, + }); + expect(manager.isAgentsEnabled()).toBe(true); + }); + }); +}); diff --git a/packages/core/src/config/experimentManager.ts b/packages/core/src/config/experimentManager.ts new file mode 100644 index 0000000000..1ed8954b30 --- /dev/null +++ b/packages/core/src/config/experimentManager.ts @@ -0,0 +1,224 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { + ExperimentFlags, + ExperimentMetadata, + getExperimentFlagName, + type ExperimentMetadataEntry, +} from '../code_assist/experiments/flagNames.js'; +import type { Experiments } from '../code_assist/experiments/experiments.js'; +import { debugLogger } from '../utils/debugLogger.js'; +import { z } from 'zod'; + +export interface ExperimentManagerOptions { + experimentalSettings?: Record; + experimentalCliArgs?: Record; + experiments?: Experiments; +} + +/** + * Manages resolution and validation of experimental flags. + */ +export class ExperimentManager { + private experimentalSettings: Record; + private readonly experimentalCliArgs: Record; + private experiments?: Experiments; + + constructor(options: ExperimentManagerOptions) { + this.experimentalSettings = options.experimentalSettings ?? {}; + this.experimentalCliArgs = options.experimentalCliArgs ?? {}; + this.experiments = options.experiments; + } + + /** + * Resolves the value of an experiment flag, applying layering logic: + * 1. Command-line argument + * 2. Local setting (settings.json) + * 3. Remote experiment (server-side) + * 4. Default value (from metadata) + */ + getExperimentValue(flagId: number): T { + const metadata = ExperimentMetadata[flagId]; + if (!metadata) { + debugLogger.warn(`Unknown experiment flag ID: ${flagId}`); + // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion + return undefined as unknown as T; + } + + const flagName = getExperimentFlagName(flagId); + + // 1. CLI Argument + if (flagName && this.experimentalCliArgs[flagName] !== undefined) { + let val: unknown = this.experimentalCliArgs[flagName]; + // Type coercion for CLI args + val = this.coerceValue(val, metadata); + + const result = metadata.schema.safeParse(val); + if (result.success) { + // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion + return result.data as T; + } + debugLogger.warn( + `Invalid CLI value for ${flagName}: ${val}. Error: ${result.error.message}`, + ); + } + + // 2. Local Setting (settings.json) + const settingKey = metadata.settingKey || flagName; + if (settingKey) { + const val = this.getNestedValue(this.experimentalSettings, settingKey); + if (val !== undefined) { + const result = metadata.schema.safeParse(val); + if (result.success) { + // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion + return result.data as T; + } + debugLogger.warn( + `Invalid local setting for ${settingKey}: ${val}. Error: ${result.error.message}`, + ); + } + } + + // 3. Remote Experiment + const remoteFlag = this.experiments?.flags[flagId]; + if (remoteFlag) { + let val: unknown = + remoteFlag.boolValue ?? + remoteFlag.floatValue ?? + (remoteFlag.intValue ? Number(remoteFlag.intValue) : undefined) ?? + remoteFlag.stringValue; + + if (val !== undefined) { + val = this.coerceValue(val, metadata); + const result = metadata.schema.safeParse(val); + if (result.success) { + // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion + return result.data as T; + } + debugLogger.warn( + `Invalid remote value for flag ${flagId}: ${val}. Error: ${result.error.message}`, + ); + } + } + + // 4. Default Value + // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion + return metadata.defaultValue as T; + } + + private coerceValue( + val: unknown, + metadata: ExperimentMetadataEntry, + ): unknown { + if (metadata.schema instanceof z.ZodNumber && typeof val === 'string') { + const num = Number(val); + if (!isNaN(num)) return num; + } + if (metadata.schema instanceof z.ZodBoolean && typeof val === 'string') { + if (val === 'true' || val === 'on') return true; + if (val === 'false' || val === 'off') return false; + } + return val; + } + + private getNestedValue(obj: Record, path: string): unknown { + const parts = path.split('.'); + let current: unknown = obj; + for (const part of parts) { + if (current === null || typeof current !== 'object') return undefined; + // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion + current = (current as Record)[part]; + } + return current; + } + + /** + * Updates the local experimental settings. + */ + updateExperimentalSettings(settings: Record): void { + this.experimentalSettings = { + ...this.experimentalSettings, + ...settings, + }; + } + + /** + * Updates remote experiments. + */ + setExperiments(experiments: Experiments): void { + this.experiments = experiments; + } + + /** + * Gets all experimental settings (for serialization/display). + */ + getExperimentalSettings(): Record { + return this.experimentalSettings; + } + + getExperiments(): Experiments | undefined { + return this.experiments; + } + + // Convenience getters for commonly used flags + + isPlanEnabled(): boolean { + return this.getExperimentValue(ExperimentFlags.PLAN); + } + + isAgentsEnabled(): boolean { + return this.getExperimentValue(ExperimentFlags.ENABLE_AGENTS); + } + + getEnableExtensionReloading(): boolean { + return this.getExperimentValue( + ExperimentFlags.EXTENSION_RELOADING, + ); + } + + getDisableLLMCorrection(): boolean { + return this.getExperimentValue( + ExperimentFlags.DISABLE_LLM_CORRECTION, + ); + } + + isModelSteeringEnabled(): boolean { + return this.getExperimentValue(ExperimentFlags.MODEL_STEERING); + } + + isJitContextEnabled(): boolean { + return this.getExperimentValue(ExperimentFlags.JIT_CONTEXT); + } + + isNumericalRoutingEnabled(): boolean { + return this.getExperimentValue( + ExperimentFlags.ENABLE_NUMERICAL_ROUTING, + ); + } + + getClassifierThreshold(): number | undefined { + return this.getExperimentValue( + ExperimentFlags.CLASSIFIER_THRESHOLD, + ); + } + + getUserCaching(): boolean { + return this.getExperimentValue(ExperimentFlags.USER_CACHING); + } + + getBannerTextNoCapacityIssues(): string { + return this.getExperimentValue( + ExperimentFlags.BANNER_TEXT_NO_CAPACITY_ISSUES, + ); + } + + getBannerTextCapacityIssues(): string { + return this.getExperimentValue( + ExperimentFlags.BANNER_TEXT_CAPACITY_ISSUES, + ); + } +}