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>
This commit is contained in:
2026-02-26 05:55:09 -05:00
parent b23e78b0ac
commit d5a5373872
@@ -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