From d5a53738727000fd4a52f090925fc2e47f14d22a Mon Sep 17 00:00:00 2001
From: Bryan Van Deusen <bryan.vandeusen@fabledsword.com>
Date: Thu, 26 Feb 2026 05:55:09 -0500
Subject: [PATCH] Add stream retry to all generation paths, not just round 0

Adds _stream_with_retry() async generator (wraps stream_chat_with_tools
with up to 2 retries on Ollama 500, 3s/6s delay). Previously only the
optimistic round 0 _fill_queue had retry logic. Two paths were still
bare: the declined-write-tool fresh stream, and the round 1+ stream.

Round 1 500s occur when tag suggestions (fire-and-forget inside
execute_tool) race the follow-up stream to the same model. The retry
waits for tag suggestions to complete before succeeding.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 .../services/generation_task.py               | 39 ++++++++++++++++++-
 1 file changed, 37 insertions(+), 2 deletions(-)

diff --git a/src/fabledassistant/services/generation_task.py b/src/fabledassistant/services/generation_task.py
index 72d7a3b..355250f 100644
--- a/src/fabledassistant/services/generation_task.py
+++ b/src/fabledassistant/services/generation_task.py
@@ -195,6 +195,41 @@ async def _drain_queue(
         yield item
 
 
+async def _stream_with_retry(
+    messages: list[dict],
+    model: str,
+    tools: list[dict],
+    think: bool,
+) -> AsyncGenerator[ChatChunk, None]:
+    """stream_chat_with_tools with automatic retry on Ollama 500 errors.
+
+    500s occur when Ollama is still loading a model or handling a concurrent
+    request (e.g. tag suggestions racing with round 1). Retries up to 2 times
+    with a short delay — by which point the model is warm and other calls done.
+    """
+    last_exc: BaseException | None = None
+    for attempt in range(3):
+        if attempt > 0:
+            delay = 3.0 * attempt
+            logger.warning(
+                "Ollama stream 500 (attempt %d/3), retrying in %.0fs", attempt, delay
+            )
+            await asyncio.sleep(delay)
+        try:
+            async for chunk in stream_chat_with_tools(messages, model, tools=tools, think=think):
+                yield chunk
+            return
+        except httpx.HTTPStatusError as exc:
+            last_exc = exc
+            if exc.response.status_code != 500:
+                break  # non-500 is not retryable
+        except BaseException as exc:
+            last_exc = exc
+            break
+    if last_exc is not None:
+        raise last_exc
+
+
 async def run_generation(
     buf: GenerationBuffer,
     history: list[dict],
@@ -467,7 +502,7 @@ async def run_generation(
 
                 if use_tool_path:
                     # Declined write tool — the optimistic stream was cancelled, start fresh.
-                    stream_source: AsyncGenerator = stream_chat_with_tools(messages, model, tools=tools, think=think)
+                    stream_source: AsyncGenerator = _stream_with_retry(messages, model, tools, think)
                 else:
                     stream_source = _drain_queue(prefetched, stream_queue)
 
@@ -546,7 +581,7 @@ async def run_generation(
 
             buf.append_event("status", {"status": "Generating response..." if _round == 0 else "Composing response..."})
             t_stream = time.monotonic()
-            async for chunk in stream_chat_with_tools(messages, model, tools=tools, think=think):
+            async for chunk in _stream_with_retry(messages, model, tools, think):
                 if buf.cancel_event.is_set():
                     cancelled = True
                     break