feat: real recovery UI for generation failures

apps/backend/src/taxflow/routers/query.py

@@ -7,6 +7,7 @@
 import logging
 import time
 
+import openai
 from fastapi import APIRouter, Depends, HTTPException
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
@@ -758,16 +759,45 @@ async def generate():
             # effective_question above); None on a first turn.
             "clarifications": parsed_clarifications,
         }
-        async for mode, chunk in research_graph.astream(
-            initial_state, stream_mode=["custom", "values"]
-        ):
-            if mode == "custom":
-                text = chunk["token"]
-                yield f"data: {json.dumps({'type': 'token', 'text': text})}\n\n"
-            elif mode == "values":
-                latest_values = chunk
-                if not first_pass_snapshot and chunk.get("answer"):
-                    first_pass_snapshot = chunk
+        # Accountant audit round three (#4): an unhandled exception here used to
+        # just crash the async generator - FastAPI closes the SSE connection
+        # with no event, EventSource fires a bare onerror, and the dashboard
+        # showed a generic "Query failed - please try again" with no
+        # explanation and no distinction between "the AI provider is briefly
+        # at capacity" (worth retrying) and a real bug (worth reporting). Any
+        # tokens already streamed stay visible; this only replaces the silent
+        # connection-drop with a real, typed event.
+        try:
+            async for mode, chunk in research_graph.astream(
+                initial_state, stream_mode=["custom", "values"]
+            ):
+                if mode == "custom":
+                    text = chunk["token"]
+                    yield f"data: {json.dumps({'type': 'token', 'text': text})}\n\n"
+                elif mode == "values":
+                    latest_values = chunk
+                    if not first_pass_snapshot and chunk.get("answer"):
+                        first_pass_snapshot = chunk
+        except Exception as e:
+            logger.warning("research graph stream failed", exc_info=True)
+            await asyncio.to_thread(
+                db.queries.update, client["id"], query_id, {"status": "failed", "error_message": str(e)}
+            )
+            # litellm normalises every provider's transient failures (rate
+            # limits, timeouts, connection drops, 5xx) onto openai's exception
+            # hierarchy (litellm.exceptions.RateLimitError subclasses
+            # openai.RateLimitError, etc.) - a real, checkable signal for
+            # "worth retrying" vs "something is actually broken" without
+            # guessing from the message text.
+            transient = isinstance(e, openai.APIError)
+            message = (
+                "The AI provider is briefly at capacity - this usually clears within a minute."
+                if transient
+                else "Something went wrong generating this answer."
+            )
+            yield f"data: {json.dumps({'type': 'error', 'transient': transient, 'message': message})}\n\n"
+            yield "data: [DONE]\n\n"
+            return
 
         final = latest_values
 

apps/backend/tests/test_feedback_and_stream.py

@@ -1,6 +1,8 @@
 """Tests for the feedback endpoint (Task C5) and stream-path metric persistence."""
+import json
 from unittest.mock import AsyncMock, MagicMock, patch
 
+import litellm
 import pytest
 
 
@@ -156,6 +158,88 @@ async def fake_astream(initial_state, stream_mode=None):
     assert chunks[-1] == "data: [DONE]\n\n"
 
 
