diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index c6458dcc1f..4bb61e17be 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -35,6 +35,7 @@ import { MemoryTool, setGeminiMdFilename } from '../tools/memoryTool.js'; import { WebSearchTool } from '../tools/web-search.js'; import { AskUserTool } from '../tools/ask-user.js'; import { ExitPlanModeTool } from '../tools/exit-plan-mode.js'; +import { EnterPlanModeTool } from '../tools/enter-plan-mode.js'; import { GeminiClient } from '../core/client.js'; import { BaseLlmClient } from '../core/baseLlmClient.js'; import type { HookDefinition, HookEventName } from '../hooks/types.js'; @@ -2155,6 +2156,7 @@ export class Config { } if (this.isPlanEnabled()) { registerCoreTool(ExitPlanModeTool, this); + registerCoreTool(EnterPlanModeTool, this); } // Register Subagents as Tools diff --git a/packages/core/src/tools/enter-plan-mode.test.ts b/packages/core/src/tools/enter-plan-mode.test.ts new file mode 100644 index 0000000000..0b1d0a37f0 --- /dev/null +++ b/packages/core/src/tools/enter-plan-mode.test.ts @@ -0,0 +1,170 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; +import { EnterPlanModeTool } from './enter-plan-mode.js'; +import { createMockMessageBus } from '../test-utils/mock-message-bus.js'; +import type { Config } from '../config/config.js'; +import type { MessageBus } from '../confirmation-bus/message-bus.js'; +import { ToolConfirmationOutcome } from './tools.js'; +import { ApprovalMode } from '../policy/types.js'; + +describe('EnterPlanModeTool', () => { + let tool: EnterPlanModeTool; + let mockMessageBus: ReturnType; + let mockConfig: Partial; + + beforeEach(() => { + mockMessageBus = createMockMessageBus(); + vi.mocked(mockMessageBus.publish).mockResolvedValue(undefined); + + mockConfig = { + setApprovalMode: vi.fn(), + storage: { + getProjectTempPlansDir: vi.fn().mockReturnValue('/mock/plans/dir'), + } as unknown as Config['storage'], + }; + tool = new EnterPlanModeTool( + mockConfig as Config, + mockMessageBus as unknown as MessageBus, + ); + }); + + afterEach(() => { + vi.restoreAllMocks(); + }); + + describe('shouldConfirmExecute', () => { + it('should return info confirmation details when policy says ASK_USER', async () => { + const invocation = tool.build({}); + + // Mock getMessageBusDecision to return ASK_USER + vi.spyOn( + invocation as unknown as { + getMessageBusDecision: () => Promise; + }, + 'getMessageBusDecision', + ).mockResolvedValue('ASK_USER'); + + const result = await invocation.shouldConfirmExecute( + new AbortController().signal, + ); + + expect(result).not.toBe(false); + if (result === false) return; + + expect(result.type).toBe('info'); + expect(result.title).toBe('Enter Plan Mode'); + if (result.type === 'info') { + expect(result.prompt).toBe( + 'This will restrict the agent to read-only tools to allow for safe planning.', + ); + } + }); + + it('should return false when policy decision is ALLOW', async () => { + const invocation = tool.build({}); + + // Mock getMessageBusDecision to return ALLOW + vi.spyOn( + invocation as unknown as { + getMessageBusDecision: () => Promise; + }, + 'getMessageBusDecision', + ).mockResolvedValue('ALLOW'); + + const result = await invocation.shouldConfirmExecute( + new AbortController().signal, + ); + + expect(result).toBe(false); + }); + + it('should throw error when policy decision is DENY', async () => { + const invocation = tool.build({}); + + // Mock getMessageBusDecision to return DENY + vi.spyOn( + invocation as unknown as { + getMessageBusDecision: () => Promise; + }, + 'getMessageBusDecision', + ).mockResolvedValue('DENY'); + + await expect( + invocation.shouldConfirmExecute(new AbortController().signal), + ).rejects.toThrow(/denied by policy/); + }); + }); + + describe('execute', () => { + it('should set approval mode to PLAN and return message', async () => { + const invocation = tool.build({}); + + const result = await invocation.execute(new AbortController().signal); + + expect(mockConfig.setApprovalMode).toHaveBeenCalledWith( + ApprovalMode.PLAN, + ); + expect(result.llmContent).toContain('Switching to Plan mode'); + expect(result.returnDisplay).toBe('Switching to Plan mode'); + }); + + it('should include optional reason in output display but not in llmContent', async () => { + const reason = 'Design new database schema'; + const invocation = tool.build({ reason }); + + const result = await invocation.execute(new AbortController().signal); + + expect(mockConfig.setApprovalMode).toHaveBeenCalledWith( + ApprovalMode.PLAN, + ); + expect(result.llmContent).toBe('Switching to Plan mode.'); + expect(result.llmContent).not.toContain(reason); + expect(result.returnDisplay).toContain(reason); + }); + + it('should not enter plan mode if cancelled', async () => { + const invocation = tool.build({}); + + // Simulate getting confirmation details + vi.spyOn( + invocation as unknown as { + getMessageBusDecision: () => Promise; + }, + 'getMessageBusDecision', + ).mockResolvedValue('ASK_USER'); + + const details = await invocation.shouldConfirmExecute( + new AbortController().signal, + ); + expect(details).not.toBe(false); + + if (details) { + // Simulate user cancelling + await details.onConfirm(ToolConfirmationOutcome.Cancel); + } + + const result = await invocation.execute(new AbortController().signal); + + expect(mockConfig.setApprovalMode).not.toHaveBeenCalled(); + expect(result.returnDisplay).toBe('Cancelled'); + expect(result.llmContent).toContain('User cancelled'); + }); + }); + + describe('validateToolParams', () => { + it('should allow empty params', () => { + const result = tool.validateToolParams({}); + expect(result).toBeNull(); + }); + + it('should allow reason param', () => { + const result = tool.validateToolParams({ reason: 'test' }); + expect(result).toBeNull(); + }); + }); +}); diff --git a/packages/core/src/tools/enter-plan-mode.ts b/packages/core/src/tools/enter-plan-mode.ts new file mode 100644 index 0000000000..89fe0cbf2f --- /dev/null +++ b/packages/core/src/tools/enter-plan-mode.ts @@ -0,0 +1,133 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { + BaseDeclarativeTool, + BaseToolInvocation, + type ToolResult, + Kind, + type ToolInfoConfirmationDetails, + ToolConfirmationOutcome, +} from './tools.js'; +import type { MessageBus } from '../confirmation-bus/message-bus.js'; +import type { Config } from '../config/config.js'; +import { ENTER_PLAN_MODE_TOOL_NAME } from './tool-names.js'; +import { ApprovalMode } from '../policy/types.js'; + +export interface EnterPlanModeParams { + reason?: string; +} + +export class EnterPlanModeTool extends BaseDeclarativeTool< + EnterPlanModeParams, + ToolResult +> { + constructor( + private config: Config, + messageBus: MessageBus, + ) { + super( + ENTER_PLAN_MODE_TOOL_NAME, + 'Enter Plan Mode', + 'Switch to Plan Mode to safely research, design, and plan complex changes using read-only tools.', + Kind.Plan, + { + type: 'object', + properties: { + reason: { + type: 'string', + description: + 'Short reason explaining why you are entering plan mode.', + }, + }, + }, + messageBus, + ); + } + + protected createInvocation( + params: EnterPlanModeParams, + messageBus: MessageBus, + toolName: string, + toolDisplayName: string, + ): EnterPlanModeInvocation { + return new EnterPlanModeInvocation( + params, + messageBus, + toolName, + toolDisplayName, + this.config, + ); + } +} + +export class EnterPlanModeInvocation extends BaseToolInvocation< + EnterPlanModeParams, + ToolResult +> { + private confirmationOutcome: ToolConfirmationOutcome | null = null; + + constructor( + params: EnterPlanModeParams, + messageBus: MessageBus, + toolName: string, + toolDisplayName: string, + private config: Config, + ) { + super(params, messageBus, toolName, toolDisplayName); + } + + getDescription(): string { + return this.params.reason || 'Initiating Plan Mode'; + } + + override async shouldConfirmExecute( + abortSignal: AbortSignal, + ): Promise { + const decision = await this.getMessageBusDecision(abortSignal); + if (decision === 'ALLOW') { + return false; + } + + if (decision === 'DENY') { + throw new Error( + `Tool execution for "${ + this._toolDisplayName || this._toolName + }" denied by policy.`, + ); + } + + // ASK_USER + return { + type: 'info', + title: 'Enter Plan Mode', + prompt: + 'This will restrict the agent to read-only tools to allow for safe planning.', + onConfirm: async (outcome: ToolConfirmationOutcome) => { + this.confirmationOutcome = outcome; + await this.publishPolicyUpdate(outcome); + }, + }; + } + + async execute(_signal: AbortSignal): Promise { + if (this.confirmationOutcome === ToolConfirmationOutcome.Cancel) { + return { + llmContent: 'User cancelled entering Plan Mode.', + returnDisplay: 'Cancelled', + }; + } + + this.config.setApprovalMode(ApprovalMode.PLAN); + + return { + llmContent: 'Switching to Plan mode.', + returnDisplay: this.params.reason + ? `Switching to Plan mode: ${this.params.reason}` + : 'Switching to Plan mode', + }; + } +} diff --git a/packages/core/src/tools/tool-names.ts b/packages/core/src/tools/tool-names.ts index 96f708fc71..5b8f89d4f5 100644 --- a/packages/core/src/tools/tool-names.ts +++ b/packages/core/src/tools/tool-names.ts @@ -26,6 +26,7 @@ export const EDIT_TOOL_NAMES = new Set([EDIT_TOOL_NAME, WRITE_FILE_TOOL_NAME]); export const ASK_USER_TOOL_NAME = 'ask_user'; export const ASK_USER_DISPLAY_NAME = 'Ask User'; export const EXIT_PLAN_MODE_TOOL_NAME = 'exit_plan_mode'; +export const ENTER_PLAN_MODE_TOOL_NAME = 'enter_plan_mode'; /** * Mapping of legacy tool names to their current names.