Add LLM retry logic and tests for providers/sessions

- Base LLMProvider now retries on 429/5xx/timeout/connect errors (2 retries, exponential backoff)
- Providers implement _generate(), base class wraps with retry in generate()
- New test_llm.py: provider selection, timeout config, retry/exhaust scenarios
- New test_sessions.py: session CRUD, enqueue, TTL expiry, GC sweep
- Test count: 42 → 63

Author: loglux

app/llm/base.py

@@ -1,6 +1,51 @@
+import asyncio
+import logging
 from typing import Any, Dict, List
 
+import httpx
+
+logger = logging.getLogger(__name__)
+
+# Retryable HTTP status codes (server errors + rate limiting)
+_RETRYABLE_STATUS = {429, 500, 502, 503, 504}
+
+LLM_MAX_RETRIES = 2
+LLM_RETRY_BASE_DELAY = 1.0  # seconds
+
 
 class LLMProvider:
     async def generate(self, messages: List[Dict[str, Any]]) -> Dict[str, Any]:
+        last_exc: Exception | None = None
+        for attempt in range(1 + LLM_MAX_RETRIES):
+            try:
+                return await self._generate(messages)
+            except httpx.HTTPStatusError as exc:
+                last_exc = exc
+                if exc.response.status_code not in _RETRYABLE_STATUS:
+                    raise
+                if attempt < LLM_MAX_RETRIES:
+                    delay = LLM_RETRY_BASE_DELAY * (2**attempt)
+                    logger.warning(
+                        "LLM request failed (%s), retrying in %.1fs (attempt %d/%d)",
+                        exc.response.status_code,
+                        delay,
+                        attempt + 1,
+                        LLM_MAX_RETRIES,
+                    )
+                    await asyncio.sleep(delay)
+            except (httpx.TimeoutException, httpx.ConnectError) as exc:
+                last_exc = exc
+                if attempt < LLM_MAX_RETRIES:
+                    delay = LLM_RETRY_BASE_DELAY * (2**attempt)
+                    logger.warning(
+                        "LLM request failed (%s), retrying in %.1fs (attempt %d/%d)",
+                        type(exc).__name__,
+                        delay,
+                        attempt + 1,
+                        LLM_MAX_RETRIES,
+                    )
+                    await asyncio.sleep(delay)
+        raise last_exc  # type: ignore[misc]
+
+    async def _generate(self, messages: List[Dict[str, Any]]) -> Dict[str, Any]:
         raise NotImplementedError

app/llm/providers/anthropic.py

