fix(tools): accept both the string and structured form of every tool argument (#957)

Co-authored-by: Ahmed Allam <ahmed39652003@gmail.com>

strix/agents/factory.py

@@ -143,6 +143,83 @@ def _with_bounded_result(tool: FunctionTool) -> FunctionTool:
     return tool
 
 
+def _schema_types(spec: dict[str, Any]) -> set[str]:
+    types: set[str] = set()
+    raw = spec.get("type")
+    if isinstance(raw, str):
+        types.add(raw)
+    elif isinstance(raw, list):
+        types.update(t for t in raw if isinstance(t, str))
+    for variant in spec.get("anyOf") or ():
+        if isinstance(variant, dict):
+            types |= _schema_types(variant)
+    types.discard("null")
+    return types
+
+
+def _decode_structured(value: str, types: set[str]) -> Any:
+    stripped = value.strip()
+    if not stripped:
+        return value
+    try:
+        decoded = json.loads(stripped)
+    except json.JSONDecodeError:
+        return value
+    wanted = list if "array" in types else dict
+    return decoded if isinstance(decoded, wanted) else value
+
+
+def _coerce_argument(value: Any, spec: dict[str, Any]) -> Any:
+    types = _schema_types(spec)
+    if not types or value is None:
+        return value
+    if isinstance(value, list | dict) and "string" in types and not types & {"array", "object"}:
+        return json.dumps(value, ensure_ascii=False)
+    if isinstance(value, str) and types & {"array", "object"} and "string" not in types:
+        return _decode_structured(value, types)
+    return value
+
+
+def _coerce_arguments(raw_input: str, schema: dict[str, Any]) -> str:
+    properties = schema.get("properties")
+    if not isinstance(properties, dict) or not properties:
+        return raw_input
+    try:
+        payload = json.loads(raw_input) if raw_input else None
+    except json.JSONDecodeError:
+        return raw_input
+    if not isinstance(payload, dict):
+        return raw_input
+
+    changed = False
+    for key, value in payload.items():
+        spec = properties.get(key)
+        if not isinstance(spec, dict):
+            continue
+        coerced = _coerce_argument(value, spec)
+        if coerced is not value:
+            payload[key] = coerced
+            changed = True
+
+    if not changed:
+        return raw_input
+    return json.dumps(payload, ensure_ascii=False)
+
+
+def _with_coerced_arguments(tool: FunctionTool) -> FunctionTool:
+    if getattr(tool, "_strix_coerced", False):
+        return tool
+    invoke_tool = tool.on_invoke_tool
+    schema = tool.params_json_schema
+
+    async def invoke(ctx: Any, raw_input: str) -> Any:
+        return await invoke_tool(ctx, _coerce_arguments(raw_input, schema))
+
+    tool.on_invoke_tool = invoke
+    tool._strix_coerced = True  # type: ignore[attr-defined]
+    return tool
+
+
 def _function_tool_with_error_result(tool: FunctionTool) -> FunctionTool:
     invoke_tool = tool.on_invoke_tool
 
@@ -212,11 +289,13 @@ def _configure_filesystem_tools(toolset: Any, *, chat_completions: bool) -> None
             if isinstance(tool, CustomTool):
                 setattr(toolset, name, _custom_tool_as_function_tool(tool))
             elif isinstance(tool, FunctionTool):
-                setattr(toolset, name, _function_tool_with_error_result(tool))
+                setattr(
+                    toolset, name, _function_tool_with_error_result(_with_coerced_arguments(tool))
+                )
         elif isinstance(tool, CustomTool):
             setattr(toolset, name, _bound_custom_tool(tool))
         elif isinstance(tool, FunctionTool):
-            setattr(toolset, name, _with_bounded_result(tool))
+            setattr(toolset, name, _with_bounded_result(_with_coerced_arguments(tool)))
 
 
 def _make_filesystem_configurator(*, chat_completions: bool) -> Any:
@@ -329,7 +408,7 @@ def _configure_shell_tools(toolset: Any, *, chat_completions: bool) -> None:
     for name, tool in vars(toolset).items():
         if not isinstance(tool, FunctionTool):
             continue
-        wrapped = tool
+        wrapped = _with_coerced_arguments(tool)
         if tool.name == "exec_command":
             wrapped = _wrap_exec_command(wrapped)
         elif tool.name == "write_stdin":
@@ -523,7 +602,10 @@ def build_strix_agent(
         tools = [*_BASE_TOOLS, *agent_tools, agent_finish]
     _ensure_unique_tool_names(tools)
     tools = [
-        _with_bounded_result(tool) if isinstance(tool, FunctionTool) else tool for tool in tools
+        _with_bounded_result(_with_coerced_arguments(tool))
+        if isinstance(tool, FunctionTool)
+        else tool
+        for tool in tools
     ]
 
     logger.info(

tests/test_agent_factory_tool_arguments.py

@@ -0,0 +1,124 @@
+"""Tests for tool-argument shape coercion in the agent factory."""
+
+from __future__ import annotations
+
+import json
+from typing import Any, cast
+
+import pytest
+from agents.tool import FunctionTool
+
+from strix.agents import factory
+
+
+def _capturing_tool(captured: dict[str, str], schema: dict[str, Any]) -> FunctionTool:
+    async def invoke(_ctx: Any, raw_input: str) -> str:
+        captured["raw_input"] = raw_input
+        return "ok"
+
+    return FunctionTool(
+        name="probe",
+        description="test tool",
+        params_json_schema={"type": "object", "properties": schema},
+        on_invoke_tool=invoke,
+    )
+
+
+async def _roundtrip(schema: dict[str, Any], payload: dict[str, Any]) -> dict[str, Any]:
+    captured: dict[str, str] = {}
+    wrapped = factory._with_coerced_arguments(_capturing_tool(captured, schema))
+    assert await wrapped.on_invoke_tool(cast("Any", None), json.dumps(payload)) == "ok"
+    return cast("dict[str, Any]", json.loads(captured["raw_input"]))
+
+
+_STRING = {"todos": {"type": "string"}}
+_ARRAY = {"tags": {"type": "array", "items": {"type": "string"}}}
+_NULLABLE_ARRAY = {
+    "tags": {"anyOf": [{"type": "array", "items": {"type": "string"}}, {"type": "null"}]}
+}
+_OBJECT = {"modifications": {"type": "object"}}
+
+
+@pytest.mark.asyncio
+async def test_structured_value_is_encoded_for_a_string_parameter() -> None:
+    parsed = await _roundtrip(_STRING, {"todos": [{"title": "Phase 1: recon"}]})
+
+    assert parsed["todos"] == '[{"title": "Phase 1: recon"}]'
+
+
+@pytest.mark.asyncio
+async def test_string_parameter_keeps_an_already_encoded_value() -> None:
+    parsed = await _roundtrip(_STRING, {"todos": '[{"title": "a"}]'})
+
+    assert parsed["todos"] == '[{"title": "a"}]'
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("schema", [_ARRAY, _NULLABLE_ARRAY])
+async def test_encoded_list_is_decoded_for_an_array_parameter(schema: dict[str, Any]) -> None:
+    parsed = await _roundtrip(schema, {"tags": '["auth", "idor"]'})
+
+    assert parsed["tags"] == ["auth", "idor"]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "value",
+    [
+        "auth, idor",
+        "auth\nidor",
+        "auth",
+        "Endpoint /admin leaks user data, and session tokens never expire",
+        '"auth"',
+        "",
+    ],
+)
+async def test_free_form_strings_are_never_split_into_an_array(value: str) -> None:
+    parsed = await _roundtrip(_ARRAY, {"tags": value})
+
+    assert parsed["tags"] == value
+
+
+@pytest.mark.asyncio
+async def test_encoded_mapping_is_decoded_for_an_object_parameter() -> None:
+    parsed = await _roundtrip(_OBJECT, {"modifications": '{"method": "POST"}'})
+
+    assert parsed["modifications"] == {"method": "POST"}
+
+
+@pytest.mark.asyncio
+async def test_a_decoded_container_of_the_wrong_kind_is_not_substituted() -> None:
+    parsed = await _roundtrip(_OBJECT, {"modifications": '["POST"]'})
+
+    assert parsed["modifications"] == '["POST"]'
+
+
+@pytest.mark.asyncio
+async def test_values_matching_the_schema_are_left_alone() -> None:
+    parsed = await _roundtrip({**_ARRAY, **_OBJECT}, {"tags": ["auth"], "modifications": {"a": 1}})
+
+    assert parsed == {"tags": ["auth"], "modifications": {"a": 1}}
+
+
+@pytest.mark.asyncio
+async def test_unknown_and_null_arguments_are_untouched() -> None:
+    parsed = await _roundtrip(_NULLABLE_ARRAY, {"tags": None, "other": ["x"]})
+
+    assert parsed == {"tags": None, "other": ["x"]}
+
+
+@pytest.mark.asyncio
+async def test_non_object_payloads_pass_through_unchanged() -> None:
+    captured: dict[str, str] = {}
+    wrapped = factory._with_coerced_arguments(_capturing_tool(captured, _ARRAY))
+
+    assert await wrapped.on_invoke_tool(cast("Any", None), "not json") == "ok"
+    assert captured["raw_input"] == "not json"
+
+
+@pytest.mark.asyncio
+async def test_coercion_is_applied_once_per_tool() -> None:
+    captured: dict[str, str] = {}
+    tool = factory._with_coerced_arguments(_capturing_tool(captured, _ARRAY))
+
+    assert factory._with_coerced_arguments(tool) is tool