Utilize pipelining of grep_search -> read_file to eliminate turns

packages/core/src/tools/diff-utils.test.ts

@@ -0,0 +1,66 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { describe, expect, it } from 'vitest';
+import { getDiffContextSnippet } from './diff-utils.js';
+
+describe('getDiffContextSnippet', () => {
+  it('should return the whole new content if originalContent is empty', () => {
+    const original = '';
+    const modified = 'line1\nline2\nline3';
+    expect(getDiffContextSnippet(original, modified)).toBe(modified);
+  });
+
+  it('should return the whole content if there are no changes', () => {
+    const content = 'line1\nline2\nline3';
+    expect(getDiffContextSnippet(content, content)).toBe(content);
+  });
+
+  it('should show added lines with context', () => {
+    const original = '1\n2\n3\n4\n5\n6\n7\n8\n9\n10';
+    const modified = '1\n2\n3\n4\n5\nadded\n6\n7\n8\n9\n10';
+    // Default context is 5 lines.
+    expect(getDiffContextSnippet(original, modified)).toBe(modified);
+  });
+
+  it('should use ellipses for changes far apart', () => {
+    const original = Array.from({ length: 20 }, (_, i) => `${i + 1}`).join(
+      '\n',
+    );
+    const modified = original
+      .replace('2\n', '2\nadded1\n')
+      .replace('19', '19\nadded2');
+    const snippet = getDiffContextSnippet(original, modified, 2);
+
+    expect(snippet).toContain('1\n2\nadded1\n3\n4');
+    expect(snippet).toContain('...');
+    expect(snippet).toContain('18\n19\nadded2\n20');
+  });
+
+  it('should respect custom contextLines', () => {
+    const original = '1\n2\n3\n4\n5\n6\n7\n8\n9\n10';
+    const modified = '1\n2\n3\n4\n5\nadded\n6\n7\n8\n9\n10';
+    const snippet = getDiffContextSnippet(original, modified, 1);
+
+    expect(snippet).toBe('...\n5\nadded\n6\n...');
+  });
+
+  it('should handle multiple changes close together by merging ranges', () => {
+    const original = '1\n2\n3\n4\n5\n6\n7\n8\n9\n10';
+    const modified = '1\nadded1\n2\nadded2\n3\n4\n5\n6\n7\n8\n9\n10';
+    const snippet = getDiffContextSnippet(original, modified, 1);
+
+    expect(snippet).toBe('1\nadded1\n2\nadded2\n3\n...');
+  });
+
+  it('should handle removals', () => {
+    const original = '1\n2\n3\n4\n5\n6\n7\n8\n9\n10';
+    const modified = '1\n2\n3\n4\n6\n7\n8\n9\n10';
+    const snippet = getDiffContextSnippet(original, modified, 1);
+
+    expect(snippet).toBe('...\n4\n6\n...');
+  });
+});

packages/core/src/tools/diff-utils.ts

@@ -0,0 +1,75 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import * as Diff from 'diff';
+
+/**
+ * Generates a snippet of the diff between two strings, including a few lines of context around the changes.
+ */
+export function getDiffContextSnippet(
+  originalContent: string,
+  newContent: string,
+  contextLines = 5,
+): string {
+  if (!originalContent) {
+    return newContent;
+  }
+
+  const changes = Diff.diffLines(originalContent, newContent);
+  const newLines = newContent.split(/\r?\n/);
+  const ranges: Array<{ start: number; end: number }> = [];
+  let newLineIdx = 0;
+
+  for (const change of changes) {
+    if (change.added) {
+      ranges.push({ start: newLineIdx, end: newLineIdx + (change.count ?? 0) });
+      newLineIdx += change.count ?? 0;
+    } else if (change.removed) {
+      ranges.push({ start: newLineIdx, end: newLineIdx });
+    } else {
+      newLineIdx += change.count ?? 0;
+    }
+  }
+
+  if (ranges.length === 0) {
+    return newContent;
+  }
+
+  const expandedRanges = ranges.map((r) => ({
+    start: Math.max(0, r.start - contextLines),
+    end: Math.min(newLines.length, r.end + contextLines),
+  }));
+  expandedRanges.sort((a, b) => a.start - b.start);
+  const mergedRanges: Array<{ start: number; end: number }> = [];
+
+  if (expandedRanges.length > 0) {
+    let current = expandedRanges[0];
+    for (let i = 1; i < expandedRanges.length; i++) {
+      const next = expandedRanges[i];
+      if (next.start <= current.end) {
+        current.end = Math.max(current.end, next.end);
+      } else {
+        mergedRanges.push(current);
+        current = next;
+      }
+    }
+    mergedRanges.push(current);
+  }
+
+  const outputParts: string[] = [];
+  let lastEnd = 0;
+
+  for (const range of mergedRanges) {
+    if (range.start > lastEnd) outputParts.push('...');
+    outputParts.push(newLines.slice(range.start, range.end).join('\n'));
+    lastEnd = range.end;
+  }
+
+  if (lastEnd < newLines.length) {
+    outputParts.push('...');
+  }
+  return outputParts.join('\n');
+}

