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:
Teknium
2026-10-06 22:25:07 -07:00
parent a4465eb21d
commit 13eb190479
5 changed files with 69 additions and 6 deletions

View File

@@ -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

View File

@@ -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.

View File

@@ -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:

View File

@@ -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."

View File

@@ -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)