fix(tools): image-bearing tool results stay inside the turn budget; spills and refusals keep their content (review)
Greptile on b17398db089: - `enforce_turn_budget` took `len(content)` of every tool message. An image-bearing result reaches it as a part list, so it counted parts, not characters: six 60K-char MCP results with screenshots totalled ~18 "chars" and nothing was spilled. Part lists now count their text characters and spill their largest text part (images stay inline). - `_persist_multimodal_text_parts` spilled every text part under the same tool_call_id, so with a small `tool_budget.mcp_result_size_chars` the notes part overwrote the tool text's spill file. Each spilled part now gets its own id. - When `vision.max_calls_per_image` refused every image, the refusal lines were appended after the handler's JSON closing brace. They now go into the envelope's `result`.
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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."
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user