From 78cc9a54c8500177589238d9d906018a7561c6db Mon Sep 17 00:00:00 2001 From: KT Date: Thu, 23 Jul 2026 18:01:56 +0800 Subject: [PATCH] fix: add configurable llm_call_timeout to bound stalled LLM calls Internal LLM calls (main loop and sub-agents) had only an httpx per-read timeout, or none on the streaming path. A backend that trickles bytes without finishing resets that per-read timer forever, so tasks hang indefinitely with no error. Add a wall-clock cap enforced with asyncio.wait_for on every provider call, configurable via agents.defaults.llm_call_timeout (default 600s) and carried on GenerationSettings: - litellm, custom, azure and codex providers all honor the timeout; the streaming paths use a per-chunk idle cap so long but progressing generations are not cut off. - classify_error now treats the builtin TimeoutError raised by wait_for as a retryable network error, so the retry and model-fallback ladder engages instead of failing as unknown. - the streaming loop converts a timeout into a structured error response and closes the stream, instead of crashing the turn. - PerModelProvider propagates generation settings to its per-model sub-providers, so knn-routed calls honor the configured timeout, temperature and max_tokens. Co-authored-by: Claude (claude-opus-4-8) --- raven/agent/loop/main.py | 53 +++++---- raven/cli/_helpers.py | 2 + raven/config/schema.py | 4 + raven/providers/azure_openai_provider.py | 8 +- raven/providers/base.py | 17 ++- raven/providers/custom_provider.py | 11 +- raven/providers/litellm_provider.py | 40 ++++--- raven/providers/openai_codex_provider.py | 31 ++++-- raven/providers/per_model_provider.py | 6 ++ tests/test_agent_loop_stream.py | 30 +++++- tests/test_azure_openai_provider.py | 56 ++++++++++ tests/test_cli_gateway_commands.py | 3 + tests/test_litellm_provider_timeout.py | 130 +++++++++++++++++++++++ tests/test_openai_codex_provider.py | 29 +++++ tests/test_per_model_provider.py | 14 +++ tests/test_provider_fallback_chain.py | 40 +++++++ 16 files changed, 422 insertions(+), 52 deletions(-) create mode 100644 tests/test_azure_openai_provider.py create mode 100644 tests/test_litellm_provider_timeout.py diff --git a/raven/agent/loop/main.py b/raven/agent/loop/main.py index eb2042d65..7776f297b 100644 --- a/raven/agent/loop/main.py +++ b/raven/agent/loop/main.py @@ -6,7 +6,7 @@ import json import re import time -from contextlib import AsyncExitStack +from contextlib import AsyncExitStack, aclosing from dataclasses import dataclass, field from datetime import datetime from pathlib import Path @@ -1295,27 +1295,36 @@ async def _llm_call_stream( tool_call_slots: list[dict[str, Any]] = [] final_usage: dict[str, Any] | None = None - async for delta in self.provider.chat_stream( - messages=messages, - tools=tools, - model=model, - ): - reasoning_delta = getattr(delta, "reasoning_content", None) - if reasoning_delta: - reasoning_buf.append(reasoning_delta) - if on_reasoning_delta is not None: - await on_reasoning_delta(reasoning_delta) - if delta.content: - content_buf.append(delta.content) - if on_token_delta is not None: - await on_token_delta(delta.content) - if delta.tool_call_delta: - _merge_tool_call_fragments( - tool_call_slots, - delta.tool_call_delta, - ) - if delta.usage is not None: - final_usage = delta.usage + # aclosing() guarantees the async generator (and its underlying stream) + # is closed when a TimeoutError from the per-chunk idle cap unwinds the + # loop, so a stalled stream terminates with a structured error instead + # of hanging or leaking the connection. The stream path has no retry + # (see docstring), so the error surfaces as the turn's response. + try: + async with aclosing(self.provider.chat_stream(messages=messages, tools=tools, model=model)) as stream: + async for delta in stream: + reasoning_delta = getattr(delta, "reasoning_content", None) + if reasoning_delta: + reasoning_buf.append(reasoning_delta) + if on_reasoning_delta is not None: + await on_reasoning_delta(reasoning_delta) + if delta.content: + content_buf.append(delta.content) + if on_token_delta is not None: + await on_token_delta(delta.content) + if delta.tool_call_delta: + _merge_tool_call_fragments( + tool_call_slots, + delta.tool_call_delta, + ) + if delta.usage is not None: + final_usage = delta.usage + except TimeoutError: + return LLMResponse( + content="".join(content_buf), + finish_reason="error", + error_classification=self.provider.classify_error(TimeoutError()), + ) tool_calls = _finalize_tool_calls(tool_call_slots) finish_reason = "tool_calls" if tool_calls else "stop" diff --git a/raven/cli/_helpers.py b/raven/cli/_helpers.py index 8c0b4e0c5..55b0110d5 100644 --- a/raven/cli/_helpers.py +++ b/raven/cli/_helpers.py @@ -135,6 +135,7 @@ def make_provider(config: Config): temperature=defaults.temperature, max_tokens=defaults.max_tokens, reasoning_effort=defaults.reasoning_effort, + timeout=defaults.llm_call_timeout, ) return provider @@ -155,6 +156,7 @@ def make_lazy_provider(config: Config): temperature=defaults.temperature, max_tokens=defaults.max_tokens, reasoning_effort=defaults.reasoning_effort, + timeout=defaults.llm_call_timeout, ), ) provider.prewarm() diff --git a/raven/config/schema.py b/raven/config/schema.py index c6dbd9645..6da93ceae 100644 --- a/raven/config/schema.py +++ b/raven/config/schema.py @@ -251,6 +251,10 @@ class AgentDefaults(Base): max_tokens: int = 8192 context_window_tokens: int = 65_536 temperature: float = 0.1 + # Per-call wall-clock cap (seconds) for every LLM request (main loop and + # sub-agents). Bounds a stalled backend that trickles bytes without ever + # finishing, which an httpx per-read timeout never catches. + llm_call_timeout: int = 600 max_tool_iterations: int = 40 # Cap on subagent VMs running at once (excess spawns queue). ge=1: a # 0/negative cap would deadlock every subagent (Semaphore(0)). diff --git a/raven/providers/azure_openai_provider.py b/raven/providers/azure_openai_provider.py index c1d8739d7..9a2e6e662 100644 --- a/raven/providers/azure_openai_provider.py +++ b/raven/providers/azure_openai_provider.py @@ -2,6 +2,7 @@ from __future__ import annotations +import asyncio import uuid from typing import Any from urllib.parse import urljoin @@ -146,8 +147,10 @@ async def chat( ) try: - async with httpx.AsyncClient(timeout=60.0, verify=True) as client: - response = await client.post(url, headers=headers, json=payload) + async with httpx.AsyncClient(timeout=self.generation.timeout, verify=True) as client: + response = await asyncio.wait_for( + client.post(url, headers=headers, json=payload), self.generation.timeout + ) if response.status_code != 200: return LLMResponse( content=f"Azure OpenAI API Error {response.status_code}: {response.text}", @@ -161,6 +164,7 @@ async def chat( return LLMResponse( content=f"Error calling Azure OpenAI: {repr(e)}", finish_reason="error", + error_classification=self.classify_error(e), ) def _parse_response(self, response: dict[str, Any]) -> LLMResponse: diff --git a/raven/providers/base.py b/raven/providers/base.py index 55c01fdb3..d8de84103 100644 --- a/raven/providers/base.py +++ b/raven/providers/base.py @@ -107,6 +107,7 @@ class GenerationSettings: temperature: float = 0.7 max_tokens: int = 4096 reasoning_effort: str | None = None + timeout: float = 600.0 class LLMProvider(ABC): @@ -355,11 +356,17 @@ def has(*needles: str) -> bool: ): return ErrorClassification("server", retryable=True, should_fallback=True) - # Timeout / connection → retry + fallback. - if {"timeout", "apitimeouterror", "apiconnectionerror"} & names or has( - "timeout", - "timed out", - "connection", + # Timeout / connection → retry + fallback. isinstance covers the builtin + # TimeoutError raised by asyncio.wait_for (its class name "timeouterror" + # and empty str() match neither the name set nor the substrings below). + if ( + isinstance(exc, TimeoutError) + or {"timeout", "apitimeouterror", "apiconnectionerror"} & names + or has( + "timeout", + "timed out", + "connection", + ) ): return ErrorClassification("network", retryable=True, should_fallback=True) diff --git a/raven/providers/custom_provider.py b/raven/providers/custom_provider.py index 5cdf03192..796854c58 100644 --- a/raven/providers/custom_provider.py +++ b/raven/providers/custom_provider.py @@ -2,6 +2,7 @@ from __future__ import annotations +import asyncio import uuid from typing import Any @@ -45,9 +46,15 @@ async def chat( if tools: kwargs.update(tools=tools, tool_choice=tool_choice or "auto") try: - return self._parse(await self._client.chat.completions.create(**kwargs)) + return self._parse( + await asyncio.wait_for(self._client.chat.completions.create(**kwargs), self.generation.timeout) + ) except Exception as e: - return LLMResponse(content=f"Error: {e}", finish_reason="error") + return LLMResponse( + content=f"Error: {e}", + finish_reason="error", + error_classification=self.classify_error(e), + ) def _parse(self, response: Any) -> LLMResponse: choice = response.choices[0] diff --git a/raven/providers/litellm_provider.py b/raven/providers/litellm_provider.py index 497123a05..498c540cb 100644 --- a/raven/providers/litellm_provider.py +++ b/raven/providers/litellm_provider.py @@ -1,5 +1,6 @@ """LiteLLM provider implementation for multi-provider support.""" +import asyncio import hashlib import os import secrets @@ -291,12 +292,11 @@ async def chat( "messages": self._sanitize_messages(self._sanitize_empty_content(messages), extra_keys=extra_msg_keys), "max_tokens": max_tokens, "temperature": temperature, - # Hard cap on a single LLM call to prevent the agent from hanging - # when an upstream endpoint half-closes or never sends FIN. LiteLLM's - # default is 6000s — too long for an interactive agent loop. 600s - # covers slow reasoning-heavy generations while still failing fast - # on stuck connections, so chat_with_retry can retry or give up. - "timeout": 600, + # Per-read httpx cap forwarded to the underlying client. This alone + # cannot bound a backend that trickles bytes forever (the read timer + # resets on every chunk), so the awaited call is also wrapped in an + # asyncio.wait_for wall-clock cap below. + "timeout": self.generation.timeout, } # Apply model-specific overrides (e.g. kimi-k2.5 temperature) @@ -327,7 +327,7 @@ async def chat( kwargs["tool_choice"] = tool_choice or "auto" try: - response = await acompletion(**kwargs) + response = await asyncio.wait_for(acompletion(**kwargs), self.generation.timeout) return self._parse_response(response) except Exception as e: # Return error as content for graceful handling, but classify the @@ -379,6 +379,7 @@ async def chat_stream( # when usage is explicitly requested; without it the stream carries # no token counts and downstream cost / context tracking sees zero. "stream_options": {"include_usage": True}, + "timeout": self.generation.timeout, } self._apply_model_overrides(model, kwargs) @@ -398,11 +399,26 @@ async def chat_stream( kwargs["tools"] = tools kwargs["tool_choice"] = tool_choice or "auto" - response = await acompletion(**kwargs) - async for chunk in response: - delta = self._normalize_stream_chunk(chunk) - if delta is not None: - yield delta + response = await asyncio.wait_for(acompletion(**kwargs), self.generation.timeout) + # Per-chunk idle cap: the timer resets on every chunk, so a long but + # steadily-progressing generation is fine while a mid-stream stall (no + # bytes for `timeout` seconds) raises TimeoutError instead of hanging. + # aclose() in finally closes the underlying HTTP stream deterministically + # on that timeout, mirroring the `async with` cleanup on the other paths. + stream = response.__aiter__() + try: + while True: + try: + chunk = await asyncio.wait_for(stream.__anext__(), self.generation.timeout) + except StopAsyncIteration: + break + delta = self._normalize_stream_chunk(chunk) + if delta is not None: + yield delta + finally: + aclose = getattr(stream, "aclose", None) + if aclose is not None: + await aclose() def _normalize_stream_chunk(self, chunk: Any) -> StreamDelta | None: """Normalize a raw provider chunk into a StreamDelta. diff --git a/raven/providers/openai_codex_provider.py b/raven/providers/openai_codex_provider.py index 649f1c30d..e232616ae 100644 --- a/raven/providers/openai_codex_provider.py +++ b/raven/providers/openai_codex_provider.py @@ -68,14 +68,19 @@ async def chat( url = DEFAULT_CODEX_URL + timeout = self.generation.timeout try: try: - content, tool_calls, finish_reason = await _request_codex(url, headers, body, verify=True) + content, tool_calls, finish_reason = await _request_codex( + url, headers, body, verify=True, timeout=timeout + ) except Exception as e: if "CERTIFICATE_VERIFY_FAILED" not in str(e): raise logger.warning("SSL certificate verification failed for Codex API; retrying with verify=False") - content, tool_calls, finish_reason = await _request_codex(url, headers, body, verify=False) + content, tool_calls, finish_reason = await _request_codex( + url, headers, body, verify=False, timeout=timeout + ) return LLMResponse( content=content, tool_calls=tool_calls, @@ -85,6 +90,7 @@ async def chat( return LLMResponse( content=f"Error calling Codex: {str(e)}", finish_reason="error", + error_classification=self.classify_error(e), ) def get_default_model(self) -> str: @@ -114,13 +120,14 @@ async def _request_codex( headers: dict[str, str], body: dict[str, Any], verify: bool, + timeout: float, ) -> tuple[str, list[ToolCallRequest], str]: - async with httpx.AsyncClient(timeout=60.0, verify=verify) as client: + async with httpx.AsyncClient(timeout=timeout, verify=verify) as client: async with client.stream("POST", url, headers=headers, json=body) as response: if response.status_code != 200: text = await response.aread() raise RuntimeError(_friendly_error(response.status_code, text.decode("utf-8", "ignore"))) - return await _consume_sse(response) + return await _consume_sse(response, timeout) def _convert_tools(tools: list[dict[str, Any]]) -> list[dict[str, Any]]: @@ -236,9 +243,17 @@ def _prompt_cache_key(messages: list[dict[str, Any]]) -> str: return hashlib.sha256(raw.encode("utf-8")).hexdigest() -async def _iter_sse(response: httpx.Response) -> AsyncGenerator[dict[str, Any], None]: +async def _iter_sse(response: httpx.Response, timeout: float) -> AsyncGenerator[dict[str, Any], None]: buffer: list[str] = [] - async for line in response.aiter_lines(): + # Per-event idle cap: aiter_lines resets httpx's read timer on every byte, + # so a trickle/keepalive stall never trips it. wait_for on each line bounds + # the silence between SSE lines without penalizing a long, progressing run. + lines = response.aiter_lines() + while True: + try: + line = await asyncio.wait_for(lines.__anext__(), timeout) + except StopAsyncIteration: + break if line == "": if buffer: data_lines = [ln[5:].strip() for ln in buffer if ln.startswith("data:")] @@ -256,13 +271,13 @@ async def _iter_sse(response: httpx.Response) -> AsyncGenerator[dict[str, Any], buffer.append(line) -async def _consume_sse(response: httpx.Response) -> tuple[str, list[ToolCallRequest], str]: +async def _consume_sse(response: httpx.Response, timeout: float) -> tuple[str, list[ToolCallRequest], str]: content = "" tool_calls: list[ToolCallRequest] = [] tool_call_buffers: dict[str, dict[str, Any]] = {} finish_reason = "stop" - async for event in _iter_sse(response): + async for event in _iter_sse(response, timeout): event_type = event.get("type") if event_type == "response.output_item.added": item = event.get("item") or {} diff --git a/raven/providers/per_model_provider.py b/raven/providers/per_model_provider.py index 5a3e35b25..45ebc7a2e 100644 --- a/raven/providers/per_model_provider.py +++ b/raven/providers/per_model_provider.py @@ -30,6 +30,12 @@ def __init__(self, models: "Sequence[ModelEndpoint]", fallback: LLMProvider): if m.model } self._default = next(iter(self._by_model), None) or fallback.get_default_model() + # Per-model sub-providers are built here without generation settings; + # inherit the fallback's (already configured from AgentDefaults) and push + # them down so routed calls honor temperature / max_tokens / timeout. + self.generation = fallback.generation + for sub in self._by_model.values(): + sub.generation = self.generation def _pick(self, model: str | None) -> LLMProvider: return self._by_model.get(model or "", self._fallback) diff --git a/tests/test_agent_loop_stream.py b/tests/test_agent_loop_stream.py index d873b5ad0..28e7501c6 100644 --- a/tests/test_agent_loop_stream.py +++ b/tests/test_agent_loop_stream.py @@ -12,7 +12,7 @@ from typing import Any from raven.agent.loop import AgentLoop -from raven.providers.base import LLMResponse, StreamDelta +from raven.providers.base import LLMProvider, LLMResponse, StreamDelta class _FakeProvider: @@ -217,6 +217,34 @@ async def on_delta(_text: str) -> None: # --------------------------------------------------------------------------- +async def test_llm_call_stream_timeout_returns_structured_error() -> None: + """A mid-stream stall (TimeoutError from the per-chunk idle cap) terminates + with a structured, retryable error response instead of propagating and + crashing the turn. Already-streamed content is preserved on the response.""" + + class _TimeoutStreamProvider: + classify_error = LLMProvider.classify_error + + async def chat_stream(self, **_kwargs: Any): + yield StreamDelta(content="partial") + raise TimeoutError + + call = _bind_helper(_TimeoutStreamProvider()) + seen: list[str] = [] + + async def on_delta(text: str) -> None: + seen.append(text) + + response = await call(messages=[], tools=None, model="m", on_token_delta=on_delta) + + assert response.finish_reason == "error" + assert response.error_classification is not None + assert response.error_classification.category == "network" + assert response.error_classification.retryable is True + assert response.content == "partial" + assert seen == ["partial"] + + async def test_llm_call_stream_empty_stream_yields_empty_content() -> None: """Provider yields zero chunks → response.content == '' + finish_reason='stop'.""" provider = _FakeProvider([]) diff --git a/tests/test_azure_openai_provider.py b/tests/test_azure_openai_provider.py new file mode 100644 index 000000000..1adf1bf19 --- /dev/null +++ b/tests/test_azure_openai_provider.py @@ -0,0 +1,56 @@ +"""Timeout behavior for AzureOpenAIProvider (issue #150). + +The Azure path uses a raw httpx client with a per-read timeout, which cannot +bound a backend that trickles bytes. A wall-clock cap wraps the awaited POST so +a stalled endpoint yields a structured, retryable error instead of hanging. +""" + +from __future__ import annotations + +import asyncio +from typing import Any + +import pytest + +from raven.providers.azure_openai_provider import AzureOpenAIProvider +from raven.providers.base import GenerationSettings + + +class _HangingClient: + """httpx.AsyncClient stand-in whose POST never returns.""" + + def __init__(self, *args: Any, **kwargs: Any) -> None: + pass + + async def __aenter__(self) -> "_HangingClient": + return self + + async def __aexit__(self, *args: Any) -> bool: + return False + + async def post(self, *args: Any, **kwargs: Any) -> Any: + await asyncio.sleep(10) + + +def _make_provider(timeout: float) -> AzureOpenAIProvider: + provider = AzureOpenAIProvider( + api_key="test-key", + api_base="https://example.openai.azure.com", + default_model="gpt-4o", + ) + provider.generation = GenerationSettings(timeout=timeout) + return provider + + +@pytest.mark.asyncio +async def test_chat_wall_clock_cap_returns_classified_error(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + "raven.providers.azure_openai_provider.httpx.AsyncClient", + _HangingClient, + ) + provider = _make_provider(timeout=0.05) + resp = await provider.chat(messages=[{"role": "user", "content": "hi"}], model="gpt-4o") + assert resp.finish_reason == "error" + assert resp.error_classification is not None + assert resp.error_classification.category == "network" + assert resp.error_classification.retryable is True diff --git a/tests/test_cli_gateway_commands.py b/tests/test_cli_gateway_commands.py index 420776bee..fe69e322e 100644 --- a/tests/test_cli_gateway_commands.py +++ b/tests/test_cli_gateway_commands.py @@ -218,12 +218,15 @@ def test_stop_dispatch_cancels_both_scheduler_and_subagents() -> None: from raven.cli.gateway_commands import build_model_routing from raven.config.schema import ModelEndpoint, RoutingConfig +from raven.providers.base import GenerationSettings from raven.providers.per_model_provider import PerModelProvider from raven.routing.knn_router import KNNModelRouter from raven.routing.router import ModelRouter class _FakeProvider: + generation = GenerationSettings() + def get_default_model(self): return "default-model" diff --git a/tests/test_litellm_provider_timeout.py b/tests/test_litellm_provider_timeout.py new file mode 100644 index 000000000..9bfc156a8 --- /dev/null +++ b/tests/test_litellm_provider_timeout.py @@ -0,0 +1,130 @@ +"""Per-call timeout for `LiteLLMProvider` (issue #150). + +Covers: +- chat() and chat_stream() forward `timeout` (= generation.timeout) to acompletion +- chat() wall-clock cap: a hung acompletion yields a structured error response + classified as retryable `network` (so chat_with_retry retries / falls back) +- chat_stream() per-chunk idle cap: a mid-stream stall raises TimeoutError + +Mocks patch `raven.providers.litellm_provider.acompletion` (imported at module +top, so patching `litellm.acompletion` post-import would not be picked up). +""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass +from typing import Any + +import pytest + +from raven.providers.base import GenerationSettings, StreamDelta +from raven.providers.litellm_provider import LiteLLMProvider + + +@dataclass +class _FakeDelta: + content: str | None = None + tool_calls: list[Any] | None = None + + +@dataclass +class _FakeChoice: + delta: _FakeDelta + finish_reason: str | None = None + index: int = 0 + + +@dataclass +class _FakeChunk: + choices: list[_FakeChoice] + usage: Any | None = None + + +def _chunk(content: str | None) -> _FakeChunk: + return _FakeChunk(choices=[_FakeChoice(delta=_FakeDelta(content=content))]) + + +def _make_provider(timeout: float = 600.0) -> LiteLLMProvider: + provider = LiteLLMProvider(api_key="test-key", default_model="openai/gpt-4o") + provider.generation = GenerationSettings(timeout=timeout) + return provider + + +class _FakeResponse: + """Non-streaming acompletion result with one text choice.""" + + def __init__(self, text: str) -> None: + self.choices = [_FakeChoice(delta=_FakeDelta(content=text), finish_reason="stop")] + self.usage = None + + +@pytest.mark.asyncio +async def test_chat_forwards_generation_timeout_to_acompletion(monkeypatch: pytest.MonkeyPatch) -> None: + captured: dict[str, Any] = {} + + async def fake_acompletion(**kwargs: Any): + captured.update(kwargs) + return _FakeResponse("hi") + + monkeypatch.setattr("raven.providers.litellm_provider.acompletion", fake_acompletion) + provider = _make_provider(timeout=123.0) + await provider.chat(messages=[{"role": "user", "content": "hi"}], model="openai/gpt-4o") + assert captured["timeout"] == 123.0 + + +@pytest.mark.asyncio +async def test_chat_stream_forwards_generation_timeout_to_acompletion(monkeypatch: pytest.MonkeyPatch) -> None: + captured: dict[str, Any] = {} + + async def fake_stream(chunks): + for ch in chunks: + yield ch + + async def fake_acompletion(**kwargs: Any): + captured.update(kwargs) + return fake_stream([_chunk("ok")]) + + monkeypatch.setattr("raven.providers.litellm_provider.acompletion", fake_acompletion) + provider = _make_provider(timeout=77.0) + async for _ in provider.chat_stream(messages=[{"role": "user", "content": "hi"}]): + pass + assert captured["timeout"] == 77.0 + + +@pytest.mark.asyncio +async def test_chat_wall_clock_cap_returns_classified_error(monkeypatch: pytest.MonkeyPatch) -> None: + """A backend that never responds is bounded by the wall-clock cap and the + result is a retryable `network` error, not an indefinite hang.""" + + async def hanging_acompletion(**_kwargs: Any): + await asyncio.sleep(10) + + monkeypatch.setattr("raven.providers.litellm_provider.acompletion", hanging_acompletion) + provider = _make_provider(timeout=0.05) + resp = await provider.chat(messages=[{"role": "user", "content": "hi"}], model="openai/gpt-4o") + assert resp.finish_reason == "error" + assert resp.error_classification is not None + assert resp.error_classification.category == "network" + assert resp.error_classification.retryable is True + + +@pytest.mark.asyncio +async def test_chat_stream_idle_cap_raises_timeout(monkeypatch: pytest.MonkeyPatch) -> None: + """A stream that stalls after the first chunk trips the per-chunk idle cap.""" + + async def one_then_hang(**_kwargs: Any): + async def gen(): + yield _chunk("a") + await asyncio.sleep(10) + yield _chunk("b") + + return gen() + + monkeypatch.setattr("raven.providers.litellm_provider.acompletion", one_then_hang) + provider = _make_provider(timeout=0.05) + seen: list[StreamDelta] = [] + with pytest.raises(TimeoutError): + async for delta in provider.chat_stream(messages=[{"role": "user", "content": "hi"}]): + seen.append(delta) + assert [d.content for d in seen] == ["a"] diff --git a/tests/test_openai_codex_provider.py b/tests/test_openai_codex_provider.py index 7dcf7cda5..a2e7b3f22 100644 --- a/tests/test_openai_codex_provider.py +++ b/tests/test_openai_codex_provider.py @@ -7,10 +7,15 @@ from __future__ import annotations +import asyncio + +import pytest + from raven.providers.openai_codex_provider import ( DEFAULT_CODEX_URL, OpenAICodexProvider, _build_headers, + _iter_sse, ) @@ -32,3 +37,27 @@ def test_provider_default_model_is_codex(): assert provider.get_default_model() == "openai-codex/gpt-5.1-codex" # OAuth-based: constructed without an API key. assert provider.api_key is None + + +class _FakeStreamResponse: + """SSE response stand-in: emits complete events, then stalls forever.""" + + def __init__(self, lines: list[str]) -> None: + self._lines = lines + + async def aiter_lines(self): + for line in self._lines: + yield line + await asyncio.sleep(10) + + +@pytest.mark.asyncio +async def test_iter_sse_per_event_idle_timeout_raises(): + """A stream that stalls after a complete event trips the per-event idle cap + instead of hanging (httpx's per-read timeout would reset on the trickle).""" + resp = _FakeStreamResponse(['data: {"type": "ping"}', ""]) + events = [] + with pytest.raises(TimeoutError): + async for event in _iter_sse(resp, timeout=0.05): + events.append(event) + assert events == [{"type": "ping"}] diff --git a/tests/test_per_model_provider.py b/tests/test_per_model_provider.py index f4bfb5e4d..b2ac94beb 100644 --- a/tests/test_per_model_provider.py +++ b/tests/test_per_model_provider.py @@ -7,6 +7,7 @@ import pytest from raven.config.schema import ModelEndpoint +from raven.providers.base import GenerationSettings from raven.providers.per_model_provider import PerModelProvider @@ -45,6 +46,19 @@ def test_default_falls_back_when_no_models(): assert p.get_default_model() == "fallback-model" +def test_generation_propagates_to_sub_providers(): + # Sub-providers are built without generation settings; they must inherit the + # fallback's (configured from AgentDefaults) so routed calls honor + # timeout / temperature / max_tokens instead of the dataclass defaults. + fb = _fallback() + fb.generation = GenerationSettings(timeout=42.0, temperature=0.2, max_tokens=111) + eps = [ModelEndpoint(model="small", api_base="http://a/v1", api_key="KA")] + p = PerModelProvider(eps, fallback=fb) + assert p.generation is fb.generation + assert p._by_model["small"].generation is fb.generation + assert p._by_model["small"].generation.timeout == 42.0 + + @pytest.mark.asyncio async def test_chat_with_retry_dispatches_by_model(): p = _provider() diff --git a/tests/test_provider_fallback_chain.py b/tests/test_provider_fallback_chain.py index 9119f7866..652e0139b 100644 --- a/tests/test_provider_fallback_chain.py +++ b/tests/test_provider_fallback_chain.py @@ -175,6 +175,46 @@ async def test_chain_exhausted_returns_last_error(): assert provider.calls == ["primary"] * 4 + ["backup"] * 4 +def test_builtin_timeout_error_classified_as_retryable_network(): + # asyncio.wait_for raises the builtin TimeoutError, whose class name + # ("timeouterror") and empty str() match neither the network name set nor + # the substring probes. Without the isinstance check it falls to "unknown" + # (not retryable), silently defeating the timeout's retry/fallback intent. + c = LLMProvider.classify_error(TimeoutError()) + assert c.category == "network" + assert c.retryable is True + assert c.should_fallback is True + + +class _TimeoutThenOkProvider(LLMProvider): + """Raises TimeoutError for the first ``fail_times`` calls, then succeeds.""" + + def __init__(self, fail_times: int): + super().__init__(api_key="test") + self._remaining = fail_times + self.calls = 0 + + async def chat(self, messages, tools=None, model=None, **kwargs): + self.calls += 1 + if self._remaining > 0: + self._remaining -= 1 + raise TimeoutError + return LLMResponse(content="ok", finish_reason="stop") + + def get_default_model(self) -> str: + return "default-model" + + +@pytest.mark.asyncio +async def test_chat_timeout_is_retried_then_succeeds(): + provider = _TimeoutThenOkProvider(fail_times=2) + provider._CHAT_RETRY_DELAYS = (0, 0, 0) + resp = await provider.chat_with_retry(messages=[], model="m1") + assert resp.finish_reason == "stop" + assert resp.content == "ok" + assert provider.calls == 3 # two timeouts retried, third succeeds + + @pytest.mark.asyncio async def test_should_fallback_classification(): # Structured classifier (string path): transient + capacity/availability