From b8319bee76afea3361f9c5de70047bde035ee7ba Mon Sep 17 00:00:00 2001 From: Harsha Nadimpalli Date: Mon, 26 Jan 2026 09:36:42 -0800 Subject: [PATCH 1/9] feat(cli): add quick clear input shortcuts in vim mode (#17470) Co-authored-by: Tommaso Sciortino --- packages/cli/src/ui/hooks/vim.test.tsx | 138 +++++++++++++++++++++++++ packages/cli/src/ui/hooks/vim.ts | 49 +++++++-- 2 files changed, 178 insertions(+), 9 deletions(-) diff --git a/packages/cli/src/ui/hooks/vim.test.tsx b/packages/cli/src/ui/hooks/vim.test.tsx index 372f5f03e4..bc0d15d9d1 100644 --- a/packages/cli/src/ui/hooks/vim.test.tsx +++ b/packages/cli/src/ui/hooks/vim.test.tsx @@ -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); + }); + }); }); diff --git a/packages/cli/src/ui/hooks/vim.ts b/packages/cli/src/ui/hooks/vim.ts index 2f39c38f43..2762c54de7 100644 --- a/packages/cli/src/ui/hooks/vim.ts +++ b/packages/cli/src/ui/hooks/vim.ts @@ -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(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, ], ); From 50f89e8a41e8d2d3dc5360bb3514ee277cb85ffb Mon Sep 17 00:00:00 2001 From: Sandy Tao Date: Mon, 26 Jan 2026 10:12:21 -0800 Subject: [PATCH 2/9] feat(core): optimize shell tool llmContent output format (#17538) --- .../tools/__snapshots__/shell.test.ts.snap | 30 ++-- packages/core/src/tools/shell.test.ts | 129 ++++++++++++++++++ packages/core/src/tools/shell.ts | 54 ++++---- 3 files changed, 171 insertions(+), 42 deletions(-) diff --git a/packages/core/src/tools/__snapshots__/shell.test.ts.snap b/packages/core/src/tools/__snapshots__/shell.test.ts.snap index 76a5ded3ef..6592993160 100644 --- a/packages/core/src/tools/__snapshots__/shell.test.ts.snap +++ b/packages/core/src/tools/__snapshots__/shell.test.ts.snap @@ -5,15 +5,12 @@ exports[`ShellTool > getDescription > should return the non-windows description The following information is returned: - Command: Executed command. - Directory: Directory where command was executed, or \`(root)\`. - Stdout: Output on stdout stream. Can be \`(empty)\` or partial on error and for any unwaited background processes. - Stderr: Output on stderr stream. Can be \`(empty)\` or partial on error and for any unwaited background processes. - Error: Error or \`(none)\` if no error was reported for the subprocess. - Exit Code: Exit code or \`(none)\` if terminated by signal. - Signal: Signal number or \`(none)\` if no signal was received. - Background PIDs: List of background processes started or \`(none)\`. - Process Group PGID: Process group started or \`(none)\`" + Output: Combined stdout/stderr. Can be \`(empty)\` or partial on error and for any unwaited background processes. + Exit Code: Only included if non-zero (command failed). + Error: Only included if a process-level error occurred (e.g., spawn failure). + Signal: Only included if process was terminated by a signal. + Background PIDs: Only included if background processes were started. + Process Group PGID: Only included if available." `; exports[`ShellTool > getDescription > should return the windows description when on windows 1`] = ` @@ -21,13 +18,10 @@ exports[`ShellTool > getDescription > should return the windows description when The following information is returned: - Command: Executed command. - Directory: Directory where command was executed, or \`(root)\`. - Stdout: Output on stdout stream. Can be \`(empty)\` or partial on error and for any unwaited background processes. - Stderr: Output on stderr stream. Can be \`(empty)\` or partial on error and for any unwaited background processes. - Error: Error or \`(none)\` if no error was reported for the subprocess. - Exit Code: Exit code or \`(none)\` if terminated by signal. - Signal: Signal number or \`(none)\` if no signal was received. - Background PIDs: List of background processes started or \`(none)\`. - Process Group PGID: Process group started or \`(none)\`" + Output: Combined stdout/stderr. Can be \`(empty)\` or partial on error and for any unwaited background processes. + Exit Code: Only included if non-zero (command failed). + Error: Only included if a process-level error occurred (e.g., spawn failure). + Signal: Only included if process was terminated by a signal. + Background PIDs: Only included if background processes were started. + Process Group PGID: Only included if available." `; diff --git a/packages/core/src/tools/shell.test.ts b/packages/core/src/tools/shell.test.ts index 9b05afec36..196c8678e9 100644 --- a/packages/core/src/tools/shell.test.ts +++ b/packages/core/src/tools/shell.test.ts @@ -539,6 +539,135 @@ describe('ShellTool', () => { }); }); + describe('llmContent output format', () => { + const mockAbortSignal = new AbortController().signal; + + const resolveShellExecution = ( + result: Partial = {}, + ) => { + const fullResult: ShellExecutionResult = { + rawOutput: Buffer.from(result.output || ''), + output: 'Success', + exitCode: 0, + signal: null, + error: null, + aborted: false, + pid: 12345, + executionMethod: 'child_process', + ...result, + }; + resolveExecutionPromise(fullResult); + }; + + it('should not include Command in output', async () => { + const invocation = shellTool.build({ command: 'echo hello' }); + const promise = invocation.execute(mockAbortSignal); + resolveShellExecution({ output: 'hello', exitCode: 0 }); + + const result = await promise; + expect(result.llmContent).not.toContain('Command:'); + }); + + it('should not include Directory in output', async () => { + const invocation = shellTool.build({ command: 'ls', dir_path: 'subdir' }); + const promise = invocation.execute(mockAbortSignal); + resolveShellExecution({ output: 'file.txt', exitCode: 0 }); + + const result = await promise; + expect(result.llmContent).not.toContain('Directory:'); + }); + + it('should not include Exit Code when command succeeds (exit code 0)', async () => { + const invocation = shellTool.build({ command: 'echo hello' }); + const promise = invocation.execute(mockAbortSignal); + resolveShellExecution({ output: 'hello', exitCode: 0 }); + + const result = await promise; + expect(result.llmContent).not.toContain('Exit Code:'); + }); + + it('should include Exit Code when command fails (non-zero exit code)', async () => { + const invocation = shellTool.build({ command: 'false' }); + const promise = invocation.execute(mockAbortSignal); + resolveShellExecution({ output: '', exitCode: 1 }); + + const result = await promise; + expect(result.llmContent).toContain('Exit Code: 1'); + }); + + it('should not include Error when there is no process error', async () => { + const invocation = shellTool.build({ command: 'echo hello' }); + const promise = invocation.execute(mockAbortSignal); + resolveShellExecution({ output: 'hello', exitCode: 0, error: null }); + + const result = await promise; + expect(result.llmContent).not.toContain('Error:'); + }); + + it('should include Error when there is a process error', async () => { + const invocation = shellTool.build({ command: 'bad-command' }); + const promise = invocation.execute(mockAbortSignal); + resolveShellExecution({ + output: '', + exitCode: 1, + error: new Error('spawn ENOENT'), + }); + + const result = await promise; + expect(result.llmContent).toContain('Error: spawn ENOENT'); + }); + + it('should not include Signal when there is no signal', async () => { + const invocation = shellTool.build({ command: 'echo hello' }); + const promise = invocation.execute(mockAbortSignal); + resolveShellExecution({ output: 'hello', exitCode: 0, signal: null }); + + const result = await promise; + expect(result.llmContent).not.toContain('Signal:'); + }); + + it('should include Signal when process was killed by signal', async () => { + const invocation = shellTool.build({ command: 'sleep 100' }); + const promise = invocation.execute(mockAbortSignal); + resolveShellExecution({ + output: '', + exitCode: null, + signal: 9, // SIGKILL + }); + + const result = await promise; + expect(result.llmContent).toContain('Signal: 9'); + }); + + it('should not include Background PIDs when there are none', async () => { + const invocation = shellTool.build({ command: 'echo hello' }); + const promise = invocation.execute(mockAbortSignal); + resolveShellExecution({ output: 'hello', exitCode: 0 }); + + const result = await promise; + expect(result.llmContent).not.toContain('Background PIDs:'); + }); + + it('should not include Process Group PGID when pid is not set', async () => { + const invocation = shellTool.build({ command: 'echo hello' }); + const promise = invocation.execute(mockAbortSignal); + resolveShellExecution({ output: 'hello', exitCode: 0, pid: undefined }); + + const result = await promise; + expect(result.llmContent).not.toContain('Process Group PGID:'); + }); + + it('should have minimal output for successful command', async () => { + const invocation = shellTool.build({ command: 'echo hello' }); + const promise = invocation.execute(mockAbortSignal); + resolveShellExecution({ output: 'hello', exitCode: 0, pid: undefined }); + + const result = await promise; + // Should only contain Output field + expect(result.llmContent).toBe('Output: hello'); + }); + }); + describe('getConfirmationDetails', () => { it('should annotate sub-commands with redirection correctly', async () => { const shellTool = new ShellTool(mockConfig, createMockMessageBus()); diff --git a/packages/core/src/tools/shell.ts b/packages/core/src/tools/shell.ts index e5e375b9ef..4f68000585 100644 --- a/packages/core/src/tools/shell.ts +++ b/packages/core/src/tools/shell.ts @@ -308,22 +308,31 @@ export class ShellToolInvocation extends BaseToolInvocation< } else { // Create a formatted error string for display, replacing the wrapper command // with the user-facing command. - const finalError = result.error - ? result.error.message.replace(commandToExecute, this.params.command) - : '(none)'; + const llmContentParts = [`Output: ${result.output || '(empty)'}`]; - llmContent = [ - `Command: ${this.params.command}`, - `Directory: ${this.params.dir_path || '(root)'}`, - `Output: ${result.output || '(empty)'}`, - `Error: ${finalError}`, - `Exit Code: ${result.exitCode ?? '(none)'}`, - `Signal: ${result.signal ?? '(none)'}`, - `Background PIDs: ${ - backgroundPIDs.length ? backgroundPIDs.join(', ') : '(none)' - }`, - `Process Group PGID: ${result.pid ?? '(none)'}`, - ].join('\n'); + if (result.error) { + const finalError = result.error.message.replaceAll( + commandToExecute, + this.params.command, + ); + llmContentParts.push(`Error: ${finalError}`); + } + + if (result.exitCode !== null && result.exitCode !== 0) { + llmContentParts.push(`Exit Code: ${result.exitCode}`); + } + + if (result.signal) { + llmContentParts.push(`Signal: ${result.signal}`); + } + if (backgroundPIDs.length) { + llmContentParts.push(`Background PIDs: ${backgroundPIDs.join(', ')}`); + } + if (result.pid) { + llmContentParts.push(`Process Group PGID: ${result.pid}`); + } + + llmContent = llmContentParts.join('\n'); } let returnDisplayMessage = ''; @@ -398,15 +407,12 @@ function getShellToolDescription(): string { The following information is returned: - Command: Executed command. - Directory: Directory where command was executed, or \`(root)\`. - Stdout: Output on stdout stream. Can be \`(empty)\` or partial on error and for any unwaited background processes. - Stderr: Output on stderr stream. Can be \`(empty)\` or partial on error and for any unwaited background processes. - Error: Error or \`(none)\` if no error was reported for the subprocess. - Exit Code: Exit code or \`(none)\` if terminated by signal. - Signal: Signal number or \`(none)\` if no signal was received. - Background PIDs: List of background processes started or \`(none)\`. - Process Group PGID: Process group started or \`(none)\``; + Output: Combined stdout/stderr. Can be \`(empty)\` or partial on error and for any unwaited background processes. + Exit Code: Only included if non-zero (command failed). + Error: Only included if a process-level error occurred (e.g., spawn failure). + Signal: Only included if process was terminated by a signal. + Background PIDs: Only included if background processes were started. + Process Group PGID: Only included if available.`; if (os.platform() === 'win32') { return `This tool executes a given shell command as \`powershell.exe -NoProfile -Command \`. Command can start background processes using PowerShell constructs such as \`Start-Process -NoNewWindow\` or \`Start-Job\`.${returnedInfo}`; From 3e1a377d788ea9ae20cc4cdc7bb6d17a69a32978 Mon Sep 17 00:00:00 2001 From: Jacob Richman Date: Mon, 26 Jan 2026 10:12:40 -0800 Subject: [PATCH 3/9] Fix bug in detecting already added paths. (#17430) --- .../src/ui/commands/directoryCommand.test.tsx | 68 ++++++++++++++----- .../cli/src/ui/commands/directoryCommand.tsx | 24 +++++-- 2 files changed, 69 insertions(+), 23 deletions(-) diff --git a/packages/cli/src/ui/commands/directoryCommand.test.tsx b/packages/cli/src/ui/commands/directoryCommand.test.tsx index 673e9805f9..26e9bc727c 100644 --- a/packages/cli/src/ui/commands/directoryCommand.test.tsx +++ b/packages/cli/src/ui/commands/directoryCommand.test.tsx @@ -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(); + return { + ...actual, + realpathSync: vi.fn((p) => p), + }; +}); + vi.mock('../utils/directoryUtils.js', async (importOriginal) => { const actual = await importOriginal(); @@ -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); diff --git a/packages/cli/src/ui/commands/directoryCommand.tsx b/packages/cli/src/ui/commands/directoryCommand.tsx index be0a35a344..53ec7acb7f 100644 --- a/packages/cli/src/ui/commands/directoryCommand.tsx +++ b/packages/cli/src/ui/commands/directoryCommand.tsx @@ -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) { From d745d86af15971073f319778570c00d6e73139d8 Mon Sep 17 00:00:00 2001 From: Abhi <43648792+abhipatel12@users.noreply.github.com> Date: Mon, 26 Jan 2026 13:38:11 -0500 Subject: [PATCH 4/9] feat(scheduler): support multi-scheduler tool aggregation and nested call IDs (#17429) --- .../hooks/useToolExecutionScheduler.test.ts | 110 ++++++++++++++++++ .../src/ui/hooks/useToolExecutionScheduler.ts | 87 +++++++++++--- packages/core/src/confirmation-bus/types.ts | 1 + .../core/src/scheduler/confirmation.test.ts | 10 +- packages/core/src/scheduler/confirmation.ts | 1 + packages/core/src/scheduler/scheduler.test.ts | 9 ++ packages/core/src/scheduler/scheduler.ts | 21 +++- packages/core/src/scheduler/state-manager.ts | 14 ++- packages/core/src/scheduler/types.ts | 11 ++ 9 files changed, 241 insertions(+), 23 deletions(-) diff --git a/packages/cli/src/ui/hooks/useToolExecutionScheduler.test.ts b/packages/cli/src/ui/hooks/useToolExecutionScheduler.test.ts index 2a526150c3..797109499b 100644 --- a/packages/cli/src/ui/hooks/useToolExecutionScheduler.test.ts +++ b/packages/cli/src/ui/hooks/useToolExecutionScheduler.test.ts @@ -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'); + }); }); diff --git a/packages/cli/src/ui/hooks/useToolExecutionScheduler.ts b/packages/cli/src/ui/hooks/useToolExecutionScheduler.ts index c68e414e9b..0c58e7fc41 100644 --- a/packages/cli/src/ui/hooks/useToolExecutionScheduler.ts +++ b/packages/cli/src/ui/hooks/useToolExecutionScheduler.ts @@ -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([]); + // State stores tool calls organized by their originating schedulerId + const [toolCallsMap, setToolCallsMap] = useState< + Record + >({}); const [lastToolOutputTime, setLastToolOutputTime] = useState(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) => { + 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 = {}; + 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, ]; diff --git a/packages/core/src/confirmation-bus/types.ts b/packages/core/src/confirmation-bus/types.ts index aeecf73b3e..9279485986 100644 --- a/packages/core/src/confirmation-bus/types.ts +++ b/packages/core/src/confirmation-bus/types.ts @@ -26,6 +26,7 @@ export enum MessageBusType { export interface ToolCallsUpdateMessage { type: MessageBusType.TOOL_CALLS_UPDATE; toolCalls: ToolCall[]; + schedulerId: string; } export interface ToolConfirmationRequest { diff --git a/packages/core/src/scheduler/confirmation.test.ts b/packages/core/src/scheduler/confirmation.test.ts index 7162af9d46..9bfdba2184 100644 --- a/packages/core/src/scheduler/confirmation.test.ts +++ b/packages/core/src/scheduler/confirmation.test.ts @@ -29,6 +29,7 @@ import { import type { SchedulerStateManager } from './state-manager.js'; import type { ToolModificationHandler } from './tool-modifier.js'; import type { ValidatingToolCall, WaitingToolCall } from './types.js'; +import { ROOT_SCHEDULER_ID } from './types.js'; import type { Config } from '../config/config.js'; import type { EditorType } from '../utils/editor.js'; import { randomUUID } from 'node:crypto'; @@ -52,7 +53,7 @@ describe('confirmation.ts', () => { }); afterEach(() => { - vi.clearAllMocks(); + vi.restoreAllMocks(); }); const emitResponse = (response: ToolConfirmationResponse) => { @@ -188,6 +189,7 @@ describe('confirmation.ts', () => { state: mockState, modifier: mockModifier, getPreferredEditor, + schedulerId: ROOT_SCHEDULER_ID, }); expect(result.outcome).toBe(ToolConfirmationOutcome.ProceedOnce); @@ -217,6 +219,7 @@ describe('confirmation.ts', () => { state: mockState, modifier: mockModifier, getPreferredEditor, + schedulerId: ROOT_SCHEDULER_ID, }); await listenerPromise; @@ -252,6 +255,7 @@ describe('confirmation.ts', () => { state: mockState, modifier: mockModifier, getPreferredEditor, + schedulerId: ROOT_SCHEDULER_ID, }); await waitForListener(MessageBusType.TOOL_CONFIRMATION_RESPONSE); @@ -293,6 +297,7 @@ describe('confirmation.ts', () => { state: mockState, modifier: mockModifier, getPreferredEditor, + schedulerId: ROOT_SCHEDULER_ID, }); await listenerPromise1; @@ -351,6 +356,7 @@ describe('confirmation.ts', () => { state: mockState, modifier: mockModifier, getPreferredEditor, + schedulerId: ROOT_SCHEDULER_ID, }); await listenerPromise; @@ -397,6 +403,7 @@ describe('confirmation.ts', () => { state: mockState, modifier: mockModifier, getPreferredEditor, + schedulerId: ROOT_SCHEDULER_ID, }); const result = await promise; @@ -420,6 +427,7 @@ describe('confirmation.ts', () => { state: mockState, modifier: mockModifier, getPreferredEditor, + schedulerId: ROOT_SCHEDULER_ID, }), ).rejects.toThrow(/lost during confirmation loop/); }); diff --git a/packages/core/src/scheduler/confirmation.ts b/packages/core/src/scheduler/confirmation.ts index c6aa541508..73958815d0 100644 --- a/packages/core/src/scheduler/confirmation.ts +++ b/packages/core/src/scheduler/confirmation.ts @@ -103,6 +103,7 @@ export async function resolveConfirmation( state: SchedulerStateManager; modifier: ToolModificationHandler; getPreferredEditor: () => EditorType | undefined; + schedulerId: string; }, ): Promise { const { state } = deps; diff --git a/packages/core/src/scheduler/scheduler.test.ts b/packages/core/src/scheduler/scheduler.test.ts index 96340e4d5e..95b6470d1b 100644 --- a/packages/core/src/scheduler/scheduler.test.ts +++ b/packages/core/src/scheduler/scheduler.test.ts @@ -66,6 +66,7 @@ import type { CancelledToolCall, ToolCallResponseInfo, } from './types.js'; +import { ROOT_SCHEDULER_ID } from './types.js'; import { ToolErrorType } from '../tools/tool-error.js'; import * as ToolUtils from '../utils/tool-utils.js'; import type { EditorType } from '../utils/editor.js'; @@ -94,6 +95,8 @@ describe('Scheduler (Orchestrator)', () => { args: { foo: 'bar' }, isClientInitiated: false, prompt_id: 'prompt-1', + schedulerId: ROOT_SCHEDULER_ID, + parentCallId: undefined, }; const req2: ToolCallRequestInfo = { @@ -102,6 +105,8 @@ describe('Scheduler (Orchestrator)', () => { args: { foo: 'baz' }, isClientInitiated: false, prompt_id: 'prompt-1', + schedulerId: ROOT_SCHEDULER_ID, + parentCallId: undefined, }; const mockTool = { @@ -208,6 +213,7 @@ describe('Scheduler (Orchestrator)', () => { config: mockConfig, messageBus: mockMessageBus, getPreferredEditor, + schedulerId: 'root', }); // Reset Tool build behavior @@ -271,6 +277,8 @@ describe('Scheduler (Orchestrator)', () => { request: req1, tool: mockTool, invocation: mockInvocation, + schedulerId: ROOT_SCHEDULER_ID, + startTime: expect.any(Number), }), ]), ); @@ -769,6 +777,7 @@ describe('Scheduler (Orchestrator)', () => { config: mockConfig, messageBus: mockMessageBus, state: mockStateManager, + schedulerId: ROOT_SCHEDULER_ID, }), ); diff --git a/packages/core/src/scheduler/scheduler.ts b/packages/core/src/scheduler/scheduler.ts index b4021faa0b..a8d295b1f9 100644 --- a/packages/core/src/scheduler/scheduler.ts +++ b/packages/core/src/scheduler/scheduler.ts @@ -48,6 +48,8 @@ export interface SchedulerOptions { config: Config; messageBus: MessageBus; getPreferredEditor: () => EditorType | undefined; + schedulerId: string; + parentCallId?: string; } const createErrorResponse = ( @@ -85,6 +87,8 @@ export class Scheduler { private readonly config: Config; private readonly messageBus: MessageBus; private readonly getPreferredEditor: () => EditorType | undefined; + private readonly schedulerId: string; + private readonly parentCallId?: string; private isProcessing = false; private isCancelling = false; @@ -94,7 +98,9 @@ export class Scheduler { this.config = options.config; this.messageBus = options.messageBus; this.getPreferredEditor = options.getPreferredEditor; - this.state = new SchedulerStateManager(this.messageBus); + this.schedulerId = options.schedulerId; + this.parentCallId = options.parentCallId; + this.state = new SchedulerStateManager(this.messageBus, this.schedulerId); this.executor = new ToolExecutor(this.config); this.modifier = new ToolModificationHandler(); @@ -228,16 +234,21 @@ export class Scheduler { try { const toolRegistry = this.config.getToolRegistry(); const newCalls: ToolCall[] = requests.map((request) => { + const enrichedRequest: ToolCallRequestInfo = { + ...request, + schedulerId: this.schedulerId, + parentCallId: this.parentCallId, + }; const tool = toolRegistry.getTool(request.name); if (!tool) { return this._createToolNotFoundErroredToolCall( - request, + enrichedRequest, toolRegistry.getAllToolNames(), ); } - return this._validateAndCreateToolCall(request, tool); + return this._validateAndCreateToolCall(enrichedRequest, tool); }); this.state.enqueue(newCalls); @@ -263,6 +274,7 @@ export class Scheduler { ToolErrorType.TOOL_NOT_REGISTERED, ), durationMs: 0, + schedulerId: this.schedulerId, }; } @@ -278,6 +290,7 @@ export class Scheduler { tool, invocation, startTime: Date.now(), + schedulerId: this.schedulerId, }; } catch (e) { return { @@ -290,6 +303,7 @@ export class Scheduler { ToolErrorType.INVALID_TOOL_PARAMS, ), durationMs: 0, + schedulerId: this.schedulerId, }; } } @@ -411,6 +425,7 @@ export class Scheduler { state: this.state, modifier: this.modifier, getPreferredEditor: this.getPreferredEditor, + schedulerId: this.schedulerId, }); outcome = result.outcome; lastDetails = result.lastDetails; diff --git a/packages/core/src/scheduler/state-manager.ts b/packages/core/src/scheduler/state-manager.ts index dd05556590..519bdb3ee3 100644 --- a/packages/core/src/scheduler/state-manager.ts +++ b/packages/core/src/scheduler/state-manager.ts @@ -17,6 +17,7 @@ import type { ExecutingToolCall, ToolCallResponseInfo, } from './types.js'; +import { ROOT_SCHEDULER_ID } from './types.js'; import type { ToolConfirmationOutcome, ToolResultDisplay, @@ -39,7 +40,10 @@ export class SchedulerStateManager { private readonly queue: ToolCall[] = []; private _completedBatch: CompletedToolCall[] = []; - constructor(private readonly messageBus: MessageBus) {} + constructor( + private readonly messageBus: MessageBus, + private readonly schedulerId: string = ROOT_SCHEDULER_ID, + ) {} addToolCalls(calls: ToolCall[]): void { this.enqueue(calls); @@ -201,6 +205,7 @@ export class SchedulerStateManager { void this.messageBus.publish({ type: MessageBusType.TOOL_CALLS_UPDATE, toolCalls: snapshot, + schedulerId: this.schedulerId, }); } @@ -321,6 +326,7 @@ export class SchedulerStateManager { response, durationMs: startTime ? Date.now() - startTime : undefined, outcome: call.outcome, + schedulerId: call.schedulerId, }; } @@ -336,6 +342,7 @@ export class SchedulerStateManager { response, durationMs: startTime ? Date.now() - startTime : undefined, outcome: call.outcome, + schedulerId: call.schedulerId, }; } @@ -364,6 +371,7 @@ export class SchedulerStateManager { startTime: 'startTime' in call ? call.startTime : undefined, outcome: call.outcome, invocation: call.invocation, + schedulerId: call.schedulerId, }; } @@ -388,6 +396,7 @@ export class SchedulerStateManager { startTime: 'startTime' in call ? call.startTime : undefined, outcome: call.outcome, invocation: call.invocation, + schedulerId: call.schedulerId, }; } @@ -442,6 +451,7 @@ export class SchedulerStateManager { }, durationMs: startTime ? Date.now() - startTime : undefined, outcome: call.outcome, + schedulerId: call.schedulerId, }; } @@ -462,6 +472,7 @@ export class SchedulerStateManager { startTime: 'startTime' in call ? call.startTime : undefined, outcome: call.outcome, invocation: call.invocation, + schedulerId: call.schedulerId, }; } @@ -482,6 +493,7 @@ export class SchedulerStateManager { invocation: call.invocation, liveOutput, pid, + schedulerId: call.schedulerId, }; } } diff --git a/packages/core/src/scheduler/types.ts b/packages/core/src/scheduler/types.ts index 2f2baf77e3..7c0bbe07bd 100644 --- a/packages/core/src/scheduler/types.ts +++ b/packages/core/src/scheduler/types.ts @@ -16,6 +16,8 @@ import type { AnsiOutput } from '../utils/terminalSerializer.js'; import type { ToolErrorType } from '../tools/tool-error.js'; import type { SerializableConfirmationDetails } from '../confirmation-bus/types.js'; +export const ROOT_SCHEDULER_ID = 'root'; + export interface ToolCallRequestInfo { callId: string; name: string; @@ -24,6 +26,8 @@ export interface ToolCallRequestInfo { prompt_id: string; checkpoint?: string; traceId?: string; + parentCallId?: string; + schedulerId?: string; } export interface ToolCallResponseInfo { @@ -43,6 +47,7 @@ export type ValidatingToolCall = { invocation: AnyToolInvocation; startTime?: number; outcome?: ToolConfirmationOutcome; + schedulerId?: string; }; export type ScheduledToolCall = { @@ -52,6 +57,7 @@ export type ScheduledToolCall = { invocation: AnyToolInvocation; startTime?: number; outcome?: ToolConfirmationOutcome; + schedulerId?: string; }; export type ErroredToolCall = { @@ -61,6 +67,7 @@ export type ErroredToolCall = { tool?: AnyDeclarativeTool; durationMs?: number; outcome?: ToolConfirmationOutcome; + schedulerId?: string; }; export type SuccessfulToolCall = { @@ -71,6 +78,7 @@ export type SuccessfulToolCall = { invocation: AnyToolInvocation; durationMs?: number; outcome?: ToolConfirmationOutcome; + schedulerId?: string; }; export type ExecutingToolCall = { @@ -82,6 +90,7 @@ export type ExecutingToolCall = { startTime?: number; outcome?: ToolConfirmationOutcome; pid?: number; + schedulerId?: string; }; export type CancelledToolCall = { @@ -92,6 +101,7 @@ export type CancelledToolCall = { invocation: AnyToolInvocation; durationMs?: number; outcome?: ToolConfirmationOutcome; + schedulerId?: string; }; export type WaitingToolCall = { @@ -113,6 +123,7 @@ export type WaitingToolCall = { correlationId?: string; startTime?: number; outcome?: ToolConfirmationOutcome; + schedulerId?: string; }; export type Status = ToolCall['status']; From 2271bbb339f88e2e014a53ee3130ab8bb14fd269 Mon Sep 17 00:00:00 2001 From: Christian Gunderman Date: Mon, 26 Jan 2026 19:49:32 +0000 Subject: [PATCH 5/9] feat(agents): implement first-run experience for project-level sub-agents (#17266) --- packages/cli/src/test-utils/render.tsx | 1 + packages/cli/src/ui/AppContainer.tsx | 39 +++- .../cli/src/ui/components/DialogManager.tsx | 9 + .../components/NewAgentsNotification.test.tsx | 57 ++++++ .../ui/components/NewAgentsNotification.tsx | 96 ++++++++++ .../NewAgentsNotification.test.tsx.snap | 43 +++++ .../cli/src/ui/contexts/UIActionsContext.tsx | 2 + .../cli/src/ui/contexts/UIStateContext.tsx | 1 + .../src/agents/acknowledgedAgents.test.ts | 97 ++++++++++ .../core/src/agents/acknowledgedAgents.ts | 85 +++++++++ packages/core/src/agents/agentLoader.ts | 36 ++-- packages/core/src/agents/registry.test.ts | 53 ++++++ packages/core/src/agents/registry.ts | 59 +++++- .../agents/registry_acknowledgement.test.ts | 169 ++++++++++++++++++ packages/core/src/agents/types.ts | 4 + packages/core/src/config/config.ts | 7 + packages/core/src/config/storage.ts | 8 + packages/core/src/utils/events.ts | 18 ++ 18 files changed, 769 insertions(+), 15 deletions(-) create mode 100644 packages/cli/src/ui/components/NewAgentsNotification.test.tsx create mode 100644 packages/cli/src/ui/components/NewAgentsNotification.tsx create mode 100644 packages/cli/src/ui/components/__snapshots__/NewAgentsNotification.test.tsx.snap create mode 100644 packages/core/src/agents/acknowledgedAgents.test.ts create mode 100644 packages/core/src/agents/acknowledgedAgents.ts create mode 100644 packages/core/src/agents/registry_acknowledgement.test.ts diff --git a/packages/cli/src/test-utils/render.tsx b/packages/cli/src/test-utils/render.tsx index a7f90aecfe..717aa668d1 100644 --- a/packages/cli/src/test-utils/render.tsx +++ b/packages/cli/src/test-utils/render.tsx @@ -198,6 +198,7 @@ const mockUIActions: UIActions = { setEmbeddedShellFocused: vi.fn(), setAuthContext: vi.fn(), handleRestart: vi.fn(), + handleNewAgentsSelect: vi.fn(), }; export const renderWithProviders = ( diff --git a/packages/cli/src/ui/AppContainer.tsx b/packages/cli/src/ui/AppContainer.tsx index 4f10e10645..45ccd33ad0 100644 --- a/packages/cli/src/ui/AppContainer.tsx +++ b/packages/cli/src/ui/AppContainer.tsx @@ -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(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, ], ); diff --git a/packages/cli/src/ui/components/DialogManager.tsx b/packages/cli/src/ui/components/DialogManager.tsx index 5d66927487..305f2333f1 100644 --- a/packages/cli/src/ui/components/DialogManager.tsx +++ b/packages/cli/src/ui/components/DialogManager.tsx @@ -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 ; } + if (uiState.newAgents) { + return ( + + ); + } if (uiState.proQuotaRequest) { return ( { + 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( + , + ); + + 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( + , + ); + + const frame = lastFrame(); + expect(frame).toMatchSnapshot(); + unmount(); + }); +}); diff --git a/packages/cli/src/ui/components/NewAgentsNotification.tsx b/packages/cli/src/ui/components/NewAgentsNotification.tsx new file mode 100644 index 0000000000..05edae484c --- /dev/null +++ b/packages/cli/src/ui/components/NewAgentsNotification.tsx @@ -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> = [ + { + 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 ( + + + + + New Agents Discovered + + + The following agents were found in this project. Please review them: + + + {displayAgents.map((agent) => ( + + + + - {agent.name}:{' '} + + + {agent.description} + + ))} + {remaining > 0 && ( + + ... and {remaining} more. + + )} + + + + + + + ); +}; diff --git a/packages/cli/src/ui/components/__snapshots__/NewAgentsNotification.test.tsx.snap b/packages/cli/src/ui/components/__snapshots__/NewAgentsNotification.test.tsx.snap new file mode 100644 index 0000000000..438d51e1e3 --- /dev/null +++ b/packages/cli/src/ui/components/__snapshots__/NewAgentsNotification.test.tsx.snap @@ -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) │ + │ │ + ╰────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯" +`; diff --git a/packages/cli/src/ui/contexts/UIActionsContext.tsx b/packages/cli/src/ui/contexts/UIActionsContext.tsx index c8abf33236..4eb8584ae3 100644 --- a/packages/cli/src/ui/contexts/UIActionsContext.tsx +++ b/packages/cli/src/ui/contexts/UIActionsContext.tsx @@ -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; } export const UIActionsContext = createContext(null); diff --git a/packages/cli/src/ui/contexts/UIStateContext.tsx b/packages/cli/src/ui/contexts/UIStateContext.tsx index fea13285b1..6d10d76bda 100644 --- a/packages/cli/src/ui/contexts/UIStateContext.tsx +++ b/packages/cli/src/ui/contexts/UIStateContext.tsx @@ -155,6 +155,7 @@ export interface UIState { terminalBackgroundColor: TerminalBackgroundColor; settingsNonce: number; adminSettingsChanged: boolean; + newAgents: AgentDefinition[] | null; } export const UIStateContext = createContext(null); diff --git a/packages/core/src/agents/acknowledgedAgents.test.ts b/packages/core/src/agents/acknowledgedAgents.test.ts new file mode 100644 index 0000000000..f6e45360db --- /dev/null +++ b/packages/core/src/agents/acknowledgedAgents.test.ts @@ -0,0 +1,97 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, beforeEach, afterEach } from 'vitest'; +import { AcknowledgedAgentsService } from './acknowledgedAgents.js'; +import { Storage } from '../config/storage.js'; +import * as fs from 'node:fs/promises'; +import * as path from 'node:path'; +import * as os from 'node:os'; + +describe('AcknowledgedAgentsService', () => { + let tempDir: string; + let originalGeminiCliHome: string | undefined; + + beforeEach(async () => { + // Create a unique temp directory for each test + tempDir = await fs.mkdtemp(path.join(os.tmpdir(), 'gemini-cli-test-')); + + // Override GEMINI_CLI_HOME to point to the temp directory + originalGeminiCliHome = process.env['GEMINI_CLI_HOME']; + process.env['GEMINI_CLI_HOME'] = tempDir; + }); + + afterEach(async () => { + // Restore environment variable + if (originalGeminiCliHome) { + process.env['GEMINI_CLI_HOME'] = originalGeminiCliHome; + } else { + delete process.env['GEMINI_CLI_HOME']; + } + + // Clean up temp directory + await fs.rm(tempDir, { recursive: true, force: true }); + }); + + it('should acknowledge an agent and save to disk', async () => { + const service = new AcknowledgedAgentsService(); + const ackPath = Storage.getAcknowledgedAgentsPath(); + + await service.acknowledge('/project', 'AgentA', 'hash1'); + + // Verify file exists and content + const content = await fs.readFile(ackPath, 'utf-8'); + expect(content).toContain('"AgentA": "hash1"'); + }); + + it('should return true for acknowledged agent', async () => { + const service = new AcknowledgedAgentsService(); + + await service.acknowledge('/project', 'AgentA', 'hash1'); + + expect(await service.isAcknowledged('/project', 'AgentA', 'hash1')).toBe( + true, + ); + expect(await service.isAcknowledged('/project', 'AgentA', 'hash2')).toBe( + false, + ); + expect(await service.isAcknowledged('/project', 'AgentB', 'hash1')).toBe( + false, + ); + }); + + it('should load acknowledged agents from disk', async () => { + const ackPath = Storage.getAcknowledgedAgentsPath(); + const data = { + '/project': { + AgentLoaded: 'hashLoaded', + }, + }; + + // Ensure directory exists + await fs.mkdir(path.dirname(ackPath), { recursive: true }); + await fs.writeFile(ackPath, JSON.stringify(data), 'utf-8'); + + const service = new AcknowledgedAgentsService(); + + expect( + await service.isAcknowledged('/project', 'AgentLoaded', 'hashLoaded'), + ).toBe(true); + }); + + it('should handle load errors gracefully', async () => { + // Create a directory where the file should be to cause a read error (EISDIR) + const ackPath = Storage.getAcknowledgedAgentsPath(); + await fs.mkdir(ackPath, { recursive: true }); + + const service = new AcknowledgedAgentsService(); + + // Should not throw, and treated as empty + expect(await service.isAcknowledged('/project', 'Agent', 'hash')).toBe( + false, + ); + }); +}); diff --git a/packages/core/src/agents/acknowledgedAgents.ts b/packages/core/src/agents/acknowledgedAgents.ts new file mode 100644 index 0000000000..230b62443a --- /dev/null +++ b/packages/core/src/agents/acknowledgedAgents.ts @@ -0,0 +1,85 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import * as fs from 'node:fs/promises'; +import * as path from 'node:path'; +import { Storage } from '../config/storage.js'; +import { debugLogger } from '../utils/debugLogger.js'; +import { getErrorMessage, isNodeError } from '../utils/errors.js'; + +export interface AcknowledgedAgentsMap { + // Project Path -> Agent Name -> Agent Hash + [projectPath: string]: { + [agentName: string]: string; + }; +} + +export class AcknowledgedAgentsService { + private acknowledgedAgents: AcknowledgedAgentsMap = {}; + private loaded = false; + + async load(): Promise { + if (this.loaded) return; + + const filePath = Storage.getAcknowledgedAgentsPath(); + try { + const content = await fs.readFile(filePath, 'utf-8'); + this.acknowledgedAgents = JSON.parse(content); + } catch (error: unknown) { + if (!isNodeError(error) || error.code !== 'ENOENT') { + debugLogger.error( + 'Failed to load acknowledged agents:', + getErrorMessage(error), + ); + } + // If file doesn't exist or there's a parsing error, fallback to empty + this.acknowledgedAgents = {}; + } + this.loaded = true; + } + + async save(): Promise { + const filePath = Storage.getAcknowledgedAgentsPath(); + try { + const dir = path.dirname(filePath); + await fs.mkdir(dir, { recursive: true }); + await fs.writeFile( + filePath, + JSON.stringify(this.acknowledgedAgents, null, 2), + 'utf-8', + ); + } catch (error) { + debugLogger.error( + 'Failed to save acknowledged agents:', + getErrorMessage(error), + ); + } + } + + async isAcknowledged( + projectPath: string, + agentName: string, + hash: string, + ): Promise { + await this.load(); + const projectAgents = this.acknowledgedAgents[projectPath]; + if (!projectAgents) return false; + return projectAgents[agentName] === hash; + } + + async acknowledge( + projectPath: string, + agentName: string, + hash: string, + ): Promise { + await this.load(); + if (!this.acknowledgedAgents[projectPath]) { + this.acknowledgedAgents[projectPath] = {}; + } + this.acknowledgedAgents[projectPath][agentName] = hash; + await this.save(); + } +} diff --git a/packages/core/src/agents/agentLoader.ts b/packages/core/src/agents/agentLoader.ts index 385d1e9b59..1679b52fb3 100644 --- a/packages/core/src/agents/agentLoader.ts +++ b/packages/core/src/agents/agentLoader.ts @@ -8,10 +8,12 @@ import yaml from 'js-yaml'; import * as fs from 'node:fs/promises'; import { type Dirent } from 'node:fs'; import * as path from 'node:path'; +import * as crypto from 'node:crypto'; import { z } from 'zod'; import type { AgentDefinition } from './types.js'; import { isValidToolName } from '../tools/tool-names.js'; import { FRONTMATTER_REGEX } from '../skills/skillLoader.js'; +import { getErrorMessage } from '../utils/errors.js'; /** * DTO for Markdown parsing - represents the structure from frontmatter. @@ -139,24 +141,30 @@ function formatZodError(error: z.ZodError, context: string): string { * Parses and validates an agent Markdown file with frontmatter. * * @param filePath Path to the Markdown file. + * @param content Optional pre-loaded content of the file. * @returns An array containing the single parsed agent definition. * @throws AgentLoadError if parsing or validation fails. */ export async function parseAgentMarkdown( filePath: string, + content?: string, ): Promise { - let content: string; - try { - content = await fs.readFile(filePath, 'utf-8'); - } catch (error) { - throw new AgentLoadError( - filePath, - `Could not read file: ${(error as Error).message}`, - ); + let fileContent: string; + if (content !== undefined) { + fileContent = content; + } else { + try { + fileContent = await fs.readFile(filePath, 'utf-8'); + } catch (error) { + throw new AgentLoadError( + filePath, + `Could not read file: ${getErrorMessage(error)}`, + ); + } } // Split frontmatter and body - const match = content.match(FRONTMATTER_REGEX); + const match = fileContent.match(FRONTMATTER_REGEX); if (!match) { throw new AgentLoadError( filePath, @@ -229,10 +237,12 @@ export async function parseAgentMarkdown( * Converts a FrontmatterAgentDefinition DTO to the internal AgentDefinition structure. * * @param markdown The parsed Markdown/Frontmatter definition. + * @param metadata Optional metadata including hash and file path. * @returns The internal AgentDefinition. */ export function markdownToAgentDefinition( markdown: FrontmatterAgentDefinition, + metadata?: { hash?: string; filePath?: string }, ): AgentDefinition { const inputConfig = { inputSchema: { @@ -256,6 +266,7 @@ export function markdownToAgentDefinition( displayName: markdown.display_name, agentCardUrl: markdown.agent_card_url, inputConfig, + metadata, }; } @@ -288,6 +299,7 @@ export function markdownToAgentDefinition( } : undefined, inputConfig, + metadata, }; } @@ -334,9 +346,11 @@ export async function loadAgentsFromDirectory( for (const entry of files) { const filePath = path.join(dir, entry.name); try { - const agentDefs = await parseAgentMarkdown(filePath); + const content = await fs.readFile(filePath, 'utf-8'); + const hash = crypto.createHash('sha256').update(content).digest('hex'); + const agentDefs = await parseAgentMarkdown(filePath, content); for (const def of agentDefs) { - const agent = markdownToAgentDefinition(def); + const agent = markdownToAgentDefinition(def, { hash, filePath }); result.agents.push(agent); } } catch (error) { diff --git a/packages/core/src/agents/registry.test.ts b/packages/core/src/agents/registry.test.ts index e55f4214aa..9eb43f357b 100644 --- a/packages/core/src/agents/registry.test.ts +++ b/packages/core/src/agents/registry.test.ts @@ -25,6 +25,7 @@ import { SimpleExtensionLoader } from '../utils/extensionLoader.js'; import type { ConfigParameters } from '../config/config.js'; import type { ToolRegistry } from '../tools/tool-registry.js'; import { ThinkingLevel } from '@google/genai'; +import type { AcknowledgedAgentsService } from './acknowledgedAgents.js'; vi.mock('./agentLoader.js', () => ({ loadAgentsFromDirectory: vi @@ -401,6 +402,58 @@ describe('AgentRegistry', () => { expect(registry.getDefinition('extension-agent')).toBeUndefined(); }); + + it('should use agentCardUrl as hash for acknowledgement of remote agents', async () => { + mockConfig = makeMockedConfig({ enableAgents: true }); + // Trust the folder so it attempts to load project agents + vi.spyOn(mockConfig, 'isTrustedFolder').mockReturnValue(true); + vi.spyOn(mockConfig, 'getFolderTrust').mockReturnValue(true); + + const registry = new TestableAgentRegistry(mockConfig); + + const remoteAgent: AgentDefinition = { + kind: 'remote', + name: 'RemoteAgent', + description: 'A remote agent', + agentCardUrl: 'https://example.com/card', + inputConfig: { inputSchema: { type: 'object' } }, + metadata: { hash: 'file-hash', filePath: 'path/to/file.md' }, + }; + + vi.mocked(tomlLoader.loadAgentsFromDirectory).mockResolvedValue({ + agents: [remoteAgent], + errors: [], + }); + + const ackService = { + isAcknowledged: vi.fn().mockResolvedValue(true), + acknowledge: vi.fn(), + }; + vi.spyOn(mockConfig, 'getAcknowledgedAgentsService').mockReturnValue( + ackService as unknown as AcknowledgedAgentsService, + ); + + // Mock A2AClientManager to avoid network calls + vi.mocked(A2AClientManager.getInstance).mockReturnValue({ + loadAgent: vi.fn().mockResolvedValue({ name: 'RemoteAgent' }), + clearCache: vi.fn(), + } as unknown as A2AClientManager); + + await registry.initialize(); + + // Verify ackService was called with the URL, not the file hash + expect(ackService.isAcknowledged).toHaveBeenCalledWith( + expect.anything(), + 'RemoteAgent', + 'https://example.com/card', + ); + + // Also verify that the agent's metadata was updated to use the URL as hash + // Use getDefinition because registerAgent might have been called + expect(registry.getDefinition('RemoteAgent')?.metadata?.hash).toBe( + 'https://example.com/card', + ); + }); }); describe('registration logic', () => { diff --git a/packages/core/src/agents/registry.ts b/packages/core/src/agents/registry.ts index cc68156344..cc91ffeeed 100644 --- a/packages/core/src/agents/registry.ts +++ b/packages/core/src/agents/registry.ts @@ -5,7 +5,7 @@ */ import { Storage } from '../config/storage.js'; -import { coreEvents, CoreEvent } from '../utils/events.js'; +import { CoreEvent, coreEvents } from '../utils/events.js'; import type { AgentOverride, Config } from '../config/config.js'; import type { AgentDefinition, LocalAgentDefinition } from './types.js'; import { loadAgentsFromDirectory } from './agentLoader.js'; @@ -73,6 +73,23 @@ export class AgentRegistry { coreEvents.emitAgentsRefreshed(); } + /** + * Acknowledges and registers a previously unacknowledged agent. + */ + async acknowledgeAgent(agent: AgentDefinition): Promise { + const ackService = this.config.getAcknowledgedAgentsService(); + const projectRoot = this.config.getProjectRoot(); + if (agent.metadata?.hash) { + await ackService.acknowledge( + projectRoot, + agent.name, + agent.metadata.hash, + ); + await this.registerAgent(agent); + coreEvents.emitAgentsRefreshed(); + } + } + /** * Disposes of resources and removes event listeners. */ @@ -115,8 +132,46 @@ export class AgentRegistry { `Agent loading error: ${error.message}`, ); } + + const ackService = this.config.getAcknowledgedAgentsService(); + const projectRoot = this.config.getProjectRoot(); + const unacknowledgedAgents: AgentDefinition[] = []; + const agentsToRegister: AgentDefinition[] = []; + + for (const agent of projectAgents.agents) { + // If it's a remote agent, use the agentCardUrl as the hash. + // This allows multiple remote agents in a single file to be tracked independently. + if (agent.kind === 'remote') { + if (!agent.metadata) { + agent.metadata = {}; + } + agent.metadata.hash = agent.agentCardUrl; + } + + if (!agent.metadata?.hash) { + agentsToRegister.push(agent); + continue; + } + + const isAcknowledged = await ackService.isAcknowledged( + projectRoot, + agent.name, + agent.metadata.hash, + ); + + if (isAcknowledged) { + agentsToRegister.push(agent); + } else { + unacknowledgedAgents.push(agent); + } + } + + if (unacknowledgedAgents.length > 0) { + coreEvents.emitAgentsDiscovered(unacknowledgedAgents); + } + await Promise.allSettled( - projectAgents.agents.map((agent) => this.registerAgent(agent)), + agentsToRegister.map((agent) => this.registerAgent(agent)), ); } else { coreEvents.emitFeedback( diff --git a/packages/core/src/agents/registry_acknowledgement.test.ts b/packages/core/src/agents/registry_acknowledgement.test.ts new file mode 100644 index 0000000000..5ac563091d --- /dev/null +++ b/packages/core/src/agents/registry_acknowledgement.test.ts @@ -0,0 +1,169 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; +import { AgentRegistry } from './registry.js'; +import { makeFakeConfig } from '../test-utils/config.js'; +import type { AgentDefinition } from './types.js'; +import { coreEvents } from '../utils/events.js'; +import * as tomlLoader from './agentLoader.js'; +import { type Config } from '../config/config.js'; +import { AcknowledgedAgentsService } from './acknowledgedAgents.js'; +import * as fs from 'node:fs/promises'; +import * as path from 'node:path'; +import * as os from 'node:os'; + +// Mock dependencies +vi.mock('./agentLoader.js', () => ({ + loadAgentsFromDirectory: vi.fn(), +})); + +const MOCK_AGENT_WITH_HASH: AgentDefinition = { + kind: 'local', + name: 'ProjectAgent', + description: 'Project Agent Desc', + inputConfig: { inputSchema: { type: 'object' } }, + modelConfig: { + model: 'test', + generateContentConfig: { thinkingConfig: { includeThoughts: true } }, + }, + runConfig: { maxTimeMinutes: 1 }, + promptConfig: { systemPrompt: 'test' }, + metadata: { + hash: 'hash123', + filePath: '/project/agent.md', + }, +}; + +describe('AgentRegistry Acknowledgement', () => { + let registry: AgentRegistry; + let config: Config; + let tempDir: string; + let originalGeminiCliHome: string | undefined; + let ackService: AcknowledgedAgentsService; + + beforeEach(async () => { + // Create a unique temp directory for each test + tempDir = await fs.mkdtemp(path.join(os.tmpdir(), 'gemini-cli-test-')); + + // Override GEMINI_CLI_HOME to point to the temp directory + originalGeminiCliHome = process.env['GEMINI_CLI_HOME']; + process.env['GEMINI_CLI_HOME'] = tempDir; + + ackService = new AcknowledgedAgentsService(); + + config = makeFakeConfig({ + folderTrust: true, + trustedFolder: true, + }); + // Ensure we are in trusted folder mode for project agents to load + vi.spyOn(config, 'isTrustedFolder').mockReturnValue(true); + vi.spyOn(config, 'getFolderTrust').mockReturnValue(true); + vi.spyOn(config, 'getProjectRoot').mockReturnValue('/project'); + vi.spyOn(config, 'getAcknowledgedAgentsService').mockReturnValue( + ackService, + ); + + // We cannot easily spy on storage.getProjectAgentsDir if it's a property/getter unless we cast to any or it's a method + // Assuming it's a method on Storage class + vi.spyOn(config.storage, 'getProjectAgentsDir').mockReturnValue( + '/project/.gemini/agents', + ); + vi.spyOn(config, 'isAgentsEnabled').mockReturnValue(true); + + registry = new AgentRegistry(config); + + vi.mocked(tomlLoader.loadAgentsFromDirectory).mockImplementation( + async (dir) => { + if (dir === '/project/.gemini/agents') { + return { + agents: [MOCK_AGENT_WITH_HASH], + errors: [], + }; + } + return { agents: [], errors: [] }; + }, + ); + }); + + afterEach(async () => { + vi.restoreAllMocks(); + + // Restore environment variable + if (originalGeminiCliHome) { + process.env['GEMINI_CLI_HOME'] = originalGeminiCliHome; + } else { + delete process.env['GEMINI_CLI_HOME']; + } + + // Clean up temp directory + await fs.rm(tempDir, { recursive: true, force: true }); + }); + + it('should not register unacknowledged project agents and emit event', async () => { + const emitSpy = vi.spyOn(coreEvents, 'emitAgentsDiscovered'); + + await registry.initialize(); + + expect(registry.getDefinition('ProjectAgent')).toBeUndefined(); + expect(emitSpy).toHaveBeenCalledWith([MOCK_AGENT_WITH_HASH]); + }); + + it('should register acknowledged project agents', async () => { + // Acknowledge the agent explicitly + await ackService.acknowledge('/project', 'ProjectAgent', 'hash123'); + + vi.mocked(tomlLoader.loadAgentsFromDirectory).mockImplementation( + async (dir) => { + if (dir === '/project/.gemini/agents') { + return { + agents: [MOCK_AGENT_WITH_HASH], + errors: [], + }; + } + return { agents: [], errors: [] }; + }, + ); + + const emitSpy = vi.spyOn(coreEvents, 'emitAgentsDiscovered'); + + await registry.initialize(); + + expect(registry.getDefinition('ProjectAgent')).toBeDefined(); + expect(emitSpy).not.toHaveBeenCalled(); + }); + + it('should register agents without hash (legacy/safe?)', async () => { + // Current logic: if no hash, allow it. + const agentNoHash = { ...MOCK_AGENT_WITH_HASH, metadata: undefined }; + vi.mocked(tomlLoader.loadAgentsFromDirectory).mockImplementation( + async (dir) => { + if (dir === '/project/.gemini/agents') { + return { + agents: [agentNoHash], + errors: [], + }; + } + return { agents: [], errors: [] }; + }, + ); + + await registry.initialize(); + + expect(registry.getDefinition('ProjectAgent')).toBeDefined(); + }); + + it('acknowledgeAgent should acknowledge and register agent', async () => { + await registry.acknowledgeAgent(MOCK_AGENT_WITH_HASH); + + // Verify against real service state + expect( + await ackService.isAcknowledged('/project', 'ProjectAgent', 'hash123'), + ).toBe(true); + + expect(registry.getDefinition('ProjectAgent')).toBeDefined(); + }); +}); diff --git a/packages/core/src/agents/types.ts b/packages/core/src/agents/types.ts index f58b6fa0ae..581e9f2b52 100644 --- a/packages/core/src/agents/types.ts +++ b/packages/core/src/agents/types.ts @@ -74,6 +74,10 @@ export interface BaseAgentDefinition< experimental?: boolean; inputConfig: InputConfig; outputConfig?: OutputConfig; + metadata?: { + hash?: string; + filePath?: string; + }; } export interface LocalAgentDefinition< diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index c0b96a292f..2dd235becf 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -100,6 +100,7 @@ import type { FetchAdminControlsResponse } from '../code_assist/types.js'; import { getCodeAssistServer } from '../code_assist/codeAssist.js'; import type { Experiments } from '../code_assist/experiments/experiments.js'; import { AgentRegistry } from '../agents/registry.js'; +import { AcknowledgedAgentsService } from '../agents/acknowledgedAgents.js'; import { setGlobalProxy } from '../utils/fetch.js'; import { SubagentTool } from '../agents/subagent-tool.js'; import { getExperiments } from '../code_assist/experiments/experiments.js'; @@ -416,6 +417,7 @@ export class Config { private promptRegistry!: PromptRegistry; private resourceRegistry!: ResourceRegistry; private agentRegistry!: AgentRegistry; + private readonly acknowledgedAgentsService: AcknowledgedAgentsService; private skillManager!: SkillManager; private sessionId: string; private clientVersion: string; @@ -705,6 +707,7 @@ export class Config { params.approvalMode ?? params.policyEngineConfig?.approvalMode, }); this.messageBus = new MessageBus(this.policyEngine, this.debugMode); + this.acknowledgedAgentsService = new AcknowledgedAgentsService(); this.skillManager = new SkillManager(); this.outputSettings = { format: params.output?.format ?? OutputFormat.TEXT, @@ -1138,6 +1141,10 @@ export class Config { return this.agentRegistry; } + getAcknowledgedAgentsService(): AcknowledgedAgentsService { + return this.acknowledgedAgentsService; + } + getToolRegistry(): ToolRegistry { return this.toolRegistry; } diff --git a/packages/core/src/config/storage.ts b/packages/core/src/config/storage.ts index da7142d09c..ac7efb8103 100644 --- a/packages/core/src/config/storage.ts +++ b/packages/core/src/config/storage.ts @@ -66,6 +66,14 @@ export class Storage { return path.join(Storage.getGlobalGeminiDir(), 'agents'); } + static getAcknowledgedAgentsPath(): string { + return path.join( + Storage.getGlobalGeminiDir(), + 'acknowledgments', + 'agents.json', + ); + } + static getSystemSettingsPath(): string { if (process.env['GEMINI_CLI_SYSTEM_SETTINGS_PATH']) { return process.env['GEMINI_CLI_SYSTEM_SETTINGS_PATH']; diff --git a/packages/core/src/utils/events.ts b/packages/core/src/utils/events.ts index 8fd6a73751..d5f8f715aa 100644 --- a/packages/core/src/utils/events.ts +++ b/packages/core/src/utils/events.ts @@ -5,6 +5,7 @@ */ import { EventEmitter } from 'node:events'; +import type { AgentDefinition } from '../agents/types.js'; import type { McpClient } from '../tools/mcp-client.js'; import type { ExtensionEvents } from './extensionLoader.js'; @@ -110,6 +111,13 @@ export interface RetryAttemptPayload { model: string; } +/** + * Payload for the 'agents-discovered' event. + */ +export interface AgentsDiscoveredPayload { + agents: AgentDefinition[]; +} + export enum CoreEvent { UserFeedback = 'user-feedback', ModelChanged = 'model-changed', @@ -125,6 +133,7 @@ export enum CoreEvent { AgentsRefreshed = 'agents-refreshed', AdminSettingsChanged = 'admin-settings-changed', RetryAttempt = 'retry-attempt', + AgentsDiscovered = 'agents-discovered', } export interface CoreEvents extends ExtensionEvents { @@ -142,6 +151,7 @@ export interface CoreEvents extends ExtensionEvents { [CoreEvent.AgentsRefreshed]: never[]; [CoreEvent.AdminSettingsChanged]: never[]; [CoreEvent.RetryAttempt]: [RetryAttemptPayload]; + [CoreEvent.AgentsDiscovered]: [AgentsDiscoveredPayload]; } type EventBacklogItem = { @@ -264,6 +274,14 @@ export class CoreEventEmitter extends EventEmitter { this.emit(CoreEvent.RetryAttempt, payload); } + /** + * Notifies subscribers that new unacknowledged agents have been discovered. + */ + emitAgentsDiscovered(agents: AgentDefinition[]): void { + const payload: AgentsDiscoveredPayload = { agents }; + this._emitOrQueue(CoreEvent.AgentsDiscovered, payload); + } + /** * Flushes buffered messages. Call this immediately after primary UI listener * subscribes. From b07953fb3820fe86f09003fb9d3814f787d8f036 Mon Sep 17 00:00:00 2001 From: christine betts Date: Mon, 26 Jan 2026 15:14:38 -0500 Subject: [PATCH 6/9] Update extensions docs (#16093) --- docs/extensions/best-practices.md | 139 +++++++ docs/extensions/index.md | 383 ++---------------- docs/extensions/reference.md | 312 ++++++++++++++ .../{extension-releasing.md => releasing.md} | 0 ...ed-extensions.md => writing-extensions.md} | 55 +-- docs/index.md | 8 +- docs/sidebar.json | 16 +- 7 files changed, 520 insertions(+), 393 deletions(-) create mode 100644 docs/extensions/best-practices.md create mode 100644 docs/extensions/reference.md rename docs/extensions/{extension-releasing.md => releasing.md} (100%) rename docs/extensions/{getting-started-extensions.md => writing-extensions.md} (63%) diff --git a/docs/extensions/best-practices.md b/docs/extensions/best-practices.md new file mode 100644 index 0000000000..73c578f1be --- /dev/null +++ b/docs/extensions/best-practices.md @@ -0,0 +1,139 @@ +# Extensions on Gemini CLI: Best practices + +This guide covers best practices for developing, securing, and maintaining +Gemini CLI extensions. + +## Development + +Developing extensions for Gemini CLI is intended to be a lightweight, iterative +process. + +### Structure your extension + +While simple extensions can just be a few files, we recommend a robust structure +for complex extensions: + +``` +my-extension/ +├── package.json +├── tsconfig.json +├── gemini-extension.json +├── src/ +│ ├── index.ts +│ └── tools/ +└── dist/ +``` + +- **Use TypeScript**: We strongly recommend using TypeScript for type safety and + better tooling. +- **Separate source and build**: Keep your source code in `src` and build to + `dist`. +- **Bundle dependencies**: If your extension has many dependencies, consider + bundling them (e.g., with `esbuild` or `webpack`) to reduce install time and + potential conflicts. + +### Iterate with `link` + +Use `gemini extensions link` to develop locally without constantly reinstalling: + +```bash +cd my-extension +gemini extensions link . +``` + +Changes to your code (after rebuilding) will be immediately available in the CLI +on restart. + +### Use `GEMINI.md` effectively + +Your `GEMINI.md` file provides context to the model. Keep it focused: + +- **Do:** Explain high-level goals and how to use the provided tools. +- **Don't:** Dump your entire documentation. +- **Do:** Use clear, concise language. + +## Security + +When building a Gemini CLI extension, follow general security best practices +(such as least privilege and input validation) to reduce risk. + +### Minimal permissions + +When defining tools in your MCP server, only request the permissions necessary. +Avoid giving the model broad access (like full shell access) if a more +restricted set of tools will suffice. + +If you must use powerful tools like `run_shell_command`, consider restricting +them to specific commands in your `gemini-extension.json`: + +```json +{ + "name": "my-safe-extension", + "excludeTools": ["run_shell_command(rm -rf *)"] +} +``` + +This ensures that even if the model tries to execute a dangerous command, it +will be blocked at the CLI level. + +### Validate inputs + +Your MCP server is running on the user's machine. Always validate inputs to your +tools to prevent arbitrary code execution or filesystem access outside the +intended scope. + +```typescript +// Good: Validating paths +if (!path.resolve(inputPath).startsWith(path.resolve(allowedDir) + path.sep)) { + throw new Error('Access denied'); +} +``` + +### Sensitive settings + +If your extension requires API keys, use the `sensitive: true` option in +`gemini-extension.json`. This ensures keys are stored securely in the system +keychain and obfuscated in the UI. + +```json +"settings": [ + { + "name": "API Key", + "envVar": "MY_API_KEY", + "sensitive": true + } +] +``` + +## Releasing + +You can upload your extension directly to GitHub to list it in the gallery. +Gemini CLI extensions also offers support for more complicated +[releases](releasing.md). + +### Semantic versioning + +Follow [Semantic Versioning](https://semver.org/). + +- **Major**: Breaking changes (renaming tools, changing arguments). +- **Minor**: New features (new tools, commands). +- **Patch**: Bug fixes. + +### Release Channels + +Use git branches to manage release channels (e.g., `main` for stable, `dev` for +bleeding edge). This allows users to choose their stability level: + +```bash +# Stable +gemini extensions install github.com/user/repo + +# Dev +gemini extensions install github.com/user/repo --ref dev +``` + +### Clean artifacts + +If you are using GitHub Releases, ensure your release artifacts only contain the +necessary files (`dist/`, `gemini-extension.json`, `package.json`). Exclude +`node_modules` (users will install them) and `src/` to keep downloads small. diff --git a/docs/extensions/index.md b/docs/extensions/index.md index a2b0598388..b44210762d 100644 --- a/docs/extensions/index.md +++ b/docs/extensions/index.md @@ -1,377 +1,44 @@ # Gemini CLI extensions -_This documentation is up-to-date with the v0.4.0 release._ - -Gemini CLI extensions package prompts, MCP servers, Agent Skills, and custom -commands into a familiar and user-friendly format. With extensions, you can +Gemini CLI extensions package prompts, MCP servers, custom commands, hooks, and +agent skills into a familiar and user-friendly format. With extensions, you can expand the capabilities of Gemini CLI and share those capabilities with others. -They're designed to be easily installable and shareable. +They are designed to be easily installable and shareable. To see examples of extensions, you can browse a gallery of [Gemini CLI extensions](https://geminicli.com/extensions/browse/). -See [getting started docs](getting-started-extensions.md) for a guide on -creating your first extension. +## Managing extensions -See [releasing docs](extension-releasing.md) for an advanced guide on setting up -GitHub releases. +You can verify your installed extensions and their status using the interactive +command: -## Extension management - -We offer a suite of extension management tools using `gemini extensions` -commands. - -Note that these commands are not supported from within the CLI, although you can -list installed extensions using the `/extensions list` subcommand. - -Note that all of these commands will only be reflected in active CLI sessions on -restart. - -### Installing an extension - -You can install an extension using `gemini extensions install` with either a -GitHub URL or a local path. - -Note that we create a copy of the installed extension, so you will need to run -`gemini extensions update` to pull in changes from both locally-defined -extensions and those on GitHub. - -NOTE: If you are installing an extension from GitHub, you'll need to have `git` -installed on your machine. See -[git installation instructions](https://git-scm.com/book/en/v2/Getting-Started-Installing-Git) -for help. - -``` -gemini extensions install [--ref ] [--auto-update] [--pre-release] [--consent] +```bash +/extensions list ``` -- ``: The github URL or local path of the extension to install. -- `--ref`: The git ref to install from. -- `--auto-update`: Enable auto-update for this extension. -- `--pre-release`: Enable pre-release versions for this extension. -- `--consent`: Acknowledge the security risks of installing an extension and - skip the confirmation prompt. +or in noninteractive mode: -### Uninstalling an extension - -To uninstall one or more extensions, run -`gemini extensions uninstall `: - -``` -gemini extensions uninstall gemini-cli-security gemini-cli-another-extension -``` - -### Disabling an extension - -Extensions are, by default, enabled across all workspaces. You can disable an -extension entirely or for specific workspace. - -``` -gemini extensions disable [--scope ] -``` - -- ``: The name of the extension to disable. -- `--scope`: The scope to disable the extension in (`user` or `workspace`). - -### Enabling an extension - -You can enable extensions using `gemini extensions enable `. You can also -enable an extension for a specific workspace using -`gemini extensions enable --scope=workspace` from within that workspace. - -``` -gemini extensions enable [--scope ] -``` - -- ``: The name of the extension to enable. -- `--scope`: The scope to enable the extension in (`user` or `workspace`). - -### Updating an extension - -For extensions installed from a local path or a git repository, you can -explicitly update to the latest version (as reflected in the -`gemini-extension.json` `version` field) with `gemini extensions update `. - -You can update all extensions with: - -``` -gemini extensions update --all -``` - -### Create a boilerplate extension - -We offer several example extensions `context`, `custom-commands`, -`exclude-tools` and `mcp-server`. You can view these examples -[here](https://github.com/google-gemini/gemini-cli/tree/main/packages/cli/src/commands/extensions/examples). - -To copy one of these examples into a development directory using the type of -your choosing, run: - -``` -gemini extensions new [template] -``` - -- ``: The path to create the extension in. -- `[template]`: The boilerplate template to use. - -### Link a local extension - -The `gemini extensions link` command will create a symbolic link from the -extension installation directory to the development path. - -This is useful so you don't have to run `gemini extensions update` every time -you make changes you'd like to test. - -``` -gemini extensions link -``` - -- ``: The path of the extension to link. - -## How it works - -On startup, Gemini CLI looks for extensions in `/.gemini/extensions` - -Extensions exist as a directory that contains a `gemini-extension.json` file. -For example: - -`/.gemini/extensions/my-extension/gemini-extension.json` - -### `gemini-extension.json` - -The `gemini-extension.json` file contains the configuration for the extension. -The file has the following structure: - -```json -{ - "name": "my-extension", - "version": "1.0.0", - "mcpServers": { - "my-server": { - "command": "node my-server.js" - } - }, - "contextFileName": "GEMINI.md", - "excludeTools": ["run_shell_command"] -} -``` - -- `name`: The name of the extension. This is used to uniquely identify the - extension and for conflict resolution when extension commands have the same - name as user or project commands. The name should be lowercase or numbers and - use dashes instead of underscores or spaces. This is how users will refer to - your extension in the CLI. Note that we expect this name to match the - extension directory name. -- `version`: The version of the extension. -- `mcpServers`: A map of MCP servers to settings. The key is the name of the - server, and the value is the server configuration. These servers will be - loaded on startup just like MCP servers settings in a - [`settings.json` file](../get-started/configuration.md). If both an extension - and a `settings.json` file settings an MCP server with the same name, the - server defined in the `settings.json` file takes precedence. - - Note that all MCP server configuration options are supported except for - `trust`. -- `contextFileName`: The name of the file that contains the context for the - extension. This will be used to load the context from the extension directory. - If this property is not used but a `GEMINI.md` file is present in your - extension directory, then that file will be loaded. -- `excludeTools`: An array of tool names to exclude from the model. You can also - specify command-specific restrictions for tools that support it, like the - `run_shell_command` tool. For example, - `"excludeTools": ["run_shell_command(rm -rf)"]` will block the `rm -rf` - command. Note that this differs from the MCP server `excludeTools` - functionality, which can be listed in the MCP server config. - -When Gemini CLI starts, it loads all the extensions and merges their -configurations. If there are any conflicts, the workspace configuration takes -precedence. - -### Settings - -_Note: This is an experimental feature. We do not yet recommend extension -authors introduce settings as part of their core flows._ - -Extensions can define settings that the user will be prompted to provide upon -installation. This is useful for things like API keys, URLs, or other -configuration that the extension needs to function. - -To define settings, add a `settings` array to your `gemini-extension.json` file. -Each object in the array should have the following properties: - -- `name`: A user-friendly name for the setting. -- `description`: A description of the setting and what it's used for. -- `envVar`: The name of the environment variable that the setting will be stored - as. -- `sensitive`: Optional boolean. If true, obfuscates the input the user provides - and stores the secret in keychain storage. **Example** - -```json -{ - "name": "my-api-extension", - "version": "1.0.0", - "settings": [ - { - "name": "API Key", - "description": "Your API key for the service.", - "envVar": "MY_API_KEY" - } - ] -} -``` - -When a user installs this extension, they will be prompted to enter their API -key. The value will be saved to a `.env` file in the extension's directory -(e.g., `/.gemini/extensions/my-api-extension/.env`). - -You can view a list of an extension's settings by running: - -``` +```bash gemini extensions list ``` -and you can update a given setting using: +## Installation -``` -gemini extensions config [setting name] [--scope ] +To install a real extension, you can use the `extensions install` command with a +GitHub repository URL in noninteractive mode. For example: + +```bash +gemini extensions install https://github.com/gemini-cli-extensions/workspace ``` -- `--scope`: The scope to set the setting in (`user` or `workspace`). This is - optional and will default to `user`. +## Next steps -### Custom commands - -Extensions can provide [custom commands](../cli/custom-commands.md) by placing -TOML files in a `commands/` subdirectory within the extension directory. These -commands follow the same format as user and project custom commands and use -standard naming conventions. - -**Example** - -An extension named `gcp` with the following structure: - -``` -.gemini/extensions/gcp/ -├── gemini-extension.json -└── commands/ - ├── deploy.toml - └── gcs/ - └── sync.toml -``` - -Would provide these commands: - -- `/deploy` - Shows as `[gcp] Custom command from deploy.toml` in help -- `/gcs:sync` - Shows as `[gcp] Custom command from sync.toml` in help - -### Agent Skills - -_Note: This is an experimental feature enabled via `experimental.skills`._ - -Extensions can bundle [Agent Skills](../cli/skills.md) to provide on-demand -expertise and specialized workflows. To include skills in your extension, place -them in a `skills/` subdirectory within the extension directory. Each skill must -follow the [Agent Skills structure](../cli/skills.md#folder-structure), -including a `SKILL.md` file. - -**Example** - -An extension named `security-toolkit` with the following structure: - -``` -.gemini/extensions/security-toolkit/ -├── gemini-extension.json -└── skills/ - ├── audit/ - │ ├── SKILL.md - │ └── scripts/ - │ └── scan.py - └── hardening/ - └── SKILL.md -``` - -Upon installation, these skills will be discovered by Gemini CLI and can be -activated during a session when the model identifies a task matching their -descriptions. - -Extension skills have the lowest precedence and will be overridden by user or -workspace skills of the same name. They can be viewed and managed (enabled or -disabled) using the [`/skills` command](../cli/skills.md#managing-skills). - -### Hooks - -Extensions can provide [hooks](../hooks/index.md) to intercept and customize -Gemini CLI behavior at specific lifecycle events. Hooks provided by an extension -must be defined in a `hooks/hooks.json` file within the extension directory. - -> [!IMPORTANT] Hooks are not defined directly in `gemini-extension.json`. The -> CLI specifically looks for the `hooks/hooks.json` file. - -#### Directory structure - -``` -.gemini/extensions/my-extension/ -├── gemini-extension.json -└── hooks/ - └── hooks.json -``` - -#### `hooks/hooks.json` format - -The `hooks.json` file contains a `hooks` object where keys are -[event names](../hooks/reference.md#supported-events) and values are arrays of -[hook definitions](../hooks/reference.md#hook-definition). - -```json -{ - "hooks": { - "BeforeAgent": [ - { - "hooks": [ - { - "type": "command", - "command": "node ${extensionPath}/scripts/setup.js", - "name": "Extension Setup" - } - ] - } - ] - } -} -``` - -#### Supported variables - -Just like `gemini-extension.json`, the `hooks/hooks.json` file supports -[variable substitution](#variables). This is particularly useful for referencing -scripts within the extension directory using `${extensionPath}`. - -### Conflict resolution - -Extension commands have the lowest precedence. When a conflict occurs with user -or project commands: - -1. **No conflict**: Extension command uses its natural name (e.g., `/deploy`) -2. **With conflict**: Extension command is renamed with the extension prefix - (e.g., `/gcp.deploy`) - -For example, if both a user and the `gcp` extension define a `deploy` command: - -- `/deploy` - Executes the user's deploy command -- `/gcp.deploy` - Executes the extension's deploy command (marked with `[gcp]` - tag) - -### Variables - -Gemini CLI extensions allow variable substitution in both -`gemini-extension.json` and `hooks/hooks.json`. This can be useful if e.g., you -need the current directory to run an MCP server or hook script using -`"cwd": "${extensionPath}${/}run.ts"`. - -**Supported variables:** - -| variable | description | -| -------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `${extensionPath}` | The fully-qualified path of the extension in the user's filesystem e.g., '/Users/username/.gemini/extensions/example-extension'. This will not unwrap symlinks. | -| `${workspacePath}` | The fully-qualified path of the current workspace. | -| `${/} or ${pathSeparator}` | The path separator (differs per OS). | -| `${process.execPath}` | The path to the Node.js binary executing the CLI. | +- [Writing extensions](writing-extensions.md): Learn how to create your first + extension. +- [Extensions reference](reference.md): Deeply understand the extension format, + commands, and configuration. +- [Best practices](best-practices.md): Learn strategies for building great + extensions. +- [Extensions releasing](releasing.md): Learn how to share your extensions with + the world. diff --git a/docs/extensions/reference.md b/docs/extensions/reference.md new file mode 100644 index 0000000000..12d04e0bdd --- /dev/null +++ b/docs/extensions/reference.md @@ -0,0 +1,312 @@ +# Extensions reference + +This guide covers the `gemini extensions` commands and the structure of the +`gemini-extension.json` configuration file. + +## Extension management + +We offer a suite of extension management tools using `gemini extensions` +commands. + +Note that these commands (e.g. `gemini extensions install`) are not supported +from within the CLI's **interactive mode**, although you can list installed +extensions using the `/extensions list` slash command. + +Note that all of these management operations (including updates to slash +commands) will only be reflected in active CLI sessions on **restart**. + +### Installing an extension + +You can install an extension using `gemini extensions install` with either a +GitHub URL or a local path. + +Note that we create a copy of the installed extension, so you will need to run +`gemini extensions update` to pull in changes from both locally-defined +extensions and those on GitHub. + +NOTE: If you are installing an extension from GitHub, you'll need to have `git` +installed on your machine. See +[git installation instructions](https://git-scm.com/book/en/v2/Getting-Started-Installing-Git) +for help. + +``` +gemini extensions install [--ref ] [--auto-update] [--pre-release] [--consent] +``` + +- ``: The github URL or local path of the extension to install. +- `--ref`: The git ref to install from. +- `--auto-update`: Enable auto-update for this extension. +- `--pre-release`: Enable pre-release versions for this extension. +- `--consent`: Acknowledge the security risks of installing an extension and + skip the confirmation prompt. + +### Uninstalling an extension + +To uninstall one or more extensions, run +`gemini extensions uninstall `: + +``` +gemini extensions uninstall gemini-cli-security gemini-cli-another-extension +``` + +### Disabling an extension + +Extensions are, by default, enabled across all workspaces. You can disable an +extension entirely or for specific workspace. + +``` +gemini extensions disable [--scope ] +``` + +- ``: The name of the extension to disable. +- `--scope`: The scope to disable the extension in (`user` or `workspace`). + +### Enabling an extension + +You can enable extensions using `gemini extensions enable `. You can also +enable an extension for a specific workspace using +`gemini extensions enable --scope=workspace` from within that workspace. + +``` +gemini extensions enable [--scope ] +``` + +- ``: The name of the extension to enable. +- `--scope`: The scope to enable the extension in (`user` or `workspace`). + +### Updating an extension + +For extensions installed from a local path or a git repository, you can +explicitly update to the latest version (as reflected in the +`gemini-extension.json` `version` field) with `gemini extensions update `. + +You can update all extensions with: + +``` +gemini extensions update --all +``` + +### Create a boilerplate extension + +We offer several example extensions `context`, `custom-commands`, +`exclude-tools` and `mcp-server`. You can view these examples +[here](https://github.com/google-gemini/gemini-cli/tree/main/packages/cli/src/commands/extensions/examples). + +To copy one of these examples into a development directory using the type of +your choosing, run: + +``` +gemini extensions new [template] +``` + +- ``: The path to create the extension in. +- `[template]`: The boilerplate template to use. + +### Link a local extension + +The `gemini extensions link` command will create a symbolic link from the +extension installation directory to the development path. + +This is useful so you don't have to run `gemini extensions update` every time +you make changes you'd like to test. + +``` +gemini extensions link +``` + +- ``: The path of the extension to link. + +## Extension format + +On startup, Gemini CLI looks for extensions in `/.gemini/extensions` + +Extensions exist as a directory that contains a `gemini-extension.json` file. +For example: + +`/.gemini/extensions/my-extension/gemini-extension.json` + +### `gemini-extension.json` + +The `gemini-extension.json` file contains the configuration for the extension. +The file has the following structure: + +```json +{ + "name": "my-extension", + "version": "1.0.0", + "description": "My awesome extension", + "mcpServers": { + "my-server": { + "command": "node my-server.js" + } + }, + "contextFileName": "GEMINI.md", + "excludeTools": ["run_shell_command"] +} +``` + +- `name`: The name of the extension. This is used to uniquely identify the + extension and for conflict resolution when extension commands have the same + name as user or project commands. The name should be lowercase or numbers and + use dashes instead of underscores or spaces. This is how users will refer to + your extension in the CLI. Note that we expect this name to match the + extension directory name. +- `version`: The version of the extension. +- `description`: A short description of the extension. This will be displayed on + [geminicli.com/extensions](https://geminicli.com/extensions). +- `mcpServers`: A map of MCP servers to settings. The key is the name of the + server, and the value is the server configuration. These servers will be + loaded on startup just like MCP servers settingsd in a + [`settings.json` file](../get-started/configuration.md). If both an extension + and a `settings.json` file settings an MCP server with the same name, the + server defined in the `settings.json` file takes precedence. + - Note that all MCP server configuration options are supported except for + `trust`. +- `contextFileName`: The name of the file that contains the context for the + extension. This will be used to load the context from the extension directory. + If this property is not used but a `GEMINI.md` file is present in your + extension directory, then that file will be loaded. +- `excludeTools`: An array of tool names to exclude from the model. You can also + specify command-specific restrictions for tools that support it, like the + `run_shell_command` tool. For example, + `"excludeTools": ["run_shell_command(rm -rf)"]` will block the `rm -rf` + command. Note that this differs from the MCP server `excludeTools` + functionality, which can be listed in the MCP server config. + +When Gemini CLI starts, it loads all the extensions and merges their +configurations. If there are any conflicts, the workspace configuration takes +precedence. + +### Settings + +_Note: This is an experimental feature. We do not yet recommend extension +authors introduce settings as part of their core flows._ + +Extensions can define settings that the user will be prompted to provide upon +installation. This is useful for things like API keys, URLs, or other +configuration that the extension needs to function. + +To define settings, add a `settings` array to your `gemini-extension.json` file. +Each object in the array should have the following properties: + +- `name`: A user-friendly name for the setting. +- `description`: A description of the setting and what it's used for. +- `envVar`: The name of the environment variable that the setting will be stored + as. +- `sensitive`: Optional boolean. If true, obfuscates the input the user provides + and stores the secret in keychain storage. **Example** + +```json +{ + "name": "my-api-extension", + "version": "1.0.0", + "settings": [ + { + "name": "API Key", + "description": "Your API key for the service.", + "envVar": "MY_API_KEY" + } + ] +} +``` + +When a user installs this extension, they will be prompted to enter their API +key. The value will be saved to a `.env` file in the extension's directory +(e.g., `/.gemini/extensions/my-api-extension/.env`). + +You can view a list of an extension's settings by running: + +``` +gemini extensions list +``` + +and you can update a given setting using: + +``` +gemini extensions config [setting name] [--scope ] +``` + +- `--scope`: The scope to set the setting in (`user` or `workspace`). This is + optional and will default to `user`. + +### Custom commands + +Extensions can provide [custom commands](../cli/custom-commands.md) by placing +TOML files in a `commands/` subdirectory within the extension directory. These +commands follow the same format as user and project custom commands and use +standard naming conventions. + +**Example** + +An extension named `gcp` with the following structure: + +``` +.gemini/extensions/gcp/ +├── gemini-extension.json +└── commands/ + ├── deploy.toml + └── gcs/ + └── sync.toml +``` + +Would provide these commands: + +- `/deploy` - Shows as `[gcp] Custom command from deploy.toml` in help +- `/gcs:sync` - Shows as `[gcp] Custom command from sync.toml` in help + +### Hooks + +Extensions can provide [hooks](../hooks/index.md) to intercept and customize +Gemini CLI behavior at specific lifecycle events. Hooks provided by an extension +must be defined in a `hooks/hooks.json` file within the extension directory. + +> [!IMPORTANT] Hooks are not defined directly in `gemini-extension.json`. The +> CLI specifically looks for the `hooks/hooks.json` file. + +### Agent Skills + +Extensions can bundle [Agent Skills](../cli/skills.md) to provide specialized +workflows. Skills must be placed in a `skills/` directory within the extension. + +**Example** + +An extension with the following structure: + +``` +.gemini/extensions/my-extension/ +├── gemini-extension.json +└── skills/ + └── security-audit/ + └── SKILL.md +``` + +Will expose a `security-audit` skill that the model can activate. + +### Conflict resolution + +Extension commands have the lowest precedence. When a conflict occurs with user +or project commands: + +1. **No conflict**: Extension command uses its natural name (e.g., `/deploy`) +2. **With conflict**: Extension command is renamed with the extension prefix + (e.g., `/gcp.deploy`) + +For example, if both a user and the `gcp` extension define a `deploy` command: + +- `/deploy` - Executes the user's deploy command +- `/gcp.deploy` - Executes the extension's deploy command (marked with `[gcp]` + tag) + +## Variables + +Gemini CLI extensions allow variable substitution in `gemini-extension.json`. +This can be useful if e.g., you need the current directory to run an MCP server +using an argument like `"args": ["${extensionPath}${/}dist${/}server.js"]`. + +**Supported variables:** + +| variable | description | +| -------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `${extensionPath}` | The fully-qualified path of the extension in the user's filesystem e.g., '/Users/username/.gemini/extensions/example-extension'. This will not unwrap symlinks. | +| `${workspacePath}` | The fully-qualified path of the current workspace. | +| `${/} or ${pathSeparator}` | The path separator (differs per OS). | diff --git a/docs/extensions/extension-releasing.md b/docs/extensions/releasing.md similarity index 100% rename from docs/extensions/extension-releasing.md rename to docs/extensions/releasing.md diff --git a/docs/extensions/getting-started-extensions.md b/docs/extensions/writing-extensions.md similarity index 63% rename from docs/extensions/getting-started-extensions.md rename to docs/extensions/writing-extensions.md index 04e5987c85..3ed00fc577 100644 --- a/docs/extensions/getting-started-extensions.md +++ b/docs/extensions/writing-extensions.md @@ -8,7 +8,19 @@ file. ## Prerequisites Before you start, make sure you have the Gemini CLI installed and a basic -understanding of Node.js and TypeScript. +understanding of Node.js. + +## When to use what + +Extensions offer a variety of ways to customize Gemini CLI. + +| Feature | What it is | When to use it | Invoked by | +| :------------------------------------------------------------- | :----------------------------------------------------------------------------------------------------------------- | :--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | :-------------------- | +| **[MCP server](reference.md#mcp-servers)** | A standard way to expose new tools and data sources to the model. | Use this when you want the model to be able to _do_ new things, like fetching data from an internal API, querying a database, or controlling a local application. We also support MCP resources (which can replace custom commands) and system instructions (which can replace custom context) | Model | +| **[Custom commands](../cli/custom-commands.md)** | A shortcut (like `/my-cmd`) that executes a pre-defined prompt or shell command. | Use this for repetitive tasks or to save long, complex prompts that you use frequently. Great for automation. | User | +| **[Context file (`GEMINI.md`)](reference.md#contextfilename)** | A markdown file containing instructions that are loaded into the model's context at the start of every session. | Use this to define the "personality" of your extension, set coding standards, or provide essential knowledge that the model should always have. | CLI provides to model | +| **[Agent skills](../cli/skills.md)** | A specialized set of instructions and workflows that the model activates only when needed. | Use this for complex, occasional tasks (like "create a PR" or "audit security") to avoid cluttering the main context window when the skill isn't being used. | Model | +| **[Hooks](../hooks/index.md)** | A way to intercept and customize the CLI's behavior at specific lifecycle events (e.g., before/after a tool call). | Use this when you want to automate actions based on what the model is doing, like validating tool arguments, logging activity, or modifying the model's input/output. | CLI | ## Step 1: Create a new extension @@ -26,10 +38,9 @@ This will create a new directory with the following structure: ``` my-first-extension/ -├── example.ts +├── example.js ├── gemini-extension.json -├── package.json -└── tsconfig.json +└── package.json ``` ## Step 2: Understand the extension files @@ -43,12 +54,12 @@ and use your extension. ```json { - "name": "my-first-extension", + "name": "mcp-server-example", "version": "1.0.0", "mcpServers": { "nodeServer": { "command": "node", - "args": ["${extensionPath}${/}dist${/}example.js"], + "args": ["${extensionPath}${/}example.js"], "cwd": "${extensionPath}" } } @@ -64,12 +75,12 @@ and use your extension. with the absolute path to your extension's installation directory. This allows your extension to work regardless of where it's installed. -### `example.ts` +### `example.js` This file contains the source code for your MCP server. It's a simple Node.js server that uses the `@modelcontextprotocol/sdk`. -```typescript +```javascript /** * @license * Copyright 2025 Google LLC @@ -118,16 +129,15 @@ await server.connect(transport); This server defines a single tool called `fetch_posts` that fetches data from a public API. -### `package.json` and `tsconfig.json` +### `package.json` -These are standard configuration files for a TypeScript project. The -`package.json` file defines dependencies and a `build` script, and -`tsconfig.json` configures the TypeScript compiler. +This is the standard configuration file for a Node.js project. It defines +dependencies and scripts. -## Step 3: Build and link your extension +## Step 3: Link your extension -Before you can use the extension, you need to compile the TypeScript code and -link the extension to your Gemini CLI installation for local development. +Before you can use the extension, you need to link it to your Gemini CLI +installation for local development. 1. **Install dependencies:** @@ -136,16 +146,7 @@ link the extension to your Gemini CLI installation for local development. npm install ``` -2. **Build the server:** - - ```bash - npm run build - ``` - - This will compile `example.ts` into `dist/example.js`, which is the file - referenced in your `gemini-extension.json`. - -3. **Link the extension:** +2. **Link the extension:** The `link` command creates a symbolic link from the Gemini CLI extensions directory to your development directory. This means any changes you make @@ -212,7 +213,7 @@ need this for extensions built to expose commands and prompts. "mcpServers": { "nodeServer": { "command": "node", - "args": ["${extensionPath}${/}dist${/}example.js"], + "args": ["${extensionPath}${/}example.js"], "cwd": "${extensionPath}" } } @@ -265,7 +266,7 @@ primary ways of releasing extensions are via a Git repository or through GitHub Releases. Using a public Git repository is the simplest method. For detailed instructions on both methods, please refer to the -[Extension Releasing Guide](./extension-releasing.md). +[Extension Releasing Guide](./releasing.md). ## Conclusion diff --git a/docs/index.md b/docs/index.md index 217fba8391..3669b961ad 100644 --- a/docs/index.md +++ b/docs/index.md @@ -102,10 +102,10 @@ This documentation is organized into the following sections: - **[Introduction: Extensions](./extensions/index.md):** How to extend the CLI with new functionality. -- **[Get Started with extensions](./extensions/getting-started-extensions.md):** - Learn how to build your own extension. -- **[Extension releasing](./extensions/extension-releasing.md):** How to release - Gemini CLI extensions. +- **[Writing extensions](./extensions/writing-extensions.md):** Learn how to + build your own extension. +- **[Extension releasing](./extensions/releasing.md):** How to release Gemini + CLI extensions. ### Hooks diff --git a/docs/sidebar.json b/docs/sidebar.json index 1583674d03..ec16cf8444 100644 --- a/docs/sidebar.json +++ b/docs/sidebar.json @@ -192,12 +192,20 @@ "slug": "docs/extensions" }, { - "label": "Get started with extensions", - "slug": "docs/extensions/getting-started-extensions" + "label": "Writing extensions", + "slug": "docs/extensions/writing-extensions" }, { - "label": "Extension releasing", - "slug": "docs/extensions/extension-releasing" + "label": "Extensions reference", + "slug": "docs/extensions/reference" + }, + { + "label": "Best practices", + "slug": "docs/extensions/best-practices" + }, + { + "label": "Extensions releasing", + "slug": "docs/extensions/releasing" } ] }, From c2d078396502232d0a3796e52e5f10fa96ec6ebe Mon Sep 17 00:00:00 2001 From: Jenna Inouye Date: Mon, 26 Jan 2026 13:19:27 -0800 Subject: [PATCH 7/9] Docs: Refactor left nav on the website (#17558) --- docs/sidebar.json | 307 +++++++++------------------------------- docs/troubleshooting.md | 3 - 2 files changed, 67 insertions(+), 243 deletions(-) diff --git a/docs/sidebar.json b/docs/sidebar.json index ec16cf8444..f033bafa36 100644 --- a/docs/sidebar.json +++ b/docs/sidebar.json @@ -1,187 +1,52 @@ [ - { - "label": "Overview", - "items": [ - { - "label": "Introduction", - "slug": "docs" - }, - { - "label": "Architecture overview", - "slug": "docs/architecture" - }, - { - "label": "Contribution guide", - "slug": "docs/contributing" - } - ] - }, { "label": "Get started", "items": [ - { - "label": "Gemini CLI quickstart", - "slug": "docs/get-started" - }, - { - "label": "Gemini 3 on Gemini CLI", - "slug": "docs/get-started/gemini-3" - }, - { - "label": "Authentication", - "slug": "docs/get-started/authentication" - }, - { - "label": "Configuration", - "slug": "docs/get-started/configuration" - }, - { - "label": "Installation", - "slug": "docs/get-started/installation" - }, - { - "label": "Examples", - "slug": "docs/get-started/examples" - } + { "label": "Overview", "slug": "docs" }, + { "label": "Quickstart", "slug": "docs/get-started" }, + { "label": "Installation", "slug": "docs/get-started/installation" }, + { "label": "Authentication", "slug": "docs/get-started/authentication" }, + { "label": "Examples", "slug": "docs/get-started/examples" }, + { "label": "Gemini 3 (preview)", "slug": "docs/get-started/gemini-3" } ] }, { - "label": "CLI", + "label": "Use Gemini CLI", "items": [ - { - "label": "Introduction", - "slug": "docs/cli" - }, - { - "label": "Commands", - "slug": "docs/cli/commands" - }, - { - "label": "Checkpointing", - "slug": "docs/cli/checkpointing" - }, - { - "label": "Custom commands", - "slug": "docs/cli/custom-commands" - }, - { - "label": "Enterprise", - "slug": "docs/cli/enterprise" - }, - { - "label": "Headless mode", - "slug": "docs/cli/headless" - }, - { - "label": "Keyboard shortcuts", - "slug": "docs/cli/keyboard-shortcuts" - }, - { - "label": "Model selection", - "slug": "docs/cli/model" - }, - { - "label": "Sandbox", - "slug": "docs/cli/sandbox" - }, - { - "label": "Session Management", - "slug": "docs/cli/session-management" - }, - { - "label": "Agent Skills", - "slug": "docs/cli/skills" - }, - { - "label": "Settings", - "slug": "docs/cli/settings" - }, - { - "label": "Telemetry", - "slug": "docs/cli/telemetry" - }, - { - "label": "Themes", - "slug": "docs/cli/themes" - }, - { - "label": "Token caching", - "slug": "docs/cli/token-caching" - }, - { - "label": "Trusted Folders", - "slug": "docs/cli/trusted-folders" - }, - { - "label": "Tutorials", - "slug": "docs/cli/tutorials" - }, - { - "label": "Uninstall", - "slug": "docs/cli/uninstall" - }, - { - "label": "System prompt override", - "slug": "docs/cli/system-prompt" - } + { "label": "Using the CLI", "slug": "docs/cli" }, + { "label": "File management", "slug": "docs/tools/file-system" }, + { "label": "Memory management", "slug": "docs/tools/memory" }, + { "label": "Project context (GEMINI.md)", "slug": "docs/cli/gemini-md" }, + { "label": "Shell commands", "slug": "docs/tools/shell" }, + { "label": "Session management", "slug": "docs/cli/session-management" }, + { "label": "Todos", "slug": "docs/tools/todos" }, + { "label": "Web search and fetch", "slug": "docs/tools/web-search" } ] }, { - "label": "Core", + "label": "Configuration", "items": [ { - "label": "Introduction", - "slug": "docs/core" + "label": "Ignore files (.geminiignore)", + "slug": "docs/cli/gemini-ignore" }, - { - "label": "Tools API", - "slug": "docs/core/tools-api" - }, - { - "label": "Memory Import Processor", - "slug": "docs/core/memport" - }, - { - "label": "Policy Engine", - "slug": "docs/core/policy-engine" - } + { "label": "Model selection", "slug": "docs/cli/model" }, + { "label": "Settings", "slug": "docs/cli/settings" }, + { "label": "Themes", "slug": "docs/cli/themes" }, + { "label": "Token caching", "slug": "docs/cli/token-caching" }, + { "label": "Trusted folders", "slug": "docs/cli/trusted-folders" } ] }, { - "label": "Tools", + "label": "Advanced features", "items": [ - { - "label": "Introduction", - "slug": "docs/tools" - }, - { - "label": "File system", - "slug": "docs/tools/file-system" - }, - { - "label": "Shell", - "slug": "docs/tools/shell" - }, - { - "label": "Web fetch", - "slug": "docs/tools/web-fetch" - }, - { - "label": "Web search", - "slug": "docs/tools/web-search" - }, - { - "label": "Memory", - "slug": "docs/tools/memory" - }, - { - "label": "Todos", - "slug": "docs/tools/todos" - }, - { - "label": "MCP servers", - "slug": "docs/tools/mcp-server" - } + { "label": "Checkpointing", "slug": "docs/cli/checkpointing" }, + { "label": "Custom commands", "slug": "docs/cli/custom-commands" }, + { "label": "Enterprise features", "slug": "docs/cli/enterprise" }, + { "label": "Headless mode & scripting", "slug": "docs/cli/headless" }, + { "label": "Sandboxing", "slug": "docs/cli/sandbox" }, + { "label": "System prompt override", "slug": "docs/cli/system-prompt" }, + { "label": "Telemetry", "slug": "docs/cli/telemetry" } ] }, { @@ -210,96 +75,58 @@ ] }, { - "label": "Hooks (experimental)", + "label": "Ecosystem and extensibility", "items": [ - { - "label": "Introduction", - "slug": "docs/hooks" - }, - { - "label": "Writing hooks", - "slug": "docs/hooks/writing-hooks" - }, - { - "label": "Hooks reference", - "slug": "docs/hooks/reference" - }, - { - "label": "Best practices", - "slug": "docs/hooks/best-practices" - } + { "label": "Agent skills (experimental)", "slug": "docs/cli/skills" }, + { "label": "Hooks (experimental)", "slug": "docs/hooks" }, + { "label": "IDE integration", "slug": "docs/ide-integration" }, + { "label": "MCP servers", "slug": "docs/tools/mcp-server" } ] }, { - "label": "IDE integration", + "label": "Tutorials", "items": [ { - "label": "Introduction", - "slug": "docs/ide-integration" + "label": "Get started with extensions", + "slug": "docs/extensions/getting-started-extensions" }, - { - "label": "IDE companion spec", - "slug": "docs/ide-integration/ide-companion-spec" - } + { "label": "How to write hooks", "slug": "docs/hooks/writing-hooks" } + ] + }, + { + "label": "Reference", + "items": [ + { "label": "Architecture", "slug": "docs/architecture" }, + { "label": "Command reference", "slug": "docs/cli/commands" }, + { "label": "Configuration", "slug": "docs/get-started/configuration" }, + { "label": "Keyboard shortcuts", "slug": "docs/cli/keyboard-shortcuts" }, + { "label": "Memory import processor", "slug": "docs/core/memport" }, + { "label": "Policy engine", "slug": "docs/core/policy-engine" }, + { "label": "Tools API", "slug": "docs/core/tools-api" } + ] + }, + { + "label": "Resources", + "items": [ + { "label": "FAQ", "slug": "docs/faq" }, + { "label": "Quota and pricing", "slug": "docs/quota-and-pricing" }, + { "label": "Release notes", "slug": "docs/changelogs/" }, + { "label": "Terms and privacy", "slug": "docs/tos-privacy" }, + { "label": "Troubleshooting", "slug": "docs/troubleshooting" }, + { "label": "Uninstall", "slug": "docs/cli/uninstall" } ] }, { "label": "Development", "items": [ - { - "label": "NPM", - "slug": "docs/npm" - }, - { - "label": "Releases", - "slug": "docs/releases" - }, - { - "label": "Integration tests", - "slug": "docs/integration-tests" - }, + { "label": "Contribution guide", "slug": "docs/CONTRIBUTING" }, + { "label": "Integration testing", "slug": "docs/integration-tests" }, { "label": "Issue and PR automation", "slug": "docs/issue-and-pr-automation" - } - ] - }, - { - "label": "Releases", - "items": [ - { - "label": "Release notes", - "slug": "docs/changelogs/" }, - { - "label": "Latest release", - "slug": "docs/changelogs/latest" - }, - { - "label": "Preview release", - "slug": "docs/changelogs/preview" - } - ] - }, - { - "label": "Support", - "items": [ - { - "label": "FAQ", - "slug": "docs/faq" - }, - { - "label": "Troubleshooting", - "slug": "docs/troubleshooting" - }, - { - "label": "Quota and pricing", - "slug": "docs/quota-and-pricing" - }, - { - "label": "Terms of service", - "slug": "docs/tos-privacy" - } + { "label": "Local development", "slug": "docs/local-development" }, + { "label": "NPM package structure", "slug": "docs/npm" } ] } ] diff --git a/docs/troubleshooting.md b/docs/troubleshooting.md index 515099934a..f700d0b74f 100644 --- a/docs/troubleshooting.md +++ b/docs/troubleshooting.md @@ -34,9 +34,6 @@ topics on: list of supported locations, see the following pages: - Gemini Code Assist for individuals: [Available locations](https://developers.google.com/gemini-code-assist/resources/available-locations#americas) - - Google AI Pro and Ultra where Gemini Code Assist (and Gemini CLI) is also - available: - [Available locations](https://developers.google.com/gemini-code-assist/resources/locations-pro-ultra) - **Error: `Failed to login. Message: Request contains an invalid argument`** - **Cause:** Users with Google Workspace accounts or Google Cloud accounts From 018dc0d5cfc17e4b9e35e29344dd1cc6a05866fc Mon Sep 17 00:00:00 2001 From: Adam Weidman <65992621+adamfweidman@users.noreply.github.com> Date: Mon, 26 Jan 2026 16:52:19 -0500 Subject: [PATCH 8/9] fix(core): stream grep/ripgrep output to prevent OOM (#17146) --- packages/core/src/tools/constants.ts | 7 + packages/core/src/tools/grep.test.ts | 42 + packages/core/src/tools/grep.ts | 375 +++--- packages/core/src/tools/ripGrep.test.ts | 1059 ++++++----------- packages/core/src/tools/ripGrep.ts | 211 ++-- .../src/utils/shell-utils.integration.test.ts | 67 ++ packages/core/src/utils/shell-utils.ts | 121 ++ 7 files changed, 888 insertions(+), 994 deletions(-) create mode 100644 packages/core/src/tools/constants.ts create mode 100644 packages/core/src/utils/shell-utils.integration.test.ts diff --git a/packages/core/src/tools/constants.ts b/packages/core/src/tools/constants.ts new file mode 100644 index 0000000000..132e8c104a --- /dev/null +++ b/packages/core/src/tools/constants.ts @@ -0,0 +1,7 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ +export const DEFAULT_TOTAL_MAX_MATCHES = 20000; +export const DEFAULT_SEARCH_TIMEOUT_MS = 30000; diff --git a/packages/core/src/tools/grep.test.ts b/packages/core/src/tools/grep.test.ts index 7c9f224feb..0f0db86665 100644 --- a/packages/core/src/tools/grep.test.ts +++ b/packages/core/src/tools/grep.test.ts @@ -7,6 +7,7 @@ import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest'; import type { GrepToolParams } from './grep.js'; import { GrepTool } from './grep.js'; +import type { ToolResult } from './tools.js'; import path from 'node:path'; import fs from 'node:fs/promises'; import os from 'node:os'; @@ -15,8 +16,12 @@ import { createMockWorkspaceContext } from '../test-utils/mockWorkspaceContext.j import { ToolErrorType } from './tool-error.js'; import * as glob from 'glob'; import { createMockMessageBus } from '../test-utils/mock-message-bus.js'; +import { execStreaming } from '../utils/shell-utils.js'; vi.mock('glob', { spy: true }); +vi.mock('../utils/shell-utils.js', () => ({ + execStreaming: vi.fn(), +})); // Mock the child_process module to control grep/git grep behavior vi.mock('child_process', () => ({ @@ -129,6 +134,14 @@ describe('GrepTool', () => { }); }); + function createLineGenerator(lines: string[]): AsyncGenerator { + return (async function* () { + for (const line of lines) { + yield line; + } + })(); + } + describe('execute', () => { it('should find matches for a simple pattern in all files', async () => { const params: GrepToolParams = { pattern: 'world' }; @@ -147,6 +160,35 @@ describe('GrepTool', () => { expect(result.returnDisplay).toBe('Found 3 matches'); }, 30000); + it('should include files that start with ".." in JS fallback', async () => { + await fs.writeFile(path.join(tempRootDir, '..env'), 'world in ..env'); + const params: GrepToolParams = { pattern: 'world' }; + const invocation = grepTool.build(params); + const result = await invocation.execute(abortSignal); + expect(result.llmContent).toContain('File: ..env'); + expect(result.llmContent).toContain('L1: world in ..env'); + }); + + it('should ignore system grep output that escapes base path', async () => { + vi.mocked(execStreaming).mockImplementationOnce(() => + createLineGenerator(['..env:1:hello', '../secret.txt:2:leak']), + ); + + const params: GrepToolParams = { pattern: 'hello' }; + const invocation = grepTool.build(params) as unknown as { + isCommandAvailable: (command: string) => Promise; + execute: (signal: AbortSignal) => Promise; + }; + invocation.isCommandAvailable = vi.fn( + async (command: string) => command === 'grep', + ); + + const result = await invocation.execute(abortSignal); + expect(result.llmContent).toContain('File: ..env'); + expect(result.llmContent).toContain('L1: hello'); + expect(result.llmContent).not.toContain('secret.txt'); + }); + it('should find matches in a specific path', async () => { const params: GrepToolParams = { pattern: 'world', dir_path: 'sub' }; const invocation = grepTool.build(params); diff --git a/packages/core/src/tools/grep.ts b/packages/core/src/tools/grep.ts index 3fbbb141d6..ed4fdcb93a 100644 --- a/packages/core/src/tools/grep.ts +++ b/packages/core/src/tools/grep.ts @@ -8,10 +8,14 @@ import type { MessageBus } from '../confirmation-bus/message-bus.js'; import fs from 'node:fs'; import fsPromises from 'node:fs/promises'; import path from 'node:path'; -import { EOL } from 'node:os'; import { spawn } from 'node:child_process'; import { globStream } from 'glob'; import type { ToolInvocation, ToolResult } from './tools.js'; +import { execStreaming } from '../utils/shell-utils.js'; +import { + DEFAULT_TOTAL_MAX_MATCHES, + DEFAULT_SEARCH_TIMEOUT_MS, +} from './constants.js'; import { BaseDeclarativeTool, BaseToolInvocation, Kind } from './tools.js'; import { makeRelative, shortenPath } from '../utils/paths.js'; import { getErrorMessage, isNodeError } from '../utils/errors.js'; @@ -111,6 +115,47 @@ class GrepToolInvocation extends BaseToolInvocation< return targetPath; } + /** + * Parses a single line of grep-like output (git grep, system grep). + * Expects format: filePath:lineNumber:lineContent + * @param {string} line The line to parse. + * @param {string} basePath The absolute directory for path resolution. + * @returns {GrepMatch | null} Parsed match or null if malformed. + */ + private parseGrepLine(line: string, basePath: string): GrepMatch | null { + if (!line.trim()) return null; + + // Use regex to locate the first occurrence of :: + // This allows filenames to contain colons, as long as they don't look like :: + // Note: This regex assumes filenames do not contain colons, or at least not followed by digits. + const match = line.match(/^(.+?):(\d+):(.*)$/); + if (!match) return null; + + const [, filePathRaw, lineNumberStr, lineContent] = match; + const lineNumber = parseInt(lineNumberStr, 10); + + if (!isNaN(lineNumber)) { + const absoluteFilePath = path.resolve(basePath, filePathRaw); + const relativeCheck = path.relative(basePath, absoluteFilePath); + if ( + relativeCheck === '..' || + relativeCheck.startsWith(`..${path.sep}`) || + path.isAbsolute(relativeCheck) + ) { + return null; + } + + const relativeFilePath = path.relative(basePath, absoluteFilePath); + + return { + filePath: relativeFilePath || path.basename(absoluteFilePath), + lineNumber, + line: lineContent, + }; + } + return null; + } + async execute(signal: AbortSignal): Promise { try { const workspaceContext = this.config.getWorkspaceContext(); @@ -129,23 +174,48 @@ class GrepToolInvocation extends BaseToolInvocation< // Collect matches from all search directories let allMatches: GrepMatch[] = []; - for (const searchDir of searchDirectories) { - const matches = await this.performGrepSearch({ - pattern: this.params.pattern, - path: searchDir, - include: this.params.include, - signal, - }); + const totalMaxMatches = DEFAULT_TOTAL_MAX_MATCHES; - // Add directory prefix if searching multiple directories - if (searchDirectories.length > 1) { - const dirName = path.basename(searchDir); - matches.forEach((match) => { - match.filePath = path.join(dirName, match.filePath); + // Create a timeout controller to prevent indefinitely hanging searches + const timeoutController = new AbortController(); + const timeoutId = setTimeout(() => { + timeoutController.abort(); + }, DEFAULT_SEARCH_TIMEOUT_MS); + + // Link the passed signal to our timeout controller + const onAbort = () => timeoutController.abort(); + if (signal.aborted) { + onAbort(); + } else { + signal.addEventListener('abort', onAbort, { once: true }); + } + + try { + for (const searchDir of searchDirectories) { + const remainingLimit = totalMaxMatches - allMatches.length; + if (remainingLimit <= 0) break; + + const matches = await this.performGrepSearch({ + pattern: this.params.pattern, + path: searchDir, + include: this.params.include, + maxMatches: remainingLimit, + signal: timeoutController.signal, }); - } - allMatches = allMatches.concat(matches); + // Add directory prefix if searching multiple directories + if (searchDirectories.length > 1) { + const dirName = path.basename(searchDir); + matches.forEach((match) => { + match.filePath = path.join(dirName, match.filePath); + }); + } + + allMatches = allMatches.concat(matches); + } + } finally { + clearTimeout(timeoutId); + signal.removeEventListener('abort', onAbort); } let searchLocationDescription: string; @@ -164,6 +234,8 @@ class GrepToolInvocation extends BaseToolInvocation< return { llmContent: noMatchMsg, returnDisplay: `No matches found` }; } + const wasTruncated = allMatches.length >= totalMaxMatches; + // Group matches by file const matchesByFile = allMatches.reduce( (acc, match) => { @@ -181,9 +253,7 @@ class GrepToolInvocation extends BaseToolInvocation< const matchCount = allMatches.length; const matchTerm = matchCount === 1 ? 'match' : 'matches'; - let llmContent = `Found ${matchCount} ${matchTerm} for pattern "${this.params.pattern}" ${searchLocationDescription}${this.params.include ? ` (filter: "${this.params.include}")` : ''}: ---- -`; + let llmContent = `Found ${matchCount} ${matchTerm} for pattern "${this.params.pattern}" ${searchLocationDescription}${this.params.include ? ` (filter: "${this.params.include}")` : ''}${wasTruncated ? ` (results limited to ${totalMaxMatches} matches for performance)` : ''}:\n---\n`; for (const filePath in matchesByFile) { llmContent += `File: ${filePath}\n`; @@ -196,7 +266,7 @@ class GrepToolInvocation extends BaseToolInvocation< return { llmContent: llmContent.trim(), - returnDisplay: `Found ${matchCount} ${matchTerm}`, + returnDisplay: `Found ${matchCount} ${matchTerm}${wasTruncated ? ' (limited)' : ''}`, }; } catch (error) { debugLogger.warn(`Error during GrepLogic execution: ${error}`); @@ -241,92 +311,6 @@ class GrepToolInvocation extends BaseToolInvocation< }); } - /** - * Parses the standard output of grep-like commands (git grep, system grep). - * Expects format: filePath:lineNumber:lineContent - * Handles colons within file paths and line content correctly. - * @param {string} output The raw stdout string. - * @param {string} basePath The absolute directory the search was run from, for relative paths. - * @returns {GrepMatch[]} Array of match objects. - */ - private parseGrepOutput(output: string, basePath: string): GrepMatch[] { - const results: GrepMatch[] = []; - if (!output) return results; - - const lines = output.split(EOL); // Use OS-specific end-of-line - - for (const line of lines) { - if (!line.trim()) continue; - - // Find the index of the first colon. - const firstColonIndex = line.indexOf(':'); - if (firstColonIndex === -1) continue; // Malformed - - // Find the index of the second colon, searching *after* the first one. - const secondColonIndex = line.indexOf(':', firstColonIndex + 1); - if (secondColonIndex === -1) continue; // Malformed - - // Extract parts based on the found colon indices - const filePathRaw = line.substring(0, firstColonIndex); - const lineNumberStr = line.substring( - firstColonIndex + 1, - secondColonIndex, - ); - const lineContent = line.substring(secondColonIndex + 1); - - const lineNumber = parseInt(lineNumberStr, 10); - - if (!isNaN(lineNumber)) { - const absoluteFilePath = path.resolve(basePath, filePathRaw); - const relativeFilePath = path.relative(basePath, absoluteFilePath); - - results.push({ - filePath: relativeFilePath || path.basename(absoluteFilePath), - lineNumber, - line: lineContent, - }); - } - } - return results; - } - - /** - * Gets a description of the grep operation - * @returns A string describing the grep - */ - getDescription(): string { - let description = `'${this.params.pattern}'`; - if (this.params.include) { - description += ` in ${this.params.include}`; - } - if (this.params.dir_path) { - const resolvedPath = path.resolve( - this.config.getTargetDir(), - this.params.dir_path, - ); - if ( - resolvedPath === this.config.getTargetDir() || - this.params.dir_path === '.' - ) { - description += ` within ./`; - } else { - const relativePath = makeRelative( - resolvedPath, - this.config.getTargetDir(), - ); - description += ` within ${shortenPath(relativePath)}`; - } - } else { - // When no path is specified, indicate searching all workspace directories - const workspaceContext = this.config.getWorkspaceContext(); - const directories = workspaceContext.getDirectories(); - if (directories.length > 1) { - description += ` across all workspace directories`; - } - } - return description; - } - /** * Performs the actual search using the prioritized strategies. * @param options Search options including pattern, absolute path, and include glob. @@ -336,9 +320,10 @@ class GrepToolInvocation extends BaseToolInvocation< pattern: string; path: string; // Expects absolute path include?: string; + maxMatches: number; signal: AbortSignal; }): Promise { - const { pattern, path: absolutePath, include } = options; + const { pattern, path: absolutePath, include, maxMatches } = options; let strategyUsed = 'none'; try { @@ -361,32 +346,23 @@ class GrepToolInvocation extends BaseToolInvocation< } try { - const output = await new Promise((resolve, reject) => { - const child = spawn('git', gitArgs, { - cwd: absolutePath, - windowsHide: true, - }); - const stdoutChunks: Buffer[] = []; - const stderrChunks: Buffer[] = []; - - child.stdout.on('data', (chunk) => stdoutChunks.push(chunk)); - child.stderr.on('data', (chunk) => stderrChunks.push(chunk)); - child.on('error', (err) => - reject(new Error(`Failed to start git grep: ${err.message}`)), - ); - child.on('close', (code) => { - const stdoutData = Buffer.concat(stdoutChunks).toString('utf8'); - const stderrData = Buffer.concat(stderrChunks).toString('utf8'); - if (code === 0) resolve(stdoutData); - else if (code === 1) - resolve(''); // No matches - else - reject( - new Error(`git grep exited with code ${code}: ${stderrData}`), - ); - }); + const generator = execStreaming('git', gitArgs, { + cwd: absolutePath, + signal: options.signal, + allowedExitCodes: [0, 1], }); - return this.parseGrepOutput(output, absolutePath); + + const results: GrepMatch[] = []; + for await (const line of generator) { + const match = this.parseGrepLine(line, absolutePath); + if (match) { + results.push(match); + if (results.length >= maxMatches) { + break; + } + } + } + return results; } catch (gitError: unknown) { debugLogger.debug( `GrepLogic: git grep failed: ${getErrorMessage( @@ -433,67 +409,31 @@ class GrepToolInvocation extends BaseToolInvocation< grepArgs.push(pattern); grepArgs.push('.'); + const results: GrepMatch[] = []; try { - const output = await new Promise((resolve, reject) => { - const child = spawn('grep', grepArgs, { - cwd: absolutePath, - windowsHide: true, - }); - const stdoutChunks: Buffer[] = []; - const stderrChunks: Buffer[] = []; - - const onData = (chunk: Buffer) => stdoutChunks.push(chunk); - const onStderr = (chunk: Buffer) => { - const stderrStr = chunk.toString(); - // Suppress common harmless stderr messages - if ( - !stderrStr.includes('Permission denied') && - !/grep:.*: Is a directory/i.test(stderrStr) - ) { - stderrChunks.push(chunk); - } - }; - const onError = (err: Error) => { - cleanup(); - reject(new Error(`Failed to start system grep: ${err.message}`)); - }; - const onClose = (code: number | null) => { - const stdoutData = Buffer.concat(stdoutChunks).toString('utf8'); - const stderrData = Buffer.concat(stderrChunks) - .toString('utf8') - .trim(); - cleanup(); - if (code === 0) resolve(stdoutData); - else if (code === 1) - resolve(''); // No matches - else { - if (stderrData) - reject( - new Error( - `System grep exited with code ${code}: ${stderrData}`, - ), - ); - else resolve(''); // Exit code > 1 but no stderr, likely just suppressed errors - } - }; - - const cleanup = () => { - child.stdout.removeListener('data', onData); - child.stderr.removeListener('data', onStderr); - child.removeListener('error', onError); - child.removeListener('close', onClose); - if (child.connected) { - child.disconnect(); - } - }; - - child.stdout.on('data', onData); - child.stderr.on('data', onStderr); - child.on('error', onError); - child.on('close', onClose); + const generator = execStreaming('grep', grepArgs, { + cwd: absolutePath, + signal: options.signal, + allowedExitCodes: [0, 1], }); - return this.parseGrepOutput(output, absolutePath); + + for await (const line of generator) { + const match = this.parseGrepLine(line, absolutePath); + if (match) { + results.push(match); + if (results.length >= maxMatches) { + break; + } + } + } + return results; } catch (grepError: unknown) { + if ( + grepError instanceof Error && + /Permission denied|Is a directory/i.test(grepError.message) + ) { + return results; + } debugLogger.debug( `GrepLogic: System grep failed: ${getErrorMessage( grepError, @@ -523,11 +463,22 @@ class GrepToolInvocation extends BaseToolInvocation< const allMatches: GrepMatch[] = []; for await (const filePath of filesStream) { + if (allMatches.length >= maxMatches) break; const fileAbsolutePath = filePath; + // security check + const relativePath = path.relative(absolutePath, fileAbsolutePath); + if ( + relativePath === '..' || + relativePath.startsWith(`..${path.sep}`) || + path.isAbsolute(relativePath) + ) + continue; + try { const content = await fsPromises.readFile(fileAbsolutePath, 'utf8'); const lines = content.split(/\r?\n/); - lines.forEach((line, index) => { + for (let index = 0; index < lines.length; index++) { + const line = lines[index]; if (regex.test(line)) { allMatches.push({ filePath: @@ -536,8 +487,9 @@ class GrepToolInvocation extends BaseToolInvocation< lineNumber: index + 1, line, }); + if (allMatches.length >= maxMatches) break; } - }); + } } catch (readError: unknown) { // Ignore errors like permission denied or file gone during read if (!isNodeError(readError) || readError.code !== 'ENOENT') { @@ -560,9 +512,40 @@ class GrepToolInvocation extends BaseToolInvocation< throw error; // Re-throw } } -} -// --- GrepLogic Class --- + getDescription(): string { + let description = `'${this.params.pattern}'`; + if (this.params.include) { + description += ` in ${this.params.include}`; + } + if (this.params.dir_path) { + const resolvedPath = path.resolve( + this.config.getTargetDir(), + this.params.dir_path, + ); + if ( + resolvedPath === this.config.getTargetDir() || + this.params.dir_path === '.' + ) { + description += ` within ./`; + } else { + const relativePath = makeRelative( + resolvedPath, + this.config.getTargetDir(), + ); + description += ` within ${shortenPath(relativePath)}`; + } + } else { + // When no path is specified, indicate searching all workspace directories + const workspaceContext = this.config.getWorkspaceContext(); + const directories = workspaceContext.getDirectories(); + if (directories.length > 1) { + description += ` across all workspace directories`; + } + } + return description; + } +} /** * Implementation of the Grep tool logic (moved from CLI) @@ -581,8 +564,7 @@ export class GrepTool extends BaseDeclarativeTool { { properties: { pattern: { - description: - "The regular expression (regex) pattern to search for within file contents (e.g., 'function\\s+myFunction', 'import\\s+\\{.*\\}\\s+from\\s+.*').", + description: `The regular expression (regex) pattern to search for within file contents (e.g., 'function\\s+myFunction', 'import\\s+\\{.*\\}\\s+from\\s+.*').`, type: 'string', }, dir_path: { @@ -591,8 +573,7 @@ export class GrepTool extends BaseDeclarativeTool { type: 'string', }, include: { - description: - "Optional: A glob pattern to filter which files are searched (e.g., '*.js', '*.{ts,tsx}', 'src/**'). If omitted, searches all files (respecting potential global ignores).", + description: `Optional: A glob pattern to filter which files are searched (e.g., '*.js', '*.{ts,tsx}', 'src/**'). If omitted, searches all files (respecting potential global ignores).`, type: 'string', }, }, diff --git a/packages/core/src/tools/ripGrep.test.ts b/packages/core/src/tools/ripGrep.test.ts index 415db097e3..4c27dde5b1 100644 --- a/packages/core/src/tools/ripGrep.test.ts +++ b/packages/core/src/tools/ripGrep.test.ts @@ -23,6 +23,8 @@ import { Storage } from '../config/storage.js'; import { createMockWorkspaceContext } from '../test-utils/mockWorkspaceContext.js'; import type { ChildProcess } from 'node:child_process'; import { spawn } from 'node:child_process'; +import { PassThrough, Readable } from 'node:stream'; +import EventEmitter from 'node:events'; import { downloadRipGrep } from '@joshua.litt/get-ripgrep'; import { createMockMessageBus } from '../test-utils/mock-message-bus.js'; // Mock dependencies for canUseRipgrep @@ -197,47 +199,43 @@ describe('ensureRgPath', () => { function createMockSpawn( options: { outputData?: string; - exitCode?: number; + exitCode?: number | null; signal?: string; } = {}, ) { const { outputData, exitCode = 0, signal } = options; return () => { - const mockProcess = { - stdout: { - on: vi.fn(), - removeListener: vi.fn(), + // strict Readable implementation + let pushed = false; + const stdout = new Readable({ + read() { + if (!pushed) { + if (outputData) { + this.push(outputData); + } + this.push(null); // EOF + pushed = true; + } }, - stderr: { - on: vi.fn(), - removeListener: vi.fn(), - }, - on: vi.fn(), - removeListener: vi.fn(), - kill: vi.fn(), - }; + }); - // Set up event listeners immediately + const stderr = new PassThrough(); + const mockProcess = new EventEmitter() as ChildProcess; + mockProcess.stdout = stdout as unknown as Readable; + mockProcess.stderr = stderr; + mockProcess.kill = vi.fn(); + // @ts-expect-error - mocking private/internal property + mockProcess.killed = false; + // @ts-expect-error - mocking private/internal property + mockProcess.exitCode = null; + + // Emulating process exit setTimeout(() => { - const stdoutDataHandler = mockProcess.stdout.on.mock.calls.find( - (call) => call[0] === 'data', - )?.[1]; + mockProcess.emit('close', exitCode, signal); + }, 10); - const closeHandler = mockProcess.on.mock.calls.find( - (call) => call[0] === 'close', - )?.[1]; - - if (stdoutDataHandler && outputData) { - stdoutDataHandler(Buffer.from(outputData)); - } - - if (closeHandler) { - closeHandler(exitCode, signal); - } - }, 0); - - return mockProcess as unknown as ChildProcess; + return mockProcess; }; } @@ -406,6 +404,40 @@ describe('RipGrepTool', () => { expect(result.returnDisplay).toBe('Found 3 matches'); }); + it('should ignore matches that escape the base path', async () => { + mockSpawn.mockImplementationOnce( + createMockSpawn({ + outputData: + JSON.stringify({ + type: 'match', + data: { + path: { text: '..env' }, + line_number: 1, + lines: { text: 'world in ..env\n' }, + }, + }) + + '\n' + + JSON.stringify({ + type: 'match', + data: { + path: { text: '../secret.txt' }, + line_number: 1, + lines: { text: 'leak\n' }, + }, + }) + + '\n', + exitCode: 0, + }), + ); + + const params: RipGrepToolParams = { pattern: 'world' }; + const invocation = grepTool.build(params); + const result = await invocation.execute(abortSignal); + expect(result.llmContent).toContain('File: ..env'); + expect(result.llmContent).toContain('L1: world in ..env'); + expect(result.llmContent).not.toContain('secret.txt'); + }); + it('should find matches in a specific path', async () => { // Setup specific mock for this test - searching in 'sub' should only return matches from that directory mockSpawn.mockImplementationOnce( @@ -471,51 +503,20 @@ describe('RipGrepTool', () => { ); // Setup specific mock for this test - searching for 'hello' in 'sub' with '*.js' filter - mockSpawn.mockImplementationOnce(() => { - const mockProcess = { - stdout: { - on: vi.fn(), - removeListener: vi.fn(), - }, - stderr: { - on: vi.fn(), - removeListener: vi.fn(), - }, - on: vi.fn(), - removeListener: vi.fn(), - kill: vi.fn(), - }; - - setTimeout(() => { - const onData = mockProcess.stdout.on.mock.calls.find( - (call) => call[0] === 'data', - )?.[1]; - const onClose = mockProcess.on.mock.calls.find( - (call) => call[0] === 'close', - )?.[1]; - - if (onData) { - // Only return match from the .js file in sub directory - onData( - Buffer.from( - JSON.stringify({ - type: 'match', - data: { - path: { text: 'another.js' }, - line_number: 1, - lines: { text: 'const greeting = "hello";\n' }, - }, - }) + '\n', - ), - ); - } - if (onClose) { - onClose(0); - } - }, 0); - - return mockProcess as unknown as ChildProcess; - }); + mockSpawn.mockImplementationOnce( + createMockSpawn({ + outputData: + JSON.stringify({ + type: 'match', + data: { + path: { text: 'another.js' }, + line_number: 1, + lines: { text: 'const greeting = "hello";\n' }, + }, + }) + '\n', + exitCode: 0, + }), + ); const params: RipGrepToolParams = { pattern: 'hello', @@ -559,59 +560,114 @@ describe('RipGrepTool', () => { const params: RipGrepToolParams = { pattern: '[[' }; const invocation = grepTool.build(params); const result = await invocation.execute(abortSignal); - expect(result.llmContent).toContain('ripgrep exited with code 2'); + expect(result.llmContent).toContain('Process exited with code 2'); expect(result.returnDisplay).toContain( - 'Error: ripgrep exited with code 2', + 'Error: Process exited with code 2', ); }); + it('should ignore invalid regex error from ripgrep when it is not a user error', async () => { + mockSpawn.mockImplementation( + createMockSpawn({ + outputData: '', + exitCode: 2, + signal: undefined, + }), + ); + + const invocation = grepTool.build({ + pattern: 'foo', + dir_path: tempRootDir, + }); + + const result = await invocation.execute(abortSignal); + expect(result.llmContent).toContain('Process exited with code 2'); + expect(result.returnDisplay).toContain( + 'Error: Process exited with code 2', + ); + }); + + it('should handle massive output by terminating early without crashing (Regression)', async () => { + const massiveOutputLines = 30000; + + // Custom mock for massive streaming + mockSpawn.mockImplementation(() => { + const stdout = new PassThrough(); + const stderr = new PassThrough(); + const mockProcess = new EventEmitter() as ChildProcess; + mockProcess.stdout = stdout; + mockProcess.stderr = stderr; + mockProcess.kill = vi.fn(); + // @ts-expect-error - mocking private/internal property + mockProcess.killed = false; + // @ts-expect-error - mocking private/internal property + mockProcess.exitCode = null; + + // Push data over time + let linesPushed = 0; + const pushInterval = setInterval(() => { + if (linesPushed >= massiveOutputLines) { + clearInterval(pushInterval); + stdout.end(); + mockProcess.emit('close', 0); + return; + } + + // Push a batch + try { + for (let i = 0; i < 2000 && linesPushed < massiveOutputLines; i++) { + const match = JSON.stringify({ + type: 'match', + data: { + path: { text: `file_${linesPushed}.txt` }, + line_number: 1, + lines: { text: `match ${linesPushed}\n` }, + }, + }); + stdout.write(match + '\n'); + linesPushed++; + } + } catch (_e) { + clearInterval(pushInterval); + } + }, 1); + + mockProcess.kill = vi.fn().mockImplementation(() => { + clearInterval(pushInterval); + stdout.end(); + // Emit close async to allow listeners to attach + setTimeout(() => mockProcess.emit('close', 0, 'SIGTERM'), 0); + return true; + }); + + return mockProcess; + }); + + const invocation = grepTool.build({ + pattern: 'test', + dir_path: tempRootDir, + }); + const result = await invocation.execute(abortSignal); + + expect(result.returnDisplay).toContain('(limited)'); + }, 10000); + it('should handle regex special characters correctly', async () => { // Setup specific mock for this test - regex pattern 'foo.*bar' should match 'const foo = "bar";' - mockSpawn.mockImplementationOnce(() => { - const mockProcess = { - stdout: { - on: vi.fn(), - removeListener: vi.fn(), - }, - stderr: { - on: vi.fn(), - removeListener: vi.fn(), - }, - on: vi.fn(), - removeListener: vi.fn(), - kill: vi.fn(), - }; - - setTimeout(() => { - const onData = mockProcess.stdout.on.mock.calls.find( - (call) => call[0] === 'data', - )?.[1]; - const onClose = mockProcess.on.mock.calls.find( - (call) => call[0] === 'close', - )?.[1]; - - if (onData) { - // Return match for the regex pattern - onData( - Buffer.from( - JSON.stringify({ - type: 'match', - data: { - path: { text: 'fileB.js' }, - line_number: 1, - lines: { text: 'const foo = "bar";\n' }, - }, - }) + '\n', - ), - ); - } - if (onClose) { - onClose(0); - } - }, 0); - - return mockProcess as unknown as ChildProcess; - }); + mockSpawn.mockImplementationOnce( + createMockSpawn({ + outputData: + JSON.stringify({ + type: 'match', + data: { + path: { text: 'fileB.js' }, + line_number: 1, + lines: { text: 'const foo = "bar";\n' }, + }, + }) + '\n', + exitCode: 0, + }), + ); const params: RipGrepToolParams = { pattern: 'foo.*bar' }; // Matches 'const foo = "bar";' const invocation = grepTool.build(params); @@ -625,61 +681,30 @@ describe('RipGrepTool', () => { it('should be case-insensitive by default (JS fallback)', async () => { // Setup specific mock for this test - case insensitive search for 'HELLO' - mockSpawn.mockImplementationOnce(() => { - const mockProcess = { - stdout: { - on: vi.fn(), - removeListener: vi.fn(), - }, - stderr: { - on: vi.fn(), - removeListener: vi.fn(), - }, - on: vi.fn(), - removeListener: vi.fn(), - kill: vi.fn(), - }; - - setTimeout(() => { - const onData = mockProcess.stdout.on.mock.calls.find( - (call) => call[0] === 'data', - )?.[1]; - const onClose = mockProcess.on.mock.calls.find( - (call) => call[0] === 'close', - )?.[1]; - - if (onData) { - // Return case-insensitive matches for 'HELLO' - onData( - Buffer.from( - JSON.stringify({ - type: 'match', - data: { - path: { text: 'fileA.txt' }, - line_number: 1, - lines: { text: 'hello world\n' }, - }, - }) + - '\n' + - JSON.stringify({ - type: 'match', - data: { - path: { text: 'fileB.js' }, - line_number: 2, - lines: { text: 'function baz() { return "hello"; }\n' }, - }, - }) + - '\n', - ), - ); - } - if (onClose) { - onClose(0); - } - }, 0); - - return mockProcess as unknown as ChildProcess; - }); + mockSpawn.mockImplementationOnce( + createMockSpawn({ + outputData: + JSON.stringify({ + type: 'match', + data: { + path: { text: 'fileA.txt' }, + line_number: 1, + lines: { text: 'hello world\n' }, + }, + }) + + '\n' + + JSON.stringify({ + type: 'match', + data: { + path: { text: 'fileB.js' }, + line_number: 2, + lines: { text: 'function baz() { return "hello"; }\n' }, + }, + }) + + '\n', + exitCode: 0, + }), + ); const params: RipGrepToolParams = { pattern: 'HELLO' }; const invocation = grepTool.build(params); @@ -742,97 +767,39 @@ describe('RipGrepTool', () => { // Setup specific mock for this test - multi-directory search for 'world' // Mock will be called twice - once for each directory - let callCount = 0; - mockSpawn.mockImplementation(() => { - callCount++; - const mockProcess = { - stdout: { - on: vi.fn(), - removeListener: vi.fn(), - }, - stderr: { - on: vi.fn(), - removeListener: vi.fn(), - }, - on: vi.fn(), - removeListener: vi.fn(), - kill: vi.fn(), - }; - setTimeout(() => { - const stdoutDataHandler = mockProcess.stdout.on.mock.calls.find( - (call) => call[0] === 'data', - )?.[1]; - - const closeHandler = mockProcess.on.mock.calls.find( - (call) => call[0] === 'close', - )?.[1]; - - let outputData = ''; - if (callCount === 1) { - // First directory (tempRootDir) - outputData = - JSON.stringify({ - type: 'match', - data: { - path: { text: 'fileA.txt' }, - line_number: 1, - lines: { text: 'hello world\n' }, - }, - }) + - '\n' + - JSON.stringify({ - type: 'match', - data: { - path: { text: 'fileA.txt' }, - line_number: 2, - lines: { text: 'second line with world\n' }, - }, - }) + - '\n' + - JSON.stringify({ - type: 'match', - data: { - path: { text: 'sub/fileC.txt' }, - line_number: 1, - lines: { text: 'another world in sub dir\n' }, - }, - }) + - '\n'; - } else if (callCount === 2) { - // Second directory (secondDir) - outputData = - JSON.stringify({ - type: 'match', - data: { - path: { text: 'other.txt' }, - line_number: 2, - lines: { text: 'world in second\n' }, - }, - }) + - '\n' + - JSON.stringify({ - type: 'match', - data: { - path: { text: 'another.js' }, - line_number: 1, - lines: { text: 'function world() { return "test"; }\n' }, - }, - }) + - '\n'; - } - - if (stdoutDataHandler && outputData) { - stdoutDataHandler(Buffer.from(outputData)); - } - - if (closeHandler) { - closeHandler(0); - } - }, 0); - - return mockProcess as unknown as ChildProcess; - }); + mockSpawn.mockImplementationOnce( + createMockSpawn({ + outputData: + JSON.stringify({ + type: 'match', + data: { + path: { text: 'fileA.txt' }, + line_number: 1, + lines: { text: 'hello world\n' }, + }, + }) + + '\n' + + JSON.stringify({ + type: 'match', + data: { + path: { text: 'fileA.txt' }, + line_number: 2, + lines: { text: 'second line with world\n' }, + }, + }) + + '\n' + + JSON.stringify({ + type: 'match', + data: { + path: { text: 'sub/fileC.txt' }, + line_number: 1, + lines: { text: 'another world in sub dir\n' }, + }, + }) + + '\n', + }), + ); const multiDirGrepTool = new RipGrepTool( multiDirConfig, @@ -886,50 +853,19 @@ describe('RipGrepTool', () => { } as unknown as Config; // Setup specific mock for this test - searching in 'sub' should only return matches from that directory - mockSpawn.mockImplementationOnce(() => { - const mockProcess = { - stdout: { - on: vi.fn(), - removeListener: vi.fn(), - }, - stderr: { - on: vi.fn(), - removeListener: vi.fn(), - }, - on: vi.fn(), - removeListener: vi.fn(), - kill: vi.fn(), - }; - - setTimeout(() => { - const onData = mockProcess.stdout.on.mock.calls.find( - (call) => call[0] === 'data', - )?.[1]; - const onClose = mockProcess.on.mock.calls.find( - (call) => call[0] === 'close', - )?.[1]; - - if (onData) { - onData( - Buffer.from( - JSON.stringify({ - type: 'match', - data: { - path: { text: 'fileC.txt' }, - line_number: 1, - lines: { text: 'another world in sub dir\n' }, - }, - }) + '\n', - ), - ); - } - if (onClose) { - onClose(0); - } - }, 0); - - return mockProcess as unknown as ChildProcess; - }); + mockSpawn.mockImplementationOnce( + createMockSpawn({ + outputData: + JSON.stringify({ + type: 'match', + data: { + path: { text: 'fileC.txt' }, + line_number: 1, + lines: { text: 'another world in sub dir\n' }, + }, + }) + '\n', + }), + ); const multiDirGrepTool = new RipGrepTool( multiDirConfig, @@ -970,35 +906,12 @@ describe('RipGrepTool', () => { it('should abort streaming search when signal is triggered', async () => { // Setup specific mock for this test - simulate process being killed due to abort - mockSpawn.mockImplementationOnce(() => { - const mockProcess = { - stdout: { - on: vi.fn(), - removeListener: vi.fn(), - }, - stderr: { - on: vi.fn(), - removeListener: vi.fn(), - }, - on: vi.fn(), - removeListener: vi.fn(), - kill: vi.fn(), - }; - - // Simulate process being aborted - use setTimeout to ensure handlers are registered first - setTimeout(() => { - const closeHandler = mockProcess.on.mock.calls.find( - (call) => call[0] === 'close', - )?.[1]; - - if (closeHandler) { - // Simulate process killed by signal (code is null, signal is SIGTERM) - closeHandler(null, 'SIGTERM'); - } - }, 0); - - return mockProcess as unknown as ChildProcess; - }); + mockSpawn.mockImplementationOnce( + createMockSpawn({ + exitCode: null, + signal: 'SIGTERM', + }), + ); const controller = new AbortController(); const params: RipGrepToolParams = { pattern: 'test' }; @@ -1008,12 +921,7 @@ describe('RipGrepTool', () => { controller.abort(); const result = await invocation.execute(controller.signal); - expect(result.llmContent).toContain( - 'Error during grep search operation: ripgrep was terminated by signal:', - ); - expect(result.returnDisplay).toContain( - 'Error: ripgrep was terminated by signal:', - ); + expect(result.returnDisplay).toContain('No matches found'); }); }); @@ -1060,50 +968,19 @@ describe('RipGrepTool', () => { ); // Setup specific mock for this test - searching for 'world' should find the file with special characters - mockSpawn.mockImplementationOnce(() => { - const mockProcess = { - stdout: { - on: vi.fn(), - removeListener: vi.fn(), - }, - stderr: { - on: vi.fn(), - removeListener: vi.fn(), - }, - on: vi.fn(), - removeListener: vi.fn(), - kill: vi.fn(), - }; - - setTimeout(() => { - const onData = mockProcess.stdout.on.mock.calls.find( - (call) => call[0] === 'data', - )?.[1]; - const onClose = mockProcess.on.mock.calls.find( - (call) => call[0] === 'close', - )?.[1]; - - if (onData) { - onData( - Buffer.from( - JSON.stringify({ - type: 'match', - data: { - path: { text: specialFileName }, - line_number: 1, - lines: { text: 'hello world with special chars\n' }, - }, - }) + '\n', - ), - ); - } - if (onClose) { - onClose(0); - } - }, 0); - - return mockProcess as unknown as ChildProcess; - }); + mockSpawn.mockImplementationOnce( + createMockSpawn({ + outputData: + JSON.stringify({ + type: 'match', + data: { + path: { text: specialFileName }, + line_number: 1, + lines: { text: 'hello world with special chars\n' }, + }, + }) + '\n', + }), + ); const params: RipGrepToolParams = { pattern: 'world' }; const invocation = grepTool.build(params); @@ -1122,50 +999,19 @@ describe('RipGrepTool', () => { ); // Setup specific mock for this test - searching for 'deep' should find the deeply nested file - mockSpawn.mockImplementationOnce(() => { - const mockProcess = { - stdout: { - on: vi.fn(), - removeListener: vi.fn(), - }, - stderr: { - on: vi.fn(), - removeListener: vi.fn(), - }, - on: vi.fn(), - removeListener: vi.fn(), - kill: vi.fn(), - }; - - setTimeout(() => { - const onData = mockProcess.stdout.on.mock.calls.find( - (call) => call[0] === 'data', - )?.[1]; - const onClose = mockProcess.on.mock.calls.find( - (call) => call[0] === 'close', - )?.[1]; - - if (onData) { - onData( - Buffer.from( - JSON.stringify({ - type: 'match', - data: { - path: { text: 'a/b/c/d/e/deep.txt' }, - line_number: 1, - lines: { text: 'content in deep directory\n' }, - }, - }) + '\n', - ), - ); - } - if (onClose) { - onClose(0); - } - }, 0); - - return mockProcess as unknown as ChildProcess; - }); + mockSpawn.mockImplementationOnce( + createMockSpawn({ + outputData: + JSON.stringify({ + type: 'match', + data: { + path: { text: 'a/b/c/d/e/deep.txt' }, + line_number: 1, + lines: { text: 'content in deep directory\n' }, + }, + }) + '\n', + }), + ); const params: RipGrepToolParams = { pattern: 'deep' }; const invocation = grepTool.build(params); @@ -1184,50 +1030,19 @@ describe('RipGrepTool', () => { ); // Setup specific mock for this test - regex pattern should match function declarations - mockSpawn.mockImplementationOnce(() => { - const mockProcess = { - stdout: { - on: vi.fn(), - removeListener: vi.fn(), - }, - stderr: { - on: vi.fn(), - removeListener: vi.fn(), - }, - on: vi.fn(), - removeListener: vi.fn(), - kill: vi.fn(), - }; - - setTimeout(() => { - const onData = mockProcess.stdout.on.mock.calls.find( - (call) => call[0] === 'data', - )?.[1]; - const onClose = mockProcess.on.mock.calls.find( - (call) => call[0] === 'close', - )?.[1]; - - if (onData) { - onData( - Buffer.from( - JSON.stringify({ - type: 'match', - data: { - path: { text: 'code.js' }, - line_number: 1, - lines: { text: 'function getName() { return "test"; }\n' }, - }, - }) + '\n', - ), - ); - } - if (onClose) { - onClose(0); - } - }, 0); - - return mockProcess as unknown as ChildProcess; - }); + mockSpawn.mockImplementationOnce( + createMockSpawn({ + outputData: + JSON.stringify({ + type: 'match', + data: { + path: { text: 'code.js' }, + line_number: 1, + lines: { text: 'function getName() { return "test"; }\n' }, + }, + }) + '\n', + }), + ); const params: RipGrepToolParams = { pattern: 'function\\s+\\w+\\s*\\(' }; const invocation = grepTool.build(params); @@ -1244,69 +1059,38 @@ describe('RipGrepTool', () => { ); // Setup specific mock for this test - case insensitive search should match all variants - mockSpawn.mockImplementationOnce(() => { - const mockProcess = { - stdout: { - on: vi.fn(), - removeListener: vi.fn(), - }, - stderr: { - on: vi.fn(), - removeListener: vi.fn(), - }, - on: vi.fn(), - removeListener: vi.fn(), - kill: vi.fn(), - }; - - setTimeout(() => { - const onData = mockProcess.stdout.on.mock.calls.find( - (call) => call[0] === 'data', - )?.[1]; - const onClose = mockProcess.on.mock.calls.find( - (call) => call[0] === 'close', - )?.[1]; - - if (onData) { - onData( - Buffer.from( - JSON.stringify({ - type: 'match', - data: { - path: { text: 'case.txt' }, - line_number: 1, - lines: { text: 'Hello World\n' }, - }, - }) + - '\n' + - JSON.stringify({ - type: 'match', - data: { - path: { text: 'case.txt' }, - line_number: 2, - lines: { text: 'hello world\n' }, - }, - }) + - '\n' + - JSON.stringify({ - type: 'match', - data: { - path: { text: 'case.txt' }, - line_number: 3, - lines: { text: 'HELLO WORLD\n' }, - }, - }) + - '\n', - ), - ); - } - if (onClose) { - onClose(0); - } - }, 0); - - return mockProcess as unknown as ChildProcess; - }); + mockSpawn.mockImplementationOnce( + createMockSpawn({ + outputData: + JSON.stringify({ + type: 'match', + data: { + path: { text: 'case.txt' }, + line_number: 1, + lines: { text: 'Hello World\n' }, + }, + }) + + '\n' + + JSON.stringify({ + type: 'match', + data: { + path: { text: 'case.txt' }, + line_number: 2, + lines: { text: 'hello world\n' }, + }, + }) + + '\n' + + JSON.stringify({ + type: 'match', + data: { + path: { text: 'case.txt' }, + line_number: 3, + lines: { text: 'HELLO WORLD\n' }, + }, + }) + + '\n', + }), + ); const params: RipGrepToolParams = { pattern: 'hello' }; const invocation = grepTool.build(params); @@ -1324,50 +1108,19 @@ describe('RipGrepTool', () => { ); // Setup specific mock for this test - escaped regex pattern should match price format - mockSpawn.mockImplementationOnce(() => { - const mockProcess = { - stdout: { - on: vi.fn(), - removeListener: vi.fn(), - }, - stderr: { - on: vi.fn(), - removeListener: vi.fn(), - }, - on: vi.fn(), - removeListener: vi.fn(), - kill: vi.fn(), - }; - - setTimeout(() => { - const onData = mockProcess.stdout.on.mock.calls.find( - (call) => call[0] === 'data', - )?.[1]; - const onClose = mockProcess.on.mock.calls.find( - (call) => call[0] === 'close', - )?.[1]; - - if (onData) { - onData( - Buffer.from( - JSON.stringify({ - type: 'match', - data: { - path: { text: 'special.txt' }, - line_number: 1, - lines: { text: 'Price: $19.99\n' }, - }, - }) + '\n', - ), - ); - } - if (onClose) { - onClose(0); - } - }, 0); - - return mockProcess as unknown as ChildProcess; - }); + mockSpawn.mockImplementationOnce( + createMockSpawn({ + outputData: + JSON.stringify({ + type: 'match', + data: { + path: { text: 'special.txt' }, + line_number: 1, + lines: { text: 'Price: $19.99\n' }, + }, + }) + '\n', + }), + ); const params: RipGrepToolParams = { pattern: '\\$\\d+\\.\\d+' }; const invocation = grepTool.build(params); @@ -1392,60 +1145,29 @@ describe('RipGrepTool', () => { await fs.writeFile(path.join(tempRootDir, 'test.txt'), 'text content'); // Setup specific mock for this test - include pattern should filter to only ts/tsx files - mockSpawn.mockImplementationOnce(() => { - const mockProcess = { - stdout: { - on: vi.fn(), - removeListener: vi.fn(), - }, - stderr: { - on: vi.fn(), - removeListener: vi.fn(), - }, - on: vi.fn(), - removeListener: vi.fn(), - kill: vi.fn(), - }; - - setTimeout(() => { - const onData = mockProcess.stdout.on.mock.calls.find( - (call) => call[0] === 'data', - )?.[1]; - const onClose = mockProcess.on.mock.calls.find( - (call) => call[0] === 'close', - )?.[1]; - - if (onData) { - onData( - Buffer.from( - JSON.stringify({ - type: 'match', - data: { - path: { text: 'test.ts' }, - line_number: 1, - lines: { text: 'typescript content\n' }, - }, - }) + - '\n' + - JSON.stringify({ - type: 'match', - data: { - path: { text: 'test.tsx' }, - line_number: 1, - lines: { text: 'tsx content\n' }, - }, - }) + - '\n', - ), - ); - } - if (onClose) { - onClose(0); - } - }, 0); - - return mockProcess as unknown as ChildProcess; - }); + mockSpawn.mockImplementationOnce( + createMockSpawn({ + outputData: + JSON.stringify({ + type: 'match', + data: { + path: { text: 'test.ts' }, + line_number: 1, + lines: { text: 'typescript content\n' }, + }, + }) + + '\n' + + JSON.stringify({ + type: 'match', + data: { + path: { text: 'test.tsx' }, + line_number: 1, + lines: { text: 'tsx content\n' }, + }, + }) + + '\n', + }), + ); const params: RipGrepToolParams = { pattern: 'content', @@ -1469,50 +1191,19 @@ describe('RipGrepTool', () => { await fs.writeFile(path.join(tempRootDir, 'other.ts'), 'other code'); // Setup specific mock for this test - include pattern should filter to only src/** files - mockSpawn.mockImplementationOnce(() => { - const mockProcess = { - stdout: { - on: vi.fn(), - removeListener: vi.fn(), - }, - stderr: { - on: vi.fn(), - removeListener: vi.fn(), - }, - on: vi.fn(), - removeListener: vi.fn(), - kill: vi.fn(), - }; - - setTimeout(() => { - const onData = mockProcess.stdout.on.mock.calls.find( - (call) => call[0] === 'data', - )?.[1]; - const onClose = mockProcess.on.mock.calls.find( - (call) => call[0] === 'close', - )?.[1]; - - if (onData) { - onData( - Buffer.from( - JSON.stringify({ - type: 'match', - data: { - path: { text: 'src/main.ts' }, - line_number: 1, - lines: { text: 'source code\n' }, - }, - }) + '\n', - ), - ); - } - if (onClose) { - onClose(0); - } - }, 0); - - return mockProcess as unknown as ChildProcess; - }); + mockSpawn.mockImplementationOnce( + createMockSpawn({ + outputData: + JSON.stringify({ + type: 'match', + data: { + path: { text: 'src/main.ts' }, + line_number: 1, + lines: { text: 'source code\n' }, + }, + }) + '\n', + }), + ); const params: RipGrepToolParams = { pattern: 'code', diff --git a/packages/core/src/tools/ripGrep.ts b/packages/core/src/tools/ripGrep.ts index 12f6d720e2..c4642ca20e 100644 --- a/packages/core/src/tools/ripGrep.ts +++ b/packages/core/src/tools/ripGrep.ts @@ -7,7 +7,6 @@ import type { MessageBus } from '../confirmation-bus/message-bus.js'; import fs from 'node:fs'; import path from 'node:path'; -import { spawn } from 'node:child_process'; import { downloadRipGrep } from '@joshua.litt/get-ripgrep'; import type { ToolInvocation, ToolResult } from './tools.js'; import { BaseDeclarativeTool, BaseToolInvocation, Kind } from './tools.js'; @@ -24,8 +23,11 @@ import { COMMON_DIRECTORY_EXCLUDES, } from '../utils/ignorePatterns.js'; import { GeminiIgnoreParser } from '../utils/geminiIgnoreParser.js'; - -const DEFAULT_TOTAL_MAX_MATCHES = 20000; +import { execStreaming } from '../utils/shell-utils.js'; +import { + DEFAULT_TOTAL_MAX_MATCHES, + DEFAULT_SEARCH_TIMEOUT_MS, +} from './constants.js'; function getRgCandidateFilenames(): readonly string[] { return process.platform === 'win32' ? ['rg.exe', 'rg'] : ['rg']; @@ -213,21 +215,38 @@ class GrepToolInvocation extends BaseToolInvocation< debugLogger.log(`[GrepTool] Total result limit: ${totalMaxMatches}`); } - let allMatches = await this.performRipgrepSearch({ - pattern: this.params.pattern, - path: searchDirAbs!, - include: this.params.include, - case_sensitive: this.params.case_sensitive, - fixed_strings: this.params.fixed_strings, - context: this.params.context, - after: this.params.after, - before: this.params.before, - no_ignore: this.params.no_ignore, - signal, - }); + // Create a timeout controller to prevent indefinitely hanging searches + const timeoutController = new AbortController(); + const timeoutId = setTimeout(() => { + timeoutController.abort(); + }, DEFAULT_SEARCH_TIMEOUT_MS); - if (allMatches.length >= totalMaxMatches) { - allMatches = allMatches.slice(0, totalMaxMatches); + // Link the passed signal to our timeout controller + const onAbort = () => timeoutController.abort(); + if (signal.aborted) { + onAbort(); + } else { + signal.addEventListener('abort', onAbort, { once: true }); + } + + let allMatches: GrepMatch[]; + try { + allMatches = await this.performRipgrepSearch({ + pattern: this.params.pattern, + path: searchDirAbs!, + include: this.params.include, + case_sensitive: this.params.case_sensitive, + fixed_strings: this.params.fixed_strings, + context: this.params.context, + after: this.params.after, + before: this.params.before, + no_ignore: this.params.no_ignore, + maxMatches: totalMaxMatches, + signal: timeoutController.signal, + }); + } finally { + clearTimeout(timeoutId); + signal.removeEventListener('abort', onAbort); } const searchLocationDescription = `in path "${searchDirDisplay}"`; @@ -254,13 +273,7 @@ class GrepToolInvocation extends BaseToolInvocation< const matchCount = allMatches.length; const matchTerm = matchCount === 1 ? 'match' : 'matches'; - let llmContent = `Found ${matchCount} ${matchTerm} for pattern "${this.params.pattern}" ${searchLocationDescription}${this.params.include ? ` (filter: "${this.params.include}")` : ''}`; - - if (wasTruncated) { - llmContent += ` (results limited to ${totalMaxMatches} matches for performance)`; - } - - llmContent += `:\n---\n`; + let llmContent = `Found ${matchCount} ${matchTerm} for pattern "${this.params.pattern}" ${searchLocationDescription}${this.params.include ? ` (filter: "${this.params.include}")` : ''}${wasTruncated ? ` (results limited to ${totalMaxMatches} matches for performance)` : ''}:\n---\n`; for (const filePath in matchesByFile) { llmContent += `File: ${filePath}\n`; @@ -271,14 +284,11 @@ class GrepToolInvocation extends BaseToolInvocation< llmContent += '---\n'; } - let displayMessage = `Found ${matchCount} ${matchTerm}`; - if (wasTruncated) { - displayMessage += ` (limited)`; - } - return { llmContent: llmContent.trim(), - returnDisplay: displayMessage, + returnDisplay: `Found ${matchCount} ${matchTerm}${ + wasTruncated ? ' (limited)' : '' + }`, }; } catch (error) { debugLogger.warn(`Error during GrepLogic execution: ${error}`); @@ -290,41 +300,6 @@ class GrepToolInvocation extends BaseToolInvocation< } } - private parseRipgrepJsonOutput( - output: string, - basePath: string, - ): GrepMatch[] { - const results: GrepMatch[] = []; - if (!output) return results; - - const lines = output.trim().split('\n'); - - for (const line of lines) { - if (!line.trim()) continue; - - try { - const json = JSON.parse(line); - if (json.type === 'match') { - const match = json.data; - // Defensive check: ensure text properties exist (skips binary/invalid encoding) - if (match.path?.text && match.lines?.text) { - const absoluteFilePath = path.resolve(basePath, match.path.text); - const relativeFilePath = path.relative(basePath, absoluteFilePath); - - results.push({ - filePath: relativeFilePath || path.basename(absoluteFilePath), - lineNumber: match.line_number, - line: match.lines.text.trimEnd(), - }); - } - } - } catch (error) { - debugLogger.warn(`Failed to parse ripgrep JSON line: ${line}`, error); - } - } - return results; - } - private async performRipgrepSearch(options: { pattern: string; path: string; @@ -335,6 +310,7 @@ class GrepToolInvocation extends BaseToolInvocation< after?: number; before?: number; no_ignore?: boolean; + maxMatches: number; signal: AbortSignal; }): Promise { const { @@ -347,6 +323,7 @@ class GrepToolInvocation extends BaseToolInvocation< after, before, no_ignore, + maxMatches, } = options; const rgArgs = ['--json']; @@ -402,64 +379,72 @@ class GrepToolInvocation extends BaseToolInvocation< rgArgs.push('--threads', '4'); rgArgs.push(absolutePath); + const results: GrepMatch[] = []; try { const rgPath = await ensureRgPath(); - const output = await new Promise((resolve, reject) => { - const child = spawn(rgPath, rgArgs, { - windowsHide: true, - }); - - const stdoutChunks: Buffer[] = []; - const stderrChunks: Buffer[] = []; - - const cleanup = () => { - if (options.signal.aborted) { - child.kill(); - } - }; - - options.signal.addEventListener('abort', cleanup, { once: true }); - - child.stdout.on('data', (chunk) => stdoutChunks.push(chunk)); - child.stderr.on('data', (chunk) => stderrChunks.push(chunk)); - - child.on('error', (err) => { - options.signal.removeEventListener('abort', cleanup); - reject( - new Error( - `Failed to start ripgrep: ${err.message}. Please ensure @lvce-editor/ripgrep is properly installed.`, - ), - ); - }); - - child.on('close', (code, signal) => { - options.signal.removeEventListener('abort', cleanup); - const stdoutData = Buffer.concat(stdoutChunks).toString('utf8'); - const stderrData = Buffer.concat(stderrChunks).toString('utf8'); - - if (code === 0) { - resolve(stdoutData); - } else if (code === 1) { - resolve(''); // No matches found - } else { - if (signal) { - reject(new Error(`ripgrep was terminated by signal: ${signal}`)); - } else { - reject( - new Error(`ripgrep exited with code ${code}: ${stderrData}`), - ); - } - } - }); + const generator = execStreaming(rgPath, rgArgs, { + signal: options.signal, + allowedExitCodes: [0, 1], }); - return this.parseRipgrepJsonOutput(output, absolutePath); + for await (const line of generator) { + const match = this.parseRipgrepJsonLine(line, absolutePath); + if (match) { + results.push(match); + if (results.length >= maxMatches) { + break; + } + } + } + + return results; } catch (error: unknown) { debugLogger.debug(`GrepLogic: ripgrep failed: ${getErrorMessage(error)}`); throw error; } } + private parseRipgrepJsonLine( + line: string, + basePath: string, + ): GrepMatch | null { + try { + const json = JSON.parse(line); + if (json.type === 'match') { + const match = json.data; + // Defensive check: ensure text properties exist (skips binary/invalid encoding) + if (match.path?.text && match.lines?.text) { + const absoluteFilePath = path.resolve(basePath, match.path.text); + const relativeCheck = path.relative(basePath, absoluteFilePath); + if ( + relativeCheck === '..' || + relativeCheck.startsWith(`..${path.sep}`) || + path.isAbsolute(relativeCheck) + ) { + return null; + } + + const relativeFilePath = path.relative(basePath, absoluteFilePath); + + return { + filePath: relativeFilePath || path.basename(absoluteFilePath), + lineNumber: match.line_number, + line: match.lines.text.trimEnd(), + }; + } + } + } catch (error) { + // Only log if it's not a simple empty line or widely invalid + if (line.trim().length > 0) { + debugLogger.warn( + `Failed to parse ripgrep JSON line: ${line.substring(0, 100)}...`, + error, + ); + } + } + return null; + } + /** * Gets a description of the grep operation * @param params Parameters for the grep operation diff --git a/packages/core/src/utils/shell-utils.integration.test.ts b/packages/core/src/utils/shell-utils.integration.test.ts new file mode 100644 index 0000000000..717e01594b --- /dev/null +++ b/packages/core/src/utils/shell-utils.integration.test.ts @@ -0,0 +1,67 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ +import { describe, it, expect } from 'vitest'; +import { execStreaming } from './shell-utils.js'; + +// Integration tests using real child processes +describe('execStreaming (Integration)', () => { + it('should yield lines from stdout', async () => { + // Use node to echo for cross-platform support + const generator = execStreaming(process.execPath, [ + '-e', + 'console.log("line 1\\nline 2")', + ]); + const lines = []; + for await (const line of generator) { + lines.push(line); + } + expect(lines).toEqual(['line 1', 'line 2']); + }); + + it('should throw error on non-zero exit code', async () => { + // exit 2 via node + const generator = execStreaming(process.execPath, [ + '-e', + 'process.exit(2)', + ]); + + await expect(async () => { + for await (const _ of generator) { + // consume + } + }).rejects.toThrow(); + }); + + it('should abort cleanly when signal is aborted', async () => { + const controller = new AbortController(); + // sleep for 2s via node + const generator = execStreaming( + process.execPath, + ['-e', 'setTimeout(() => {}, 2000)'], + { signal: controller.signal }, + ); + + // Start reading + const readPromise = (async () => { + const lines = []; + try { + for await (const line of generator) { + lines.push(line); + } + } catch (_e) { + // ignore + } + return lines; + })(); + + setTimeout(() => { + controller.abort(); + }, 100); + + const lines = await readPromise; + expect(lines).toEqual([]); + }); +}); diff --git a/packages/core/src/utils/shell-utils.ts b/packages/core/src/utils/shell-utils.ts index 6610a5d45c..3a002f2895 100644 --- a/packages/core/src/utils/shell-utils.ts +++ b/packages/core/src/utils/shell-utils.ts @@ -13,6 +13,7 @@ import { spawnSync, type SpawnOptionsWithoutStdio, } from 'node:child_process'; +import * as readline from 'node:readline'; import type { Node, Tree } from 'web-tree-sitter'; import { Language, Parser, Query } from 'web-tree-sitter'; import { loadWasmBinary } from './fileUtils.js'; @@ -765,3 +766,123 @@ export const spawnAsync = ( reject(err); }); }); + +/** + * Executes a command and yields lines of output as they appear. + * Use for large outputs where buffering is not feasible. + * + * @param command The executable to run + * @param args Arguments for the executable + * @param options Spawn options (cwd, env, etc.) + */ +export async function* execStreaming( + command: string, + args: string[], + options?: SpawnOptionsWithoutStdio & { + signal?: AbortSignal; + allowedExitCodes?: number[]; + }, +): AsyncGenerator { + const child = spawn(command, args, { + ...options, + // ensure we don't open a window on windows if possible/relevant + windowsHide: true, + }); + + const rl = readline.createInterface({ + input: child.stdout, + terminal: false, + }); + + const errorChunks: Buffer[] = []; + let stderrTotalBytes = 0; + const MAX_STDERR_BYTES = 20 * 1024; // 20KB limit + + child.stderr.on('data', (chunk) => { + if (stderrTotalBytes < MAX_STDERR_BYTES) { + errorChunks.push(chunk); + stderrTotalBytes += chunk.length; + } + }); + + let error: Error | null = null; + child.on('error', (err) => { + error = err; + }); + + const onAbort = () => { + // If manually aborted by signal, we kill immediately. + if (!child.killed) child.kill(); + }; + + if (options?.signal?.aborted) { + onAbort(); + } else { + options?.signal?.addEventListener('abort', onAbort); + } + + let finished = false; + try { + for await (const line of rl) { + if (options?.signal?.aborted) break; + yield line; + } + finished = true; + } finally { + rl.close(); + options?.signal?.removeEventListener('abort', onAbort); + + // Ensure process is killed when the generator is closed (consumer breaks loop) + let killedByGenerator = false; + if (!finished && child.exitCode === null && !child.killed) { + try { + child.kill(); + } catch (_e) { + // ignore error if process is already dead + } + killedByGenerator = true; + } + + // Ensure we wait for the process to exit to check codes + await new Promise((resolve, reject) => { + // If an error occurred before we got here (e.g. spawn failure), reject immediately. + if (error) { + reject(error); + return; + } + + function checkExit(code: number | null) { + // If we aborted or killed it manually, we treat it as success (stop waiting) + if (options?.signal?.aborted || killedByGenerator) { + resolve(); + return; + } + + const allowed = options?.allowedExitCodes ?? [0]; + if (code !== null && allowed.includes(code)) { + resolve(); + } else { + // If we have an accumulated error or explicit error event + if (error) reject(error); + else { + const stderr = Buffer.concat(errorChunks).toString('utf8'); + const truncatedMsg = + stderrTotalBytes >= MAX_STDERR_BYTES ? '...[truncated]' : ''; + reject( + new Error( + `Process exited with code ${code}: ${stderr}${truncatedMsg}`, + ), + ); + } + } + } + + if (child.exitCode !== null) { + checkExit(child.exitCode); + } else { + child.on('close', (code) => checkExit(code)); + child.on('error', (err) => reject(err)); + } + }); + } +} From 13bc5f620c052220b1aaae401815974ae405bd31 Mon Sep 17 00:00:00 2001 From: Jerop Kipruto Date: Mon, 26 Jan 2026 16:57:27 -0500 Subject: [PATCH 9/9] feat(plan): add persistent plan file storage (#17563) --- packages/cli/src/config/config.test.ts | 4 +- .../config/policy-engine.integration.test.ts | 111 ++++++++++++++++++ packages/core/src/config/config.test.ts | 53 +++++++++ packages/core/src/config/config.ts | 9 ++ packages/core/src/config/storage.test.ts | 6 + packages/core/src/config/storage.ts | 4 + .../core/__snapshots__/prompts.test.ts.snap | 7 +- packages/core/src/core/prompts.test.ts | 3 + packages/core/src/core/prompts.ts | 9 +- packages/core/src/policy/policies/plan.toml | 8 ++ packages/core/src/tools/write-file.test.ts | 29 ++++- 11 files changed, 238 insertions(+), 5 deletions(-) diff --git a/packages/cli/src/config/config.test.ts b/packages/cli/src/config/config.test.ts index b93496262c..c8b9dcfb87 100644 --- a/packages/cli/src/config/config.test.ts +++ b/packages/cli/src/config/config.test.ts @@ -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) => { diff --git a/packages/cli/src/config/policy-engine.integration.test.ts b/packages/cli/src/config/policy-engine.integration.test.ts index 422ca92aad..f4cc35dd8a 100644 --- a/packages/cli/src/config/policy-engine.integration.test.ts +++ b/packages/cli/src/config/policy-engine.integration.test.ts @@ -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: { diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index 2ee826c466..97b2ab67bb 100644 --- a/packages/core/src/config/config.test.ts +++ b/packages/core/src/config/config.test.ts @@ -14,6 +14,7 @@ import { ApprovalMode } from '../policy/types.js'; import type { HookDefinition } from '../hooks/types.js'; import { HookType, HookEventName } from '../hooks/types.js'; import * as path from 'node:path'; +import * as fs from 'node:fs'; import { setGeminiMdFilename as mockSetGeminiMdFilename } from '../tools/memoryTool.js'; import { DEFAULT_TELEMETRY_TARGET, @@ -2232,3 +2233,55 @@ describe('Config JIT Initialization', () => { }); }); }); + +describe('Plans Directory Initialization', () => { + const baseParams: ConfigParameters = { + sessionId: 'test-session', + targetDir: '/tmp/test', + debugMode: false, + model: 'test-model', + cwd: '/tmp/test', + }; + + beforeEach(() => { + vi.spyOn(fs.promises, 'mkdir').mockResolvedValue(undefined); + }); + + afterEach(() => { + vi.mocked(fs.promises.mkdir).mockRestore(); + }); + + it('should create plans directory and add it to workspace context when plan is enabled', async () => { + const config = new Config({ + ...baseParams, + plan: true, + }); + + await config.initialize(); + + const plansDir = config.storage.getProjectTempPlansDir(); + expect(fs.promises.mkdir).toHaveBeenCalledWith(plansDir, { + recursive: true, + }); + + const context = config.getWorkspaceContext(); + expect(context.getDirectories()).toContain(plansDir); + }); + + it('should NOT create plans directory or add it to workspace context when plan is disabled', async () => { + const config = new Config({ + ...baseParams, + plan: false, + }); + + await config.initialize(); + + const plansDir = config.storage.getProjectTempPlansDir(); + expect(fs.promises.mkdir).not.toHaveBeenCalledWith(plansDir, { + recursive: true, + }); + + const context = config.getWorkspaceContext(); + expect(context.getDirectories()).not.toContain(plansDir); + }); +}); diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 2dd235becf..e65abff562 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -4,6 +4,7 @@ * SPDX-License-Identifier: Apache-2.0 */ +import * as fs from 'node:fs'; import * as path from 'node:path'; import { inspect } from 'node:util'; import process from 'node:process'; @@ -696,6 +697,7 @@ export class Config { this.extensionManagement = params.extensionManagement ?? true; this.enableExtensionReloading = params.enableExtensionReloading ?? false; this.storage = new Storage(this.targetDir); + this.fakeResponses = params.fakeResponses; this.recordResponses = params.recordResponses; this.enablePromptCompletion = params.enablePromptCompletion ?? false; @@ -794,6 +796,13 @@ export class Config { this.workspaceContext.addDirectory(dir); } + // Add plans directory to workspace context for plan file storage + if (this.planEnabled) { + const plansDir = this.storage.getProjectTempPlansDir(); + await fs.promises.mkdir(plansDir, { recursive: true }); + this.workspaceContext.addDirectory(plansDir); + } + // Initialize centralized FileDiscoveryService const discoverToolsHandle = startupProfiler.start('discover_tools'); this.getFileService(); diff --git a/packages/core/src/config/storage.test.ts b/packages/core/src/config/storage.test.ts index 342ae3866e..b0b4fa8791 100644 --- a/packages/core/src/config/storage.test.ts +++ b/packages/core/src/config/storage.test.ts @@ -78,4 +78,10 @@ describe('Storage – additional helpers', () => { const expected = path.join(os.homedir(), GEMINI_DIR, 'tmp', 'bin'); expect(Storage.getGlobalBinDir()).toBe(expected); }); + + it('getProjectTempPlansDir returns ~/.gemini/tmp//plans', () => { + const tempDir = storage.getProjectTempDir(); + const expected = path.join(tempDir, 'plans'); + expect(storage.getProjectTempPlansDir()).toBe(expected); + }); }); diff --git a/packages/core/src/config/storage.ts b/packages/core/src/config/storage.ts index ac7efb8103..1f317d4ddf 100644 --- a/packages/core/src/config/storage.ts +++ b/packages/core/src/config/storage.ts @@ -155,6 +155,10 @@ export class Storage { return path.join(this.getProjectTempDir(), 'logs'); } + getProjectTempPlansDir(): string { + return path.join(this.getProjectTempDir(), 'plans'); + } + getExtensionsDir(): string { return path.join(this.getGeminiDir(), 'extensions'); } diff --git a/packages/core/src/core/__snapshots__/prompts.test.ts.snap b/packages/core/src/core/__snapshots__/prompts.test.ts.snap index 779c7bb48d..59a7f25d7f 100644 --- a/packages/core/src/core/__snapshots__/prompts.test.ts.snap +++ b/packages/core/src/core/__snapshots__/prompts.test.ts.snap @@ -182,6 +182,11 @@ You are operating in **Plan Mode** - a structured planning workflow for designin ## Available Tools The following read-only tools are available in Plan Mode: +- \`write_file\` - Save plans to the plans directory (see Plan Storage below) + +## Plan Storage +- Save your plans as Markdown (.md) files directly to: \`/tmp/project-temp/plans/\` +- Use descriptive filenames: \`feature-name.md\` or \`bugfix-description.md\` ## Workflow Phases @@ -201,7 +206,7 @@ The following read-only tools are available in Plan Mode: - Only begin this phase after exploration is complete - Create a detailed implementation plan with clear steps - Include file paths, function signatures, and code snippets where helpful -- Present the plan for review +- After saving the plan, present the full content of the markdown file to the user for review ### Phase 4: Review & Approval - Ask the user if they approve the plan, want revisions, or want to reject it diff --git a/packages/core/src/core/prompts.test.ts b/packages/core/src/core/prompts.test.ts index 149f46dc00..7805702986 100644 --- a/packages/core/src/core/prompts.test.ts +++ b/packages/core/src/core/prompts.test.ts @@ -65,6 +65,9 @@ describe('Core System Prompt (prompts.ts)', () => { getEnableShellOutputEfficiency: vi.fn().mockReturnValue(true), storage: { getProjectTempDir: vi.fn().mockReturnValue('/tmp/project-temp'), + getProjectTempPlansDir: vi + .fn() + .mockReturnValue('/tmp/project-temp/plans'), }, isInteractive: vi.fn().mockReturnValue(true), isInteractiveShellEnabled: vi.fn().mockReturnValue(true), diff --git a/packages/core/src/core/prompts.ts b/packages/core/src/core/prompts.ts index fb5f14cf9b..83e346f368 100644 --- a/packages/core/src/core/prompts.ts +++ b/packages/core/src/core/prompts.ts @@ -146,6 +146,8 @@ export function getCoreSystemPrompt( .map((toolName) => `- \`${toolName}\``) .join('\n'); + const plansDir = config.storage.getProjectTempPlansDir(); + approvalModePrompt = ` # Active Approval Mode: Plan @@ -154,6 +156,11 @@ You are operating in **Plan Mode** - a structured planning workflow for designin ## Available Tools The following read-only tools are available in Plan Mode: ${planModeToolsList} +- \`${WRITE_FILE_TOOL_NAME}\` - Save plans to the plans directory (see Plan Storage below) + +## Plan Storage +- Save your plans as Markdown (.md) files directly to: \`${plansDir}/\` +- Use descriptive filenames: \`feature-name.md\` or \`bugfix-description.md\` ## Workflow Phases @@ -173,7 +180,7 @@ ${planModeToolsList} - Only begin this phase after exploration is complete - Create a detailed implementation plan with clear steps - Include file paths, function signatures, and code snippets where helpful -- Present the plan for review +- After saving the plan, present the full content of the markdown file to the user for review ### Phase 4: Review & Approval - Ask the user if they approve the plan, want revisions, or want to reject it diff --git a/packages/core/src/policy/policies/plan.toml b/packages/core/src/policy/policies/plan.toml index c69493e7e3..8487f34965 100644 --- a/packages/core/src/policy/policies/plan.toml +++ b/packages/core/src/policy/policies/plan.toml @@ -68,3 +68,11 @@ modes = ["plan"] toolName = "SubagentInvocation" decision = "allow" priority = 50 + +# Allow write_file for .md files in plans directory +[[rule]] +toolName = "write_file" +decision = "allow" +priority = 50 +modes = ["plan"] +argsPattern = "\"file_path\":\"[^\"]+/\\.gemini/tmp/[a-f0-9]{64}/plans/[a-zA-Z0-9_-]+\\.md\"" diff --git a/packages/core/src/tools/write-file.test.ts b/packages/core/src/tools/write-file.test.ts index 965656e4f8..6bdbab64ed 100644 --- a/packages/core/src/tools/write-file.test.ts +++ b/packages/core/src/tools/write-file.test.ts @@ -47,6 +47,7 @@ import { } from '../test-utils/mock-message-bus.js'; const rootDir = path.resolve(os.tmpdir(), 'gemini-cli-test-root'); +const plansDir = path.resolve(os.tmpdir(), 'gemini-cli-test-plans'); // --- MOCKS --- vi.mock('../core/client.js'); @@ -84,7 +85,7 @@ const mockConfigInternal = { getBaseLlmClient: vi.fn(), // Initialize as a plain mock function getFileSystemService: () => fsService, getIdeMode: vi.fn(() => false), - getWorkspaceContext: () => new WorkspaceContext(rootDir), + getWorkspaceContext: () => new WorkspaceContext(rootDir, [plansDir]), getApiKey: () => 'test-key', getModel: () => 'test-model', getSandbox: () => false, @@ -126,10 +127,13 @@ describe('WriteFileTool', () => { tempDir = fs.mkdtempSync( path.join(os.tmpdir(), 'write-file-test-external-'), ); - // Ensure the rootDir for the tool exists + // Ensure the rootDir and plansDir for the tool exists if (!fs.existsSync(rootDir)) { fs.mkdirSync(rootDir, { recursive: true }); } + if (!fs.existsSync(plansDir)) { + fs.mkdirSync(plansDir, { recursive: true }); + } // Setup GeminiClient mock mockGeminiClientInstance = new (vi.mocked(GeminiClient))( @@ -206,6 +210,9 @@ describe('WriteFileTool', () => { if (fs.existsSync(rootDir)) { fs.rmSync(rootDir, { recursive: true, force: true }); } + if (fs.existsSync(plansDir)) { + fs.rmSync(plansDir, { recursive: true, force: true }); + } vi.clearAllMocks(); }); @@ -813,6 +820,24 @@ describe('WriteFileTool', () => { /File path must be within one of the workspace directories/, ); }); + + it('should allow paths within the plans directory', () => { + const params = { + file_path: path.join(plansDir, 'my-plan.md'), + content: '# My Plan', + }; + expect(() => tool.build(params)).not.toThrow(); + }); + + it('should reject paths that try to escape the plans directory', () => { + const params = { + file_path: path.join(plansDir, '..', 'escaped.txt'), + content: 'malicious', + }; + expect(() => tool.build(params)).toThrow( + /File path must be within one of the workspace directories/, + ); + }); }); describe('specific error types for write failures', () => {