RomanNum3ral
/
cai
mirror of https://github.com/aliasrobotics/cai.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fix: retry transient streaming provider disconnects

This commit is contained in:
avir4er 2026-07-03 12:47:28 +08:00
parent 3cfbf06313
commit 0ac1c715d0
5 changed files with 201 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
src/cai/cli_headless.py
View File

@ -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

21
src/cai/sdk/agents/models/chatcompletions/litellm_adapter.py
View File

@ -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())

87
src/cai/sdk/agents/models/openai_chatcompletions.py
View File

@ -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:

29
tests/cli/test_cli_headless_cancellation.py
View File

@ -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=[]))

61
tests/core/test_openai_chatcompletions_stream.py
View File

@ -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: