diff --git a/agent/tool_executor.py b/agent/tool_executor.py index 50a9e62599..ad27f8396a 100644 --- a/agent/tool_executor.py +++ b/agent/tool_executor.py @@ -1166,15 +1166,18 @@ def _persist_multimodal_text_parts(result: dict, tool_name: str, tool_call_id: s entirely and ride every later request inline. Image parts are left untouched (their size is governed by the vision embed budget); a fresh dict is returned so history is never mutated.""" parts = result.get("content") or [] - bounded_parts, first_replacement = [], None + bounded_parts, first_replacement, spilled = [], None, 0 for part in parts: text = part.get("text") if isinstance(part, dict) and part.get("type") == "text" else None if isinstance(text, str): - replaced = maybe_persist_tool_result(content=text, tool_name=tool_name, tool_use_id=tool_call_id, + # One spill file per part: a second oversized part under the same id would overwrite the first's file. + part_id = tool_call_id if not spilled else f"{tool_call_id}_part{spilled}" + replaced = maybe_persist_tool_result(content=text, tool_name=tool_name, tool_use_id=part_id, env=env, config=budget) if replaced != text: part = {**part, "text": replaced} first_replacement = first_replacement or replaced + spilled += 1 bounded_parts.append(part) if first_replacement is None: return result diff --git a/tests/tools/test_mcp_image_content.py b/tests/tools/test_mcp_image_content.py index f32583d649..b5bf211c4b 100644 --- a/tests/tools/test_mcp_image_content.py +++ b/tests/tools/test_mcp_image_content.py @@ -17,6 +17,7 @@ images natively. from __future__ import annotations import base64 +import json from types import SimpleNamespace @@ -190,6 +191,13 @@ class TestNativeImageAttach: assert any("multiply any coordinates you report by 2.00" in t for t in texts) # ... the map was not assert bounded["content"][-1]["type"] == "image_url" + # A tiny tool_budget.mcp_result_size_chars spills BOTH text parts: each gets its own file, so the notes' + # spill cannot overwrite the tool text's. + import re + both = _persist_multimodal_text_parts(out, "mcp__srv__snap", "call-2", None, BudgetConfig(mcp_result_size=50)) + paths = {re.search(r"saved to: (\S+)", p["text"]).group(1) for p in both["content"] if p.get("type") == "text"} + assert len(paths) == 2 + def test_image_prep_sees_the_callers_runtime(self, tmp_path, monkeypatch): """The prep pool runs in the caller's context: a managed local runtime (stb_image, no WebP) set on the calling session converts a small WebP instead of attaching bytes its server silently cannot decode.""" @@ -237,6 +245,7 @@ class TestNativeImageAttach: assert isinstance(self._call(monkeypatch, tmp_path, capped), dict) repeat = self._call(monkeypatch, tmp_path, capped) assert isinstance(repeat, str) and "already in context" in repeat and "MEDIA:" in repeat + assert "already in context" in json.loads(repeat)["result"] # still the handler's JSON envelope import io from PIL import Image # A valid JPEG header over a truncated pixel stream passes the cache's and the sniff's header checks. diff --git a/tests/tools/test_tool_result_storage.py b/tests/tools/test_tool_result_storage.py index f7b8c812da..3576365f42 100644 --- a/tests/tools/test_tool_result_storage.py +++ b/tests/tools/test_tool_result_storage.py @@ -341,6 +341,17 @@ class TestEnforceTurnBudget: result = enforce_turn_budget([], env=None, config=BudgetConfig(turn_budget=200_000)) assert result == [] + def test_multimodal_results_count_their_text_and_spill_it(self): + """Image-bearing results are part lists: the aggregate counts their text characters (not the number of + parts) and spills the big text part, keeping images and the short notes inline.""" + img = {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,AAAA"}} + msgs = [{"role": "tool", "tool_call_id": f"mm{i}", + "content": [{"type": "text", "text": "x" * 60_000}, {"type": "text", "text": "scale note"}, img]} + for i in range(6)] + enforce_turn_budget(msgs, env=None, config=BudgetConfig(turn_budget=200_000)) + assert sum(PERSISTED_OUTPUT_TAG in m["content"][0]["text"] for m in msgs) >= 2 + assert all(m["content"][1]["text"] == "scale note" and m["content"][2] == img for m in msgs) + # ── Per-tool threshold integration ──────────────────────────────────── class TestPerToolThresholds: diff --git a/tools/mcp_tool_content.py b/tools/mcp_tool_content.py index 21ac723e5c..bce0aa7e42 100644 --- a/tools/mcp_tool_content.py +++ b/tools/mcp_tool_content.py @@ -3,6 +3,7 @@ capping, _meta filtering, image/audio caching to MEDIA tags, resource links and embedded resources.""" import base64 +import json import logging import mimetypes from typing import Any, Dict, List, Optional, Tuple @@ -217,7 +218,15 @@ def _mcp_result_with_native_images(text: str, image_paths: List[str]) -> Any: if note: notes += f"\n- MEDIA:{path}: {note}" if not attached: - return text + notes + if not notes: + return text + # Keep the result valid JSON: the refusal lines ride inside the envelope, not after its closing brace. + try: + payload = json.loads(text) + payload["result"] = f"{payload.get('result') or ''}{notes}" + return json.dumps(payload, ensure_ascii=False) + except (TypeError, ValueError, AttributeError): + return text + notes # The header and the scale notes are their own short part: an oversized tool text gets spilled to a file and # replaced by its head, and the coordinate map must survive that next to the resized screenshot it describes. header = "The image(s) from this call are attached — inspect them with your native vision." diff --git a/tools/tool_result_storage.py b/tools/tool_result_storage.py index 658eb69352..2b9aba720d 100644 --- a/tools/tool_result_storage.py +++ b/tools/tool_result_storage.py @@ -334,14 +334,38 @@ def maybe_persist_tool_result(content: str, tool_name: str, tool_use_id: str, en "Full output could not be saved to sandbox.]") +def _content_text_size(content) -> int: + """Characters of TEXT a tool message carries: a string's length, or the summed text parts of a multimodal + part list (image parts are governed by the vision embed budget, and a list's ``len`` counts parts).""" + if isinstance(content, list): + return sum(len(p.get("text") or "") for p in content if isinstance(p, dict) and p.get("type") == "text") + return len(content or "") + + +def _spill_largest_text_part(parts: list, tool_use_id: str, env, config: BudgetConfig) -> tuple[list, int]: + """Persist the largest non-persisted text part of a multimodal list; ``(new parts, chars saved)``.""" + texts = [(i, p["text"]) for i, p in enumerate(parts) if isinstance(p, dict) and p.get("type") == "text" + and isinstance(p.get("text"), str) and PERSISTED_OUTPUT_TAG not in p["text"]] + if not texts: + return parts, 0 + idx, text = max(texts, key=lambda t: len(t[1])) + replacement = maybe_persist_tool_result(content=text, tool_name=_BUDGET_TOOL_NAME, + tool_use_id=f"{tool_use_id}_part{idx}", env=env, config=config, threshold=0) + if replacement == text: + return parts, 0 + return [*parts[:idx], {**parts[idx], "text": replacement}, *parts[idx + 1:]], len(text) - len(replacement) + + def enforce_turn_budget(tool_messages: list[dict], env=None, config: BudgetConfig = DEFAULT_BUDGET) -> list[dict]: """Layer 3: persist the largest non-persisted results first until the turn's aggregate is - under budget. Mutates the list in-place and returns it.""" - sizes = [len(msg.get("content", "")) for msg in tool_messages] + under budget. Mutates the list in-place and returns it. A multimodal result (part list) counts its + text characters and spills its largest text part; its images stay inline.""" + sizes = [_content_text_size(msg.get("content", "")) for msg in tool_messages] total_size = sum(sizes) candidates = [(i, size) for i, size in enumerate(sizes) - if PERSISTED_OUTPUT_TAG not in tool_messages[i].get("content", "")] + if isinstance(tool_messages[i].get("content"), list) + or PERSISTED_OUTPUT_TAG not in tool_messages[i].get("content", "")] if total_size <= config.turn_budget: return tool_messages for idx, size in sorted(candidates, key=lambda x: x[1], reverse=True): @@ -349,6 +373,13 @@ def enforce_turn_budget(tool_messages: list[dict], env=None, break content = tool_messages[idx]["content"] tool_use_id = tool_messages[idx].get("tool_call_id", f"budget_{idx}") + if isinstance(content, list): + parts, saved = _spill_largest_text_part(content, tool_use_id, env, config) + if saved: + total_size -= saved + tool_messages[idx]["content"] = parts + logger.info("Budget enforcement: persisted a text part of tool result %s (%d chars)", tool_use_id, saved) + continue replacement = maybe_persist_tool_result( content=content, tool_name=_BUDGET_TOOL_NAME, tool_use_id=tool_use_id, env=env, config=config, threshold=0)