feat: add MedGemma readiness gate to prevent cold-start pipeline failures

src/backend/app/api/health.py

@@ -23,4 +23,18 @@ async def config_check():
         "medgemma_model_id": settings.medgemma_model_id,
         "hf_token_set": bool(settings.hf_token),
         "medgemma_max_tokens": settings.medgemma_max_tokens,
+    }
+
+
+@router.get("/api/health/model")
+async def model_readiness():
+    """Check if the MedGemma endpoint is warm and accepting requests."""
+    from app.services.medgemma import MedGemmaService
+
+    service = MedGemmaService()
+    ready = await service.check_readiness()
+    return {
+        "ready": ready,
+        "model_id": settings.medgemma_model_id,
+        "base_url_set": bool(settings.medgemma_base_url),
     }

src/backend/app/api/ws.py

@@ -11,12 +11,15 @@ from __future__ import annotations
 
 import asyncio
 import json
+import logging
 
 from fastapi import APIRouter, WebSocket, WebSocketDisconnect
 
 from app.agent.orchestrator import Orchestrator
 from app.models.schemas import CaseSubmission
+from app.services.medgemma import MedGemmaService
 
+logger = logging.getLogger(__name__)
 router = APIRouter()
 
 
@@ -46,7 +49,37 @@ async def agent_websocket(websocket: WebSocket):
         # Send acknowledgment
         await websocket.send_json({
             "type": "ack",
-            "message": "Case received. Starting agent pipeline...",
+            "message": "Case received. Checking model readiness...",
+        })
+
+        # ── Readiness gate: wait for MedGemma to be warm ──
+        medgemma = MedGemmaService()
+
+        async def _send_warming(elapsed: float, message: str):
+            """Stream warm-up progress to client."""
+            try:
+                await websocket.send_json({
+                    "type": "warming_up",
+                    "message": message,
+                    "elapsed_seconds": int(elapsed),
+                })
+            except Exception:
+                pass  # client may have disconnected
+
+        ready = await medgemma.wait_until_ready(on_waiting=_send_warming)
+        if not ready:
+            await websocket.send_json({
+                "type": "error",
+                "message": (
+                    "MedGemma model did not become ready within the timeout. "
+                    "The endpoint may be starting up — please try again in a minute."
+                ),
+            })
+            return
+
+        await websocket.send_json({
+            "type": "model_ready",
+            "message": "MedGemma is ready. Starting agent pipeline...",
         })
 
         # Run the orchestrator and stream updates

src/backend/app/services/medgemma.py

@@ -27,6 +27,10 @@ T = TypeVar("T", bound=BaseModel)
 MAX_API_RETRIES = 3
 RETRY_BASE_DELAY = 5.0  # seconds, doubles on each retry
 
