diff --git a/strix/agents/factory.py b/strix/agents/factory.py index b2fcbf08..613c1f1a 100644 --- a/strix/agents/factory.py +++ b/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, @@ -260,6 +261,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 diff --git a/strix/config/models.py b/strix/config/models.py index d8444ee1..427fc0c9 100644 --- a/strix/config/models.py +++ b/strix/config/models.py @@ -38,6 +38,7 @@ from openai.types.shared import Reasoning 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 @@ -251,6 +252,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. @@ -303,7 +310,7 @@ class _TurnGuardModel(Model): conversation_id: str | None, prompt: ResponsePromptParam | None, ) -> ModelResponse: - sanitized = dedupe_input(input) + sanitized = _sanitize_input(input) rewriter = TurnCallIdRewriter(sanitized) response = await self._inner.get_response( system_instructions, @@ -336,7 +343,7 @@ class _TurnGuardModel(Model): conversation_id: str | None, prompt: ResponsePromptParam | None, ) -> AsyncIterator[TResponseStreamEvent]: - sanitized = dedupe_input(input) + sanitized = _sanitize_input(input) rewriter = TurnCallIdRewriter(sanitized) limiter = self._limiter() stream = self._inner.stream_response( @@ -358,6 +365,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)) + + async def _aclose(stream: AsyncIterator[TResponseStreamEvent]) -> None: if isinstance(stream, AsyncGenerator): with contextlib.suppress(Exception): diff --git a/strix/config/tool_call_arguments.py b/strix/config/tool_call_arguments.py new file mode 100644 index 00000000..9169dcdc --- /dev/null +++ b/strix/config/tool_call_arguments.py @@ -0,0 +1,89 @@ +"""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 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 = json.loads(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 = json.loads(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 diff --git a/tests/test_agent_factory_tool_arguments.py b/tests/test_agent_factory_tool_arguments.py index d7503944..4330ca4a 100644 --- a/tests/test_agent_factory_tool_arguments.py +++ b/tests/test_agent_factory_tool_arguments.py @@ -38,6 +38,20 @@ 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"}]} @@ -136,12 +150,14 @@ async def test_unknown_and_null_arguments_are_untouched() -> None: @pytest.mark.asyncio -async def test_non_object_payloads_pass_through_unchanged() -> None: +async def test_non_object_payloads_are_reported_instead_of_invoked() -> 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), "not json") + + assert result.startswith("probe: the tool call was not executed") + assert "raw_input" not in captured @pytest.mark.asyncio diff --git a/tests/test_tool_call_arguments.py b/tests/test_tool_call_arguments.py new file mode 100644 index 00000000..55631a78 --- /dev/null +++ b/tests/test_tool_call_arguments.py @@ -0,0 +1,261 @@ +"""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, + describe_malformed_arguments, + repair_arguments, + repair_history_arguments, + repair_input, +) + + +if TYPE_CHECKING: + from collections.abc import Iterator + + +TRUNCATED = '{"n": 1' + + +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]}}']) +def test_repair_arguments_leaves_json_objects_alone(arguments: str) -> None: + assert repair_arguments(arguments) is None + + +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 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" + + +@pytest.mark.parametrize("arguments", ['{"cmd": "ls -la', "[1, 2]", "null", "not json"]) +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