From cdb6bce77fb6430739ec294368924710d449288f Mon Sep 17 00:00:00 2001
From: avir4er <avir4er@protonmail.com>
Date: Fri, 3 Jul 2026 11:32:18 +0800
Subject: [PATCH] fix: bound LiteLLM model wait time

---
 .../models/chatcompletions/litellm_adapter.py | 40 +++++++++++++
 tests/sdk/test_litellm_adapter_streaming.py   | 59 +++++++++++++++++++
 2 files changed, 99 insertions(+)

diff --git a/src/cai/sdk/agents/models/chatcompletions/litellm_adapter.py b/src/cai/sdk/agents/models/chatcompletions/litellm_adapter.py
index 8533bb68..9555a6f0 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 os
 import time
 from typing import TYPE_CHECKING, Any, Literal, cast
 
@@ -28,6 +29,41 @@ if TYPE_CHECKING:
     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) -> dict:
+    """Add bounded model 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 _build_response_obj(
     model: str,
     model_settings: "ModelSettings",
@@ -66,6 +102,8 @@ 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:
             stream_obj = await litellm.acompletion(**kwargs)
@@ -148,6 +186,8 @@ async def fetch_response_litellm_ollama(
 
     api_base = get_ollama_api_base()
 
+    ollama_kwargs = _apply_litellm_timeouts(ollama_kwargs, stream=stream)
+
     if stream:
         response = _build_response_obj(model_name, model_settings, tool_choice, parallel_tool_calls)
         stream_obj = await litellm.acompletion(
diff --git a/tests/sdk/test_litellm_adapter_streaming.py b/tests/sdk/test_litellm_adapter_streaming.py
index 7c05dde4..e7c6fe15 100644
--- a/tests/sdk/test_litellm_adapter_streaming.py
+++ b/tests/sdk/test_litellm_adapter_streaming.py
@@ -84,3 +84,62 @@ async def test_litellm_streaming_tool_call_id_retry_opens_one_retry_stream(monke
     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]