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

This commit is contained in:
2026-09-02 10:52:38 -04:00
parent f04ddbe1f8
commit 88293ed02f
44 changed files with 2488 additions and 56 deletions
+9
View File
@@ -98,6 +98,15 @@ def app_server(mock_llm: int) -> Iterator[str]:
# The production default stays 0.62 (re-tuned against the real
# `embed` model's 0.41–0.84 cosine range, PLAN A8).
env["BOR_RELEVANCE_THRESHOLD"] = "0.30"
# Phase 67 (LLM retry): the e2e pins the retry MECHANISM with instant
# waits (BOR_LLM_RETRY_DELAY=0 — the 5 s default is unit-pinned via
# tests/unit/test_config.py). BOR_LLM_RETRIES is forced to the code
# default (derived from the class field, never drifts from
# app/config.py) so the exhaustion test relies on the REAL budget and
# an operator's local (gitignored) .env cannot leak a different one
# into the app under test (the phase-61 leak-guard pattern).
env["BOR_LLM_RETRY_DELAY"] = "0"
env["BOR_LLM_RETRIES"] = str(_Settings.model_fields["llm_retries"].default)
env.setdefault("BOR_DATABASE_URL", "postgresql+psycopg://reese:reese@localhost:5432/brain_of_reese")
# Phase 16: admin auth must be set or create_app() refuses to boot.
env["BOR_ADMIN_PASSWORD"] = ADMIN_PASSWORD
+123 -2
View File
@@ -112,6 +112,36 @@ Implements just enough of the aipi surface:
answer; the E2E asks it against an on-topic fixture (HIGH gate) and
asserts non-deflection.
Failure injection (phase 67, LLM retry, TODO.md L3) — deterministic
dead-endpoint behavior for the retry E2E suite (``tests/e2e/
test_llm_retry.py``). The mock is single-conversation per e2e server, so
the sequences are driven by module-level counters that reset per
trigger phrase after the success they guard (a second question with
the same trigger re-drives the sequence from zero):
- user message containing ``fail then answer`` (``RETRY_TRIGGER``):
the first ``RETRY_DEAD_ATTEMPTS`` (2) app-level streaming attempts
respond 500 (JSON body, like a dead proxy) and the third streams
the normal composed answer — 2 = 1 original attempt + 1 retry under
the default ``BOR_LLM_RETRIES=3``, so a suite exercises a real
retry without waiting for the 4-attempt exhaustion. Counted in
APP-LEVEL attempts, not raw HTTP POSTs: while the endpoint stays
dead, the openai SDK's default policy (max_retries=2 — the app's
``LLMClient`` keeps it) re-POSTs a 500'd streaming request twice
before surfacing the error, so each dead attempt costs exactly 3
POSTs (``_HTTPS_PER_DEAD_ATTEMPT``).
- user message containing ``always fail``
(``ALWAYS_FAIL_TRIGGER``): EVERY streaming chat/completions request
responds 500 — the retry-budget exhaustion path (the terminal
error banner in the UI).
- embeddings request whose input contains ``embed fail once``
(``EMBED_FAIL_TRIGGER``): the FIRST such request responds 500, the
next returns the normal bag-of-words vector — the endpoint's
pre-stream embedding retry loop. Raw httpx on the client side (no
SDK-level retries), so one POST per attempt: the counter is per
POST here, unlike the chat counter above.
Non-streaming requests (document summaries, KB overview) never 500 —
the retry scope is the chat turn only (owner-locked A1).
``max_tokens`` is honored deterministically (token ≈ whitespace word),
like a real endpoint: an answer longer than the cap is truncated. This
is what makes the phase-11 truncation regression observable.
@@ -129,7 +159,7 @@ import uuid
from typing import Any
from fastapi import FastAPI
from fastapi.responses import StreamingResponse
from fastapi.responses import JSONResponse, StreamingResponse
app = FastAPI()
@@ -265,6 +295,76 @@ TABLE_ANSWER = (
"| value-one | value-two | value-three | value-four | value-five |"
)
# ---------------------------------------------------------------------------
# Phase 67 (LLM retry, TODO.md L3): deterministic failure injection
# ---------------------------------------------------------------------------
#: A user message containing this substring (case-insensitive) gets
#: ``RETRY_DEAD_ATTEMPTS`` dead streaming attempts (500, JSON body) before
#: the normal composed answer streams — 1 original attempt + 1 retry under
#: the default ``BOR_LLM_RETRIES=3`` (see the module docstring).
RETRY_TRIGGER = "fail then answer"
#: App-level attempts the endpoint stays dead for before the answer.
RETRY_DEAD_ATTEMPTS = 2
#: A user message containing this substring (case-insensitive) makes
#: EVERY streaming chat/completions request respond 500 — the
#: retry-budget exhaustion path (the terminal error banner in the UI).
ALWAYS_FAIL_TRIGGER = "always fail"
#: An embeddings request whose input contains this substring
#: (case-insensitive) 500s on its FIRST POST; the next returns the normal
#: bag-of-words vector — the endpoint's pre-stream embedding retry loop.
EMBED_FAIL_TRIGGER = "embed fail once"
#: One DEAD app-level chat attempt costs exactly this many HTTP POSTs
#: while the endpoint stays down: the openai SDK's default policy
#: (max_retries=2 — the app's ``LLMClient`` keeps it) re-POSTs a 500'd
#: streaming request twice before surfacing the error to
#: ``chat_stream_retried``. The failure counters below therefore count
#: app-level attempts (groups of this size), not raw POSTs — the visible
#: sequence (one SSE ``retry`` frame after each dead attempt, the answer
#: on the third) stays deterministic regardless of the SDK's internal
#: backoff pacing.
_HTTPS_PER_DEAD_ATTEMPT = 3
#: Module-level failure counters — the mock is single-conversation per
#: e2e server. Keyed by trigger phrase (reset per trigger): the number
#: of matching POSTs served so far. Each sequence resets after the
#: success it guards, so a second question carrying the same trigger
#: re-drives the failure sequence from zero.
_fail_posts: dict[str, int] = {}
def _llm_500(why: str) -> JSONResponse:
"""A dead-proxy 500 with a JSON error body (phase 67 injection)."""
return JSONResponse(
status_code=500,
content={
"error": {
"message": f"upstream connection reset ({why})",
"type": "proxy_error",
}
},
)
def _bump_fail(key: str) -> int:
n = _fail_posts.get(key, 0) + 1
_fail_posts[key] = n
return n
def _chat_dead(key: str, dead_attempts: int) -> bool:
"""Bump *key*'s counter; True while the endpoint stays dead.
Counted in app-level attempts (see ``_HTTPS_PER_DEAD_ATTEMPT``): the
first ``dead_attempts * _HTTPS_PER_DEAD_ATTEMPT`` POSTs 500 and the
next attempt's first POST streams (the caller resets the counter on
the success).
"""
return _bump_fail(key) <= dead_attempts * _HTTPS_PER_DEAD_ATTEMPT
#: The agent's ``read_document`` tool-result prefix (app.rag.agent
#: ``_execute_tool``): ``"Document <source/path>:\n<content>"``.
@@ -656,11 +756,22 @@ def models() -> dict[str, Any]:
@app.post("/v1/embeddings")
def embeddings(body: dict[str, Any]) -> dict[str, Any]:
def embeddings(body: dict[str, Any]) -> Any: # dict, or a 500 (phase 67)
raw = body.get("input")
if isinstance(raw, str):
raw = [raw]
inputs: list[Any] = list(raw) if isinstance(raw, list) else []
# Phase 67 (embedding retry): the first embeddings request whose
# input carries the marker 500s; the next returns the normal
# bag-of-words vector (see the module docstring). Raw httpx on the
# client side — no SDK-level retries — so one POST per app attempt:
# the counter is per POST here (unlike the chat counter below).
joined = " ".join(str(t) for t in inputs if isinstance(t, str)).lower()
if EMBED_FAIL_TRIGGER in joined:
n = _bump_fail(EMBED_FAIL_TRIGGER)
if n == 1:
return _llm_500(EMBED_FAIL_TRIGGER)
_fail_posts[EMBED_FAIL_TRIGGER] = 0 # the vector went out — restart
data = [
{"object": "embedding", "index": i, "embedding": embed_text(t)}
for i, t in enumerate(inputs)
@@ -827,6 +938,16 @@ def chat_completions(body: dict[str, Any]) -> Any:
# the flow handles streaming requests; a non-streaming marker request
# (never issued by the app) falls through to the regular answer.
if body.get("stream"):
# Phase 67 (LLM retry): deterministic failure injection — see
# the module docstring. Checked before the marker tool flow: the
# injection markers never combine with the tool-flow markers in
# any suite, and a dead endpoint answers nothing (no flow).
if ALWAYS_FAIL_TRIGGER in user_lower:
return _llm_500(ALWAYS_FAIL_TRIGGER)
if RETRY_TRIGGER in user_lower:
if _chat_dead(RETRY_TRIGGER, RETRY_DEAD_ATTEMPTS):
return _llm_500(RETRY_TRIGGER)
_fail_posts[RETRY_TRIGGER] = 0 # the answer streamed — restart
flow = _tool_flow(body)
if flow is not None:
if flow[0] == "list":
+506
View File
@@ -0,0 +1,506 @@
"""Phase 67 E2E (Playwright): LLM retry with live "Trying Again" feedback.
Source: ``TODO.md`` L3 — "Add a .env configurable retry in case the LLM
server fails to respond. Allow 3 retries by default, with 5 seconds
between each retry. Update the user interface to show 'communication
interrupted, trying again' … if the LLM server stops communicating."
(TODO-derived phase — no story file.)
Run in isolation (DB must be up: ``podman compose up -d db``):
uv run pytest tests/e2e/test_llm_retry.py -v --no-cov
MOCK-ONLY suite: ``E2E_REAL_LLM=1`` is not supported — the retry gate is
the deterministic failure injection in ``tests/e2e/mock_llm.py``:
* ``fail then answer`` (``RETRY_TRIGGER``): the first 2 app-level
streaming attempts 500 (JSON body, like a dead proxy) and the third
streams the normal composed answer — 1 original attempt + 1 retry
under the default ``BOR_LLM_RETRIES=3``. (Each dead app attempt costs
3 HTTP POSTs — the openai SDK's default 1 + 2 internal retries — so
the mock counts attempts, not POSTs; see the mock docstring.)
* ``always fail`` (``ALWAYS_FAIL_TRIGGER``): every streaming request
500s — the retry-budget exhaustion path.
* ``embed fail once`` (``EMBED_FAIL_TRIGGER``): the first embeddings
request 500s, the next returns the normal vector — the endpoint's
pre-stream embedding retry loop.
The e2e app server boots with ``BOR_LLM_RETRY_DELAY=0`` (conftest) so
the retry waits are instant — the suite pins the MECHANISM; the 5 s
default is unit-pinned via ``tests/unit/test_config.py``.
``BOR_LLM_RETRIES`` is forced to its real default (3, conftest) — the
exhaustion test relies on the real budget: 4 attempts total, so the
last ``retry`` frame reads "(4 of 4)".
KB seed (phase-37 direct-seed pattern): ONE fixture document
(``homelab/kubernetes.md``), one chunk carrying the mock's own
bag-of-words embedding. The marker questions were verified against this
exact seed (``plan_turn``, E2E threshold 0.30):
* the "sourdough" questions (DEFLECT_Q / EXHAUST_Q) share no FTS token
with the document (cosine 0.000, fts 0) → LOW → the DEFLECTED path
(``chat_stream_retried`` directly, no agent round);
* the "kubernetes cluster" questions (GROUNDED_Q / EMBED_Q) FTS-hit the
document (cosine 0.171, fts 1) → HIGH → the GROUNDED path (the agent
loop's per-round retry).
Test → source mapping:
1. ``test_dead_then_recovered_deflected`` — LOW turn: the recorded
``#send-status`` values contain "Communication interrupted —
retrying (2 of 4)…" and "(3 of 4)…" (the owner-locked A4 copy), the
wire carries the two ``retry`` frames ahead of the first delta, the
deflected answer settles, and no error banner appears.
2. ``test_dead_then_recovered_grounded`` — HIGH turn: the same status +
wire assertions; the grounded answer completes with the source chip
(the agent round retried, the turn is intact).
3. ``test_embedding_retry_completes`` — ``embed fail once``: the
pre-stream embedding retry is visible to the UI (a ``retry`` status
before any answer frame) and the turn completes normally.
4. ``test_exhaustion_lands_on_the_error_banner`` — ``always fail``:
after the 4th dead attempt the EXISTING terminal error banner
appears (role=alert, the "dropped the connection" copy), the last
retrying status is the highest attempt — "(4 of 4)…" — and the send
button re-enables (the banner path settles the state machine).
"""
from __future__ import annotations
import hashlib
import json
import re
import time
from datetime import UTC, datetime
from pathlib import Path
from playwright.sync_api import Page, expect
from sqlalchemy import select, text
from sqlalchemy.orm import Session
from app.db import SessionLocal
from app.models import Chunk, Document, QueryLog
from tests.e2e.mock_llm import embed_text
REPO = Path(__file__).resolve().parents[2]
# --------------------------------------------------------------------------
# Seed + questions (see the module docstring for the gate verification)
# --------------------------------------------------------------------------
SEED_SOURCE = "docs"
SEED_PATH = "homelab/kubernetes.md"
SEED_SP = f"{SEED_SOURCE}/{SEED_PATH}"
KUB_CONTENT = (REPO / "tests" / "fixtures" / "docs" / SEED_PATH).read_text()
#: LOW turn (deflected path) + the retry trigger: no FTS overlap with the
#: seeded document, cosine 0.000 → the honesty gate deflects.
DEFLECT_Q = "How do I bake sourdough bread? fail then answer"
#: HIGH turn (grounded path) + the retry trigger: FTS-hits the seeded
#: document → the agent loop runs (and its single round is retried).
GROUNDED_Q = "How is my Kubernetes cluster set up? fail then answer"
#: HIGH turn + the embedding trigger: the pre-stream embedding 500s once.
EMBED_Q = "How is my Kubernetes cluster set up? embed fail once"
#: LOW turn + the exhaustion trigger: every streaming request 500s.
EXHAUST_Q = "How do I bake sourdough bread? always fail"
#: The conftest forces BOR_LLM_RETRIES to its default (3) → 4 attempts
#: total; the attempt math in the assertions is fixed by that budget.
MAX_ATTEMPTS = 4
def _retry_status(attempt: int) -> str:
"""The owner-locked A4 copy for *attempt* (1-based) of MAX_ATTEMPTS."""
return f"Communication interrupted — retrying ({attempt} of {MAX_ATTEMPTS})…"
MOCK_ANSWER_MARKER = "Deterministic mock answer for E2E"
DEFLECT_PHRASE = r"haven't done anything like that"
ERROR_COPY = "The chat model dropped the connection — try again?"
# --------------------------------------------------------------------------
# DB seeding (TRUNCATE-then-seed, cf. test_agent_document_tools.py)
# --------------------------------------------------------------------------
def _seed(db: Session) -> None:
"""The single fixture document (see the module docstring)."""
md = Document(
source=SEED_SOURCE,
path=SEED_PATH,
full_path=f"/tmp/{SEED_PATH}",
title="Kubernetes",
content=KUB_CONTENT,
content_hash=hashlib.sha256(KUB_CONTENT.encode()).hexdigest(),
indexed_at=datetime.now(UTC),
)
db.add(md)
db.flush()
# One chunk carrying the mock's own embedding → genuine token overlap
# for the grounded questions (the FTS path carries them to HIGH).
db.add(
Chunk(
document_id=md.id,
position=0,
content=KUB_CONTENT,
embedding=embed_text(KUB_CONTENT),
)
)
def _reset_db() -> None:
"""Truncate the KB (plus the prompt-shaping tables), then re-seed.
``steering_notes`` / ``kb_overview`` are truncated too, so the
prompts are exactly ``<relevance>`` + ``<documents>`` (+ ``<tools>``
on HIGH) regardless of leftovers from other suites — byte-stable
prompts, byte-stable answers.
"""
with SessionLocal() as db:
db.execute(
text("TRUNCATE chunks, documents, query_log, steering_notes, kb_overview")
)
db.commit()
_seed(db)
db.commit()
def _last_query_log() -> QueryLog:
with SessionLocal() as db:
rows = db.scalars(select(QueryLog)).all()
assert len(rows) == 1, f"expected exactly one query_log row, got {len(rows)}"
return rows[0]
# --------------------------------------------------------------------------
# Page hooks (the #send-status recorder + SSE capture, the phase-37
# pattern from test_agent_document_tools.py)
# --------------------------------------------------------------------------
#: Records every value #send-status takes during the turn (a
#: MutationObserver on the element), so the in-flight status sequence —
#: including the transient phase-67 "retrying" states — is captured
#: deterministically (no polling race).
STATUS_RECORDER = """
() => {
if (window.__statusesInstalled) return;
window.__statusesInstalled = true;
window.__statuses = [];
const el = document.querySelector('#send-status');
if (!el) return;
const rec = (v) => {
const l = window.__statuses;
if (!l.length || l[l.length - 1] !== v) l.push(v);
};
rec(el.textContent);
new MutationObserver(() => rec(el.textContent)).observe(el, {
childList: true,
subtree: true,
});
}
"""
#: Captures the raw SSE ``data:`` payloads of the /api/chat stream
#: (a response clone read in the background) — wire-level assertions
#: for the ``retry`` frames, independent of the UI rendering.
SSE_HOOK = """
() => {
if (window.__sseInstalled) return;
window.__sseInstalled = true;
window.__sseFrames = [];
const origFetch = window.fetch;
window.fetch = async function (...args) {
const res = await origFetch.apply(this, args);
try {
const url = typeof args[0] === 'string' ? args[0] : args[0].url;
if (url.includes('/api/chat')) {
res.clone().text().then((bodyText) => {
for (const block of bodyText.split('\\n\\n')) {
const line = block.trim();
if (line.startsWith('data: ')) {
window.__sseFrames.push(line.slice(6));
}
}
});
}
} catch (e) { /* non-clonable responses: ignored */ }
return res;
};
}
"""
def _install_page_hooks(page: Page) -> None:
"""Install both hooks on the loaded page (post-goto, pre-submit).
The fetch wrapper only needs to be in place before the turn's
``fetch("/api/chat")`` call; the observer needs the rendered
``#send-status``. (``add_init_script`` would not do — it binds to
the NEXT navigation, and the story page is navigated exactly once.)
"""
page.evaluate(SSE_HOOK)
page.evaluate(STATUS_RECORDER)
def _frames(page: Page, terminal: str = "done") -> list[dict]:
"""The captured SSE frames, once the *terminal* frame lands.
The hook reads ``res.clone().text()`` in a background promise that
resolves right after the stream closes — poll briefly until the
terminal (``done``, or ``error`` for the exhaustion test) lands.
"""
deadline = time.monotonic() + 30.0
while True:
raw = page.evaluate("() => window.__sseFrames || []")
parsed = [json.loads(line) for line in raw if line]
if any(f.get("type") == terminal for f in parsed):
return parsed
if time.monotonic() > deadline:
raise AssertionError(
f"SSE hook captured no `{terminal}` frame (frames so far: "
f"{len(parsed)}) — hook install failed?"
)
time.sleep(0.05)
def _retry_frames(frames: list[dict]) -> list[dict]:
return [f for f in frames if f.get("type") == "retry"]
def _assert_retries_before_first_delta(frames: list[dict], retries: list[dict]) -> None:
"""The wire contract (locked A2): every retry frame precedes the
first answer delta — a retry can only restart a request that never
streamed a frame."""
assert retries, "no retry frames on the wire"
first_delta = next(i for i, f in enumerate(frames) if f.get("type") == "delta")
assert all(
i < first_delta for i, f in enumerate(frames) if f.get("type") == "retry"
)
def _assert_no_error_frames(frames: list[dict]) -> None:
assert not [f for f in frames if f.get("type") == "error"]
def _submit(page: Page, question: str) -> None:
page.fill("#message-input", question)
page.click("#send-btn")
# The user bubble lands synchronously with the submit handler.
expect(page.locator(".msg.user .bubble").last).to_contain_text(question)
def _wait_settled(page: Page) -> None:
"""The turn is complete: answer text in the bubble, button recovered.
Phase 48: the label assertion carries the settle wait with an
explicit timeout — the in-flight button is the enabled Stop control
(never disabled), so ``to_be_enabled`` no longer blocks until the
turn settles, and Playwright expect's default (5s) does not inherit
the page default."""
expect(page.locator(".msg.brain .bubble").last).not_to_have_text("", timeout=30_000)
expect(page.locator("#send-btn")).to_be_enabled(timeout=30_000)
expect(page.locator("#send-label")).to_have_text("Send", timeout=30_000)
def _assert_no_error_banner(page: Page) -> None:
"""A retried turn settles through the normal done path — never the
red role=alert error banner (the KB-offline banner is a separate,
health-driven state the db_ready fixture keeps away)."""
banner = page.locator("#kb-banner")
expect(banner).to_be_hidden()
expect(banner).not_to_have_attribute("role", "alert")
expect(banner).not_to_have_class(re.compile(r"is-error"))
# --------------------------------------------------------------------------
# 1. Dead-then-recovered, DEFLECTED path (LOW turn): the status line
# shows the live "retrying" copy and the answer still completes
# --------------------------------------------------------------------------
def test_dead_then_recovered_deflected(
page: Page, app_url: str, mock_llm: int, db_ready: None
) -> None:
_reset_db()
page.set_default_timeout(30_000)
page.goto(app_url)
_install_page_hooks(page)
_submit(page, DEFLECT_Q)
_wait_settled(page)
# The live status (owner-locked A4 copy) was recorded for BOTH
# restarts: attempt 2 (after the original attempt died) and attempt
# 3 (after the first retry died) — attempt 4 never needed to start.
statuses = page.evaluate("() => window.__statuses")
assert _retry_status(2) in statuses, statuses
assert _retry_status(3) in statuses, statuses
# Wire: exactly the two retry frames (attempt = the attempt about to
# be tried, 1-based; max_attempts = the forced-default budget of 4),
# both ahead of the first delta, and no error frame anywhere.
frames = _frames(page)
retries = _retry_frames(frames)
assert retries == [
{"type": "retry", "attempt": 2, "max_attempts": MAX_ATTEMPTS},
{"type": "retry", "attempt": 3, "max_attempts": MAX_ATTEMPTS},
], retries
_assert_retries_before_first_delta(frames, retries)
_assert_no_error_frames(frames)
done = next(f for f in frames if f.get("type") == "done")
assert done["deflected"] is True
# The deflected answer settled normally — no error banner.
last = page.locator(".msg.brain").last
expect(last).to_have_class(re.compile(r"is-deflected"))
expect(last.locator(".bubble")).to_contain_text(
re.compile(DEFLECT_PHRASE, re.IGNORECASE)
)
_assert_no_error_banner(page)
row = _last_query_log()
assert row.question == DEFLECT_Q
assert row.deflected is True
# --------------------------------------------------------------------------
# 2. Dead-then-recovered, GROUNDED path (HIGH turn): the agent round
# retried and the answer completes with its sources
# --------------------------------------------------------------------------
def test_dead_then_recovered_grounded(
page: Page, app_url: str, mock_llm: int, db_ready: None
) -> None:
_reset_db()
page.set_default_timeout(30_000)
page.goto(app_url)
_install_page_hooks(page)
_submit(page, GROUNDED_Q)
_wait_settled(page)
# Same status sequence as the deflected path — the per-round retry
# (agent loop, task 03) surfaces through the same status line.
statuses = page.evaluate("() => window.__statuses")
assert _retry_status(2) in statuses, statuses
assert _retry_status(3) in statuses, statuses
frames = _frames(page)
retries = _retry_frames(frames)
assert retries == [
{"type": "retry", "attempt": 2, "max_attempts": MAX_ATTEMPTS},
{"type": "retry", "attempt": 3, "max_attempts": MAX_ATTEMPTS},
], retries
_assert_retries_before_first_delta(frames, retries)
_assert_no_error_frames(frames)
done = next(f for f in frames if f.get("type") == "done")
assert done["deflected"] is False
assert [(s["source"], s["path"]) for s in done["sources"]] == [
(SEED_SOURCE, SEED_PATH)
]
# The grounded answer completed with the source chip — the agent
# round retried and the turn is intact.
bubble = page.locator(".msg.brain .bubble").last
expect(bubble).to_contain_text(GROUNDED_Q)
expect(bubble).to_contain_text(MOCK_ANSWER_MARKER)
chips = page.locator(".msg.brain .source-chip")
expect(chips).to_have_count(1)
expect(chips.first).to_contain_text(SEED_SP)
_assert_no_error_banner(page)
row = _last_query_log()
assert row.question == GROUNDED_Q
assert row.deflected is False
assert row.sources == SEED_SP
# --------------------------------------------------------------------------
# 3. Embedding retry: the pre-stream embedding loop is visible to the
# UI (a retry status before any answer) and the turn completes
# --------------------------------------------------------------------------
def test_embedding_retry_completes(
page: Page, app_url: str, mock_llm: int, db_ready: None
) -> None:
_reset_db()
page.set_default_timeout(30_000)
page.goto(app_url)
_install_page_hooks(page)
_submit(page, EMBED_Q)
_wait_settled(page)
# The embedding step runs BEFORE the honesty gate and the answer
# stream, so its single retry is the ONLY retry frame of the turn —
# and the UI saw it as the live status (no answer frame had landed).
statuses = page.evaluate("() => window.__statuses")
assert _retry_status(2) in statuses, statuses
frames = _frames(page)
retries = _retry_frames(frames)
assert retries == [
{"type": "retry", "attempt": 2, "max_attempts": MAX_ATTEMPTS},
], retries
_assert_retries_before_first_delta(frames, retries)
_assert_no_error_frames(frames)
done = next(f for f in frames if f.get("type") == "done")
assert done["deflected"] is False
# The turn completed normally with the grounded answer + chip.
bubble = page.locator(".msg.brain .bubble").last
expect(bubble).to_contain_text(MOCK_ANSWER_MARKER)
chips = page.locator(".msg.brain .source-chip")
expect(chips).to_have_count(1)
expect(chips.first).to_contain_text(SEED_SP)
_assert_no_error_banner(page)
# --------------------------------------------------------------------------
# 4. Exhaustion: a dead endpoint burns the whole budget, then the
# EXISTING terminal error banner lands and the composer recovers
# --------------------------------------------------------------------------
def test_exhaustion_lands_on_the_error_banner(
page: Page, app_url: str, mock_llm: int, db_ready: None
) -> None:
_reset_db()
page.set_default_timeout(30_000)
page.goto(app_url)
_install_page_hooks(page)
_submit(page, EXHAUST_Q)
# After the 4th dead attempt the turn dies: the existing terminal
# error banner (role=alert) with the existing copy, and the send
# button re-enabled (the banner path settles the state machine).
expect(page.locator("#kb-banner")).to_have_attribute(
"role", "alert", timeout=60_000
)
expect(page.locator("#kb-banner")).to_contain_text(ERROR_COPY)
expect(page.locator("#send-btn")).to_be_enabled(timeout=30_000)
expect(page.locator("#send-label")).to_have_text("Send", timeout=30_000)
# The live status climbed the whole budget: the LAST retrying status
# is the highest attempt — "(4 of 4)…" (no attempt 5 exists).
statuses = page.evaluate("() => window.__statuses")
retry_statuses = [s for s in statuses if "retrying" in s]
assert retry_statuses, statuses
assert retry_statuses[-1] == _retry_status(MAX_ATTEMPTS), statuses
# Wire: the three retry frames (attempts 2, 3, 4 of 4), then the
# terminal error frame as the LAST event — no done, no delta.
frames = _frames(page, terminal="error")
retries = _retry_frames(frames)
assert retries == [
{"type": "retry", "attempt": a, "max_attempts": MAX_ATTEMPTS}
for a in (2, 3, 4)
], retries
assert frames[-1]["type"] == "error"
assert ERROR_COPY in frames[-1]["detail"]
assert not [f for f in frames if f.get("type") == "done"]
assert not [f for f in frames if f.get("type") == "delta"]
# No answer bubble was ever rendered (no frame ever streamed a
# token) — the user bubble is the only message in the DOM.
expect(page.locator("#messages > .msg.brain")).to_have_count(0)
expect(page.locator(".msg.user .bubble")).to_have_count(1)
+200 -6
View File
@@ -61,6 +61,8 @@ class FakeRagLLM:
stream_error: Exception | None = None,
fail_mid_stream: bool = False,
tool_script: list[list[StreamPiece | ToolCallPiece]] | None = None,
embed_fail_count: int = 0,
stream_fail_count: int = 0,
) -> None:
self.settings = Settings(_env_file=None) # pyright: ignore[reportCallIssue]
self.embed_batches = 0
@@ -69,6 +71,14 @@ class FakeRagLLM:
self.embed_error = embed_error
self.stream_error = stream_error
self.fail_mid_stream = fail_mid_stream
#: Phase 67: the first N ``embed_one`` calls raise an
#: ``EmbeddingError`` (then succeed) — a dead-then-recovered
#: embeddings endpoint for the retry loop.
self.embed_fail_count = embed_fail_count
#: Phase 67: the first N ``chat_stream`` requests die with an
#: ``LLMError`` BEFORE any piece (then succeed) — a dead-then-
#: recovered answer endpoint for the pre-first-piece retry rule.
self.stream_fail_count = stream_fail_count
self.question_embeds: list[str] = []
self.seen_messages: list[list[dict[str, str]]] = []
#: Every request's ``tools`` value (phase 37) — ``None`` is the
@@ -100,6 +110,10 @@ class FakeRagLLM:
async def embed_one(self, text: str) -> list[float]:
if self.embed_error is not None:
raise self.embed_error
if self.embed_fail_count > 0:
self.embed_fail_count -= 1
self.question_embeds.append(text)
raise EmbeddingError("simulated embeddings endpoint failure")
self.question_embeds.append(text)
return _token_vec(text)
@@ -118,6 +132,9 @@ class FakeRagLLM:
self.seen_tools.append(tools)
if self.stream_error is not None:
raise self.stream_error
if self.stream_fail_count > 0:
self.stream_fail_count -= 1
raise LLMError("simulated pre-piece endpoint failure")
if tools is not None and self.tool_script:
for piece in self.tool_script.pop(0):
yield piece
@@ -388,17 +405,33 @@ def test_chat_empty_kb_streams_empty_sources(client, db) -> None:
assert row.sources == ""
def test_chat_embed_failure_yields_error_event(client, db, seeded_kb: FakeRagLLM) -> None:
def test_chat_embed_failure_yields_error_event(
client, db, seeded_kb: FakeRagLLM, monkeypatch: pytest.MonkeyPatch
) -> None:
"""Phase 67: a dead embeddings endpoint retries on the configured
budget — one ``retry`` frame per restart (the attempt about to be
tried, 1-based) — and settles on the existing terminal error frame;
no query_log row. Zero delay keeps the exhaustion path fast."""
broken = FakeRagLLM(embed_error=EmbeddingError("embeddings endpoint down"))
live = get_settings()
monkeypatch.setattr(
chat_api, "get_settings", lambda: _retry_settings(live)
)
fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: broken
try:
_, _, frames = _stream_chat(client, QUESTION)
finally:
fastapi_app.dependency_overrides.clear()
assert len(frames) == 1
assert frames[0]["type"] == "error"
assert "embedding" in frames[0]["detail"]
retries = live.llm_retries
assert [f["type"] for f in frames] == ["retry"] * retries + ["error"]
assert [f["attempt"] for f in frames if f["type"] == "retry"] == list(
range(2, retries + 2)
)
assert all(
f["max_attempts"] == retries + 1 for f in frames if f["type"] == "retry"
)
assert "embedding" in frames[-1]["detail"]
assert db.scalars(select(QueryLog)).all() == []
@@ -416,11 +449,19 @@ def test_chat_mid_stream_failure_yields_error_after_partial_deltas(client, db) -
assert db.scalars(select(QueryLog)).all() == []
def test_error_event_matches_contract_shape(client, db, seeded_kb) -> None:
def test_error_event_matches_contract_shape(
client, db, seeded_kb, monkeypatch: pytest.MonkeyPatch
) -> None:
"""The SSE error event (PLAN §4) is exactly ``{type, detail}`` — the
client's loading-feedback state machine (phase 06) keys off this shape
to flip to the error state and re-enable the send button."""
to flip to the error state and re-enable the send button.
``llm_retries=0`` keeps this a single-attempt turn: the contract under
test is the error frame itself, not the phase-67 retry loop."""
broken = FakeRagLLM(embed_error=EmbeddingError("embeddings endpoint down"))
live = get_settings()
monkeypatch.setattr(
chat_api, "get_settings", lambda: _retry_settings(live, llm_retries=0)
)
fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: broken
try:
_, _, frames = _stream_chat(client, QUESTION)
@@ -718,3 +759,156 @@ def test_tool_execution_db_failure_yields_error_event(
assert frames[0]["name"] == "list_documents"
assert "offline mid-question" in frames[1]["detail"]
assert db.scalars(select(QueryLog)).all() == [] # no row for a failed turn
# ---------- phase 67: LLM retries before the first token ----------
def _retry_settings(live: Settings, **overrides: Any) -> Settings:
"""Settings for the retry tests: the live (mock-calibrated) threshold
plus the phase-67 knobs, with a ZERO delay so the suite never sleeps.
(The 5 s default is unit-pinned in ``tests/unit/test_config.py``.)"""
kwargs: dict[str, Any] = {
"relevance_threshold": live.relevance_threshold,
"llm_retry_delay": 0.0,
}
kwargs.update(overrides)
return Settings(_env_file=None, **kwargs) # pyright: ignore[reportCallIssue]
def test_embed_failure_retries_then_turn_completes(
client,
db,
seeded_kb: FakeRagLLM,
monkeypatch: pytest.MonkeyPatch,
caplog: pytest.LogCaptureFixture,
) -> None:
"""A dead-then-recovered embeddings endpoint: one SSE ``retry`` frame
(the attempt about to be tried, 1-based) ahead of the normal answer
frames; the turn completes and the per-turn log line counts the
retry (``retries=1``)."""
flaky = FakeRagLLM(embed_fail_count=1)
live = get_settings()
monkeypatch.setattr(
chat_api, "get_settings", lambda: _retry_settings(live, llm_retries=1)
)
fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: flaky
try:
caplog.set_level(logging.INFO, logger="app.chat")
_, _, frames = _stream_chat(client, QUESTION)
finally:
fastapi_app.dependency_overrides.clear()
assert frames[0] == {"type": "retry", "attempt": 2, "max_attempts": 2}
assert not any(f["type"] == "error" for f in frames)
deltas = [f for f in frames if f["type"] == "delta"]
assert len(deltas) >= 2
assert "".join(d["text"] for d in deltas) == flaky.answer
assert frames[-1]["type"] == "done"
lines = [r.getMessage() for r in caplog.records if "question=" in r.getMessage()]
assert lines and "retries=1" in lines[-1]
def test_embed_failure_exhausts_retries_then_terminal_error(
client, db, seeded_kb: FakeRagLLM, monkeypatch: pytest.MonkeyPatch
) -> None:
"""A dead embeddings endpoint (``llm_retries=2`` → 3 attempts): one
``retry`` frame per restart (attempts 2 and 3 of 3), then the
EXISTING terminal error frame — the copy is unchanged, no query_log
row."""
dead = FakeRagLLM(embed_fail_count=99) # every attempt fails
live = get_settings()
monkeypatch.setattr(
chat_api, "get_settings", lambda: _retry_settings(live, llm_retries=2)
)
fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: dead
try:
_, _, frames = _stream_chat(client, QUESTION)
finally:
fastapi_app.dependency_overrides.clear()
assert [f["type"] for f in frames] == ["retry", "retry", "error"]
assert [f["attempt"] for f in frames if f["type"] == "retry"] == [2, 3]
assert all(f["max_attempts"] == 3 for f in frames if f["type"] == "retry")
assert "embedding" in frames[-1]["detail"]
assert db.scalars(select(QueryLog)).all() == []
def test_deflected_stream_retries_before_the_first_piece(
client,
db,
seeded_kb: FakeRagLLM,
monkeypatch: pytest.MonkeyPatch,
caplog: pytest.LogCaptureFixture,
) -> None:
"""Deflected answer stream: the first attempt dies before any piece,
the restart streams — a ``retry`` frame ahead of the deltas, the
request restarted with the same messages (no tools key), and the
per-turn log line counts the retry."""
flaky = FakeRagLLM(stream_fail_count=1)
live = get_settings()
monkeypatch.setattr(
chat_api, "get_settings", lambda: _retry_settings(live, llm_retries=2)
)
fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: flaky
try:
caplog.set_level(logging.INFO, logger="app.chat")
_, _, frames = _stream_chat(client, OFF_TOPIC)
finally:
fastapi_app.dependency_overrides.clear()
assert frames[0] == {"type": "retry", "attempt": 2, "max_attempts": 3}
rest = frames[1:]
assert all(f["type"] in ("delta", "done") for f in rest)
assert "".join(f["text"] for f in rest if f["type"] == "delta") == flaky.answer
assert rest[-1]["type"] == "done" and rest[-1]["deflected"] is True
assert len(flaky.seen_messages) == 2 # the request was restarted
assert flaky.seen_tools == [None, None] # …byte-identical (no tools key)
lines = [r.getMessage() for r in caplog.records if "question=" in r.getMessage()]
assert lines and "retries=1" in lines[-1]
def test_deflected_stream_failure_after_first_frame_is_terminal(
client, db, seeded_kb: FakeRagLLM, monkeypatch: pytest.MonkeyPatch
) -> None:
"""Locked A2: a stream failure AFTER the first output frame is
terminal — no ``retry`` frame, the existing error copy, no row (a
partial answer is never redone)."""
broken = FakeRagLLM(fail_mid_stream=True)
live = get_settings()
monkeypatch.setattr(
chat_api, "get_settings", lambda: _retry_settings(live, llm_retries=3)
)
fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: broken
try:
_, _, frames = _stream_chat(client, OFF_TOPIC)
finally:
fastapi_app.dependency_overrides.clear()
assert [f["type"] for f in frames] == ["delta", "error"]
assert not any(f["type"] == "retry" for f in frames)
assert "dropped the connection" in frames[1]["detail"]
assert db.scalars(select(QueryLog)).all() == []
def test_zero_retries_keep_the_pre_phase_wire_shape(
client, db, seeded_kb: FakeRagLLM, monkeypatch: pytest.MonkeyPatch
) -> None:
"""The ``BOR_LLM_RETRIES=0`` kill switch: one attempt, the existing
terminal error frame, no ``retry`` frames — the pre-phase-67
byte-identical wire shape."""
broken = FakeRagLLM(embed_error=EmbeddingError("embeddings endpoint down"))
live = get_settings()
monkeypatch.setattr(
chat_api, "get_settings", lambda: _retry_settings(live, llm_retries=0)
)
fastapi_app.dependency_overrides[chat_api.get_llm] = lambda: broken
try:
_, _, frames = _stream_chat(client, QUESTION)
finally:
fastapi_app.dependency_overrides.clear()
assert len(frames) == 1
assert frames[0]["type"] == "error"
assert "embedding" in frames[0]["detail"]
assert not any(f["type"] == "retry" for f in frames)
+307 -7
View File
@@ -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) ----------
+35
View File
@@ -112,6 +112,41 @@ def test_max_output_tokens_env_override(monkeypatch) -> None:
assert s.max_output_tokens == 1234
def test_llm_retry_defaults(monkeypatch: pytest.MonkeyPatch) -> None:
"""Phase 67: a failed LLM request is retried by default — 3 retries
with a flat 5 s delay (the TODO-locked values, no backoff)."""
monkeypatch.delenv("BOR_LLM_RETRIES", raising=False)
monkeypatch.delenv("BOR_LLM_RETRY_DELAY", raising=False)
s = _settings()
assert s.llm_retries == 3
assert s.llm_retry_delay == 5.0
def test_llm_retry_env_overrides(monkeypatch: pytest.MonkeyPatch) -> None:
"""``BOR_LLM_RETRIES`` / ``BOR_LLM_RETRY_DELAY`` override the defaults;
``0`` retries is the no-retry kill switch (pre-phase-67 behavior)."""
monkeypatch.setenv("BOR_LLM_RETRIES", "0")
monkeypatch.setenv("BOR_LLM_RETRY_DELAY", "1.5")
s = _settings()
assert s.llm_retries == 0
assert s.llm_retry_delay == 1.5
def test_llm_retries_rejects_negative(monkeypatch: pytest.MonkeyPatch) -> None:
"""``0`` is the kill switch — a negative value is a typo, so the
validator fails loudly at startup (the ``agent_max_rounds`` pattern)."""
monkeypatch.setenv("BOR_LLM_RETRIES", "-1")
with pytest.raises(ValidationError, match="llm_retries"):
_settings()
def test_llm_retry_delay_rejects_negative(monkeypatch: pytest.MonkeyPatch) -> None:
"""A negative delay is a typo — fail loudly at startup."""
monkeypatch.setenv("BOR_LLM_RETRY_DELAY", "-0.5")
with pytest.raises(ValidationError, match="llm_retry_delay"):
_settings()
def test_agent_max_rounds_default_and_env_override(monkeypatch: pytest.MonkeyPatch) -> None:
"""Phase 45: the per-tool budgets are gone — ``BOR_AGENT_MAX_ROUNDS``
(default 10) is the single agent-loop knob; ``0`` is the no-tools
+53
View File
@@ -605,3 +605,56 @@ def test_index_messages_comment_documents_the_meta_actions() -> None:
assert "every visitor" in comment.lower() or "everyone" in comment.lower(), (
"Retry is documented as available to all visitors"
)
# ---------- llm retry status (phase 67, task 04) ----------
def test_retry_frame_is_a_first_class_branch_between_tool_and_delta() -> None:
"""Phase 67 (owner-locked A4, task 02's contract): a `retry` SSE
frame (the server restarted the LLM request before its first piece
— locked A2) is a first-class branch in runTurn's handler, ordered
BETWEEN the `tool` and `delta` branches. It clears the 120s guard
(a frame arrived), resolves n/N from the frame, and writes the
owner-locked copy literal onto the EXISTING channels only — the
#send-status live region + the typing-indicator aria-label. No DOM
of its own: no addMessage, no appendToolLine, no error banner, no
UI-state change. The gate covers BOTH live states: a later agent
round may restart while the UI already streams."""
js = _js()
tool_idx = js.find('ev.type === "tool"')
retry_idx = js.find('ev.type === "retry"')
delta_idx = js.find('ev.type === "delta"')
assert -1 < tool_idx < retry_idx < delta_idx, (
"the turn handler must branch on retry frames, between tool and delta"
)
assert js.count('ev.type === "retry"') == 1, "exactly one retry branch"
branch = js[retry_idx:delta_idx]
assert "clearTurnTimeout()" in branch, "a frame arrived — the 120s guard clears"
assert "Number(ev.attempt)" in branch, "n = the attempt about to be tried"
assert "Number(ev.max_attempts)" in branch, "N = the configured total"
assert (
"`Communication interrupted — retrying (${attempt} of ${max})…`" in branch
), "the owner-locked copy literal (A4)"
assert "sendStatus.textContent = retryStatus" in branch, (
"the existing #send-status live region carries the status"
)
assert "#typing-indicator .bubble" in branch, ("the typing-indicator is reused")
assert 'setAttribute("aria-label", retryStatus)' in branch
assert "UI_STATE.thinking" in branch and "UI_STATE.streaming" in branch, (
"the gate covers BOTH live states (a later round may restart mid-stream)"
)
for forbidden in ("addMessage", "appendToolLine", "showErrorBanner", "setUiState"):
assert forbidden not in branch, f"transient status only — no {forbidden}"
def test_header_inventory_documents_the_retry_frame() -> None:
"""The app.js file-header doc comment inventories every SSE frame
type (house convention); phase 67 adds the `retry` frame there, with
the owner-locked copy quoted, so the A4 literal has exactly two
homes: the header doc and the handler branch."""
js = _js()
header = js[: js.find("import {")]
assert "`retry`" in header, "the header must list the retry frame"
assert "Communication interrupted — retrying" in header
assert header.count("Communication interrupted — retrying") == 1
+274
View File
@@ -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
+26 -1
View File
@@ -4,7 +4,12 @@ from __future__ import annotations
import json
from app.api.chat import sse_event
from app.schemas import ChatErrorEvent, ChatThinkingEvent, ChatToolEvent
from app.schemas import (
ChatErrorEvent,
ChatRetryEvent,
ChatThinkingEvent,
ChatToolEvent,
)
def _payload(frame: str) -> dict:
@@ -100,3 +105,23 @@ def test_tool_event_shape_is_type_name_argument_only() -> None:
dumped = ChatToolEvent(name="read_document", argument="S/p.md").model_dump()
assert set(dumped.keys()) == {"type", "name", "argument"}
assert dumped["type"] == "tool" # default — call sites never spell it out
def test_retry_frame_serializes_exactly() -> None:
"""Phase 67: the ``retry`` frame is exactly ``{type: "retry",
attempt: int, max_attempts: int}`` — a transient status the client's
readSSE handler branches on ("Communication interrupted — retrying
(n of N)…"), never an error. ``attempt`` is the attempt the server is
about to try next (1-based); the contract the frontend branch and the
E2E suite key off, locked here ahead of the JS implementation."""
frame = sse_event(ChatRetryEvent(attempt=2, max_attempts=4).model_dump())
assert frame == 'data: {"type": "retry", "attempt": 2, "max_attempts": 4}\n\n'
assert _payload(frame) == {"type": "retry", "attempt": 2, "max_attempts": 4}
def test_retry_event_shape_is_type_attempt_max_attempts_only() -> None:
dumped = ChatRetryEvent(attempt=2, max_attempts=4).model_dump()
assert set(dumped.keys()) == {"type", "attempt", "max_attempts"}
assert dumped["type"] == "retry" # default — call sites never spell it out
assert isinstance(dumped["attempt"], int)
assert isinstance(dumped["max_attempts"], int)