fix queuing system to break less

app.py

@@ -160,14 +160,37 @@ let chars = 0;
 let htmlChunks = [];
 let htmlSize = 0;
 let mode = 'idle';
+let requestId = null;
+let heartbeatInterval = null;
+let waitingConfirmStart = null;
+let statusStart = null;
 const interval = setInterval(() => {
   if (mode === 'generating') {
     stats.textContent = 'Characters: ' + chars.toLocaleString() + ' | ' + ((Date.now() - start) / 1000).toFixed(1) + 's';
   }
+  if (mode === 'waiting_confirmation' && waitingConfirmStart && Date.now() - waitingConfirmStart > 10000) {
+    es.close();
+    clearInterval(interval);
+    clearInterval(heartbeatInterval);
+    clearTimeout(timeout);
+    titleEl.textContent = 'Connection lost';
+    document.querySelector('.subtitle').textContent = 'Could not connect to worker. Please refresh.';
+    retryBtn.style.display = 'inline-block';
+  }
+  if (mode === 'status' && statusStart && Date.now() - statusStart > 60000) {
+    es.close();
+    clearInterval(interval);
+    clearInterval(heartbeatInterval);
+    clearTimeout(timeout);
+    titleEl.textContent = 'Connection lost';
+    document.querySelector('.subtitle').textContent = 'Could not connect to worker. Please refresh.';
+    retryBtn.style.display = 'inline-block';
+  }
 }, 100);
 const timeout = setTimeout(() => {
   es.close();
   clearInterval(interval);
+  clearInterval(heartbeatInterval);
   titleEl.textContent = 'Request timed out';
   document.querySelector('.subtitle').textContent = 'Generation took too long. Please try again.';
   retryBtn.style.display = 'inline-block';
@@ -175,8 +198,12 @@ const timeout = setTimeout(() => {
 window.__wotf_html = '';
 es.onmessage = (e) => {
   if (e.data === '[DONE]') {
+    clearInterval(heartbeatInterval);
+    heartbeatInterval = null;
     titleEl.textContent = 'Validating HTML...';
   } else if (e.data === '[RENDER]') {
+    clearInterval(heartbeatInterval);
+    heartbeatInterval = null;
     clearTimeout(timeout);
     clearInterval(interval);
     es.close();
@@ -191,13 +218,35 @@ es.onmessage = (e) => {
     htmlChunks = [];
     htmlSize = 0;
     chars = 0;
+  } else if (e.data.startsWith('[ID:')) {
+    requestId = e.data.slice(4, -1);
   } else if (e.data.startsWith('[STATUS:')) {
     const msg = e.data.slice(8);
-    if (mode !== 'generating') {
-      stats.textContent = msg;
+    const cleanMsg = msg.endsWith(']') ? msg.slice(0, -1) : msg;
+    if (cleanMsg === 'waiting_confirmation') {
+      mode = 'waiting_confirmation';
+      waitingConfirmStart = Date.now();
+      if (requestId && !window.__confirmed) {
+        window.__confirmed = true;
+        fetch('/confirm/' + requestId, { method: 'POST' }).catch(() => {
+          es.close();
+          clearInterval(interval);
+          clearInterval(heartbeatInterval);
+          clearTimeout(timeout);
+          titleEl.textContent = 'Connection to worker failed';
+          document.querySelector('.subtitle').textContent = 'Could not connect to generation worker. Please refresh.';
+          retryBtn.style.display = 'inline-block';
+        });
+      }
+    } else if (mode !== 'generating') {
+      stats.textContent = cleanMsg;
       mode = 'status';
+      waitingConfirmStart = null;
+      statusStart = Date.now();
     }
   } else if (e.data === '[ERROR]') {
+    clearInterval(heartbeatInterval);
+    heartbeatInterval = null;
     clearTimeout(timeout);
     clearInterval(interval);
     es.close();
@@ -206,6 +255,12 @@ es.onmessage = (e) => {
     retryBtn.style.display = 'inline-block';
   } else {
     mode = 'generating';
+    statusStart = null;
+    if (!heartbeatInterval && requestId) {
+      heartbeatInterval = setInterval(() => {
+        fetch('/heartbeat/' + requestId, { method: 'POST' });
+      }, 3000);
+    }
     const bytes = Uint8Array.from(atob(e.data), c => c.charCodeAt(0));
     const decoded = new TextDecoder().decode(bytes);
     htmlChunks.push(decoded);
@@ -214,6 +269,8 @@ es.onmessage = (e) => {
   }
 };
 es.onerror = () => {
+  clearInterval(heartbeatInterval);
+  heartbeatInterval = null;
   clearTimeout(timeout);
   clearInterval(interval);
   es.close();
@@ -589,9 +646,42 @@ def worker_loop() -> None:
                 time.sleep(0.5)
                 continue
 
-            r.set(f"wotf:req:{request_id}:status", "processing")  # type: ignore[union-attr]
+            r.set(f"wotf:req:{request_id}:status", "waiting_confirmation")  # type: ignore[union-attr]
             r.delete(f"wotf:stream:{request_id}:chunks")  # type: ignore[union-attr]
 
+            r.rpush(f"wotf:stream:{request_id}:chunks", f"[ID:{request_id}]")  # type: ignore[union-attr]
+            r.rpush(f"wotf:stream:{request_id}:chunks", "[STATUS:waiting_confirmation]")  # type: ignore[union-attr]
+
+            r.delete("wotf:processing")  # type: ignore[union-attr]
+            r.delete(hb_key)  # type: ignore[union-attr]
+
+            confirmed = False
+            for _ in range(150):
+                time.sleep(0.1)
+                if r.get(f"wotf:req:{request_id}:confirmed"):  # type: ignore[union-attr]
+                    confirmed = True
+                    break
+
+            if not confirmed:
+                logger.info(f"[worker] No confirmation for {request_id}, marking as error")
+                r.set(f"wotf:req:{request_id}:status", "error")  # type: ignore[union-attr]
+                r.set(f"wotf:req:{request_id}:error", "Worker did not receive confirmation from client")  # type: ignore[union-attr]
+                r.rpush(f"wotf:stream:{request_id}:chunks", "[STATUS:Connection to worker failed]")  # type: ignore[union-attr]
+                r.rpush(f"wotf:stream:{request_id}:chunks", "[ERROR]")  # type: ignore[union-attr]
+                r.delete(f"wotf:req:{request_id}:confirmed")  # type: ignore[union-attr]
+                r.expire(f"wotf:req:{request_id}:ctx", 30)  # type: ignore[union-attr]
+                time.sleep(0.5)
+                continue
+
+            acquired = r.set("wotf:processing", worker_id, ex=120, nx=True)  # type: ignore[union-attr]
+            if not acquired:
+                time.sleep(0.5)
+                continue
+
+            r.set(hb_key, worker_id, ex=15)  # type: ignore[union-attr]
+            r.set(f"wotf:req:{request_id}:status", "processing")  # type: ignore[union-attr]
+            r.set(f"wotf:req:{request_id}:last_heartbeat", str(time.time()))  # type: ignore[union-attr]
+
             ctx_raw = r.hgetall(f"wotf:req:{request_id}:ctx")  # type: ignore[union-attr]
             ctx: dict[str, Any] = {}
             if isinstance(ctx_raw, dict):
@@ -605,12 +695,37 @@ def worker_loop() -> None:
 
             start = time.time()
             hb_start = time.time()
+            hb_check_start = time.time()
             html_buffer: list[str] = []
+            aborted = False
             try:
                 for chunk in stream_llm(ctx):
                     if time.time() - hb_start >= 5:
                         r.set(hb_key, worker_id, ex=15)  # type: ignore[union-attr]
                         hb_start = time.time()
+
+                    if time.time() - hb_check_start >= 2:
+                        hb_check_start = time.time()
+                        last_hb_raw = r.get(f"wotf:req:{request_id}:last_heartbeat")  # type: ignore[union-attr]
+                        if last_hb_raw:
+                            last_hb = float(last_hb_raw)  # type: ignore[arg-type]
+                            if time.time() - last_hb > 10:
+                                logger.info(f"[worker] Client disconnected (heartbeat stale), aborting {request_id}")
+                                r.set(f"wotf:req:{request_id}:status", "aborted")  # type: ignore[union-attr]
+                                r.rpush(f"wotf:stream:{request_id}:chunks", "[ERROR]")  # type: ignore[union-attr]
+                                r.delete("wotf:processing")  # type: ignore[union-attr]
+                                r.delete(hb_key)  # type: ignore[union-attr]
+                                aborted = True
+                                break
+                        else:
+                            logger.info(f"[worker] No heartbeat from client, aborting {request_id}")
+                            r.set(f"wotf:req:{request_id}:status", "aborted")  # type: ignore[union-attr]
+                            r.rpush(f"wotf:stream:{request_id}:chunks", "[ERROR]")  # type: ignore[union-attr]
+                            r.delete("wotf:processing")  # type: ignore[union-attr]
+                            r.delete(hb_key)  # type: ignore[union-attr]
+                            aborted = True
+                            break
+
                     html_buffer.append(chunk)
                     encoded = base64.b64encode(chunk.encode("utf-8")).decode("ascii")
                     r.rpush(f"wotf:stream:{request_id}:chunks", encoded)  # type: ignore[union-attr]
@@ -632,6 +747,9 @@ def worker_loop() -> None:
                 r.delete(hb_key)  # type: ignore[union-attr]
                 time.sleep(1)
                 continue
+            if aborted:
+                time.sleep(1)
+                continue
             html = "".join(html_buffer)
             r.rpush(f"wotf:stream:{request_id}:chunks", "[STATUS:Validating HTML structure...]")  # type: ignore[union-attr]
             errors = validate_html(html)
@@ -770,7 +888,7 @@ def stream_llm(context: dict[str, Any]) -> Generator[str]:
         elapsed = time.time() - start
         logger.error(f"[{req_id}] HTTP error in {elapsed:.1f}s: {e.response.status_code} {e.response.text[:200]}")
         yield f"\n\n<!-- Error: {e.response.status_code} -->"
-    except BaseException as e:  # noqa: BLE001
+    except Exception as e:  # noqa: BLE001
         elapsed = time.time() - start
         logger.error(f"[{req_id}] Stream failed in {elapsed:.1f}s: {e}")
         yield f"\n\n<!-- Error: {e} -->"
@@ -842,6 +960,7 @@ def stream() -> Response:
     def generate() -> Generator[str]:
         yield "data: [CLEAR]\n\n"
 
+        last_idx = 0
         while True:
             status_raw = r.get(f"wotf:req:{request_id}:status")  # type: ignore[union-attr]
             status = status_raw if isinstance(status_raw, str) else "pending"
@@ -853,18 +972,36 @@ def stream() -> Response:
                 yield "data: [ERROR]\n\n"
                 return
 
-            if status == "done":
-                break
-
-            if status == "processing":
+            if status == "waiting_confirmation":
                 chunks_raw = r.lrange(f"wotf:stream:{request_id}:chunks", 0, -1)  # type: ignore[union-attr]
                 chunks: list[str] = list(chunks_raw) if isinstance(chunks_raw, list) else []  # type: ignore[arg-type]
                 if chunks:
-                    break
-                yield "data: [STATUS:Generating your page...\n\n"
-                time.sleep(0.5)
+                    for c in chunks:
+                        yield f"data: {c}\n\n"
+                else:
+                    yield "data: [STATUS:Waiting for your connection...\n\n"
+                time.sleep(0.2)
                 continue
 
+            if status == "processing":
+                chunks_raw = r.lrange(f"wotf:stream:{request_id}:chunks", last_idx, -1)  # type: ignore[union-attr]
+                chunks: list[str] = list(chunks_raw) if isinstance(chunks_raw, list) else []  # type: ignore[arg-type]
+                if chunks:
+                    for c in chunks:
+                        yield f"data: {c}\n\n"
+                    last_idx += len(chunks)
+                else:
+                    yield "data: [STATUS:Generating your page...\n\n"
+                time.sleep(0.1)
+                continue
+
+            if status == "done":
+                chunks_raw = r.lrange(f"wotf:stream:{request_id}:chunks", last_idx, -1)  # type: ignore[union-attr]
+                chunks: list[str] = list(chunks_raw) if isinstance(chunks_raw, list) else []  # type: ignore[arg-type]
+                for c in chunks:
+                    yield f"data: {c}\n\n"
+                break
+
             info = get_queue_info(request_id)
             pos = info["position"]
             eta = info["eta"]
@@ -874,7 +1011,6 @@ def stream() -> Response:
                 yield f"data: [STATUS:Position {pos} in line | Est. wait: ~{eta:.0f}s\n\n"
             time.sleep(1)
 
-        last_idx = 0
         while True:
             chunks_raw = r.lrange(f"wotf:stream:{request_id}:chunks", last_idx, -1)  # type: ignore[union-attr]
             chunks: list[str] = list(chunks_raw) if isinstance(chunks_raw, list) else []  # type: ignore[arg-type]
@@ -890,6 +1026,11 @@ def stream() -> Response:
                 status_raw = r.get(f"wotf:req:{request_id}:status")  # type: ignore[union-attr]
                 status = status_raw if isinstance(status_raw, str) else "pending"
                 if status == "done":
+                    chunks_raw = r.lrange(f"wotf:stream:{request_id}:chunks", last_idx, -1)  # type: ignore[union-attr]
+                    chunks: list[str] = list(chunks_raw) if isinstance(chunks_raw, list) else []  # type: ignore[arg-type]
+                    if chunks:
+                        for chunk in chunks:
+                            yield f"data: {chunk}\n\n"
                     break
                 if status == "error":
                     error_raw = r.get(f"wotf:req:{request_id}:error")  # type: ignore[union-attr]
@@ -897,6 +1038,10 @@ def stream() -> Response:
                     yield f"data: [STATUS:Error: {error_msg}\n\n"
                     yield "data: [ERROR]\n\n"
                     return
+                if status == "aborted":
+                    yield "data: [STATUS:Generation aborted (connection lost)\n\n"
+                    yield "data: [ERROR]\n\n"
+                    return
                 time.sleep(0.2)
 
     return Response(
@@ -910,6 +1055,24 @@ def stream() -> Response:
     )
 
 
+@app.route("/confirm/<request_id>", methods=["POST"])
+def confirm(request_id: str) -> Response:
+    ctx = r.exists(f"wotf:req:{request_id}:ctx")  # type: ignore[union-attr]
+    if not ctx:
+        return Response("not found", status=404)
+    r.set(f"wotf:req:{request_id}:confirmed", "1")  # type: ignore[union-attr]
+    return Response("ok", status=200)
+
+
+@app.route("/heartbeat/<request_id>", methods=["POST"])
+def heartbeat(request_id: str) -> Response:
+    status = r.get(f"wotf:req:{request_id}:status")  # type: ignore[union-attr]
+    if status is None:
+        return Response("not found", status=404)
+    r.set(f"wotf:req:{request_id}:last_heartbeat", str(time.time()))  # type: ignore[union-attr]
+    return Response("ok", status=200)
+
+
 def start_worker_if_needed() -> None:
     global worker_running, worker_thread
     if worker_running: