diff --git a/strix/agents/factory.py b/strix/agents/factory.py index 77e6c474..459b554f 100644 --- a/strix/agents/factory.py +++ b/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( diff --git a/tests/test_agent_factory_tool_arguments.py b/tests/test_agent_factory_tool_arguments.py new file mode 100644 index 00000000..49dabf5d --- /dev/null +++ b/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