mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-07-24 00:30:59 -07:00
Compare commits
7 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 matt korwel
|
221849db32 | ||
|
Matt Korwel
|
cbe727352c | ||
|
matt korwel
|
53068fad5c | ||
|
mkorwel
|
ab717bdb24 | ||
|
mkorwel
|
173cb56643 | ||
|
mkorwel
|
c7367f485a | ||
|
mkorwel
|
31ebfb496a |
@@ -30,6 +30,7 @@ they appear in the UI.
|
||||
| Enable Notifications | `general.enableNotifications` | Enable run-event notifications for action-required prompts and session completion. Currently macOS only. | `false` |
|
||||
| Plan Directory | `general.plan.directory` | The directory where planning artifacts are stored. If not specified, defaults to the system temporary directory. | `undefined` |
|
||||
| Enable Prompt Completion | `general.enablePromptCompletion` | Enable AI-powered prompt completion suggestions while typing. | `false` |
|
||||
| Tool Preselection | `general.toolPreselection` | Exclude unneeded tools from context to save tokens and improve performance. | `true` |
|
||||
| Debug Keystroke Logging | `general.debugKeystrokeLogging` | Enable debug logging of keystrokes to the console. | `false` |
|
||||
| Enable Session Cleanup | `general.sessionRetention.enabled` | Enable automatic session cleanup | `false` |
|
||||
| Keep chat history | `general.sessionRetention.maxAge` | Automatically delete chats older than this time period (e.g., "30d", "7d", "24h", "1w") | `undefined` |
|
||||
|
||||
@@ -148,6 +148,11 @@ their corresponding top-level category object in your `settings.json` file.
|
||||
request" errors.
|
||||
- **Default:** `false`
|
||||
|
||||
- **`general.toolPreselection`** (boolean):
|
||||
- **Description:** Exclude unneeded tools from context to save tokens and
|
||||
improve performance.
|
||||
- **Default:** `true`
|
||||
|
||||
- **`general.debugKeystrokeLogging`** (boolean):
|
||||
- **Description:** Enable debug logging of keystrokes to the console.
|
||||
- **Default:** `false`
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
{"method":"generateContent","response":{"candidates":[{"content":{"parts":[{"text":"{\"relevant_tools\": [\"list_directory\"]}"}]}}],"usageMetadata":{"promptTokenCount":100,"totalTokenCount":110}}}
|
||||
{"method":"generateContentStream","response":[{"candidates":[{"content":{"parts":[{"functionCall":{"name":"list_directory","args":{"dir_path":"."}}}]},"index":0}],"usageMetadata":{"promptTokenCount":100,"totalTokenCount":110}}]}
|
||||
{"method":"generateContentStream","response":[{"candidates":[{"finishReason":"STOP","content":{"parts":[{"text":"I listed the files."}]},"index":0}],"usageMetadata":{"promptTokenCount":200,"totalTokenCount":210}}]}
|
||||
@@ -0,0 +1,107 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { expect, describe, it, beforeEach, afterEach } from 'vitest';
|
||||
import { TestRig } from './test-helper.js';
|
||||
import { join } from 'node:path';
|
||||
|
||||
describe('Tool Preselection Integration', () => {
|
||||
let rig: TestRig;
|
||||
|
||||
beforeEach(() => {
|
||||
rig = new TestRig();
|
||||
});
|
||||
|
||||
afterEach(async () => {
|
||||
if (rig) {
|
||||
await rig.cleanup();
|
||||
}
|
||||
});
|
||||
|
||||
it('should perform tool pre-selection correctly', async () => {
|
||||
rig.setup('tool-preselection-v2', {
|
||||
fakeResponsesPath: join(
|
||||
import.meta.dirname,
|
||||
'tool-preselection.responses',
|
||||
),
|
||||
settings: {
|
||||
general: {
|
||||
toolPreselection: true,
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
const result = await rig.run({
|
||||
args: 'Please list the files in the current directory.',
|
||||
});
|
||||
|
||||
// Verify it called list_directory as mocked
|
||||
expect(result).toContain('I listed the files.');
|
||||
|
||||
// Wait for telemetry to flush
|
||||
await rig.waitForTelemetryEvent('api_request');
|
||||
|
||||
const logs = rig.readTelemetryLogs();
|
||||
|
||||
// 1st request: Tool pre-selection (classifier model)
|
||||
// 2nd request: Main agent call (with filtered tools)
|
||||
// 3rd request: Final response
|
||||
|
||||
const apiRequests = logs.filter(
|
||||
(l) => l.attributes?.['event.name'] === 'gemini_cli.api_request',
|
||||
);
|
||||
|
||||
// Find the request from the main agent loop (not the classifier)
|
||||
// Classifier request will have prompt_id: 'tool-preselection'
|
||||
const agentRequest = apiRequests.find(
|
||||
(l) =>
|
||||
l.attributes?.prompt_id?.includes('########') &&
|
||||
l.attributes?.prompt_id?.includes('agent'),
|
||||
);
|
||||
|
||||
if (agentRequest) {
|
||||
// The prompt text is available in agentRequest.attributes.request_text
|
||||
// In the real code, tools are sent in the GenerateContentConfig, but
|
||||
// ApiRequestEvent logs the whole contents which might not show tools.
|
||||
// Wait, let's look at ApiRequestEvent constructor again.
|
||||
// It takes GenAIPromptDetails which has generate_content_config.
|
||||
// And toLogRecord puts prompt_id, request_text in attributes.
|
||||
}
|
||||
|
||||
// Since we can't easily see the tool definitions in ApiRequestEvent's request_text (which is just 'contents')
|
||||
// and prompt.generate_content_config is not directly in attributes (it is in StartSessionEvent though?),
|
||||
// wait, ApiRequestEvent.toLogRecord:
|
||||
/*
|
||||
const attributes: LogAttributes = {
|
||||
...getCommonAttributes(config),
|
||||
'event.name': EVENT_API_REQUEST,
|
||||
'event.timestamp': this['event.timestamp'],
|
||||
model: this.model,
|
||||
prompt_id: this.prompt.prompt_id,
|
||||
request_text: this.request_text,
|
||||
};
|
||||
*/
|
||||
// It doesn't seem to log the tools in the flat telemetry log.
|
||||
|
||||
// However, if ToolPreselectionService selected ONLY list_directory,
|
||||
// and the agent tried to call something else, it would fail or not have it.
|
||||
// Our mock responses are tailored:
|
||||
// 1. Classifier returns {relevant_tools: ["list_directory"]}
|
||||
// 2. Agent response calls list_directory.
|
||||
// This works. If tool preselection DIDN'T work, and our mock for turn 2 called say 'write_file',
|
||||
// it would still work because the mock doesn't care about what's in the prompt.
|
||||
|
||||
// To truly verify pre-selection in E2E, we'd need to see the tools in the request.
|
||||
// Given the current telemetry, maybe we can look for the 'tool-preselection' prompt itself.
|
||||
const preselectionRequest = apiRequests.find(
|
||||
(l) => l.attributes?.prompt_id === 'tool-preselection',
|
||||
);
|
||||
expect(preselectionRequest).toBeDefined();
|
||||
expect(preselectionRequest?.attributes?.request_text).toContain(
|
||||
'select only the tools that are strictly necessary',
|
||||
);
|
||||
});
|
||||
});
|
||||
@@ -859,6 +859,7 @@ export async function loadCliConfig(
|
||||
fakeResponses: argv.fakeResponses,
|
||||
recordResponses: argv.recordResponses,
|
||||
retryFetchErrors: settings.general?.retryFetchErrors,
|
||||
toolPreselection: settings.general?.toolPreselection,
|
||||
ptyInfo: ptyInfo?.name,
|
||||
disableLLMCorrection: settings.tools?.disableLLMCorrection,
|
||||
rawOutput: argv.rawOutput,
|
||||
|
||||
@@ -307,6 +307,16 @@ const SETTINGS_SCHEMA = {
|
||||
'Retry on "exception TypeError: fetch failed sending request" errors.',
|
||||
showInDialog: false,
|
||||
},
|
||||
toolPreselection: {
|
||||
type: 'boolean',
|
||||
label: 'Tool Preselection',
|
||||
category: 'General',
|
||||
requiresRestart: false,
|
||||
default: true,
|
||||
description:
|
||||
'Exclude unneeded tools from context to save tokens and improve performance.',
|
||||
showInDialog: true,
|
||||
},
|
||||
debugKeystrokeLogging: {
|
||||
type: 'boolean',
|
||||
label: 'Debug Keystroke Logging',
|
||||
|
||||
@@ -58,8 +58,8 @@ describe('AppRig', () => {
|
||||
|
||||
// Resolve and finish. Also removes read_file breakpoint.
|
||||
await rig.resolveTool('read_file');
|
||||
await rig.waitForOutput('Task complete.', 100000);
|
||||
});
|
||||
await rig.waitForOutput('Task complete.', 120000);
|
||||
}, 120000);
|
||||
|
||||
it('should render the app and handle a simple message', async () => {
|
||||
const fakeResponsesPath = path.join(
|
||||
@@ -83,6 +83,6 @@ describe('AppRig', () => {
|
||||
await rig.pressEnter();
|
||||
|
||||
// Wait for model response
|
||||
await rig.waitForOutput('Hello! How can I help you today?');
|
||||
});
|
||||
await rig.waitForOutput('Hello! How can I help you today?', 30000);
|
||||
}, 30000);
|
||||
});
|
||||
|
||||
@@ -139,6 +139,7 @@ export class AppRig {
|
||||
approvalMode,
|
||||
policyEngineConfig,
|
||||
enableEventDrivenScheduler: true,
|
||||
toolPreselection: false,
|
||||
extensionLoader: new MockExtensionManager(),
|
||||
excludeTools: this.options.configOverrides?.excludeTools,
|
||||
...this.options.configOverrides,
|
||||
@@ -178,41 +179,29 @@ export class AppRig {
|
||||
}
|
||||
|
||||
private createRigSettings(): LoadedSettings {
|
||||
const userSettings = {
|
||||
security: {
|
||||
auth: {
|
||||
selectedType: AuthType.USE_GEMINI,
|
||||
useExternal: true,
|
||||
},
|
||||
folderTrust: {
|
||||
enabled: true,
|
||||
},
|
||||
},
|
||||
ide: {
|
||||
enabled: false,
|
||||
hasSeenNudge: true,
|
||||
},
|
||||
};
|
||||
|
||||
return createMockSettings({
|
||||
user: {
|
||||
path: path.join(this.testDir, '.gemini', 'user_settings.json'),
|
||||
settings: {
|
||||
security: {
|
||||
auth: {
|
||||
selectedType: AuthType.USE_GEMINI,
|
||||
useExternal: true,
|
||||
},
|
||||
folderTrust: {
|
||||
enabled: true,
|
||||
},
|
||||
},
|
||||
ide: {
|
||||
enabled: false,
|
||||
hasSeenNudge: true,
|
||||
},
|
||||
},
|
||||
originalSettings: {},
|
||||
},
|
||||
merged: {
|
||||
security: {
|
||||
auth: {
|
||||
selectedType: AuthType.USE_GEMINI,
|
||||
useExternal: true,
|
||||
},
|
||||
folderTrust: {
|
||||
enabled: true,
|
||||
},
|
||||
},
|
||||
ide: {
|
||||
enabled: false,
|
||||
hasSeenNudge: true,
|
||||
},
|
||||
settings: userSettings,
|
||||
originalSettings: structuredClone(userSettings),
|
||||
},
|
||||
merged: userSettings,
|
||||
});
|
||||
}
|
||||
|
||||
|
||||
@@ -52,17 +52,29 @@ export const createMockSettings = (
|
||||
|
||||
const loaded = new LoadedSettings(
|
||||
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
|
||||
(system as any) || { path: '', settings: {}, originalSettings: {} },
|
||||
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
|
||||
(systemDefaults as any) || { path: '', settings: {}, originalSettings: {} },
|
||||
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
|
||||
(user as any) || {
|
||||
path: '',
|
||||
settings: settingsOverrides,
|
||||
originalSettings: settingsOverrides,
|
||||
(system as any) || {
|
||||
path: '/tmp/gemini_system_settings.json',
|
||||
settings: {},
|
||||
originalSettings: {},
|
||||
},
|
||||
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
|
||||
(workspace as any) || { path: '', settings: {}, originalSettings: {} },
|
||||
(systemDefaults as any) || {
|
||||
path: '/tmp/gemini_system_defaults.json',
|
||||
settings: {},
|
||||
originalSettings: {},
|
||||
},
|
||||
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
|
||||
(user as any) || {
|
||||
path: '/tmp/gemini_user_settings.json',
|
||||
settings: settingsOverrides,
|
||||
originalSettings: structuredClone(settingsOverrides),
|
||||
},
|
||||
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
|
||||
(workspace as any) || {
|
||||
path: '/tmp/gemini_workspace_settings.json',
|
||||
settings: {},
|
||||
originalSettings: {},
|
||||
},
|
||||
isTrusted ?? true,
|
||||
errors || [],
|
||||
);
|
||||
|
||||
@@ -78,6 +78,13 @@ exports[`InputPrompt > mouse interaction > should toggle paste expansion on doub
|
||||
"
|
||||
`;
|
||||
|
||||
exports[`InputPrompt > mouse interaction > should toggle paste expansion on double-click 4`] = `
|
||||
"▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀
|
||||
> [Pasted Text: 10 lines]
|
||||
▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄
|
||||
"
|
||||
`;
|
||||
|
||||
exports[`InputPrompt > snapshots > should not show inverted cursor when shell is focused 1`] = `
|
||||
"▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀
|
||||
> Type your message or @path/to/file
|
||||
|
||||
@@ -28,12 +28,12 @@ exports[`SettingsDialog > Initial Rendering > should render settings list with v
|
||||
│ Enable Prompt Completion false │
|
||||
│ Enable AI-powered prompt completion suggestions while typing. │
|
||||
│ │
|
||||
│ Tool Preselection true │
|
||||
│ Exclude unneeded tools from context to save tokens and improve performance. │
|
||||
│ │
|
||||
│ Debug Keystroke Logging false │
|
||||
│ Enable debug logging of keystrokes to the console. │
|
||||
│ │
|
||||
│ Enable Session Cleanup false │
|
||||
│ Enable automatic session cleanup │
|
||||
│ │
|
||||
│ ▼ │
|
||||
│ │
|
||||
│ Apply To │
|
||||
@@ -75,12 +75,12 @@ exports[`SettingsDialog > Snapshot Tests > should render 'accessibility settings
|
||||
│ Enable Prompt Completion false │
|
||||
│ Enable AI-powered prompt completion suggestions while typing. │
|
||||
│ │
|
||||
│ Tool Preselection true │
|
||||
│ Exclude unneeded tools from context to save tokens and improve performance. │
|
||||
│ │
|
||||
│ Debug Keystroke Logging false │
|
||||
│ Enable debug logging of keystrokes to the console. │
|
||||
│ │
|
||||
│ Enable Session Cleanup false │
|
||||
│ Enable automatic session cleanup │
|
||||
│ │
|
||||
│ ▼ │
|
||||
│ │
|
||||
│ Apply To │
|
||||
@@ -122,12 +122,12 @@ exports[`SettingsDialog > Snapshot Tests > should render 'all boolean settings d
|
||||
│ Enable Prompt Completion false* │
|
||||
│ Enable AI-powered prompt completion suggestions while typing. │
|
||||
│ │
|
||||
│ Tool Preselection true │
|
||||
│ Exclude unneeded tools from context to save tokens and improve performance. │
|
||||
│ │
|
||||
│ Debug Keystroke Logging false* │
|
||||
│ Enable debug logging of keystrokes to the console. │
|
||||
│ │
|
||||
│ Enable Session Cleanup false │
|
||||
│ Enable automatic session cleanup │
|
||||
│ │
|
||||
│ ▼ │
|
||||
│ │
|
||||
│ Apply To │
|
||||
@@ -169,12 +169,12 @@ exports[`SettingsDialog > Snapshot Tests > should render 'default state' correct
|
||||
│ Enable Prompt Completion false │
|
||||
│ Enable AI-powered prompt completion suggestions while typing. │
|
||||
│ │
|
||||
│ Tool Preselection true │
|
||||
│ Exclude unneeded tools from context to save tokens and improve performance. │
|
||||
│ │
|
||||
│ Debug Keystroke Logging false │
|
||||
│ Enable debug logging of keystrokes to the console. │
|
||||
│ │
|
||||
│ Enable Session Cleanup false │
|
||||
│ Enable automatic session cleanup │
|
||||
│ │
|
||||
│ ▼ │
|
||||
│ │
|
||||
│ Apply To │
|
||||
@@ -216,12 +216,12 @@ exports[`SettingsDialog > Snapshot Tests > should render 'file filtering setting
|
||||
│ Enable Prompt Completion false │
|
||||
│ Enable AI-powered prompt completion suggestions while typing. │
|
||||
│ │
|
||||
│ Tool Preselection true │
|
||||
│ Exclude unneeded tools from context to save tokens and improve performance. │
|
||||
│ │
|
||||
│ Debug Keystroke Logging false │
|
||||
│ Enable debug logging of keystrokes to the console. │
|
||||
│ │
|
||||
│ Enable Session Cleanup false │
|
||||
│ Enable automatic session cleanup │
|
||||
│ │
|
||||
│ ▼ │
|
||||
│ │
|
||||
│ Apply To │
|
||||
@@ -263,12 +263,12 @@ exports[`SettingsDialog > Snapshot Tests > should render 'focused on scope selec
|
||||
│ Enable Prompt Completion false │
|
||||
│ Enable AI-powered prompt completion suggestions while typing. │
|
||||
│ │
|
||||
│ Tool Preselection true │
|
||||
│ Exclude unneeded tools from context to save tokens and improve performance. │
|
||||
│ │
|
||||
│ Debug Keystroke Logging false │
|
||||
│ Enable debug logging of keystrokes to the console. │
|
||||
│ │
|
||||
│ Enable Session Cleanup false │
|
||||
│ Enable automatic session cleanup │
|
||||
│ │
|
||||
│ ▼ │
|
||||
│ │
|
||||
│ > Apply To │
|
||||
@@ -310,12 +310,12 @@ exports[`SettingsDialog > Snapshot Tests > should render 'mixed boolean and numb
|
||||
│ Enable Prompt Completion false │
|
||||
│ Enable AI-powered prompt completion suggestions while typing. │
|
||||
│ │
|
||||
│ Tool Preselection true │
|
||||
│ Exclude unneeded tools from context to save tokens and improve performance. │
|
||||
│ │
|
||||
│ Debug Keystroke Logging false │
|
||||
│ Enable debug logging of keystrokes to the console. │
|
||||
│ │
|
||||
│ Enable Session Cleanup false │
|
||||
│ Enable automatic session cleanup │
|
||||
│ │
|
||||
│ ▼ │
|
||||
│ │
|
||||
│ Apply To │
|
||||
@@ -357,12 +357,12 @@ exports[`SettingsDialog > Snapshot Tests > should render 'tools and security set
|
||||
│ Enable Prompt Completion false │
|
||||
│ Enable AI-powered prompt completion suggestions while typing. │
|
||||
│ │
|
||||
│ Tool Preselection true │
|
||||
│ Exclude unneeded tools from context to save tokens and improve performance. │
|
||||
│ │
|
||||
│ Debug Keystroke Logging false │
|
||||
│ Enable debug logging of keystrokes to the console. │
|
||||
│ │
|
||||
│ Enable Session Cleanup false │
|
||||
│ Enable automatic session cleanup │
|
||||
│ │
|
||||
│ ▼ │
|
||||
│ │
|
||||
│ Apply To │
|
||||
@@ -404,12 +404,12 @@ exports[`SettingsDialog > Snapshot Tests > should render 'various boolean settin
|
||||
│ Enable Prompt Completion true* │
|
||||
│ Enable AI-powered prompt completion suggestions while typing. │
|
||||
│ │
|
||||
│ Tool Preselection true │
|
||||
│ Exclude unneeded tools from context to save tokens and improve performance. │
|
||||
│ │
|
||||
│ Debug Keystroke Logging true* │
|
||||
│ Enable debug logging of keystrokes to the console. │
|
||||
│ │
|
||||
│ Enable Session Cleanup false │
|
||||
│ Enable automatic session cleanup │
|
||||
│ │
|
||||
│ ▼ │
|
||||
│ │
|
||||
│ Apply To │
|
||||
|
||||
@@ -50,6 +50,7 @@ export const GeneralistAgent = (
|
||||
const tools = config.getToolRegistry().getAllToolNames();
|
||||
return {
|
||||
tools,
|
||||
preselectTools: true,
|
||||
};
|
||||
},
|
||||
get promptConfig() {
|
||||
|
||||
@@ -30,6 +30,7 @@ import {
|
||||
} from '../core/geminiChat.js';
|
||||
import {
|
||||
type FunctionCall,
|
||||
type FunctionDeclaration,
|
||||
type Part,
|
||||
type GenerateContentResponse,
|
||||
type Content,
|
||||
@@ -71,6 +72,7 @@ import type {
|
||||
import type { AgentRegistry } from './registry.js';
|
||||
import { getModelConfigAlias } from './registry.js';
|
||||
import type { ModelRouterService } from '../routing/modelRouterService.js';
|
||||
import { ToolPreselectionService } from '../services/toolPreselectionService.js';
|
||||
|
||||
const {
|
||||
mockSendMessageStream,
|
||||
@@ -552,6 +554,70 @@ describe('LocalAgentExecutor', () => {
|
||||
});
|
||||
|
||||
describe('run (Execution Loop and Logic)', () => {
|
||||
it('should pre-select tools when preselectTools is enabled', async () => {
|
||||
const definition = createTestDefinition([
|
||||
LS_TOOL_NAME,
|
||||
READ_FILE_TOOL_NAME,
|
||||
]);
|
||||
definition.toolConfig!.preselectTools = true;
|
||||
definition.promptConfig.query = '${goal}';
|
||||
|
||||
// Mock tool pre-selection to only keep LS
|
||||
const selectToolsSpy = vi
|
||||
.spyOn(ToolPreselectionService.prototype, 'selectTools')
|
||||
.mockResolvedValue([LS_TOOL_NAME]);
|
||||
|
||||
// Mock a response to terminate the loop
|
||||
mockSendMessageStream.mockImplementation(async () =>
|
||||
(async function* () {
|
||||
yield {
|
||||
type: StreamEventType.CHUNK,
|
||||
value: createMockResponseChunk(
|
||||
[],
|
||||
[
|
||||
{
|
||||
name: TASK_COMPLETE_TOOL_NAME,
|
||||
args: { result: 'done' },
|
||||
id: 'call1',
|
||||
},
|
||||
],
|
||||
),
|
||||
} as StreamEvent;
|
||||
})(),
|
||||
);
|
||||
|
||||
const executor = await LocalAgentExecutor.create(
|
||||
definition,
|
||||
mockConfig,
|
||||
onActivity,
|
||||
);
|
||||
|
||||
const inputs = { goal: 'Test pre-selection' };
|
||||
await executor.run(inputs, signal);
|
||||
|
||||
// Verify ToolPreselectionService was called
|
||||
expect(selectToolsSpy).toHaveBeenCalledWith(
|
||||
expect.stringContaining('Test pre-selection'),
|
||||
expect.arrayContaining([
|
||||
expect.objectContaining({ name: LS_TOOL_NAME }),
|
||||
expect.objectContaining({ name: READ_FILE_TOOL_NAME }),
|
||||
]),
|
||||
expect.any(AbortSignal),
|
||||
);
|
||||
|
||||
// Verify GeminiChat was initialized with ONLY LS and complete_task
|
||||
const chatConstructorArgs = MockedGeminiChat.mock.calls[0];
|
||||
const tools = chatConstructorArgs[2]![0]
|
||||
.functionDeclarations as FunctionDeclaration[];
|
||||
const toolNames = tools.map((t) => t.name);
|
||||
|
||||
expect(toolNames).toContain(LS_TOOL_NAME);
|
||||
expect(toolNames).toContain(TASK_COMPLETE_TOOL_NAME);
|
||||
expect(toolNames).not.toContain(READ_FILE_TOOL_NAME);
|
||||
|
||||
selectToolsSpy.mockRestore();
|
||||
});
|
||||
|
||||
it('should log AgentFinish with error if run throws', async () => {
|
||||
const definition = createTestDefinition();
|
||||
// Make the definition invalid to cause an error during run
|
||||
|
||||
@@ -16,6 +16,7 @@ import type {
|
||||
Schema,
|
||||
} from '@google/genai';
|
||||
import { ToolRegistry } from '../tools/tool-registry.js';
|
||||
import { ToolPreselectionService } from '../services/toolPreselectionService.js';
|
||||
import {
|
||||
DiscoveredMCPTool,
|
||||
MCP_QUALIFIED_NAME_SEPARATOR,
|
||||
@@ -460,11 +461,41 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
|
||||
};
|
||||
|
||||
tools = this.prepareToolsList();
|
||||
chat = await this.createChatObject(augmentedInputs, tools);
|
||||
|
||||
const query = this.definition.promptConfig.query
|
||||
? templateString(this.definition.promptConfig.query, augmentedInputs)
|
||||
: DEFAULT_QUERY_STRING;
|
||||
|
||||
// Tool Pre-selection
|
||||
if (
|
||||
this.definition.toolConfig?.preselectTools &&
|
||||
this.runtimeContext.isToolPreselectionEnabled()
|
||||
) {
|
||||
const preselector = new ToolPreselectionService(this.runtimeContext);
|
||||
const selectedToolNames = await preselector.selectTools(
|
||||
query,
|
||||
tools,
|
||||
combinedSignal,
|
||||
);
|
||||
|
||||
// Filter the agent's tool registry to only include selected tools.
|
||||
// We MUST always include the task completion tool.
|
||||
const toolsToKeep = new Set(selectedToolNames);
|
||||
toolsToKeep.add(TASK_COMPLETE_TOOL_NAME);
|
||||
|
||||
const allAgentTools = this.toolRegistry.getAllToolNames();
|
||||
for (const toolName of allAgentTools) {
|
||||
if (!toolsToKeep.has(toolName)) {
|
||||
this.toolRegistry.unregisterTool(toolName);
|
||||
}
|
||||
}
|
||||
|
||||
// Re-prepare the tools list after filtering the registry.
|
||||
tools = this.prepareToolsList();
|
||||
}
|
||||
|
||||
chat = await this.createChatObject(augmentedInputs, tools);
|
||||
|
||||
const pendingHintsQueue: string[] = [];
|
||||
const hintListener = (hint: string) => {
|
||||
pendingHintsQueue.push(hint);
|
||||
|
||||
@@ -410,6 +410,13 @@ export class AgentRegistry {
|
||||
);
|
||||
}
|
||||
|
||||
if (overrides.preselectTools !== undefined) {
|
||||
merged.toolConfig = {
|
||||
...(definition.toolConfig ?? { tools: [] }),
|
||||
preselectTools: overrides.preselectTools,
|
||||
};
|
||||
}
|
||||
|
||||
return merged;
|
||||
}
|
||||
|
||||
|
||||
@@ -158,6 +158,11 @@ export interface PromptConfig {
|
||||
*/
|
||||
export interface ToolConfig {
|
||||
tools: Array<string | FunctionDeclaration | AnyDeclarativeTool>;
|
||||
/**
|
||||
* Whether to pre-select a subset of tools based on the user request.
|
||||
* This can help reduce context size and improve performance.
|
||||
*/
|
||||
preselectTools?: boolean;
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
@@ -193,6 +193,7 @@ export interface AgentRunConfig {
|
||||
export interface AgentOverride {
|
||||
modelConfig?: ModelConfig;
|
||||
runConfig?: AgentRunConfig;
|
||||
preselectTools?: boolean;
|
||||
enabled?: boolean;
|
||||
}
|
||||
|
||||
@@ -474,6 +475,7 @@ export interface ConfigParameters {
|
||||
continueOnFailedApiCall?: boolean;
|
||||
retryFetchErrors?: boolean;
|
||||
enableShellOutputEfficiency?: boolean;
|
||||
toolPreselection?: boolean;
|
||||
shellToolInactivityTimeout?: number;
|
||||
fakeResponses?: string;
|
||||
recordResponses?: string;
|
||||
@@ -653,6 +655,7 @@ export class Config {
|
||||
private readonly outputSettings: OutputSettings;
|
||||
private readonly continueOnFailedApiCall: boolean;
|
||||
private readonly retryFetchErrors: boolean;
|
||||
private readonly toolPreselection: boolean;
|
||||
private readonly enableShellOutputEfficiency: boolean;
|
||||
private readonly shellToolInactivityTimeout: number;
|
||||
readonly fakeResponses?: string;
|
||||
@@ -875,6 +878,7 @@ export class Config {
|
||||
format: params.output?.format ?? OutputFormat.TEXT,
|
||||
};
|
||||
this.retryFetchErrors = params.retryFetchErrors ?? false;
|
||||
this.toolPreselection = params.toolPreselection ?? true;
|
||||
this.disableYoloMode = params.disableYoloMode ?? false;
|
||||
this.rawOutput = params.rawOutput ?? false;
|
||||
this.acceptRawOutputRisk = params.acceptRawOutputRisk ?? false;
|
||||
@@ -2445,6 +2449,10 @@ export class Config {
|
||||
);
|
||||
}
|
||||
|
||||
isToolPreselectionEnabled(): boolean {
|
||||
return this.toolPreselection;
|
||||
}
|
||||
|
||||
getNextCompressionTruncationId(): number {
|
||||
return ++this.compressionTruncationCounter;
|
||||
}
|
||||
|
||||
@@ -216,6 +216,7 @@ describe('Gemini Client (client.ts)', () => {
|
||||
isJitContextEnabled: vi.fn().mockReturnValue(false),
|
||||
getToolOutputMaskingEnabled: vi.fn().mockReturnValue(false),
|
||||
getDisableLoopDetection: vi.fn().mockReturnValue(false),
|
||||
isToolPreselectionEnabled: vi.fn().mockReturnValue(false),
|
||||
|
||||
getSessionId: vi.fn().mockReturnValue('test-session-id'),
|
||||
getProxy: vi.fn().mockReturnValue(undefined),
|
||||
|
||||
@@ -35,6 +35,7 @@ import type {
|
||||
import type { ContentGenerator } from './contentGenerator.js';
|
||||
import { LoopDetectionService } from '../services/loopDetectionService.js';
|
||||
import { ChatCompressionService } from '../services/chatCompressionService.js';
|
||||
import { ToolPreselectionService } from '../services/toolPreselectionService.js';
|
||||
import { ideContextStore } from '../ide/ideContext.js';
|
||||
import {
|
||||
logContentRetryFailure,
|
||||
@@ -259,18 +260,36 @@ export class GeminiClient {
|
||||
|
||||
private lastUsedModelId?: string;
|
||||
|
||||
async setTools(modelId?: string): Promise<void> {
|
||||
async setTools(
|
||||
modelId?: string,
|
||||
query?: string,
|
||||
signal?: AbortSignal,
|
||||
): Promise<void> {
|
||||
if (!this.chat) {
|
||||
return;
|
||||
}
|
||||
|
||||
if (modelId && modelId === this.lastUsedModelId) {
|
||||
if (modelId && modelId === this.lastUsedModelId && !query) {
|
||||
return;
|
||||
}
|
||||
this.lastUsedModelId = modelId;
|
||||
|
||||
const toolRegistry = this.config.getToolRegistry();
|
||||
const toolDeclarations = toolRegistry.getFunctionDeclarations(modelId);
|
||||
let toolDeclarations = toolRegistry.getFunctionDeclarations(modelId);
|
||||
|
||||
if (query && signal && this.config.isToolPreselectionEnabled()) {
|
||||
const preselector = new ToolPreselectionService(this.config);
|
||||
const selectedNames = await preselector.selectTools(
|
||||
query,
|
||||
toolDeclarations,
|
||||
signal,
|
||||
);
|
||||
const selectedSet = new Set(selectedNames);
|
||||
toolDeclarations = toolDeclarations.filter((t) =>
|
||||
t.name ? selectedSet.has(t.name) : false,
|
||||
);
|
||||
}
|
||||
|
||||
const tools: Tool[] = [{ functionDeclarations: toolDeclarations }];
|
||||
this.getChat().setTools(tools);
|
||||
}
|
||||
@@ -677,7 +696,8 @@ export class GeminiClient {
|
||||
this.currentSequenceModel = modelToUse;
|
||||
|
||||
// Update tools with the final modelId to ensure model-dependent descriptions are used.
|
||||
await this.setTools(modelToUse);
|
||||
// Also perform tool pre-selection if enabled.
|
||||
await this.setTools(modelToUse, partToString(request), signal);
|
||||
|
||||
const resultStream = turn.run(
|
||||
modelConfigKey,
|
||||
|
||||
@@ -116,6 +116,7 @@ export * from './services/chatRecordingService.js';
|
||||
export * from './services/fileSystemService.js';
|
||||
export * from './services/sessionSummaryUtils.js';
|
||||
export * from './services/contextManager.js';
|
||||
export * from './services/toolPreselectionService.js';
|
||||
export * from './skills/skillManager.js';
|
||||
export * from './skills/skillLoader.js';
|
||||
|
||||
|
||||
@@ -0,0 +1,172 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect, vi, beforeEach, type Mock } from 'vitest';
|
||||
import { ToolPreselectionService } from './toolPreselectionService.js';
|
||||
import type { Config } from '../config/config.js';
|
||||
import type { FunctionDeclaration } from '@google/genai';
|
||||
|
||||
describe('ToolPreselectionService', () => {
|
||||
let mockConfig: Config;
|
||||
let mockLlmClient: Record<string, Mock>;
|
||||
let service: ToolPreselectionService;
|
||||
|
||||
beforeEach(() => {
|
||||
mockLlmClient = {
|
||||
generateJson: vi.fn(),
|
||||
};
|
||||
|
||||
mockConfig = {
|
||||
getBaseLlmClient: vi.fn().mockReturnValue(mockLlmClient),
|
||||
} as unknown as Config;
|
||||
|
||||
service = new ToolPreselectionService(mockConfig);
|
||||
});
|
||||
|
||||
it('returns all tools if count is below threshold', async () => {
|
||||
const tools: FunctionDeclaration[] = [
|
||||
{ name: 'tool1', description: 'desc1' },
|
||||
{ name: 'tool2', description: 'desc2' },
|
||||
];
|
||||
const result = await service.selectTools(
|
||||
'query',
|
||||
tools,
|
||||
new AbortController().signal,
|
||||
);
|
||||
expect(result).toEqual(['tool1', 'tool2']);
|
||||
expect(mockLlmClient['generateJson']).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('calls LLM for pre-selection if count is above threshold', async () => {
|
||||
const tools: FunctionDeclaration[] = [
|
||||
{ name: 'tool1', description: 'desc1' },
|
||||
{ name: 'tool2', description: 'desc2' },
|
||||
{ name: 'tool3', description: 'desc3' },
|
||||
{ name: 'tool4', description: 'desc4' },
|
||||
{ name: 'tool5', description: 'desc5' },
|
||||
{ name: 'tool6', description: 'desc6' },
|
||||
];
|
||||
|
||||
mockLlmClient['generateJson'].mockResolvedValue({
|
||||
relevant_tools: ['tool1', 'tool3'],
|
||||
});
|
||||
|
||||
const result = await service.selectTools(
|
||||
'my query',
|
||||
tools,
|
||||
new AbortController().signal,
|
||||
);
|
||||
|
||||
expect(result).toEqual(['tool1', 'tool3']);
|
||||
expect(mockLlmClient['generateJson']).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
contents: [
|
||||
{
|
||||
role: 'user',
|
||||
parts: [
|
||||
{
|
||||
text: expect.stringContaining('my query'),
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('respects maxTools option', async () => {
|
||||
const tools: FunctionDeclaration[] = [
|
||||
{ name: 'tool1', description: 'desc1' },
|
||||
{ name: 'tool2', description: 'desc2' },
|
||||
{ name: 'tool3', description: 'desc3' },
|
||||
];
|
||||
|
||||
mockLlmClient['generateJson'].mockResolvedValue({
|
||||
relevant_tools: ['tool1'],
|
||||
});
|
||||
|
||||
const result = await service.selectTools(
|
||||
'query',
|
||||
tools,
|
||||
new AbortController().signal,
|
||||
{ maxTools: 2 },
|
||||
);
|
||||
expect(result).toEqual(['tool1']);
|
||||
expect(mockLlmClient['generateJson']).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('falls back to all tools if LLM call fails', async () => {
|
||||
const tools: FunctionDeclaration[] = [
|
||||
{ name: 'tool1', description: 'desc1' },
|
||||
{ name: 'tool2', description: 'desc2' },
|
||||
{ name: 'tool3', description: 'desc3' },
|
||||
{ name: 'tool4', description: 'desc4' },
|
||||
{ name: 'tool5', description: 'desc5' },
|
||||
{ name: 'tool6', description: 'desc6' },
|
||||
];
|
||||
|
||||
mockLlmClient['generateJson'].mockRejectedValue(new Error('LLM error'));
|
||||
|
||||
const result = await service.selectTools(
|
||||
'query',
|
||||
tools,
|
||||
new AbortController().signal,
|
||||
);
|
||||
expect(result).toEqual([
|
||||
'tool1',
|
||||
'tool2',
|
||||
'tool3',
|
||||
'tool4',
|
||||
'tool5',
|
||||
'tool6',
|
||||
]);
|
||||
});
|
||||
|
||||
it('falls back to all tools if LLM returns invalid format', async () => {
|
||||
const tools: FunctionDeclaration[] = [
|
||||
{ name: 'tool1', description: 'desc1' },
|
||||
{ name: 'tool2', description: 'desc2' },
|
||||
{ name: 'tool3', description: 'desc3' },
|
||||
{ name: 'tool4', description: 'desc4' },
|
||||
{ name: 'tool5', description: 'desc5' },
|
||||
{ name: 'tool6', description: 'desc6' },
|
||||
];
|
||||
|
||||
// Missing 'relevant_tools' key
|
||||
mockLlmClient['generateJson'].mockResolvedValue({
|
||||
something_else: ['tool1'],
|
||||
});
|
||||
|
||||
const result = await service.selectTools(
|
||||
'query',
|
||||
tools,
|
||||
new AbortController().signal,
|
||||
);
|
||||
expect(result).toEqual([
|
||||
'tool1',
|
||||
'tool2',
|
||||
'tool3',
|
||||
'tool4',
|
||||
'tool5',
|
||||
'tool6',
|
||||
]);
|
||||
});
|
||||
|
||||
it('safely handles tools with undefined names', async () => {
|
||||
const tools: FunctionDeclaration[] = [
|
||||
{ name: 'tool1', description: 'desc1' },
|
||||
{ name: undefined, description: 'desc2' },
|
||||
];
|
||||
|
||||
const result = await service.selectTools(
|
||||
'query',
|
||||
tools,
|
||||
new AbortController().signal,
|
||||
{ maxTools: 10 }, // Force threshold bypass
|
||||
);
|
||||
expect(result).toEqual(['tool1']);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,108 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { FunctionDeclaration } from '@google/genai';
|
||||
import type { Config } from '../config/config.js';
|
||||
import { LlmRole } from '../telemetry/types.js';
|
||||
import { debugLogger } from '../utils/debugLogger.js';
|
||||
|
||||
export interface ToolPreselectionOptions {
|
||||
maxTools?: number;
|
||||
modelConfigKey?: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Service to pre-select a relevant subset of tools for a given query.
|
||||
* This helps reduce context size by excluding unneeded tool descriptions.
|
||||
*/
|
||||
export class ToolPreselectionService {
|
||||
constructor(private readonly config: Config) {}
|
||||
|
||||
/**
|
||||
* Selects relevant tools for a query.
|
||||
*
|
||||
* @param query The user's query or task description.
|
||||
* @param tools The full list of available function declarations.
|
||||
* @param signal AbortSignal for the request.
|
||||
* @param options Optional configuration for pre-selection.
|
||||
* @returns A list of tool names that are considered relevant.
|
||||
*/
|
||||
async selectTools(
|
||||
query: string,
|
||||
tools: FunctionDeclaration[],
|
||||
signal: AbortSignal,
|
||||
options: ToolPreselectionOptions = {},
|
||||
): Promise<string[]> {
|
||||
if (tools.length === 0) {
|
||||
return [];
|
||||
}
|
||||
|
||||
// Threshold below which we don't bother with pre-selection.
|
||||
const threshold = options.maxTools ?? 5;
|
||||
if (tools.length <= threshold) {
|
||||
return tools.filter((t) => !!t.name).map((t) => t.name!);
|
||||
}
|
||||
|
||||
const schema = {
|
||||
type: 'object',
|
||||
properties: {
|
||||
relevant_tools: {
|
||||
type: 'array',
|
||||
items: { type: 'string' },
|
||||
description:
|
||||
'The names of the tools that are relevant to the user request.',
|
||||
},
|
||||
},
|
||||
required: ['relevant_tools'],
|
||||
};
|
||||
|
||||
const toolsList = tools
|
||||
.map((_t) => `- ${_t.name}: ${_t.description}`)
|
||||
.join('\n');
|
||||
|
||||
const prompt = `Given the following user request and a list of available tools, select only the tools that are strictly necessary to solve the request.
|
||||
Return the result as a JSON array of tool names.
|
||||
|
||||
Request: ${query}
|
||||
|
||||
Available Tools:
|
||||
${toolsList}`;
|
||||
|
||||
try {
|
||||
const llmClient = this.config.getBaseLlmClient();
|
||||
const modelConfigKey = options.modelConfigKey || 'classifier';
|
||||
|
||||
const result = await llmClient.generateJson({
|
||||
modelConfigKey: { model: modelConfigKey },
|
||||
contents: [{ role: 'user', parts: [{ text: prompt }] }],
|
||||
schema: schema as Record<string, unknown>,
|
||||
abortSignal: signal,
|
||||
promptId: 'tool-preselection',
|
||||
role: LlmRole.UTILITY_TOOL,
|
||||
});
|
||||
|
||||
const selectedTools = result['relevant_tools'];
|
||||
|
||||
if (
|
||||
!Array.isArray(selectedTools) ||
|
||||
!selectedTools.every((item): item is string => typeof item === 'string')
|
||||
) {
|
||||
throw new Error(
|
||||
'Tool preselection returned invalid data format. Expected an array of strings.',
|
||||
);
|
||||
}
|
||||
|
||||
debugLogger.debug(
|
||||
`ToolPreselectionService: Selected ${selectedTools.length} tools out of ${tools.length} for query: "${query.substring(0, 50)}..."`,
|
||||
);
|
||||
return selectedTools;
|
||||
} catch (error) {
|
||||
debugLogger.error('ToolPreselectionService failed:', error);
|
||||
// Fallback: return all tools if pre-selection fails.
|
||||
return tools.filter((t) => !!t.name).map((t) => t.name!);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -76,6 +76,7 @@ export class GeminiCliSession {
|
||||
fakeResponses: options.fakeResponses,
|
||||
skillsSupport: true,
|
||||
adminSkillsEnabled: true,
|
||||
toolPreselection: false,
|
||||
policyEngineConfig: {
|
||||
// TODO: Revisit this default when we have a mechanism for wiring up approvals
|
||||
defaultDecision: PolicyDecision.ALLOW,
|
||||
|
||||
@@ -431,6 +431,7 @@ export class TestRig {
|
||||
// Nightly releases sometimes becomes out of sync with local code and
|
||||
// triggers auto-update, which causes tests to fail.
|
||||
disableAutoUpdate: true,
|
||||
toolPreselection: false,
|
||||
},
|
||||
telemetry: {
|
||||
enabled: true,
|
||||
@@ -963,6 +964,10 @@ export class TestRig {
|
||||
);
|
||||
}
|
||||
|
||||
readTelemetryLogs(): any[] {
|
||||
return this._readAndParseTelemetryLog();
|
||||
}
|
||||
|
||||
async waitForToolCall(
|
||||
toolName: string,
|
||||
timeout?: number,
|
||||
|
||||
@@ -135,6 +135,13 @@
|
||||
"default": false,
|
||||
"type": "boolean"
|
||||
},
|
||||
"toolPreselection": {
|
||||
"title": "Tool Preselection",
|
||||
"description": "Exclude unneeded tools from context to save tokens and improve performance.",
|
||||
"markdownDescription": "Exclude unneeded tools from context to save tokens and improve performance.\n\n- Category: `General`\n- Requires restart: `no`\n- Default: `true`",
|
||||
"default": true,
|
||||
"type": "boolean"
|
||||
},
|
||||
"debugKeystrokeLogging": {
|
||||
"title": "Debug Keystroke Logging",
|
||||
"description": "Enable debug logging of keystrokes to the console.",
|
||||
|
||||
Reference in New Issue
Block a user