diff --git a/packages/core/src/code_assist/experiments/flagNames.ts b/packages/core/src/code_assist/experiments/flagNames.ts index ba26b68cc2..03b6aaac0a 100644 --- a/packages/core/src/code_assist/experiments/flagNames.ts +++ b/packages/core/src/code_assist/experiments/flagNames.ts @@ -13,6 +13,9 @@ export const ExperimentFlags = { ENABLE_NUMERICAL_ROUTING: 45750526, CLASSIFIER_THRESHOLD: 45750527, ENABLE_ADMIN_CONTROLS: 45752213, + MASKING_PROTECTION_THRESHOLD: 45758817, + MASKING_PRUNABLE_THRESHOLD: 45758818, + MASKING_PROTECT_LATEST_TURN: 45758819, } as const; export type ExperimentFlagName = diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 48f81d081f..4df65f51a2 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -1433,8 +1433,39 @@ export class Config { return this.toolOutputMasking.enabled; } - getToolOutputMaskingConfig(): ToolOutputMaskingConfig { - return this.toolOutputMasking; + 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 { + enabled: this.toolOutputMasking.enabled, + toolProtectionThreshold: + parsedProtection !== undefined && !isNaN(parsedProtection) + ? parsedProtection + : this.toolOutputMasking.toolProtectionThreshold, + minPrunableTokensThreshold: + parsedPrunable !== undefined && !isNaN(parsedPrunable) + ? parsedPrunable + : this.toolOutputMasking.minPrunableTokensThreshold, + protectLatestTurn: + remoteProtectLatest ?? this.toolOutputMasking.protectLatestTurn, + }; } getGeminiMdFileCount(): number { diff --git a/packages/core/src/services/toolOutputMaskingService.test.ts b/packages/core/src/services/toolOutputMaskingService.test.ts index 08d8187ff3..1187a28ae1 100644 --- a/packages/core/src/services/toolOutputMaskingService.test.ts +++ b/packages/core/src/services/toolOutputMaskingService.test.ts @@ -46,7 +46,7 @@ describe('ToolOutputMaskingService', () => { getSessionId: () => 'mock-session', getUsageStatisticsEnabled: () => false, getToolOutputMaskingEnabled: () => true, - getToolOutputMaskingConfig: () => ({ + getToolOutputMaskingConfig: async () => ({ enabled: true, toolProtectionThreshold: 50000, minPrunableTokensThreshold: 30000, @@ -63,6 +63,44 @@ describe('ToolOutputMaskingService', () => { } }); + it('should respect remote configuration overrides', async () => { + mockConfig.getToolOutputMaskingConfig = async () => ({ + enabled: true, + toolProtectionThreshold: 100, // Very low threshold + minPrunableTokensThreshold: 50, + protectLatestTurn: false, + }); + + const history: Content[] = [ + { + role: 'user', + parts: [ + { + functionResponse: { + name: 'test_tool', + response: { output: 'A'.repeat(200) }, + }, + }, + ], + }, + ]; + + mockedEstimateTokenCountSync.mockImplementation((parts) => { + const resp = parts[0].functionResponse?.response as Record< + string, + unknown + >; + const content = (resp?.['output'] as string) ?? JSON.stringify(resp); + return content.includes(MASKING_INDICATOR_TAG) ? 10 : 200; + }); + + const result = await service.mask(history, mockConfig); + + // With low thresholds and protectLatestTurn=false, it should mask even the latest turn + expect(result.maskedCount).toBe(1); + expect(result.tokensSaved).toBeGreaterThan(0); + }); + it('should not mask if total tool tokens are below protection threshold', async () => { const history: Content[] = [ { diff --git a/packages/core/src/services/toolOutputMaskingService.ts b/packages/core/src/services/toolOutputMaskingService.ts index 53804a1909..5c7ff3500b 100644 --- a/packages/core/src/services/toolOutputMaskingService.ts +++ b/packages/core/src/services/toolOutputMaskingService.ts @@ -68,7 +68,8 @@ export interface MaskingResult { */ export class ToolOutputMaskingService { async mask(history: Content[], config: Config): Promise { - if (history.length === 0) { + const maskingConfig = await config.getToolOutputMaskingConfig(); + if (!maskingConfig.enabled || history.length === 0) { return { newHistory: history, maskedCount: 0, tokensSaved: 0 }; } @@ -85,8 +86,6 @@ export class ToolOutputMaskingService { originalPart: Part; }> = []; - const maskingConfig = config.getToolOutputMaskingConfig(); - // Decide where to start scanning. // If PROTECT_LATEST_TURN is true, we skip the most recent message (index history.length - 1). const scanStartIdx = maskingConfig.protectLatestTurn