fix: stabilize DeepSeek streaming output

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

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

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

src/cai/sdk/agents/models/chatcompletions/litellm_adapter.py

@@ -8,13 +8,16 @@
 
 from __future__ import annotations
 
+import asyncio
+import os
 import time
 from typing import TYPE_CHECKING, Any, Literal, cast
 
 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
 
@@ -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)