+# Readiness probe configuration
+READINESS_TIMEOUT = 180  # max seconds to wait for model warm-up
+READINESS_POLL_INTERVAL = 5  # seconds between readiness checks
+
 
 class MedGemmaService:
     """
@@ -58,6 +62,71 @@ class MedGemmaService:
                 )
         return self._client
 
+    async def check_readiness(self) -> bool:
+        """
+        Lightweight probe to check if the MedGemma endpoint is warm and
+        accepting requests.  Sends a tiny 1-token generate call.
+
+        Returns True if the model responds, False on any transient error.
+        """
+        if self._mode != "api":
+            return True  # local mode is always "ready"
+        try:
+            client = await self._get_client()
+            response = await client.chat.completions.create(
+                model=settings.medgemma_model_id,
+                messages=[{"role": "user", "content": "ping"}],
+                max_tokens=1,
+                temperature=0.0,
+            )
+            return bool(response.choices)
+        except Exception as e:
+            logger.debug(f"Readiness probe failed: {e}")
+            return False
+
+    async def wait_until_ready(
+        self,
+        timeout: float = READINESS_TIMEOUT,
+        poll_interval: float = READINESS_POLL_INTERVAL,
+        on_waiting: Optional[Any] = None,
+    ) -> bool:
+        """
+        Poll check_readiness() until the model is warm or timeout expires.
+
+        Args:
+            timeout: Maximum seconds to wait.
+            poll_interval: Seconds between probes.
+            on_waiting: Optional async callback(elapsed_seconds, message) invoked
+                        each time we're still waiting — used to stream status to
+                        the client.
+
+        Returns:
+            True if the model became ready, False if timeout was reached.
+        """
+        import time
+        start = time.monotonic()
+        attempt = 0
+        while True:
+            attempt += 1
+            if await self.check_readiness():
+                logger.info("MedGemma readiness probe succeeded (%.1fs)", time.monotonic() - start)
+                return True
+
+            elapsed = time.monotonic() - start
+            if elapsed >= timeout:
+                logger.error("MedGemma readiness timeout after %.0fs", elapsed)
+                return False
+
+            msg = (
+                f"Warming up MedGemma model... "
+                f"({int(elapsed)}s elapsed, attempt {attempt})"
+            )
+            logger.info(msg)
+            if on_waiting:
+                await on_waiting(elapsed, msg)
+
+            await asyncio.sleep(poll_interval)
+
     async def generate(
         self,
         prompt: str,

src/frontend/src/app/page.tsx

@@ -7,7 +7,7 @@ import { CDSReport } from "@/components/CDSReport";
 import { useAgentWebSocket } from "@/hooks/useAgentWebSocket";
 
 export default function Home() {
-  const { steps, report, isRunning, error, submitCase } = useAgentWebSocket();
+  const { steps, report, isRunning, isWarmingUp, warmUpMessage, error, submitCase } = useAgentWebSocket();
   const [hasSubmitted, setHasSubmitted] = useState(false);
 
   const handleSubmit = (patientText: string) => {
@@ -72,6 +72,21 @@ export default function Home() {
             <div className="lg:col-span-2">
               {report ? (
                 <CDSReport report={report} />
+              ) : isWarmingUp ? (
+                <div className="flex items-center justify-center h-64 text-amber-600">
+                  <div className="text-center">
+                    <div className="animate-pulse w-10 h-10 rounded-full bg-amber-100 flex items-center justify-center mx-auto mb-4">
+                      <span className="text-xl">&#9881;</span>
+                    </div>
+                    <p className="font-medium">Model Warming Up</p>
+                    <p className="text-sm text-amber-500 mt-1">
+                      {warmUpMessage || "Waiting for MedGemma endpoint..."}
+                    </p>
+                    <p className="text-xs text-gray-400 mt-2">
+                      This happens when the model scales from zero. Usually takes 1-2 minutes.
+                    </p>
+                  </div>
+                </div>
               ) : isRunning ? (
                 <div className="flex items-center justify-center h-64 text-gray-400">
                   <div className="text-center">

src/frontend/src/hooks/useAgentWebSocket.ts

@@ -22,6 +22,8 @@ interface UseAgentWebSocketReturn {
   steps: Step[];
   report: any | null;
   isRunning: boolean;
+  isWarmingUp: boolean;
+  warmUpMessage: string | null;
   error: string | null;
   submitCase: (submission: CaseSubmission) => void;
 }
@@ -45,6 +47,8 @@ export function useAgentWebSocket(): UseAgentWebSocketReturn {
   const [steps, setSteps] = useState<Step[]>([]);
   const [report, setReport] = useState<any | null>(null);
   const [isRunning, setIsRunning] = useState(false);
+  const [isWarmingUp, setIsWarmingUp] = useState(false);
+  const [warmUpMessage, setWarmUpMessage] = useState<string | null>(null);
   const [error, setError] = useState<string | null>(null);
   const wsRef = useRef<WebSocket | null>(null);
 
@@ -54,6 +58,8 @@ export function useAgentWebSocket(): UseAgentWebSocketReturn {
     setReport(null);
     setError(null);
     setIsRunning(true);
+    setIsWarmingUp(false);
+    setWarmUpMessage(null);
 
     // Close existing connection
     if (wsRef.current) {
@@ -75,6 +81,16 @@ export function useAgentWebSocket(): UseAgentWebSocketReturn {
           // Pipeline acknowledged
           break;
 
+        case "warming_up":
+          setIsWarmingUp(true);
+          setWarmUpMessage(data.message);
+          break;
+
+        case "model_ready":
+          setIsWarmingUp(false);
+          setWarmUpMessage(null);
+          break;
+
         case "step_update":
           setSteps((prev) => {
             const existing = prev.findIndex(
@@ -114,5 +130,5 @@ export function useAgentWebSocket(): UseAgentWebSocketReturn {
     };
   }, []);
 
-  return { steps, report, isRunning, error, submitCase };
+  return { steps, report, isRunning, isWarmingUp, warmUpMessage, error, submitCase };
 }