Fix unit tests with a server-side agent session

Replace the stateless /fix_unittests_issue call with an agent loop: the first
failure starts a session via /agent/start, the client executes the tool calls
the model asks for (read/grep/ls/edit/write/delete files, run the unit tests)
and posts results to /agent/continue until the agent calls submit_fix. The
state machine then re-runs the unit tests as before; a repeated failure is fed
back into the same session as the submit_fix result, so earlier attempts stay
in context instead of being re-diagnosed from scratch.

- render_machine/agent_tools.py: tool implementations confined to the build
  folder (writes) and project root (reads); changed files are tracked in the
  unit tests running context so they reach the FRID commit as before.
- FixUnitTests: start/continue the session, per-attempt turn cap, fresh
  session when the previous one ended without a submission.
- UnitTestsRunningContext carries the session id and the pending submit_fix
  call; it is recreated per unit-test loop so a session spans one FRID.

Author: VitjanZ

codeplain_REST_api.py

@@ -227,38 +227,6 @@ def render_functional_requirement(
 
         return self.post_request(endpoint_url, headers, payload, run_state)
 
-    def fix_unittests_issue(
-        self,
-        frid,
-        plain_source_tree,
-        linked_resources,
-        existing_files_content,
-        module_name: str,
-        required_modules,
-        unittests_issue,
-        run_state: RunState,
-        conformance_tests_fixes: list[dict] | None = None,
-    ):
-        endpoint_url = f"{self.api_url}/fix_unittests_issue"
-        headers = {"X-API-Key": self.api_key, "Content-Type": "application/json"}
-
-        payload = {
-            "frid": frid,
-            "plain_source_tree": plain_source_tree,
-            "linked_resources": linked_resources,
-            "existing_files_content": existing_files_content,
-            "module_name": module_name,
-            "required_modules": required_modules,
-            "unittests_issue": unittests_issue,
-            "unittest_batch_id": run_state.unittest_batch_id,
-        }
-        # Implementation code changes made by the conformance tests fixer right before this unit tests run;
-        # only sent when unit tests are processed inside the conformance tests phase.
-        if conformance_tests_fixes is not None:
-            payload["conformance_tests_fixes"] = conformance_tests_fixes
-
-        return self.post_request(endpoint_url, headers, payload, run_state)
-
     def distill_conformance_test_memory(
         self,
         frid,
@@ -532,3 +500,29 @@ def summarize_finished_conformance_tests(
         }
 
         return self.post_request(endpoint_url, headers, payload, run_state)
+
+    def agent_start(self, task_type: str, task_params: dict, frid: str, module_name: str, run_state: RunState):
+        """Start a server-side agent session; returns the first turn (tool calls or completion)."""
+        endpoint_url = f"{self.api_url}/agent/start"
+        headers = {"X-API-Key": self.api_key, "Content-Type": "application/json"}
+        payload = {
+            "task_type": task_type,
+            "task_params": task_params,
+            "frid": frid,
+            "module_name": module_name,
+        }
+        return self.post_request(endpoint_url, headers, payload, run_state)
+
+    def agent_continue(
+        self, session_id: str, tool_results: list[dict], frid: str, module_name: str, run_state: RunState
+    ):
+        """Feed tool results into an agent session and run its next turn."""
+        endpoint_url = f"{self.api_url}/agent/continue"
+        headers = {"X-API-Key": self.api_key, "Content-Type": "application/json"}
+        payload = {
+            "session_id": session_id,
+            "tool_results": tool_results,
+            "frid": frid,
+            "module_name": module_name,
+        }
+        return self.post_request(endpoint_url, headers, payload, run_state)

render_machine/actions/fix_unit_tests.py

@@ -1,79 +1,139 @@
 from typing import Any
 
 import file_utils
-import render_machine.render_utils as render_utils
+import plain_spec
 from plain2code_console import console
 from plain2code_exceptions import InternalClientError
+from render_machine import agent_tools
 from render_machine.actions.base_action import BaseAction
-from render_machine.implementation_code_helpers import ImplementationCodeHelpers
 from render_machine.render_context import RenderContext
 
-MAX_ISSUE_LENGTH = 10000
+TASK_TYPE = "fix_unit_tests"
+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
 
 
 class FixUnitTests(BaseAction):
+    """Fix failing unit tests with a server-side agent whose session spans all attempts for the FRID.
+
+    The first failure starts a session; the agent then drives read/grep/edit/run tool calls
+    (executed here) until it calls submit_fix. The state machine re-runs the unit tests and,
+    if they still fail, the next execution of this action answers that submit_fix call with
+    the new failure output inside the same session, so earlier attempts stay in context.
+    """
+
     SUCCESSFUL_OUTCOME = "unit_tests_fix_generated"
 
     def execute(self, render_context: RenderContext, previous_action_payload: Any | None):
-        if not previous_action_payload.get("previous_unittests_issue"):
+        if not previous_action_payload or not previous_action_payload.get("previous_unittests_issue"):
             raise InternalClientError(
                 "Internal client error: Previous action payload does not contain previous unit tests issue."
             )
-        previous_unittests_issue = previous_action_payload["previous_unittests_issue"]
-
-        if previous_unittests_issue and len(previous_unittests_issue) > MAX_ISSUE_LENGTH:
-            console.debug(
-                f"Unit tests issue text is too long and will be smartly truncated to {MAX_ISSUE_LENGTH} characters."
+        unittests_issue = previous_action_payload["previous_unittests_issue"]
+        context = render_context.unit_tests_running_context
+        api = render_context.codeplain_api
+        frid, module_name = render_context.frid_context.frid, render_context.module_name
+        changed_files_before = set(context.changed_files)
+
+        if context.agent_session_id is None:
+            console.info("Starting an agent session to fix the unit tests.")
+            response = api.agent_start(
+                TASK_TYPE,
+                self._build_task_params(render_context, unittests_issue),
+                frid,
+                module_name,
+                render_context.run_state,
             )
-
-        existing_files, existing_files_content = ImplementationCodeHelpers.fetch_existing_files(
-            render_context.build_folder
-        )
-
-        render_utils.print_inputs(render_context, existing_files_content, "Files sent as input to unit tests fixing:")
-
-        conformance_tests_fixes = self._get_conformance_tests_fixes(render_context)
-        if conformance_tests_fixes:
-            console.info(
-                f"Unit tests are fixed while preserving {len(conformance_tests_fixes)} implementation code change(s) "
-                "made to fix the conformance tests."
+            context.agent_session_id = response["session_id"]
+        else:
+            console.info(f"Continuing agent session {context.agent_session_id} with the new unit tests failure.")
+            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}",
+                }
+            ]
+            context.pending_tool_results, context.pending_submit_call_id = [], None
+            response = api.agent_continue(
+                context.agent_session_id, tool_results, frid, module_name, render_context.run_state
             )
 
-        response_files = render_context.codeplain_api.fix_unittests_issue(
-            render_context.frid_context.frid,
-            render_context.plain_source_tree,
-            render_context.frid_context.linked_resources,
-            existing_files_content,
-            render_context.module_name,
-            render_context.get_required_modules_functionalities(),
-            previous_unittests_issue,
-            run_state=render_context.run_state,
-            conformance_tests_fixes=conformance_tests_fixes,
-        )
+        submitted = False
+        turns = 0
+        while response.get("status") == "tool_calls" and turns < MAX_AGENT_TURNS_PER_ATTEMPT:
+            turns += 1
+            calls = response["calls"]
+            submit_call = next((call for call in calls if call["name"] == SUBMIT_FIX_TOOL), None)
+            tool_results = agent_tools.execute_calls(
+                [call for call in calls if call is not submit_call], render_context
+            )
+            if submit_call is not None:
+                context.pending_submit_call_id = submit_call["id"]
+                context.pending_tool_results = tool_results
+                submitted = True
+                console.info(f"Agent submitted a fix: {submit_call['args'].get('changes_made', '')}")
+                break
+            response = api.agent_continue(
+                context.agent_session_id, tool_results, frid, module_name, render_context.run_state
+            )
 
-        _, changed_files = file_utils.update_build_folder_with_rendered_files(
-            render_context.build_folder, existing_files, response_files
+        if not submitted:
+            # The session ended without a submission (finished in text, failed, or used up this
+            # attempt's turn budget) — a fresh session is started if the tests still fail.
+            status = response.get("status")
+            if status == "failed":
+                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.agent_session_id, context.pending_submit_call_id, context.pending_tool_results = None, None, []
+
+        console.print_files(
+            "Files changed while fixing unit tests:",
+            render_context.build_folder,
+            {path: "" for path in sorted(context.changed_files - changed_files_before)},
+            style=console.OUTPUT_STYLE,
         )
-
-        render_context.unit_tests_running_context.changed_files.update(changed_files)
-
-        console.print_files("Files fixed:", render_context.build_folder, response_files, style=console.OUTPUT_STYLE)
-
         return self.SUCCESSFUL_OUTCOME, None
 
     @staticmethod
-    def _get_conformance_tests_fixes(render_context: RenderContext) -> list[dict] | None:
-        """Implementation code changes the conformance tests fixer made before these unit tests were run.
-
-        Only present when unit tests are processed inside the conformance tests phase - the implementation
-        and refactoring unit test passes have no conformance tests running context.
-        """
-        conformance_tests_running_context = getattr(render_context, "conformance_tests_running_context", None)
-        if conformance_tests_running_context is None:
-            return None
+    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 {
+            "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),
+            "linked_resources": render_context.frid_context.linked_resources,
+            "build_folder": render_context.build_folder,
+            "module_name": render_context.module_name,
+            "unittests_script_content": FixUnitTests._read_script(render_context.unittests_script),
+            "unittests_issue": unittests_issue,
+        }
 
-        implementation_code_fixes = getattr(conformance_tests_running_context, "implementation_code_fixes", None)
-        if not implementation_code_fixes:
-            return None
+    @staticmethod
+    def _functional_requirements_section(render_context: RenderContext, specifications: dict) -> str:
+        sections = []
+        for module_name, functionalities in render_context.get_required_modules_functionalities().items():
+            sections.append(
+                f"### Module: {module_name} (Already Implemented, for context)\n" + "\n".join(functionalities)
+            )
+        current = specifications.get(plain_spec.FUNCTIONAL_REQUIREMENTS, [])
+        if len(current) > 1:
+            sections.append(
+                f"### Module: {render_context.module_name} (Already Implemented, for context)\n"
+                + "\n".join(current[:-1])
+            )
+        if current:
+            sections.append(f"### Module: {render_context.module_name} (Currently Being Implemented)\n{current[-1]}")
+        return "\n\n".join(sections)
 
-        return list(implementation_code_fixes)
+    @staticmethod
+    def _read_script(script: str | None) -> str:
+        if not script:
+            return ""
+        try:
+            with open(file_utils.add_current_path_if_no_path(script), "r", encoding="utf-8") as f:
+                return f.read()
+        except OSError:
+            return ""