fix(tools): bound tool error bodies at the dispatch boundary

tool_error() caps its message at _MAX_TOOL_ERROR_CHARS (2048), logging
the full body at DEBUG before trimming the context-bound copy.

Handlers that serialize exceptions directly -- json.dumps({"error":
str(exc), ...}) -- bypass that helper, so _normalize_handler_result
also runs every string result through _bound_json_error_result: if it
parses as a JSON object with an oversized string error field, only
that field is trimmed and the payload re-serialized. Non-error
results, non-JSON strings, and multimodal envelopes pass untouched.
This commit is contained in:
Erosika 2026-06-02 17:24:49 -04:00 committed by kshitij
parent 350f366a81
commit 2181d2e7c2
2 changed files with 126 additions and 3 deletions

View File

@ -6,7 +6,13 @@ import threading
from pathlib import Path
from unittest.mock import patch
from tools.registry import ToolRegistry, _module_registers_tools, discover_builtin_tools
from tools.registry import (
ToolRegistry,
_MAX_TOOL_ERROR_CHARS,
_module_registers_tools,
discover_builtin_tools,
tool_error,
)
def _dummy_handler(args, **kwargs):
@ -138,6 +144,87 @@ class TestUnknownToolDispatch:
assert "Unknown tool" in result["error"]
class TestToolErrorBounding:
def test_short_message_unchanged(self):
result = json.loads(tool_error("Missing required parameter: query"))
assert result["error"] == "Missing required parameter: query"
def test_extra_kwargs_preserved(self):
result = json.loads(tool_error("bad input", success=False))
assert result["error"] == "bad input"
assert result["success"] is False
def test_oversized_body_truncated(self):
result = json.loads(tool_error("boom: " + "X" * 5000))
assert result["error"].endswith("… [truncated]")
assert len(result["error"]) <= _MAX_TOOL_ERROR_CHARS + len("… [truncated]")
def test_at_limit_not_truncated(self):
msg = "Y" * _MAX_TOOL_ERROR_CHARS
result = json.loads(tool_error(msg))
assert result["error"] == msg
def test_full_body_preserved_in_logs_when_truncated(self, caplog):
import logging
body = "boom: " + "Z" * 5000
with caplog.at_level(logging.DEBUG, logger="tools.registry"):
json.loads(tool_error(body))
assert any(body in rec.getMessage() for rec in caplog.records)
class TestDispatchBoundsDirectErrorResults:
"""Handlers that bypass tool_error() and serialize errors directly are
still bounded at the dispatch boundary."""
@staticmethod
def _register(reg, name, handler):
reg.register(
name=name,
toolset="core",
schema=_make_schema(name),
handler=handler,
)
def test_direct_json_error_result_truncated(self):
reg = ToolRegistry()
self._register(reg, "direct", lambda args, **kw: json.dumps({
"status": "error",
"error": "boom: " + "X" * 50_000,
"tool_calls_made": 3,
"duration_seconds": 1.2,
}, ensure_ascii=False))
result = json.loads(reg.dispatch("direct", {}))
assert result["error"].endswith("… [truncated]")
assert len(result["error"]) <= _MAX_TOOL_ERROR_CHARS + len("… [truncated]")
assert result["status"] == "error"
assert result["tool_calls_made"] == 3
assert result["duration_seconds"] == 1.2
def test_small_error_result_unchanged(self):
reg = ToolRegistry()
payload = json.dumps({"error": "not found", "success": False})
self._register(reg, "small", lambda args, **kw: payload)
assert reg.dispatch("small", {}) == payload
def test_oversized_non_error_result_unchanged(self):
reg = ToolRegistry()
payload = json.dumps({"data": "D" * 50_000})
self._register(reg, "big_data", lambda args, **kw: payload)
assert reg.dispatch("big_data", {}) == payload
def test_oversized_non_json_result_unchanged(self):
reg = ToolRegistry()
payload = "plain text " * 10_000
self._register(reg, "plain", lambda args, **kw: payload)
assert reg.dispatch("plain", {}) == payload
def test_non_string_error_value_unchanged(self):
reg = ToolRegistry()
payload = json.dumps({"error": {"detail": "E" * 5_000}})
self._register(reg, "nested", lambda args, **kw: payload)
assert reg.dispatch("nested", {}) == payload
class TestToolsetAvailability:
def test_no_check_fn_is_available(self):
reg = ToolRegistry()

View File

@ -26,6 +26,41 @@ from typing import Callable, Dict, List, Optional, Set
logger = logging.getLogger(__name__)
# Cap on a tool error body; only trims runaway interpolated exceptions (static msgs are ~115 chars).
_MAX_TOOL_ERROR_CHARS = 2048
_TOOL_ERROR_TRUNCATION_MARKER = "… [truncated]"
def _bound_error_text(text: str) -> str:
"""Bound an error body destined for model context; full body stays in logs."""
if len(text) <= _MAX_TOOL_ERROR_CHARS:
return text
logger.debug("tool error body truncated for context (%d chars): %s", len(text), text)
return text[:_MAX_TOOL_ERROR_CHARS] + _TOOL_ERROR_TRUNCATION_MARKER
def _bound_json_error_result(result: str) -> str:
"""Trim an oversized ``error`` field in a JSON string result.
Handlers that serialize exceptions directly ``json.dumps({"error":
str(exc), ...})`` instead of ``tool_error()`` bypass the cap in
``tool_error``. Applied at the dispatch boundary so no registered tool
can return an unbounded error body that stacks across retries.
"""
if len(result) <= _MAX_TOOL_ERROR_CHARS or '"error"' not in result:
return result
try:
payload = json.loads(result)
except ValueError:
return result
if not isinstance(payload, dict):
return result
error = payload.get("error")
if not isinstance(error, str) or len(error) <= _MAX_TOOL_ERROR_CHARS:
return result
payload["error"] = _bound_error_text(error)
return json.dumps(payload, ensure_ascii=False)
def _is_registry_register_call(node: ast.AST) -> bool:
"""Return True when *node* is a ``registry.register(...)`` call expression."""
@ -736,7 +771,7 @@ class ToolRegistry:
persistence from receiving values they cannot safely slice or size.
"""
if isinstance(result, str):
return result
return _bound_json_error_result(result)
if (
isinstance(result, dict)
and result.get("_multimodal") is True
@ -935,7 +970,8 @@ def tool_error(message, **extra) -> str:
>>> tool_error("bad input", success=False)
'{"error": "bad input", "success": false}'
"""
result = {"error": str(message)}
# Bound the context-bound copy so a raw exception can't bloat history across retries.
result = {"error": _bound_error_text(str(message))}
if extra:
result.update(extra)
return json.dumps(result, ensure_ascii=False)