diff --git a/src/cai/cli_headless.py b/src/cai/cli_headless.py index 24ef2d84..bbe22a4f 100644 --- a/src/cai/cli_headless.py +++ b/src/cai/cli_headless.py @@ -77,6 +77,10 @@ from cai.errors import ( 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 @@ def _run_streamed(agent, conversation_input, console, force_until_flag, ctf_glob 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/sdk/agents/models/chatcompletions/litellm_adapter.py b/src/cai/sdk/agents/models/chatcompletions/litellm_adapter.py index 6c7eb9b6..d6c5034e 100644 --- a/src/cai/sdk/agents/models/chatcompletions/litellm_adapter.py +++ b/src/cai/sdk/agents/models/chatcompletions/litellm_adapter.py @@ -66,6 +66,27 @@ def apply_litellm_timeouts(kwargs: dict, *, stream: bool = False) -> dict: 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()) diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index c1f0b9ff..d7baf272 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -133,6 +133,8 @@ 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, @@ -154,7 +156,7 @@ from cai.util.llm_api_base import ( 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, @@ -960,9 +962,11 @@ class OpenAIChatCompletionsModel(Model): 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]}") @@ -977,8 +981,11 @@ class OpenAIChatCompletionsModel(Model): 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( @@ -1785,8 +1792,13 @@ class OpenAIChatCompletionsModel(Model): # 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 @@ -1814,8 +1826,58 @@ class OpenAIChatCompletionsModel(Model): # 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 @@ -3067,8 +3129,15 @@ class OpenAIChatCompletionsModel(Model): 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: diff --git a/tests/cli/test_cli_headless_cancellation.py b/tests/cli/test_cli_headless_cancellation.py index 054194ee..0630c4eb 100644 --- a/tests/cli/test_cli_headless_cancellation.py +++ b/tests/cli/test_cli_headless_cancellation.py @@ -4,9 +4,11 @@ from types import SimpleNamespace 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 test_streamed_cancelled_error_uses_interrupt_flow(monkeypatch): ) +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 4c197c03..c6788d4b 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 test_stream_response_yields_events_for_text_content(monkeypatch) -> No 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: