Merge branch 'main' into adibakm/disable-esc-esc-rewind

This commit is contained in:
Adib234
2026-01-26 17:16:34 -05:00
committed by GitHub
58 changed files with 3142 additions and 1748 deletions
+3 -1
View File
@@ -57,7 +57,9 @@ vi.mock('fs', async (importOriginal) => {
return {
...actualFs,
mkdirSync: vi.fn(),
mkdirSync: vi.fn((p) => {
mockPaths.add(p.toString());
}),
writeFileSync: vi.fn(),
existsSync: vi.fn((p) => mockPaths.has(p.toString())),
statSync: vi.fn((p) => {
@@ -324,6 +324,117 @@ describe('Policy Engine Integration Tests', () => {
).toBe(PolicyDecision.DENY);
});
it('should allow write_file to plans directory in Plan mode', async () => {
const settings: Settings = {};
const config = await createPolicyEngineConfig(
settings,
ApprovalMode.PLAN,
);
const engine = new PolicyEngine(config);
// Valid plan file path (64-char hex hash, .md extension, safe filename)
const validPlanPath =
'/home/user/.gemini/tmp/a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2/plans/my-plan.md';
expect(
(
await engine.check(
{ name: 'write_file', args: { file_path: validPlanPath } },
undefined,
)
).decision,
).toBe(PolicyDecision.ALLOW);
// Valid plan with underscore in filename
const validPlanPath2 =
'/home/user/.gemini/tmp/a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2/plans/feature_auth.md';
expect(
(
await engine.check(
{ name: 'write_file', args: { file_path: validPlanPath2 } },
undefined,
)
).decision,
).toBe(PolicyDecision.ALLOW);
});
it('should deny write_file outside plans directory in Plan mode', async () => {
const settings: Settings = {};
const config = await createPolicyEngineConfig(
settings,
ApprovalMode.PLAN,
);
const engine = new PolicyEngine(config);
// Write to workspace (not plans dir) should be denied
expect(
(
await engine.check(
{ name: 'write_file', args: { file_path: '/project/src/file.ts' } },
undefined,
)
).decision,
).toBe(PolicyDecision.DENY);
// Write to plans dir but wrong extension should be denied
const wrongExtPath =
'/home/user/.gemini/tmp/a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2/plans/script.js';
expect(
(
await engine.check(
{ name: 'write_file', args: { file_path: wrongExtPath } },
undefined,
)
).decision,
).toBe(PolicyDecision.DENY);
// Path traversal attempt should be denied (filename contains /)
const traversalPath =
'/home/user/.gemini/tmp/a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2/plans/../../../etc/passwd.md';
expect(
(
await engine.check(
{ name: 'write_file', args: { file_path: traversalPath } },
undefined,
)
).decision,
).toBe(PolicyDecision.DENY);
// Invalid hash length should be denied
const shortHashPath = '/home/user/.gemini/tmp/abc123/plans/plan.md';
expect(
(
await engine.check(
{ name: 'write_file', args: { file_path: shortHashPath } },
undefined,
)
).decision,
).toBe(PolicyDecision.DENY);
});
it('should deny write_file to subdirectories in Plan mode', async () => {
const settings: Settings = {};
const config = await createPolicyEngineConfig(
settings,
ApprovalMode.PLAN,
);
const engine = new PolicyEngine(config);
// Write to subdirectory should be denied
const subdirPath =
'/home/user/.gemini/tmp/a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2/plans/subdir/plan.md';
expect(
(
await engine.check(
{ name: 'write_file', args: { file_path: subdirPath } },
undefined,
)
).decision,
).toBe(PolicyDecision.DENY);
});
it('should verify priority ordering works correctly in practice', async () => {
const settings: Settings = {
tools: {
+1
View File
@@ -198,6 +198,7 @@ const mockUIActions: UIActions = {
setEmbeddedShellFocused: vi.fn(),
setAuthContext: vi.fn(),
handleRestart: vi.fn(),
handleNewAgentsSelect: vi.fn(),
};
export const renderWithProviders = (
+37 -2
View File
@@ -63,6 +63,7 @@ import {
SessionStartSource,
SessionEndReason,
generateSummary,
type AgentsDiscoveredPayload,
ChangeAuthRequestedError,
} from '@google/gemini-cli-core';
import { validateAuthMethod } from '../config/auth.js';
@@ -133,6 +134,7 @@ import {
QUEUE_ERROR_DISPLAY_DURATION_MS,
} from './constants.js';
import { LoginWithGoogleRestartDialog } from './auth/LoginWithGoogleRestartDialog.js';
import { NewAgentsChoice } from './components/NewAgentsNotification.js';
import { isSlashCommand } from './utils/commandUtils.js';
function isToolExecuting(pendingHistoryItems: HistoryItemWithoutId[]) {
@@ -218,6 +220,8 @@ export const AppContainer = (props: AppContainerProps) => {
null,
);
const [newAgents, setNewAgents] = useState<AgentDefinition[] | null>(null);
const [defaultBannerText, setDefaultBannerText] = useState('');
const [warningBannerText, setWarningBannerText] = useState('');
const [bannerVisible, setBannerVisible] = useState(true);
@@ -414,14 +418,20 @@ export const AppContainer = (props: AppContainerProps) => {
setAdminSettingsChanged(true);
};
const handleAgentsDiscovered = (payload: AgentsDiscoveredPayload) => {
setNewAgents(payload.agents);
};
coreEvents.on(CoreEvent.SettingsChanged, handleSettingsChanged);
coreEvents.on(CoreEvent.AdminSettingsChanged, handleAdminSettingsChanged);
coreEvents.on(CoreEvent.AgentsDiscovered, handleAgentsDiscovered);
return () => {
coreEvents.off(CoreEvent.SettingsChanged, handleSettingsChanged);
coreEvents.off(
CoreEvent.AdminSettingsChanged,
handleAdminSettingsChanged,
);
coreEvents.off(CoreEvent.AgentsDiscovered, handleAgentsDiscovered);
};
}, []);
@@ -1564,8 +1574,8 @@ Logging in with Google... Restarting Gemini CLI to continue.
!!proQuotaRequest ||
!!validationRequest ||
isSessionBrowserOpen ||
isAuthDialogOpen ||
authState === AuthState.AwaitingApiKeyInput;
authState === AuthState.AwaitingApiKeyInput ||
!!newAgents;
const pendingHistoryItems = useMemo(
() => [...pendingSlashCommandHistoryItems, ...pendingGeminiHistoryItems],
@@ -1728,6 +1738,7 @@ Logging in with Google... Restarting Gemini CLI to continue.
terminalBackgroundColor: config.getTerminalBackground(),
settingsNonce,
adminSettingsChanged,
newAgents,
}),
[
isThemeDialogOpen,
@@ -1828,6 +1839,7 @@ Logging in with Google... Restarting Gemini CLI to continue.
config,
settingsNonce,
adminSettingsChanged,
newAgents,
],
);
@@ -1879,6 +1891,26 @@ Logging in with Google... Restarting Gemini CLI to continue.
await runExitCleanup();
process.exit(RELAUNCH_EXIT_CODE);
},
handleNewAgentsSelect: async (choice: NewAgentsChoice) => {
if (newAgents && choice === NewAgentsChoice.ACKNOWLEDGE) {
const registry = config.getAgentRegistry();
try {
await Promise.all(
newAgents.map((agent) => registry.acknowledgeAgent(agent)),
);
} catch (error) {
debugLogger.error('Failed to acknowledge agents:', error);
historyManager.addItem(
{
type: MessageType.ERROR,
text: `Failed to acknowledge agents: ${getErrorMessage(error)}`,
},
Date.now(),
);
}
}
setNewAgents(null);
},
}),
[
handleThemeSelect,
@@ -1918,6 +1950,9 @@ Logging in with Google... Restarting Gemini CLI to continue.
setBannerVisible,
setEmbeddedShellFocused,
setAuthContext,
newAgents,
config,
historyManager,
],
);
@@ -17,9 +17,18 @@ import type { CommandContext, OpenCustomDialogActionReturn } from './types.js';
import { MessageType } from '../types.js';
import * as os from 'node:os';
import * as path from 'node:path';
import * as fs from 'node:fs';
import * as trustedFolders from '../../config/trustedFolders.js';
import type { LoadedTrustedFolders } from '../../config/trustedFolders.js';
vi.mock('node:fs', async (importOriginal) => {
const actual = await importOriginal<typeof import('node:fs')>();
return {
...actual,
realpathSync: vi.fn((p) => p),
};
});
vi.mock('../utils/directoryUtils.js', async (importOriginal) => {
const actual =
await importOriginal<typeof import('../utils/directoryUtils.js')>();
@@ -42,13 +51,14 @@ describe('directoryCommand', () => {
beforeEach(() => {
mockWorkspaceContext = {
targetDir: path.resolve('/test/dir'),
addDirectory: vi.fn(),
addDirectories: vi.fn().mockReturnValue({ added: [], failed: [] }),
getDirectories: vi
.fn()
.mockReturnValue([
path.normalize('/home/user/project1'),
path.normalize('/home/user/project2'),
path.resolve('/home/user/project1'),
path.resolve('/home/user/project2'),
]),
} as unknown as WorkspaceContext;
@@ -58,7 +68,7 @@ describe('directoryCommand', () => {
getGeminiClient: vi.fn().mockReturnValue({
addDirectoryContext: vi.fn(),
}),
getWorkingDir: () => '/test/dir',
getWorkingDir: () => path.resolve('/test/dir'),
shouldLoadMemoryFromIncludeDirectories: () => false,
getDebugMode: () => false,
getFileService: () => ({}),
@@ -91,9 +101,9 @@ describe('directoryCommand', () => {
expect(mockContext.ui.addItem).toHaveBeenCalledWith(
expect.objectContaining({
type: MessageType.INFO,
text: `Current workspace directories:\n- ${path.normalize(
text: `Current workspace directories:\n- ${path.resolve(
'/home/user/project1',
)}\n- ${path.normalize('/home/user/project2')}`,
)}\n- ${path.resolve('/home/user/project2')}`,
}),
);
});
@@ -125,7 +135,7 @@ describe('directoryCommand', () => {
});
it('should call addDirectory and show a success message for a single path', async () => {
const newPath = path.normalize('/home/user/new-project');
const newPath = path.resolve('/home/user/new-project');
vi.mocked(mockWorkspaceContext.addDirectories).mockReturnValue({
added: [newPath],
failed: [],
@@ -144,8 +154,8 @@ describe('directoryCommand', () => {
});
it('should call addDirectory for each path and show a success message for multiple paths', async () => {
const newPath1 = path.normalize('/home/user/new-project1');
const newPath2 = path.normalize('/home/user/new-project2');
const newPath1 = path.resolve('/home/user/new-project1');
const newPath2 = path.resolve('/home/user/new-project2');
vi.mocked(mockWorkspaceContext.addDirectories).mockReturnValue({
added: [newPath1, newPath2],
failed: [],
@@ -166,7 +176,7 @@ describe('directoryCommand', () => {
it('should show an error if addDirectory throws an exception', async () => {
const error = new Error('Directory does not exist');
const newPath = path.normalize('/home/user/invalid-project');
const newPath = path.resolve('/home/user/invalid-project');
vi.mocked(mockWorkspaceContext.addDirectories).mockReturnValue({
added: [],
failed: [{ path: newPath, error }],
@@ -184,7 +194,7 @@ describe('directoryCommand', () => {
it('should add directory directly when folder trust is disabled', async () => {
if (!addCommand?.action) throw new Error('No action');
vi.spyOn(trustedFolders, 'isFolderTrustEnabled').mockReturnValue(false);
const newPath = path.normalize('/home/user/new-project');
const newPath = path.resolve('/home/user/new-project');
vi.mocked(mockWorkspaceContext.addDirectories).mockReturnValue({
added: [newPath],
failed: [],
@@ -198,7 +208,7 @@ describe('directoryCommand', () => {
});
it('should show an info message for an already added directory', async () => {
const existingPath = path.normalize('/home/user/project1');
const existingPath = path.resolve('/home/user/project1');
if (!addCommand?.action) throw new Error('No action');
await addCommand.action(mockContext, existingPath);
expect(mockContext.ui.addItem).toHaveBeenCalledWith(
@@ -212,9 +222,33 @@ describe('directoryCommand', () => {
);
});
it('should show an info message for an already added directory specified as a relative path', async () => {
const existingPath = path.resolve('/home/user/project1');
const relativePath = './project1';
const absoluteRelativePath = path.resolve(
path.resolve('/test/dir'),
relativePath,
);
vi.mocked(fs.realpathSync).mockImplementation((p) => {
if (p === absoluteRelativePath) return existingPath;
return p as string;
});
if (!addCommand?.action) throw new Error('No action');
await addCommand.action(mockContext, relativePath);
expect(mockContext.ui.addItem).toHaveBeenCalledWith(
expect.objectContaining({
type: MessageType.INFO,
text: `The following directories are already in the workspace:\n- ${relativePath}`,
}),
);
});
it('should handle a mix of successful and failed additions', async () => {
const validPath = path.normalize('/home/user/valid-project');
const invalidPath = path.normalize('/home/user/invalid-project');
const validPath = path.resolve('/home/user/valid-project');
const invalidPath = path.resolve('/home/user/invalid-project');
const error = new Error('Directory does not exist');
vi.mocked(mockWorkspaceContext.addDirectories).mockReturnValue({
added: [validPath],
@@ -318,7 +352,7 @@ describe('directoryCommand', () => {
it('should add a trusted directory', async () => {
if (!addCommand?.action) throw new Error('No action');
mockIsPathTrusted.mockReturnValue(true);
const newPath = path.normalize('/home/user/trusted-project');
const newPath = path.resolve('/home/user/trusted-project');
vi.mocked(mockWorkspaceContext.addDirectories).mockReturnValue({
added: [newPath],
failed: [],
@@ -334,7 +368,7 @@ describe('directoryCommand', () => {
it('should return a custom dialog for an explicitly untrusted directory (upgrade flow)', async () => {
if (!addCommand?.action) throw new Error('No action');
mockIsPathTrusted.mockReturnValue(false); // DO_NOT_TRUST
const newPath = path.normalize('/home/user/untrusted-project');
const newPath = path.resolve('/home/user/untrusted-project');
const result = await addCommand.action(mockContext, newPath);
@@ -357,7 +391,7 @@ describe('directoryCommand', () => {
it('should return a custom dialog for a directory with undefined trust', async () => {
if (!addCommand?.action) throw new Error('No action');
mockIsPathTrusted.mockReturnValue(undefined);
const newPath = path.normalize('/home/user/undefined-trust-project');
const newPath = path.resolve('/home/user/undefined-trust-project');
const result = await addCommand.action(mockContext, newPath);
@@ -385,7 +419,7 @@ describe('directoryCommand', () => {
source: 'file',
});
mockIsPathTrusted.mockReturnValue(undefined);
const newPath = path.normalize('/home/user/new-project');
const newPath = path.resolve('/home/user/new-project');
const result = await addCommand.action(mockContext, newPath);
@@ -20,6 +20,7 @@ import {
} from '../utils/directoryUtils.js';
import type { Config } from '@google/gemini-cli-core';
import * as path from 'node:path';
import * as fs from 'node:fs';
async function finishAddingDirectories(
config: Config,
@@ -100,7 +101,7 @@ export const directoryCommand: SlashCommand = {
const workspaceContext =
context.services.config.getWorkspaceContext();
const existingDirs = new Set(
workspaceContext.getDirectories().map((dir) => path.normalize(dir)),
workspaceContext.getDirectories().map((dir) => path.resolve(dir)),
);
filteredSuggestions = suggestions.filter((s) => {
@@ -172,12 +173,23 @@ export const directoryCommand: SlashCommand = {
const pathsToProcess: string[] = [];
for (const pathToAdd of pathsToAdd) {
const expandedPath = expandHomeDir(pathToAdd.trim());
if (currentWorkspaceDirs.includes(expandedPath)) {
alreadyAdded.push(pathToAdd.trim());
} else {
pathsToProcess.push(pathToAdd.trim());
const trimmedPath = pathToAdd.trim();
const expandedPath = expandHomeDir(trimmedPath);
try {
const absolutePath = path.resolve(
workspaceContext.targetDir,
expandedPath,
);
const resolvedPath = fs.realpathSync(absolutePath);
if (currentWorkspaceDirs.includes(resolvedPath)) {
alreadyAdded.push(trimmedPath);
continue;
}
} catch (_e) {
// Path might not exist or be inaccessible.
// We'll let batchAddDirectories handle it later.
}
pathsToProcess.push(trimmedPath);
}
if (alreadyAdded.length > 0) {
@@ -32,6 +32,7 @@ import process from 'node:process';
import { type UseHistoryManagerReturn } from '../hooks/useHistoryManager.js';
import { AdminSettingsChangedDialog } from './AdminSettingsChangedDialog.js';
import { IdeTrustChangeDialog } from './IdeTrustChangeDialog.js';
import { NewAgentsNotification } from './NewAgentsNotification.js';
import { AgentConfigDialog } from './AgentConfigDialog.js';
interface DialogManagerProps {
@@ -58,6 +59,14 @@ export const DialogManager = ({
if (uiState.showIdeRestartPrompt) {
return <IdeTrustChangeDialog reason={uiState.ideTrustRestartReason} />;
}
if (uiState.newAgents) {
return (
<NewAgentsNotification
agents={uiState.newAgents}
onSelect={uiActions.handleNewAgentsSelect}
/>
);
}
if (uiState.proQuotaRequest) {
return (
<ProQuotaDialog
@@ -0,0 +1,57 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { describe, it, expect, vi } from 'vitest';
import { renderWithProviders as render } from '../../test-utils/render.js';
import { NewAgentsNotification } from './NewAgentsNotification.js';
describe('NewAgentsNotification', () => {
const mockAgents = [
{
name: 'Agent A',
description: 'Description A',
kind: 'remote' as const,
agentCardUrl: '',
inputConfig: { inputSchema: {} },
},
{
name: 'Agent B',
description: 'Description B',
kind: 'remote' as const,
agentCardUrl: '',
inputConfig: { inputSchema: {} },
},
];
const onSelect = vi.fn();
it('renders agent list', () => {
const { lastFrame, unmount } = render(
<NewAgentsNotification agents={mockAgents} onSelect={onSelect} />,
);
const frame = lastFrame();
expect(frame).toMatchSnapshot();
unmount();
});
it('truncates list if more than 5 agents', () => {
const manyAgents = Array.from({ length: 7 }, (_, i) => ({
name: `Agent ${i}`,
description: `Description ${i}`,
kind: 'remote' as const,
agentCardUrl: '',
inputConfig: { inputSchema: {} },
}));
const { lastFrame, unmount } = render(
<NewAgentsNotification agents={manyAgents} onSelect={onSelect} />,
);
const frame = lastFrame();
expect(frame).toMatchSnapshot();
unmount();
});
});
@@ -0,0 +1,96 @@
/**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { Box, Text } from 'ink';
import { type AgentDefinition } from '@google/gemini-cli-core';
import { theme } from '../semantic-colors.js';
import {
RadioButtonSelect,
type RadioSelectItem,
} from './shared/RadioButtonSelect.js';
export enum NewAgentsChoice {
ACKNOWLEDGE = 'acknowledge',
IGNORE = 'ignore',
}
interface NewAgentsNotificationProps {
agents: AgentDefinition[];
onSelect: (choice: NewAgentsChoice) => void;
}
export const NewAgentsNotification = ({
agents,
onSelect,
}: NewAgentsNotificationProps) => {
const options: Array<RadioSelectItem<NewAgentsChoice>> = [
{
label: 'Acknowledge and Enable',
value: NewAgentsChoice.ACKNOWLEDGE,
key: 'acknowledge',
},
{
label: 'Do not enable (Ask again next time)',
value: NewAgentsChoice.IGNORE,
key: 'ignore',
},
];
// Limit display to 5 agents to avoid overflow, show count for rest
const MAX_DISPLAYED_AGENTS = 5;
const displayAgents = agents.slice(0, MAX_DISPLAYED_AGENTS);
const remaining = agents.length - MAX_DISPLAYED_AGENTS;
return (
<Box flexDirection="column" width="100%">
<Box
flexDirection="column"
borderStyle="round"
borderColor={theme.status.warning}
padding={1}
marginLeft={1}
marginRight={1}
>
<Box flexDirection="column" marginBottom={1}>
<Text bold color={theme.text.primary}>
New Agents Discovered
</Text>
<Text color={theme.text.primary}>
The following agents were found in this project. Please review them:
</Text>
<Box
flexDirection="column"
marginTop={1}
borderStyle="single"
padding={1}
>
{displayAgents.map((agent) => (
<Box key={agent.name}>
<Box flexShrink={0}>
<Text bold color={theme.text.primary}>
- {agent.name}:{' '}
</Text>
</Box>
<Text color={theme.text.secondary}> {agent.description}</Text>
</Box>
))}
{remaining > 0 && (
<Text color={theme.text.secondary}>
... and {remaining} more.
</Text>
)}
</Box>
</Box>
<RadioButtonSelect
items={options}
onSelect={onSelect}
isFocused={true}
/>
</Box>
</Box>
);
};
@@ -0,0 +1,43 @@
// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
exports[`NewAgentsNotification > renders agent list 1`] = `
" ╭────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╮
│ │
│ New Agents Discovered │
│ The following agents were found in this project. Please review them: │
│ │
│ ┌────────────────────────────────────────────────────────────────────────────────────────────────────────────────┐ │
│ │ │ │
│ │ - Agent A: Description A │ │
│ │ - Agent B: Description B │ │
│ │ │ │
│ └────────────────────────────────────────────────────────────────────────────────────────────────────────────────┘ │
│ │
│ ● 1. Acknowledge and Enable │
│ 2. Do not enable (Ask again next time) │
│ │
╰────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯"
`;
exports[`NewAgentsNotification > truncates list if more than 5 agents 1`] = `
" ╭────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╮
│ │
│ New Agents Discovered │
│ The following agents were found in this project. Please review them: │
│ │
│ ┌────────────────────────────────────────────────────────────────────────────────────────────────────────────────┐ │
│ │ │ │
│ │ - Agent 0: Description 0 │ │
│ │ - Agent 1: Description 1 │ │
│ │ - Agent 2: Description 2 │ │
│ │ - Agent 3: Description 3 │ │
│ │ - Agent 4: Description 4 │ │
│ │ ... and 2 more. │ │
│ │ │ │
│ └────────────────────────────────────────────────────────────────────────────────────────────────────────────────┘ │
│ │
│ ● 1. Acknowledge and Enable │
│ 2. Do not enable (Ask again next time) │
│ │
╰────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯"
`;
@@ -17,6 +17,7 @@ import { type LoadableSettingScope } from '../../config/settings.js';
import type { AuthState } from '../types.js';
import { type PermissionsDialogProps } from '../components/PermissionsModifyTrustDialog.js';
import type { SessionInfo } from '../../utils/sessionUtils.js';
import { type NewAgentsChoice } from '../components/NewAgentsNotification.js';
export interface UIActions {
handleThemeSelect: (themeName: string, scope: LoadableSettingScope) => void;
@@ -69,6 +70,7 @@ export interface UIActions {
setEmbeddedShellFocused: (value: boolean) => void;
setAuthContext: (context: { requiresRestart?: boolean }) => void;
handleRestart: () => void;
handleNewAgentsSelect: (choice: NewAgentsChoice) => Promise<void>;
}
export const UIActionsContext = createContext<UIActions | null>(null);
@@ -155,6 +155,7 @@ export interface UIState {
terminalBackgroundColor: TerminalBackgroundColor;
settingsNonce: number;
adminSettingsChanged: boolean;
newAgents: AgentDefinition[] | null;
}
export const UIStateContext = createContext<UIState | null>(null);
@@ -19,6 +19,7 @@ import {
type ToolCallsUpdateMessage,
type AnyDeclarativeTool,
type AnyToolInvocation,
ROOT_SCHEDULER_ID,
} from '@google/gemini-cli-core';
import { createMockMessageBus } from '@google/gemini-cli-core/src/test-utils/mock-message-bus.js';
@@ -73,6 +74,10 @@ describe('useToolExecutionScheduler', () => {
} as unknown as Config;
});
afterEach(() => {
vi.clearAllMocks();
});
it('initializes with empty tool calls', () => {
const { result } = renderHook(() =>
useToolExecutionScheduler(
@@ -112,6 +117,7 @@ describe('useToolExecutionScheduler', () => {
void mockMessageBus.publish({
type: MessageBusType.TOOL_CALLS_UPDATE,
toolCalls: [mockToolCall],
schedulerId: ROOT_SCHEDULER_ID,
} as ToolCallsUpdateMessage);
});
@@ -156,6 +162,7 @@ describe('useToolExecutionScheduler', () => {
void mockMessageBus.publish({
type: MessageBusType.TOOL_CALLS_UPDATE,
toolCalls: [mockToolCall],
schedulerId: ROOT_SCHEDULER_ID,
} as ToolCallsUpdateMessage);
});
@@ -212,6 +219,7 @@ describe('useToolExecutionScheduler', () => {
void mockMessageBus.publish({
type: MessageBusType.TOOL_CALLS_UPDATE,
toolCalls: [mockToolCall],
schedulerId: ROOT_SCHEDULER_ID,
} as ToolCallsUpdateMessage);
});
@@ -274,6 +282,7 @@ describe('useToolExecutionScheduler', () => {
void mockMessageBus.publish({
type: MessageBusType.TOOL_CALLS_UPDATE,
toolCalls: [mockToolCall],
schedulerId: ROOT_SCHEDULER_ID,
} as ToolCallsUpdateMessage);
});
@@ -290,6 +299,7 @@ describe('useToolExecutionScheduler', () => {
void mockMessageBus.publish({
type: MessageBusType.TOOL_CALLS_UPDATE,
toolCalls: [mockToolCall],
schedulerId: ROOT_SCHEDULER_ID,
} as ToolCallsUpdateMessage);
});
@@ -326,6 +336,7 @@ describe('useToolExecutionScheduler', () => {
invocation: createMockInvocation(),
},
],
schedulerId: ROOT_SCHEDULER_ID,
} as ToolCallsUpdateMessage);
});
@@ -412,4 +423,103 @@ describe('useToolExecutionScheduler', () => {
expect(completedResult).toEqual([completedToolCall]);
expect(onComplete).toHaveBeenCalledWith([completedToolCall]);
});
it('setToolCallsForDisplay re-groups tools by schedulerId (Multi-Scheduler support)', () => {
const { result } = renderHook(() =>
useToolExecutionScheduler(
vi.fn().mockResolvedValue(undefined),
mockConfig,
() => undefined,
),
);
const callRoot = {
status: 'success' as const,
request: {
callId: 'call-root',
name: 'test',
args: {},
isClientInitiated: false,
prompt_id: 'p1',
},
tool: createMockTool(),
invocation: createMockInvocation(),
response: {
callId: 'call-root',
responseParts: [],
resultDisplay: 'OK',
error: undefined,
errorType: undefined,
},
schedulerId: ROOT_SCHEDULER_ID,
};
const callSub = {
...callRoot,
request: { ...callRoot.request, callId: 'call-sub' },
schedulerId: 'subagent-1',
};
// 1. Populate state with multiple schedulers
act(() => {
void mockMessageBus.publish({
type: MessageBusType.TOOL_CALLS_UPDATE,
toolCalls: [callRoot],
schedulerId: ROOT_SCHEDULER_ID,
} as ToolCallsUpdateMessage);
void mockMessageBus.publish({
type: MessageBusType.TOOL_CALLS_UPDATE,
toolCalls: [callSub],
schedulerId: 'subagent-1',
} as ToolCallsUpdateMessage);
});
let [toolCalls] = result.current;
expect(toolCalls).toHaveLength(2);
expect(
toolCalls.find((t) => t.request.callId === 'call-root')?.schedulerId,
).toBe(ROOT_SCHEDULER_ID);
expect(
toolCalls.find((t) => t.request.callId === 'call-sub')?.schedulerId,
).toBe('subagent-1');
// 2. Call setToolCallsForDisplay (e.g., simulate a manual update or clear)
act(() => {
const [, , , setToolCalls] = result.current;
setToolCalls((prev) =>
prev.map((t) => ({ ...t, responseSubmittedToGemini: true })),
);
});
// 3. Verify that tools are still present and maintain their scheduler IDs
// The internal map should have been re-grouped.
[toolCalls] = result.current;
expect(toolCalls).toHaveLength(2);
expect(toolCalls.every((t) => t.responseSubmittedToGemini)).toBe(true);
const updatedRoot = toolCalls.find((t) => t.request.callId === 'call-root');
const updatedSub = toolCalls.find((t) => t.request.callId === 'call-sub');
expect(updatedRoot?.schedulerId).toBe(ROOT_SCHEDULER_ID);
expect(updatedSub?.schedulerId).toBe('subagent-1');
// 4. Verify that a subsequent update to ONE scheduler doesn't wipe the other
act(() => {
void mockMessageBus.publish({
type: MessageBusType.TOOL_CALLS_UPDATE,
toolCalls: [{ ...callRoot, status: 'executing' }],
schedulerId: ROOT_SCHEDULER_ID,
} as ToolCallsUpdateMessage);
});
[toolCalls] = result.current;
expect(toolCalls).toHaveLength(2);
expect(
toolCalls.find((t) => t.request.callId === 'call-root')?.status,
).toBe('executing');
expect(
toolCalls.find((t) => t.request.callId === 'call-sub')?.schedulerId,
).toBe('subagent-1');
});
});
@@ -16,6 +16,7 @@ import {
Scheduler,
type EditorType,
type ToolCallsUpdateMessage,
ROOT_SCHEDULER_ID,
} from '@google/gemini-cli-core';
import { useCallback, useState, useMemo, useEffect, useRef } from 'react';
@@ -54,8 +55,10 @@ export function useToolExecutionScheduler(
CancelAllFn,
number,
] {
// State stores Core objects, not Display objects
const [toolCalls, setToolCalls] = useState<TrackedToolCall[]>([]);
// State stores tool calls organized by their originating schedulerId
const [toolCallsMap, setToolCallsMap] = useState<
Record<string, TrackedToolCall[]>
>({});
const [lastToolOutputTime, setLastToolOutputTime] = useState<number>(0);
const messageBus = useMemo(() => config.getMessageBus(), [config]);
@@ -76,6 +79,7 @@ export function useToolExecutionScheduler(
config,
messageBus,
getPreferredEditor: () => getPreferredEditorRef.current(),
schedulerId: ROOT_SCHEDULER_ID,
}),
[config, messageBus],
);
@@ -88,15 +92,21 @@ export function useToolExecutionScheduler(
useEffect(() => {
const handler = (event: ToolCallsUpdateMessage) => {
setToolCalls((prev) => {
const adapted = internalAdaptToolCalls(event.toolCalls, prev);
// Update output timer for UI spinners (Side Effect)
if (event.toolCalls.some((tc) => tc.status === 'executing')) {
setLastToolOutputTime(Date.now());
}
// Update output timer for UI spinners
if (event.toolCalls.some((tc) => tc.status === 'executing')) {
setLastToolOutputTime(Date.now());
}
setToolCallsMap((prev) => {
const adapted = internalAdaptToolCalls(
event.toolCalls,
prev[event.schedulerId] ?? [],
);
return adapted;
return {
...prev,
[event.schedulerId]: adapted,
};
});
};
@@ -109,12 +119,14 @@ export function useToolExecutionScheduler(
const schedule: ScheduleFn = useCallback(
async (request, signal) => {
// Clear state for new run
setToolCalls([]);
setToolCallsMap({});
// 1. Await Core Scheduler directly
const results = await scheduler.schedule(request, signal);
// 2. Trigger legacy reinjection logic (useGeminiStream loop)
// Since this hook instance owns the "root" scheduler, we always trigger
// onComplete when it finishes its batch.
await onCompleteRef.current(results);
return results;
@@ -131,13 +143,52 @@ export function useToolExecutionScheduler(
const markToolsAsSubmitted: MarkToolsAsSubmittedFn = useCallback(
(callIdsToMark: string[]) => {
setToolCalls((prevCalls) =>
prevCalls.map((tc) =>
callIdsToMark.includes(tc.request.callId)
? { ...tc, responseSubmittedToGemini: true }
: tc,
),
);
setToolCallsMap((prevMap) => {
const nextMap = { ...prevMap };
for (const [sid, calls] of Object.entries(nextMap)) {
nextMap[sid] = calls.map((tc) =>
callIdsToMark.includes(tc.request.callId)
? { ...tc, responseSubmittedToGemini: true }
: tc,
);
}
return nextMap;
});
},
[],
);
// Flatten the map for the UI components that expect a single list of tools.
const toolCalls = useMemo(
() => Object.values(toolCallsMap).flat(),
[toolCallsMap],
);
// Provide a setter that maintains compatibility with legacy [].
const setToolCallsForDisplay = useCallback(
(action: React.SetStateAction<TrackedToolCall[]>) => {
setToolCallsMap((prev) => {
const currentFlattened = Object.values(prev).flat();
const nextFlattened =
typeof action === 'function' ? action(currentFlattened) : action;
if (nextFlattened.length === 0) {
return {};
}
// Re-group by schedulerId to preserve multi-scheduler state
const nextMap: Record<string, TrackedToolCall[]> = {};
for (const call of nextFlattened) {
// All tool calls should have a schedulerId from the core.
// Default to ROOT_SCHEDULER_ID as a safeguard.
const sid = call.schedulerId ?? ROOT_SCHEDULER_ID;
if (!nextMap[sid]) {
nextMap[sid] = [];
}
nextMap[sid].push(call);
}
return nextMap;
});
},
[],
);
@@ -146,7 +197,7 @@ export function useToolExecutionScheduler(
toolCalls,
schedule,
markToolsAsSubmitted,
setToolCalls,
setToolCallsForDisplay,
cancelAll,
lastToolOutputTime,
];
+138
View File
@@ -89,6 +89,7 @@ const TEST_SEQUENCES = {
LINE_START: createKey({ sequence: '0' }),
LINE_END: createKey({ sequence: '$' }),
REPEAT: createKey({ sequence: '.' }),
CTRL_C: createKey({ sequence: '\x03', name: 'c', ctrl: true }),
} as const;
describe('useVim hook', () => {
@@ -1614,4 +1615,141 @@ describe('useVim hook', () => {
},
);
});
describe('double-escape to clear buffer', () => {
beforeEach(() => {
mockBuffer = createMockBuffer('hello world');
mockVimContext.vimEnabled = true;
mockVimContext.vimMode = 'NORMAL';
mockHandleFinalSubmit = vi.fn();
vi.useFakeTimers();
});
afterEach(() => {
vi.useRealTimers();
});
it('should clear buffer on double-escape in NORMAL mode', async () => {
const { result } = renderHook(() =>
useVim(mockBuffer as TextBuffer, mockHandleFinalSubmit),
);
// First escape - should pass through (return false)
let handled: boolean;
await act(async () => {
handled = result.current.handleInput(TEST_SEQUENCES.ESCAPE);
});
expect(handled!).toBe(false);
// Second escape within timeout - should clear buffer (return true)
await act(async () => {
handled = result.current.handleInput(TEST_SEQUENCES.ESCAPE);
});
expect(handled!).toBe(true);
expect(mockBuffer.setText).toHaveBeenCalledWith('');
});
it('should clear buffer on double-escape in INSERT mode', async () => {
mockVimContext.vimMode = 'INSERT';
const { result } = renderHook(() =>
useVim(mockBuffer as TextBuffer, mockHandleFinalSubmit),
);
// First escape - switches to NORMAL mode
let handled: boolean;
await act(async () => {
handled = result.current.handleInput(TEST_SEQUENCES.ESCAPE);
});
expect(handled!).toBe(true);
expect(mockBuffer.vimEscapeInsertMode).toHaveBeenCalled();
// Second escape within timeout - should clear buffer
await act(async () => {
handled = result.current.handleInput(TEST_SEQUENCES.ESCAPE);
});
expect(handled!).toBe(true);
expect(mockBuffer.setText).toHaveBeenCalledWith('');
});
it('should NOT clear buffer if escapes are too slow', async () => {
const { result } = renderHook(() =>
useVim(mockBuffer as TextBuffer, mockHandleFinalSubmit),
);
// First escape
await act(async () => {
result.current.handleInput(TEST_SEQUENCES.ESCAPE);
});
// Wait longer than timeout (500ms)
await act(async () => {
vi.advanceTimersByTime(600);
});
// Second escape - should NOT clear buffer because timeout expired
let handled: boolean;
await act(async () => {
handled = result.current.handleInput(TEST_SEQUENCES.ESCAPE);
});
// First escape of new sequence, passes through
expect(handled!).toBe(false);
expect(mockBuffer.setText).not.toHaveBeenCalled();
});
it('should clear escape history when clearing pending operator', async () => {
const { result } = renderHook(() =>
useVim(mockBuffer as TextBuffer, mockHandleFinalSubmit),
);
// First escape
await act(async () => {
result.current.handleInput(TEST_SEQUENCES.ESCAPE);
});
// Type 'd' to set pending operator
await act(async () => {
result.current.handleInput(TEST_SEQUENCES.DELETE);
});
// Escape to clear pending operator
await act(async () => {
result.current.handleInput(TEST_SEQUENCES.ESCAPE);
});
// Another escape - should NOT clear buffer (history was reset)
let handled: boolean;
await act(async () => {
handled = result.current.handleInput(TEST_SEQUENCES.ESCAPE);
});
expect(handled!).toBe(false);
expect(mockBuffer.setText).not.toHaveBeenCalled();
});
it('should pass Ctrl+C through to InputPrompt in NORMAL mode', async () => {
const { result } = renderHook(() =>
useVim(mockBuffer as TextBuffer, mockHandleFinalSubmit),
);
let handled: boolean;
await act(async () => {
handled = result.current.handleInput(TEST_SEQUENCES.CTRL_C);
});
// Should return false to let InputPrompt handle it
expect(handled!).toBe(false);
});
it('should pass Ctrl+C through to InputPrompt in INSERT mode', async () => {
mockVimContext.vimMode = 'INSERT';
const { result } = renderHook(() =>
useVim(mockBuffer as TextBuffer, mockHandleFinalSubmit),
);
let handled: boolean;
await act(async () => {
handled = result.current.handleInput(TEST_SEQUENCES.CTRL_C);
});
// Should return false to let InputPrompt handle it
expect(handled!).toBe(false);
});
});
});
+40 -9
View File
@@ -4,11 +4,12 @@
* SPDX-License-Identifier: Apache-2.0
*/
import { useCallback, useReducer, useEffect } from 'react';
import { useCallback, useReducer, useEffect, useRef } from 'react';
import type { Key } from './useKeypress.js';
import type { TextBuffer } from '../components/shared/text-buffer.js';
import { useVimMode } from '../contexts/VimModeContext.js';
import { debugLogger } from '@google/gemini-cli-core';
import { keyMatchers, Command } from '../keyMatchers.js';
export type VimMode = 'NORMAL' | 'INSERT';
@@ -16,6 +17,7 @@ export type VimMode = 'NORMAL' | 'INSERT';
const DIGIT_MULTIPLIER = 10;
const DEFAULT_COUNT = 1;
const DIGIT_1_TO_9 = /^[1-9]$/;
const DOUBLE_ESCAPE_TIMEOUT_MS = 500; // Timeout for double-escape to clear input
// Command types
const CMD_TYPES = {
@@ -130,6 +132,9 @@ export function useVim(buffer: TextBuffer, onSubmit?: (value: string) => void) {
const { vimEnabled, vimMode, setVimMode } = useVimMode();
const [state, dispatch] = useReducer(vimReducer, initialVimState);
// Track last escape timestamp for double-escape detection
const lastEscapeTimestampRef = useRef<number>(0);
// Sync vim mode from context to local state
useEffect(() => {
dispatch({ type: 'SET_MODE', mode: vimMode });
@@ -150,6 +155,19 @@ export function useVim(buffer: TextBuffer, onSubmit?: (value: string) => void) {
[state.count],
);
// Returns true if two escapes occurred within DOUBLE_ESCAPE_TIMEOUT_MS.
const checkDoubleEscape = useCallback((): boolean => {
const now = Date.now();
const lastEscape = lastEscapeTimestampRef.current;
lastEscapeTimestampRef.current = now;
if (now - lastEscape <= DOUBLE_ESCAPE_TIMEOUT_MS) {
lastEscapeTimestampRef.current = 0;
return true;
}
return false;
}, []);
/** Executes common commands to eliminate duplication in dot (.) repeat command */
const executeCommand = useCallback(
(cmdType: string, count: number) => {
@@ -247,9 +265,9 @@ export function useVim(buffer: TextBuffer, onSubmit?: (value: string) => void) {
*/
const handleInsertModeInput = useCallback(
(normalizedKey: Key): boolean => {
// Handle escape key immediately - switch to NORMAL mode on any escape
if (normalizedKey.name === 'escape') {
// Vim behavior: move cursor left when exiting insert mode (unless at beginning of line)
if (keyMatchers[Command.ESCAPE](normalizedKey)) {
// Record for double-escape detection (clearing happens in NORMAL mode)
checkDoubleEscape();
buffer.vimEscapeInsertMode();
dispatch({ type: 'ESCAPE_TO_NORMAL' });
updateMode('NORMAL');
@@ -298,7 +316,7 @@ export function useVim(buffer: TextBuffer, onSubmit?: (value: string) => void) {
buffer.handleInput(normalizedKey);
return true; // Handled by vim
},
[buffer, dispatch, updateMode, onSubmit],
[buffer, dispatch, updateMode, onSubmit, checkDoubleEscape],
);
/**
@@ -401,6 +419,11 @@ export function useVim(buffer: TextBuffer, onSubmit?: (value: string) => void) {
return false;
}
// Let InputPrompt handle Ctrl+C for clearing input (works in all modes)
if (keyMatchers[Command.CLEAR_INPUT](normalizedKey)) {
return false;
}
// Handle INSERT mode
if (state.mode === 'INSERT') {
return handleInsertModeInput(normalizedKey);
@@ -408,14 +431,21 @@ export function useVim(buffer: TextBuffer, onSubmit?: (value: string) => void) {
// Handle NORMAL mode
if (state.mode === 'NORMAL') {
// If in NORMAL mode, allow escape to pass through to other handlers
// if there's no pending operation.
if (normalizedKey.name === 'escape') {
if (keyMatchers[Command.ESCAPE](normalizedKey)) {
if (state.pendingOperator) {
dispatch({ type: 'CLEAR_PENDING_STATES' });
lastEscapeTimestampRef.current = 0;
return true; // Handled by vim
}
return false; // Pass through to other handlers
// Check for double-escape to clear buffer
if (checkDoubleEscape()) {
buffer.setText('');
return true;
}
// First escape in NORMAL mode - pass through for UI feedback
return false;
}
// Handle count input (numbers 1-9, and 0 if count > 0)
@@ -776,6 +806,7 @@ export function useVim(buffer: TextBuffer, onSubmit?: (value: string) => void) {
buffer,
executeCommand,
updateMode,
checkDoubleEscape,
],
);