Compare commits

...
Author SHA1 Message Date
Ahmed AllamandAhmed Allam 60b1d72642 fix(llm): abandon a model stream that stops producing events 2026-08-06 00:07:01 +03:00
8bd6c8e87a 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 <ahmed39652003@gmail.com>
2026-08-06 00:06:59 +03:00
8 changed files with 528 additions and 28 deletions
+2
View File
@@ -237,6 +237,8 @@ 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_stream_idle_timeout.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
+102 -20
View File
@@ -2,10 +2,13 @@
from __future__ import annotations
import asyncio
import contextlib
import inspect
import logging
import os
import time
from collections.abc import AsyncGenerator
from typing import TYPE_CHECKING, Any, cast
from agents import (
@@ -36,6 +39,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 +58,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 +242,46 @@ 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.
Stalled streams: a turn that emits a few tokens and then goes silent is
not covered by the request timeout, which resets on any byte (keepalives
included). ``LLM_STREAM_IDLE_TIMEOUT`` bounds the gap between events so the
turn fails instead of hanging, and the existing retry path replays it.
"""
def __init__(self, inner: Model) -> None:
def __init__(
self,
inner: Model,
*,
max_tool_calls_per_turn: int = 0,
stream_idle_timeout: float = 0.0,
) -> None:
self._inner = inner
self._max_tool_calls_per_turn = max_tool_calls_per_turn
self._stream_idle_timeout = stream_idle_timeout
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 +317,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 +338,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),
@@ -313,21 +351,55 @@ class _UniqueToolCallIdModel(Model):
conversation_id=conversation_id,
prompt=prompt,
)
async for event in _with_idle_timeout(stream, self._stream_idle_timeout):
guarded = _guard_event(event, rewriter, limiter)
if guarded is not None:
yield guarded
self._log_dropped(limiter)
async def _aclose(stream: AsyncIterator[TResponseStreamEvent]) -> None:
if isinstance(stream, AsyncGenerator):
with contextlib.suppress(Exception):
await stream.aclose()
async def _with_idle_timeout(
stream: AsyncIterator[TResponseStreamEvent], timeout: float
) -> AsyncIterator[TResponseStreamEvent]:
if timeout <= 0:
async for event in stream:
yield _rewrite_event_call_ids(event, rewriter)
yield event
return
iterator = stream.__aiter__()
while True:
try:
event = await asyncio.wait_for(iterator.__anext__(), timeout)
except StopAsyncIteration:
return
except TimeoutError:
await _aclose(stream)
message = f"model stream produced no event for {timeout:.0f}s"
logger.warning("%s; abandoning the turn", message)
raise TimeoutError(message) from None
yield event
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})}
)
@@ -399,19 +471,29 @@ class StrixProvider(MultiProvider):
def get_model(self, model_name: str | None) -> Model:
llm = load_settings().llm
slug = codex.subscription_model(model_name)
idle_timeout = float(llm.stream_idle_timeout)
if slug:
# 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)
# The wrapper emits its single event only once the whole request
# is done, so an idle gap is meaningless here; the request
# timeout bounds it instead.
idle_timeout = 0.0
return _TurnGuardModel(
model,
max_tool_calls_per_turn=llm.max_tool_calls_per_turn,
stream_idle_timeout=idle_timeout,
)
DEFAULT_MODEL_RETRY = ModelRetrySettings(
+6
View File
@@ -57,6 +57,12 @@ class LlmSettings(BaseSettings):
alias="LLM_DISABLE_STREAMING",
)
timeout: int = Field(default=300, alias="LLM_TIMEOUT")
stream_idle_timeout: int = Field(default=300, ge=0, alias="LLM_STREAM_IDLE_TIMEOUT")
max_tool_calls_per_turn: int = Field(
default=32,
ge=0,
alias="LLM_MAX_TOOL_CALLS_PER_TURN",
)
class DedupeSettings(BaseSettings):
+46
View File
@@ -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)]
+8 -6
View File
@@ -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)
+173
View File
@@ -0,0 +1,173 @@
"""Tests for the model-stream idle watchdog.
A turn that streams a few tokens and then goes silent is not covered by the
request timeout: the read timeout resets on every byte, keepalives included.
The watchdog bounds the gap between events so the turn fails and can be
retried instead of parking the agent forever.
"""
from __future__ import annotations
import asyncio
import json
import threading
import time
from http.server import BaseHTTPRequestHandler, HTTPServer
from typing import TYPE_CHECKING, Any
import pytest
from agents.model_settings import ModelSettings
from agents.models.interface import Model, ModelTracing
from agents.models.openai_chatcompletions import OpenAIChatCompletionsModel
from openai import AsyncOpenAI
from strix.config import loader
from strix.config.loader import load_settings
from strix.config.models import StrixProvider, _TurnGuardModel, _with_idle_timeout
if TYPE_CHECKING:
from collections.abc import AsyncIterator, Iterator
_STALL_SECONDS = 30.0
def _chunk(text: str) -> bytes:
payload = {
"id": "chatcmpl-1",
"object": "chat.completion.chunk",
"created": 0,
"model": "gw-model",
"choices": [{"index": 0, "delta": {"content": text}, "finish_reason": None}],
}
return b"data: " + json.dumps(payload).encode() + b"\n\n"
class _StallingHandler(BaseHTTPRequestHandler):
"""Streams a couple of tokens, then stops producing anything."""
stop = threading.Event()
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)
self.send_response(200)
self.send_header("Content-Type", "text/event-stream")
self.end_headers()
self.wfile.write(_chunk("Now"))
self.wfile.write(_chunk(" spawning"))
self.wfile.flush()
self.stop.wait(_STALL_SECONDS)
@pytest.fixture
def stalling_gateway() -> Iterator[str]:
_StallingHandler.stop.clear()
server = HTTPServer(("127.0.0.1", 0), _StallingHandler)
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:
_StallingHandler.stop.set()
server.shutdown()
server.server_close()
def _stream(base_url: str, *, idle_timeout: float) -> AsyncIterator[Any]:
client = AsyncOpenAI(api_key="tok", base_url=base_url, max_retries=0, timeout=_STALL_SECONDS)
inner: Model = OpenAIChatCompletionsModel(model="gw-model", openai_client=client)
guarded = _TurnGuardModel(inner, stream_idle_timeout=idle_timeout)
return guarded.stream_response(
None,
"go",
ModelSettings(),
[],
None,
[],
ModelTracing.DISABLED,
previous_response_id=None,
conversation_id=None,
prompt=None,
)
async def _drain(base_url: str, *, idle_timeout: float) -> list[Any]:
return [event async for event in _stream(base_url, idle_timeout=idle_timeout)]
@pytest.mark.asyncio
async def test_stalled_stream_hangs_without_the_watchdog(stalling_gateway: str) -> None:
# Repro: tokens arrive, then nothing. Un-watched, the turn just sits there;
# the request timeout is far away and would reset on any keepalive byte.
with pytest.raises(TimeoutError):
await asyncio.wait_for(_drain(stalling_gateway, idle_timeout=0), timeout=2)
@pytest.mark.asyncio
async def test_stalled_stream_is_abandoned_by_the_watchdog(stalling_gateway: str) -> None:
started = time.monotonic()
with pytest.raises(TimeoutError, match="produced no event"):
await _drain(stalling_gateway, idle_timeout=1)
assert time.monotonic() - started < _STALL_SECONDS
@pytest.mark.asyncio
async def test_events_keep_flowing_while_the_stream_is_alive() -> None:
async def _live() -> AsyncIterator[Any]:
for i in range(5):
await asyncio.sleep(0.05)
yield f"event-{i}"
seen: list[Any] = [event async for event in _with_idle_timeout(_live(), 1.0)]
assert seen == [f"event-{i}" for i in range(5)]
@pytest.fixture
def _reset_settings(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]:
for key in ("STRIX_LLM", "LLM_DISABLE_STREAMING", "LLM_STREAM_IDLE_TIMEOUT"):
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_idle_timeout_is_configurable(
monkeypatch: pytest.MonkeyPatch, _reset_settings: None
) -> None:
monkeypatch.setattr("strix.config.models.MultiProvider.get_model", lambda *_: _DummyModel())
monkeypatch.setenv("LLM_STREAM_IDLE_TIMEOUT", "45")
load_settings()
model = StrixProvider().get_model("openai/gpt-4o-mini")
assert isinstance(model, _TurnGuardModel)
assert model._stream_idle_timeout == 45
def test_idle_timeout_is_off_without_streaming(
monkeypatch: pytest.MonkeyPatch, _reset_settings: None
) -> None:
# LLM_DISABLE_STREAMING turns the whole request into one event, so an idle
# gap would just be the request duration — the request timeout bounds that.
monkeypatch.setattr("strix.config.models.MultiProvider.get_model", lambda *_: _DummyModel())
monkeypatch.setenv("LLM_STREAM_IDLE_TIMEOUT", "45")
monkeypatch.setenv("LLM_DISABLE_STREAMING", "true")
load_settings()
model = StrixProvider().get_model("openai/gpt-4o-mini")
assert isinstance(model, _TurnGuardModel)
assert model._stream_idle_timeout == 0
+2 -2
View File
@@ -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(
+189
View File
@@ -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