Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
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
2 changes: 2 additions & 0 deletions packages/core/src/config/config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -2144,6 +2145,7 @@ export class Config {
}
if (this.isPlanEnabled()) {
registerCoreTool(ExitPlanModeTool, this);
registerCoreTool(EnterPlanModeTool, this);
}

// Register Subagents as Tools
Expand Down
170 changes: 170 additions & 0 deletions packages/core/src/tools/enter-plan-mode.test.ts
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();
});
});
});
132 changes: 132 additions & 0 deletions packages/core/src/tools/enter-plan-mode.ts
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.',
},
},
},
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',
};
}
}
1 change: 1 addition & 0 deletions packages/core/src/tools/tool-names.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
Loading