-
Notifications
You must be signed in to change notification settings - Fork 13k
feat(core): add enter_plan_mode tool
#18324
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
Merged
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit
Hold shift + click to select a range
File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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<typeof createMockMessageBus>; | ||
| let mockConfig: Partial<Config>; | ||
|
|
||
| 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<string>; | ||
| }, | ||
| '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<string>; | ||
| }, | ||
| '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<string>; | ||
| }, | ||
| '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<string>; | ||
| }, | ||
| '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(); | ||
| }); | ||
| }); | ||
| }); |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,132 @@ | ||
| /** | ||
| * @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 before starting implementation.', | ||
| Kind.Plan, | ||
| { | ||
| type: 'object', | ||
| properties: { | ||
| reason: { | ||
| type: 'string', | ||
| description: 'Optional reason for entering plan mode.', | ||
jerop marked this conversation as resolved.
Outdated
Show resolved
Hide resolved
|
||
| }, | ||
| }, | ||
| }, | ||
| 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<ToolInfoConfirmationDetails | false> { | ||
| 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<ToolResult> { | ||
| 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', | ||
| }; | ||
| } | ||
| } | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.