From 5e83c8a56df402e2e67bbd4cded56079ce142187 Mon Sep 17 00:00:00 2001
From: Bryan Van Deusen <bryan.vandeusen@fabledsword.com>
Date: Thu, 26 Feb 2026 22:49:06 -0500
Subject: [PATCH] Add explicit warm-wait before generation starts
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Instead of relying solely on retry-on-500, poll /api/ps before starting
any LLM stream so the main model has time to fully load into VRAM.

- llm.py: add wait_for_model_loaded(model, timeout=90s) — polls /api/ps
  every 2s, returns True when model appears in loaded list
- generation_task.py: launch model_load_task in parallel with build_context
  and classify_intent (both use fast/small-model ops that don't need the
  main model); after context is built, await the load task — shows
  "Loading model..." status only if the user actually has to wait;
  logs a warning and proceeds if 90s timeout elapses

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 .../services/generation_task.py               | 19 +++++++++++---
 src/fabledassistant/services/llm.py           | 25 +++++++++++++++++++
 2 files changed, 40 insertions(+), 4 deletions(-)

diff --git a/src/fabledassistant/services/generation_task.py b/src/fabledassistant/services/generation_task.py
index 93793a1..4edff59 100644
--- a/src/fabledassistant/services/generation_task.py
+++ b/src/fabledassistant/services/generation_task.py
@@ -22,7 +22,7 @@ from fabledassistant.services.generation_buffer import (
     GenerationBuffer,
     GenerationState,
 )
-from fabledassistant.services.llm import ChatChunk, build_context, generate_completion, stream_chat, stream_chat_with_tools, summarize_history_for_context
+from fabledassistant.services.llm import ChatChunk, build_context, generate_completion, stream_chat, stream_chat_with_tools, summarize_history_for_context, wait_for_model_loaded
 from fabledassistant.services.chat import update_conversation_title
 from fabledassistant.services.intent import IntentResult, classify_intent
 from fabledassistant.services.logging import log_generation
@@ -210,9 +210,12 @@ async def run_generation(
         buf.append_event("status", {"status": "Summarizing conversation history..."})
         history_to_use, history_summary = await summarize_history_for_context(history, intent_model)
 
-    # Phase 3: Build context and start intent classification in parallel.
-    # We block on context (need messages to stream) — intent is consumed
-    # after context is ready, at the start of round 0.
+    # Phase 3: Build context, classify intent, and wait for model — all in parallel.
+    # build_context is fast DB/search ops that don't need the main model.
+    # classify_intent uses the small intent model, not the main model.
+    # wait_for_model_loaded polls /api/ps so the main stream starts without 500 errors.
+    model_load_task = asyncio.create_task(wait_for_model_loaded(model, timeout=90.0))
+
     context_task = asyncio.create_task(build_context(
         user_id, history_to_use, context_note_id, user_content,
         history_summary=history_summary,
@@ -235,6 +238,14 @@ async def run_generation(
     # Emit context event
     buf.append_event("context", {"context": context_meta})
 
+    # Wait for main model to be loaded before starting any generation.
+    # If it's already loaded (common case), this returns immediately.
+    if not model_load_task.done():
+        buf.append_event("status", {"status": "Loading model..."})
+    loaded = await model_load_task
+    if not loaded:
+        logger.warning("Model %s did not load within 90s — proceeding anyway", model)
+
     t_start = time.monotonic()
     timing: dict = {
         "intent_ms": None,
diff --git a/src/fabledassistant/services/llm.py b/src/fabledassistant/services/llm.py
index a9a7a84..1564aee 100644
--- a/src/fabledassistant/services/llm.py
+++ b/src/fabledassistant/services/llm.py
@@ -2,6 +2,7 @@ import asyncio
 import json
 import logging
 import re
+import time
 from collections.abc import AsyncGenerator
 from dataclasses import dataclass, field
 from typing import Literal
@@ -75,6 +76,30 @@ async def ensure_model(model: str) -> None:
         logger.warning("Failed to pull model '%s' — chat may not work", model, exc_info=True)
 
 
+async def wait_for_model_loaded(model: str, timeout: float = 90.0) -> bool:
+    """Poll /api/ps every 2s until the model appears in Ollama's loaded-model list.
+
+    Returns True when the model is loaded, False if timeout elapses first.
+    Used before generation to avoid streaming 500s during cold model loads.
+    """
+    base = model.removesuffix(":latest")
+    deadline = time.monotonic() + timeout
+    while True:
+        try:
+            async with httpx.AsyncClient(timeout=5.0) as client:
+                resp = await client.get(f"{Config.OLLAMA_URL}/api/ps")
+                resp.raise_for_status()
+                loaded = {m["name"] for m in resp.json().get("models", [])}
+                if model in loaded or f"{base}:latest" in loaded or base in loaded:
+                    return True
+        except Exception:
+            pass  # Ollama may still be starting up
+        remaining = deadline - time.monotonic()
+        if remaining <= 0:
+            return False
+        await asyncio.sleep(min(2.0, remaining))
+
+
 async def stream_chat(
     messages: list[dict],
     model: str,