diff --git a/apps/backend/src/taxflow/routers/query.py b/apps/backend/src/taxflow/routers/query.py index 3d99fde..8bd1f64 100644 --- a/apps/backend/src/taxflow/routers/query.py +++ b/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 diff --git a/apps/backend/tests/test_feedback_and_stream.py b/apps/backend/tests/test_feedback_and_stream.py index f32d935..28362d3 100644 --- a/apps/backend/tests/test_feedback_and_stream.py +++ b/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) diff --git a/apps/dashboard/app/dashboard/query/page.tsx b/apps/dashboard/app/dashboard/query/page.tsx index 0044238..e002542 100644 --- a/apps/dashboard/app/dashboard/query/page.tsx +++ b/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(() => crypto.randomUUID()); const [loading, setLoading] = useState(false); const [error, setError] = useState(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(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() { )} + {generationError && ( +
+ +
+

{generationError.message}

+

+ {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."} +

+ +
+
+ )} {error &&

{error}

}