Skip to content
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions .gemini/settings.json
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
{
"experimental": {
"toolOutputMasking": {
"enabled": true
}
}
}
1 change: 1 addition & 0 deletions packages/cli/src/config/config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -806,6 +806,7 @@ export async function loadCliConfig(
skillsSupport: settings.skills?.enabled ?? true,
disabledSkills: settings.skills?.disabled,
experimentalJitContext: settings.experimental?.jitContext,
toolOutputMasking: settings.experimental?.toolOutputMasking,
noBrowser: !!process.env['NO_BROWSER'],
summarizeToolOutput: settings.model?.summarizeToolOutput,
ideMode,
Expand Down
52 changes: 52 additions & 0 deletions packages/cli/src/config/settingsSchema.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1462,6 +1462,58 @@ const SETTINGS_SCHEMA = {
description: 'Setting to enable experimental features',
showInDialog: false,
properties: {
toolOutputMasking: {
type: 'object',
label: 'Tool Output Masking',
category: 'Experimental',
requiresRestart: true,
ignoreInDocs: true,
default: {},
description:
'Advanced settings for tool output masking to manage context window efficiency.',
showInDialog: false,
properties: {
enabled: {
type: 'boolean',
label: 'Enable Tool Output Masking',
category: 'Experimental',
requiresRestart: true,
default: false,
description: 'Enables tool output masking to save tokens.',
showInDialog: false,
},
toolProtectionThreshold: {
type: 'number',
label: 'Tool Protection Threshold',
category: 'Experimental',
requiresRestart: true,
default: 50000,
description:
'Minimum number of tokens to protect from masking (most recent tool outputs).',
showInDialog: false,
},
minPrunableTokensThreshold: {
type: 'number',
label: 'Min Prunable Tokens Threshold',
category: 'Experimental',
requiresRestart: true,
default: 30000,
description:
'Minimum prunable tokens required to trigger a masking pass.',
showInDialog: false,
},
protectLatestTurn: {
type: 'boolean',
label: 'Protect Latest Turn',
category: 'Experimental',
requiresRestart: true,
default: true,
description:
'Ensures the absolute latest turn is never masked, regardless of token count.',
showInDialog: false,
},
},
},
enableAgents: {
type: 'boolean',
label: 'Enable Agents',
Expand Down
34 changes: 34 additions & 0 deletions packages/core/src/config/config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -149,6 +149,13 @@ export interface OutputSettings {
format?: OutputFormat;
}

export interface ToolOutputMaskingConfig {
enabled: boolean;
toolProtectionThreshold: number;
minPrunableTokensThreshold: number;
protectLatestTurn: boolean;
}

export interface ExtensionSetting {
name: string;
description: string;
Expand Down Expand Up @@ -273,6 +280,11 @@ import {
DEFAULT_FILE_FILTERING_OPTIONS,
DEFAULT_MEMORY_FILE_FILTERING_OPTIONS,
} from './constants.js';
import {
DEFAULT_TOOL_PROTECTION_THRESHOLD,
DEFAULT_MIN_PRUNABLE_TOKENS_THRESHOLD,
DEFAULT_PROTECT_LATEST_TURN,
} from '../services/toolOutputMaskingService.js';

import {
type ExtensionLoader,
Expand Down Expand Up @@ -462,6 +474,7 @@ export interface ConfigParameters {
disabledSkills?: string[];
adminSkillsEnabled?: boolean;
experimentalJitContext?: boolean;
toolOutputMasking?: Partial<ToolOutputMaskingConfig>;
disableLLMCorrection?: boolean;
plan?: boolean;
onModelChange?: (model: string) => void;
Expand Down Expand Up @@ -599,6 +612,7 @@ export class Config {
private pendingIncludeDirectories: string[];
private readonly enableHooks: boolean;
private readonly enableHooksUI: boolean;
private readonly toolOutputMasking: ToolOutputMaskingConfig;
private hooks: { [K in HookEventName]?: HookDefinition[] } | undefined;
private projectHooks:
| ({ [K in HookEventName]?: HookDefinition[] } & { disabled?: string[] })
Expand Down Expand Up @@ -721,6 +735,18 @@ export class Config {
this.modelAvailabilityService = new ModelAvailabilityService();
this.previewFeatures = params.previewFeatures ?? undefined;
this.experimentalJitContext = params.experimentalJitContext ?? false;
this.toolOutputMasking = {
enabled: params.toolOutputMasking?.enabled ?? false,
toolProtectionThreshold:
params.toolOutputMasking?.toolProtectionThreshold ??
DEFAULT_TOOL_PROTECTION_THRESHOLD,
minPrunableTokensThreshold:
params.toolOutputMasking?.minPrunableTokensThreshold ??
DEFAULT_MIN_PRUNABLE_TOKENS_THRESHOLD,
protectLatestTurn:
params.toolOutputMasking?.protectLatestTurn ??
DEFAULT_PROTECT_LATEST_TURN,
};
this.maxSessionTurns = params.maxSessionTurns ?? -1;
this.experimentalZedIntegration =
params.experimentalZedIntegration ?? false;
Expand Down Expand Up @@ -1445,6 +1471,14 @@ export class Config {
return this.experimentalJitContext;
}

getToolOutputMaskingEnabled(): boolean {
return this.toolOutputMasking.enabled;
}

getToolOutputMaskingConfig(): ToolOutputMaskingConfig {
return this.toolOutputMasking;
}

getGeminiMdFileCount(): number {
if (this.experimentalJitContext && this.contextManager) {
return this.contextManager.getLoadedPaths().size;
Expand Down
1 change: 1 addition & 0 deletions packages/core/src/core/client.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -213,6 +213,7 @@ describe('Gemini Client (client.ts)', () => {
getGlobalMemory: vi.fn().mockReturnValue(''),
getEnvironmentMemory: vi.fn().mockReturnValue(''),
isJitContextEnabled: vi.fn().mockReturnValue(false),
getToolOutputMaskingEnabled: vi.fn().mockReturnValue(false),
getDisableLoopDetection: vi.fn().mockReturnValue(false),

getSessionId: vi.fn().mockReturnValue('test-session-id'),
Expand Down
21 changes: 21 additions & 0 deletions packages/core/src/core/client.ts
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@ import { handleFallback } from '../fallback/handler.js';
import type { RoutingContext } from '../routing/routingStrategy.js';
import { debugLogger } from '../utils/debugLogger.js';
import type { ModelConfigKey } from '../services/modelConfigService.js';
import { ToolOutputMaskingService } from '../services/toolOutputMaskingService.js';
import { calculateRequestTokenCount } from '../utils/tokenCalculation.js';
import {
applyModelSelection,
Expand Down Expand Up @@ -84,6 +85,7 @@ export class GeminiClient {

private readonly loopDetector: LoopDetectionService;
private readonly compressionService: ChatCompressionService;
private readonly toolOutputMaskingService: ToolOutputMaskingService;
private lastPromptId: string;
private currentSequenceModel: string | null = null;
private lastSentIdeContext: IdeContext | undefined;
Expand All @@ -98,6 +100,7 @@ export class GeminiClient {
constructor(private readonly config: Config) {
this.loopDetector = new LoopDetectionService(config);
this.compressionService = new ChatCompressionService();
this.toolOutputMaskingService = new ToolOutputMaskingService();
this.lastPromptId = this.config.getSessionId();

coreEvents.on(CoreEvent.ModelChanged, this.handleModelChanged);
Expand Down Expand Up @@ -562,6 +565,8 @@ export class GeminiClient {
const remainingTokenCount =
tokenLimit(modelForLimitCheck) - this.getChat().getLastPromptTokenCount();

await this.tryMaskToolOutputs(this.getHistory());

// Estimate tokens. For text-only requests, we estimate based on character length.
// For requests with non-text parts (like images, tools), we use the countTokens API.
const estimatedRequestTokenCount = await calculateRequestTokenCount(
Expand Down Expand Up @@ -1056,4 +1061,20 @@ export class GeminiClient {

return info;
}

/**
* Masks bulky tool outputs to save context window space.
*/
private async tryMaskToolOutputs(history: Content[]): Promise<void> {
if (!this.config.getToolOutputMaskingEnabled()) {
return;
}
const result = await this.toolOutputMaskingService.mask(
history,
this.config,
);
if (result.maskedCount > 0) {
this.getChat().setHistory(result.newHistory);
}
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html

exports[`ToolOutputMaskingService > should match the expected snapshot for a masked tool output 1`] = `
"<tool_output_masked>
Line
Line
Line
Line
Line
Line
Line
Line
Line
Line

... [6 lines omitted] ...

Line
Line
Line
Line
Line
Line
Line
Line
Line


Output too large. Full output available at: /mock/history/tool-outputs/run_shell_command_deterministic.txt
</tool_output_masked>"
`;
Loading
Loading