packages/core/src/tools/edit.ts

@@ -30,6 +30,7 @@ import { ApprovalMode } from '../policy/types.js';
 import { CoreToolCallStatus } from '../scheduler/types.js';
 
 import { DEFAULT_DIFF_OPTIONS, getDiffStat } from './diffOptions.js';
+import { getDiffContextSnippet } from './diff-utils.js';
 import {
   type ModifiableDeclarativeTool,
   type ModifyContext,
@@ -876,6 +877,16 @@ class EditToolInvocation
           ? `Created new file: ${this.params.file_path} with provided content.`
           : `Successfully modified file: ${this.params.file_path} (${editData.occurrences} replacements).`,
       ];
+
+      // Return a diff of the file before and after the write so that the agent
+      // can avoid the need to spend a turn doing a verification read.
+      const snippet = getDiffContextSnippet(
+        editData.currentContent ?? '',
+        finalContent,
+        5,
+      );
+      llmSuccessMessageParts.push(`Here is the updated code:
+${snippet}`);
       const fuzzyFeedback = getFuzzyMatchFeedback(editData);
       if (fuzzyFeedback) {
         llmSuccessMessageParts.push(fuzzyFeedback);

packages/core/src/tools/grep-utils.ts

@@ -0,0 +1,212 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import fsPromises from 'node:fs/promises';
+import { debugLogger } from '../utils/debugLogger.js';
+
+/**
+ * Result object for a single grep match
+ */
+export interface GrepMatch {
+  filePath: string;
+  absolutePath: string;
+  lineNumber: number;
+  line: string;
+  isContext?: boolean;
+}
+
+/**
+ * Groups matches by their file path and ensures they are sorted by line number.
+ */
+export function groupMatchesByFile(
+  allMatches: GrepMatch[],
+): Record<string, GrepMatch[]> {
+  const groups: Record<string, GrepMatch[]> = {};
+
+  for (const match of allMatches) {
+    if (!groups[match.filePath]) {
+      groups[match.filePath] = [];
+    }
+    groups[match.filePath].push(match);
+  }
+
+  for (const filePath in groups) {
+    groups[filePath].sort((a, b) => a.lineNumber - b.lineNumber);
+  }
+
+  return groups;
+}
+
+/**
+ * Reads the content of a file and splits it into lines.
+ * Returns null if the file cannot be read.
+ */
+export async function readFileLines(
+  absolutePath: string,
+): Promise<string[] | null> {
+  try {
+    const content = await fsPromises.readFile(absolutePath, 'utf8');
+    return content.split(/\r?\n/);
+  } catch (err) {
+    debugLogger.warn(`Failed to read file for context: ${absolutePath}`, err);
+    return null;
+  }
+}
+
+/**
+ * Automatically enriches grep results with surrounding context if the match count is low
+ * and no specific context was requested. This optimization can enable the agent
+ * to skip turns that would be spent reading files after grep calls.
+ */
+export async function enrichWithAutoContext(
+  matchesByFile: Record<string, GrepMatch[]>,
+  matchCount: number,
+  params: {
+    names_only?: boolean;
+    context?: number;
+    before?: number;
+    after?: number;
+  },
+): Promise<void> {
+  const { names_only, context, before, after } = params;
+
+  if (
+    matchCount >= 1 &&
+    matchCount <= 3 &&
+    !names_only &&
+    context === undefined &&
+    before === undefined &&
+    after === undefined
+  ) {
+    const contextLines = matchCount === 1 ? 50 : 15;
+    for (const filePath in matchesByFile) {
+      const fileMatches = matchesByFile[filePath];
+      if (fileMatches.length === 0) continue;
+
+      const fileLines = await readFileLines(fileMatches[0].absolutePath);
+
+      if (fileLines) {
+        const newFileMatches: GrepMatch[] = [];
+        const seenLines = new Set<number>();
+
+        // Sort matches to process them in order
+        fileMatches.sort((a, b) => a.lineNumber - b.lineNumber);
+
+        for (const match of fileMatches) {
+          const startLine = Math.max(0, match.lineNumber - 1 - contextLines);
+          const endLine = Math.min(
+            fileLines.length,
+            match.lineNumber - 1 + contextLines + 1,
+          );
+
+          for (let i = startLine; i < endLine; i++) {
+            const lineNum = i + 1;
+            if (!seenLines.has(lineNum)) {
+              newFileMatches.push({
+                absolutePath: match.absolutePath,
+                filePath: match.filePath,
+                lineNumber: lineNum,
+                line: fileLines[i],
+                isContext: lineNum !== match.lineNumber,
+              });
+              seenLines.add(lineNum);
+            } else if (lineNum === match.lineNumber) {
+              const existing = newFileMatches.find(
+                (m) => m.lineNumber === lineNum,
+              );
+              if (existing) {
+                existing.isContext = false;
+              }
+            }
+          }
+        }
+        matchesByFile[filePath] = newFileMatches.sort(
+          (a, b) => a.lineNumber - b.lineNumber,
+        );
+      }
+    }
+  }
+}
+
+/**
+ * Formats the grep results for the LLM, including optional context.
+ */
+export async function formatGrepResults(
+  allMatches: GrepMatch[],
+  params: {
+    pattern: string;
+    names_only?: boolean;
+    include?: string;
+    // Context params to determine if auto-context should be skipped
+    context?: number;
+    before?: number;
+    after?: number;
+  },
+  searchLocationDescription: string,
+  totalMaxMatches: number,
+): Promise<{ llmContent: string; returnDisplay: string }> {
+  const { pattern, names_only, include } = params;
+
+  if (allMatches.length === 0) {
+    const noMatchMsg = `No matches found for pattern "${pattern}" ${searchLocationDescription}${include ? ` (filter: "${include}")` : ''}.`;
+    return { llmContent: noMatchMsg, returnDisplay: `No matches found` };
+  }
+
+  const matchesByFile = groupMatchesByFile(allMatches);
+
+  const matchesOnly = allMatches.filter((m) => !m.isContext);
+  const matchCount = matchesOnly.length; // Count actual matches, not context lines
+  const matchTerm = matchCount === 1 ? 'match' : 'matches';
+
+  // If the result count is low and Gemini didn't request before/after lines of context
+  // add a small amount anyways to enable the agent to avoid one or more extra turns
+  // reading the matched files. This optimization reduces turns count by ~10% in SWEBench.
+  await enrichWithAutoContext(matchesByFile, matchCount, params);
+
+  const wasTruncated = matchCount >= totalMaxMatches;
+
+  if (names_only) {
+    const filePaths = Object.keys(matchesByFile).sort();
+    let llmContent = `Found ${filePaths.length} files with matches for pattern "${pattern}" ${searchLocationDescription}${
+      include ? ` (filter: "${include}")` : ''
+    }${
+      wasTruncated
+        ? ` (results limited to ${totalMaxMatches} matches for performance)`
+        : ''
+    }:\n`;
+    llmContent += filePaths.join('\n');
+    return {
+      llmContent: llmContent.trim(),
+      returnDisplay: `Found ${filePaths.length} files${wasTruncated ? ' (limited)' : ''}`,
+    };
+  }
+
+  let llmContent = `Found ${matchCount} ${matchTerm} for pattern "${pattern}" ${searchLocationDescription}${include ? ` (filter: "${include}")` : ''}`;
+
+  if (wasTruncated) {
+    llmContent += ` (results limited to ${totalMaxMatches} matches for performance)`;
+  }
+
+  llmContent += `:\n---\n`;
+
+  for (const filePath in matchesByFile) {
+    llmContent += `File: ${filePath}\n`;
+    matchesByFile[filePath].forEach((match) => {
+      // If isContext is undefined, assume it's a match (false)
+      const separator = match.isContext ? '-' : ':';
+      // trimEnd to avoid double newlines if line has them, but we want to preserve indentation
+      llmContent += `L${match.lineNumber}${separator} ${match.line.trimEnd()}\n`;
+    });
+    llmContent += '---\n';
+  }
+
+  return {
+    llmContent: llmContent.trim(),
+    returnDisplay: `Found ${matchCount} ${matchTerm}${
+      wasTruncated ? ' (limited)' : ''
+    }`,
+  };
+}