Speed up the agentic unit-test fixer

- Send raw test output as test_output (run_unit_tests and the submit_fix
  follow-up) so the server condenses it, instead of head/tail truncation that
  can drop the root cause.
- Make the full test logs readable by read_file/grep and point the agent at
  them; grep gains context_lines and include.
- Seed the first turn with the build folder's file tree and the files changed
  for the current FRID.
- Skip the harness unit-test run when the agent's own run passed and no file
  changed since.
- Start a new session after an abandoned one with previous_session_id so the
  server can pass on what was tried.
- edit_file returns the edited region; repeated read-only calls with no file
  change in between return a pointer instead of the same output again.

Author: VitjanZ

render_machine/actions/fix_unit_tests.py

@@ -1,3 +1,4 @@
+import os
 from typing import Any
 
 import file_utils
@@ -12,6 +13,9 @@
 SUBMIT_FIX_TOOL = "submit_fix"
 # Upper bound on LLM turns spent on one fix attempt; the server bounds the whole session.
 MAX_AGENT_TURNS_PER_ATTEMPT = 40
+# Seeding the first turn: the build folder's file list and the files changed for the FRID.
+MAX_FILE_TREE_ENTRIES = 500
+MAX_RELEVANT_FILES_CHARS = 60_000
 
 
 class FixUnitTests(BaseAction):
