feat(rag): retry a failed LLM request before the first token lands — BOR_LLM_RETRIES/BOR_LLM_RETRY_DELAY with a live 'retrying' status

tests/unit/test_llm_client.py

@@ -10,6 +10,7 @@ from __future__ import annotations
 
 import asyncio
 import json
+from collections.abc import AsyncGenerator
 from types import SimpleNamespace
 from typing import Any, cast
 
@@ -21,8 +22,10 @@ from app.rag.llm import (
     EmbeddingError,
     LLMClient,
     LLMError,
+    RetryPiece,
     StreamPiece,
     ToolCallPiece,
+    chat_stream_retried,
 )
 
 
@@ -800,3 +803,274 @@ def test_chat_whitespace_only_content_raises_llm_error() -> None:
     llm, _ = _make_chat_client(_FakeCompletion("   \n\t "))
     with pytest.raises(LLMError, match="empty content"):
         asyncio.run(llm.chat([{"role": "user", "content": "q"}]))
+
+
+# ---------- chat_stream_retried (phase 67, task 01) ----------
+
+_RETRY_MSGS: list[dict[str, str]] = [{"role": "user", "content": "q"}]
+
+
+class _ScriptedClient(LLMClient):
+    """An LLMClient whose ``chat_stream`` is scripted per attempt — no
+    endpoint. ``attempts`` scripts the Nth call: ``(pieces, error)`` — the
+    stream yields *pieces*, then raises *error* if not None (an error after
+    zero pieces = "the endpoint died before the first token"; after some
+    pieces = a mid-stream drop). Records every request's messages/tools
+    and every attempt's stream teardown (the phase-48 close analog)."""
+
+    def __init__(
+        self, attempts: list[tuple[list[StreamPiece | ToolCallPiece], Exception | None]]
+    ) -> None:
+        super().__init__(_settings())
+        self.attempts = list(attempts)
+        self.request_args: list[
+            tuple[list[dict[str, str]], list[dict[str, Any]] | None]
+        ] = []
+        #: Indices of attempts whose stream teardown has run.
+        self.closed: list[int] = []
+
+    def chat_stream(
+        self,
+        messages: list[dict[str, str]],
+        tools: list[dict[str, Any]] | None = None,
+    ) -> AsyncGenerator[StreamPiece | ToolCallPiece, None]:
+        index = len(self.request_args)
+        pieces, error = (
+            self.attempts[index]
+            if index < len(self.attempts)
+            else ([], LLMError("script exhausted"))
+        )
+        self.request_args.append(
+            (list(messages), list(tools) if tools is not None else None)
+        )
+        return self._attempt(index, pieces, error)
+
+    async def _attempt(
+        self, index: int, pieces: list[StreamPiece | ToolCallPiece],
+        error: Exception | None,
+    ) -> AsyncGenerator[StreamPiece | ToolCallPiece, None]:
+        try:
+            for piece in pieces:
+                yield piece
+            if error is not None:
+                raise error
+        finally:
+            self.closed.append(index)
+
+
+def _record_sleeps(monkeypatch: pytest.MonkeyPatch) -> list[float]:
+    """Monkeypatch ``asyncio.sleep`` (what ``chat_stream_retried`` awaits
+    for the flat pre-retry delay) and record every awaited delay."""
+    sleeps: list[float] = []
+
+    async def fake_sleep(seconds: float) -> None:
+        sleeps.append(seconds)
+
+    monkeypatch.setattr(asyncio, "sleep", fake_sleep)
+    return sleeps
+
+
+def _collect_retried(
+    client: _ScriptedClient,
+    messages: list[dict[str, str]],
+    *,
+    tools: list[dict[str, Any]] | None = None,
+    retries: int,
+    delay: float,
+) -> list[StreamPiece | ToolCallPiece | RetryPiece]:
+    async def run() -> list[StreamPiece | ToolCallPiece | RetryPiece]:
+        return [
+            p
+            async for p in chat_stream_retried(
+                client, messages, tools=tools, retries=retries, delay=delay
+            )
+        ]
+
+    return asyncio.run(run())
+
+
+def test_retried_retries_a_dead_attempt_before_the_first_piece(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Failure on attempt 1, success on attempt 2 → [RetryPiece(2, N)]
+    (the attempt about to be tried, 1-based) then the answer pieces; the
+    request is restarted byte-identical and the flat delay is awaited
+    exactly once."""
+    answer: list[StreamPiece | ToolCallPiece] = [
+        StreamPiece("content", "A "),
+        StreamPiece("content", "B"),
+    ]
+    client = _ScriptedClient([([], LLMError("connection refused")), (answer, None)])
+    sleeps = _record_sleeps(monkeypatch)
+    pieces = _collect_retried(client, _RETRY_MSGS, retries=3, delay=2.5)
+    assert pieces == [RetryPiece(2, 4), *answer]
+    assert len(client.request_args) == 2
+    # The restart is byte-identical: same messages, same (absent) tools.
+    assert client.request_args[0] == client.request_args[1]
+    assert client.request_args[0][1] is None
+    assert sleeps == [2.5]
+
+
+def test_retried_exhaustion_yields_all_retries_then_raises(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """retries=2 → 3 attempts; each pre-first-piece failure yields a
+    RetryPiece naming the attempt about to be tried (attempts 2 and 3 of
+    3), the final failure raises the terminal LLMError, and no sleep
+    follows the last attempt."""
+    client = _ScriptedClient(
+        [([], LLMError("down 1")), ([], LLMError("down 2")), ([], LLMError("down 3"))]
+    )
+    sleeps = _record_sleeps(monkeypatch)
+
+    async def run() -> list[RetryPiece]:
+        out: list[RetryPiece] = []
+        with pytest.raises(LLMError, match="down 3"):
+            async for p in chat_stream_retried(
+                client, _RETRY_MSGS, retries=2, delay=0.5
+            ):
+                assert isinstance(p, RetryPiece)
+                out.append(p)
+        return out
+
+    out = asyncio.run(run())
+    assert out == [RetryPiece(2, 3), RetryPiece(3, 3)]
+    assert len(client.request_args) == 3
+    assert sleeps == [0.5, 0.5]
+
+
+def test_retried_failure_after_first_piece_is_terminal(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Locked A2: a piece has already flowed → the LLMError is re-raised
+    immediately — no RetryPiece, no sleep, no second call (a partial
+    answer is never redone)."""
+    client = _ScriptedClient(
+        [([StreamPiece("content", "partial ")], LLMError("mid-stream drop"))],
+    )
+    sleeps = _record_sleeps(monkeypatch)
+
+    async def run() -> list[StreamPiece | ToolCallPiece | RetryPiece]:
+        out: list[StreamPiece | ToolCallPiece | RetryPiece] = []
+        with pytest.raises(LLMError, match="mid-stream drop"):
+            async for p in chat_stream_retried(
+                client, _RETRY_MSGS, retries=3, delay=5.0
+            ):
+                out.append(p)
+        return out
+
+    out = asyncio.run(run())
+    assert out == [StreamPiece("content", "partial ")]
+    assert not any(isinstance(p, RetryPiece) for p in out)
+    assert len(client.request_args) == 1
+    assert sleeps == []
+
+
+def test_retried_zero_retries_is_one_attempt_no_retry_pieces(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """The kill-switch path (retries=0): exactly one attempt, the error
+    propagates, no RetryPiece, no sleep — the pre-phase-67 behavior."""
+    client = _ScriptedClient([([], LLMError("connection refused"))])
+    sleeps = _record_sleeps(monkeypatch)
+    with pytest.raises(LLMError, match="connection refused"):
+        _collect_retried(client, _RETRY_MSGS, retries=0, delay=5.0)
+    assert len(client.request_args) == 1
+    assert sleeps == []
+
+
+def test_retried_healthy_stream_is_untouched(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """No failure → exactly one attempt, every piece kind (thinking / tool
+    call / content) passes through unchanged, no RetryPiece, no sleep —
+    a healthy turn is byte-identical to the plain chat_stream."""
+    answer = [
+        StreamPiece("thinking", "hmm"),
+        ToolCallPiece(id="call_1", name="list_documents", arguments={}),
+        StreamPiece("content", "Talos."),
+    ]
+    client = _ScriptedClient([(answer, None)])
+    sleeps = _record_sleeps(monkeypatch)
+    tools = [
+        {
+            "type": "function",
+            "function": {"name": "list_documents", "parameters": {}},
+        }
+    ]
+    pieces = _collect_retried(
+        client, _RETRY_MSGS, tools=tools, retries=3, delay=5.0
+    )
+    assert pieces == answer
+    assert client.request_args == [(_RETRY_MSGS, tools)]
+    assert sleeps == []
+
+
+def test_retried_zero_delay_still_notifies(monkeypatch: pytest.MonkeyPatch) -> None:
+    """The e2e fast path (BOR_LLM_RETRY_DELAY=0): the RetryPiece is still
+    emitted and the (zero) sleep is still awaited."""
+    client = _ScriptedClient(
+        [([], LLMError("down")), ([StreamPiece("content", "ok")], None)]
+    )
+    sleeps = _record_sleeps(monkeypatch)
+    pieces = _collect_retried(client, _RETRY_MSGS, retries=1, delay=0)
+    assert pieces == [RetryPiece(2, 2), StreamPiece("content", "ok")]
+    assert sleeps == [0]
+
+
+def test_retried_abandon_mid_attempt_closes_the_attempt_stream() -> None:
+    """Consumer abandon at a mid-attempt piece (the stop-generation path,
+    phase 48): no exception leaks and the attempt's stream is torn down
+    through the wrapper's explicit close."""
+    client = _ScriptedClient(
+        [([StreamPiece("content", "A "), StreamPiece("content", "B ")], None)]
+    )
+
+    async def run() -> None:
+        gen = chat_stream_retried(client, _RETRY_MSGS, retries=3, delay=1.0)
+        first = await gen.__anext__()
+        assert first == StreamPiece("content", "A ")
+        await gen.aclose()  # the consumer stops after the first piece
+
+    asyncio.run(run())
+    assert client.closed == [0]  # attempt 1's stream was closed
+    assert len(client.request_args) == 1  # no second attempt
+
+
+def test_retried_abandon_during_retry_sleep_leaks_nothing(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Consumer abandon while the generator is parked in the pre-retry
+    sleep (client disconnect): the driving task is cancelled cleanly,
+    the phase-48 teardown ``aclose()`` on the generator does not raise,
+    and the retry never starts."""
+    entered = asyncio.Event()
+
+    async def parking_sleep(seconds: float) -> None:
+        entered.set()
+        await asyncio.Event().wait()  # park until the abandon arrives
+
+    monkeypatch.setattr(asyncio, "sleep", parking_sleep)
+    client = _ScriptedClient(
+        [([], LLMError("endpoint down")), ([StreamPiece("content", "never")], None)]
+    )
+
+    async def run() -> None:
+        gen = chat_stream_retried(client, _RETRY_MSGS, retries=3, delay=1.5)
+
+        async def consumer() -> list[StreamPiece | ToolCallPiece | RetryPiece]:
+            return [p async for p in gen]
+
+        task = asyncio.ensure_future(consumer())
+        await entered.wait()  # the generator is inside the pre-retry sleep
+        assert not task.done()
+        task.cancel()  # client disconnect: the driving task is cancelled
+        with pytest.raises(asyncio.CancelledError):
+            await task
+        # Production teardown (phase 48 pattern): the endpoint's finally
+        # closes the stream generator — must not raise.
+        await gen.aclose()
+
+    asyncio.run(run())
+    assert len(client.request_args) == 1  # the retry never started
+    assert client.closed == [0]  # attempt 1's stream was torn down