From f8361d29c8e2a2be6ba9ada32f1d694bf47a4b6a Mon Sep 17 00:00:00 2001
From: Teknium <127238744+teknium1@users.noreply.github.com>
Date: Thu, 9 Jul 2026 21:32:01 -0700
Subject: [PATCH] fix(tools): enforce registry result contract (#61787)

---
 tests/tools/test_registry.py | 64 ++++++++++++++++++++++++++++++++++++
 tools/registry.py            | 41 +++++++++++++++++++++--
 2 files changed, 102 insertions(+), 3 deletions(-)

diff --git a/tests/tools/test_registry.py b/tests/tools/test_registry.py
index 4d8b56556d7..49d1193cf67 100644
--- a/tests/tools/test_registry.py
+++ b/tests/tools/test_registry.py
@@ -47,6 +47,70 @@ class TestRegisterAndDispatch:
         result = json.loads(reg.dispatch("echo", {"msg": "hi"}))
         assert result == {"msg": "hi"}
 
+    def test_dispatch_preserves_supported_multimodal_result(self):
+        reg = ToolRegistry()
+        multimodal = {
+            "_multimodal": True,
+            "content": [{"type": "text", "text": "captured"}],
+            "text_summary": "captured",
+        }
+        reg.register(
+            name="capture",
+            toolset="computer_use",
+            schema=_make_schema("capture"),
+            handler=lambda args, **kw: multimodal,
+        )
+
+        assert reg.dispatch("capture", {}) is multimodal
+
+    def test_dispatch_rejects_unsupported_handler_results_with_structured_error(self):
+        invalid_results = ({"ok": True}, b"bytes", None, 42)
+
+        for invalid in invalid_results:
+            reg = ToolRegistry()
+            reg.register(
+                name="bad_result",
+                toolset="core",
+                schema=_make_schema("bad_result"),
+                handler=lambda args, _invalid=invalid, **kw: _invalid,
+            )
+
+            raw = reg.dispatch("bad_result", {})
+            result = json.loads(raw)
+
+            assert isinstance(raw, str)
+            assert result["error_type"] == "tool_result_contract"
+            assert result["tool"] == "bad_result"
+            assert result["result_type"] == type(invalid).__name__
+            assert "unsupported result type" in result["error"]
+
+    def test_handler_contract_error_survives_model_tools_pipeline(self):
+        from model_tools import handle_function_call, registry
+
+        name = "test_invalid_registry_result"
+        registry.register(
+            name=name,
+            toolset="core",
+            schema=_make_schema(name),
+            handler=lambda args, **kw: None,
+        )
+        try:
+            raw = handle_function_call(
+                name,
+                {},
+                task_id="contract-test",
+                skip_pre_tool_call_hook=True,
+            )
+        finally:
+            registry.deregister(name)
+
+        result = json.loads(raw)
+        assert len(raw) > 0  # downstream sizing/logging remains safe
+        assert json.loads(json.dumps({"content": raw}))["content"] == raw
+        assert result["error_type"] == "tool_result_contract"
+        assert result["tool"] == name
+        assert result["result_type"] == "NoneType"
+
 
 class TestGetDefinitions:
     def test_returns_openai_format(self):
diff --git a/tools/registry.py b/tools/registry.py
index 35589bd2c84..9b6611fb407 100644
--- a/tools/registry.py
+++ b/tools/registry.py
@@ -571,10 +571,43 @@ class ToolRegistry:
     # Dispatch
     # ------------------------------------------------------------------
 
-    def dispatch(self, name: str, args: dict, **kwargs) -> str:
+    @staticmethod
+    def _normalize_handler_result(name: str, result):
+        """Enforce the result shapes supported by the agent tool pipeline.
+
+        Normal tool results are strings.  The sole structured exception is the
+        multimodal envelope consumed by the agent executor.  Returning every
+        other value as a string error keeps logging, hooks, budgeting, and
+        persistence from receiving values they cannot safely slice or size.
+        """
+        if isinstance(result, str):
+            return result
+        if (
+            isinstance(result, dict)
+            and result.get("_multimodal") is True
+            and isinstance(result.get("content"), list)
+        ):
+            return result
+
+        result_type = type(result).__name__
+        logger.error(
+            "Tool %s handler returned unsupported result type: %s",
+            name,
+            result_type,
+        )
+        return json.dumps({
+            "error": f"Tool handler returned unsupported result type: {result_type}",
+            "error_type": "tool_result_contract",
+            "tool": name,
+            "result_type": result_type,
+        }, ensure_ascii=False)
+
+    def dispatch(self, name: str, args: dict, **kwargs) -> str | dict:
         """Execute a tool handler by name.
 
         * Async handlers are bridged automatically via ``_run_async()``.
+        * Handler results are normalized to a string or supported multimodal
+          envelope before leaving the registry.
         * All exceptions are caught and returned as ``{"error": "..."}``
           for consistent error format.
         """
@@ -584,8 +617,10 @@ class ToolRegistry:
         try:
             if entry.is_async:
                 from model_tools import _run_async
-                return _run_async(entry.handler(args, **kwargs))
-            return entry.handler(args, **kwargs)
+                result = _run_async(entry.handler(args, **kwargs))
+            else:
+                result = entry.handler(args, **kwargs)
+            return self._normalize_handler_result(name, result)
         except Exception as e:
             logger.exception("Tool %s dispatch error: %s", name, e)
             # Route through the sanitizer so framing tokens / CDATA / fences