@@ -35,9 +39,14 @@ def execute(self, render_context: RenderContext, previous_action_payload: Any |
         api = render_context.codeplain_api
         frid, module_name = render_context.frid_context.frid, render_context.module_name
         changed_files_before = set(context.changed_files)
+        log_path = render_context.script_execution_history.latest_unit_test_output_path
+        if log_path:
+            agent_tools.register_log_path(log_path, render_context)
 
         if context.agent_session_id is None:
             console.info("Starting an agent session to fix the unit tests.")
+            # Cached read results point at earlier turns, which a new session does not have.
+            context.tool_result_cache.clear()
             response = api.agent_start(
                 TASK_TYPE,
                 self._build_task_params(render_context, unittests_issue),
@@ -51,7 +60,9 @@ def execute(self, render_context: RenderContext, previous_action_payload: Any |
             tool_results = context.pending_tool_results + [
                 {
                     "call_id": context.pending_submit_call_id,
-                    "output": f"The fix was applied, but the unit tests still fail:\n{unittests_issue}",
+                    "output": "The fix was applied, but the unit tests still fail."
+                    + agent_tools.full_log_pointer(log_path),
+                    "test_output": unittests_issue,
                 }
             ]
             context.pending_tool_results, context.pending_submit_call_id = [], None
@@ -86,6 +97,7 @@ def execute(self, render_context: RenderContext, previous_action_payload: Any |
                 console.warning(f"Agent session failed: {response.get('error', 'unknown error')}")
             elif status == "tool_calls":
                 console.warning(f"Agent used {MAX_AGENT_TURNS_PER_ATTEMPT} turns without submitting a fix.")
+            context.previous_session_id = context.agent_session_id
             context.agent_session_id, context.pending_submit_call_id, context.pending_tool_results = None, None, []
 
         console.print_files(
@@ -100,7 +112,8 @@ def execute(self, render_context: RenderContext, previous_action_payload: Any |
     def _build_task_params(render_context: RenderContext, unittests_issue: str) -> dict:
         frid = render_context.frid_context.frid
         specifications, _ = plain_spec.get_specifications_for_frid(render_context.plain_source_tree, frid)
-        return {
+        context = render_context.unit_tests_running_context
+        task_params = {
             "definitions": "\n".join(specifications.get(plain_spec.DEFINITIONS, [])),
             "non_functional_requirements": "\n".join(specifications.get(plain_spec.NON_FUNCTIONAL_REQUIREMENTS, [])),
             "functional_requirements": FixUnitTests._functional_requirements_section(render_context, specifications),
@@ -109,7 +122,42 @@ def _build_task_params(render_context: RenderContext, unittests_issue: str) -> d
             "module_name": render_context.module_name,
             "unittests_script_content": FixUnitTests._read_script(render_context.unittests_script),
             "unittests_issue": unittests_issue,
+            "unittests_log_path": render_context.script_execution_history.latest_unit_test_output_path,
+            "file_tree": FixUnitTests._file_tree(render_context.build_folder),
+            "relevant_files": FixUnitTests._relevant_files(
+                render_context.build_folder, render_context.frid_context.changed_files | context.changed_files
+            ),
         }
+        if context.previous_session_id:
+            task_params["previous_session_id"] = context.previous_session_id
+        return task_params
+
+    @staticmethod
+    def _file_tree(build_folder: str) -> str:
+        paths = []
+        for root, dirs, files in os.walk(build_folder):
+            dirs[:] = sorted(d for d in dirs if d not in agent_tools.GREP_EXCLUDED_DIRS and not d.startswith("."))
+            paths.extend(os.path.relpath(os.path.join(root, name), build_folder) for name in sorted(files))
+            if len(paths) > MAX_FILE_TREE_ENTRIES:
+                return "\n".join(paths[:MAX_FILE_TREE_ENTRIES]) + "\n... [more files not listed]"
+        return "\n".join(paths)
+
+    @staticmethod
+    def _relevant_files(build_folder: str, file_names: set[str]) -> dict[str, str]:
+        """Contents of the given build-relative files, smallest first, within MAX_RELEVANT_FILES_CHARS."""
+        contents = {}
+        for name in file_names:
+            full_path = os.path.join(build_folder, name)
+            if os.path.isfile(full_path):
+                with open(full_path, "r", encoding="utf-8", errors="replace") as f:
+                    contents[name] = f.read()
+        relevant, total = {}, 0
+        for name in sorted(contents, key=lambda n: (len(contents[n]), n)):
+            if total + len(contents[name]) > MAX_RELEVANT_FILES_CHARS:
+                break
+            relevant[name] = contents[name]
+            total += len(contents[name])
+        return dict(sorted(relevant.items()))
 
     @staticmethod
     def _functional_requirements_section(render_context: RenderContext, specifications: dict) -> str:

render_machine/actions/run_unit_tests.py

@@ -16,6 +16,15 @@ class RunUnitTests(BaseAction):
     UNRECOVERABLE_ERROR_OUTCOME = "unrecoverable_error_occurred"
 
     def execute(self, render_context: RenderContext, _previous_action_payload: Any | None):
+        context = render_context.unit_tests_running_context
+        if context.verified_passing:
+            # The fixing agent's own run passed and no file changed since; running again is redundant.
+            context.verified_passing = False
+            console.info("Unit tests already passed in the fixing agent's last run; not running them again.")
+            render_context.script_execution_history.latest_unit_test_output_path = context.verified_passing_log_path
+            render_context.script_execution_history.should_update_script_outputs = True
+            return self.SUCCESSFUL_OUTCOME, None
+
         unittests_script = os.path.normpath(render_context.unittests_script)
 
         console.info(

render_machine/agent_tools.py

@@ -7,6 +7,7 @@
 """
 
 import glob
+import json
 import os
 import subprocess
 import tempfile
@@ -19,6 +20,11 @@
 DEFAULT_READ_LIMIT = 200
 MAX_LINE_CHARS = 10_000
 MAX_OUTPUT_CHARS = 30_000
+EDIT_SNIPPET_CONTEXT_LINES = 3
+MAX_EDIT_SNIPPET_LINES = 60
+MAX_GREP_CONTEXT_LINES = 20
+# Tools without side effects; repeating one while no file changed returns a pointer instead.
+READ_ONLY_TOOLS = ("read_file", "grep", "ls_files")
 GREP_EXCLUDED_DIRS = (".git", "__pycache__", "node_modules", ".venv", "target", "dist", "build")
 
 
@@ -37,7 +43,16 @@ def _within(path: str, folder: str) -> bool:
 
 
 def _readable(path: str, render_context: RenderContext) -> bool:
-    return _within(path, _build_folder(render_context)) or _within(path, os.path.normpath(os.getcwd()))
+    return (
+        _within(path, _build_folder(render_context))
+        or _within(path, os.path.normpath(os.getcwd()))
+        or path in render_context.unit_tests_running_context.readable_log_paths
+    )
+
+
+def register_log_path(log_path: str, render_context: RenderContext) -> None:
+    """Allow read_file/grep on a full test log the agent is pointed to."""
+    render_context.unit_tests_running_context.readable_log_paths.add(os.path.normpath(os.path.abspath(log_path)))
 
 
 def _writable(path: str, render_context: RenderContext) -> bool:
@@ -58,8 +73,10 @@ def _bound(text: str) -> str:
 
 
 def _track_change(full_path: str, render_context: RenderContext) -> None:
-    relative_path = os.path.relpath(full_path, _build_folder(render_context))
-    render_context.unit_tests_running_context.changed_files.add(relative_path)
+    context = render_context.unit_tests_running_context
+    context.changed_files.add(os.path.relpath(full_path, _build_folder(render_context)))
+    context.verified_passing, context.verified_passing_log_path = False, None
+    context.tool_result_cache.clear()
 
 
 def read_file(args: dict, render_context: RenderContext) -> str:
@@ -98,7 +115,13 @@ def grep(args: dict, render_context: RenderContext) -> str:
     # is the form the other tools accept.
     build_folder = _build_folder(render_context)
     cwd = build_folder if _within(target, build_folder) else os.getcwd()
-    command = ["grep", "-rnI", *[f"--exclude-dir={d}" for d in GREP_EXCLUDED_DIRS], "-e", pattern, "--"]
+    options = [f"--exclude-dir={d}" for d in GREP_EXCLUDED_DIRS]
+    context_lines = min(max(int(args.get("context_lines") or 0), 0), MAX_GREP_CONTEXT_LINES)
+    if context_lines:
+        options.append(f"-C{context_lines}")
+    if args.get("include"):
+        options.append(f"--include={args['include']}")
+    command = ["grep", "-rnI", *options, "-e", pattern, "--"]
     command.append(os.path.relpath(target, cwd) if _within(target, cwd) else target)
     result = subprocess.run(command, capture_output=True, text=True, cwd=cwd)
     if result.returncode == 1:
@@ -118,7 +141,7 @@ def ls_files(args: dict, render_context: RenderContext) -> str:
         listing = [entry + "/" if os.path.isdir(os.path.join(target, entry)) else entry for entry in entries]
         return f"{target}:\n" + ("\n".join(listing) if listing else "(empty)")
     matches = sorted(glob.glob(target, recursive=True))
-    return "\n".join(matches) if matches else f"No files match '{target}'."
+    return _bound("\n".join(matches)) if matches else f"No files match '{target}'."
 
 
 def edit_file(args: dict, render_context: RenderContext) -> str:
@@ -141,7 +164,18 @@ def edit_file(args: dict, render_context: RenderContext) -> str:
     with open(full_path, "w", encoding="utf-8") as f:
         f.write(content.replace(search, replace, 1))
     _track_change(full_path, render_context)
-    return f"Edited '{full_path}'."
+    return f"Edited '{full_path}'. The edited region now reads:\n" + _edit_snippet(content, search, replace)
+
+
+def _edit_snippet(content: str, search: str, replace: str) -> str:
+    """Numbered lines of the replacement plus a little context, so the edit needs no re-read."""
+    lines = content.replace(search, replace, 1).split("\n")
+    first = content[: content.index(search)].count("\n")
+    last = first + replace.count("\n")
+    start = max(first - EDIT_SNIPPET_CONTEXT_LINES, 0)
+    end = min(last + EDIT_SNIPPET_CONTEXT_LINES + 1, len(lines), start + MAX_EDIT_SNIPPET_LINES)
+    snippet = "\n".join(f"{start + i + 1}: {line}" for i, line in enumerate(lines[start:end]))
+    return _bound(snippet)
 
 
 def write_file(args: dict, render_context: RenderContext) -> str:
@@ -166,25 +200,37 @@ def delete_file(args: dict, render_context: RenderContext) -> str:
     return f"Deleted '{full_path}'."
 
 
-def run_unit_tests(_args: dict, render_context: RenderContext) -> str:
+def full_log_pointer(log_file_path: str | None) -> str:
+    return f" Full log: {log_file_path} (search it with grep, passing that path as file_path)." if log_file_path else ""
+
+
+def run_unit_tests(_args: dict, render_context: RenderContext) -> dict:
+    """Returns a result dict: a short header in `output` and the raw failure output in
+    `test_output`, which the server condenses (summarizing it when it is long)."""
     exit_code, output, log_file_path = render_utils.execute_script(
         os.path.normpath(render_context.unittests_script),
         [render_context.build_folder],
         "Unit Tests",
         timeout=render_context.test_script_timeout,
         stop_event=render_context.stop_event,
     )
+    context = render_context.unit_tests_running_context
     if exit_code == 0:
-        return "All unit tests passed."
+        context.verified_passing, context.verified_passing_log_path = True, log_file_path
+        return {"output": "All unit tests passed."}
     if not log_file_path and output:
         with tempfile.NamedTemporaryFile("w", encoding="utf-8", delete=False, suffix=".unittest_output") as f:
             f.write(output)
             log_file_path = f.name
-    pointer = f" Full output: {log_file_path} (use read_file)." if log_file_path else ""
-    return f"Unit tests failed (exit code {exit_code}).{pointer}\n{_bound(output)}"
+    if log_file_path:
+        register_log_path(log_file_path, render_context)
+    return {
+        "output": f"Unit tests failed (exit code {exit_code}).{full_log_pointer(log_file_path)}",
+        "test_output": output,
+    }
 
 
-TOOLS: dict[str, Callable[[dict, RenderContext], str]] = {
+TOOLS: dict[str, Callable[[dict, RenderContext], str | dict]] = {
     "read_file": read_file,
     "grep": grep,
     "ls_files": ls_files,
@@ -196,17 +242,31 @@ def run_unit_tests(_args: dict, render_context: RenderContext) -> str:
 
 
 def execute_calls(calls: list[dict], render_context: RenderContext) -> list[dict]:
-    """Execute the agent's tool calls in order; every call gets a result, errors included."""
+    """Execute the agent's tool calls in order; every call gets a result, errors included.
+
+    A repeated read-only call (same tool and arguments, no file changed in between) is not
+    re-executed: its result is already in the conversation, so a short pointer is returned."""
+    cache = render_context.unit_tests_running_context.tool_result_cache
     results = []
     for call in calls:
         tool = TOOLS.get(call["name"])
+        cache_key = json.dumps([call["name"], call.get("args") or {}], sort_keys=True)
+        result: dict
         if tool is None:
-            output = f"Error: unknown tool '{call['name']}'."
+            result = {"output": f"Error: unknown tool '{call['name']}'."}
+        elif call["name"] in READ_ONLY_TOOLS and cache_key in cache:
+            result = {
+                "output": "Same call as an earlier one and no file has changed since; "
+                "its result is unchanged (see the earlier result above)."
+            }
         else:
             try:
                 output = tool(call.get("args") or {}, render_context)
+                result = output if isinstance(output, dict) else {"output": output}
             except Exception as e:
-                output = f"Error: tool '{call['name']}' failed: {type(e).__name__}: {e}"
-        console.debug(f"Agent tool {call['name']}({call.get('args')}) -> {output[:200]!r}")
-        results.append({"call_id": call["id"], "output": output})
+                result = {"output": f"Error: tool '{call['name']}' failed: {type(e).__name__}: {e}"}
+            if call["name"] in READ_ONLY_TOOLS and not result["output"].startswith("Error"):
+                cache[cache_key] = result["output"]
+        console.debug(f"Agent tool {call['name']}({call.get('args')}) -> {result['output'][:200]!r}")
+        results.append({"call_id": call["id"], **result})
     return results

render_machine/render_types.py

@@ -60,6 +60,18 @@ class UnitTestsRunningContext:
     # run's outcome, plus results of any tool calls made in the same turn as submit_fix.
     pending_submit_call_id: Optional[str] = None
     pending_tool_results: list[dict] = field(default_factory=list)
+    # Session abandoned without a submission (turn budget used up, LLM failure); the next session
+    # starts with a digest of what it tried.
+    previous_session_id: Optional[str] = None
+    # Set when the agent's own run_unit_tests passed and no file changed since, so the harness
+    # can accept the fix without running the suite again.
+    verified_passing: bool = False
+    verified_passing_log_path: Optional[str] = None
+    # Full test logs the agent was pointed to; readable by read_file/grep although outside the
+    # build folder.
+    readable_log_paths: set[str] = field(default_factory=set)
+    # Results of read-only tool calls, keyed by call; cleared whenever a file changes.
+    tool_result_cache: dict[str, str] = field(default_factory=dict)
 
 
 class ConformanceTestsRunningContext:

tests/test_agent_tools.py

@@ -6,6 +6,7 @@
 import pytest
 
 from render_machine import agent_tools
+from render_machine.render_types import UnitTestsRunningContext
 
 
 @pytest.fixture
@@ -17,7 +18,7 @@ def project(tmp_path, monkeypatch):
     monkeypatch.chdir(tmp_path)
     render_context = SimpleNamespace(
         build_folder=str(build),
-        unit_tests_running_context=SimpleNamespace(changed_files=set()),
+        unit_tests_running_context=UnitTestsRunningContext(fix_attempts=1),
         unittests_script=None,
         test_script_timeout=None,
         stop_event=None,
@@ -97,14 +98,56 @@ def boom(_args, _rc):
 
 
 def test_run_unit_tests_reports_pass_and_failure(project, monkeypatch):
-    outcomes = iter([(0, "", None), (1, "FAILED test_x", "/tmp/log.txt")])
+    log = project.root.parent / "log.txt"
+    log.write_text("FAILED test_x\nCaused by: boom\n")
+    outcomes = iter([(0, "", "/tmp/pass.txt"), (1, "FAILED test_x", str(log))])
     monkeypatch.setattr(agent_tools.render_utils, "execute_script", lambda *a, **k: next(outcomes))
     project.rc.unittests_script = "run_tests.sh"
-    assert agent_tools.run_unit_tests({}, project.rc) == "All unit tests passed."
+    context = project.rc.unit_tests_running_context
+
+    assert agent_tools.run_unit_tests({}, project.rc) == {"output": "All unit tests passed."}
+    assert context.verified_passing and context.verified_passing_log_path == "/tmp/pass.txt"
+
     failure = agent_tools.run_unit_tests({}, project.rc)
-    assert (
-        failure.startswith("Unit tests failed (exit code 1). Full output: /tmp/log.txt") and "FAILED test_x" in failure
+    assert failure["output"].startswith(f"Unit tests failed (exit code 1). Full log: {log}")
+    # raw output goes to the server for condensing, not truncated here
+    assert failure["test_output"] == "FAILED test_x"
+    # the full log is outside the build folder and project root but greppable
+    assert agent_tools.grep({"pattern": "Caused by", "file_path": str(log)}, project.rc).endswith("Caused by: boom")
+
+
+def test_file_change_invalidates_verified_pass_and_read_cache(project):
+    context = project.rc.unit_tests_running_context
+    context.verified_passing = True
+    first = agent_tools.execute_calls([{"id": "1", "name": "read_file", "args": {"file_path": "app.py"}}], project.rc)
+    repeat = agent_tools.execute_calls([{"id": "2", "name": "read_file", "args": {"file_path": "app.py"}}], project.rc)
+    assert first[0]["output"].startswith("1: def add")
+    assert repeat[0]["output"].startswith("Same call as an earlier one")
+
+    agent_tools.write_file({"file_path": "other.py", "content": "y = 1\n"}, project.rc)
+    assert context.verified_passing is False
+    again = agent_tools.execute_calls([{"id": "3", "name": "read_file", "args": {"file_path": "app.py"}}], project.rc)
+    assert again[0]["output"].startswith("1: def add")
+
+
+def test_edit_file_returns_the_edited_region(project):
+    out = agent_tools.edit_file(
+        {
+            "file_path": "app.py",
+            "search": "def sub(a, b):\n    return a - b",
+            "replace": "def sub(a, b):\n    return b",
+        },
+        project.rc,
     )
+    assert out.startswith("Edited") and "5: def sub(a, b):\n6:     return b" in out
+    assert "2:     return a - b" in out and "1: def add" not in out
+
+
+def test_grep_context_lines_and_include(project):
+    (project.build / "notes.txt").write_text("def sub is documented here\n")
+    hits = agent_tools.grep({"pattern": "def sub", "context_lines": 1, "include": "*.py"}, project.rc)
+    assert "notes.txt" not in hits
+    assert "app.py-4-" in hits and "app.py:5:def sub(a, b):" in hits and "app.py-6-" in hits
 
 
 def test_bound_truncates_long_lines_and_large_output():