From 9ee5eadac0d1763e502be33682e9f34bdfe58454 Mon Sep 17 00:00:00 2001 From: Allen Hutchison Date: Wed, 28 May 2025 17:08:05 -0700 Subject: [PATCH] fix(cli): Support multiple @file references in atCommandProcessor (#590) --- .../src/ui/hooks/atCommandProcessor.test.ts | 378 ++++++++++----- .../cli/src/ui/hooks/atCommandProcessor.ts | 447 +++++++++++------- 2 files changed, 554 insertions(+), 271 deletions(-) diff --git a/packages/cli/src/ui/hooks/atCommandProcessor.test.ts b/packages/cli/src/ui/hooks/atCommandProcessor.test.ts index 966fd0fa..ed37ec14 100644 --- a/packages/cli/src/ui/hooks/atCommandProcessor.test.ts +++ b/packages/cli/src/ui/hooks/atCommandProcessor.test.ts @@ -6,22 +6,20 @@ import { describe, it, expect, vi, beforeEach, afterEach, Mock } from 'vitest'; import { handleAtCommand } from './atCommandProcessor.js'; -import { Config, ToolResult } from '@gemini-code/server'; -import { ToolCallStatus } from '../types.js'; // Adjusted import -import { /* PartListUnion, */ Part } from '@google/genai'; // Removed PartListUnion +import { Config } from '@gemini-code/server'; +import { ToolCallStatus } from '../types.js'; import { UseHistoryManagerReturn } from './useHistoryManager.js'; -import * as fsPromises from 'fs/promises'; // Import for mocking stat -import type { Stats } from 'fs'; // Import Stats type for mocking +import * as fsPromises from 'fs/promises'; +import type { Stats } from 'fs'; -// Mock Config and ToolRegistry const mockGetToolRegistry = vi.fn(); const mockGetTargetDir = vi.fn(); const mockConfig = { getToolRegistry: mockGetToolRegistry, getTargetDir: mockGetTargetDir, + isSandboxed: vi.fn(() => false), } as unknown as Config; -// Mock read_many_files tool const mockReadManyFilesExecute = vi.fn(); const mockReadManyFilesTool = { name: 'read_many_files', @@ -31,7 +29,14 @@ const mockReadManyFilesTool = { getDescription: vi.fn((params) => `Read files: ${params.paths.join(', ')}`), }; -// Mock addItem from useHistoryManager +const mockGlobExecute = vi.fn(); +const mockGlobTool = { + name: 'glob', + displayName: 'Glob Tool', + execute: mockGlobExecute, + getDescription: vi.fn(() => 'Glob tool description'), +}; + const mockAddItem: Mock = vi.fn(); const mockOnDebugMessage: Mock<(message: string) => void> = vi.fn(); @@ -39,7 +44,7 @@ vi.mock('fs/promises', async () => { const actual = await vi.importActual('fs/promises'); return { ...actual, - stat: vi.fn(), // Mock stat here + stat: vi.fn(), }; }); @@ -52,20 +57,26 @@ describe('handleAtCommand', () => { mockGetTargetDir.mockReturnValue('/test/dir'); mockGetToolRegistry.mockReturnValue({ getTool: vi.fn((toolName: string) => { - if (toolName === 'read_many_files') { - return mockReadManyFilesTool; - } + if (toolName === 'read_many_files') return mockReadManyFilesTool; + if (toolName === 'glob') return mockGlobTool; return undefined; }), }); - // Default mock for fs.stat if not overridden by a specific test vi.mocked(fsPromises.stat).mockResolvedValue({ isDirectory: () => false, - } as unknown as Stats); + } as Stats); + mockReadManyFilesExecute.mockResolvedValue({ + llmContent: '', + returnDisplay: '', + }); + mockGlobExecute.mockResolvedValue({ + llmContent: 'No files found', + returnDisplay: '', + }); }); afterEach(() => { - abortController.abort(); // Ensure any pending operations are cancelled + abortController.abort(); }); it('should pass through query if no @ command is present', async () => { @@ -78,7 +89,6 @@ describe('handleAtCommand', () => { messageId: 123, signal: abortController.signal, }); - expect(mockAddItem).toHaveBeenCalledWith( { type: 'user', text: query }, 123, @@ -88,28 +98,24 @@ describe('handleAtCommand', () => { expect(mockReadManyFilesExecute).not.toHaveBeenCalled(); }); - it('should pass through query if only a lone @ symbol is present', async () => { + it('should pass through original query if only a lone @ symbol is present', async () => { const queryWithSpaces = ' @ '; - // const trimmedQuery = queryWithSpaces.trim(); // Not needed for addItem expectation here const result = await handleAtCommand({ - query: queryWithSpaces, // Pass the version with spaces to the function + query: queryWithSpaces, config: mockConfig, addItem: mockAddItem, onDebugMessage: mockOnDebugMessage, messageId: 124, signal: abortController.signal, }); - - // For a lone '@', addItem is called with the *original untrimmed* query expect(mockAddItem).toHaveBeenCalledWith( { type: 'user', text: queryWithSpaces }, 124, ); - // processedQuery should also be the original untrimmed version for lone @ expect(result.processedQuery).toEqual([{ text: queryWithSpaces }]); expect(result.shouldProceed).toBe(true); expect(mockOnDebugMessage).toHaveBeenCalledWith( - 'Lone @ detected, passing directly to LLM.', + 'Lone @ detected, will be treated as text in the modified query.', ); }); @@ -118,10 +124,11 @@ describe('handleAtCommand', () => { const query = `@${filePath}`; const fileContent = 'This is the file content.'; mockReadManyFilesExecute.mockResolvedValue({ - llmContent: fileContent, + llmContent: ` +--- ${filePath} --- +${fileContent}`, returnDisplay: 'Read 1 file.', - } as ToolResult); - // fs.stat will use the default mock (isDirectory: false) + }); const result = await handleAtCommand({ query, @@ -131,7 +138,6 @@ describe('handleAtCommand', () => { messageId: 125, signal: abortController.signal, }); - expect(mockAddItem).toHaveBeenCalledWith( { type: 'user', text: query }, 125, @@ -143,20 +149,16 @@ describe('handleAtCommand', () => { expect(mockAddItem).toHaveBeenCalledWith( expect.objectContaining({ type: 'tool_group', - tools: expect.arrayContaining([ - expect.objectContaining({ - name: 'Read Many Files', - status: ToolCallStatus.Success, - resultDisplay: 'Read 1 file.', - }), - ]), + tools: [expect.objectContaining({ status: ToolCallStatus.Success })], }), 125, ); expect(result.processedQuery).toEqual([ - '\n--- Content from: ${contentLabel} ---\n', - fileContent, - '\n--- End of content ---\n', + { text: `@${filePath}` }, + { text: '\n--- Content from referenced files ---' }, + { text: `\nContent from @${filePath}:\n` }, + { text: fileContent }, + { text: '\n--- End of content ---' }, ]); expect(result.shouldProceed).toBe(true); }); @@ -164,19 +166,17 @@ describe('handleAtCommand', () => { it('should process a valid directory path and convert to glob', async () => { const dirPath = 'path/to/dir'; const query = `@${dirPath}`; - const dirContent = [ - 'Content of file 1.', - 'Content of file 2.', - 'Content of file 3.', - ]; + const resolvedGlob = `${dirPath}/**`; + const fileContent = 'Directory content.'; vi.mocked(fsPromises.stat).mockResolvedValue({ isDirectory: () => true, - } as unknown as Stats); - + } as Stats); mockReadManyFilesExecute.mockResolvedValue({ - llmContent: dirContent, + llmContent: ` +--- ${resolvedGlob} --- +${fileContent}`, returnDisplay: 'Read directory contents.', - } as ToolResult); + }); const result = await handleAtCommand({ query, @@ -186,41 +186,40 @@ describe('handleAtCommand', () => { messageId: 126, signal: abortController.signal, }); - expect(mockAddItem).toHaveBeenCalledWith( { type: 'user', text: query }, 126, ); expect(mockReadManyFilesExecute).toHaveBeenCalledWith( - { paths: [`${dirPath}/**`] }, // Expect glob pattern + { paths: [resolvedGlob] }, abortController.signal, ); expect(mockOnDebugMessage).toHaveBeenCalledWith( - `Path resolved to directory, using glob: ${dirPath}/**`, - ); - expect(mockAddItem).toHaveBeenCalledWith( - expect.objectContaining({ type: 'tool_group' }), - 126, + `Path ${dirPath} resolved to directory, using glob: ${resolvedGlob}`, ); expect(result.processedQuery).toEqual([ - '\n--- Content from: ${contentLabel} ---\n', - ...dirContent, - '\n--- End of content ---\n', + { text: `@${resolvedGlob}` }, + { text: '\n--- Content from referenced files ---' }, + { text: `\nContent from @${resolvedGlob}:\n` }, + { text: fileContent }, + { text: '\n--- End of content ---' }, ]); expect(result.shouldProceed).toBe(true); }); - it('should process a valid image file path', async () => { + it('should process a valid image file path (as text content for now)', async () => { const imagePath = 'path/to/image.png'; const query = `@${imagePath}`; - const imageData: Part = { - inlineData: { mimeType: 'image/png', data: 'base64imagedata' }, - }; + // For @-commands, read_many_files is expected to return text or structured text. + // If it were to return actual image Part, the test and handling would be different. + // Current implementation of read_many_files for images returns base64 in text. + const imageFileTextContent = '[base64 image data for path/to/image.png]'; mockReadManyFilesExecute.mockResolvedValue({ - llmContent: [imageData], + llmContent: ` +--- ${imagePath} --- +${imageFileTextContent}`, returnDisplay: 'Read 1 image.', - } as ToolResult); - // fs.stat will use the default mock (isDirectory: false) + }); const result = await handleAtCommand({ query, @@ -230,47 +229,28 @@ describe('handleAtCommand', () => { messageId: 127, signal: abortController.signal, }); - - expect(mockAddItem).toHaveBeenCalledWith( - { type: 'user', text: query }, - 127, - ); - expect(mockReadManyFilesExecute).toHaveBeenCalledWith( - { paths: [imagePath] }, - abortController.signal, - ); - expect(mockAddItem).toHaveBeenCalledWith( - expect.objectContaining({ - type: 'tool_group', - tools: expect.arrayContaining([ - expect.objectContaining({ - name: 'Read Many Files', - status: ToolCallStatus.Success, - resultDisplay: 'Read 1 image.', - }), - ]), - }), - 127, - ); expect(result.processedQuery).toEqual([ - '\n--- Content from: ${contentLabel} ---\n', - imageData, - '\n--- End of content ---\n', + { text: `@${imagePath}` }, + { text: '\n--- Content from referenced files ---' }, + { text: `\nContent from @${imagePath}:\n` }, + { text: imageFileTextContent }, + { text: '\n--- End of content ---' }, ]); expect(result.shouldProceed).toBe(true); }); it('should handle query with text before and after @command', async () => { - const textBefore = 'Explain this:'; + const textBefore = 'Explain this: '; const filePath = 'doc.md'; - const textAfter = 'in detail.'; - const query = `${textBefore} @${filePath} ${textAfter}`; + const textAfter = ' in detail.'; + const query = `${textBefore}@${filePath}${textAfter}`; const fileContent = 'Markdown content.'; mockReadManyFilesExecute.mockResolvedValue({ - llmContent: fileContent, + llmContent: ` +--- ${filePath} --- +${fileContent}`, returnDisplay: 'Read 1 doc.', - } as ToolResult); - // fs.stat will use the default mock (isDirectory: false) + }); const result = await handleAtCommand({ query, @@ -280,17 +260,16 @@ describe('handleAtCommand', () => { messageId: 128, signal: abortController.signal, }); - expect(mockAddItem).toHaveBeenCalledWith( - { type: 'user', text: query }, // Expect original query for addItem + { type: 'user', text: query }, 128, ); expect(result.processedQuery).toEqual([ - { text: textBefore }, - '\n--- Content from: ${contentLabel} ---\n', - fileContent, - '\n--- End of content ---\n', - { text: textAfter }, + { text: `${textBefore}@${filePath}${textAfter}` }, + { text: '\n--- Content from referenced files ---' }, + { text: `\nContent from @${filePath}:\n` }, + { text: fileContent }, + { text: '\n--- End of content ---' }, ]); expect(result.shouldProceed).toBe(true); }); @@ -301,10 +280,11 @@ describe('handleAtCommand', () => { const query = `@${rawPath}`; const fileContent = 'Content of file with space.'; mockReadManyFilesExecute.mockResolvedValue({ - llmContent: fileContent, + llmContent: ` +--- ${unescapedPath} --- +${fileContent}`, returnDisplay: 'Read 1 file.', - } as ToolResult); - // fs.stat will use the default mock (isDirectory: false) + }); await handleAtCommand({ query, @@ -314,10 +294,196 @@ describe('handleAtCommand', () => { messageId: 129, signal: abortController.signal, }); - expect(mockReadManyFilesExecute).toHaveBeenCalledWith( - { paths: [unescapedPath] }, // Expect unescaped path + { paths: [unescapedPath] }, abortController.signal, ); }); + + it('should handle multiple @file references', async () => { + const file1 = 'file1.txt'; + const content1 = 'Content file1'; + const file2 = 'file2.md'; + const content2 = 'Content file2'; + const query = `@${file1} @${file2}`; + + mockReadManyFilesExecute.mockResolvedValue({ + llmContent: ` +--- ${file1} --- +${content1} +--- ${file2} --- +${content2}`, + returnDisplay: 'Read 2 files.', + }); + + const result = await handleAtCommand({ + query, + config: mockConfig, + addItem: mockAddItem, + onDebugMessage: mockOnDebugMessage, + messageId: 130, + signal: abortController.signal, + }); + expect(mockReadManyFilesExecute).toHaveBeenCalledWith( + { paths: [file1, file2] }, + abortController.signal, + ); + expect(result.processedQuery).toEqual([ + { text: `@${file1} @${file2}` }, + { text: '\n--- Content from referenced files ---' }, + { text: `\nContent from @${file1}:\n` }, + { text: content1 }, + { text: `\nContent from @${file2}:\n` }, + { text: content2 }, + { text: '\n--- End of content ---' }, + ]); + expect(result.shouldProceed).toBe(true); + }); + + it('should handle multiple @file references with interleaved text', async () => { + const text1 = 'Check '; + const file1 = 'f1.txt'; + const content1 = 'C1'; + const text2 = ' and '; + const file2 = 'f2.md'; + const content2 = 'C2'; + const text3 = ' please.'; + const query = `${text1}@${file1}${text2}@${file2}${text3}`; + + mockReadManyFilesExecute.mockResolvedValue({ + llmContent: ` +--- ${file1} --- +${content1} +--- ${file2} --- +${content2}`, + returnDisplay: 'Read 2 files.', + }); + + const result = await handleAtCommand({ + query, + config: mockConfig, + addItem: mockAddItem, + onDebugMessage: mockOnDebugMessage, + messageId: 131, + signal: abortController.signal, + }); + expect(mockReadManyFilesExecute).toHaveBeenCalledWith( + { paths: [file1, file2] }, + abortController.signal, + ); + expect(result.processedQuery).toEqual([ + { text: `${text1}@${file1}${text2}@${file2}${text3}` }, + { text: '\n--- Content from referenced files ---' }, + { text: `\nContent from @${file1}:\n` }, + { text: content1 }, + { text: `\nContent from @${file2}:\n` }, + { text: content2 }, + { text: '\n--- End of content ---' }, + ]); + expect(result.shouldProceed).toBe(true); + }); + + it('should handle a mix of valid, invalid, and lone @ references', async () => { + const file1 = 'valid1.txt'; + const content1 = 'Valid content 1'; + const invalidFile = 'nonexistent.txt'; + const query = `Look at @${file1} then @${invalidFile} and also just @ symbol, then @valid2.glob`; + const file2Glob = 'valid2.glob'; + const resolvedFile2 = 'resolved/valid2.actual'; + const content2 = 'Globbed content'; + + // Mock fs.stat for file1 (valid) + vi.mocked(fsPromises.stat).mockImplementation(async (p) => { + if (p.toString().endsWith(file1)) + return { isDirectory: () => false } as Stats; + if (p.toString().endsWith(invalidFile)) + throw Object.assign(new Error('ENOENT'), { code: 'ENOENT' }); + // For valid2.glob, stat will fail, triggering glob + if (p.toString().endsWith(file2Glob)) + throw Object.assign(new Error('ENOENT'), { code: 'ENOENT' }); + return { isDirectory: () => false } as Stats; // Default + }); + + // Mock glob to find resolvedFile2 for valid2.glob + mockGlobExecute.mockImplementation(async (params) => { + if (params.pattern.includes('valid2.glob')) { + return { + llmContent: `Found files:\n${mockGetTargetDir()}/${resolvedFile2}`, + returnDisplay: 'Found 1 file', + }; + } + return { llmContent: 'No files found', returnDisplay: '' }; + }); + + mockReadManyFilesExecute.mockResolvedValue({ + llmContent: ` +--- ${file1} --- +${content1} +--- ${resolvedFile2} --- +${content2}`, + returnDisplay: 'Read 2 files.', + }); + + const result = await handleAtCommand({ + query, + config: mockConfig, + addItem: mockAddItem, + onDebugMessage: mockOnDebugMessage, + messageId: 132, + signal: abortController.signal, + }); + + expect(mockReadManyFilesExecute).toHaveBeenCalledWith( + { paths: [file1, resolvedFile2] }, + abortController.signal, + ); + expect(result.processedQuery).toEqual([ + // Original query has @nonexistent.txt and @, but resolved has @resolved/valid2.actual + { + text: `Look at @${file1} then @${invalidFile} and also just @ symbol, then @${resolvedFile2}`, + }, + { text: '\n--- Content from referenced files ---' }, + { text: `\nContent from @${file1}:\n` }, + { text: content1 }, + { text: `\nContent from @${resolvedFile2}:\n` }, + { text: content2 }, + { text: '\n--- End of content ---' }, + ]); + expect(result.shouldProceed).toBe(true); + expect(mockOnDebugMessage).toHaveBeenCalledWith( + `Path ${invalidFile} not found directly, attempting glob search.`, + ); + expect(mockOnDebugMessage).toHaveBeenCalledWith( + `Glob search for '**/*${invalidFile}*' found no files or an error. Path ${invalidFile} will be skipped.`, + ); + expect(mockOnDebugMessage).toHaveBeenCalledWith( + 'Lone @ detected, will be treated as text in the modified query.', + ); + }); + + it('should return original query if all @paths are invalid or lone @', async () => { + const query = 'Check @nonexistent.txt and @ also'; + vi.mocked(fsPromises.stat).mockRejectedValue( + Object.assign(new Error('ENOENT'), { code: 'ENOENT' }), + ); + mockGlobExecute.mockResolvedValue({ + llmContent: 'No files found', + returnDisplay: '', + }); + + const result = await handleAtCommand({ + query, + config: mockConfig, + addItem: mockAddItem, + onDebugMessage: mockOnDebugMessage, + messageId: 133, + signal: abortController.signal, + }); + expect(mockReadManyFilesExecute).not.toHaveBeenCalled(); + // The modified query string will be "Check @nonexistent.txt and @ also" because no paths were resolved for reading. + expect(result.processedQuery).toEqual([ + { text: 'Check @nonexistent.txt and @ also' }, + ]); + expect(result.shouldProceed).toBe(true); + }); }); diff --git a/packages/cli/src/ui/hooks/atCommandProcessor.ts b/packages/cli/src/ui/hooks/atCommandProcessor.ts index dd97a0d6..4416ef82 100644 --- a/packages/cli/src/ui/hooks/atCommandProcessor.ts +++ b/packages/cli/src/ui/hooks/atCommandProcessor.ts @@ -34,53 +34,82 @@ interface HandleAtCommandResult { shouldProceed: boolean; } -/** - * Parses a query string to find the first '@' command, - * handling \ escaped spaces within the path. - */ -function parseAtCommand( - query: string, -): { textBefore: string; atPath: string; textAfter: string } | null { - let atIndex = -1; - for (let i = 0; i < query.length; i++) { - if (query[i] === '@' && (i === 0 || query[i - 1] !== '\\')) { - atIndex = i; - break; - } - } - - if (atIndex === -1) { - return null; - } - - const textBefore = query.substring(0, atIndex).trim(); - let pathEndIndex = atIndex + 1; - let inEscape = false; - - while (pathEndIndex < query.length) { - const char = query[pathEndIndex]; - if (inEscape) { - inEscape = false; - } else if (char === '\\') { - inEscape = true; - } else if (/\s/.test(char)) { - break; - } - pathEndIndex++; - } - - const rawAtPath = query.substring(atIndex, pathEndIndex); - const textAfter = query.substring(pathEndIndex).trim(); - const atPath = unescapePath(rawAtPath); - - return { textBefore, atPath, textAfter }; +interface AtCommandPart { + type: 'text' | 'atPath'; + content: string; } /** - * Processes user input potentially containing an '@' command. - * If found, it attempts to read the specified file/directory using the - * 'read_many_files' tool, adds the user query and tool result/error to history, - * and prepares the content for the LLM. + * Parses a query string to find all '@' commands and text segments. + * Handles \ escaped spaces within paths. + */ +function parseAllAtCommands(query: string): AtCommandPart[] { + const parts: AtCommandPart[] = []; + let currentIndex = 0; + + while (currentIndex < query.length) { + let atIndex = -1; + let nextSearchIndex = currentIndex; + // Find next unescaped '@' + while (nextSearchIndex < query.length) { + if ( + query[nextSearchIndex] === '@' && + (nextSearchIndex === 0 || query[nextSearchIndex - 1] !== '\\') + ) { + atIndex = nextSearchIndex; + break; + } + nextSearchIndex++; + } + + if (atIndex === -1) { + // No more @ + if (currentIndex < query.length) { + parts.push({ type: 'text', content: query.substring(currentIndex) }); + } + break; + } + + // Add text before @ + if (atIndex > currentIndex) { + parts.push({ + type: 'text', + content: query.substring(currentIndex, atIndex), + }); + } + + // Parse @path + let pathEndIndex = atIndex + 1; + let inEscape = false; + while (pathEndIndex < query.length) { + const char = query[pathEndIndex]; + if (inEscape) { + inEscape = false; + } else if (char === '\\') { + inEscape = true; + } else if (/\s/.test(char)) { + // Path ends at first whitespace not escaped + break; + } + pathEndIndex++; + } + const rawAtPath = query.substring(atIndex, pathEndIndex); + // unescapePath expects the @ symbol to be present, and will handle it. + const atPath = unescapePath(rawAtPath); + parts.push({ type: 'atPath', content: atPath }); + currentIndex = pathEndIndex; + } + // Filter out empty text parts that might result from consecutive @paths or leading/trailing spaces + return parts.filter( + (part) => !(part.type === 'text' && part.content.trim() === ''), + ); +} + +/** + * Processes user input potentially containing one or more '@' commands. + * If found, it attempts to read the specified files/directories using the + * 'read_many_files' tool. The user query is modified to include resolved paths, + * and the content of the files is appended in a structured block. * * @returns An object indicating whether the main hook should proceed with an * LLM call and the processed query parts (including file content). @@ -93,42 +122,25 @@ export async function handleAtCommand({ messageId: userMessageTimestamp, signal, }: HandleAtCommandParams): Promise { - const trimmedQuery = query.trim(); - const parsedCommand = parseAtCommand(trimmedQuery); + const commandParts = parseAllAtCommands(query); + const atPathCommandParts = commandParts.filter( + (part) => part.type === 'atPath', + ); - // If no @ command, add user query normally and proceed to LLM - if (!parsedCommand) { + if (atPathCommandParts.length === 0) { addItem({ type: 'user', text: query }, userMessageTimestamp); return { processedQuery: [{ text: query }], shouldProceed: true }; } - const { textBefore, atPath, textAfter } = parsedCommand; - - // Add the original user query to history first addItem({ type: 'user', text: query }, userMessageTimestamp); - // If the atPath is just "@", pass the original query to the LLM - if (atPath === '@') { - onDebugMessage('Lone @ detected, passing directly to LLM.'); - return { processedQuery: [{ text: query }], shouldProceed: true }; - } - - const pathPart = atPath.substring(1); // Remove leading '@' - - // This error condition is for cases where pathPart becomes empty *after* the initial "@" check, - // which is unlikely with the current parser but good for robustness. - if (!pathPart) { - addItem( - { type: 'error', text: 'Error: No valid path specified after @ symbol.' }, - userMessageTimestamp, - ); - return { processedQuery: null, shouldProceed: false }; - } - - const contentLabel = pathPart; + const pathSpecsToRead: string[] = []; + const atPathToResolvedSpecMap = new Map(); + const contentLabelsForDisplay: string[] = []; const toolRegistry = config.getToolRegistry(); const readManyFilesTool = toolRegistry.getTool('read_many_files'); + const globTool = toolRegistry.getTool('glob'); if (!readManyFilesTool) { addItem( @@ -138,127 +150,237 @@ export async function handleAtCommand({ return { processedQuery: null, shouldProceed: false }; } - // Determine path spec (file or directory glob) - let pathSpec = pathPart; - try { - const absolutePath = path.resolve(config.getTargetDir(), pathPart); - const stats = await fs.stat(absolutePath); - if (stats.isDirectory()) { - pathSpec = pathPart.endsWith('/') ? `${pathPart}**` : `${pathPart}/**`; - onDebugMessage(`Path resolved to directory, using glob: ${pathSpec}`); - } else { - onDebugMessage(`Path resolved to file: ${pathSpec}`); - } - } catch (error) { - if (isNodeError(error) && error.code === 'ENOENT') { + for (const atPathPart of atPathCommandParts) { + const originalAtPath = atPathPart.content; // e.g., "@file.txt" or "@" + + if (originalAtPath === '@') { onDebugMessage( - `Path ${pathPart} not found directly, attempting glob search.`, + 'Lone @ detected, will be treated as text in the modified query.', ); - const globTool = toolRegistry.getTool('glob'); - if (globTool) { - try { - const globResult = await globTool.execute( - { - pattern: `**/*${pathPart}*`, - path: config.getTargetDir(), // Ensure glob searches from the root - }, - signal, - ); - // Assuming llmContent contains the list of files or a "no files found" message. - // And that paths are absolute. - if ( - globResult.llmContent && - typeof globResult.llmContent === 'string' && - !globResult.llmContent.startsWith('No files found') && - !globResult.llmContent.startsWith('Error:') - ) { - // Extract the first line after the header - const lines = globResult.llmContent.split('\n'); - if (lines.length > 1 && lines[1]) { - const firstMatchAbsolute = lines[1].trim(); - // Convert absolute path from glob to relative path for read_many_files - pathSpec = path.relative( - config.getTargetDir(), - firstMatchAbsolute, - ); - onDebugMessage( - `Glob search found ${firstMatchAbsolute}, using relative path: ${pathSpec}`, - ); + continue; + } + + const pathName = originalAtPath.substring(1); + if (!pathName) { + // This case should ideally not be hit if parseAllAtCommands ensures content after @ + // but as a safeguard: + addItem( + { + type: 'error', + text: `Error: Invalid @ command '${originalAtPath}'. No path specified.`, + }, + userMessageTimestamp, + ); + // Decide if this is a fatal error for the whole command or just skip this @ part + // For now, let's be strict and fail the command if one @path is malformed. + return { processedQuery: null, shouldProceed: false }; + } + + let currentPathSpec = pathName; + let resolvedSuccessfully = false; + + try { + const absolutePath = path.resolve(config.getTargetDir(), pathName); + const stats = await fs.stat(absolutePath); + if (stats.isDirectory()) { + currentPathSpec = pathName.endsWith('/') + ? `${pathName}**` + : `${pathName}/**`; + onDebugMessage( + `Path ${pathName} resolved to directory, using glob: ${currentPathSpec}`, + ); + } else { + onDebugMessage(`Path ${pathName} resolved to file: ${currentPathSpec}`); + } + resolvedSuccessfully = true; + } catch (error) { + if (isNodeError(error) && error.code === 'ENOENT') { + onDebugMessage( + `Path ${pathName} not found directly, attempting glob search.`, + ); + if (globTool) { + try { + const globResult = await globTool.execute( + { pattern: `**/*${pathName}*`, path: config.getTargetDir() }, + signal, + ); + if ( + globResult.llmContent && + typeof globResult.llmContent === 'string' && + !globResult.llmContent.startsWith('No files found') && + !globResult.llmContent.startsWith('Error:') + ) { + const lines = globResult.llmContent.split('\n'); + if (lines.length > 1 && lines[1]) { + const firstMatchAbsolute = lines[1].trim(); + currentPathSpec = path.relative( + config.getTargetDir(), + firstMatchAbsolute, + ); + onDebugMessage( + `Glob search for ${pathName} found ${firstMatchAbsolute}, using relative path: ${currentPathSpec}`, + ); + resolvedSuccessfully = true; + } else { + onDebugMessage( + `Glob search for '**/*${pathName}*' did not return a usable path. Path ${pathName} will be skipped.`, + ); + } } else { onDebugMessage( - `Glob search for '**/*${pathPart}*' did not return a usable path. Proceeding with original: ${pathPart}`, + `Glob search for '**/*${pathName}*' found no files or an error. Path ${pathName} will be skipped.`, ); - // pathSpec remains pathPart } - } else { - onDebugMessage( - `Glob search for '**/*${pathPart}*' found no files or an error occurred. Proceeding with original: ${pathPart}`, + } catch (globError) { + console.error( + `Error during glob search for ${pathName}: ${getErrorMessage(globError)}`, + ); + onDebugMessage( + `Error during glob search for ${pathName}. Path ${pathName} will be skipped.`, ); - // pathSpec remains pathPart } - } catch (globError) { - console.error( - `Error during glob search: ${getErrorMessage(globError)}`, - ); + } else { onDebugMessage( - `Error during glob search. Proceeding with original: ${pathPart}`, + `Glob tool not found. Path ${pathName} will be skipped.`, ); - // pathSpec remains pathPart } } else { - onDebugMessage( - 'Glob tool not found. Proceeding with original path: ${pathPart}', + console.error( + `Error stating path ${pathName}: ${getErrorMessage(error)}`, + ); + onDebugMessage( + `Error stating path ${pathName}. Path ${pathName} will be skipped.`, ); - // pathSpec remains pathPart } - } else { - console.error( - `Error stating path ${pathPart}: ${getErrorMessage(error)}`, - ); - onDebugMessage( - `Error stating path, proceeding with original: ${pathSpec}`, - ); + } + + if (resolvedSuccessfully) { + pathSpecsToRead.push(currentPathSpec); + atPathToResolvedSpecMap.set(originalAtPath, currentPathSpec); + contentLabelsForDisplay.push(pathName); } } - const toolArgs = { paths: [pathSpec] }; + // Construct the initial part of the query for the LLM + let initialQueryText = ''; + for (let i = 0; i < commandParts.length; i++) { + const part = commandParts[i]; + if (part.type === 'text') { + initialQueryText += part.content; + } else { + // type === 'atPath' + const resolvedSpec = atPathToResolvedSpecMap.get(part.content); + if ( + i > 0 && + initialQueryText.length > 0 && + !initialQueryText.endsWith(' ') && + resolvedSpec + ) { + // Add space if previous part was text and didn't end with space, or if previous was @path + const prevPart = commandParts[i - 1]; + if ( + prevPart.type === 'text' || + (prevPart.type === 'atPath' && + atPathToResolvedSpecMap.has(prevPart.content)) + ) { + initialQueryText += ' '; + } + } + if (resolvedSpec) { + initialQueryText += `@${resolvedSpec}`; + } else { + // If not resolved for reading (e.g. lone @ or invalid path that was skipped), + // add the original @-string back, ensuring spacing if it's not the first element. + if ( + i > 0 && + initialQueryText.length > 0 && + !initialQueryText.endsWith(' ') && + !part.content.startsWith(' ') + ) { + initialQueryText += ' '; + } + initialQueryText += part.content; + } + } + } + initialQueryText = initialQueryText.trim(); + + // Fallback for lone "@" or completely invalid @-commands resulting in empty initialQueryText + if (pathSpecsToRead.length === 0) { + onDebugMessage('No valid file paths found in @ commands to read.'); + if (initialQueryText === '@' && query.trim() === '@') { + // If the only thing was a lone @, pass original query (which might have spaces) + return { processedQuery: [{ text: query }], shouldProceed: true }; + } else if (!initialQueryText && query) { + // If all @-commands were invalid and no surrounding text, pass original query + return { processedQuery: [{ text: query }], shouldProceed: true }; + } + // Otherwise, proceed with the (potentially modified) query text that doesn't involve file reading + return { + processedQuery: [{ text: initialQueryText || query }], + shouldProceed: true, + }; + } + + const processedQueryParts: PartUnion[] = [{ text: initialQueryText }]; + + const toolArgs = { paths: pathSpecsToRead }; let toolCallDisplay: IndividualToolCallDisplay; try { const result = await readManyFilesTool.execute(toolArgs, signal); - toolCallDisplay = { callId: `client-read-${userMessageTimestamp}`, name: readManyFilesTool.displayName, description: readManyFilesTool.getDescription(toolArgs), status: ToolCallStatus.Success, - resultDisplay: result.returnDisplay, + resultDisplay: + result.returnDisplay || + `Successfully read: ${contentLabelsForDisplay.join(', ')}`, confirmationDetails: undefined, }; - // Prepare the query parts for the LLM - const processedQueryParts: PartUnion[] = []; - if (textBefore) { - processedQueryParts.push({ text: textBefore }); - } - - // Process the result from the tool - processedQueryParts.push('\n--- Content from: ${contentLabel} ---\n'); - if (Array.isArray(result.llmContent)) { - for (const part of result.llmContent) { - processedQueryParts.push(part); + if ( + result.llmContent && + typeof result.llmContent === 'string' && + result.llmContent.trim() !== '' + ) { + processedQueryParts.push({ + text: '\n--- Content from referenced files ---', + }); + const fileContentRegex = + /\n--- (.*?) ---\n([\s\S]*?)(?=\n--- .*? ---\n|$)/g; + let match; + const foundContentForSpecs = new Set(); + while ((match = fileContentRegex.exec(result.llmContent)) !== null) { + const filePathSpecInContent = match[1]; // This is a resolved pathSpec + const fileActualContent = match[2].trim(); + if (pathSpecsToRead.includes(filePathSpecInContent)) { + // Ensure we only add content for paths we requested + processedQueryParts.push({ + text: `\nContent from @${filePathSpecInContent}:\n`, + }); + processedQueryParts.push({ text: fileActualContent }); + foundContentForSpecs.add(filePathSpecInContent); + } } + // Check if any requested pathSpecs didn't yield content in the parsed block, could indicate an issue. + for (const requestedSpec of pathSpecsToRead) { + if (!foundContentForSpecs.has(requestedSpec)) { + onDebugMessage( + `Content for @${requestedSpec} was expected but not found in read_many_files output.`, + ); + // Optionally add a note about missing content for this spec + // processedQueryParts.push({ text: `\nContent for @${requestedSpec} not found or empty.\n` }); + } + } + processedQueryParts.push({ text: '\n--- End of content ---' }); } else { - processedQueryParts.push(result.llmContent); + onDebugMessage( + 'read_many_files tool returned no content or empty content.', + ); } - processedQueryParts.push('\n--- End of content ---\n'); - if (textAfter) { - processedQueryParts.push({ text: textAfter }); - } - const processedQuery: PartListUnion = processedQueryParts; - - // Add the successful tool result to history addItem( { type: 'tool_group', tools: [toolCallDisplay] } as Omit< HistoryItem, @@ -266,20 +388,16 @@ export async function handleAtCommand({ >, userMessageTimestamp, ); - - return { processedQuery, shouldProceed: true }; + return { processedQuery: processedQueryParts, shouldProceed: true }; } catch (error: unknown) { - // Handle errors during tool execution toolCallDisplay = { callId: `client-read-${userMessageTimestamp}`, name: readManyFilesTool.displayName, description: readManyFilesTool.getDescription(toolArgs), status: ToolCallStatus.Error, - resultDisplay: `Error reading ${contentLabel}: ${getErrorMessage(error)}`, + resultDisplay: `Error reading files (${contentLabelsForDisplay.join(', ')}): ${getErrorMessage(error)}`, confirmationDetails: undefined, }; - - // Add the error tool result to history addItem( { type: 'tool_group', tools: [toolCallDisplay] } as Omit< HistoryItem, @@ -287,7 +405,6 @@ export async function handleAtCommand({ >, userMessageTimestamp, ); - return { processedQuery: null, shouldProceed: false }; } }