MediaMetz/gemini-cli
1
0
Fork 0
mirror of https://github.com/google-gemini/gemini-cli.git synced 2026-05-14 13:53:02 -07:00
Code Issues Packages Projects Releases Wiki Activity

refactor(core): replace manual syncPlanModeTools with declarative policy rules (#20596)

Browse Source
This commit is contained in:
Jerop Kipruto
2026-03-02 17:30:50 -05:00
committed by GitHub
parent e43b1cff58
commit d05ba11a31
11 changed files with 198 additions and 226 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/core/src/config/config.test.ts
+22 -127
View File
@@ -223,8 +223,6 @@ import type {
ModelConfigService,
ModelConfigServiceConfig,
} from '../services/modelConfigService.js';
import { ExitPlanModeTool } from '../tools/exit-plan-mode.js';
import { EnterPlanModeTool } from '../tools/enter-plan-mode.js';
import { LocalLiteRtLmClient } from '../core/localLiteRtLmClient.js';
vi.mock('../core/baseLlmClient.js');
@@ -1204,6 +1202,28 @@ describe('Server Config (config.ts)', () => {
expect(SubAgentToolMock).not.toHaveBeenCalled();
});
it('should register EnterPlanModeTool and ExitPlanModeTool when plan is enabled', async () => {
const params: ConfigParameters = {
...baseParams,
plan: true,
};
const config = new Config(params);
await config.initialize();
const registerToolMock = (
(await vi.importMock('../tools/tool-registry')) as {
ToolRegistry: { prototype: { registerTool: Mock } };
}
).ToolRegistry.prototype.registerTool;
const registeredTools = registerToolMock.mock.calls.map(
(call) => call[0].constructor.name,
);
expect(registeredTools).toContain('EnterPlanModeTool');
expect(registeredTools).toContain('ExitPlanModeTool');
});
describe('with minified tool class names', () => {
beforeEach(() => {
Object.defineProperty(
@@ -2961,131 +2981,6 @@ describe('Plans Directory Initialization', () => {
expect(fs.promises.mkdir).not.toHaveBeenCalledWith(plansDir, {
recursive: true,
});
const context = config.getWorkspaceContext();
expect(context.getDirectories()).not.toContain(plansDir);
});
});
describe('syncPlanModeTools', () => {
const baseParams: ConfigParameters = {
sessionId: 'test-session',
targetDir: '.',
debugMode: false,
model: 'test-model',
cwd: '.',
};
it('should register ExitPlanModeTool and unregister EnterPlanModeTool when in PLAN mode', async () => {
const config = new Config({
...baseParams,
approvalMode: ApprovalMode.PLAN,
});
const registry = new ToolRegistry(config, config.getMessageBus());
vi.spyOn(config, 'getToolRegistry').mockReturnValue(registry);
const registerSpy = vi.spyOn(registry, 'registerTool');
const unregisterSpy = vi.spyOn(registry, 'unregisterTool');
const getToolSpy = vi.spyOn(registry, 'getTool');
getToolSpy.mockImplementation((name) => {
if (name === 'enter_plan_mode')
return new EnterPlanModeTool(config, config.getMessageBus());
return undefined;
});
config.syncPlanModeTools();
expect(unregisterSpy).toHaveBeenCalledWith('enter_plan_mode');
expect(registerSpy).toHaveBeenCalledWith(expect.anything());
const registeredTool = registerSpy.mock.calls[0][0];
const { ExitPlanModeTool } = await import('../tools/exit-plan-mode.js');
expect(registeredTool).toBeInstanceOf(ExitPlanModeTool);
});
it('should register EnterPlanModeTool and unregister ExitPlanModeTool when NOT in PLAN mode and experimental.plan is enabled', async () => {
const config = new Config({
...baseParams,
approvalMode: ApprovalMode.DEFAULT,
plan: true,
});
const registry = new ToolRegistry(config, config.getMessageBus());
vi.spyOn(config, 'getToolRegistry').mockReturnValue(registry);
const registerSpy = vi.spyOn(registry, 'registerTool');
const unregisterSpy = vi.spyOn(registry, 'unregisterTool');
const getToolSpy = vi.spyOn(registry, 'getTool');
getToolSpy.mockImplementation((name) => {
if (name === 'exit_plan_mode')
return new ExitPlanModeTool(config, config.getMessageBus());
return undefined;
});
config.syncPlanModeTools();
expect(unregisterSpy).toHaveBeenCalledWith('exit_plan_mode');
expect(registerSpy).toHaveBeenCalledWith(expect.anything());
const registeredTool = registerSpy.mock.calls[0][0];
const { EnterPlanModeTool } = await import('../tools/enter-plan-mode.js');
expect(registeredTool).toBeInstanceOf(EnterPlanModeTool);
});
it('should NOT register EnterPlanModeTool when experimental.plan is disabled', async () => {
const config = new Config({
...baseParams,
approvalMode: ApprovalMode.DEFAULT,
plan: false,
});
const registry = new ToolRegistry(config, config.getMessageBus());
vi.spyOn(config, 'getToolRegistry').mockReturnValue(registry);
const registerSpy = vi.spyOn(registry, 'registerTool');
vi.spyOn(registry, 'getTool').mockReturnValue(undefined);
config.syncPlanModeTools();
const { EnterPlanModeTool } = await import('../tools/enter-plan-mode.js');
const registeredTool = registerSpy.mock.calls.find(
(call) => call[0] instanceof EnterPlanModeTool,
);
expect(registeredTool).toBeUndefined();
});
it('should NOT register EnterPlanModeTool when in YOLO mode, even if plan is enabled', async () => {
const config = new Config({
...baseParams,
approvalMode: ApprovalMode.YOLO,
plan: true,
});
const registry = new ToolRegistry(config, config.getMessageBus());
vi.spyOn(config, 'getToolRegistry').mockReturnValue(registry);
const registerSpy = vi.spyOn(registry, 'registerTool');
vi.spyOn(registry, 'getTool').mockReturnValue(undefined);
config.syncPlanModeTools();
const { EnterPlanModeTool } = await import('../tools/enter-plan-mode.js');
const registeredTool = registerSpy.mock.calls.find(
(call) => call[0] instanceof EnterPlanModeTool,
);
expect(registeredTool).toBeUndefined();
});
it('should call geminiClient.setTools if initialized', async () => {
const config = new Config(baseParams);
const registry = new ToolRegistry(config, config.getMessageBus());
vi.spyOn(config, 'getToolRegistry').mockReturnValue(registry);
const client = config.getGeminiClient();
vi.spyOn(client, 'isInitialized').mockReturnValue(true);
const setToolsSpy = vi
.spyOn(client, 'setTools')
.mockResolvedValue(undefined);
config.syncPlanModeTools();
expect(setToolsSpy).toHaveBeenCalled();
});
});
packages/core/src/config/config.ts
+5 -47
View File
@@ -370,10 +370,6 @@ import { McpClientManager } from '../tools/mcp-client-manager.js';
import { type McpContext } from '../tools/mcp-client.js';
import type { EnvironmentSanitizationConfig } from '../services/environmentSanitization.js';
import { getErrorMessage } from '../utils/errors.js';
import {
ENTER_PLAN_MODE_TOOL_NAME,
EXIT_PLAN_MODE_TOOL_NAME,
} from '../tools/tool-names.js';
export type { FileFilteringOptions };
export {
@@ -1172,7 +1168,6 @@ export class Config implements McpContext {
}
await this.geminiClient.initialize();
this.syncPlanModeTools();
this.initialized = true;
}
@@ -1998,52 +1993,15 @@ export class Config implements McpContext {
(currentMode === ApprovalMode.YOLO || mode === ApprovalMode.YOLO);
if (isPlanModeTransition || isYoloModeTransition) {
this.syncPlanModeTools();
if (this.geminiClient?.isInitialized()) {
this.geminiClient.setTools().catch((err) => {
debugLogger.error('Failed to update tools', err);
});
}
this.updateSystemInstructionIfInitialized();
}
}
/**
* 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.planEnabled && !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.
*/