Bound the output of grep_files

graphcore/tools/vfs.py

@@ -125,6 +125,41 @@ def _get_file(cont: str | None, range: FileRange | None) -> str:
 
 
 
+# Bounds on one grep result, which goes into the prompt whole: past the model's
+# context window the next call fails. 100k characters is a few tens of thousands
+# of tokens. One line of a minified or generated file can be megabytes long.
+GREP_MAX_OUTPUT_CHARS = 100_000
+GREP_MAX_LINE_CHARS = 500
+
+
+def _clip_line(line: str) -> str:
+    if len(line) <= GREP_MAX_LINE_CHARS:
+        return line
+    rest = len(line) - GREP_MAX_LINE_CHARS
+    return f"{line[:GREP_MAX_LINE_CHARS]} [... line cut, {rest} more characters]"
+
+
+def _bounded_join(entries: list[str], unit: str) -> str:
+    """Join *entries* one per line, stopping before the total passes
+    :data:`GREP_MAX_OUTPUT_CHARS`. When entries are left out, a last line says
+    how many and tells the agent how to narrow the search."""
+    kept: list[str] = []
+    size = 0
+    for entry in entries:
+        # +1 for the newline that joins it to the previous entry.
+        if size + len(entry) + 1 > GREP_MAX_OUTPUT_CHARS:
+            break
+        kept.append(entry)
+        size += len(entry) + 1
+    if len(kept) == len(entries):
+        return "\n".join(kept)
+    kept.append(
+        f"[Output cut: showing {len(kept)} of {len(entries)} {unit}. "
+        "Use a more specific search_string, or restrict the search with match_in.]"
+    )
+    return "\n".join(kept)
+
+
 def _grep_impl(
     search_string: str,
     matching_lines: bool,
@@ -137,11 +172,15 @@ def _grep_impl(
 
     Args:
         search_string: Regex pattern to search for
+        matching_lines: If true, return matching lines; otherwise matching file names
         file_contents: Iterator of (filename, content) tuples
-        check_allowed: Filter function for allowed filenames
+        get_content: Reads a file's content by name, for the matching-lines pass
+        match_in: If set, only these paths are searched
 
     Returns:
-        Newline-separated list of matching filenames, or error message
+        Newline-separated matching file names or ``file:lineno:line`` entries,
+        bounded by :data:`GREP_MAX_OUTPUT_CHARS` and :data:`GREP_MAX_LINE_CHARS`,
+        or an error message
     """
     comp: re.Pattern
     try:
@@ -164,7 +203,7 @@ def _grep_impl(
             matches.append(k)
 
     if not matching_lines:
-        return "\n".join(matches)
+        return _bounded_join(matches, "matching files")
 
     matched_lines = []
 
@@ -174,9 +213,9 @@ def _grep_impl(
         cont = cont_s.splitlines()
         for (lno, l) in enumerate(cont, start=1):
             if comp.search(l):
-                matched_lines.append(f"{match_name}:{lno}:{l}")
+                matched_lines.append(f"{match_name}:{lno}:{_clip_line(l)}")
 
-    return "\n".join(matched_lines)
+    return _bounded_join(matched_lines, "matching lines")
 
 
 
@@ -261,6 +300,11 @@ class _GrepFileSchemaBase(BaseModel):
 
     In both modes, the paths to search can be restricted with `match_in`.
 
+    The output is bounded. A very long matching line is cut, and when there are more
+    matches than fit, the output stops and its last line says how many were left out.
+    If that happens, use a more specific search string or restrict the search with
+    `match_in`.
+
     IMPORTANT: This tool searches file *contents*, NOT their names. DO NOT use this tool to search for files
     whose name make a regex, it will not work.
     """

tests/test_fs_layered.py

@@ -16,6 +16,7 @@
 import pytest
 
 from graphcore.tools.vfs import (
+    GREP_MAX_OUTPUT_CHARS,
     DirBackend,
     _LayeredMaterializer,
     _make_global_include_pred,
@@ -276,6 +277,61 @@ async def test_grep_matching_lines_includes_line_number(self):
         assert "x.sol:2:" in resp
 
 
+# ---------------------------------------------------------------------------
+# fs_tools_layered — grep output bounds
+# ---------------------------------------------------------------------------
+
+
+class TestGrepOutputBounds:
+    pytestmark = pytest.mark.asyncio
+
+    async def test_many_matching_lines_are_cut_with_a_notice(self):
+        # 2,000 files of 200 matching lines each: about 20 MB if returned whole.
+        body = "\n".join(f"function f{i}() external {{}}" for i in range(200))
+        b = InMemoryBackend({f"lib/dep{n}/F.sol": body for n in range(2000)})
+        tools, _ = _tools_by_name([b])
+        resp = tools["grep_files"].invoke({
+            "search_string": "function",
+            "matching_lines": True,
+        })
+        assert len(resp) <= GREP_MAX_OUTPUT_CHARS + 300
+        assert "showing" in resp and "of 400000 matching lines" in resp
+        assert "match_in" in resp.splitlines()[-1]
+        # Every line before the notice is a complete ``file:lineno:line`` entry.
+        for entry in resp.splitlines()[:-1]:
+            assert entry.startswith("lib/dep") and "function f" in entry
+
+    async def test_long_matching_line_is_clipped(self):
+        b = InMemoryBackend({"dist/app.sol": "x" * 3_000_000 + "NEEDLE" + "y" * 10})
+        tools, _ = _tools_by_name([b])
+        resp = tools["grep_files"].invoke({
+            "search_string": "NEEDLE",
+            "matching_lines": True,
+        })
+        assert resp.startswith("dist/app.sol:1:")
+        assert len(resp) < 1_000
+        assert "line cut" in resp
+
+    async def test_many_matching_files_are_cut_with_a_notice(self):
+        b = InMemoryBackend({f"{'d' * 80}/f{n}.sol": "hit" for n in range(5000)})
+        tools, _ = _tools_by_name([b])
+        resp = tools["grep_files"].invoke({
+            "search_string": "hit",
+            "matching_lines": False,
+        })
+        assert len(resp) <= GREP_MAX_OUTPUT_CHARS + 300
+        assert "of 5000 matching files" in resp.splitlines()[-1]
+
+    async def test_output_under_the_bound_has_no_notice(self):
+        b = InMemoryBackend({"a.sol": "hit\nhit", "b.sol": "hit"})
+        tools, _ = _tools_by_name([b])
+        resp = tools["grep_files"].invoke({
+            "search_string": "hit",
+            "matching_lines": True,
+        })
+        assert resp == "a.sol:1:hit\na.sol:2:hit\nb.sol:1:hit"
+
+
 # ---------------------------------------------------------------------------
 # fs_tools_layered — forbidden_read filtering
 # ---------------------------------------------------------------------------

tests/test_vfs.py

@@ -6,7 +6,7 @@
 
 from langgraph.graph import MessagesState
 
-from graphcore.tools.vfs import vfs_tools, VFSState, VFSToolConfig
+from graphcore.tools.vfs import vfs_tools, VFSState, VFSToolConfig, GREP_MAX_OUTPUT_CHARS
 from graphcore.testing import Scenario, tool_call_raw, ToolCallDict
 
 pytestmark = pytest.mark.asyncio
@@ -175,6 +175,16 @@ async def test_grep_matching_lines(self):
         ).run_last_single_tool(_GREP)
         assert "src/Foo.sol:2:" in resp
 
+    async def test_grep_matching_lines_output_is_bounded(self):
+        body = "\n".join(f"uint v{i};" for i in range(500))
+        resp = await _scenario({
+            f"lib/dep{n}/V.sol": body for n in range(500)
+        }).turn(
+            _grep("uint", matching_lines=True),
+        ).run_last_single_tool(_GREP)
+        assert len(resp) <= GREP_MAX_OUTPUT_CHARS + 300
+        assert "of 250000 matching lines" in resp.splitlines()[-1]
+
 
 # =========================================================================
 # Forbidden write enforcement