@@ -20,7 +20,7 @@ def __init__(self, api_key: str, model: str, base_url: str | None, timeout: floa
         self.base_url = (base_url or "https://api.anthropic.com").rstrip("/")
         self.timeout = timeout
 
-    async def generate(self, messages: List[Dict[str, Any]]) -> Dict[str, Any]:
+    async def _generate(self, messages: List[Dict[str, Any]]) -> Dict[str, Any]:
         if not self.api_key:
             raise ValueError("LLM_API_KEY is required for Anthropic provider")
 

app/llm/providers/chat_completions.py

@@ -18,7 +18,7 @@ def __init__(self, api_key: str, model: str, base_url: str | None, timeout: floa
         self.base_url = (base_url or "https://api.openai.com").rstrip("/")
         self.timeout = timeout
 
-    async def generate(self, messages: List[Dict[str, Any]]) -> Dict[str, Any]:
+    async def _generate(self, messages: List[Dict[str, Any]]) -> Dict[str, Any]:
         url = f"{self.base_url}/v1/chat/completions"
         payload: Dict[str, Any] = {
             "model": self.model,

app/llm/providers/ollama.py

@@ -11,7 +11,7 @@ def __init__(self, model: str, base_url: str | None, timeout: float = 60) -> Non
         self.base_url = (base_url or "http://localhost:11434").rstrip("/")
         self.timeout = timeout
 
-    async def generate(self, messages: List[Dict[str, Any]]) -> Dict[str, Any]:
+    async def _generate(self, messages: List[Dict[str, Any]]) -> Dict[str, Any]:
         url = f"{self.base_url}/api/chat"
         payload = {
             "model": self.model,

tests/test_llm.py

@@ -0,0 +1,158 @@
+"""Tests for LLM provider selection and retry logic."""
+
+from unittest.mock import patch
+
+import httpx
+import pytest
+from app.assistant.service import _build_provider
+from app.config import Config
+from app.llm.base import LLM_MAX_RETRIES, LLMProvider
+from app.llm.providers.anthropic import AnthropicProvider
+from app.llm.providers.chat_completions import ChatCompletionsProvider
+from app.llm.providers.ollama import OllamaProvider
+
+pytestmark = pytest.mark.asyncio
+
+
+def _make_config(**overrides) -> Config:
+    defaults = dict(
+        db_url="sqlite:///:memory:",
+        mode="read-only",
+        limit_default=100,
+        timeout_ms=5000,
+        enable_ui=True,
+        enable_explanations=True,
+        allowed_origins=["http://localhost:8000"],
+        allow_destructive=False,
+        llm_provider="openai",
+        llm_api_key="test-key",
+        llm_model="test-model",
+        llm_base_url=None,
+        openai_api_mode="chat",
+        llm_timeout_ms=60000,
+        chat_history_enabled=True,
+        chat_history_limit=10,
+    )
+    defaults.update(overrides)
+    return Config(**defaults)
+
+
+# ── Provider selection ─────────────────────────────────────────────────────
+
+
+def test_build_provider_openai() -> None:
+    cfg = _make_config(llm_provider="openai")
+    provider = _build_provider(cfg)
+    assert isinstance(provider, ChatCompletionsProvider)
+    assert provider.base_url == "https://api.openai.com"
+
+
+def test_build_provider_anthropic() -> None:
+    cfg = _make_config(llm_provider="anthropic")
+    provider = _build_provider(cfg)
+    assert isinstance(provider, AnthropicProvider)
+
+
+def test_build_provider_ollama() -> None:
+    cfg = _make_config(llm_provider="ollama")
+    provider = _build_provider(cfg)
+    assert isinstance(provider, OllamaProvider)
+    assert provider.base_url == "http://localhost:11434"
+
+
+def test_build_provider_deepseek() -> None:
+    cfg = _make_config(llm_provider="deepseek")
+    provider = _build_provider(cfg)
+    assert isinstance(provider, ChatCompletionsProvider)
+    assert provider.base_url == "https://api.deepseek.com"
+
+
+def test_build_provider_gemini() -> None:
+    cfg = _make_config(llm_provider="gemini")
+    provider = _build_provider(cfg)
+    assert isinstance(provider, ChatCompletionsProvider)
+    assert "generativelanguage" in provider.base_url
+
+
+def test_build_provider_custom_base_url() -> None:
+    cfg = _make_config(llm_provider="openai", llm_base_url="https://my-proxy.example.com")
+    provider = _build_provider(cfg)
+    assert isinstance(provider, ChatCompletionsProvider)
+    assert provider.base_url == "https://my-proxy.example.com"
+
+
+def test_build_provider_timeout_passed() -> None:
+    cfg = _make_config(llm_timeout_ms=30000)
+    provider = _build_provider(cfg)
+    assert provider.timeout == 30.0
+
+
+# ── Retry logic ────────────────────────────────────────────────────────────
+
+
+class _FlakyProvider(LLMProvider):
+    """Provider that fails N times then succeeds."""
+
+    def __init__(self, fail_times: int, exc: Exception) -> None:
+        self.fail_times = fail_times
+        self.exc = exc
+        self.attempts = 0
+
+    async def _generate(self, messages):
+        self.attempts += 1
+        if self.attempts <= self.fail_times:
+            raise self.exc
+        return {"text": "ok", "raw": {}}
+
+
+def _make_http_error(status: int) -> httpx.HTTPStatusError:
+    response = httpx.Response(status_code=status)
+    return httpx.HTTPStatusError(
+        message=f"{status}", request=httpx.Request("POST", "http://x"), response=response
+    )
+
+
+async def test_retry_on_500() -> None:
+    provider = _FlakyProvider(fail_times=1, exc=_make_http_error(500))
+    with patch("app.llm.base.LLM_RETRY_BASE_DELAY", 0):
+        result = await provider.generate([])
+    assert result["text"] == "ok"
+    assert provider.attempts == 2
+
+
+async def test_retry_on_429() -> None:
+    provider = _FlakyProvider(fail_times=1, exc=_make_http_error(429))
+    with patch("app.llm.base.LLM_RETRY_BASE_DELAY", 0):
+        result = await provider.generate([])
+    assert result["text"] == "ok"
+    assert provider.attempts == 2
+
+
+async def test_no_retry_on_400() -> None:
+    provider = _FlakyProvider(fail_times=1, exc=_make_http_error(400))
+    with pytest.raises(httpx.HTTPStatusError):
+        await provider.generate([])
+    assert provider.attempts == 1  # no retry for client errors
+
+
+async def test_retry_on_timeout() -> None:
+    provider = _FlakyProvider(fail_times=1, exc=httpx.TimeoutException("timeout"))
+    with patch("app.llm.base.LLM_RETRY_BASE_DELAY", 0):
+        result = await provider.generate([])
+    assert result["text"] == "ok"
+    assert provider.attempts == 2
+
+
+async def test_retry_on_connect_error() -> None:
+    provider = _FlakyProvider(fail_times=1, exc=httpx.ConnectError("refused"))
+    with patch("app.llm.base.LLM_RETRY_BASE_DELAY", 0):
+        result = await provider.generate([])
+    assert result["text"] == "ok"
+
+
+async def test_retry_exhausted_raises() -> None:
+    provider = _FlakyProvider(fail_times=10, exc=_make_http_error(503))
+    with patch("app.llm.base.LLM_RETRY_BASE_DELAY", 0):
+        with pytest.raises(httpx.HTTPStatusError):
+            await provider.generate([])
+    assert provider.attempts == 1 + LLM_MAX_RETRIES

tests/test_sessions.py

@@ -0,0 +1,108 @@
+"""Tests for SSE session management (create, get, expire, GC)."""
+
+import asyncio
+import time
+from unittest.mock import patch
+
+import pytest
+from app import main as main_module
+from app.main import (
+    SESSION_TTL,
+    _create_session,
+    _enqueue,
+    _get_session,
+    _remove_session,
+    _sessions,
+    _sessions_lock,
+)
+
+pytestmark = pytest.mark.asyncio
+
+
+async def _clear_sessions() -> None:
+    async with _sessions_lock:
+        _sessions.clear()
+
+
+async def test_create_and_get_session() -> None:
+    await _clear_sessions()
+    sid = await _create_session()
+    assert sid
+    queue = await _get_session(sid)
+    assert queue is not None
+
+
+async def test_get_nonexistent_session() -> None:
+    await _clear_sessions()
+    queue = await _get_session("nonexistent-id")
+    assert queue is None
+
+
+async def test_remove_session() -> None:
+    await _clear_sessions()
+    sid = await _create_session()
+    await _remove_session(sid)
+    queue = await _get_session(sid)
+    assert queue is None
+
+
+async def test_remove_nonexistent_session() -> None:
+    await _clear_sessions()
+    # Should not raise
+    await _remove_session("nonexistent-id")
+
+
+async def test_enqueue_and_dequeue() -> None:
+    await _clear_sessions()
+    sid = await _create_session()
+    ok = await _enqueue(sid, {"test": "payload"})
+    assert ok
+    queue = await _get_session(sid)
+    message = queue.get_nowait()
+    assert '"test"' in message
+
+
+async def test_enqueue_nonexistent_session() -> None:
+    await _clear_sessions()
+    ok = await _enqueue("nonexistent-id", {"test": "payload"})
+    assert not ok
+
+
+async def test_get_session_updates_last_seen() -> None:
+    await _clear_sessions()
+    sid = await _create_session()
+    async with _sessions_lock:
+        _, ts1 = _sessions[sid]
+    await asyncio.sleep(0.01)
+    await _get_session(sid)
+    async with _sessions_lock:
+        _, ts2 = _sessions[sid]
+    assert ts2 >= ts1
+
+
+async def test_gc_removes_expired_sessions() -> None:
+    await _clear_sessions()
+    sid = await _create_session()
+
+    # Manually set last_seen to the past
+    async with _sessions_lock:
+        queue, _ = _sessions[sid]
+        _sessions[sid] = (queue, time.time() - SESSION_TTL - 10)
+
+    # Run one GC cycle (patch sleep to break after first iteration)
+    call_count = 0
+
+    async def _fake_sleep(seconds):
+        nonlocal call_count
+        call_count += 1
+        if call_count > 1:
+            raise asyncio.CancelledError
+
+    with patch("app.main.asyncio.sleep", side_effect=_fake_sleep):
+        try:
+            await main_module._gc_sessions()
+        except asyncio.CancelledError:
+            pass
+
+    queue = await _get_session(sid)
+    assert queue is None  # expired session was removed