Skip to content
Merged
Show file tree
Hide file tree
Changes from all 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
3 changes: 3 additions & 0 deletions packages/core/src/code_assist/experiments/flagNames.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 =
Expand Down
35 changes: 33 additions & 2 deletions packages/core/src/config/config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1433,8 +1433,39 @@ export class Config {
return this.toolOutputMasking.enabled;
}

getToolOutputMaskingConfig(): ToolOutputMaskingConfig {
return this.toolOutputMasking;
async getToolOutputMaskingConfig(): Promise<ToolOutputMaskingConfig> {
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 {
Expand Down
40 changes: 39 additions & 1 deletion packages/core/src/services/toolOutputMaskingService.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,7 @@ describe('ToolOutputMaskingService', () => {
getSessionId: () => 'mock-session',
getUsageStatisticsEnabled: () => false,
getToolOutputMaskingEnabled: () => true,
getToolOutputMaskingConfig: () => ({
getToolOutputMaskingConfig: async () => ({
enabled: true,
toolProtectionThreshold: 50000,
minPrunableTokensThreshold: 30000,
Expand All @@ -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[] = [
{
Expand Down
5 changes: 2 additions & 3 deletions packages/core/src/services/toolOutputMaskingService.ts
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,8 @@ export interface MaskingResult {
*/
export class ToolOutputMaskingService {
async mask(history: Content[], config: Config): Promise<MaskingResult> {
if (history.length === 0) {
const maskingConfig = await config.getToolOutputMaskingConfig();
if (!maskingConfig.enabled || history.length === 0) {
return { newHistory: history, maskedCount: 0, tokensSaved: 0 };
}

Expand All @@ -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
Expand Down
Loading