+# --- accountant audit round three, #4: a stream failure gets a real event,
+# not a silently-dropped connection --------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_stream_provider_error_emits_transient_error_event():
+    """A litellm/openai APIError (rate limit, timeout, provider 5xx) - the class
+    of failure litellm normalises every provider's transient errors onto - must
+    surface as a typed, retryable `error` event, not crash the generator."""
+    import taxflow.routers.query as q
+
+    fake_client = {"id": "client-1", "email": "a@b.com.au"}
+
+    captured_update = {}
+    mock_db = MagicMock()
+    mock_db.queries.insert.return_value = {"id": "query-1"}
+    mock_db.queries.update.side_effect = lambda cid, qid, payload: captured_update.update(payload)
+
+    async def fake_astream(initial_state, stream_mode=None):
+        yield ("custom", {"token": "Partial answer before it broke"})
+        raise litellm.exceptions.RateLimitError(
+            message="rate limited", llm_provider="openrouter", model="deepseek"
+        )
+        yield  # pragma: no cover - unreachable, makes this a generator
+
+    with patch.object(
+        q, "embed", new=AsyncMock(return_value=[0.0] * 1536)
+    ), patch.object(q, "increment_usage", new=AsyncMock()), patch.object(
+        q.research_graph, "astream", new=fake_astream
+    ), patch.object(
+        q.answer_cache, "get_cached_answer", new=AsyncMock(return_value=None)
+    ):
+        response = await q.stream_query(question="q", client=fake_client, _trial=fake_client, db=mock_db)
+        chunks = [c async for c in response.body_iterator]
+
+    types = [_event_type(c) for c in chunks]
+    assert types == ["token", "error", None]
+    assert chunks[-1] == "data: [DONE]\n\n"
+
+    error_payload = json.loads(chunks[1].removeprefix("data: ").strip())
+    assert error_payload["transient"] is True
+    assert "capacity" in error_payload["message"].lower()
+
+    # The already-partially-processed query row is marked failed, not left
+    # stuck in "processing" forever.
+    assert captured_update["status"] == "failed"
+    assert "rate limited" in captured_update["error_message"]
+
+
+@pytest.mark.asyncio
+async def test_stream_unexpected_error_emits_non_transient_error_event():
+    """A non-provider exception (e.g. a bug in graph code) must still surface
+    as a typed error event, but marked non-transient - retrying the identical
+    question is not expected to help the way it might for a capacity error."""
+    import taxflow.routers.query as q
+
+    fake_client = {"id": "client-1", "email": "a@b.com.au"}
+    mock_db = MagicMock()
+    mock_db.queries.insert.return_value = {"id": "query-1"}
+    mock_db.queries.update.return_value = None
+
+    async def fake_astream(initial_state, stream_mode=None):
+        raise ValueError("something in the graph broke")
+        yield  # pragma: no cover - unreachable, makes this a generator
+
+    with patch.object(
+        q, "embed", new=AsyncMock(return_value=[0.0] * 1536)
+    ), patch.object(q, "increment_usage", new=AsyncMock()), patch.object(
+        q.research_graph, "astream", new=fake_astream
+    ), patch.object(
+        q.answer_cache, "get_cached_answer", new=AsyncMock(return_value=None)
+    ):
+        response = await q.stream_query(question="q", client=fake_client, _trial=fake_client, db=mock_db)
+        chunks = [c async for c in response.body_iterator]
+
+    types = [_event_type(c) for c in chunks]
+    assert types == ["error", None]
+    error_payload = json.loads(chunks[0].removeprefix("data: ").strip())
+    assert error_payload["transient"] is False
+    assert "capacity" not in error_payload["message"].lower()
+
+
 @pytest.mark.asyncio
 async def test_stream_observability_failure_is_best_effort():
     """Task 1b: if the observability add-on (check_citation_validity / run_cost)

apps/dashboard/app/dashboard/query/page.tsx

@@ -4,6 +4,7 @@ import { memo, useCallback, useEffect, useRef, useState } from "react";
 import Link from "next/link";
 import { toast } from "sonner";
 import {
+  AlertTriangle,
   BookOpen,
   CheckCircle2,
   Copy,
@@ -14,6 +15,7 @@ import {
   MessageSquare,
   MessagesSquare,
   Pencil,
+  RotateCw,
   ThumbsUp,
   ThumbsDown,
 } from "lucide-react";
@@ -869,6 +871,17 @@ export default function QueryPage() {
   const [sessionId, setSessionId] = useState<string>(() => crypto.randomUUID());
   const [loading, setLoading] = useState(false);
   const [error, setError] = useState<string | null>(null);
+  // Accountant audit round three (#4): a real generation failure now gets its
+  // own typed state - message + whether retrying is likely to help + the
+  // exact question that failed - instead of collapsing into the generic
+  // `error` string every other unrelated save/edit/promote flow also uses.
+  // Lets the answer pane render an actual "Try again" action instead of a
+  // one-line red toast with no next step.
+  const [generationError, setGenerationError] = useState<{
+    message: string;
+    transient: boolean;
+    question: string;
+  } | null>(null);
   const [result, setResult] = useState<QueryResult | null>(null);
   // True only once the answer is authoritative: a live stream has emitted
   // [DONE] (after any correction/regeneration), or a persisted conversation was
@@ -1023,6 +1036,7 @@ export default function QueryPage() {
     setSavedDocId(null);
     setDocType("advice_memo");
     setError(null);
+    setGenerationError(null);
     // Phase 4: clear any prior clarify card / follow-up chips.
     setClarifyQuestions(null);
     setClarifyAskedQuestion("");
@@ -1289,6 +1303,8 @@ export default function QueryPage() {
             session?: AnswerTrace["session"];
             re_retrieval?: AnswerTrace["re_retrieval"];
             passes?: AnswerTrace["passes"];
+            message?: string;
+            transient?: boolean;
           } = JSON.parse(event.data);
 
           if (parsed.type === "token" && parsed.text) {
@@ -1350,6 +1366,18 @@ export default function QueryPage() {
               re_retrieval: parsed.re_retrieval ?? null,
               passes: parsed.passes ?? null,
             });
+          } else if (parsed.type === "error") {
+            // A real, typed failure from the backend (provider capacity,
+            // timeout, or an actual bug) - not a dropped connection. The
+            // backend still sends [DONE] right after this, so the promise
+            // below resolves normally instead of hitting the generic
+            // onerror/catch path.
+            setVerifying(false);
+            setGenerationError({
+              message: parsed.message ?? "Something went wrong generating this answer.",
+              transient: parsed.transient ?? false,
+              question: askedQuestion,
+            });
           }
         };
         source.onerror = () => {
@@ -1358,8 +1386,14 @@ export default function QueryPage() {
         };
       });
     } catch {
-      setError("Query failed - please try again");
-      toast.error("Query failed - please try again");
+      // The connection itself dropped before the backend could send a typed
+      // `error` event (e.g. a network blip) - still give a real retry action,
+      // not just a toast with nowhere to go.
+      setGenerationError({
+        message: "The connection dropped before this finished.",
+        transient: true,
+        question: askedQuestion,
+      });
     } finally {
       setLoading(false);
     }
@@ -1654,6 +1688,28 @@ export default function QueryPage() {
             </div>
           )}
 
+          {generationError && (
+            <div className="flex items-start gap-3 rounded-lg border border-destructive/30 bg-destructive/5 p-4">
+              <AlertTriangle className="mt-0.5 size-4 shrink-0 text-destructive" />
+              <div className="min-w-0 flex-1 space-y-2">
+                <p className="text-sm text-foreground">{generationError.message}</p>
+                <p className="text-xs text-muted-foreground">
+                  {generationError.transient
+                    ? "This is usually a brief provider issue, not something wrong with your question."
+                    : "If this keeps happening on the same question, let us know."}
+                </p>
+                <Button
+                  size="sm"
+                  variant="outline"
+                  disabled={loading}
+                  onClick={() => handleSubmit({ questionOverride: generationError.question })}
+                >
+                  <RotateCw className="size-3.5" />
+                  Try again
+                </Button>
+              </div>
+            </div>
+          )}
           {error && <p className="text-sm text-destructive">{error}</p>}
         </div>