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_agent.py

@@ -8,15 +8,21 @@ removed the per-tool budgets, the assistant/tool message history), the
 kill switch (``agent_max_rounds=0`` single-call path), the round cap
 forcing a final no-tools answer (an always-calling stream and an
 always-rejected stream), re-lists and multi-reads executing without
-budgets, dedupe, unknown tool / missing args / unknown path, and the
-``<tools>`` prompt section (HIGH only).
+budgets, dedupe, unknown tool / missing args / unknown path, the
+``<tools>`` prompt section (HIGH only), and the phase-67 per-round
+retries (a dead-then-recovered round restarts before its first piece
+with a ``RetryPiece``; a mid-stream drop stays terminal — locked A2;
+the forced final no-tools call retries too; ``llm_retries=0`` is one
+plain attempt; retries are invisible to the round cap; consumer abandon
+mid-retry-sleep leaks nothing).
 """
 from __future__ import annotations
 
 import asyncio
 import json
+import logging
 import uuid
-from collections.abc import AsyncIterator
+from collections.abc import AsyncGenerator, AsyncIterator
 from copy import deepcopy
 from typing import Any, cast
 
@@ -31,7 +37,7 @@ from app.rag.agent import (
     AgentHolder,
     run_agent,
 )
-from app.rag.llm import LLMClient, StreamPiece, ToolCallPiece
+from app.rag.llm import LLMClient, LLMError, RetryPiece, StreamPiece, ToolCallPiece
 from app.rag.prompts import TOOLS_SECTION, _base, build_deflect_prompt, build_high_prompt
 
 
@@ -73,12 +79,12 @@ class ScriptedLLM:
 
 
 async def _run(
-    llm: ScriptedLLM,
+    llm: ScriptedLLM | FailingLLM,
     holder: AgentHolder,
     settings: Settings,
     seed_docs: list[Document] | None = None,
-) -> list[StreamPiece | ToolCallPiece]:
-    out: list[StreamPiece | ToolCallPiece] = []
+) -> list[StreamPiece | ToolCallPiece | RetryPiece]:
+    out: list[StreamPiece | ToolCallPiece | RetryPiece] = []
     async for piece in run_agent(
         cast("LLMClient", llm),
         cast("Session", None),
@@ -541,6 +547,300 @@ def test_read_document_missing_arguments_refused(
     assert llm.requests[1][1] == AGENT_TOOLS
 
 
+# ---------- retries inside the agent loop (phase 67, locked A2) ----------
+
+
+class FailingLLM:
+    """A scripted fake whose Nth ``chat_stream`` call yields pieces and
+    then raises (phase 67): ``attempts`` is a list of ``(pieces, error)``
+    — 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 the indices of the attempts whose stream
+    teardown ran (``closed``)."""
+
+    def __init__(
+        self,
+        attempts: list[tuple[list[StreamPiece | ToolCallPiece], Exception | None]],
+    ) -> None:
+        self.attempts = list(attempts)
+        self.requests: list[tuple[list[dict[str, Any]], 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,
+    ) -> AsyncIterator[StreamPiece | ToolCallPiece]:
+        index = len(self.requests)
+        pieces, error = (
+            self.attempts[index]
+            if index < len(self.attempts)
+            else ([], LLMError("script exhausted"))
+        )
+        self.requests.append(
+            (deepcopy(messages), deepcopy(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,
+    ) -> AsyncIterator[StreamPiece | ToolCallPiece]:
+        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 test_round_retried_before_first_piece(
+    monkeypatch: pytest.MonkeyPatch,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """A tool round that dies before its first piece is restarted with the
+    same messages: the stream carries a RetryPiece BEFORE the tool call,
+    the tool executes, the final answer streams, and the per-call log line
+    is still emitted exactly once (retries are invisible to the loop)."""
+    monkeypatch.setattr(agent, "list_catalog", lambda db: [("S", "a.md", "A")])
+    holder = AgentHolder()
+    llm = FailingLLM(
+        [
+            ([], LLMError("connection refused")),
+            ([ToolCallPiece(id="call_1", name="list_documents", arguments={})], None),
+            ([StreamPiece("content", "Done!")], None),
+        ]
+    )
+    sleeps = _record_sleeps(monkeypatch)
+    with caplog.at_level(logging.INFO, logger="app.agent"):
+        pieces = asyncio.run(
+            _run(llm, holder, _settings(agent_max_rounds=2, llm_retry_delay=2.5))
+        )
+    assert pieces == [
+        RetryPiece(2, 4),  # default llm_retries=3 → 4 attempts
+        ToolCallPiece(id="call_1", name="list_documents", arguments={}),
+        StreamPiece("content", "Done!"),
+    ]
+    assert holder.tool_calls == 1
+    assert holder.read_docs == []
+    # The restart is byte-identical: same messages, same tools offered.
+    assert len(llm.requests) == 3
+    assert llm.requests[0] == llm.requests[1]
+    assert llm.requests[0][1] == AGENT_TOOLS
+    assert llm.requests[2][1] == AGENT_TOOLS  # the answer round still offered
+    # The flat delay was awaited exactly once, before the retry.
+    assert sleeps == [2.5]
+    tool_logs = [r for r in caplog.records if r.getMessage().startswith("agent tool=")]
+    assert len(tool_logs) == 1  # the retry did not re-run the tool or log
+    assert tool_logs[0].getMessage() == "agent tool=list_documents args={} round=1/2"
+
+
+def test_round_failure_after_first_piece_is_terminal(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Locked A2: a round that already streamed a piece fails the turn —
+    the LLMError propagates out of ``run_agent``, no RetryPiece, no
+    sleep, no second request, and the holder is untouched (the tool
+    never ran)."""
+    monkeypatch.setattr(agent, "list_catalog", lambda db: [])
+    holder = AgentHolder()
+    llm = FailingLLM(
+        [([StreamPiece("content", "partial ")], LLMError("mid-stream drop"))]
+    )
+    sleeps = _record_sleeps(monkeypatch)
+
+    async def drain() -> list[StreamPiece | ToolCallPiece | RetryPiece]:
+        out: list[StreamPiece | ToolCallPiece | RetryPiece] = []
+        with pytest.raises(LLMError, match="mid-stream drop"):
+            async for piece in run_agent(
+                cast("LLMClient", llm),
+                cast("Session", None),
+                system_prompt="SYSTEM_PROMPT",
+                user_message="QUESTION",
+                seed_docs=[],
+                settings=_settings(),
+                holder=holder,
+            ):
+                out.append(piece)
+        return out
+
+    out = asyncio.run(drain())
+    assert out == [StreamPiece("content", "partial ")]  # no RetryPiece
+    assert len(llm.requests) == 1  # no retry
+    assert sleeps == []
+    assert holder.read_docs == [] and holder.tool_calls == 0
+
+
+def test_forced_final_no_tools_call_is_retried(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """The forced final request (round cap reached) goes through the same
+    retry rule: a failure before its first piece yields a RetryPiece and
+    restarts with ``tools=None``; the answer from the retry streams."""
+    monkeypatch.setattr(agent, "list_catalog", lambda db: [("S", "a.md", "A")])
+    holder = AgentHolder()
+    llm = FailingLLM(
+        [
+            ([ToolCallPiece(id="call_1", name="list_documents", arguments={})], None),
+            ([ToolCallPiece(id="call_2", name="list_documents", arguments={})], None),
+            ([], LLMError("down at the cap")),
+            ([StreamPiece("content", "forced answer")], None),
+        ]
+    )
+    pieces = asyncio.run(_run(llm, holder, _settings(agent_max_rounds=2)))
+    assert [type(p) for p in pieces] == [
+        ToolCallPiece,
+        ToolCallPiece,
+        RetryPiece,
+        StreamPiece,
+    ]
+    assert pieces[2] == RetryPiece(2, 4)
+    assert pieces[3] == StreamPiece("content", "forced answer")
+    assert len(llm.requests) == 4  # 2 tool rounds + the final + its retry
+    # The forced final (and its retry) carry no tools, whatever is left.
+    assert llm.requests[2][1] is None
+    assert llm.requests[3][1] is None
+    # …and the restart is byte-identical.
+    assert llm.requests[2][0] == llm.requests[3][0]
+    assert holder.tool_calls == 2
+
+
+def test_zero_retries_is_one_plain_attempt(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """The kill-switch path (``llm_retries=0``): a dead round raises
+    immediately — one request, no RetryPiece, no sleep (pre-phase-67
+    behavior)."""
+    monkeypatch.setattr(agent, "list_catalog", lambda db: [])
+    holder = AgentHolder()
+    llm = FailingLLM([([], LLMError("connection refused"))])
+    sleeps = _record_sleeps(monkeypatch)
+
+    async def drain() -> list[StreamPiece | ToolCallPiece | RetryPiece]:
+        out: list[StreamPiece | ToolCallPiece | RetryPiece] = []
+        with pytest.raises(LLMError, match="connection refused"):
+            async for piece in run_agent(
+                cast("LLMClient", llm),
+                cast("Session", None),
+                system_prompt="SYSTEM_PROMPT",
+                user_message="QUESTION",
+                seed_docs=[],
+                settings=_settings(llm_retries=0),
+                holder=holder,
+            ):
+                out.append(piece)
+        return out
+
+    out = asyncio.run(drain())
+    assert out == []  # nothing streamed, no RetryPiece
+    assert len(llm.requests) == 1
+    assert sleeps == []
+    assert holder.read_docs == [] and holder.tool_calls == 0
+
+
+def test_abandon_mid_retry_sleep_leaks_nothing(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Consumer abandon while a retried round is parked in the pre-retry
+    sleep (client disconnect): the driving task is cancelled cleanly, the
+    production teardown ``aclose()`` on ``run_agent`` does not raise, the
+    inner attempt's stream was torn down, 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)
+    monkeypatch.setattr(agent, "list_catalog", lambda db: [])
+    holder = AgentHolder()
+    llm = FailingLLM(
+        [([], LLMError("endpoint down")), ([StreamPiece("content", "never")], None)]
+    )
+
+    async def run() -> None:
+        gen = run_agent(
+            cast("LLMClient", llm),
+            cast("Session", None),
+            system_prompt="SYSTEM_PROMPT",
+            user_message="QUESTION",
+            seed_docs=[],
+            settings=_settings(),
+            holder=holder,
+        )
+
+        async def consumer() -> list[StreamPiece | ToolCallPiece | RetryPiece]:
+            return [p async for p in gen]
+
+        task = asyncio.ensure_future(consumer())
+        await entered.wait()  # the round's retry is parked in the 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): must not raise. ``run_agent``
+        # is an async generator despite its AsyncIterator annotation.
+        await cast(
+            "AsyncGenerator[StreamPiece | ToolCallPiece | RetryPiece, None]", gen
+        ).aclose()
+
+    asyncio.run(run())
+    assert len(llm.requests) == 1  # the retry never started
+    assert llm.closed == [0]  # attempt 1's inner stream was torn down
+    assert holder.read_docs == [] and holder.tool_calls == 0
+
+
+def test_retries_are_invisible_to_the_round_cap(
+    monkeypatch: pytest.MonkeyPatch,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """A failing-then-succeeding round consumes ONE round: with a cap of
+    2, the retried first round and the second tool round fill the cap —
+    the forced final follows the SECOND call, and the log lines read
+    round=1/2 and round=2/2."""
+    monkeypatch.setattr(agent, "list_catalog", lambda db: [("S", "a.md", "A")])
+    holder = AgentHolder()
+    llm = FailingLLM(
+        [
+            ([], LLMError("down")),
+            ([ToolCallPiece(id="call_1", name="list_documents", arguments={})], None),
+            ([ToolCallPiece(id="call_2", name="list_documents", arguments={})], None),
+            ([StreamPiece("content", "forced answer")], None),
+        ]
+    )
+    with caplog.at_level(logging.INFO, logger="app.agent"):
+        pieces = asyncio.run(_run(llm, holder, _settings(agent_max_rounds=2)))
+    assert [type(p) for p in pieces] == [
+        RetryPiece,
+        ToolCallPiece,
+        ToolCallPiece,
+        StreamPiece,
+    ]
+    assert len(llm.requests) == 4  # 2 (round 1 + its retry) + 1 + the forced final
+    assert llm.requests[3][1] is None  # the forced final, after round 2
+    assert holder.tool_calls == 2
+    msgs = [r.getMessage() for r in caplog.records]
+    assert "agent tool=list_documents args={} round=1/2" in msgs
+    assert "agent tool=list_documents args={} round=2/2" in msgs
+    assert any("round cap reached (rounds=2)" in m for m in msgs)
+
+
 # ---------- prompts: <tools> section (HIGH only) ----------