fix(core): fix three JIT context bugs in read_file, read_many_files, and memoryDiscovery (#22679)

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];
+}

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<Record<string, unknown>>;
+      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.',
+      );
+    });
   });
 });

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 {

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);
+    });
   });
 });

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}`,

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 () => {

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: [] };