Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
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();
});
});
});
133 changes: 133 additions & 0 deletions packages/core/src/tools/enter-plan-mode.ts
Original file line number Diff line number Diff line change
@@ -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<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