From 3cfbf06313f760dceecbb462c27f9def762638c8 Mon Sep 17 00:00:00 2001 From: avir4er Date: Fri, 3 Jul 2026 11:57:22 +0800 Subject: [PATCH] fix: enforce LiteLLM completion timeout --- src/cai/continuation.py | 7 +- src/cai/ctr/digest.py | 5 +- .../models/chatcompletions/litellm_adapter.py | 112 +++++++++++++++--- .../agents/models/openai_chatcompletions.py | 17 ++- src/cai/tui/components/agent_creator_panel.py | 16 ++- tests/sdk/test_litellm_adapter_streaming.py | 68 +++++++++++ 6 files changed, 193 insertions(+), 32 deletions(-) diff --git a/src/cai/continuation.py b/src/cai/continuation.py index df544564..2de6717f 100644 --- a/src/cai/continuation.py +++ b/src/cai/continuation.py @@ -11,6 +11,7 @@ from typing import List, Dict, Any, Optional from rich.console import Console from cai.config import get_config +from cai.sdk.agents.models.chatcompletions.litellm_adapter import acompletion_with_timeout logger = logging.getLogger(__name__) @@ -108,9 +109,7 @@ Generate a specific, actionable continuation prompt that: IMPORTANT: Respond with ONLY the continuation prompt. No explanations, no "Here's a prompt:", just the direct instruction.""" try: - # Use litellm directly, which is how the rest of the codebase handles API calls - import litellm - + # Use LiteLLM through CAI's timeout wrapper. # Enable debug logging for litellm if in debug mode if logger.isEnabledFor(logging.DEBUG): logger.debug(f"Generating continuation advice with model: {model_name}") @@ -138,7 +137,7 @@ IMPORTANT: Respond with ONLY the continuation prompt. No explanations, no "Here' # Make the API call logger.debug(f"Making API call with kwargs: {kwargs.get('model')}, provider: {kwargs.get('custom_llm_provider', 'default')}") - response = await litellm.acompletion(**kwargs) + response = await acompletion_with_timeout(kwargs, stream=False, model_name=model_name) # Extract content safely continuation_prompt = None diff --git a/src/cai/ctr/digest.py b/src/cai/ctr/digest.py index 04a3555d..059d1141 100644 --- a/src/cai/ctr/digest.py +++ b/src/cai/ctr/digest.py @@ -19,6 +19,8 @@ from pathlib import Path from typing import Dict, List, Tuple, Optional from rich.console import Console +from cai.sdk.agents.models.chatcompletions.litellm_adapter import acompletion_with_timeout + console = Console() # Global cache for digest results (per-run caching) @@ -519,7 +521,6 @@ OUTPUT REQUIREMENTS: - Maximum 350 words""" # Use LiteLLM for model compatibility (handles alias1, OpenRouter, etc.) - import litellm model = os.getenv("CAI_CTR_DIGEST_MODEL", "alias1") @@ -549,7 +550,7 @@ OUTPUT REQUIREMENTS: kwargs["api_key"] = os.getenv("ALIAS_API_KEY", "sk-alias-1234567890") try: - response = await litellm.acompletion(**kwargs) + response = await acompletion_with_timeout(kwargs, stream=False, model_name=model) # Extract content (handle reasoning models like alias1/o1 that use reasoning_content) message = response.choices[0].message diff --git a/src/cai/sdk/agents/models/chatcompletions/litellm_adapter.py b/src/cai/sdk/agents/models/chatcompletions/litellm_adapter.py index 9555a6f0..6c7eb9b6 100644 --- a/src/cai/sdk/agents/models/chatcompletions/litellm_adapter.py +++ b/src/cai/sdk/agents/models/chatcompletions/litellm_adapter.py @@ -8,6 +8,7 @@ Extracted from openai_chatcompletions.py [F] to reduce monolith size. from __future__ import annotations +import asyncio import os import time from typing import TYPE_CHECKING, Any, Literal, cast @@ -16,6 +17,7 @@ import litellm from openai import NOT_GIVEN, NotGiven from openai.types.responses import Response +from cai.errors import LLMTimeout from cai.util import get_ollama_api_base from ..fake_id import FAKE_RESPONSES_ID @@ -32,7 +34,7 @@ if TYPE_CHECKING: _DEFAULT_MODEL_TIMEOUT = 180.0 -def _configured_model_timeout() -> float | None: +def configured_model_timeout() -> float | None: """Return CAI's LiteLLM request timeout in seconds. ``CAI_MODEL_TIMEOUT`` is the public name. ``CAI_LLM_TIMEOUT`` is accepted @@ -53,9 +55,9 @@ def _configured_model_timeout() -> float | None: return timeout -def _apply_litellm_timeouts(kwargs: dict, *, stream: bool) -> dict: - """Add bounded model request timeouts unless the caller already set them.""" - timeout = _configured_model_timeout() +def apply_litellm_timeouts(kwargs: dict, *, stream: bool = False) -> dict: + """Add bounded LiteLLM request timeouts unless the caller already set them.""" + timeout = configured_model_timeout() if timeout is None: return kwargs kwargs.setdefault("timeout", timeout) @@ -64,6 +66,84 @@ def _apply_litellm_timeouts(kwargs: dict, *, stream: bool) -> dict: return kwargs +def _timeout_from_kwargs(kwargs: dict) -> float | None: + """Return the effective numeric timeout for CAI's outer asyncio guard.""" + raw_timeout = kwargs.get("timeout", configured_model_timeout()) + if raw_timeout is None: + return None + try: + timeout = float(raw_timeout) + except (TypeError, ValueError): + return configured_model_timeout() + if timeout <= 0: + return None + return timeout + + +def wrap_stream_with_idle_timeout(stream_obj: Any, *, model_name: str, timeout: float | None = None) -> Any: + """Bound waits for each streamed chunk. + + Some LiteLLM/provider combinations return the stream object quickly, then + stall while the caller awaits the next SSE chunk. ``timeout``/ + ``stream_timeout`` do not consistently protect that phase, so CAI wraps the + async iterator itself. Non-async-iterable test doubles are returned as-is. + """ + if timeout is None: + timeout = configured_model_timeout() + if timeout is None or not hasattr(stream_obj, "__aiter__"): + return stream_obj + + async def _iter_with_timeout(): + iterator = stream_obj.__aiter__() + while True: + try: + chunk = await asyncio.wait_for(iterator.__anext__(), timeout=timeout) + except StopAsyncIteration: + return + except asyncio.TimeoutError as exc: + raise LLMTimeout( + f"Timed out waiting for streamed model chunk after {timeout:g}s " + f"[{model_name}]" + ) from exc + yield chunk + + return _iter_with_timeout() + + +async def acompletion_with_timeout( + kwargs: dict, + *, + stream: bool = False, + model_name: str | None = None, +) -> Any: + """Call LiteLLM with CAI request and stream-idle timeouts applied.""" + kwargs = apply_litellm_timeouts(kwargs, stream=stream) + timeout = _timeout_from_kwargs(kwargs) + model_label = str(model_name or kwargs.get("model") or "unknown model") + + completion_coro = litellm.acompletion(**kwargs) + try: + if timeout is None: + result = await completion_coro + else: + result = await asyncio.wait_for(completion_coro, timeout=timeout) + except asyncio.TimeoutError as exc: + raise LLMTimeout( + f"Timed out waiting for model response after {timeout:g}s [{model_label}]" + ) from exc + + if stream: + return wrap_stream_with_idle_timeout(result, model_name=model_label, timeout=timeout) + return result + + +# Backward-compatible private aliases for local/internal imports. +_configured_model_timeout = configured_model_timeout +_apply_litellm_timeouts = apply_litellm_timeouts +_wrap_stream_with_idle_timeout = wrap_stream_with_idle_timeout +_acompletion_with_timeout = acompletion_with_timeout + + def _build_response_obj( model: str, model_settings: "ModelSettings", @@ -106,10 +186,10 @@ async def fetch_response_litellm_openai( try: if stream: - stream_obj = await litellm.acompletion(**kwargs) + stream_obj = await acompletion_with_timeout(kwargs, stream=True, model_name=model_name) return _build_response_obj(model_name, model_settings, tool_choice, parallel_tool_calls), stream_obj else: - return await litellm.acompletion(**kwargs) + return await acompletion_with_timeout(kwargs, stream=False, model_name=model_name) except Exception as e: error_msg = str(e) if ( @@ -139,10 +219,10 @@ async def fetch_response_litellm_openai( kwargs["messages"] = messages if stream: - stream_obj = await litellm.acompletion(**kwargs) + stream_obj = await acompletion_with_timeout(kwargs, stream=True, model_name=model_name) return _build_response_obj(model_name, model_settings, tool_choice, parallel_tool_calls), stream_obj else: - return await litellm.acompletion(**kwargs) + return await acompletion_with_timeout(kwargs, stream=False, model_name=model_name) else: raise @@ -188,15 +268,15 @@ async def fetch_response_litellm_ollama( ollama_kwargs = _apply_litellm_timeouts(ollama_kwargs, stream=stream) + call_kwargs = { + **ollama_kwargs, + "api_base": api_base, + "custom_llm_provider": "openai", + } + if stream: response = _build_response_obj(model_name, model_settings, tool_choice, parallel_tool_calls) - stream_obj = await litellm.acompletion( - **ollama_kwargs, api_base=api_base, custom_llm_provider="openai" - ) + stream_obj = await acompletion_with_timeout(call_kwargs, stream=True, model_name=model_name) return response, stream_obj else: - return await litellm.acompletion( - **ollama_kwargs, - api_base=api_base, - custom_llm_provider="openai", - ) + return await acompletion_with_timeout(call_kwargs, stream=False, model_name=model_name) diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index abd933b6..c1f0b9ff 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -130,6 +130,7 @@ from .chatcompletions.httpx_client import ( verbose_http_retries, ) from .chatcompletions.litellm_adapter import ( + acompletion_with_timeout, fetch_response_litellm_openai as _fetch_litellm_openai_impl, fetch_response_litellm_ollama as _fetch_litellm_ollama_impl, ) @@ -4116,10 +4117,14 @@ class OpenAIChatCompletionsModel(Model): tools=[], parallel_tool_calls=parallel_tool_calls or False, ) - stream_obj = await litellm.acompletion(**retry_kwargs) + stream_obj = await acompletion_with_timeout( + retry_kwargs, stream=True, model_name=str(self.model) + ) return response, stream_obj else: - ret = await litellm.acompletion(**retry_kwargs) + ret = await acompletion_with_timeout( + retry_kwargs, stream=False, model_name=str(self.model) + ) return ret except Exception: # If retry also fails, raise the original error @@ -4168,11 +4173,15 @@ class OpenAIChatCompletionsModel(Model): tools=[], parallel_tool_calls=parallel_tool_calls or False, ) - stream_obj = await litellm.acompletion(**qwen_params) + stream_obj = await acompletion_with_timeout( + qwen_params, stream=True, model_name=str(self.model) + ) return response, stream_obj else: # Non-streaming case - ret = await litellm.acompletion(**qwen_params) + ret = await acompletion_with_timeout( + qwen_params, stream=False, model_name=str(self.model) + ) return ret except Exception as direct_e: # All approaches failed, log and raise the original error diff --git a/src/cai/tui/components/agent_creator_panel.py b/src/cai/tui/components/agent_creator_panel.py index c70f1e91..465edc86 100644 --- a/src/cai/tui/components/agent_creator_panel.py +++ b/src/cai/tui/components/agent_creator_panel.py @@ -15,7 +15,7 @@ import json import os from typing import Dict, List, Any, Optional from cai.agents.agent_builder import AgentBuilder -import litellm +from cai.sdk.agents.models.chatcompletions.litellm_adapter import acompletion_with_timeout class AgentCreationConfirmed(Message): @@ -514,14 +514,18 @@ IMPORTANT: The "tools" field in your response must be exactly: {json.dumps(selec """ # Use litellm to generate the configuration - response = await litellm.acompletion( - model=os.getenv("CAI_MODEL", "gpt-4"), - messages=[ + model_name = os.getenv("CAI_MODEL", "gpt-4") + kwargs = { + "model": model_name, + "messages": [ {"role": "system", "content": "You are an AI agent configuration generator. Always respond with valid JSON only."}, {"role": "user", "content": meta_prompt} ], - temperature=0.7, - max_tokens=2000 + "temperature": 0.7, + "max_tokens": 2000, + } + response = await acompletion_with_timeout( + kwargs, stream=False, model_name=model_name ) # Parse the response diff --git a/tests/sdk/test_litellm_adapter_streaming.py b/tests/sdk/test_litellm_adapter_streaming.py index e7c6fe15..a55f88f7 100644 --- a/tests/sdk/test_litellm_adapter_streaming.py +++ b/tests/sdk/test_litellm_adapter_streaming.py @@ -1,6 +1,10 @@ +import asyncio + import pytest from openai import NOT_GIVEN +from cai.errors import LLMTimeout + from cai.sdk.agents.model_settings import ModelSettings from cai.sdk.agents.models.chatcompletions.litellm_adapter import ( fetch_response_litellm_openai, @@ -143,3 +147,67 @@ async def test_litellm_model_timeout_uses_env_override(monkeypatch): assert response is sentinel_response assert calls[0]["timeout"] == 45.0 assert "stream_timeout" not in calls[0] + + +@pytest.mark.asyncio +async def test_litellm_nonstream_times_out_when_completion_call_stalls(monkeypatch): + monkeypatch.setenv("CAI_MODEL_TIMEOUT", "0.01") + + async def fake_acompletion(**_kwargs): + await asyncio.sleep(60) + return object() + + monkeypatch.setattr( + "cai.sdk.agents.models.chatcompletions.litellm_adapter.litellm.acompletion", + fake_acompletion, + ) + + with pytest.raises(LLMTimeout, match="Timed out waiting for model response"): + await asyncio.wait_for( + fetch_response_litellm_openai( + kwargs={ + "model": "deepseek/deepseek-v4-pro", + "messages": [], + "stream": False, + }, + model_name="deepseek/deepseek-v4-pro", + model_settings=ModelSettings(), + tool_choice=NOT_GIVEN, + stream=False, + parallel_tool_calls=False, + ), + timeout=0.2, + ) + + +@pytest.mark.asyncio +async def test_litellm_streaming_times_out_when_next_chunk_stalls(monkeypatch): + monkeypatch.setenv("CAI_MODEL_TIMEOUT", "0.01") + + class StalledStream: + def __aiter__(self): + return self + + async def __anext__(self): + await asyncio.sleep(60) + return object() + + async def fake_acompletion(**_kwargs): + return StalledStream() + + monkeypatch.setattr( + "cai.sdk.agents.models.chatcompletions.litellm_adapter.litellm.acompletion", + fake_acompletion, + ) + + _response, stream = await fetch_response_litellm_openai( + kwargs={"model": "deepseek/deepseek-v4-pro", "messages": [], "stream": True}, + model_name="deepseek/deepseek-v4-pro", + model_settings=ModelSettings(), + tool_choice=NOT_GIVEN, + stream=True, + parallel_tool_calls=False, + ) + + with pytest.raises(LLMTimeout, match="Timed out waiting for streamed model chunk"): + await stream.__anext__()