diff --git a/packages/core/src/tools/jit-context.ts b/packages/core/src/tools/jit-context.ts index 4697cb6389..f8ee4be6dc 100644 --- a/packages/core/src/tools/jit-context.ts +++ b/packages/core/src/tools/jit-context.ts @@ -4,6 +4,7 @@ * SPDX-License-Identifier: Apache-2.0 */ +import type { Part, PartListUnion, PartUnion } from '@google/genai'; import type { Config } from '../config/config.js'; /** @@ -63,3 +64,24 @@ export function appendJitContext( } return `${llmContent}${JIT_CONTEXT_PREFIX}${jitContext}${JIT_CONTEXT_SUFFIX}`; } + +/** + * Appends JIT context to non-string tool content (e.g., images, PDFs) by + * wrapping both the original content and the JIT context into a Part array. + * + * @param llmContent - The original non-string tool output content. + * @param jitContext - The discovered JIT context string. + * @returns A Part array containing the original content and JIT context. + */ +export function appendJitContextToParts( + llmContent: PartListUnion, + jitContext: string, +): PartUnion[] { + const jitPart: Part = { + text: `${JIT_CONTEXT_PREFIX}${jitContext}${JIT_CONTEXT_SUFFIX}`, + }; + const existingParts: PartUnion[] = Array.isArray(llmContent) + ? llmContent + : [llmContent]; + return [...existingParts, jitPart]; +} diff --git a/packages/core/src/tools/read-file.test.ts b/packages/core/src/tools/read-file.test.ts index 85981ff80b..fa7a0669d6 100644 --- a/packages/core/src/tools/read-file.test.ts +++ b/packages/core/src/tools/read-file.test.ts @@ -30,6 +30,15 @@ vi.mock('./jit-context.js', () => ({ if (!context) return content; return `${content}\n\n--- Newly Discovered Project Context ---\n${context}\n--- End Project Context ---`; }), + appendJitContextToParts: vi.fn().mockImplementation((content, context) => { + const jitPart = { + text: `\n\n--- Newly Discovered Project Context ---\n${context}\n--- End Project Context ---`, + }; + const existing = Array.isArray(content) ? content : [content]; + return [...existing, jitPart]; + }), + JIT_CONTEXT_PREFIX: '\n\n--- Newly Discovered Project Context ---\n', + JIT_CONTEXT_SUFFIX: '\n--- End Project Context ---', })); describe('ReadFileTool', () => { @@ -637,5 +646,43 @@ describe('ReadFileTool', () => { 'Newly Discovered Project Context', ); }); + + it('should append JIT context as Part array for non-string llmContent (binary files)', async () => { + const { discoverJitContext } = await import('./jit-context.js'); + vi.mocked(discoverJitContext).mockResolvedValue( + 'Auth rules: use httpOnly cookies.', + ); + + // Create a minimal valid PNG file (1x1 pixel) + const pngHeader = Buffer.from([ + 0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a, 0x00, 0x00, 0x00, 0x0d, + 0x49, 0x48, 0x44, 0x52, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, + 0x08, 0x02, 0x00, 0x00, 0x00, 0x90, 0x77, 0x53, 0xde, 0x00, 0x00, 0x00, + 0x0c, 0x49, 0x44, 0x41, 0x54, 0x08, 0xd7, 0x63, 0xf8, 0xcf, 0xc0, 0x00, + 0x00, 0x00, 0x02, 0x00, 0x01, 0xe2, 0x21, 0xbc, 0x33, 0x00, 0x00, 0x00, + 0x00, 0x49, 0x45, 0x4e, 0x44, 0xae, 0x42, 0x60, 0x82, + ]); + const filePath = path.join(tempRootDir, 'test-image.png'); + await fsp.writeFile(filePath, pngHeader); + + const invocation = tool.build({ file_path: filePath }); + const result = await invocation.execute(abortSignal); + + expect(discoverJitContext).toHaveBeenCalled(); + // Result should be an array containing both the image part and JIT context + expect(Array.isArray(result.llmContent)).toBe(true); + const parts = result.llmContent as Array>; + const jitTextPart = parts.find( + (p) => + typeof p['text'] === 'string' && p['text'].includes('Auth rules'), + ); + expect(jitTextPart).toBeDefined(); + expect(jitTextPart!['text']).toContain( + 'Newly Discovered Project Context', + ); + expect(jitTextPart!['text']).toContain( + 'Auth rules: use httpOnly cookies.', + ); + }); }); }); diff --git a/packages/core/src/tools/read-file.ts b/packages/core/src/tools/read-file.ts index c2f2157869..69f9e0274b 100644 --- a/packages/core/src/tools/read-file.ts +++ b/packages/core/src/tools/read-file.ts @@ -20,7 +20,7 @@ import { import { ToolErrorType } from './tool-error.js'; import { buildFilePathArgsPattern } from '../policy/utils.js'; -import type { PartUnion } from '@google/genai'; +import type { PartListUnion } from '@google/genai'; import { processSingleFileContent, getSpecificMimeType, @@ -34,7 +34,11 @@ import { READ_FILE_TOOL_NAME, READ_FILE_DISPLAY_NAME } from './tool-names.js'; import { FileDiscoveryService } from '../services/fileDiscoveryService.js'; import { READ_FILE_DEFINITION } from './definitions/coreTools.js'; import { resolveToolDeclaration } from './definitions/resolver.js'; -import { discoverJitContext, appendJitContext } from './jit-context.js'; +import { + discoverJitContext, + appendJitContext, + appendJitContextToParts, +} from './jit-context.js'; /** * Parameters for the ReadFile tool @@ -135,7 +139,7 @@ class ReadFileToolInvocation extends BaseToolInvocation< }; } - let llmContent: PartUnion; + let llmContent: PartListUnion; if (result.isTruncated) { const [start, end] = result.linesShown!; const total = result.originalLineCount!; @@ -173,8 +177,12 @@ ${result.llmContent}`; // Discover JIT subdirectory context for the accessed file path const jitContext = await discoverJitContext(this.config, this.resolvedPath); - if (jitContext && typeof llmContent === 'string') { - llmContent = appendJitContext(llmContent, jitContext); + if (jitContext) { + if (typeof llmContent === 'string') { + llmContent = appendJitContext(llmContent, jitContext); + } else { + llmContent = appendJitContextToParts(llmContent, jitContext); + } } return { diff --git a/packages/core/src/tools/read-many-files.test.ts b/packages/core/src/tools/read-many-files.test.ts index b2f7ff2f7d..6a526d2b62 100644 --- a/packages/core/src/tools/read-many-files.test.ts +++ b/packages/core/src/tools/read-many-files.test.ts @@ -860,5 +860,62 @@ Content of file[1] : String(result.llmContent); expect(llmContent).not.toContain('Newly Discovered Project Context'); }); + + it('should discover JIT context sequentially to avoid duplicate shared parent context', async () => { + const { discoverJitContext } = await import('./jit-context.js'); + + // Simulate two subdirectories sharing a parent GEMINI.md. + // Sequential execution means the second call sees the parent already + // loaded, so it only returns its own leaf context. + const callOrder: string[] = []; + let firstCallDone = false; + vi.mocked(discoverJitContext).mockImplementation(async (_config, dir) => { + callOrder.push(dir); + if (!firstCallDone) { + // First call (whichever dir) loads the shared parent + its own leaf + firstCallDone = true; + return 'Parent context\nFirst leaf context'; + } + // Second call only returns its own leaf (parent already loaded) + return 'Second leaf context'; + }); + + // Create files in two sibling subdirectories + fs.mkdirSync(path.join(tempRootDir, 'subA'), { recursive: true }); + fs.mkdirSync(path.join(tempRootDir, 'subB'), { recursive: true }); + fs.writeFileSync( + path.join(tempRootDir, 'subA', 'a.ts'), + 'const a = 1;', + 'utf8', + ); + fs.writeFileSync( + path.join(tempRootDir, 'subB', 'b.ts'), + 'const b = 2;', + 'utf8', + ); + + const invocation = tool.build({ include: ['subA/a.ts', 'subB/b.ts'] }); + const result = await invocation.execute(new AbortController().signal); + + // Verify both directories were discovered (order depends on Set iteration) + expect(callOrder).toHaveLength(2); + expect(callOrder).toEqual( + expect.arrayContaining([ + expect.stringContaining('subA'), + expect.stringContaining('subB'), + ]), + ); + + const llmContent = Array.isArray(result.llmContent) + ? result.llmContent.join('') + : String(result.llmContent); + expect(llmContent).toContain('Parent context'); + expect(llmContent).toContain('First leaf context'); + expect(llmContent).toContain('Second leaf context'); + + // Parent context should appear only once (from the first call), not duplicated + const parentMatches = llmContent.match(/Parent context/g); + expect(parentMatches).toHaveLength(1); + }); }); }); diff --git a/packages/core/src/tools/read-many-files.ts b/packages/core/src/tools/read-many-files.ts index 34a2def596..e2a283c726 100644 --- a/packages/core/src/tools/read-many-files.ts +++ b/packages/core/src/tools/read-many-files.ts @@ -416,14 +416,19 @@ ${finalExclusionPatternsForDescription } } - // Discover JIT subdirectory context for all unique directories of processed files + // Discover JIT subdirectory context for all unique directories of processed files. + // Run sequentially so each call sees paths marked as loaded by the previous + // one, preventing shared parent GEMINI.md files from being injected twice. const uniqueDirs = new Set( Array.from(filesToConsider).map((f) => path.dirname(f)), ); - const jitResults = await Promise.all( - Array.from(uniqueDirs).map((dir) => discoverJitContext(this.config, dir)), - ); - const jitParts = jitResults.filter(Boolean); + const jitParts: string[] = []; + for (const dir of uniqueDirs) { + const ctx = await discoverJitContext(this.config, dir); + if (ctx) { + jitParts.push(ctx); + } + } if (jitParts.length > 0) { contentParts.push( `${JIT_CONTEXT_PREFIX}${jitParts.join('\n')}${JIT_CONTEXT_SUFFIX}`, diff --git a/packages/core/src/utils/memoryDiscovery.test.ts b/packages/core/src/utils/memoryDiscovery.test.ts index c2b865dad1..9cb9942747 100644 --- a/packages/core/src/utils/memoryDiscovery.test.ts +++ b/packages/core/src/utils/memoryDiscovery.test.ts @@ -1155,6 +1155,60 @@ included directory memory // Ensure outer memory is NOT loaded expect(result.files.find((f) => f.path === outerMemory)).toBeUndefined(); }); + + it('should resolve file target to its parent directory for traversal', async () => { + const rootDir = await createEmptyDir( + path.join(testRootDir, 'jit_file_resolve'), + ); + const subDir = await createEmptyDir(path.join(rootDir, 'src')); + + // Create the target file so fs.stat can identify it as a file + const targetFile = await createTestFile( + path.join(subDir, 'app.ts'), + 'const x = 1;', + ); + + const subDirMemory = await createTestFile( + path.join(subDir, DEFAULT_CONTEXT_FILENAME), + 'Src context rules', + ); + + const result = await loadJitSubdirectoryMemory( + targetFile, + [rootDir], + new Set(), + ); + + // Should find the GEMINI.md in the same directory as the file + expect(result.files).toHaveLength(1); + expect(result.files[0].path).toBe(subDirMemory); + expect(result.files[0].content).toBe('Src context rules'); + }); + + it('should handle non-existent file target by using parent directory', async () => { + const rootDir = await createEmptyDir( + path.join(testRootDir, 'jit_nonexistent'), + ); + const subDir = await createEmptyDir(path.join(rootDir, 'src')); + + // Target file does NOT exist (e.g. write_file creating a new file) + const targetFile = path.join(subDir, 'new-file.ts'); + + const subDirMemory = await createTestFile( + path.join(subDir, DEFAULT_CONTEXT_FILENAME), + 'Rules for new files', + ); + + const result = await loadJitSubdirectoryMemory( + targetFile, + [rootDir], + new Set(), + ); + + expect(result.files).toHaveLength(1); + expect(result.files[0].path).toBe(subDirMemory); + expect(result.files[0].content).toBe('Rules for new files'); + }); }); it('refreshServerHierarchicalMemory should refresh memory and update config', async () => { diff --git a/packages/core/src/utils/memoryDiscovery.ts b/packages/core/src/utils/memoryDiscovery.ts index 2d7de3327c..f772394d79 100644 --- a/packages/core/src/utils/memoryDiscovery.ts +++ b/packages/core/src/utils/memoryDiscovery.ts @@ -767,8 +767,24 @@ export async function loadJitSubdirectoryMemory( `(Trusted root: ${bestRoot})`, ); - // Traverse from target up to the trusted root - const potentialPaths = await findUpwardGeminiFiles(resolvedTarget, bestRoot); + // Resolve the target to a directory before traversing upward. + // When the target is a file (e.g. /app/src/file.ts), start from its + // parent directory to avoid a wasted fs.access check on a nonsensical + // path like /app/src/file.ts/GEMINI.md. + let startDir = resolvedTarget; + try { + const stat = await fs.stat(resolvedTarget); + if (stat.isFile()) { + startDir = normalizePath(path.dirname(resolvedTarget)); + } + } catch { + // If stat fails (e.g. file doesn't exist yet for write_file), + // assume it's a file path and use its parent directory. + startDir = normalizePath(path.dirname(resolvedTarget)); + } + + // Traverse from the resolved directory up to the trusted root + const potentialPaths = await findUpwardGeminiFiles(startDir, bestRoot); if (potentialPaths.length === 0) { return { files: [], fileIdentities: [] };