diff --git a/src/cai/cli_headless.py b/src/cai/cli_headless.py index 24ef2d849..bbe22a4ff 100644 --- a/src/cai/cli_headless.py +++ b/src/cai/cli_headless.py @@ -77,6 +77,10 @@ LLMTimeout, ) from cai.sdk.agents.models.chatcompletions.httpx_client import verbose_http_retries +from cai.sdk.agents.models.chatcompletions.litellm_adapter import ( + is_transient_litellm_provider_error, + provider_error_summary, +) from cai.continuation import generate_continuation_advice, should_continue_automatically from litellm.exceptions import RateLimitError, Timeout @@ -1497,6 +1501,14 @@ async def process_streamed_response(agent, conversation_input): except Exception: pass logger = logging.getLogger(__name__) + if isinstance(e, (LLMProviderUnavailable, LLMTimeout, LLMRateLimited)): + raise + if is_transient_litellm_provider_error(e): + summary = provider_error_summary(e) + logger.warning("Streaming provider error: %s", summary) + raise LLMProviderUnavailable( + f"Model provider disconnected during streaming: {summary}" + ) from e logger.error(f"Error occurred during streaming: {str(e)}", exc_info=True) if _get_config().debug == 2: import traceback diff --git a/src/cai/continuation.py b/src/cai/continuation.py index df5445645..2de6717f4 100644 --- a/src/cai/continuation.py +++ b/src/cai/continuation.py @@ -11,6 +11,7 @@ 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 @@ async def generate_continuation_advice( 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 @@ async def generate_continuation_advice( # 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 04a3555d4..059d1141b 100644 --- a/src/cai/ctr/digest.py +++ b/src/cai/ctr/digest.py @@ -19,6 +19,8 @@ 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 @@ async def generate_llm_digest(ctr_dir: str) -> str: - 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 @@ async def generate_llm_digest(ctr_dir: str) -> str: 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 795d65a01..d6c5034ec 100644 --- a/src/cai/sdk/agents/models/chatcompletions/litellm_adapter.py +++ b/src/cai/sdk/agents/models/chatcompletions/litellm_adapter.py @@ -8,6 +8,8 @@ from __future__ import annotations +import asyncio +import os import time from typing import TYPE_CHECKING, Any, Literal, cast @@ -15,6 +17,7 @@ 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 @@ -28,6 +31,140 @@ from ...model_settings import ModelSettings +_DEFAULT_MODEL_TIMEOUT = 180.0 + + +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 + as a compatibility alias for local configs/scripts. Values <= 0 disable the + injected timeout and defer entirely to LiteLLM/provider defaults. + """ + raw = os.getenv("CAI_MODEL_TIMEOUT") + if raw is None: + raw = os.getenv("CAI_LLM_TIMEOUT") + if raw is None or raw == "": + return _DEFAULT_MODEL_TIMEOUT + try: + timeout = float(raw) + except (TypeError, ValueError): + return _DEFAULT_MODEL_TIMEOUT + if timeout <= 0: + return None + return 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) + if stream: + kwargs.setdefault("stream_timeout", timeout) + return kwargs + + +def is_transient_litellm_provider_error(exc: BaseException) -> bool: + """Return True for provider/proxy failures that are safe to retry.""" + return isinstance( + exc, + ( + litellm.exceptions.APIConnectionError, + litellm.exceptions.BadGatewayError, + litellm.exceptions.InternalServerError, + litellm.exceptions.ServiceUnavailableError, + ), + ) + + +def provider_error_summary(exc: BaseException, *, limit: int = 220) -> str: + """Compact provider error text for user-facing typed exceptions.""" + message = " ".join(str(exc).split()) + if len(message) > limit: + message = f"{message[:limit]}..." + return f"{type(exc).__name__}: {message}" + + +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", @@ -66,13 +203,14 @@ async def fetch_response_litellm_openai( too long, truncate all tool_call ids in the messages to 40 characters and retry once silently. """ + kwargs = _apply_litellm_timeouts(kwargs, stream=stream) + try: if stream: - ret = await litellm.acompletion(**kwargs) - 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 ( @@ -102,11 +240,10 @@ async def fetch_response_litellm_openai( kwargs["messages"] = messages if stream: - ret = await litellm.acompletion(**kwargs) - 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 @@ -150,15 +287,17 @@ async def fetch_response_litellm_ollama( api_base = get_ollama_api_base() + 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 47d871b0b..d7baf272e 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -130,8 +130,11 @@ 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, + is_transient_litellm_provider_error, + provider_error_summary, ) from .chatcompletions.model import ( ACTIVE_MODEL_INSTANCES, @@ -153,7 +156,7 @@ resolve_llm_openai_compatible_base, resolve_llm_openai_compatible_api_key, ) -from cai.errors import LLMEmptyAssistantError, LLMRateLimited, LLMTimeout +from cai.errors import LLMEmptyAssistantError, LLMProviderUnavailable, LLMRateLimited, LLMTimeout from cai.util.gateway_rate_limiter import ( COMPLETION_BUDGET_TOKENS, get_gateway_rate_limiter, @@ -959,9 +962,11 @@ async def get_response( return result except ( + litellm.exceptions.APIConnectionError, litellm.exceptions.BadGatewayError, litellm.exceptions.ServiceUnavailableError, litellm.exceptions.InternalServerError, + LLMProviderUnavailable, ) as e: # Transient server errors (502, 503, 500): retry with backoff self.logger.warning(f"Server error (high-level recovery): {str(e)[:200]}") @@ -976,8 +981,11 @@ async def get_response( if self._high_level_retry_count > 3: self._high_level_retry_count = 0 if verbose_http_retries(): - print(f"\n❌ Server error after 3 recovery attempts [{self.model}]") - raise + print(f"\n❌ Provider error after 3 recovery attempts [{self.model}]") + raise LLMProviderUnavailable( + f"Provider unavailable after 3 recovery attempts " + f"[{self.model}]: {provider_error_summary(e)}" + ) from e wait_secs = 10 * self._high_level_retry_count # 10s, 20s, 30s self.logger.warning( @@ -1784,8 +1792,13 @@ async def stream_response( # Clean retry: same input, no "continue" in history async for event in self.stream_response( - system_instructions, input, model_settings, - tools, output_schema, handoffs, tracing, + system_instructions, + input, + model_settings, + tools, + output_schema, + handoffs, + tracing, ): yield event self._high_level_retry_count = 0 @@ -1813,8 +1826,58 @@ async def stream_response( # Clean retry: same input, no "continue" in history async for event in self.stream_response( - system_instructions, input, model_settings, - tools, output_schema, handoffs, tracing, + system_instructions, + input, + model_settings, + tools, + output_schema, + handoffs, + tracing, + ): + yield event + self._high_level_retry_count = 0 + return + + except ( + litellm.exceptions.APIConnectionError, + litellm.exceptions.BadGatewayError, + litellm.exceptions.ServiceUnavailableError, + litellm.exceptions.InternalServerError, + LLMProviderUnavailable, + ) as e: + await stream_wait_hints.stop() + self.logger.warning( + "Transient provider error in stream_response [%s]: %s", + self.model, + provider_error_summary(e), + ) + stop_active_timer() + start_idle_timer() + + if not hasattr(self, "_high_level_retry_count"): + self._high_level_retry_count = 0 + self._high_level_retry_count += 1 + + if self._high_level_retry_count > 3: + self._high_level_retry_count = 0 + raise LLMProviderUnavailable( + f"Provider unavailable after 3 attempts " + f"[{self.model}]: {provider_error_summary(e)}" + ) from e + + await self._retry_with_backoff( + self._high_level_retry_count - 1, "Provider error" + ) + + # Clean retry: same input, no "continue" in history + async for event in self.stream_response( + system_instructions, + input, + model_settings, + tools, + output_schema, + handoffs, + tracing, ): yield event self._high_level_retry_count = 0 @@ -1994,10 +2057,7 @@ def next_sequence_number() -> int: # Check if model supports reasoning (Claude or DeepSeek) model_str_lower = str(self.model).lower() - if ( - detect_claude_thinking_in_stream(str(self.model)) - or "deepseek" in model_str_lower - ): + if detect_claude_thinking_in_stream(str(self.model)): print_claude_reasoning_simple( reasoning_content, self.agent_name, str(self.model) ) @@ -3069,8 +3129,15 @@ def next_sequence_number() -> int: raise except Exception as e: - # Handle other exceptions - logger.error(f"Error in stream_response: {e}") + # Provider/proxy errors are already retried above; keep logs concise. + if isinstance(e, (LLMProviderUnavailable, LLMTimeout, LLMRateLimited)) or ( + is_transient_litellm_provider_error(e) + ): + logger.warning( + "Model provider error in stream_response: %s", provider_error_summary(e) + ) + else: + logger.error(f"Error in stream_response: {e}") raise finally: @@ -4119,10 +4186,14 @@ async def stream_response(): 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 @@ -4171,11 +4242,15 @@ async def stream_response(): 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 c70f1e918..465edc86e 100644 --- a/src/cai/tui/components/agent_creator_panel.py +++ b/src/cai/tui/components/agent_creator_panel.py @@ -15,7 +15,7 @@ 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 @@ async def _use_meta_agent_for_creation(self, description: str, selected_tools: L """ # 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/src/cai/util/streaming.py b/src/cai/util/streaming.py index ead368db2..f37065b85 100644 --- a/src/cai/util/streaming.py +++ b/src/cai/util/streaming.py @@ -4208,7 +4208,7 @@ def create_claude_thinking_context(agent_name, counter, model): context = { "thinking_id": thinking_id, "live": live, - "panel": panel, + "panel": None, "header": header, "thinking_content": thinking_content, "timestamp": timestamp, @@ -4358,10 +4358,29 @@ def finish_claude_thinking_display(context): return False +def _raw_reasoning_display_enabled() -> bool: + """Return whether raw provider reasoning text should be displayed. + + DeepSeek-compatible gateways often stream ``reasoning_content`` in tiny + deltas. Rendering those deltas by default floods the terminal and can leak + model-internal scratch text. Keep it opt-in for DeepSeek via + ``CAI_SHOW_REASONING=true`` (or ``CAI_SHOW_THINKING=true`` for older local + configs). + """ + raw = os.getenv("CAI_SHOW_REASONING") + if raw is None: + raw = os.getenv("CAI_SHOW_THINKING") + if raw is None: + return False + return raw.strip().lower() in ("1", "true", "yes", "on") + + def detect_claude_thinking_in_stream(model_name): """ Detect if a model should show thinking/reasoning display. - Applies to Claude and DeepSeek models with reasoning capability. + + Claude keeps the historical default. DeepSeek raw reasoning is opt-in + because providers commonly stream it token-by-token as ``reasoning_content``. Args: model_name: The model name to check @@ -4389,17 +4408,10 @@ def detect_claude_thinking_in_stream(model_name): or "thinking" in model_str ) - # Check for DeepSeek models with reasoning capability - has_deepseek_reasoning = "deepseek" in model_str and ( - # DeepSeek reasoner models - "reasoner" in model_str - or - # DeepSeek chat models also support reasoning - "chat" in model_str - or - # Generic deepseek models likely support it - "/" in model_str # e.g., deepseek/deepseek-chat - ) + # DeepSeek reasoning display is intentionally opt-in. The text is still + # accumulated internally by the model stream handler so empty-response + # detection and accounting keep working; it is just not printed by default. + has_deepseek_reasoning = "deepseek" in model_str and _raw_reasoning_display_enabled() return has_claude_reasoning or has_deepseek_reasoning diff --git a/tests/cli/test_cli_headless_cancellation.py b/tests/cli/test_cli_headless_cancellation.py index 054194ee1..0630c4ebc 100644 --- a/tests/cli/test_cli_headless_cancellation.py +++ b/tests/cli/test_cli_headless_cancellation.py @@ -4,9 +4,11 @@ from unittest.mock import Mock import pytest +import litellm from cai import cli_headless from cai import parallel_worker +from cai.errors import LLMProviderUnavailable def test_non_streamed_cancelled_error_uses_interrupt_flow(monkeypatch): @@ -41,6 +43,33 @@ def cancelled_asyncio_run(*_args, **_kwargs): ) +def test_streamed_provider_disconnect_raises_typed_provider_error(monkeypatch): + class DummyResult: + async def stream_events(self): + raise litellm.exceptions.InternalServerError( + message="DeepseekException - Server disconnected", + llm_provider="deepseek", + model="deepseek/deepseek-v4-pro", + ) + yield # pragma: no cover + + def _cleanup_tasks(self): + pass + + monkeypatch.setattr( + cli_headless.Runner, "run_streamed", lambda *_args, **_kwargs: DummyResult() + ) + + with pytest.raises(LLMProviderUnavailable, match="Server disconnected"): + cli_headless._run_streamed( + SimpleNamespace(model=SimpleNamespace(message_history=[])), + "input", + Mock(), + False, + None, + ) + + def test_simple_parallel_cancelled_error_uses_interrupt_flow(monkeypatch): dummy_agent = SimpleNamespace(model=SimpleNamespace(model="test-model", message_history=[])) diff --git a/tests/core/test_openai_chatcompletions_stream.py b/tests/core/test_openai_chatcompletions_stream.py index 4c197c034..c6788d4b3 100644 --- a/tests/core/test_openai_chatcompletions_stream.py +++ b/tests/core/test_openai_chatcompletions_stream.py @@ -1,6 +1,7 @@ from collections.abc import AsyncIterator import pytest +import litellm from openai.types.chat.chat_completion_chunk import ( ChatCompletionChunk, Choice, @@ -121,6 +122,66 @@ async def patched_fetch_response(self, *args, **kwargs): assert completed_resp.usage.total_tokens == 12 +@pytest.mark.allow_call_model_methods +@pytest.mark.asyncio +async def test_stream_response_retries_transient_provider_disconnect(monkeypatch) -> None: + chunk = ChatCompletionChunk( + id="chunk-id", + created=1, + model="fake", + object="chat.completion.chunk", + choices=[Choice(index=0, delta=ChoiceDelta(content="ok"))], + usage=CompletionUsage(completion_tokens=1, prompt_tokens=1, total_tokens=2), + ) + + async def fake_stream() -> AsyncIterator[ChatCompletionChunk]: + yield chunk + + calls = {"count": 0} + + async def patched_fetch_response(self, *args, **kwargs): + calls["count"] += 1 + if calls["count"] == 1: + raise litellm.exceptions.InternalServerError( + message="DeepseekException - Server disconnected", + llm_provider="deepseek", + model="deepseek/deepseek-v4-pro", + ) + resp = Response( + id="resp-id", + created_at=0, + model="fake-model", + object="response", + output=[], + tool_choice="none", + tools=[], + parallel_tool_calls=False, + ) + return resp, fake_stream() + + async def no_sleep_retry(self, *_args, **_kwargs): + return None + + monkeypatch.setattr(OpenAIChatCompletionsModel, "_fetch_response", patched_fetch_response) + monkeypatch.setattr(OpenAIChatCompletionsModel, "_retry_with_backoff", no_sleep_retry) + + model = OpenAIProvider(use_responses=False).get_model(cai_model) + events = [] + async for event in model.stream_response( + system_instructions=None, + input="", + model_settings=ModelSettings(), + tools=[], + output_schema=None, + handoffs=[], + tracing=ModelTracing.DISABLED, + ): + events.append(event) + + assert calls["count"] == 2 + assert any(getattr(event, "delta", None) == "ok" for event in events) + + @pytest.mark.allow_call_model_methods @pytest.mark.asyncio async def test_stream_response_yields_events_for_refusal_content(monkeypatch) -> None: diff --git a/tests/sdk/test_litellm_adapter_streaming.py b/tests/sdk/test_litellm_adapter_streaming.py new file mode 100644 index 000000000..a55f88f7b --- /dev/null +++ b/tests/sdk/test_litellm_adapter_streaming.py @@ -0,0 +1,213 @@ +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, +) + + +@pytest.mark.asyncio +async def test_litellm_streaming_fetch_opens_one_completion(monkeypatch): + calls = [] + sentinel_stream = object() + + async def fake_acompletion(**kwargs): + calls.append(kwargs.copy()) + return sentinel_stream + + 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, + ) + + assert stream is sentinel_stream + assert response.model == "deepseek/deepseek-v4-pro" + assert len(calls) == 1 + assert calls[0]["stream"] is True + + +@pytest.mark.asyncio +async def test_litellm_streaming_tool_call_id_retry_opens_one_retry_stream(monkeypatch): + calls = [] + sentinel_stream = object() + + async def fake_acompletion(**kwargs): + calls.append(kwargs.copy()) + if len(calls) == 1: + raise Exception("Invalid 'messages': tool_call_id string too long maximum length") + return sentinel_stream + + monkeypatch.setattr( + "cai.sdk.agents.models.chatcompletions.litellm_adapter.litellm.acompletion", + fake_acompletion, + ) + + long_id = "call_" + "x" * 80 + kwargs = { + "model": "deepseek/deepseek-v4-pro", + "messages": [ + {"role": "tool", "tool_call_id": long_id, "content": "ok"}, + { + "role": "assistant", + "tool_calls": [ + { + "id": long_id, + "type": "function", + "function": {"name": "probe", "arguments": "{}"}, + } + ], + }, + ], + "stream": True, + } + + _response, stream = await fetch_response_litellm_openai( + kwargs=kwargs, + model_name="deepseek/deepseek-v4-pro", + model_settings=ModelSettings(), + tool_choice=NOT_GIVEN, + stream=True, + parallel_tool_calls=False, + ) + + assert stream is sentinel_stream + assert len(calls) == 2 + assert kwargs["messages"][0]["tool_call_id"] == long_id[:40] + assert kwargs["messages"][1]["tool_calls"][0]["id"] == long_id[:40] + + +@pytest.mark.asyncio +async def test_litellm_streaming_applies_default_model_timeout(monkeypatch): + monkeypatch.delenv("CAI_MODEL_TIMEOUT", raising=False) + monkeypatch.delenv("CAI_LLM_TIMEOUT", raising=False) + calls = [] + sentinel_stream = object() + + async def fake_acompletion(**kwargs): + calls.append(kwargs.copy()) + return sentinel_stream + + 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, + ) + + assert stream is sentinel_stream + assert calls[0]["timeout"] == 180.0 + assert calls[0]["stream_timeout"] == 180.0 + + +@pytest.mark.asyncio +async def test_litellm_model_timeout_uses_env_override(monkeypatch): + monkeypatch.setenv("CAI_MODEL_TIMEOUT", "45") + calls = [] + sentinel_response = object() + + async def fake_acompletion(**kwargs): + calls.append(kwargs.copy()) + return sentinel_response + + monkeypatch.setattr( + "cai.sdk.agents.models.chatcompletions.litellm_adapter.litellm.acompletion", + fake_acompletion, + ) + + response = await 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, + ) + + 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__() diff --git a/tests/util/test_thinking_display.py b/tests/util/test_thinking_display.py new file mode 100644 index 000000000..ea2065588 --- /dev/null +++ b/tests/util/test_thinking_display.py @@ -0,0 +1,39 @@ +from cai.util import streaming + + +def test_create_claude_thinking_context_for_deepseek_does_not_reference_missing_panel(capsys): + streaming._CLAUDE_THINKING_PANELS.clear() + + context = streaming.create_claude_thinking_context( + "Web App Pentester", + 1, + "deepseek/deepseek-v4-pro", + ) + + captured = capsys.readouterr() + assert context is not None + assert "Error creating DeepSeek thinking context" not in captured.out + assert context["model_display"] == "DeepSeek" + assert context["accumulated_thinking"] == "" + assert context["is_started"] is False + + thinking_id = context["thinking_id"] + assert streaming._CLAUDE_THINKING_PANELS[thinking_id] is context + streaming._CLAUDE_THINKING_PANELS.clear() + + +def test_deepseek_reasoning_display_is_opt_in(monkeypatch): + monkeypatch.delenv("CAI_SHOW_REASONING", raising=False) + monkeypatch.delenv("CAI_SHOW_THINKING", raising=False) + + assert streaming.detect_claude_thinking_in_stream("deepseek/deepseek-v4-pro") is False + + monkeypatch.setenv("CAI_SHOW_REASONING", "true") + assert streaming.detect_claude_thinking_in_stream("deepseek/deepseek-v4-pro") is True + + +def test_claude_reasoning_display_keeps_historical_default(monkeypatch): + monkeypatch.delenv("CAI_SHOW_REASONING", raising=False) + monkeypatch.delenv("CAI_SHOW_THINKING", raising=False) + + assert streaming.detect_claude_thinking_in_stream("claude-sonnet-4-20250514") is True