diff --git a/agent/tool_executor.py b/agent/tool_executor.py index 81bc34a6d0..65cc84e268 100644 --- a/agent/tool_executor.py +++ b/agent/tool_executor.py @@ -53,10 +53,10 @@ from tools.tool_result_storage import ( extract_persisted_path, ) from tools.budget_config import BudgetConfig, DEFAULT_BUDGET, budget_for_context_window -from hermes_cli.mem_trim import trim_memory -# A tool result this large (raw stdout, file dumps) is the biggest allocation a turn ever drops; -# once spilled and flushed it is the natural point to hand allocator pages back (#70684). +# A tool result this large (raw stdout, file dumps) is the biggest allocation a turn ever drops. +# The commit only flags it: the string is still referenced by the publish frames here, so the +# trim runs once the whole batch has unwound (AIAgent._execute_tool_calls) (#70684). _LARGE_TOOL_RESULT_TRIM_CHARS = 1_000_000 logger = logging.getLogger(__name__) @@ -1104,7 +1104,7 @@ def _commit_tool_result( "tool.completed", function_name, None, None, duration=tool_duration, is_error=is_error, result=function_result, ) if isinstance(function_result, str) and len(function_result) >= _LARGE_TOOL_RESULT_TRIM_CHARS: - trim_memory(reason="large tool result") + agent._trim_after_tool_batch = True return persisted_result, function_result, tool_message.get("_tool_output_risk") diff --git a/run_agent.py b/run_agent.py index 9fe1c58e88..37fafd5bef 100644 --- a/run_agent.py +++ b/run_agent.py @@ -1333,6 +1333,12 @@ class AIAgent( return execute_tool_calls_segmented(self, *args, segments=segments) finally: self._executing_tools = False + if getattr(self, "_trim_after_tool_batch", False): + # Every executor frame that held a >=1 MB raw result has unwound; only the + # spilled preview lives in ``messages`` now (agent/tool_executor.py, #70684). + self._trim_after_tool_batch = False + from hermes_cli.mem_trim import trim_memory + trim_memory(reason="large tool result") def _dispatch_delegate_task(self, function_args: dict) -> str: """Single call site for delegate_task dispatch; new DELEGATE_TASK_SCHEMA fields are added only here.""" diff --git a/tests/agent/test_tool_result_memory_trim.py b/tests/agent/test_tool_result_memory_trim.py index d35eea1970..22073779cd 100644 --- a/tests/agent/test_tool_result_memory_trim.py +++ b/tests/agent/test_tool_result_memory_trim.py @@ -1,8 +1,10 @@ -"""Publishing a >=1 MB tool result hands allocator pages back via ``trim_memory`` (#70684). +"""A >=1 MB tool result hands allocator pages back via ``trim_memory`` — after the batch (#70684). -Compaction already trims after it frees the compressed-away messages; a huge tool -result (raw stdout, file dumps) is the other allocation a turn drops, and both publish -paths (sequential and concurrent) commit through the same point. +Compaction already trims after it frees the compressed-away messages; a huge tool result +(raw stdout, file dumps) is the other allocation a turn drops. The commit point only flags +it, because the string is still referenced by the publish frames there; the trim runs once +``AIAgent._execute_tool_calls`` has unwound every executor frame, so ``gc.collect`` + +``malloc_trim`` actually see the allocation as garbage. """ from unittest.mock import MagicMock @@ -24,33 +26,43 @@ def _agent_returning(monkeypatch, payload): return agent -def test_large_sequential_result_trims_memory_once(monkeypatch): - import agent.tool_executor as te +def _trim_recorder(monkeypatch, agent): + seen = [] + + def trim(*, reason): + seen.append((reason, agent._executing_tools)) + return True + + monkeypatch.setattr("hermes_cli.mem_trim.trim_memory", trim) + return seen + + +def test_large_result_flags_the_batch_and_small_does_not(monkeypatch): + big = _agent_returning(monkeypatch, "x" * 1_000_000) + small = _agent_returning(monkeypatch, "x" * 999_999) + for agent in (big, small): + seen = _trim_recorder(monkeypatch, agent) + messages: list = [] + agent._execute_tool_calls_concurrent(_FakeAssistantMsg([_FakeToolCall("terminal", "tc")]), messages, "task") + assert [m["role"] for m in messages] == ["tool"] + assert seen == [] # the commit never trims in-frame: the raw result is still referenced here + assert big._trim_after_tool_batch is True + assert getattr(small, "_trim_after_tool_batch", False) is False + + +def test_execute_tool_calls_trims_once_after_every_executor_frame_unwound(monkeypatch): + import run_agent as _ra - trim = MagicMock(return_value=True) - monkeypatch.setattr(te, "trim_memory", trim) agent = _agent_returning(monkeypatch, "x" * 1_000_000) + agent._execute_tool_calls = _ra.AIAgent._execute_tool_calls.__get__(agent) + # stand-in for the sequential executor: publishes through the real concurrent commit path + agent._execute_tool_calls_sequential = agent._execute_tool_calls_concurrent + seen = _trim_recorder(monkeypatch, agent) messages: list = [] - ref = te._ToolCallRef("terminal", {"command": "cat big.log"}, "task", "tc_big", []) - managed = te._ManagedToolResult("x" * 1_000_000, ref.args, [], blocked=False, dispatched=True) - assert te._publish_sequential_result( - agent, messages, ref, managed, tool_duration=0.1, index=1, budget=te.DEFAULT_BUDGET, - ) + agent._execute_tool_calls(_FakeAssistantMsg([_FakeToolCall("terminal", "tc_big")]), messages, "task") assert [m["role"] for m in messages] == ["tool"] - trim.assert_called_once_with(reason="large tool result") - - -def test_small_concurrent_result_does_not_trim(monkeypatch): - import agent.tool_executor as te - - trim = MagicMock(return_value=True) - monkeypatch.setattr(te, "trim_memory", trim) - agent = _agent_returning(monkeypatch, "x" * 999_999) - - messages: list = [] - agent._execute_tool_calls_concurrent(_FakeAssistantMsg([_FakeToolCall("terminal", "tc_small")]), messages, "task") - - assert [m["role"] for m in messages] == ["tool"] - trim.assert_not_called() + # one trim per batch, issued only after the tool-execution scope closed (frames holding the 1 MB str are gone) + assert seen == [("large tool result", False)] + assert agent._trim_after_tool_batch is False