From 8bd6c8e87a35fa9d211587756ffcb8c32626dc9d Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 6 Aug 2026 00:06:59 +0300 Subject: [PATCH] fix(llm): cap the tool calls one assistant response may queue (#977) * fix(llm): cap the tool calls one assistant response may queue * fix(llm): cap the subscription backend's responses too --------- Co-authored-by: Ahmed Allam --- pyproject.toml | 1 + strix/config/models.py | 71 ++++++++---- strix/config/settings.py | 5 + strix/config/tool_call_limits.py | 46 ++++++++ tests/test_disable_streaming.py | 14 ++- tests/test_tool_call_ids.py | 4 +- tests/test_tool_call_limits.py | 189 +++++++++++++++++++++++++++++++ 7 files changed, 302 insertions(+), 28 deletions(-) create mode 100644 strix/config/tool_call_limits.py create mode 100644 tests/test_tool_call_limits.py diff --git a/pyproject.toml b/pyproject.toml index d264dcca..a35ea357 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -237,6 +237,7 @@ ignore = [ "tests/test_codex_streaming.py" = ["N802"] "tests/test_disable_streaming.py" = ["N802"] "tests/test_tool_call_ids.py" = ["N802"] +"tests/test_tool_call_limits.py" = ["N802", "SLF001"] "tests/test_unknown_tool_recovery.py" = ["N802"] "tests/test_report_pdf.py" = ["S105", "S106"] # Stdlib HTTP handler overrides (do_GET/do_POST) and lazy imports that avoid a diff --git a/strix/config/models.py b/strix/config/models.py index 8685a2fb..f14abfcf 100644 --- a/strix/config/models.py +++ b/strix/config/models.py @@ -4,6 +4,7 @@ from __future__ import annotations import contextlib import inspect +import logging import os import time from typing import TYPE_CHECKING, Any, cast @@ -36,6 +37,7 @@ from openai.types.shared import Reasoning from strix.config import codex from strix.config.loader import load_settings from strix.config.tool_call_ids import TurnCallIdRewriter, dedupe_input +from strix.config.tool_call_limits import TurnToolCallLimiter if TYPE_CHECKING: @@ -54,6 +56,9 @@ if TYPE_CHECKING: from strix.config.settings import LlmSettings, ReasoningEffort, Settings +logger = logging.getLogger(__name__) + + def request_timeout_extra_args(timeout_s: float | None) -> dict[str, float] | None: """Per-request model timeout; a plain float so ``ModelSettings.to_json_dict()`` stays serializable.""" # noqa: E501 if not timeout_s or timeout_s <= 0: @@ -235,18 +240,34 @@ class _NonStreamingModel(Model): yield _completed_stream_event(response, getattr(self._inner, "model", None)) -class _UniqueToolCallIdModel(Model): - """Keep tool-call ids unique so a recycled id can't invalidate the history. +class _TurnGuardModel(Model): + """Keep one turn from corrupting the conversation or running away. - Providers that number tool calls per turn (``exec_command:0``, ...) restart - the counter each turn, so the same id eventually appears twice in one - conversation and strict providers reject every subsequent request. 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 ids: providers that number calls per turn (``exec_command:0``, + ...) restart the counter each turn, so the same id eventually appears twice + in one conversation and strict providers reject every subsequent request. + 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 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. """ - def __init__(self, inner: Model) -> None: + def __init__(self, inner: Model, *, max_tool_calls_per_turn: int = 0) -> None: self._inner = inner + self._max_tool_calls_per_turn = max_tool_calls_per_turn + + def _limiter(self) -> TurnToolCallLimiter: + return TurnToolCallLimiter(self._max_tool_calls_per_turn) + + def _log_dropped(self, limiter: TurnToolCallLimiter) -> None: + if limiter.dropped: + logger.warning( + "dropped %d tool call(s) past the per-response limit of %d", + limiter.dropped, + self._max_tool_calls_per_turn, + ) async def close(self) -> None: await self._inner.close() @@ -282,7 +303,9 @@ class _UniqueToolCallIdModel(Model): conversation_id=conversation_id, prompt=prompt, ) - response.output = rewriter.rewrite_items(list(response.output)) + limiter = self._limiter() + response.output = limiter.filter_items(rewriter.rewrite_items(list(response.output))) + self._log_dropped(limiter) return response async def stream_response( @@ -301,6 +324,7 @@ class _UniqueToolCallIdModel(Model): ) -> AsyncIterator[TResponseStreamEvent]: sanitized = dedupe_input(input) rewriter = TurnCallIdRewriter(sanitized) + limiter = self._limiter() stream = self._inner.stream_response( system_instructions, cast("str | list[TResponseInputItem]", sanitized), @@ -314,20 +338,26 @@ class _UniqueToolCallIdModel(Model): prompt=prompt, ) async for event in stream: - yield _rewrite_event_call_ids(event, rewriter) + guarded = _guard_event(event, rewriter, limiter) + if guarded is not None: + yield guarded + self._log_dropped(limiter) -def _rewrite_event_call_ids( - event: TResponseStreamEvent, rewriter: TurnCallIdRewriter -) -> TResponseStreamEvent: +def _guard_event( + event: TResponseStreamEvent, rewriter: TurnCallIdRewriter, limiter: TurnToolCallLimiter +) -> TResponseStreamEvent | None: if isinstance(event, ResponseOutputItemAddedEvent | ResponseOutputItemDoneEvent): rewritten = rewriter.rewrite_item(event.item) + if not limiter.allow(rewritten): + return None if rewritten is not event.item: return event.model_copy(update={"item": rewritten}) return event if isinstance(event, ResponseCompletedEvent): - output = rewriter.rewrite_items(list(event.response.output)) - if output != list(event.response.output): + original = list(event.response.output) + output = limiter.filter_items(rewriter.rewrite_items(original)) + if output != original: return event.model_copy( update={"response": event.response.model_copy(update={"output": output})} ) @@ -403,15 +433,16 @@ class StrixProvider(MultiProvider): # The ChatGPT subscription backend is always streamed; it has no # non-streaming mode to fall back to, so LLM_DISABLE_STREAMING # does not apply here. - return _CodexResponsesModel( + model: Model = _CodexResponsesModel( slug, codex.get_subscription_client(), reasoning_effort=llm.reasoning_effort, ) - model = super().get_model(model_name) - if llm.disable_streaming: - model = _NonStreamingModel(model) - return _UniqueToolCallIdModel(model) + else: + model = super().get_model(model_name) + if llm.disable_streaming: + model = _NonStreamingModel(model) + return _TurnGuardModel(model, max_tool_calls_per_turn=llm.max_tool_calls_per_turn) DEFAULT_MODEL_RETRY = ModelRetrySettings( diff --git a/strix/config/settings.py b/strix/config/settings.py index eda4ebce..a4a18b78 100644 --- a/strix/config/settings.py +++ b/strix/config/settings.py @@ -57,6 +57,11 @@ class LlmSettings(BaseSettings): alias="LLM_DISABLE_STREAMING", ) timeout: int = Field(default=300, alias="LLM_TIMEOUT") + max_tool_calls_per_turn: int = Field( + default=32, + ge=0, + alias="LLM_MAX_TOOL_CALLS_PER_TURN", + ) class DedupeSettings(BaseSettings): diff --git a/strix/config/tool_call_limits.py b/strix/config/tool_call_limits.py new file mode 100644 index 00000000..a014e8db --- /dev/null +++ b/strix/config/tool_call_limits.py @@ -0,0 +1,46 @@ +"""Bound how many tool calls one assistant response may queue. + +A degenerate generation can emit hundreds or thousands of tool calls in a +single response — typically a poll/wait loop the model writes out ahead of +time instead of issuing one call and yielding. The run loop honours all of +them, so the agent stops reacting to anything for hours. Keeping only the +first ``limit`` calls of a response bounds that blast radius; the model sees +their results on the next turn and can reconsider. +""" + +from __future__ import annotations + +from typing import Any + +from openai.types.responses import ResponseFunctionToolCall + + +class TurnToolCallLimiter: + """Decide, once per call, whether a turn's tool call is within the limit.""" + + def __init__(self, limit: int) -> None: + self._limit = limit + self._decisions: dict[str, bool] = {} + self._kept = 0 + self.dropped = 0 + + @property + def enabled(self) -> bool: + return self._limit > 0 + + def allow(self, item: Any) -> bool: + if not self.enabled or not isinstance(item, ResponseFunctionToolCall): + return True + decided = self._decisions.get(item.call_id) + if decided is not None: + return decided + allowed = self._kept < self._limit + if allowed: + self._kept += 1 + else: + self.dropped += 1 + self._decisions[item.call_id] = allowed + return allowed + + def filter_items(self, items: list[Any]) -> list[Any]: + return [item for item in items if self.allow(item)] diff --git a/tests/test_disable_streaming.py b/tests/test_disable_streaming.py index d8438a54..00e99bcc 100644 --- a/tests/test_disable_streaming.py +++ b/tests/test_disable_streaming.py @@ -31,7 +31,7 @@ from openai.types.responses import ( from strix.config import codex, loader from strix.config.loader import load_settings -from strix.config.models import StrixProvider, _NonStreamingModel, _UniqueToolCallIdModel +from strix.config.models import StrixProvider, _NonStreamingModel, _TurnGuardModel if TYPE_CHECKING: @@ -299,7 +299,7 @@ def test_get_model_wraps_when_disabled( load_settings() model = StrixProvider().get_model("openai/gpt-4o-mini") - assert isinstance(model, _UniqueToolCallIdModel) + assert isinstance(model, _TurnGuardModel) assert isinstance(model._inner, _NonStreamingModel) @@ -311,18 +311,20 @@ def test_get_model_keeps_streaming_by_default( load_settings() model = StrixProvider().get_model("openai/gpt-4o-mini") - assert isinstance(model, _UniqueToolCallIdModel) + assert isinstance(model, _TurnGuardModel) assert model._inner is inner -def test_get_model_does_not_wrap_subscription_model( +def test_get_model_guards_subscription_model_but_keeps_it_streaming( monkeypatch: pytest.MonkeyPatch, _reset_settings: None ) -> None: - # Subscription (ChatGPT) models are always streamed and must not be wrapped. + # Subscription (ChatGPT) models are always streamed, so LLM_DISABLE_STREAMING + # must not apply — but a runaway response needs capping there too. monkeypatch.setattr(codex, "subscription_model", lambda *_: "gpt-5.5") monkeypatch.setattr(codex, "get_subscription_client", lambda: AsyncOpenAI(api_key="x")) monkeypatch.setenv("LLM_DISABLE_STREAMING", "true") load_settings() model = StrixProvider().get_model("gpt-5.5") - assert not isinstance(model, _NonStreamingModel) + assert isinstance(model, _TurnGuardModel) + assert not isinstance(model._inner, _NonStreamingModel) diff --git a/tests/test_tool_call_ids.py b/tests/test_tool_call_ids.py index a4033f4b..bfb03598 100644 --- a/tests/test_tool_call_ids.py +++ b/tests/test_tool_call_ids.py @@ -23,7 +23,7 @@ from agents.run import RunConfig from openai import AsyncOpenAI from openai.types.responses import ResponseFunctionToolCall -from strix.config.models import _NonStreamingModel, _UniqueToolCallIdModel +from strix.config.models import _NonStreamingModel, _TurnGuardModel from strix.config.tool_call_ids import TurnCallIdRewriter, dedupe_history_call_ids @@ -153,7 +153,7 @@ async def _run_agent(base_url: str, *, wrap: bool) -> Any: class _Provider(ModelProvider): def get_model(self, model_name: str | None) -> Model: # noqa: ARG002 model = _model(base_url) - return _UniqueToolCallIdModel(model) if wrap else model + 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( diff --git a/tests/test_tool_call_limits.py b/tests/test_tool_call_limits.py new file mode 100644 index 00000000..2c34a213 --- /dev/null +++ b/tests/test_tool_call_limits.py @@ -0,0 +1,189 @@ +"""Tests for the per-response tool-call cap. + +A degenerate generation can emit hundreds of tool calls in one assistant +response — a wait/poll loop the model writes out ahead of time. The run loop +honours every one of them, so the agent stops reacting for hours. The cap +keeps the first N calls of a response and drops the tail. +""" + +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 + +from strix.config import loader +from strix.config.loader import load_settings +from strix.config.models import StrixProvider, _NonStreamingModel, _TurnGuardModel + + +if TYPE_CHECKING: + from collections.abc import Iterator + + +_RUNAWAY_CALLS = 200 +_CAP = 32 + + +def _runaway_completion() -> 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": f"call_{i}", + "type": "function", + "function": {"name": "wait_for_message", "arguments": "{}"}, + } + for i in range(_RUNAWAY_CALLS) + ], + }, + } + ], + "usage": {"prompt_tokens": 5, "completion_tokens": 2, "total_tokens": 7}, + } + + +def _text_completion() -> 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": "done"}, + } + ], + "usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8}, + } + + +_TURNS: list[int] = [] + + +class _RunawayHandler(BaseHTTPRequestHandler): + """First turn queues a huge poll loop; the next turn ends the run.""" + + def log_message(self, *args: Any) -> None: + pass + + def do_POST(self) -> None: + length = int(self.headers.get("Content-Length", 0)) + self.rfile.read(length) + _TURNS.append(1) + payload = _runaway_completion() if len(_TURNS) == 1 else _text_completion() + encoded = json.dumps(payload).encode() + self.send_response(200) + 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 runaway_gateway() -> Iterator[str]: + _TURNS.clear() + server = HTTPServer(("127.0.0.1", 0), _RunawayHandler) + 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, *, cap: int) -> list[int]: + executed: list[int] = [] + + @function_tool + def wait_for_message() -> str: + executed.append(1) + return "nothing new" + + class _Provider(ModelProvider): + def get_model(self, model_name: str | None) -> Model: # noqa: ARG002 + return _TurnGuardModel(_model(base_url), max_tool_calls_per_turn=cap) + + agent = Agent(name="t", instructions="orchestrate", tools=[wait_for_message], model="gw-model") + result = Runner.run_streamed( + agent, input="go", run_config=RunConfig(model_provider=_Provider()) + ) + async for _ in result.stream_events(): + pass + assert result.final_output == "done" + return executed + + +@pytest.mark.asyncio +async def test_runaway_response_runs_every_queued_call_when_uncapped(runaway_gateway: str) -> None: + # Repro: one response queues 200 calls and the run loop honours all of them. + executed = await _run_agent(runaway_gateway, cap=0) + + assert len(executed) == _RUNAWAY_CALLS + + +@pytest.mark.asyncio +async def test_runaway_response_is_capped(runaway_gateway: str) -> None: + executed = await _run_agent(runaway_gateway, cap=_CAP) + + assert len(executed) == _CAP + + +@pytest.mark.asyncio +async def test_response_below_the_cap_is_untouched(runaway_gateway: str) -> None: + executed = await _run_agent(runaway_gateway, cap=_RUNAWAY_CALLS + 1) + + assert len(executed) == _RUNAWAY_CALLS + + +@pytest.fixture +def _reset_settings(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + for key in ("STRIX_LLM", "LLM_DISABLE_STREAMING", "LLM_MAX_TOOL_CALLS_PER_TURN"): + monkeypatch.delenv(key, raising=False) + monkeypatch.setattr(loader, "_cached", None) + monkeypatch.setattr(loader, "_override", None) + yield + + +class _DummyModel(Model): + async def get_response(self, *args: Any, **kwargs: Any) -> Any: + raise NotImplementedError + + def stream_response(self, *args: Any, **kwargs: Any) -> Any: + raise NotImplementedError + + +def test_cap_is_configurable(monkeypatch: pytest.MonkeyPatch, _reset_settings: None) -> None: + monkeypatch.setattr("strix.config.models.MultiProvider.get_model", lambda *_: _DummyModel()) + monkeypatch.setenv("LLM_MAX_TOOL_CALLS_PER_TURN", "7") + load_settings() + + model = StrixProvider().get_model("openai/gpt-4o-mini") + assert isinstance(model, _TurnGuardModel) + assert model._max_tool_calls_per_turn == 7