fix(models): replay malformed tool-call arguments as valid JSON (#1327)

* fix(models): replay malformed tool-call arguments as valid JSON

A model that emits a tool call with truncated or otherwise invalid JSON
arguments fails that one call, but the raw string is recorded in the
session. Strict OpenAI-compatible servers (vLLM, SGLang) validate every
assistant tool call in the request, so each later turn is rejected with
"Assistant tool call function.arguments must be valid JSON" and the agent
never recovers. The turn guard now rewrites such arguments on the way out
to a JSON object carrying the original text, alongside the existing
duplicate call-id repair.

* fix(agents): tell the model to re-issue tool calls whose arguments are not valid JSON

A streamed turn that ends mid-tool-call leaves the client with a JSON prefix
as the arguments. Instead of the SDK's generic parse-error result, the tool
wrapper now short-circuits with a clear 'not executed, re-issue with complete
arguments' message so the model recovers in one turn.

* fix(models): treat NaN/Infinity tool-call arguments as malformed

* test(cost): pin model resolution in the provider-prefixed cost test

* fix(models): treat over-nested tool-call arguments as malformed

strix/agents/factory.py

@@ -18,6 +18,7 @@ from pydantic import ValidationError
 
 from strix.agents.prompt import render_system_prompt
 from strix.config import load_settings
+from strix.config.tool_call_arguments import describe_malformed_arguments
 from strix.tools.agents_graph.tools import (
     agent_finish,
     create_agent,
@@ -267,6 +268,10 @@ def _with_coerced_arguments(tool: FunctionTool) -> FunctionTool:
     nullish = tool.name.startswith(_QUERY_TOOL_PREFIXES)
 
     async def invoke(ctx: Any, raw_input: str) -> Any:
+        malformed = describe_malformed_arguments(tool.name, raw_input)
+        if malformed is not None:
+            logger.debug("Tool %s got malformed arguments; asking the model to re-issue", tool.name)
+            return malformed
         return await invoke_tool(ctx, _coerce_arguments(raw_input, schema, nullish=nullish))
 
     tool.on_invoke_tool = invoke

strix/config/models.py

@@ -41,6 +41,7 @@ from openai.types.shared import Reasoning
 from strix.agents.prompt import CACHE_POINT
 from strix.config import codex
 from strix.config.loader import load_settings
+from strix.config.tool_call_arguments import repair_input
 from strix.config.tool_call_ids import TurnCallIdRewriter, dedupe_input
 from strix.config.tool_call_limits import TurnToolCallLimiter
 from strix.llm import request_log
@@ -255,6 +256,12 @@ class _TurnGuardModel(Model):
     Ids that collide with the history are rewritten before the turn is
     recorded, and already-corrupted histories are repaired on the way out.
 
+    Tool-call arguments: a turn whose ``arguments`` are not valid JSON fails
+    that one call with a parse error, but the raw string is recorded and
+    strict providers then reject every request that replays it. Such
+    arguments are rewritten to a valid JSON object on the way out (see
+    :mod:`strix.config.tool_call_arguments`).
+
     Tool-call volume: a degenerate response can queue hundreds of calls that
     the run loop then honours one by one. Only the first
     ``LLM_MAX_TOOL_CALLS_PER_TURN`` calls of a response are kept.
@@ -313,7 +320,7 @@ class _TurnGuardModel(Model):
         conversation_id: str | None,
         prompt: ResponsePromptParam | None,
     ) -> ModelResponse:
-        sanitized = dedupe_input(input)
+        sanitized = _sanitize_input(input)
         rewriter = TurnCallIdRewriter(sanitized)
         system_instructions, sanitized = _split_cached_prefix(
             system_instructions, sanitized, model_settings
@@ -349,7 +356,7 @@ class _TurnGuardModel(Model):
         conversation_id: str | None,
         prompt: ResponsePromptParam | None,
     ) -> AsyncIterator[TResponseStreamEvent]:
-        sanitized = dedupe_input(input)
+        sanitized = _sanitize_input(input)
         rewriter = TurnCallIdRewriter(sanitized)
         system_instructions, sanitized = _split_cached_prefix(
             system_instructions, sanitized, model_settings
@@ -379,6 +386,10 @@ class _TurnGuardModel(Model):
         self._log_dropped(limiter)
 
 
+def _sanitize_input(model_input: str | list[TResponseInputItem]) -> str | list[Any]:
+    return repair_input(dedupe_input(model_input))
+
+
 def _split_cached_prefix(
     system_instructions: str | None,
     model_input: str | list[Any],

strix/config/tool_call_arguments.py

@@ -0,0 +1,100 @@
+"""Keep replayed tool-call arguments valid JSON.
+
+A model occasionally emits a tool call whose ``arguments`` string is not a
+JSON object (truncated, unbalanced, empty). The call itself fails safely: the
+tool reports the parse error and the run continues. But the raw string is
+recorded in the session as-is, and strict OpenAI-compatible servers (vLLM,
+SGLang, ...) validate every assistant tool call in the request, so from then
+on each turn is rejected with ``Assistant tool call function.arguments must
+be valid JSON`` and the agent can never recover. Rewriting the replayed
+arguments to a JSON object that carries the original text keeps the history
+valid while the model still sees what it sent.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import Any
+
+from openai.types.responses import ResponseFunctionToolCall
+
+
+MALFORMED_ARGUMENTS_KEY = "malformed_arguments"
+
+
+def _reject_constant(name: str) -> Any:
+    raise ValueError(f"non-standard JSON constant {name}")
+
+
+def _loads_strict(text: str) -> Any:
+    try:
+        return json.loads(text, parse_constant=_reject_constant)
+    except RecursionError as exc:
+        raise ValueError("arguments are nested too deeply to parse") from exc
+
+
+def describe_malformed_arguments(tool_name: str, arguments: str) -> str | None:
+    """Return a model-facing recovery message if ``arguments`` is not a JSON object.
+
+    A tool call whose arguments do not parse is almost always one the stream cut
+    off (the server flushed a prefix of the JSON, or the client closed early),
+    so instead of the SDK's generic parse-error result the model is told the
+    call never ran and must be re-issued whole.
+    """
+    if not arguments.strip():
+        return None
+    try:
+        parsed = _loads_strict(arguments)
+    except ValueError as exc:
+        detail = str(exc)
+    else:
+        if isinstance(parsed, dict):
+            return None
+        detail = f"expected a JSON object, got {type(parsed).__name__}"
+    return (
+        f"{tool_name}: the tool call was not executed because its arguments were "
+        f"truncated or otherwise not valid JSON ({detail}). The response was likely "
+        "cut off mid-call. Re-issue the call with complete, valid JSON arguments."
+    )
+
+
+def repair_arguments(arguments: object) -> str | None:
+    """Return replacement arguments a strict server accepts, or ``None`` if already valid."""
+    if not isinstance(arguments, str) or not arguments.strip():
+        return "{}"
+    try:
+        parsed = _loads_strict(arguments)
+    except ValueError:
+        parsed = None
+    if isinstance(parsed, dict):
+        return None
+    return json.dumps({MALFORMED_ARGUMENTS_KEY: arguments}, ensure_ascii=False)
+
+
+def repair_history_arguments(items: list[Any]) -> tuple[list[Any], bool]:
+    """Rewrite function calls in a conversation history whose arguments are not a JSON object."""
+    rebuilt: list[Any] = []
+    changed = False
+
+    for item in items:
+        if isinstance(item, dict):
+            if item.get("type") == "function_call":
+                repaired = repair_arguments(item.get("arguments"))
+                if repaired is not None:
+                    item = {**item, "arguments": repaired}  # noqa: PLW2901
+                    changed = True
+        elif isinstance(item, ResponseFunctionToolCall):
+            repaired = repair_arguments(item.arguments)
+            if repaired is not None:
+                item = item.model_copy(update={"arguments": repaired})  # noqa: PLW2901
+                changed = True
+        rebuilt.append(item)
+
+    return rebuilt, changed
+
+
+def repair_input(model_input: str | list[Any]) -> str | list[Any]:
+    if isinstance(model_input, str):
+        return model_input
+    rebuilt, changed = repair_history_arguments(model_input)
+    return rebuilt if changed else model_input

tests/test_agent_factory_tool_arguments.py

@@ -38,11 +38,26 @@ async def _roundtrip(
 
 
 _STRING = {"todos": {"type": "string"}}
+
+
+@pytest.mark.asyncio
+async def test_truncated_arguments_short_circuit_with_a_reissue_message() -> None:
+    captured: dict[str, str] = {}
+    wrapped = factory._with_coerced_arguments(_capturing_tool(captured, _STRING))
+
+    result = await wrapped.on_invoke_tool(cast("Any", None), '{"todos": "a, b')
+
+    assert "not executed" in result
+    assert "Re-issue the call" in result
+    assert "raw_input" not in captured
+
+
 _ARRAY = {"tags": {"type": "array", "items": {"type": "string"}}}
 _NULLABLE_ARRAY = {
     "tags": {"anyOf": [{"type": "array", "items": {"type": "string"}}, {"type": "null"}]}
 }
 _OBJECT = {"modifications": {"type": "object"}}
+OVER_NESTED = '{"a":' + "[" * 100_000
 
 
 @pytest.mark.asyncio
@@ -136,12 +151,15 @@ async def test_unknown_and_null_arguments_are_untouched() -> None:
 
 
 @pytest.mark.asyncio
-async def test_non_object_payloads_pass_through_unchanged() -> None:
+@pytest.mark.parametrize("arguments", ["not json", '{"n": NaN}', OVER_NESTED])
+async def test_non_object_payloads_are_reported_instead_of_invoked(arguments: str) -> 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"
+    result = await wrapped.on_invoke_tool(cast("Any", None), arguments)
+
+    assert result.startswith("probe: the tool call was not executed")
+    assert "raw_input" not in captured
 
 
 @pytest.mark.asyncio

tests/test_cost_tracking.py

@@ -136,6 +136,7 @@ def test_cost_callback_estimates_cost_with_provider_prefixed_model() -> None:
 
     with (
         patch("strix.report.state.get_global_report_state", return_value=report_state),
+        patch("strix.report.state.resolve_litellm_model", side_effect=lambda model: model),
         patch("litellm.completion_cost", side_effect=fake_completion_cost),
     ):
         litellm_cost_callback(kwargs, response)

tests/test_tool_call_arguments.py

@@ -0,0 +1,300 @@
+"""Tests for keeping replayed tool-call arguments valid JSON.
+
+A model that emits a tool call with malformed ``arguments`` fails that one
+call, but the raw string is recorded in the session. Strict OpenAI-compatible
+servers validate every assistant tool call in the request, so each later turn
+is rejected and the agent can never recover. A gateway that validates
+arguments the way those servers do proves both the failure and the fix.
+"""
+
+from __future__ import annotations
+
+import json
+import threading
+from http.server import BaseHTTPRequestHandler, HTTPServer
+from typing import TYPE_CHECKING, Any
+
+import pytest
+from agents import Agent, Runner, function_tool
+from agents.models.interface import Model, ModelProvider
+from agents.models.openai_chatcompletions import OpenAIChatCompletionsModel
+from agents.run import RunConfig
+from openai import AsyncOpenAI, BadRequestError
+from openai.types.responses import ResponseFunctionToolCall
+
+from strix.config.models import _NonStreamingModel, _TurnGuardModel
+from strix.config.tool_call_arguments import (
+    MALFORMED_ARGUMENTS_KEY,
+    _loads_strict,
+    describe_malformed_arguments,
+    repair_arguments,
+    repair_history_arguments,
+    repair_input,
+)
+
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+
+TRUNCATED = '{"n": 1'
+OVER_NESTED = '{"a":' + "[" * 100_000
+
+
+def _tool_call_completion(arguments: str) -> dict[str, Any]:
+    return {
+        "id": "chatcmpl-1",
+        "object": "chat.completion",
+        "created": 0,
+        "model": "gw-model",
+        "choices": [
+            {
+                "index": 0,
+                "finish_reason": "tool_calls",
+                "message": {
+                    "role": "assistant",
+                    "content": None,
+                    "tool_calls": [
+                        {
+                            "id": "call_1",
+                            "type": "function",
+                            "function": {"name": "do_thing", "arguments": arguments},
+                        }
+                    ],
+                },
+            }
+        ],
+        "usage": {"prompt_tokens": 5, "completion_tokens": 2, "total_tokens": 7},
+    }
+
+
+def _text_completion(text: str) -> dict[str, Any]:
+    return {
+        "id": "chatcmpl-2",
+        "object": "chat.completion",
+        "created": 0,
+        "model": "gw-model",
+        "choices": [
+            {"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": text}}
+        ],
+        "usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8},
+    }
+
+
+_REQUESTS: list[list[dict[str, Any]]] = []
+
+
+def _assistant_arguments(messages: list[dict[str, Any]]) -> list[str]:
+    return [
+        str(call["function"]["arguments"])
+        for message in messages
+        for call in message.get("tool_calls") or []
+    ]
+
+
+def _tool_results(messages: list[dict[str, Any]]) -> list[str]:
+    return [str(m.get("content")) for m in messages if m.get("role") == "tool"]
+
+
+class _StrictHandler(BaseHTTPRequestHandler):
+    """Gateway that rejects malformed assistant tool-call arguments, like vLLM/SGLang do."""
+
+    def log_message(self, *args: Any) -> None:
+        pass
+
+    def do_POST(self) -> None:
+        length = int(self.headers.get("Content-Length", 0))
+        body = json.loads(self.rfile.read(length) or b"{}")
+        messages = body.get("messages", [])
+        _REQUESTS.append(messages)
+
+        for arguments in _assistant_arguments(messages):
+            try:
+                json.loads(arguments)
+            except ValueError:
+                self._respond(
+                    400,
+                    {
+                        "object": "error",
+                        "message": "Assistant tool call function.arguments must be valid JSON.",
+                        "type": "BadRequest",
+                        "param": None,
+                        "code": 400,
+                    },
+                )
+                return
+
+        if len(_REQUESTS) == 1:
+            self._respond(200, _tool_call_completion(TRUNCATED))
+        else:
+            self._respond(200, _text_completion("all done"))
+
+    def _respond(self, status: int, payload: dict[str, Any]) -> None:
+        encoded = json.dumps(payload).encode()
+        self.send_response(status)
+        self.send_header("Content-Type", "application/json")
+        self.send_header("Content-Length", str(len(encoded)))
+        self.end_headers()
+        self.wfile.write(encoded)
+
+
+@pytest.fixture
+def strict_gateway() -> Iterator[str]:
+    _REQUESTS.clear()
+    server = HTTPServer(("127.0.0.1", 0), _StrictHandler)
+    thread = threading.Thread(target=server.serve_forever, daemon=True)
+    thread.start()
+    try:
+        yield f"http://127.0.0.1:{server.server_address[1]}/v1"
+    finally:
+        server.shutdown()
+        server.server_close()
+
+
+def _model(base_url: str) -> Model:
+    client = AsyncOpenAI(api_key="tok", base_url=base_url, max_retries=0)
+    return _NonStreamingModel(OpenAIChatCompletionsModel(model="gw-model", openai_client=client))
+
+
+async def _run_agent(base_url: str, *, wrap: bool) -> Any:
+    @function_tool
+    def do_thing(n: int) -> str:
+        return f"did {n}"
+
+    class _Provider(ModelProvider):
+        def get_model(self, model_name: str | None) -> Model:  # noqa: ARG002
+            model = _model(base_url)
+            return _TurnGuardModel(model) if wrap else model
+
+    agent = Agent(name="t", instructions="use the tool", tools=[do_thing], model="gw-model")
+    result = Runner.run_streamed(
+        agent, input="please", run_config=RunConfig(model_provider=_Provider())
+    )
+    async for _ in result.stream_events():
+        pass
+    return result
+
+
+@pytest.mark.asyncio
+async def test_malformed_arguments_poison_every_later_turn_without_the_wrapper(
+    strict_gateway: str,
+) -> None:
+    # Repro: the model truncates a tool call's arguments once. The tool fails
+    # that call gracefully, but the next request replays the raw string and
+    # the provider rejects the whole conversation from then on.
+    with pytest.raises(BadRequestError, match="must be valid JSON"):
+        await _run_agent(strict_gateway, wrap=False)
+
+    assert _assistant_arguments(_REQUESTS[-1]) == [TRUNCATED]
+
+
+@pytest.mark.asyncio
+async def test_malformed_arguments_are_replayed_as_valid_json(strict_gateway: str) -> None:
+    result = await _run_agent(strict_gateway, wrap=True)
+
+    assert result.final_output == "all done"
+    (arguments,) = _assistant_arguments(_REQUESTS[-1])
+    assert json.loads(arguments) == {MALFORMED_ARGUMENTS_KEY: TRUNCATED}
+    (tool_result,) = _tool_results(_REQUESTS[-1])
+    assert "JSON" in tool_result
+
+
+@pytest.mark.parametrize(
+    ("arguments", "expected"),
+    [
+        ('{"cmd": "ls', json.dumps({MALFORMED_ARGUMENTS_KEY: '{"cmd": "ls'})),
+        ("[1, 2]", json.dumps({MALFORMED_ARGUMENTS_KEY: "[1, 2]"})),
+        ("", "{}"),
+        ("  ", "{}"),
+        (None, "{}"),
+    ],
+)
+def test_repair_arguments_rewrites_anything_but_a_json_object(
+    arguments: str | None, expected: str
+) -> None:
+    assert repair_arguments(arguments) == expected
+
+
+@pytest.mark.parametrize(
+    "arguments", ["{}", '{"cmd": "ls -la"}', '{"nested": {"a": [1]}}', '{"n": "NaN"}']
+)
+def test_repair_arguments_leaves_json_objects_alone(arguments: str) -> None:
+    assert repair_arguments(arguments) is None
+
+
+@pytest.mark.parametrize("arguments", ['{"n": NaN}', '{"n": Infinity}', '{"n": -Infinity}'])
+def test_repair_arguments_wraps_non_standard_json_constants(arguments: str) -> None:
+    repaired = repair_arguments(arguments)
+
+    assert repaired is not None
+    assert _loads_strict(repaired) == {MALFORMED_ARGUMENTS_KEY: arguments}
+
+
+def test_repair_arguments_wraps_over_nested_json() -> None:
+    repaired = repair_arguments(OVER_NESTED)
+
+    assert repaired is not None
+    assert _loads_strict(repaired) == {MALFORMED_ARGUMENTS_KEY: OVER_NESTED}
+
+
+def test_history_repair_rewrites_only_malformed_calls() -> None:
+    items = [
+        {"role": "user", "content": "go"},
+        {"type": "function_call", "call_id": "a", "name": "x", "arguments": '{"cmd": "ls'},
+        {"type": "function_call_output", "call_id": "a", "output": "invalid JSON"},
+        {"type": "function_call", "call_id": "b", "name": "y", "arguments": '{"ok": true}'},
+        ResponseFunctionToolCall(call_id="c", name="z", arguments="", type="function_call"),
+    ]
+
+    rebuilt, changed = repair_history_arguments(items)
+
+    assert changed
+    assert rebuilt[0] is items[0]
+    assert json.loads(rebuilt[1]["arguments"]) == {MALFORMED_ARGUMENTS_KEY: '{"cmd": "ls'}
+    assert rebuilt[1]["call_id"] == "a"
+    assert rebuilt[2] is items[2]
+    assert rebuilt[3] is items[3]
+    assert isinstance(rebuilt[4], ResponseFunctionToolCall)
+    assert rebuilt[4].arguments == "{}"
+    assert isinstance(items[1], dict)
+    assert items[1]["arguments"] == '{"cmd": "ls'
+
+
+def test_repair_input_returns_same_object_when_nothing_changes() -> None:
+    items = [{"type": "function_call", "call_id": "a", "name": "x", "arguments": "{}"}]
+
+    assert repair_input(items) is items
+    assert repair_input("plain prompt") == "plain prompt"
+
+
+def test_repair_input_wraps_over_nested_function_call_arguments() -> None:
+    items = [
+        {
+            "type": "function_call",
+            "call_id": "a",
+            "name": "x",
+            "arguments": OVER_NESTED,
+        }
+    ]
+
+    repaired = repair_input(items)
+
+    assert isinstance(repaired, list)
+    assert _loads_strict(repaired[0]["arguments"]) == {MALFORMED_ARGUMENTS_KEY: OVER_NESTED}
+
+
+@pytest.mark.parametrize(
+    "arguments",
+    ['{"cmd": "ls -la', "[1, 2]", "null", "not json", '{"n": NaN}', OVER_NESTED],
+)
+def test_describe_malformed_arguments_tells_the_model_to_reissue(arguments: str) -> None:
+    message = describe_malformed_arguments("exec_command", arguments)
+
+    assert message is not None
+    assert message.startswith("exec_command: the tool call was not executed")
+    assert "Re-issue the call" in message
+
+
+@pytest.mark.parametrize("arguments", ["", "   ", "{}", '{"cmd": "ls"}'])
+def test_describe_malformed_arguments_accepts_objects_and_empty_input(arguments: str) -> None:
+    assert describe_malformed_arguments("exec_command", arguments) is None