diff --git a/docs/cli/commands.md b/docs/cli/commands.md index beeecb88..9d45257b 100644 --- a/docs/cli/commands.md +++ b/docs/cli/commands.md @@ -271,7 +271,7 @@ When a custom command attempts to execute a shell command, Gemini CLI will now p 1. **Inject Commands:** Use the `!{...}` syntax. 2. **Argument Substitution:** If `{{args}}` is present inside the block, it is automatically shell-escaped (see [Context-Aware Injection](#1-context-aware-injection-with-args) above). -3. **Robust Parsing:** The parser correctly handles complex shell commands that include nested braces, such as JSON payloads. +3. **Robust Parsing:** The parser correctly handles complex shell commands that include nested braces, such as JSON payloads. **Note:** The content inside `!{...}` must have balanced braces (`{` and `}`). If you need to execute a command containing unbalanced braces, consider wrapping it in an external script file and calling the script within the `!{...}` block. 4. **Security Check and Confirmation:** The CLI performs a security check on the final, resolved command (after arguments are escaped and substituted). A dialog will appear showing the exact command(s) to be executed. 5. **Execution and Error Reporting:** The command is executed. If the command fails, the output injected into the prompt will include the error messages (stderr) followed by a status line, e.g., `[Shell command exited with code 1]`. This helps the model understand the context of the failure. @@ -299,6 +299,41 @@ Please generate a Conventional Commit message based on the following git diff: When you run `/git:commit`, the CLI first executes `git diff --staged`, then replaces `!{git diff --staged}` with the output of that command before sending the final, complete prompt to the model. +##### 4. Injecting File Content with `@{...}` + +You can directly embed the content of a file or a directory listing into your prompt using the `@{...}` syntax. This is useful for creating commands that operate on specific files. + +**How It Works:** + +- **File Injection**: `@{path/to/file.txt}` is replaced by the content of `file.txt`. +- **Multimodal Support**: If the path points to a supported image (e.g., PNG, JPEG), PDF, audio, or video file, it will be correctly encoded and injected as multimodal input. Other binary files are handled gracefully and skipped. +- **Directory Listing**: `@{path/to/dir}` is traversed and each file present within the directory and all subdirectories are inserted into the prompt. This respects `.gitignore` and `.geminiignore` if enabled. +- **Workspace-Aware**: The command searches for the path in the current directory and any other workspace directories. Absolute paths are allowed if they are within the workspace. +- **Processing Order**: File content injection with `@{...}` is processed _before_ shell commands (`!{...}`) and argument substitution (`{{args}}`). +- **Parsing**: The parser requires the content inside `@{...}` (the path) to have balanced braces (`{` and `}`). + +**Example (`review.toml`):** + +This command injects the content of a _fixed_ best practices file (`docs/best-practices.md`) and uses the user's arguments to provide context for the review. + +```toml +# In: /.gemini/commands/review.toml +# Invoked via: /review FileCommandLoader.ts + +description = "Reviews the provided context using a best practice guide." +prompt = """ +You are an expert code reviewer. + +Your task is to review {{args}}. + +Use the following best practices when providing your review: + +@{docs/best-practices.md} +""" +``` + +When you run `/review FileCommandLoader.ts`, the `@{docs/best-practices.md}` placeholder is replaced by the content of that file, and `{{args}}` is replaced by the text you provided, before the final prompt is sent to the model. + --- #### Example: A "Pure Function" Refactoring Command diff --git a/packages/cli/src/services/FileCommandLoader.test.ts b/packages/cli/src/services/FileCommandLoader.test.ts index 0b65f6c7..37093a36 100644 --- a/packages/cli/src/services/FileCommandLoader.test.ts +++ b/packages/cli/src/services/FileCommandLoader.test.ts @@ -14,6 +14,7 @@ import { createMockCommandContext } from '../test-utils/mockCommandContext.js'; import { SHELL_INJECTION_TRIGGER, SHORTHAND_ARGS_PLACEHOLDER, + type PromptPipelineContent, } from './prompt-processors/types.js'; import { ConfirmationRequiredError, @@ -21,8 +22,15 @@ import { } from './prompt-processors/shellProcessor.js'; import { DefaultArgumentProcessor } from './prompt-processors/argumentProcessor.js'; import type { CommandContext } from '../ui/commands/types.js'; +import { AtFileProcessor } from './prompt-processors/atFileProcessor.js'; const mockShellProcess = vi.hoisted(() => vi.fn()); +const mockAtFileProcess = vi.hoisted(() => vi.fn()); +vi.mock('./prompt-processors/atFileProcessor.js', () => ({ + AtFileProcessor: vi.fn().mockImplementation(() => ({ + process: mockAtFileProcess, + })), +})); vi.mock('./prompt-processors/shellProcessor.js', () => ({ ShellProcessor: vi.fn().mockImplementation(() => ({ process: mockShellProcess, @@ -68,15 +76,28 @@ describe('FileCommandLoader', () => { beforeEach(() => { vi.clearAllMocks(); mockShellProcess.mockImplementation( - (prompt: string, context: CommandContext) => { + (prompt: PromptPipelineContent, context: CommandContext) => { const userArgsRaw = context?.invocation?.args || ''; - const processedPrompt = prompt.replaceAll( + // This is a simplified mock. A real implementation would need to iterate + // through all parts and process only the text parts. + const firstTextPart = prompt.find( + (p) => typeof p === 'string' || 'text' in p, + ); + let textContent = ''; + if (typeof firstTextPart === 'string') { + textContent = firstTextPart; + } else if (firstTextPart && 'text' in firstTextPart) { + textContent = firstTextPart.text ?? ''; + } + + const processedText = textContent.replaceAll( SHORTHAND_ARGS_PLACEHOLDER, userArgsRaw, ); - return Promise.resolve(processedPrompt); + return Promise.resolve([{ text: processedText }]); }, ); + mockAtFileProcess.mockImplementation(async (prompt: string) => prompt); }); afterEach(() => { @@ -110,7 +131,7 @@ describe('FileCommandLoader', () => { '', ); if (result?.type === 'submit_prompt') { - expect(result.content).toBe('This is a test prompt'); + expect(result.content).toEqual([{ text: 'This is a test prompt' }]); } else { assert.fail('Incorrect action type'); } @@ -203,7 +224,7 @@ describe('FileCommandLoader', () => { const mockConfig = { getProjectRoot: vi.fn(() => '/path/to/project'), getExtensions: vi.fn(() => []), - } as Config; + } as unknown as Config; const loader = new FileCommandLoader(mockConfig); const commands = await loader.loadCommands(signal); expect(commands).toHaveLength(1); @@ -246,7 +267,7 @@ describe('FileCommandLoader', () => { const mockConfig = { getProjectRoot: vi.fn(() => process.cwd()), getExtensions: vi.fn(() => []), - } as Config; + } as unknown as Config; const loader = new FileCommandLoader(mockConfig); const commands = await loader.loadCommands(signal); @@ -262,7 +283,7 @@ describe('FileCommandLoader', () => { '', ); if (userResult?.type === 'submit_prompt') { - expect(userResult.content).toBe('User prompt'); + expect(userResult.content).toEqual([{ text: 'User prompt' }]); } else { assert.fail('Incorrect action type for user command'); } @@ -277,7 +298,7 @@ describe('FileCommandLoader', () => { '', ); if (projectResult?.type === 'submit_prompt') { - expect(projectResult.content).toBe('Project prompt'); + expect(projectResult.content).toEqual([{ text: 'Project prompt' }]); } else { assert.fail('Incorrect action type for project command'); } @@ -446,6 +467,54 @@ describe('FileCommandLoader', () => { expect(ShellProcessor).toHaveBeenCalledTimes(1); expect(DefaultArgumentProcessor).not.toHaveBeenCalled(); }); + + it('instantiates AtFileProcessor and DefaultArgumentProcessor if @{} is present', async () => { + const userCommandsDir = Storage.getUserCommandsDir(); + mock({ + [userCommandsDir]: { + 'at-file.toml': `prompt = "Context: @{./my-file.txt}"`, + }, + }); + + const loader = new FileCommandLoader(null as unknown as Config); + await loader.loadCommands(signal); + + expect(AtFileProcessor).toHaveBeenCalledTimes(1); + expect(ShellProcessor).not.toHaveBeenCalled(); + expect(DefaultArgumentProcessor).toHaveBeenCalledTimes(1); + }); + + it('instantiates ShellProcessor and AtFileProcessor if !{} and @{} are present', async () => { + const userCommandsDir = Storage.getUserCommandsDir(); + mock({ + [userCommandsDir]: { + 'shell-and-at.toml': `prompt = "Run !{cmd} with @{file.txt}"`, + }, + }); + + const loader = new FileCommandLoader(null as unknown as Config); + await loader.loadCommands(signal); + + expect(ShellProcessor).toHaveBeenCalledTimes(1); + expect(AtFileProcessor).toHaveBeenCalledTimes(1); + expect(DefaultArgumentProcessor).toHaveBeenCalledTimes(1); // because no {{args}} + }); + + it('instantiates only ShellProcessor and AtFileProcessor if {{args}} and @{} are present', async () => { + const userCommandsDir = Storage.getUserCommandsDir(); + mock({ + [userCommandsDir]: { + 'args-and-at.toml': `prompt = "Run {{args}} with @{file.txt}"`, + }, + }); + + const loader = new FileCommandLoader(null as unknown as Config); + await loader.loadCommands(signal); + + expect(ShellProcessor).toHaveBeenCalledTimes(1); + expect(AtFileProcessor).toHaveBeenCalledTimes(1); + expect(DefaultArgumentProcessor).not.toHaveBeenCalled(); + }); }); describe('Extension Command Loading', () => { @@ -487,7 +556,7 @@ describe('FileCommandLoader', () => { path: extensionDir, }, ]), - } as Config; + } as unknown as Config; const loader = new FileCommandLoader(mockConfig); const commands = await loader.loadCommands(signal); @@ -538,7 +607,7 @@ describe('FileCommandLoader', () => { path: extensionDir, }, ]), - } as Config; + } as unknown as Config; const loader = new FileCommandLoader(mockConfig); const commands = await loader.loadCommands(signal); @@ -559,7 +628,7 @@ describe('FileCommandLoader', () => { ); expect(result0?.type).toBe('submit_prompt'); if (result0?.type === 'submit_prompt') { - expect(result0.content).toBe('User deploy command'); + expect(result0.content).toEqual([{ text: 'User deploy command' }]); } expect(commands[1].name).toBe('deploy'); @@ -576,7 +645,7 @@ describe('FileCommandLoader', () => { ); expect(result1?.type).toBe('submit_prompt'); if (result1?.type === 'submit_prompt') { - expect(result1.content).toBe('Project deploy command'); + expect(result1.content).toEqual([{ text: 'Project deploy command' }]); } expect(commands[2].name).toBe('deploy'); @@ -594,7 +663,7 @@ describe('FileCommandLoader', () => { ); expect(result2?.type).toBe('submit_prompt'); if (result2?.type === 'submit_prompt') { - expect(result2.content).toBe('Extension deploy command'); + expect(result2.content).toEqual([{ text: 'Extension deploy command' }]); } }); @@ -645,7 +714,7 @@ describe('FileCommandLoader', () => { path: extensionDir2, }, ]), - } as Config; + } as unknown as Config; const loader = new FileCommandLoader(mockConfig); const commands = await loader.loadCommands(signal); @@ -681,7 +750,7 @@ describe('FileCommandLoader', () => { path: extensionDir, }, ]), - } as Config; + } as unknown as Config; const loader = new FileCommandLoader(mockConfig); const commands = await loader.loadCommands(signal); expect(commands).toHaveLength(0); @@ -713,7 +782,7 @@ describe('FileCommandLoader', () => { getExtensions: vi.fn(() => [ { name: 'a', version: '1.0.0', isActive: true, path: extensionDir }, ]), - } as Config; + } as unknown as Config; const loader = new FileCommandLoader(mockConfig); const commands = await loader.loadCommands(signal); @@ -737,7 +806,9 @@ describe('FileCommandLoader', () => { '', ); if (result?.type === 'submit_prompt') { - expect(result.content).toBe('Nested command from extension a'); + expect(result.content).toEqual([ + { text: 'Nested command from extension a' }, + ]); } else { assert.fail('Incorrect action type'); } @@ -771,7 +842,9 @@ describe('FileCommandLoader', () => { ); expect(result?.type).toBe('submit_prompt'); if (result?.type === 'submit_prompt') { - expect(result.content).toBe('The user wants to: do something cool'); + expect(result.content).toEqual([ + { text: 'The user wants to: do something cool' }, + ]); } }); }); @@ -805,7 +878,7 @@ describe('FileCommandLoader', () => { if (result?.type === 'submit_prompt') { const expectedContent = 'This is the instruction.\n\n/model_led 1.2.0 added "a feature"'; - expect(result.content).toBe(expectedContent); + expect(result.content).toEqual([{ text: expectedContent }]); } }); }); @@ -859,7 +932,7 @@ describe('FileCommandLoader', () => { 'shell.toml': `prompt = "Run !{echo 'hello'}"`, }, }); - mockShellProcess.mockResolvedValue('Run hello'); + mockShellProcess.mockResolvedValue([{ text: 'Run hello' }]); const loader = new FileCommandLoader(null as unknown as Config); const commands = await loader.loadCommands(signal); @@ -875,7 +948,7 @@ describe('FileCommandLoader', () => { expect(result?.type).toBe('submit_prompt'); if (result?.type === 'submit_prompt') { - expect(result.content).toBe('Run hello'); + expect(result.content).toEqual([{ text: 'Run hello' }]); } }); @@ -938,23 +1011,36 @@ describe('FileCommandLoader', () => { ), ).rejects.toThrow('Something else went wrong'); }); - it('assembles the processor pipeline in the correct order (Shell -> Default)', async () => { + it('assembles the processor pipeline in the correct order (AtFile -> Shell -> Default)', async () => { const userCommandsDir = Storage.getUserCommandsDir(); mock({ [userCommandsDir]: { - // This prompt uses !{} but NOT {{args}}, so both processors should be active. + // This prompt uses !{}, @{}, but NOT {{args}}, so all processors should be active. 'pipeline.toml': ` - prompt = "Shell says: ${SHELL_INJECTION_TRIGGER}echo foo}." + prompt = "Shell says: !{echo foo}. File says: @{./bar.txt}" `, }, + './bar.txt': 'bar content', }); const defaultProcessMock = vi .fn() - .mockImplementation((p) => Promise.resolve(`${p}-default-processed`)); + .mockImplementation((p: PromptPipelineContent) => + Promise.resolve([ + { text: `${(p[0] as { text: string }).text}-default-processed` }, + ]), + ); - mockShellProcess.mockImplementation((p) => - Promise.resolve(`${p}-shell-processed`), + mockShellProcess.mockImplementation((p: PromptPipelineContent) => + Promise.resolve([ + { text: `${(p[0] as { text: string }).text}-shell-processed` }, + ]), + ); + + mockAtFileProcess.mockImplementation((p: PromptPipelineContent) => + Promise.resolve([ + { text: `${(p[0] as { text: string }).text}-at-file-processed` }, + ]), ); vi.mocked(DefaultArgumentProcessor).mockImplementation( @@ -972,35 +1058,115 @@ describe('FileCommandLoader', () => { const result = await command!.action!( createMockCommandContext({ invocation: { - raw: '/pipeline bar', + raw: '/pipeline baz', name: 'pipeline', - args: 'bar', + args: 'baz', }, }), - 'bar', + 'baz', ); + expect(mockAtFileProcess.mock.invocationCallOrder[0]).toBeLessThan( + mockShellProcess.mock.invocationCallOrder[0], + ); expect(mockShellProcess.mock.invocationCallOrder[0]).toBeLessThan( defaultProcessMock.mock.invocationCallOrder[0], ); // Verify the flow of the prompt through the processors - // 1. Shell processor runs first - expect(mockShellProcess).toHaveBeenCalledWith( - expect.stringContaining(SHELL_INJECTION_TRIGGER), + // 1. AtFile processor runs first + expect(mockAtFileProcess).toHaveBeenCalledWith( + [{ text: expect.stringContaining('@{./bar.txt}') }], expect.any(Object), ); - // 2. Default processor runs second + // 2. Shell processor runs second + expect(mockShellProcess).toHaveBeenCalledWith( + [{ text: expect.stringContaining('-at-file-processed') }], + expect.any(Object), + ); + // 3. Default processor runs third expect(defaultProcessMock).toHaveBeenCalledWith( - expect.stringContaining('-shell-processed'), + [{ text: expect.stringContaining('-shell-processed') }], expect.any(Object), ); if (result?.type === 'submit_prompt') { - expect(result.content).toContain('-shell-processed-default-processed'); + const contentAsArray = Array.isArray(result.content) + ? result.content + : [result.content]; + expect(contentAsArray.length).toBeGreaterThan(0); + const firstPart = contentAsArray[0]; + + if (typeof firstPart === 'object' && firstPart && 'text' in firstPart) { + expect(firstPart.text).toContain( + '-at-file-processed-shell-processed-default-processed', + ); + } else { + assert.fail( + 'First part of content is not a text part or is a string', + ); + } } else { assert.fail('Incorrect action type'); } }); }); + + describe('@-file Processor Integration', () => { + it('correctly processes a command with @{file}', async () => { + const userCommandsDir = Storage.getUserCommandsDir(); + mock({ + [userCommandsDir]: { + 'at-file.toml': + 'prompt = "Context from file: @{./test.txt}"\ndescription = "@-file test"', + }, + './test.txt': 'file content', + }); + + mockAtFileProcess.mockImplementation( + async (prompt: PromptPipelineContent) => { + // A simplified mock of AtFileProcessor's behavior + const textContent = (prompt[0] as { text: string }).text; + if (textContent.includes('@{./test.txt}')) { + return [ + { + text: textContent.replace('@{./test.txt}', 'file content'), + }, + ]; + } + return prompt; + }, + ); + + // Prevent default processor from interfering + vi.mocked(DefaultArgumentProcessor).mockImplementation( + () => + ({ + process: (p: PromptPipelineContent) => Promise.resolve(p), + }) as unknown as DefaultArgumentProcessor, + ); + + const loader = new FileCommandLoader(null as unknown as Config); + const commands = await loader.loadCommands(signal); + const command = commands.find((c) => c.name === 'at-file'); + expect(command).toBeDefined(); + + const result = await command!.action?.( + createMockCommandContext({ + invocation: { + raw: '/at-file', + name: 'at-file', + args: '', + }, + }), + '', + ); + expect(result?.type).toBe('submit_prompt'); + if (result?.type === 'submit_prompt') { + expect(result.content).toEqual([ + { text: 'Context from file: file content' }, + ]); + } + }); + }); }); diff --git a/packages/cli/src/services/FileCommandLoader.ts b/packages/cli/src/services/FileCommandLoader.ts index 294d35f9..8b06d0e7 100644 --- a/packages/cli/src/services/FileCommandLoader.ts +++ b/packages/cli/src/services/FileCommandLoader.ts @@ -19,15 +19,20 @@ import type { } from '../ui/commands/types.js'; import { CommandKind } from '../ui/commands/types.js'; import { DefaultArgumentProcessor } from './prompt-processors/argumentProcessor.js'; -import type { IPromptProcessor } from './prompt-processors/types.js'; +import type { + IPromptProcessor, + PromptPipelineContent, +} from './prompt-processors/types.js'; import { SHORTHAND_ARGS_PLACEHOLDER, SHELL_INJECTION_TRIGGER, + AT_FILE_INJECTION_TRIGGER, } from './prompt-processors/types.js'; import { ConfirmationRequiredError, ShellProcessor, } from './prompt-processors/shellProcessor.js'; +import { AtFileProcessor } from './prompt-processors/atFileProcessor.js'; interface CommandDirectory { path: string; @@ -224,16 +229,25 @@ export class FileCommandLoader implements ICommandLoader { const usesShellInjection = validDef.prompt.includes( SHELL_INJECTION_TRIGGER, ); + const usesAtFileInjection = validDef.prompt.includes( + AT_FILE_INJECTION_TRIGGER, + ); - // Interpolation (Shell Execution and Argument Injection) - // If the prompt uses either shell injection OR argument placeholders, - // we must use the ShellProcessor. + // 1. @-File Injection (Security First). + // This runs first to ensure we're not executing shell commands that + // could dynamically generate malicious @-paths. + if (usesAtFileInjection) { + processors.push(new AtFileProcessor(baseCommandName)); + } + + // 2. Argument and Shell Injection. + // This runs after file content has been safely injected. if (usesShellInjection || usesArgs) { processors.push(new ShellProcessor(baseCommandName)); } - // Default Argument Handling - // If NO explicit argument injection ({{args}}) was used, we append the raw invocation. + // 3. Default Argument Handling. + // Appends the raw invocation if no explicit {{args}} are used. if (!usesArgs) { processors.push(new DefaultArgumentProcessor()); } @@ -253,19 +267,24 @@ export class FileCommandLoader implements ICommandLoader { ); return { type: 'submit_prompt', - content: validDef.prompt, // Fallback to unprocessed prompt + content: [{ text: validDef.prompt }], // Fallback to unprocessed prompt }; } try { - let processedPrompt = validDef.prompt; + let processedContent: PromptPipelineContent = [ + { text: validDef.prompt }, + ]; for (const processor of processors) { - processedPrompt = await processor.process(processedPrompt, context); + processedContent = await processor.process( + processedContent, + context, + ); } return { type: 'submit_prompt', - content: processedPrompt, + content: processedContent, }; } catch (e) { // Check if it's our specific error type diff --git a/packages/cli/src/services/prompt-processors/argumentProcessor.test.ts b/packages/cli/src/services/prompt-processors/argumentProcessor.test.ts index 1a4c0c6b..80bde128 100644 --- a/packages/cli/src/services/prompt-processors/argumentProcessor.test.ts +++ b/packages/cli/src/services/prompt-processors/argumentProcessor.test.ts @@ -13,7 +13,7 @@ describe('Argument Processors', () => { const processor = new DefaultArgumentProcessor(); it('should append the full command if args are provided', async () => { - const prompt = 'Parse the command.'; + const prompt = [{ text: 'Parse the command.' }]; const context = createMockCommandContext({ invocation: { raw: '/mycommand arg1 "arg two"', @@ -22,11 +22,13 @@ describe('Argument Processors', () => { }, }); const result = await processor.process(prompt, context); - expect(result).toBe('Parse the command.\n\n/mycommand arg1 "arg two"'); + expect(result).toEqual([ + { text: 'Parse the command.\n\n/mycommand arg1 "arg two"' }, + ]); }); it('should NOT append the full command if no args are provided', async () => { - const prompt = 'Parse the command.'; + const prompt = [{ text: 'Parse the command.' }]; const context = createMockCommandContext({ invocation: { raw: '/mycommand', @@ -35,7 +37,7 @@ describe('Argument Processors', () => { }, }); const result = await processor.process(prompt, context); - expect(result).toBe('Parse the command.'); + expect(result).toEqual([{ text: 'Parse the command.' }]); }); }); }); diff --git a/packages/cli/src/services/prompt-processors/argumentProcessor.ts b/packages/cli/src/services/prompt-processors/argumentProcessor.ts index 6fa9b47b..b5e5b38a 100644 --- a/packages/cli/src/services/prompt-processors/argumentProcessor.ts +++ b/packages/cli/src/services/prompt-processors/argumentProcessor.ts @@ -4,7 +4,8 @@ * SPDX-License-Identifier: Apache-2.0 */ -import type { IPromptProcessor } from './types.js'; +import { appendToLastTextPart } from '@google/gemini-cli-core'; +import type { IPromptProcessor, PromptPipelineContent } from './types.js'; import type { CommandContext } from '../../ui/commands/types.js'; /** @@ -14,9 +15,12 @@ import type { CommandContext } from '../../ui/commands/types.js'; * This processor is only used if the prompt does NOT contain {{args}}. */ export class DefaultArgumentProcessor implements IPromptProcessor { - async process(prompt: string, context: CommandContext): Promise { - if (context.invocation!.args) { - return `${prompt}\n\n${context.invocation!.raw}`; + async process( + prompt: PromptPipelineContent, + context: CommandContext, + ): Promise { + if (context.invocation?.args) { + return appendToLastTextPart(prompt, context.invocation.raw); } return prompt; } diff --git a/packages/cli/src/services/prompt-processors/atFileProcessor.test.ts b/packages/cli/src/services/prompt-processors/atFileProcessor.test.ts new file mode 100644 index 00000000..3f492481 --- /dev/null +++ b/packages/cli/src/services/prompt-processors/atFileProcessor.test.ts @@ -0,0 +1,221 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, beforeEach, vi } from 'vitest'; +import { createMockCommandContext } from '../../test-utils/mockCommandContext.js'; +import { type CommandContext } from '../../ui/commands/types.js'; +import { AtFileProcessor } from './atFileProcessor.js'; +import { MessageType } from '../../ui/types.js'; +import type { Config } from '@google/gemini-cli-core'; +import type { PartUnion } from '@google/genai'; + +// Mock the core dependency +const mockReadPathFromWorkspace = vi.hoisted(() => vi.fn()); +vi.mock('@google/gemini-cli-core', async (importOriginal) => { + const original = await importOriginal(); + return { + ...original, + readPathFromWorkspace: mockReadPathFromWorkspace, + }; +}); + +describe('AtFileProcessor', () => { + let context: CommandContext; + let mockConfig: Config; + + beforeEach(() => { + vi.clearAllMocks(); + + mockConfig = { + // The processor only passes the config through, so we don't need a full mock. + } as unknown as Config; + + context = createMockCommandContext({ + services: { + config: mockConfig, + }, + }); + + // Default mock success behavior: return content wrapped in a text part. + mockReadPathFromWorkspace.mockImplementation( + async (path: string): Promise => [ + { text: `content of ${path}` }, + ], + ); + }); + + it('should not change the prompt if no @{ trigger is present', async () => { + const processor = new AtFileProcessor(); + const prompt: PartUnion[] = [{ text: 'This is a simple prompt.' }]; + const result = await processor.process(prompt, context); + expect(result).toEqual(prompt); + expect(mockReadPathFromWorkspace).not.toHaveBeenCalled(); + }); + + it('should not change the prompt if config service is missing', async () => { + const processor = new AtFileProcessor(); + const prompt: PartUnion[] = [{ text: 'Analyze @{file.txt}' }]; + const contextWithoutConfig = createMockCommandContext({ + services: { + config: null, + }, + }); + const result = await processor.process(prompt, contextWithoutConfig); + expect(result).toEqual(prompt); + expect(mockReadPathFromWorkspace).not.toHaveBeenCalled(); + }); + + describe('Parsing Logic', () => { + it('should replace a single valid @{path/to/file.txt} placeholder', async () => { + const processor = new AtFileProcessor(); + const prompt: PartUnion[] = [ + { text: 'Analyze this file: @{path/to/file.txt}' }, + ]; + const result = await processor.process(prompt, context); + expect(mockReadPathFromWorkspace).toHaveBeenCalledWith( + 'path/to/file.txt', + mockConfig, + ); + expect(result).toEqual([ + { text: 'Analyze this file: ' }, + { text: 'content of path/to/file.txt' }, + ]); + }); + + it('should replace multiple different @{...} placeholders', async () => { + const processor = new AtFileProcessor(); + const prompt: PartUnion[] = [ + { text: 'Compare @{file1.js} with @{file2.js}' }, + ]; + const result = await processor.process(prompt, context); + expect(mockReadPathFromWorkspace).toHaveBeenCalledTimes(2); + expect(mockReadPathFromWorkspace).toHaveBeenCalledWith( + 'file1.js', + mockConfig, + ); + expect(mockReadPathFromWorkspace).toHaveBeenCalledWith( + 'file2.js', + mockConfig, + ); + expect(result).toEqual([ + { text: 'Compare ' }, + { text: 'content of file1.js' }, + { text: ' with ' }, + { text: 'content of file2.js' }, + ]); + }); + + it('should handle placeholders at the beginning, middle, and end', async () => { + const processor = new AtFileProcessor(); + const prompt: PartUnion[] = [ + { text: '@{start.txt} in the @{middle.txt} and @{end.txt}' }, + ]; + const result = await processor.process(prompt, context); + expect(result).toEqual([ + { text: 'content of start.txt' }, + { text: ' in the ' }, + { text: 'content of middle.txt' }, + { text: ' and ' }, + { text: 'content of end.txt' }, + ]); + }); + + it('should correctly parse paths that contain balanced braces', async () => { + const processor = new AtFileProcessor(); + const prompt: PartUnion[] = [ + { text: 'Analyze @{path/with/{braces}/file.txt}' }, + ]; + const result = await processor.process(prompt, context); + expect(mockReadPathFromWorkspace).toHaveBeenCalledWith( + 'path/with/{braces}/file.txt', + mockConfig, + ); + expect(result).toEqual([ + { text: 'Analyze ' }, + { text: 'content of path/with/{braces}/file.txt' }, + ]); + }); + + it('should throw an error if the prompt contains an unclosed trigger', async () => { + const processor = new AtFileProcessor(); + const prompt: PartUnion[] = [{ text: 'Hello @{world' }]; + // The new parser throws an error for unclosed injections. + await expect(processor.process(prompt, context)).rejects.toThrow( + /Unclosed injection/, + ); + }); + }); + + describe('Integration and Error Handling', () => { + it('should leave the placeholder unmodified if readPathFromWorkspace throws', async () => { + const processor = new AtFileProcessor(); + const prompt: PartUnion[] = [ + { text: 'Analyze @{not-found.txt} and @{good-file.txt}' }, + ]; + mockReadPathFromWorkspace.mockImplementation(async (path: string) => { + if (path === 'not-found.txt') { + throw new Error('File not found'); + } + return [{ text: `content of ${path}` }]; + }); + + const result = await processor.process(prompt, context); + expect(result).toEqual([ + { text: 'Analyze ' }, + { text: '@{not-found.txt}' }, // Placeholder is preserved as a text part + { text: ' and ' }, + { text: 'content of good-file.txt' }, + ]); + }); + }); + + describe('UI Feedback', () => { + it('should call ui.addItem with an ERROR on failure', async () => { + const processor = new AtFileProcessor(); + const prompt: PartUnion[] = [{ text: 'Analyze @{bad-file.txt}' }]; + mockReadPathFromWorkspace.mockRejectedValue(new Error('Access denied')); + + await processor.process(prompt, context); + + expect(context.ui.addItem).toHaveBeenCalledTimes(1); + expect(context.ui.addItem).toHaveBeenCalledWith( + { + type: MessageType.ERROR, + text: "Failed to inject content for '@{bad-file.txt}': Access denied", + }, + expect.any(Number), + ); + }); + + it('should call ui.addItem with a WARNING if the file was ignored', async () => { + const processor = new AtFileProcessor(); + const prompt: PartUnion[] = [{ text: 'Analyze @{ignored.txt}' }]; + // Simulate an ignored file by returning an empty array. + mockReadPathFromWorkspace.mockResolvedValue([]); + + const result = await processor.process(prompt, context); + + // The placeholder should be removed, resulting in only the prefix. + expect(result).toEqual([{ text: 'Analyze ' }]); + + expect(context.ui.addItem).toHaveBeenCalledTimes(1); + expect(context.ui.addItem).toHaveBeenCalledWith( + { + type: MessageType.INFO, + text: "File '@{ignored.txt}' was ignored by .gitignore or .geminiignore and was not included in the prompt.", + }, + expect.any(Number), + ); + }); + + it('should NOT call ui.addItem on success', async () => { + const processor = new AtFileProcessor(); + const prompt: PartUnion[] = [{ text: 'Analyze @{good-file.txt}' }]; + await processor.process(prompt, context); + expect(context.ui.addItem).not.toHaveBeenCalled(); + }); + }); +}); diff --git a/packages/cli/src/services/prompt-processors/atFileProcessor.ts b/packages/cli/src/services/prompt-processors/atFileProcessor.ts new file mode 100644 index 00000000..bff19fb6 --- /dev/null +++ b/packages/cli/src/services/prompt-processors/atFileProcessor.ts @@ -0,0 +1,96 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { + flatMapTextParts, + readPathFromWorkspace, +} from '@google/gemini-cli-core'; +import type { CommandContext } from '../../ui/commands/types.js'; +import { MessageType } from '../../ui/types.js'; +import { + AT_FILE_INJECTION_TRIGGER, + type IPromptProcessor, + type PromptPipelineContent, +} from './types.js'; +import { extractInjections } from './injectionParser.js'; + +export class AtFileProcessor implements IPromptProcessor { + constructor(private readonly commandName?: string) {} + + async process( + input: PromptPipelineContent, + context: CommandContext, + ): Promise { + const config = context.services.config; + if (!config) { + return input; + } + + return flatMapTextParts(input, async (text) => { + if (!text.includes(AT_FILE_INJECTION_TRIGGER)) { + return [{ text }]; + } + + const injections = extractInjections( + text, + AT_FILE_INJECTION_TRIGGER, + this.commandName, + ); + if (injections.length === 0) { + return [{ text }]; + } + + const output: PromptPipelineContent = []; + let lastIndex = 0; + + for (const injection of injections) { + const prefix = text.substring(lastIndex, injection.startIndex); + if (prefix) { + output.push({ text: prefix }); + } + + const pathStr = injection.content; + try { + const fileContentParts = await readPathFromWorkspace(pathStr, config); + if (fileContentParts.length === 0) { + const uiMessage = `File '@{${pathStr}}' was ignored by .gitignore or .geminiignore and was not included in the prompt.`; + context.ui.addItem( + { type: MessageType.INFO, text: uiMessage }, + Date.now(), + ); + } + output.push(...fileContentParts); + } catch (error) { + const message = + error instanceof Error ? error.message : String(error); + const uiMessage = `Failed to inject content for '@{${pathStr}}': ${message}`; + + console.error( + `[AtFileProcessor] ${uiMessage}. Leaving placeholder in prompt.`, + ); + context.ui.addItem( + { type: MessageType.ERROR, text: uiMessage }, + Date.now(), + ); + + const placeholder = text.substring( + injection.startIndex, + injection.endIndex, + ); + output.push({ text: placeholder }); + } + lastIndex = injection.endIndex; + } + + const suffix = text.substring(lastIndex); + if (suffix) { + output.push({ text: suffix }); + } + + return output; + }); + } +} diff --git a/packages/cli/src/services/prompt-processors/injectionParser.test.ts b/packages/cli/src/services/prompt-processors/injectionParser.test.ts new file mode 100644 index 00000000..5ce0f8f7 --- /dev/null +++ b/packages/cli/src/services/prompt-processors/injectionParser.test.ts @@ -0,0 +1,223 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect } from 'vitest'; +import { extractInjections } from './injectionParser.js'; + +describe('extractInjections', () => { + const SHELL_TRIGGER = '!{'; + const AT_FILE_TRIGGER = '@{'; + + describe('Basic Functionality', () => { + it('should return an empty array if no trigger is present', () => { + const prompt = 'This is a simple prompt without injections.'; + const result = extractInjections(prompt, SHELL_TRIGGER); + expect(result).toEqual([]); + }); + + it('should extract a single, simple injection', () => { + const prompt = 'Run this command: !{ls -la}'; + const result = extractInjections(prompt, SHELL_TRIGGER); + expect(result).toEqual([ + { + content: 'ls -la', + startIndex: 18, + endIndex: 27, + }, + ]); + }); + + it('should extract multiple injections', () => { + const prompt = 'First: !{cmd1}, Second: !{cmd2}'; + const result = extractInjections(prompt, SHELL_TRIGGER); + expect(result).toHaveLength(2); + expect(result[0]).toEqual({ + content: 'cmd1', + startIndex: 7, + endIndex: 14, + }); + expect(result[1]).toEqual({ + content: 'cmd2', + startIndex: 24, + endIndex: 31, + }); + }); + + it('should handle different triggers (e.g., @{)', () => { + const prompt = 'Read this file: @{path/to/file.txt}'; + const result = extractInjections(prompt, AT_FILE_TRIGGER); + expect(result).toEqual([ + { + content: 'path/to/file.txt', + startIndex: 16, + endIndex: 35, + }, + ]); + }); + }); + + describe('Positioning and Edge Cases', () => { + it('should handle injections at the start and end of the prompt', () => { + const prompt = '!{start} middle text !{end}'; + const result = extractInjections(prompt, SHELL_TRIGGER); + expect(result).toHaveLength(2); + expect(result[0]).toEqual({ + content: 'start', + startIndex: 0, + endIndex: 8, + }); + expect(result[1]).toEqual({ + content: 'end', + startIndex: 21, + endIndex: 27, + }); + }); + + it('should handle adjacent injections', () => { + const prompt = '!{A}!{B}'; + const result = extractInjections(prompt, SHELL_TRIGGER); + expect(result).toHaveLength(2); + expect(result[0]).toEqual({ content: 'A', startIndex: 0, endIndex: 4 }); + expect(result[1]).toEqual({ content: 'B', startIndex: 4, endIndex: 8 }); + }); + + it('should handle empty injections', () => { + const prompt = 'Empty: !{}'; + const result = extractInjections(prompt, SHELL_TRIGGER); + expect(result).toEqual([ + { + content: '', + startIndex: 7, + endIndex: 10, + }, + ]); + }); + + it('should trim whitespace within the content', () => { + const prompt = '!{ \n command with space \t }'; + const result = extractInjections(prompt, SHELL_TRIGGER); + expect(result).toEqual([ + { + content: 'command with space', + startIndex: 0, + endIndex: 29, + }, + ]); + }); + + it('should ignore similar patterns that are not the exact trigger', () => { + const prompt = 'Not a trigger: !(cmd) or {cmd} or ! {cmd}'; + const result = extractInjections(prompt, SHELL_TRIGGER); + expect(result).toEqual([]); + }); + + it('should ignore extra closing braces before the trigger', () => { + const prompt = 'Ignore this } then !{run}'; + const result = extractInjections(prompt, SHELL_TRIGGER); + expect(result).toEqual([ + { + content: 'run', + startIndex: 19, + endIndex: 25, + }, + ]); + }); + + it('should stop parsing at the first balanced closing brace (non-greedy)', () => { + // This tests that the parser doesn't greedily consume extra closing braces + const prompt = 'Run !{ls -l}} extra braces'; + const result = extractInjections(prompt, SHELL_TRIGGER); + expect(result).toEqual([ + { + content: 'ls -l', + startIndex: 4, + endIndex: 12, + }, + ]); + }); + }); + + describe('Nested Braces (Balanced)', () => { + it('should correctly parse content with simple nested braces (e.g., JSON)', () => { + const prompt = `Send JSON: !{curl -d '{"key": "value"}'}`; + const result = extractInjections(prompt, SHELL_TRIGGER); + expect(result).toHaveLength(1); + expect(result[0].content).toBe(`curl -d '{"key": "value"}'`); + }); + + it('should correctly parse content with shell constructs (e.g., awk)', () => { + const prompt = `Process text: !{awk '{print $1}' file.txt}`; + const result = extractInjections(prompt, SHELL_TRIGGER); + expect(result).toHaveLength(1); + expect(result[0].content).toBe(`awk '{print $1}' file.txt`); + }); + + it('should correctly parse multiple levels of nesting', () => { + const prompt = `!{level1 {level2 {level3}} suffix}`; + const result = extractInjections(prompt, SHELL_TRIGGER); + expect(result).toHaveLength(1); + expect(result[0].content).toBe(`level1 {level2 {level3}} suffix`); + expect(result[0].endIndex).toBe(prompt.length); + }); + + it('should correctly parse paths containing balanced braces', () => { + const prompt = 'Analyze @{path/with/{braces}/file.txt}'; + const result = extractInjections(prompt, AT_FILE_TRIGGER); + expect(result).toHaveLength(1); + expect(result[0].content).toBe('path/with/{braces}/file.txt'); + }); + + it('should correctly handle an injection containing the trigger itself', () => { + // This works because the parser counts braces, it doesn't look for the trigger again until the current one is closed. + const prompt = '!{echo "The trigger is !{ confusing }"}'; + const expectedContent = 'echo "The trigger is !{ confusing }"'; + const result = extractInjections(prompt, SHELL_TRIGGER); + expect(result).toHaveLength(1); + expect(result[0].content).toBe(expectedContent); + }); + }); + + describe('Error Handling (Unbalanced/Unclosed)', () => { + it('should throw an error for a simple unclosed injection', () => { + const prompt = 'This prompt has !{an unclosed trigger'; + expect(() => extractInjections(prompt, SHELL_TRIGGER)).toThrow( + /Invalid syntax: Unclosed injection starting at index 16 \('!{'\)/, + ); + }); + + it('should throw an error if the prompt ends inside a nested block', () => { + const prompt = 'This fails: !{outer {inner'; + expect(() => extractInjections(prompt, SHELL_TRIGGER)).toThrow( + /Invalid syntax: Unclosed injection starting at index 12 \('!{'\)/, + ); + }); + + it('should include the context name in the error message if provided', () => { + const prompt = 'Failing !{command'; + const contextName = 'test-command'; + expect(() => + extractInjections(prompt, SHELL_TRIGGER, contextName), + ).toThrow( + /Invalid syntax in command 'test-command': Unclosed injection starting at index 8/, + ); + }); + + it('should throw if content contains unbalanced braces (e.g., missing closing)', () => { + // This is functionally the same as an unclosed injection from the parser's perspective. + const prompt = 'Analyze @{path/with/braces{example.txt}'; + expect(() => extractInjections(prompt, AT_FILE_TRIGGER)).toThrow( + /Invalid syntax: Unclosed injection starting at index 8 \('@{'\)/, + ); + }); + + it('should clearly state that unbalanced braces in content are not supported in the error', () => { + const prompt = 'Analyze @{path/with/braces{example.txt}'; + expect(() => extractInjections(prompt, AT_FILE_TRIGGER)).toThrow( + /Paths or commands with unbalanced braces are not supported directly/, + ); + }); + }); +}); diff --git a/packages/cli/src/services/prompt-processors/injectionParser.ts b/packages/cli/src/services/prompt-processors/injectionParser.ts new file mode 100644 index 00000000..52d3226d --- /dev/null +++ b/packages/cli/src/services/prompt-processors/injectionParser.ts @@ -0,0 +1,89 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Represents a single detected injection site in a prompt string. + */ +export interface Injection { + /** The content extracted from within the braces (e.g., the command or path), trimmed. */ + content: string; + /** The starting index of the injection (inclusive, points to the start of the trigger). */ + startIndex: number; + /** The ending index of the injection (exclusive, points after the closing '}'). */ + endIndex: number; +} + +/** + * Iteratively parses a prompt string to extract injections (e.g., !{...} or @{...}), + * correctly handling nested braces within the content. + * + * This parser relies on simple brace counting and does not support escaping. + * + * @param prompt The prompt string to parse. + * @param trigger The opening trigger sequence (e.g., '!{', '@{'). + * @param contextName Optional context name (e.g., command name) for error messages. + * @returns An array of extracted Injection objects. + * @throws Error if an unclosed injection is found. + */ +export function extractInjections( + prompt: string, + trigger: string, + contextName?: string, +): Injection[] { + const injections: Injection[] = []; + let index = 0; + + while (index < prompt.length) { + const startIndex = prompt.indexOf(trigger, index); + + if (startIndex === -1) { + break; + } + + let currentIndex = startIndex + trigger.length; + let braceCount = 1; + let foundEnd = false; + + while (currentIndex < prompt.length) { + const char = prompt[currentIndex]; + + if (char === '{') { + braceCount++; + } else if (char === '}') { + braceCount--; + if (braceCount === 0) { + const injectionContent = prompt.substring( + startIndex + trigger.length, + currentIndex, + ); + const endIndex = currentIndex + 1; + + injections.push({ + content: injectionContent.trim(), + startIndex, + endIndex, + }); + + index = endIndex; + foundEnd = true; + break; + } + } + currentIndex++; + } + + // Check if the inner loop finished without finding the closing brace. + if (!foundEnd) { + const contextInfo = contextName ? ` in command '${contextName}'` : ''; + // Enforce strict parsing (Comment 1) and clarify limitations (Comment 2). + throw new Error( + `Invalid syntax${contextInfo}: Unclosed injection starting at index ${startIndex} ('${trigger}'). Ensure braces are balanced. Paths or commands with unbalanced braces are not supported directly.`, + ); + } + } + + return injections; +} diff --git a/packages/cli/src/services/prompt-processors/shellProcessor.test.ts b/packages/cli/src/services/prompt-processors/shellProcessor.test.ts index 5e33a8a5..3b2418e1 100644 --- a/packages/cli/src/services/prompt-processors/shellProcessor.test.ts +++ b/packages/cli/src/services/prompt-processors/shellProcessor.test.ts @@ -12,10 +12,11 @@ import type { Config } from '@google/gemini-cli-core'; import { ApprovalMode } from '@google/gemini-cli-core'; import os from 'node:os'; import { quote } from 'shell-quote'; +import { createPartFromText } from '@google/genai'; +import type { PromptPipelineContent } from './types.js'; // Helper function to determine the expected escaped string based on the current OS, -// mirroring the logic in the actual `escapeShellArg` implementation. This makes -// our tests robust and platform-agnostic. +// mirroring the logic in the actual `escapeShellArg` implementation. function getExpectedEscapedArgForPlatform(arg: string): string { if (os.platform() === 'win32') { const comSpec = (process.env['ComSpec'] || 'cmd.exe').toLowerCase(); @@ -32,6 +33,11 @@ function getExpectedEscapedArgForPlatform(arg: string): string { } } +// Helper to create PromptPipelineContent +function createPromptPipelineContent(text: string): PromptPipelineContent { + return [createPartFromText(text)]; +} + const mockCheckCommandPermissions = vi.hoisted(() => vi.fn()); const mockShellExecute = vi.hoisted(() => vi.fn()); @@ -93,7 +99,7 @@ describe('ShellProcessor', () => { it('should throw an error if config is missing', async () => { const processor = new ShellProcessor('test-command'); - const prompt = '!{ls}'; + const prompt: PromptPipelineContent = createPromptPipelineContent('!{ls}'); const contextWithoutConfig = createMockCommandContext({ services: { config: null, @@ -107,15 +113,19 @@ describe('ShellProcessor', () => { it('should not change the prompt if no shell injections are present', async () => { const processor = new ShellProcessor('test-command'); - const prompt = 'This is a simple prompt with no injections.'; + const prompt: PromptPipelineContent = createPromptPipelineContent( + 'This is a simple prompt with no injections.', + ); const result = await processor.process(prompt, context); - expect(result).toBe(prompt); + expect(result).toEqual(prompt); expect(mockShellExecute).not.toHaveBeenCalled(); }); it('should process a single valid shell injection if allowed', async () => { const processor = new ShellProcessor('test-command'); - const prompt = 'The current status is: !{git status}'; + const prompt: PromptPipelineContent = createPromptPipelineContent( + 'The current status is: !{git status}', + ); mockCheckCommandPermissions.mockReturnValue({ allAllowed: true, disallowedCommands: [], @@ -138,12 +148,14 @@ describe('ShellProcessor', () => { expect.any(Object), false, ); - expect(result).toBe('The current status is: On branch main'); + expect(result).toEqual([{ text: 'The current status is: On branch main' }]); }); it('should process multiple valid shell injections if all are allowed', async () => { const processor = new ShellProcessor('test-command'); - const prompt = '!{git status} in !{pwd}'; + const prompt: PromptPipelineContent = createPromptPipelineContent( + '!{git status} in !{pwd}', + ); mockCheckCommandPermissions.mockReturnValue({ allAllowed: true, disallowedCommands: [], @@ -164,12 +176,14 @@ describe('ShellProcessor', () => { expect(mockCheckCommandPermissions).toHaveBeenCalledTimes(2); expect(mockShellExecute).toHaveBeenCalledTimes(2); - expect(result).toBe('On branch main in /usr/home'); + expect(result).toEqual([{ text: 'On branch main in /usr/home' }]); }); it('should throw ConfirmationRequiredError if a command is not allowed in default mode', async () => { const processor = new ShellProcessor('test-command'); - const prompt = 'Do something dangerous: !{rm -rf /}'; + const prompt: PromptPipelineContent = createPromptPipelineContent( + 'Do something dangerous: !{rm -rf /}', + ); mockCheckCommandPermissions.mockReturnValue({ allAllowed: false, disallowedCommands: ['rm -rf /'], @@ -182,7 +196,9 @@ describe('ShellProcessor', () => { it('should NOT throw ConfirmationRequiredError if a command is not allowed but approval mode is YOLO', async () => { const processor = new ShellProcessor('test-command'); - const prompt = 'Do something dangerous: !{rm -rf /}'; + const prompt: PromptPipelineContent = createPromptPipelineContent( + 'Do something dangerous: !{rm -rf /}', + ); mockCheckCommandPermissions.mockReturnValue({ allAllowed: false, disallowedCommands: ['rm -rf /'], @@ -203,12 +219,14 @@ describe('ShellProcessor', () => { expect.any(Object), false, ); - expect(result).toBe('Do something dangerous: deleted'); + expect(result).toEqual([{ text: 'Do something dangerous: deleted' }]); }); it('should still throw an error for a hard-denied command even in YOLO mode', async () => { const processor = new ShellProcessor('test-command'); - const prompt = 'Do something forbidden: !{reboot}'; + const prompt: PromptPipelineContent = createPromptPipelineContent( + 'Do something forbidden: !{reboot}', + ); mockCheckCommandPermissions.mockReturnValue({ allAllowed: false, disallowedCommands: ['reboot'], @@ -228,7 +246,9 @@ describe('ShellProcessor', () => { it('should throw ConfirmationRequiredError with the correct command', async () => { const processor = new ShellProcessor('test-command'); - const prompt = 'Do something dangerous: !{rm -rf /}'; + const prompt: PromptPipelineContent = createPromptPipelineContent( + 'Do something dangerous: !{rm -rf /}', + ); mockCheckCommandPermissions.mockReturnValue({ allAllowed: false, disallowedCommands: ['rm -rf /'], @@ -250,7 +270,9 @@ describe('ShellProcessor', () => { it('should throw ConfirmationRequiredError with multiple commands if multiple are disallowed', async () => { const processor = new ShellProcessor('test-command'); - const prompt = '!{cmd1} and !{cmd2}'; + const prompt: PromptPipelineContent = createPromptPipelineContent( + '!{cmd1} and !{cmd2}', + ); mockCheckCommandPermissions.mockImplementation((cmd) => { if (cmd === 'cmd1') { return { allAllowed: false, disallowedCommands: ['cmd1'] }; @@ -275,7 +297,9 @@ describe('ShellProcessor', () => { it('should not execute any commands if at least one requires confirmation', async () => { const processor = new ShellProcessor('test-command'); - const prompt = 'First: !{echo "hello"}, Second: !{rm -rf /}'; + const prompt: PromptPipelineContent = createPromptPipelineContent( + 'First: !{echo "hello"}, Second: !{rm -rf /}', + ); mockCheckCommandPermissions.mockImplementation((cmd) => { if (cmd.includes('rm')) { @@ -294,7 +318,9 @@ describe('ShellProcessor', () => { it('should only request confirmation for disallowed commands in a mixed prompt', async () => { const processor = new ShellProcessor('test-command'); - const prompt = 'Allowed: !{ls -l}, Disallowed: !{rm -rf /}'; + const prompt: PromptPipelineContent = createPromptPipelineContent( + 'Allowed: !{ls -l}, Disallowed: !{rm -rf /}', + ); mockCheckCommandPermissions.mockImplementation((cmd) => ({ allAllowed: !cmd.includes('rm'), @@ -314,7 +340,9 @@ describe('ShellProcessor', () => { it('should execute all commands if they are on the session allowlist', async () => { const processor = new ShellProcessor('test-command'); - const prompt = 'Run !{cmd1} and !{cmd2}'; + const prompt: PromptPipelineContent = createPromptPipelineContent( + 'Run !{cmd1} and !{cmd2}', + ); // Add commands to the session allowlist context.session.sessionShellAllowlist = new Set(['cmd1', 'cmd2']); @@ -346,12 +374,14 @@ describe('ShellProcessor', () => { context.session.sessionShellAllowlist, ); expect(mockShellExecute).toHaveBeenCalledTimes(2); - expect(result).toBe('Run output1 and output2'); + expect(result).toEqual([{ text: 'Run output1 and output2' }]); }); it('should trim whitespace from the command inside the injection before interpolation', async () => { const processor = new ShellProcessor('test-command'); - const prompt = 'Files: !{ ls {{args}} -l }'; + const prompt: PromptPipelineContent = createPromptPipelineContent( + 'Files: !{ ls {{args}} -l }', + ); const rawArgs = context.invocation!.args; @@ -385,7 +415,8 @@ describe('ShellProcessor', () => { it('should handle an empty command inside the injection gracefully (skips execution)', async () => { const processor = new ShellProcessor('test-command'); - const prompt = 'This is weird: !{}'; + const prompt: PromptPipelineContent = + createPromptPipelineContent('This is weird: !{}'); const result = await processor.process(prompt, context); @@ -393,77 +424,14 @@ describe('ShellProcessor', () => { expect(mockShellExecute).not.toHaveBeenCalled(); // It replaces !{} with an empty string. - expect(result).toBe('This is weird: '); - }); - - describe('Robust Parsing (Balanced Braces)', () => { - it('should correctly parse commands containing nested braces (e.g., awk)', async () => { - const processor = new ShellProcessor('test-command'); - const command = "awk '{print $1}' file.txt"; - const prompt = `Output: !{${command}}`; - mockShellExecute.mockReturnValue({ - result: Promise.resolve({ ...SUCCESS_RESULT, output: 'result' }), - }); - - const result = await processor.process(prompt, context); - - expect(mockCheckCommandPermissions).toHaveBeenCalledWith( - command, - expect.any(Object), - context.session.sessionShellAllowlist, - ); - expect(mockShellExecute).toHaveBeenCalledWith( - command, - expect.any(String), - expect.any(Function), - expect.any(Object), - false, - ); - expect(result).toBe('Output: result'); - }); - - it('should handle deeply nested braces correctly', async () => { - const processor = new ShellProcessor('test-command'); - const command = "echo '{{a},{b}}'"; - const prompt = `!{${command}}`; - mockShellExecute.mockReturnValue({ - result: Promise.resolve({ ...SUCCESS_RESULT, output: '{{a},{b}}' }), - }); - - const result = await processor.process(prompt, context); - expect(mockShellExecute).toHaveBeenCalledWith( - command, - expect.any(String), - expect.any(Function), - expect.any(Object), - false, - ); - expect(result).toBe('{{a},{b}}'); - }); - - it('should throw an error for unclosed shell injections', async () => { - const processor = new ShellProcessor('test-command'); - const prompt = 'This prompt is broken: !{ls -l'; - - await expect(processor.process(prompt, context)).rejects.toThrow( - /Unclosed shell injection/, - ); - }); - - it('should throw an error for unclosed nested braces', async () => { - const processor = new ShellProcessor('test-command'); - const prompt = 'Broken: !{echo {a}'; - - await expect(processor.process(prompt, context)).rejects.toThrow( - /Unclosed shell injection/, - ); - }); + expect(result).toEqual([{ text: 'This is weird: ' }]); }); describe('Error Reporting', () => { it('should append exit code and command name on failure', async () => { const processor = new ShellProcessor('test-command'); - const prompt = '!{cmd}'; + const prompt: PromptPipelineContent = + createPromptPipelineContent('!{cmd}'); mockShellExecute.mockReturnValue({ result: Promise.resolve({ ...SUCCESS_RESULT, @@ -475,14 +443,17 @@ describe('ShellProcessor', () => { const result = await processor.process(prompt, context); - expect(result).toBe( - "some error output\n[Shell command 'cmd' exited with code 1]", - ); + expect(result).toEqual([ + { + text: "some error output\n[Shell command 'cmd' exited with code 1]", + }, + ]); }); it('should append signal info and command name if terminated by signal', async () => { const processor = new ShellProcessor('test-command'); - const prompt = '!{cmd}'; + const prompt: PromptPipelineContent = + createPromptPipelineContent('!{cmd}'); mockShellExecute.mockReturnValue({ result: Promise.resolve({ ...SUCCESS_RESULT, @@ -495,14 +466,17 @@ describe('ShellProcessor', () => { const result = await processor.process(prompt, context); - expect(result).toBe( - "output\n[Shell command 'cmd' terminated by signal SIGTERM]", - ); + expect(result).toEqual([ + { + text: "output\n[Shell command 'cmd' terminated by signal SIGTERM]", + }, + ]); }); it('should throw a detailed error if the shell fails to spawn', async () => { const processor = new ShellProcessor('test-command'); - const prompt = '!{bad-command}'; + const prompt: PromptPipelineContent = + createPromptPipelineContent('!{bad-command}'); const spawnError = new Error('spawn EACCES'); mockShellExecute.mockReturnValue({ result: Promise.resolve({ @@ -522,7 +496,9 @@ describe('ShellProcessor', () => { it('should report abort status with command name if aborted', async () => { const processor = new ShellProcessor('test-command'); - const prompt = '!{long-running-command}'; + const prompt: PromptPipelineContent = createPromptPipelineContent( + '!{long-running-command}', + ); const spawnError = new Error('Aborted'); mockShellExecute.mockReturnValue({ result: Promise.resolve({ @@ -536,9 +512,11 @@ describe('ShellProcessor', () => { }); const result = await processor.process(prompt, context); - expect(result).toBe( - "partial output\n[Shell command 'long-running-command' aborted]", - ); + expect(result).toEqual([ + { + text: "partial output\n[Shell command 'long-running-command' aborted]", + }, + ]); }); }); @@ -552,29 +530,35 @@ describe('ShellProcessor', () => { it('should perform raw replacement if no shell injections are present (optimization path)', async () => { const processor = new ShellProcessor('test-command'); - const prompt = 'The user said: {{args}}'; + const prompt: PromptPipelineContent = createPromptPipelineContent( + 'The user said: {{args}}', + ); const result = await processor.process(prompt, context); - expect(result).toBe(`The user said: ${rawArgs}`); + expect(result).toEqual([{ text: `The user said: ${rawArgs}` }]); expect(mockShellExecute).not.toHaveBeenCalled(); }); it('should perform raw replacement outside !{} blocks', async () => { const processor = new ShellProcessor('test-command'); - const prompt = 'Outside: {{args}}. Inside: !{echo "hello"}'; + const prompt: PromptPipelineContent = createPromptPipelineContent( + 'Outside: {{args}}. Inside: !{echo "hello"}', + ); mockShellExecute.mockReturnValue({ result: Promise.resolve({ ...SUCCESS_RESULT, output: 'hello' }), }); const result = await processor.process(prompt, context); - expect(result).toBe(`Outside: ${rawArgs}. Inside: hello`); + expect(result).toEqual([{ text: `Outside: ${rawArgs}. Inside: hello` }]); }); it('should perform escaped replacement inside !{} blocks', async () => { const processor = new ShellProcessor('test-command'); - const prompt = 'Command: !{grep {{args}} file.txt}'; + const prompt: PromptPipelineContent = createPromptPipelineContent( + 'Command: !{grep {{args}} file.txt}', + ); mockShellExecute.mockReturnValue({ result: Promise.resolve({ ...SUCCESS_RESULT, output: 'match found' }), }); @@ -592,12 +576,14 @@ describe('ShellProcessor', () => { false, ); - expect(result).toBe('Command: match found'); + expect(result).toEqual([{ text: 'Command: match found' }]); }); it('should handle both raw (outside) and escaped (inside) injection simultaneously', async () => { const processor = new ShellProcessor('test-command'); - const prompt = 'User "({{args}})" requested search: !{search {{args}}}'; + const prompt: PromptPipelineContent = createPromptPipelineContent( + 'User "({{args}})" requested search: !{search {{args}}}', + ); mockShellExecute.mockReturnValue({ result: Promise.resolve({ ...SUCCESS_RESULT, output: 'results' }), }); @@ -614,12 +600,15 @@ describe('ShellProcessor', () => { false, ); - expect(result).toBe(`User "(${rawArgs})" requested search: results`); + expect(result).toEqual([ + { text: `User "(${rawArgs})" requested search: results` }, + ]); }); it('should perform security checks on the final, resolved (escaped) command', async () => { const processor = new ShellProcessor('test-command'); - const prompt = '!{rm {{args}}}'; + const prompt: PromptPipelineContent = + createPromptPipelineContent('!{rm {{args}}}'); const expectedEscapedArgs = getExpectedEscapedArgForPlatform(rawArgs); const expectedResolvedCommand = `rm ${expectedEscapedArgs}`; @@ -642,7 +631,8 @@ describe('ShellProcessor', () => { it('should report the resolved command if a hard denial occurs', async () => { const processor = new ShellProcessor('test-command'); - const prompt = '!{rm {{args}}}'; + const prompt: PromptPipelineContent = + createPromptPipelineContent('!{rm {{args}}}'); const expectedEscapedArgs = getExpectedEscapedArgForPlatform(rawArgs); const expectedResolvedCommand = `rm ${expectedEscapedArgs}`; mockCheckCommandPermissions.mockReturnValue({ @@ -662,7 +652,9 @@ describe('ShellProcessor', () => { const processor = new ShellProcessor('test-command'); const multilineArgs = 'first line\nsecond line'; context.invocation!.args = multilineArgs; - const prompt = 'Commit message: !{git commit -m {{args}}}'; + const prompt: PromptPipelineContent = createPromptPipelineContent( + 'Commit message: !{git commit -m {{args}}}', + ); const expectedEscapedArgs = getExpectedEscapedArgForPlatform(multilineArgs); @@ -691,7 +683,8 @@ describe('ShellProcessor', () => { ])('should safely escape args containing $name', async ({ input }) => { const processor = new ShellProcessor('test-command'); context.invocation!.args = input; - const prompt = '!{echo {{args}}}'; + const prompt: PromptPipelineContent = + createPromptPipelineContent('!{echo {{args}}}'); const expectedEscapedArgs = getExpectedEscapedArgForPlatform(input); const expectedCommand = `echo ${expectedEscapedArgs}`; diff --git a/packages/cli/src/services/prompt-processors/shellProcessor.ts b/packages/cli/src/services/prompt-processors/shellProcessor.ts index cfa72292..d45e8f8a 100644 --- a/packages/cli/src/services/prompt-processors/shellProcessor.ts +++ b/packages/cli/src/services/prompt-processors/shellProcessor.ts @@ -10,14 +10,16 @@ import { escapeShellArg, getShellConfiguration, ShellExecutionService, + flatMapTextParts, } from '@google/gemini-cli-core'; import type { CommandContext } from '../../ui/commands/types.js'; -import type { IPromptProcessor } from './types.js'; +import type { IPromptProcessor, PromptPipelineContent } from './types.js'; import { SHELL_INJECTION_TRIGGER, SHORTHAND_ARGS_PLACEHOLDER, } from './types.js'; +import { extractInjections, type Injection } from './injectionParser.js'; export class ConfirmationRequiredError extends Error { constructor( @@ -30,15 +32,10 @@ export class ConfirmationRequiredError extends Error { } /** - * Represents a single detected shell injection site in the prompt. + * Represents a single detected shell injection site in the prompt, + * after resolution of arguments. Extends the base Injection interface. */ -interface ShellInjection { - /** The shell command extracted from within !{...}, trimmed. */ - command: string; - /** The starting index of the injection (inclusive, points to '!'). */ - startIndex: number; - /** The ending index of the injection (exclusive, points after '}'). */ - endIndex: number; +interface ResolvedShellInjection extends Injection { /** The command after {{args}} has been escaped and substituted. */ resolvedCommand?: string; } @@ -56,11 +53,25 @@ interface ShellInjection { export class ShellProcessor implements IPromptProcessor { constructor(private readonly commandName: string) {} - async process(prompt: string, context: CommandContext): Promise { + async process( + prompt: PromptPipelineContent, + context: CommandContext, + ): Promise { + return flatMapTextParts(prompt, (text) => + this.processString(text, context), + ); + } + + private async processString( + prompt: string, + context: CommandContext, + ): Promise { const userArgsRaw = context.invocation?.args || ''; if (!prompt.includes(SHELL_INJECTION_TRIGGER)) { - return prompt.replaceAll(SHORTHAND_ARGS_PLACEHOLDER, userArgsRaw); + return [ + { text: prompt.replaceAll(SHORTHAND_ARGS_PLACEHOLDER, userArgsRaw) }, + ]; } const config = context.services.config; @@ -71,26 +82,37 @@ export class ShellProcessor implements IPromptProcessor { } const { sessionShellAllowlist } = context.session; - const injections = this.extractInjections(prompt); + const injections = extractInjections( + prompt, + SHELL_INJECTION_TRIGGER, + this.commandName, + ); + // If extractInjections found no closed blocks (and didn't throw), treat as raw. if (injections.length === 0) { - return prompt.replaceAll(SHORTHAND_ARGS_PLACEHOLDER, userArgsRaw); + return [ + { text: prompt.replaceAll(SHORTHAND_ARGS_PLACEHOLDER, userArgsRaw) }, + ]; } const { shell } = getShellConfiguration(); const userArgsEscaped = escapeShellArg(userArgsRaw, shell); - const resolvedInjections = injections.map((injection) => { - if (injection.command === '') { - return injection; - } - // Replace {{args}} inside the command string with the escaped version. - const resolvedCommand = injection.command.replaceAll( - SHORTHAND_ARGS_PLACEHOLDER, - userArgsEscaped, - ); - return { ...injection, resolvedCommand }; - }); + const resolvedInjections: ResolvedShellInjection[] = injections.map( + (injection) => { + const command = injection.content; + + if (command === '') { + return { ...injection, resolvedCommand: undefined }; + } + + const resolvedCommand = command.replaceAll( + SHORTHAND_ARGS_PLACEHOLDER, + userArgsEscaped, + ); + return { ...injection, resolvedCommand }; + }, + ); const commandsToConfirm = new Set(); for (const injection of resolvedInjections) { @@ -180,69 +202,6 @@ export class ShellProcessor implements IPromptProcessor { userArgsRaw, ); - return processedPrompt; - } - - /** - * Iteratively parses the prompt string to extract shell injections (!{...}), - * correctly handling nested braces within the command. - * - * @param prompt The prompt string to parse. - * @returns An array of extracted ShellInjection objects. - * @throws Error if an unclosed injection (`!{`) is found. - */ - private extractInjections(prompt: string): ShellInjection[] { - const injections: ShellInjection[] = []; - let index = 0; - - while (index < prompt.length) { - const startIndex = prompt.indexOf(SHELL_INJECTION_TRIGGER, index); - - if (startIndex === -1) { - break; - } - - let currentIndex = startIndex + SHELL_INJECTION_TRIGGER.length; - let braceCount = 1; - let foundEnd = false; - - while (currentIndex < prompt.length) { - const char = prompt[currentIndex]; - - // We count literal braces. This parser does not interpret shell quoting/escaping. - if (char === '{') { - braceCount++; - } else if (char === '}') { - braceCount--; - if (braceCount === 0) { - const commandContent = prompt.substring( - startIndex + SHELL_INJECTION_TRIGGER.length, - currentIndex, - ); - const endIndex = currentIndex + 1; - - injections.push({ - command: commandContent.trim(), - startIndex, - endIndex, - }); - - index = endIndex; - foundEnd = true; - break; - } - } - currentIndex++; - } - - // Check if the inner loop finished without finding the closing brace. - if (!foundEnd) { - throw new Error( - `Invalid syntax in command '${this.commandName}': Unclosed shell injection starting at index ${startIndex} ('!{'). Ensure braces are balanced.`, - ); - } - } - - return injections; + return [{ text: processedPrompt }]; } } diff --git a/packages/cli/src/services/prompt-processors/types.ts b/packages/cli/src/services/prompt-processors/types.ts index 94c2e717..c6876574 100644 --- a/packages/cli/src/services/prompt-processors/types.ts +++ b/packages/cli/src/services/prompt-processors/types.ts @@ -5,6 +5,12 @@ */ import type { CommandContext } from '../../ui/commands/types.js'; +import type { PartUnion } from '@google/genai'; + +/** + * Defines the input/output type for prompt processors. + */ +export type PromptPipelineContent = PartUnion[]; /** * Defines the interface for a prompt processor, a module that can transform @@ -13,12 +19,8 @@ import type { CommandContext } from '../../ui/commands/types.js'; */ export interface IPromptProcessor { /** - * Processes a prompt string, applying a specific transformation as part of a pipeline. - * - * Each processor in a command's pipeline receives the output of the previous - * processor. This method provides the full command context, allowing for - * complex transformations that may require access to invocation details, - * application services, or UI state. + * Processes a prompt input (which may contain text and multi-modal parts), + * applying a specific transformation as part of a pipeline. * * @param prompt The current state of the prompt string. This may have been * modified by previous processors in the pipeline. @@ -28,7 +30,10 @@ export interface IPromptProcessor { * @returns A promise that resolves to the transformed prompt string, which * will be passed to the next processor or, if it's the last one, sent to the model. */ - process(prompt: string, context: CommandContext): Promise; + process( + prompt: PromptPipelineContent, + context: CommandContext, + ): Promise; } /** @@ -42,3 +47,8 @@ export const SHORTHAND_ARGS_PLACEHOLDER = '{{args}}'; * The trigger string for shell command injection in custom commands. */ export const SHELL_INJECTION_TRIGGER = '!{'; + +/** + * The trigger string for at file injection in custom commands. + */ +export const AT_FILE_INJECTION_TRIGGER = '@{'; diff --git a/packages/cli/src/ui/commands/types.ts b/packages/cli/src/ui/commands/types.ts index d701419e..32f52a12 100644 --- a/packages/cli/src/ui/commands/types.ts +++ b/packages/cli/src/ui/commands/types.ts @@ -5,7 +5,7 @@ */ import { type ReactNode } from 'react'; -import type { Content } from '@google/genai'; +import type { Content, PartListUnion } from '@google/genai'; import type { HistoryItemWithoutId, HistoryItem } from '../types.js'; import type { Config, GitService, Logger } from '@google/gemini-cli-core'; import type { LoadedSettings } from '../../config/settings.js'; @@ -122,7 +122,7 @@ export interface LoadHistoryActionReturn { */ export interface SubmitPromptActionReturn { type: 'submit_prompt'; - content: string; + content: PartListUnion; } /** diff --git a/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts b/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts index 7099e2e6..b19cd21d 100644 --- a/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts +++ b/packages/cli/src/ui/hooks/slashCommandProcessor.test.ts @@ -491,7 +491,7 @@ describe('useSlashCommandProcessor', () => { description: 'A command from a file', action: async () => ({ type: 'submit_prompt', - content: 'The actual prompt from the TOML file.', + content: [{ text: 'The actual prompt from the TOML file.' }], }), }, CommandKind.FILE, @@ -507,7 +507,7 @@ describe('useSlashCommandProcessor', () => { expect(actionResult).toEqual({ type: 'submit_prompt', - content: 'The actual prompt from the TOML file.', + content: [{ text: 'The actual prompt from the TOML file.' }], }); expect(mockAddItem).toHaveBeenCalledWith( @@ -523,7 +523,7 @@ describe('useSlashCommandProcessor', () => { description: 'A command from mcp', action: async () => ({ type: 'submit_prompt', - content: 'The actual prompt from the mcp command.', + content: [{ text: 'The actual prompt from the mcp command.' }], }), }, CommandKind.MCP_PROMPT, @@ -539,7 +539,7 @@ describe('useSlashCommandProcessor', () => { expect(actionResult).toEqual({ type: 'submit_prompt', - content: 'The actual prompt from the mcp command.', + content: [{ text: 'The actual prompt from the mcp command.' }], }); expect(mockAddItem).toHaveBeenCalledWith( diff --git a/packages/cli/src/ui/types.ts b/packages/cli/src/ui/types.ts index d453bec9..5461b521 100644 --- a/packages/cli/src/ui/types.ts +++ b/packages/cli/src/ui/types.ts @@ -9,6 +9,7 @@ import type { ToolCallConfirmationDetails, ToolResultDisplay, } from '@google/gemini-cli-core'; +import type { PartListUnion } from '@google/genai'; // Only defining the state enum needed by the UI export enum StreamingState { @@ -239,7 +240,7 @@ export interface ConsoleMessageItem { */ export interface SubmitPromptResult { type: 'submit_prompt'; - content: string; + content: PartListUnion; } /** diff --git a/packages/core/index.ts b/packages/core/index.ts index 4b273417..447560d4 100644 --- a/packages/core/index.ts +++ b/packages/core/index.ts @@ -18,3 +18,4 @@ export { IdeConnectionType, } from './src/telemetry/types.js'; export { makeFakeConfig } from './src/test-utils/config.js'; +export * from './src/utils/pathReader.js'; diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index 7778da6b..7f2abd69 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -22,7 +22,7 @@ import { Turn, GeminiEventType } from './turn.js'; import type { Config } from '../config/config.js'; import type { UserTierId } from '../code_assist/types.js'; import { getCoreSystemPrompt, getCompressionPrompt } from './prompts.js'; -import { getResponseText } from '../utils/generateContentResponseUtilities.js'; +import { getResponseText } from '../utils/partUtils.js'; import { checkNextSpeaker } from '../utils/nextSpeakerChecker.js'; import { reportError } from '../utils/errorReporting.js'; import { GeminiChat } from './geminiChat.js'; diff --git a/packages/core/src/core/turn.ts b/packages/core/src/core/turn.ts index 24db4940..ab46ee00 100644 --- a/packages/core/src/core/turn.ts +++ b/packages/core/src/core/turn.ts @@ -18,7 +18,7 @@ import type { ToolResultDisplay, } from '../tools/tools.js'; import type { ToolErrorType } from '../tools/tool-error.js'; -import { getResponseText } from '../utils/generateContentResponseUtilities.js'; +import { getResponseText } from '../utils/partUtils.js'; import { reportError } from '../utils/errorReporting.js'; import { getErrorMessage, diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index c8a2dda1..be437ec1 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -44,7 +44,9 @@ export * from './utils/formatters.js'; export * from './utils/generateContentResponseUtilities.js'; export * from './utils/filesearch/fileSearch.js'; export * from './utils/errorParsing.js'; +export * from './utils/workspaceContext.js'; export * from './utils/ignorePatterns.js'; +export * from './utils/partUtils.js'; // Export services export * from './services/fileDiscoveryService.js'; diff --git a/packages/core/src/tools/web-fetch.ts b/packages/core/src/tools/web-fetch.ts index f4d83410..586ce362 100644 --- a/packages/core/src/tools/web-fetch.ts +++ b/packages/core/src/tools/web-fetch.ts @@ -19,7 +19,7 @@ import { ToolErrorType } from './tool-error.js'; import { getErrorMessage } from '../utils/errors.js'; import type { Config } from '../config/config.js'; import { ApprovalMode } from '../config/config.js'; -import { getResponseText } from '../utils/generateContentResponseUtilities.js'; +import { getResponseText } from '../utils/partUtils.js'; import { fetchWithTimeout, isPrivateIp } from '../utils/fetch.js'; import { convert } from 'html-to-text'; import { ProxyAgent, setGlobalDispatcher } from 'undici'; diff --git a/packages/core/src/tools/web-search.ts b/packages/core/src/tools/web-search.ts index 24bc40bf..b5f57e4b 100644 --- a/packages/core/src/tools/web-search.ts +++ b/packages/core/src/tools/web-search.ts @@ -11,7 +11,7 @@ import { ToolErrorType } from './tool-error.js'; import { getErrorMessage } from '../utils/errors.js'; import type { Config } from '../config/config.js'; -import { getResponseText } from '../utils/generateContentResponseUtilities.js'; +import { getResponseText } from '../utils/partUtils.js'; interface GroundingChunkWeb { uri?: string; diff --git a/packages/core/src/utils/generateContentResponseUtilities.test.ts b/packages/core/src/utils/generateContentResponseUtilities.test.ts index f081e38c..5f99fa8d 100644 --- a/packages/core/src/utils/generateContentResponseUtilities.test.ts +++ b/packages/core/src/utils/generateContentResponseUtilities.test.ts @@ -6,7 +6,6 @@ import { describe, it, expect } from 'vitest'; import { - getResponseText, getResponseTextFromParts, getFunctionCalls, getFunctionCallsFromParts, @@ -69,45 +68,6 @@ const minimalMockResponse = ( }); describe('generateContentResponseUtilities', () => { - describe('getResponseText', () => { - it('should return undefined for no candidates', () => { - expect(getResponseText(minimalMockResponse(undefined))).toBeUndefined(); - }); - it('should return undefined for empty candidates array', () => { - expect(getResponseText(minimalMockResponse([]))).toBeUndefined(); - }); - it('should return undefined for no parts', () => { - const response = mockResponse([]); - expect(getResponseText(response)).toBeUndefined(); - }); - it('should extract text from a single text part', () => { - const response = mockResponse([mockTextPart('Hello')]); - expect(getResponseText(response)).toBe('Hello'); - }); - it('should concatenate text from multiple text parts', () => { - const response = mockResponse([ - mockTextPart('Hello '), - mockTextPart('World'), - ]); - expect(getResponseText(response)).toBe('Hello World'); - }); - it('should ignore function call parts', () => { - const response = mockResponse([ - mockTextPart('Hello '), - mockFunctionCallPart('testFunc'), - mockTextPart('World'), - ]); - expect(getResponseText(response)).toBe('Hello World'); - }); - it('should return undefined if only function call parts exist', () => { - const response = mockResponse([ - mockFunctionCallPart('testFunc'), - mockFunctionCallPart('anotherFunc'), - ]); - expect(getResponseText(response)).toBeUndefined(); - }); - }); - describe('getResponseTextFromParts', () => { it('should return undefined for no parts', () => { expect(getResponseTextFromParts([])).toBeUndefined(); diff --git a/packages/core/src/utils/generateContentResponseUtilities.ts b/packages/core/src/utils/generateContentResponseUtilities.ts index 136b35ed..8d5bec9f 100644 --- a/packages/core/src/utils/generateContentResponseUtilities.ts +++ b/packages/core/src/utils/generateContentResponseUtilities.ts @@ -9,23 +9,7 @@ import type { Part, FunctionCall, } from '@google/genai'; - -export function getResponseText( - response: GenerateContentResponse, -): string | undefined { - const parts = response.candidates?.[0]?.content?.parts; - if (!parts) { - return undefined; - } - const textSegments = parts - .map((part) => part.text) - .filter((text): text is string => typeof text === 'string'); - - if (textSegments.length === 0) { - return undefined; - } - return textSegments.join(''); -} +import { getResponseText } from './partUtils.js'; export function getResponseTextFromParts(parts: Part[]): string | undefined { if (!parts) { diff --git a/packages/core/src/utils/partUtils.test.ts b/packages/core/src/utils/partUtils.test.ts index 5504f85c..d5301072 100644 --- a/packages/core/src/utils/partUtils.test.ts +++ b/packages/core/src/utils/partUtils.test.ts @@ -5,8 +5,13 @@ */ import { describe, it, expect } from 'vitest'; -import { partToString, getResponseText } from './partUtils.js'; -import type { GenerateContentResponse, Part } from '@google/genai'; +import { + partToString, + getResponseText, + flatMapTextParts, + appendToLastTextPart, +} from './partUtils.js'; +import type { GenerateContentResponse, Part, PartUnion } from '@google/genai'; const mockResponse = ( parts?: Array<{ text?: string; functionCall?: unknown }>, @@ -162,5 +167,135 @@ describe('partUtils', () => { const result = mockResponse([]); expect(getResponseText(result)).toBeNull(); }); + + it('should return null if the first candidate has no content property', () => { + const response: GenerateContentResponse = { + candidates: [ + { + index: 0, + }, + ], + promptFeedback: { safetyRatings: [] }, + text: undefined, + data: undefined, + functionCalls: undefined, + executableCode: undefined, + codeExecutionResult: undefined, + }; + expect(getResponseText(response)).toBeNull(); + }); + }); + + describe('flatMapTextParts', () => { + // A simple async transform function that splits a string into character parts. + const splitCharsTransform = async (text: string): Promise => + text.split('').map((char) => ({ text: char })); + + it('should return an empty array for empty input', async () => { + const result = await flatMapTextParts([], splitCharsTransform); + expect(result).toEqual([]); + }); + + it('should transform a simple string input', async () => { + const result = await flatMapTextParts('hi', splitCharsTransform); + expect(result).toEqual([{ text: 'h' }, { text: 'i' }]); + }); + + it('should transform a single text part object', async () => { + const result = await flatMapTextParts( + { text: 'cat' }, + splitCharsTransform, + ); + expect(result).toEqual([{ text: 'c' }, { text: 'a' }, { text: 't' }]); + }); + + it('should transform an array of text parts and flatten the result', async () => { + // A transform that duplicates the text to test the "flatMap" behavior. + const duplicateTransform = async (text: string): Promise => [ + { text: `${text}` }, + { text: `${text}` }, + ]; + const parts = [{ text: 'a' }, { text: 'b' }]; + const result = await flatMapTextParts(parts, duplicateTransform); + expect(result).toEqual([ + { text: 'a' }, + { text: 'a' }, + { text: 'b' }, + { text: 'b' }, + ]); + }); + + it('should pass through non-text parts unmodified', async () => { + const nonTextPart: Part = { functionCall: { name: 'do_stuff' } }; + const result = await flatMapTextParts(nonTextPart, splitCharsTransform); + expect(result).toEqual([nonTextPart]); + }); + + it('should handle a mix of text and non-text parts in an array', async () => { + const nonTextPart: Part = { + inlineData: { mimeType: 'image/jpeg', data: '' }, + }; + const parts: PartUnion[] = [{ text: 'go' }, nonTextPart, ' stop']; + const result = await flatMapTextParts(parts, splitCharsTransform); + expect(result).toEqual([ + { text: 'g' }, + { text: 'o' }, + nonTextPart, // Should be passed through + { text: ' ' }, + { text: 's' }, + { text: 't' }, + { text: 'o' }, + { text: 'p' }, + ]); + }); + + it('should handle a transform that returns an empty array', async () => { + const removeTransform = async (_text: string): Promise => []; + const parts: PartUnion[] = [ + { text: 'remove' }, + { functionCall: { name: 'keep' } }, + ]; + const result = await flatMapTextParts(parts, removeTransform); + expect(result).toEqual([{ functionCall: { name: 'keep' } }]); + }); + }); + + describe('appendToLastTextPart', () => { + it('should append to an empty prompt', () => { + const prompt: PartUnion[] = []; + const result = appendToLastTextPart(prompt, 'new text'); + expect(result).toEqual([{ text: 'new text' }]); + }); + + it('should append to a prompt with a string as the last part', () => { + const prompt: PartUnion[] = ['first part']; + const result = appendToLastTextPart(prompt, 'new text'); + expect(result).toEqual(['first part\n\nnew text']); + }); + + it('should append to a prompt with a text part object as the last part', () => { + const prompt: PartUnion[] = [{ text: 'first part' }]; + const result = appendToLastTextPart(prompt, 'new text'); + expect(result).toEqual([{ text: 'first part\n\nnew text' }]); + }); + + it('should append a new text part if the last part is not a text part', () => { + const nonTextPart: Part = { functionCall: { name: 'do_stuff' } }; + const prompt: PartUnion[] = [nonTextPart]; + const result = appendToLastTextPart(prompt, 'new text'); + expect(result).toEqual([nonTextPart, { text: '\n\nnew text' }]); + }); + + it('should not append anything if the text to append is empty', () => { + const prompt: PartUnion[] = ['first part']; + const result = appendToLastTextPart(prompt, ''); + expect(result).toEqual(['first part']); + }); + + it('should use a custom separator', () => { + const prompt: PartUnion[] = ['first part']; + const result = appendToLastTextPart(prompt, 'new text', '---'); + expect(result).toEqual(['first part---new text']); + }); }); }); diff --git a/packages/core/src/utils/partUtils.ts b/packages/core/src/utils/partUtils.ts index ce42d3b0..f5195c11 100644 --- a/packages/core/src/utils/partUtils.ts +++ b/packages/core/src/utils/partUtils.ts @@ -8,6 +8,7 @@ import type { GenerateContentResponse, PartListUnion, Part, + PartUnion, } from '@google/genai'; /** @@ -87,3 +88,82 @@ export function getResponseText( } return null; } + +/** + * Asynchronously maps over a PartListUnion, applying a transformation function + * to the text content of each text-based part. + * + * @param parts The PartListUnion to process. + * @param transform A function that takes a string of text and returns a Promise + * resolving to an array of new PartUnions. + * @returns A Promise that resolves to a new array of PartUnions with the + * transformations applied. + */ +export async function flatMapTextParts( + parts: PartListUnion, + transform: (text: string) => Promise, +): Promise { + const result: PartUnion[] = []; + const partArray = Array.isArray(parts) + ? parts + : typeof parts === 'string' + ? [{ text: parts }] + : [parts]; + + for (const part of partArray) { + let textToProcess: string | undefined; + if (typeof part === 'string') { + textToProcess = part; + } else if ('text' in part) { + textToProcess = part.text; + } + + if (textToProcess !== undefined) { + const transformedParts = await transform(textToProcess); + result.push(...transformedParts); + } else { + // Pass through non-text parts unmodified. + result.push(part); + } + } + return result; +} + +/** + * Appends a string of text to the last text part of a prompt, or adds a new + * text part if the last part is not a text part. + * + * @param prompt The prompt to modify. + * @param textToAppend The text to append to the prompt. + * @param separator The separator to add between existing text and the new text. + * @returns The modified prompt. + */ +export function appendToLastTextPart( + prompt: PartUnion[], + textToAppend: string, + separator = '\n\n', +): PartUnion[] { + if (!textToAppend) { + return prompt; + } + + if (prompt.length === 0) { + return [{ text: textToAppend }]; + } + + const newPrompt = [...prompt]; + const lastPart = newPrompt.at(-1); + + if (typeof lastPart === 'string') { + newPrompt[newPrompt.length - 1] = `${lastPart}${separator}${textToAppend}`; + } else if (lastPart && 'text' in lastPart) { + newPrompt[newPrompt.length - 1] = { + ...lastPart, + text: `${lastPart.text}${separator}${textToAppend}`, + }; + } else { + newPrompt.push({ text: `${separator}${textToAppend}` }); + } + + return newPrompt; +} diff --git a/packages/core/src/utils/pathReader.test.ts b/packages/core/src/utils/pathReader.test.ts new file mode 100644 index 00000000..45229a67 --- /dev/null +++ b/packages/core/src/utils/pathReader.test.ts @@ -0,0 +1,407 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, afterEach, vi } from 'vitest'; +import mock from 'mock-fs'; +import * as path from 'node:path'; +import { WorkspaceContext } from './workspaceContext.js'; +import { readPathFromWorkspace } from './pathReader.js'; +import type { Config } from '../config/config.js'; +import { StandardFileSystemService } from '../services/fileSystemService.js'; +import type { FileDiscoveryService } from '../services/fileDiscoveryService.js'; + +// --- Helper for creating a mock Config object --- +// We use the actual implementations of WorkspaceContext and FileSystemService +// to test the integration against mock-fs. +const createMockConfig = ( + cwd: string, + otherDirs: string[] = [], + mockFileService?: FileDiscoveryService, +): Config => { + const workspace = new WorkspaceContext(cwd, otherDirs); + const fileSystemService = new StandardFileSystemService(); + return { + getWorkspaceContext: () => workspace, + // TargetDir is used by processSingleFileContent to generate relative paths in errors/output + getTargetDir: () => cwd, + getFileSystemService: () => fileSystemService, + getFileService: () => mockFileService, + } as unknown as Config; +}; + +describe('readPathFromWorkspace', () => { + const CWD = path.resolve('/test/cwd'); + const OTHER_DIR = path.resolve('/test/other'); + const OUTSIDE_DIR = path.resolve('/test/outside'); + + afterEach(() => { + mock.restore(); + vi.resetAllMocks(); + }); + + it('should read a text file from the CWD', async () => { + mock({ + [CWD]: { + 'file.txt': 'hello from cwd', + }, + }); + const mockFileService = { + filterFiles: vi.fn((files) => files), + } as unknown as FileDiscoveryService; + const config = createMockConfig(CWD, [], mockFileService); + const result = await readPathFromWorkspace('file.txt', config); + // Expect [string] for text content + expect(result).toEqual(['hello from cwd']); + expect(mockFileService.filterFiles).toHaveBeenCalled(); + }); + + it('should read a file from a secondary workspace directory', async () => { + mock({ + [CWD]: {}, + [OTHER_DIR]: { + 'file.txt': 'hello from other dir', + }, + }); + const mockFileService = { + filterFiles: vi.fn((files) => files), + } as unknown as FileDiscoveryService; + const config = createMockConfig(CWD, [OTHER_DIR], mockFileService); + const result = await readPathFromWorkspace('file.txt', config); + expect(result).toEqual(['hello from other dir']); + }); + + it('should prioritize CWD when file exists in both CWD and secondary dir', async () => { + mock({ + [CWD]: { + 'file.txt': 'hello from cwd', + }, + [OTHER_DIR]: { + 'file.txt': 'hello from other dir', + }, + }); + const mockFileService = { + filterFiles: vi.fn((files) => files), + } as unknown as FileDiscoveryService; + const config = createMockConfig(CWD, [OTHER_DIR], mockFileService); + const result = await readPathFromWorkspace('file.txt', config); + expect(result).toEqual(['hello from cwd']); + }); + + it('should read an image file and return it as inlineData (Part object)', async () => { + // Use a real PNG header for robustness + const imageData = Buffer.from([ + 0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a, + ]); + mock({ + [CWD]: { + 'image.png': imageData, + }, + }); + const mockFileService = { + filterFiles: vi.fn((files) => files), + } as unknown as FileDiscoveryService; + const config = createMockConfig(CWD, [], mockFileService); + const result = await readPathFromWorkspace('image.png', config); + // Expect [Part] for image content + expect(result).toEqual([ + { + inlineData: { + mimeType: 'image/png', + data: imageData.toString('base64'), + }, + }, + ]); + }); + + it('should read a generic binary file and return an info string', async () => { + // Data that is clearly binary (null bytes) + const binaryData = Buffer.from([0x00, 0x01, 0x02, 0x03]); + mock({ + [CWD]: { + 'data.bin': binaryData, + }, + }); + const mockFileService = { + filterFiles: vi.fn((files) => files), + } as unknown as FileDiscoveryService; + const config = createMockConfig(CWD, [], mockFileService); + const result = await readPathFromWorkspace('data.bin', config); + // Expect [string] containing the skip message from fileUtils + expect(result).toEqual(['Cannot display content of binary file: data.bin']); + }); + + it('should read a file from an absolute path if within workspace', async () => { + const absPath = path.join(OTHER_DIR, 'abs.txt'); + mock({ + [CWD]: {}, + [OTHER_DIR]: { + 'abs.txt': 'absolute content', + }, + }); + const mockFileService = { + filterFiles: vi.fn((files) => files), + } as unknown as FileDiscoveryService; + const config = createMockConfig(CWD, [OTHER_DIR], mockFileService); + const result = await readPathFromWorkspace(absPath, config); + expect(result).toEqual(['absolute content']); + }); + + describe('Directory Expansion', () => { + it('should expand a directory and read the content of its files', async () => { + mock({ + [CWD]: { + 'my-dir': { + 'file1.txt': 'content of file 1', + 'file2.md': 'content of file 2', + }, + }, + }); + const mockFileService = { + filterFiles: vi.fn((files) => files), + } as unknown as FileDiscoveryService; + const config = createMockConfig(CWD, [], mockFileService); + const result = await readPathFromWorkspace('my-dir', config); + + // Convert to a single string for easier, order-independent checking + const resultText = result + .map((p) => { + if (typeof p === 'string') return p; + if (typeof p === 'object' && p && 'text' in p) return p.text; + // This part is important for handling binary/image data which isn't just text + if (typeof p === 'object' && p && 'inlineData' in p) return ''; + return p; + }) + .join(''); + + expect(resultText).toContain( + '--- Start of content for directory: my-dir ---', + ); + expect(resultText).toContain('--- file1.txt ---'); + expect(resultText).toContain('content of file 1'); + expect(resultText).toContain('--- file2.md ---'); + expect(resultText).toContain('content of file 2'); + expect(resultText).toContain( + '--- End of content for directory: my-dir ---', + ); + }); + + it('should recursively expand a directory and read all nested files', async () => { + mock({ + [CWD]: { + 'my-dir': { + 'file1.txt': 'content of file 1', + 'sub-dir': { + 'nested.txt': 'nested content', + }, + }, + }, + }); + const mockFileService = { + filterFiles: vi.fn((files) => files), + } as unknown as FileDiscoveryService; + const config = createMockConfig(CWD, [], mockFileService); + const result = await readPathFromWorkspace('my-dir', config); + + const resultText = result + .map((p) => { + if (typeof p === 'string') return p; + if (typeof p === 'object' && p && 'text' in p) return p.text; + return ''; + }) + .join(''); + + expect(resultText).toContain('content of file 1'); + expect(resultText).toContain('nested content'); + expect(resultText).toContain( + `--- ${path.join('sub-dir', 'nested.txt')} ---`, + ); + }); + + it('should handle mixed content and include files from subdirectories', async () => { + const imageData = Buffer.from([ + 0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a, + ]); + mock({ + [CWD]: { + 'mixed-dir': { + 'info.txt': 'some text', + 'photo.png': imageData, + 'sub-dir': { + 'nested.txt': 'this should be included', + }, + 'empty-sub-dir': {}, + }, + }, + }); + const mockFileService = { + filterFiles: vi.fn((files) => files), + } as unknown as FileDiscoveryService; + const config = createMockConfig(CWD, [], mockFileService); + const result = await readPathFromWorkspace('mixed-dir', config); + + // Check for the text part + const textContent = result + .map((p) => { + if (typeof p === 'string') return p; + if (typeof p === 'object' && p && 'text' in p) return p.text; + return ''; // Ignore non-text parts for this assertion + }) + .join(''); + expect(textContent).toContain('some text'); + expect(textContent).toContain('this should be included'); + + // Check for the image part + const imagePart = result.find( + (p) => typeof p === 'object' && 'inlineData' in p, + ); + expect(imagePart).toEqual({ + inlineData: { + mimeType: 'image/png', + data: imageData.toString('base64'), + }, + }); + }); + + it('should handle an empty directory', async () => { + mock({ + [CWD]: { + 'empty-dir': {}, + }, + }); + const mockFileService = { + filterFiles: vi.fn((files) => files), + } as unknown as FileDiscoveryService; + const config = createMockConfig(CWD, [], mockFileService); + const result = await readPathFromWorkspace('empty-dir', config); + expect(result).toEqual([ + { text: '--- Start of content for directory: empty-dir ---\n' }, + { text: '--- End of content for directory: empty-dir ---' }, + ]); + }); + }); + + describe('File Ignoring', () => { + it('should return an empty array for an ignored file', async () => { + mock({ + [CWD]: { + 'ignored.txt': 'ignored content', + }, + }); + const mockFileService = { + filterFiles: vi.fn(() => []), // Simulate the file being filtered out + } as unknown as FileDiscoveryService; + const config = createMockConfig(CWD, [], mockFileService); + const result = await readPathFromWorkspace('ignored.txt', config); + expect(result).toEqual([]); + expect(mockFileService.filterFiles).toHaveBeenCalledWith( + ['ignored.txt'], + { + respectGitIgnore: true, + respectGeminiIgnore: true, + }, + ); + }); + + it('should not read ignored files when expanding a directory', async () => { + mock({ + [CWD]: { + 'my-dir': { + 'not-ignored.txt': 'visible', + 'ignored.log': 'invisible', + }, + }, + }); + const mockFileService = { + filterFiles: vi.fn((files: string[]) => + files.filter((f) => !f.endsWith('ignored.log')), + ), + } as unknown as FileDiscoveryService; + const config = createMockConfig(CWD, [], mockFileService); + const result = await readPathFromWorkspace('my-dir', config); + const resultText = result + .map((p) => { + if (typeof p === 'string') return p; + if (typeof p === 'object' && p && 'text' in p) return p.text; + return ''; + }) + .join(''); + + expect(resultText).toContain('visible'); + expect(resultText).not.toContain('invisible'); + expect(mockFileService.filterFiles).toHaveBeenCalled(); + }); + }); + + it('should throw an error for an absolute path outside the workspace', async () => { + const absPath = path.join(OUTSIDE_DIR, 'secret.txt'); + mock({ + [CWD]: {}, + [OUTSIDE_DIR]: { + 'secret.txt': 'secrets', + }, + }); + // OUTSIDE_DIR is not added to the config's workspace + const config = createMockConfig(CWD); + await expect(readPathFromWorkspace(absPath, config)).rejects.toThrow( + `Absolute path is outside of the allowed workspace: ${absPath}`, + ); + }); + + it('should throw an error if a relative path is not found anywhere', async () => { + mock({ + [CWD]: {}, + [OTHER_DIR]: {}, + }); + const config = createMockConfig(CWD, [OTHER_DIR]); + await expect( + readPathFromWorkspace('not-found.txt', config), + ).rejects.toThrow('Path not found in workspace: not-found.txt'); + }); + + // mock-fs permission simulation is unreliable on Windows. + it.skipIf(process.platform === 'win32')( + 'should return an error string if reading a file with no permissions', + async () => { + mock({ + [CWD]: { + 'unreadable.txt': mock.file({ + content: 'you cannot read me', + mode: 0o222, // Write-only + }), + }, + }); + const mockFileService = { + filterFiles: vi.fn((files) => files), + } as unknown as FileDiscoveryService; + const config = createMockConfig(CWD, [], mockFileService); + // processSingleFileContent catches the error and returns an error string. + const result = await readPathFromWorkspace('unreadable.txt', config); + const textResult = result[0] as string; + + // processSingleFileContent formats errors using the relative path from the target dir (CWD). + expect(textResult).toContain('Error reading file unreadable.txt'); + expect(textResult).toMatch(/(EACCES|permission denied)/i); + }, + ); + + it('should return an error string for files exceeding the size limit', async () => { + // Mock a file slightly larger than the 20MB limit defined in fileUtils.ts + const largeContent = 'a'.repeat(21 * 1024 * 1024); // 21MB + mock({ + [CWD]: { + 'large.txt': largeContent, + }, + }); + const mockFileService = { + filterFiles: vi.fn((files) => files), + } as unknown as FileDiscoveryService; + const config = createMockConfig(CWD, [], mockFileService); + const result = await readPathFromWorkspace('large.txt', config); + const textResult = result[0] as string; + // The error message comes directly from processSingleFileContent + expect(textResult).toBe('File size exceeds the 20MB limit.'); + }); +}); diff --git a/packages/core/src/utils/pathReader.ts b/packages/core/src/utils/pathReader.ts new file mode 100644 index 00000000..1b177848 --- /dev/null +++ b/packages/core/src/utils/pathReader.ts @@ -0,0 +1,118 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { promises as fs } from 'node:fs'; +import path from 'node:path'; +import { glob } from 'glob'; +import type { PartUnion } from '@google/genai'; +import { processSingleFileContent } from './fileUtils.js'; +import type { Config } from '../config/config.js'; + +/** + * Reads the content of a file or recursively expands a directory from + * within the workspace, returning content suitable for LLM input. + * + * @param pathStr The path to read (can be absolute or relative). + * @param config The application configuration, providing workspace context and services. + * @returns A promise that resolves to an array of PartUnion (string | Part). + * @throws An error if the path is not found or is outside the workspace. + */ +export async function readPathFromWorkspace( + pathStr: string, + config: Config, +): Promise { + const workspace = config.getWorkspaceContext(); + const fileService = config.getFileService(); + let absolutePath: string | null = null; + + if (path.isAbsolute(pathStr)) { + if (!workspace.isPathWithinWorkspace(pathStr)) { + throw new Error( + `Absolute path is outside of the allowed workspace: ${pathStr}`, + ); + } + absolutePath = pathStr; + } else { + // Prioritized search for relative paths. + const searchDirs = workspace.getDirectories(); + for (const dir of searchDirs) { + const potentialPath = path.resolve(dir, pathStr); + try { + await fs.access(potentialPath); + absolutePath = potentialPath; + break; // Found the first match. + } catch { + // Not found, continue to the next directory. + } + } + } + + if (!absolutePath) { + throw new Error(`Path not found in workspace: ${pathStr}`); + } + + const stats = await fs.stat(absolutePath); + if (stats.isDirectory()) { + const allParts: PartUnion[] = []; + allParts.push({ + text: `--- Start of content for directory: ${pathStr} ---\n`, + }); + + // Use glob to recursively find all files within the directory. + const files = await glob('**/*', { + cwd: absolutePath, + nodir: true, // We only want files + dot: true, // Include dotfiles + absolute: true, + }); + + const relativeFiles = files.map((p) => + path.relative(config.getTargetDir(), p), + ); + const filteredFiles = fileService.filterFiles(relativeFiles, { + respectGitIgnore: true, + respectGeminiIgnore: true, + }); + const finalFiles = filteredFiles.map((p) => + path.resolve(config.getTargetDir(), p), + ); + + for (const filePath of finalFiles) { + const relativePathForDisplay = path.relative(absolutePath, filePath); + allParts.push({ text: `--- ${relativePathForDisplay} ---\n` }); + const result = await processSingleFileContent( + filePath, + config.getTargetDir(), + config.getFileSystemService(), + ); + allParts.push(result.llmContent); + allParts.push({ text: '\n' }); // Add a newline for separation + } + + allParts.push({ text: `--- End of content for directory: ${pathStr} ---` }); + return allParts; + } else { + // It's a single file, check if it's ignored. + const relativePath = path.relative(config.getTargetDir(), absolutePath); + const filtered = fileService.filterFiles([relativePath], { + respectGitIgnore: true, + respectGeminiIgnore: true, + }); + + if (filtered.length === 0) { + // File is ignored, return empty array to silently skip. + return []; + } + + // It's a single file, process it directly. + const result = await processSingleFileContent( + absolutePath, + config.getTargetDir(), + config.getFileSystemService(), + ); + return [result.llmContent]; + } +}