diff --git a/tests/tools/test_registry.py b/tests/tools/test_registry.py index 40f6900b17c59..bc1ea0f424fd5 100644 --- a/tests/tools/test_registry.py +++ b/tests/tools/test_registry.py @@ -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() diff --git a/tools/registry.py b/tools/registry.py index c92f8a0b9bcac..0684cb4da0bcf 100644 --- a/tools/registry.py +++ b/tools/registry.py @@ -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)