diff --git a/model_tools.py b/model_tools.py index c37007c41..6582f0b54 100644 --- a/model_tools.py +++ b/model_tools.py @@ -540,6 +540,29 @@ def handle_function_call( except Exception: pass + # Poka-yoke: validate tool handler return type. + # Handlers MUST return a JSON string. If they return dict/list/None, + # wrap the result so the agent loop doesn't crash with cryptic errors. + if not isinstance(result, str): + logger.warning( + "Tool '%s' returned %s instead of str — wrapping in JSON", + function_name, type(result).__name__, + ) + result = json.dumps( + {"output": str(result), "_type_warning": f"Tool returned {type(result).__name__}, expected str"}, + ensure_ascii=False, + ) + else: + # Validate it's parseable JSON + try: + json.loads(result) + except (json.JSONDecodeError, TypeError): + logger.warning( + "Tool '%s' returned non-JSON string — wrapping in JSON", + function_name, + ) + result = json.dumps({"output": result}, ensure_ascii=False) + return result except Exception as e: diff --git a/tests/test_model_tools.py b/tests/test_model_tools.py index 5e3b1d6ce..810babb47 100644 --- a/tests/test_model_tools.py +++ b/tests/test_model_tools.py @@ -137,3 +137,78 @@ class TestBackwardCompat: def test_tool_to_toolset_map(self): assert isinstance(TOOL_TO_TOOLSET_MAP, dict) assert len(TOOL_TO_TOOLSET_MAP) > 0 + + +class TestToolReturnTypeValidation: + """Poka-yoke: tool handlers must return JSON strings.""" + + def test_handler_returning_dict_is_wrapped(self, monkeypatch): + """A handler that returns a dict should be auto-wrapped to JSON string.""" + from tools.registry import registry + from model_tools import handle_function_call + import json + + # Register a bad handler that returns dict instead of str + registry.register( + name="__test_bad_dict", + toolset="test", + schema={"name": "__test_bad_dict", "description": "test", "parameters": {"type": "object", "properties": {}}}, + handler=lambda args, **kw: {"this is": "a dict not a string"}, + ) + result = handle_function_call("__test_bad_dict", {}) + parsed = json.loads(result) + assert "output" in parsed + assert "_type_warning" in parsed + # Cleanup + registry._tools.pop("__test_bad_dict", None) + + def test_handler_returning_none_is_wrapped(self, monkeypatch): + """A handler that returns None should be auto-wrapped.""" + from tools.registry import registry + from model_tools import handle_function_call + import json + + registry.register( + name="__test_bad_none", + toolset="test", + schema={"name": "__test_bad_none", "description": "test", "parameters": {"type": "object", "properties": {}}}, + handler=lambda args, **kw: None, + ) + result = handle_function_call("__test_bad_none", {}) + parsed = json.loads(result) + assert "_type_warning" in parsed + registry._tools.pop("__test_bad_none", None) + + def test_handler_returning_non_json_string_is_wrapped(self): + """A handler returning a plain string (not JSON) should be wrapped.""" + from tools.registry import registry + from model_tools import handle_function_call + import json + + registry.register( + name="__test_bad_plain", + toolset="test", + schema={"name": "__test_bad_plain", "description": "test", "parameters": {"type": "object", "properties": {}}}, + handler=lambda args, **kw: "just a plain string, not json", + ) + result = handle_function_call("__test_bad_plain", {}) + parsed = json.loads(result) + assert "output" in parsed + registry._tools.pop("__test_bad_plain", None) + + def test_handler_returning_valid_json_passes_through(self): + """A handler returning valid JSON string passes through unchanged.""" + from tools.registry import registry + from model_tools import handle_function_call + import json + + registry.register( + name="__test_good", + toolset="test", + schema={"name": "__test_good", "description": "test", "parameters": {"type": "object", "properties": {}}}, + handler=lambda args, **kw: json.dumps({"status": "ok", "data": [1, 2, 3]}), + ) + result = handle_function_call("__test_good", {}) + parsed = json.loads(result) + assert parsed == {"status": "ok", "data": [1, 2, 3]} + registry._tools.pop("__test_good", None)