perf(tool_executor): trim memory after publishing a >=1 MB tool result
Post-compression already calls trim_memory (#77356); a huge tool result (raw stdout, file dumps) is the other allocation a turn drops and was published with no collection. _commit_tool_result is the one point both the sequential and concurrent publish paths go through, so the trim lives there, after the spill + session flush, measured on the string already in hand (multimodal dicts are never re-serialised). trim_memory's own cooldown/kill-switch apply. Salvages the intent of #80974 without its bare gc.collect(), re-serialisation and 186 LOC. Closes #70684 (tool-result half). Co-authored-by: Christopher-Schulze <210261288+Christopher-Schulze@users.noreply.github.com>
This commit is contained in:
@@ -53,6 +53,11 @@ 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).
|
||||
_LARGE_TOOL_RESULT_TRIM_CHARS = 1_000_000
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -1098,6 +1103,8 @@ def _commit_tool_result(
|
||||
agent.tool_progress_callback, "Tool progress",
|
||||
"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")
|
||||
return persisted_result, function_result, tool_message.get("_tool_output_risk")
|
||||
|
||||
|
||||
|
||||
56
tests/agent/test_tool_result_memory_trim.py
Normal file
56
tests/agent/test_tool_result_memory_trim.py
Normal file
@@ -0,0 +1,56 @@
|
||||
"""Publishing a >=1 MB tool result hands allocator pages back via ``trim_memory`` (#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.
|
||||
"""
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from tests.agent.test_start_order_gate import ( # noqa: F401 — autouse fixture rides along
|
||||
_FakeAssistantMsg,
|
||||
_FakeToolCall,
|
||||
_isolate_hermes,
|
||||
_make_agent,
|
||||
)
|
||||
|
||||
|
||||
def _agent_returning(monkeypatch, payload):
|
||||
agent = _make_agent(monkeypatch)
|
||||
agent._tool_guardrails = MagicMock()
|
||||
agent._tool_guardrails.before_call = lambda name, args: MagicMock(allows_execution=True)
|
||||
agent._invoke_tool = MagicMock(return_value=payload)
|
||||
agent._append_guardrail_observation = lambda name, args, result, *a, **kw: result
|
||||
return agent
|
||||
|
||||
|
||||
def test_large_sequential_result_trims_memory_once(monkeypatch):
|
||||
import agent.tool_executor as te
|
||||
|
||||
trim = MagicMock(return_value=True)
|
||||
monkeypatch.setattr(te, "trim_memory", trim)
|
||||
agent = _agent_returning(monkeypatch, "x" * 1_000_000)
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
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()
|
||||
Reference in New Issue
Block a user