fix heartbeat issue

README.md

@@ -35,6 +35,7 @@ Loaded from `.env` via `python-dotenv`:
 | `OPENAI_API_KEY` | Yes | — | API key for the LLM provider |
 | `OPENAI_API_URL` | No | `https://api.openai.com/v1/chat/completions` | OpenAI-compatible endpoint |
 | `MODEL` | No | `gpt-4o-mini` | Model name to use |
+| `VALKEY_URL` | No | `redis://localhost:6379/0` | Valkey/Redis URL |
 
 ## Site types
 

app.py

@@ -1,10 +1,11 @@
 import base64
-import importlib
 import json
 import logging
 import os
 import random
 import re
+import signal
+import sys
 import threading
 import time
 import uuid
@@ -21,7 +22,7 @@ if os.getenv("GUNICORN_WORKER_ID"):
 
     monkey.patch_all()
 
-lxml_html: Any = importlib.import_module("lxml.html")
+from lxml import html as lxml_html  # type: ignore[import-untyped]
 
 logging.basicConfig(
     level=logging.INFO,
@@ -40,6 +41,68 @@ MODEL = os.getenv("MODEL", "gpt-4o-mini")
 VALKEY_URL = os.getenv("VALKEY_URL", "redis://localhost:6379/0")
 r = valkey.Valkey.from_url(VALKEY_URL, decode_responses=True)  # type: ignore[assignment]
 
+
+class ValkeyClient:
+    def __init__(self, client: valkey.Valkey) -> None:
+        self._c = client
+
+    def get(self, key: str) -> str | None:
+        v = self._c.get(key)  # type: ignore[union-attr]
+        return str(v) if isinstance(v, bytes) else v  # type: ignore[return-value]
+
+    def set(self, key: str, value: str, ex: int | None = None, nx: bool = False, xx: bool = False) -> bool:
+        return bool(self._c.set(key, value, ex=ex, nx=nx, xx=xx))  # type: ignore[union-attr]
+
+    def delete(self, *keys: str) -> int:
+        return int(self._c.delete(*keys))  # type: ignore[union-attr]
+
+    def exists(self, key: str) -> bool:
+        return bool(self._c.exists(key))  # type: ignore[union-attr]
+
+    def rpush(self, key: str, *values: str) -> int:
+        return int(self._c.rpush(key, *values))  # type: ignore[union-attr]
+
+    def lpop(self, key: str) -> str | None:
+        v = self._c.lpop(key)  # type: ignore[union-attr]
+        return str(v) if isinstance(v, bytes) else v  # type: ignore[return-value]
+
+    def lrange(self, key: str, start: int, end: int) -> list[str]:
+        v = self._c.lrange(key, start, end)  # type: ignore[union-attr]
+        return [str(x) for x in v]  # type: ignore[arg-type]
+
+    def llen(self, key: str) -> int:
+        return int(self._c.llen(key))  # type: ignore[union-attr]
+
+    def ttl(self, key: str) -> int:
+        return int(self._c.ttl(key))  # type: ignore[union-attr]
+
+    def hset(self, key: str, mapping: dict[str, str]) -> int:
+        return int(self._c.hset(key, mapping=mapping))  # type: ignore[union-attr]
+
+    def hgetall(self, key: str) -> dict[str, str]:
+        v = self._c.hgetall(key)  # type: ignore[union-attr]
+        return {str(k): str(vv) for k, vv in v.items()}  # type: ignore[arg-type]
+
+    def incr(self, key: str) -> int:
+        return int(self._c.incr(key))  # type: ignore[union-attr]
+
+    def expire(self, key: str, seconds: int) -> bool:
+        return bool(self._c.expire(key, seconds))  # type: ignore[union-attr]
+
+
+rv = ValkeyClient(r)
+
+CHUNK_SIZE = 1000
+CONFIRMATION_TIMEOUT_ITERATIONS = 150
+CONFIRMATION_POLL_INTERVAL = 0.1
+HEARTBEAT_TIMEOUT_S = 10
+HEARTBEAT_CHECK_INTERVAL_S = 2
+WORKER_LOCK_TTL_S = 300
+WORKER_LOCK_HEARTBEAT_INTERVAL_S = 60
+PROCESSING_LOCK_TTL_S = 120
+HB_KEY_TTL_S = 15
+REQUEST_CTX_TTL_S = 30
+
 worker_running = False
 worker_thread: threading.Thread | None = None
 
@@ -427,16 +490,16 @@ def validate_html(html: str) -> list[str]:
         errors.append("Missing DOCTYPE")
 
     try:
-        doc: Any = lxml_html.fromstring(html)
+        doc: Any = lxml_html.fromstring(html)  # type: ignore[union-attr]
-        html_tags: list[Any] = doc.xpath("//html")
+        html_tags: list[Any] = doc.xpath("//html")  # type: ignore[assignment,union-attr]
-        body_tags: list[Any] = doc.xpath("//body")
+        body_tags: list[Any] = doc.xpath("//body")  # type: ignore[assignment,union-attr]
-        head_tags: list[Any] = doc.xpath("//head")
+        head_tags: list[Any] = doc.xpath("//head")  # type: ignore[assignment,union-attr]
-        if len(html_tags) != 1:
+        if len(html_tags) != 1:  # type: ignore[arg-type]
-            errors.append(f"Expected 1 <html> tag, found {len(html_tags)}")
+            errors.append(f"Expected 1 <html> tag, found {len(html_tags)}")  # type: ignore[arg-type]
-        if len(body_tags) != 1:
+        if len(body_tags) != 1:  # type: ignore[arg-type]
-            errors.append(f"Expected 1 <body> tag, found {len(body_tags)}")
+            errors.append(f"Expected 1 <body> tag, found {len(body_tags)}")  # type: ignore[arg-type]
-        if len(head_tags) != 1:
+        if len(head_tags) != 1:  # type: ignore[arg-type]
-            errors.append(f"Expected 1 <head> tag, found {len(head_tags)}")
+            errors.append(f"Expected 1 <head> tag, found {len(head_tags)}")  # type: ignore[arg-type]
     except BaseException as e:  # noqa: BLE001
         errors.append(f"HTML parse error: {e}")
 
@@ -455,15 +518,18 @@ def validate_html(html: str) -> list[str]:
             errors.append(f"JS block {i + 1}: mismatched brackets")
 
     if script_blocks:
-        js_text = "\n".join(script_blocks)
+        try:
-        review = call_llm(
+            js_text = "\n".join(script_blocks)
-            f"Review this JavaScript for runtime errors (undefined variables, wrong selectors, syntax issues). List each issue concisely. If no issues, say 'No issues found.'\n\n<code>\n{js_text}\n</code>"
+            review = call_llm(
-        )
+                f"Review this JavaScript for runtime errors (undefined variables, wrong selectors, syntax issues). List each issue concisely. If no issues, say 'No issues found.'\n\n<code>\n{js_text}\n</code>"
-        if review and "no issue" not in review.lower():
+            )
-            for line in review.strip().split("\n"):
+            if review and "no issue" not in review.lower():
-                line = line.strip()
+                for line in review.strip().split("\n"):
-                if line and not line.startswith("<"):
+                    line = line.strip()
-                    errors.append(f"JS review: {line}")
+                    if line and not line.startswith("<"):
+                        errors.append(f"JS review: {line}")
+        except BaseException as e:  # noqa: BLE001
+            logger.warning(f"JS review failed: {e}")
 
     return errors
 
@@ -580,117 +646,111 @@ THEMES = [
 def enqueue_request(context: dict[str, Any]) -> tuple[str, int]:
     request_id = uuid.uuid4().hex[:12]
     ctx_key = f"wotf:req:{request_id}:ctx"
-    r.hset(ctx_key, mapping={k: json.dumps(v) for k, v in context.items()})  # type: ignore[union-attr]
+    rv.hset(ctx_key, mapping={k: json.dumps(v) for k, v in context.items()})
-    r.set(f"wotf:req:{request_id}:status", "pending")  # type: ignore[union-attr]
+    rv.set(f"wotf:req:{request_id}:status", "pending")
-    r.rpush("wotf:queue", request_id)  # type: ignore[union-attr]
+    rv.rpush("wotf:queue", request_id)
-    position_raw = r.llen("wotf:queue")  # type: ignore[union-attr]
+    position = rv.llen("wotf:queue")
-    position: int = int(position_raw) if isinstance(position_raw, (int, float)) else 0
     return (request_id, position)
 
 
 def get_queue_info(request_id: str) -> dict[str, Any]:
-    queue_items_raw = r.lrange("wotf:queue", 0, -1)  # type: ignore[union-attr]
+    queue_items = rv.lrange("wotf:queue", 0, -1)
-    queue_items: list[str] = list(queue_items_raw) if isinstance(queue_items_raw, list) else []  # type: ignore[arg-type]
     if request_id in queue_items:
         position = queue_items.index(request_id) + 1
     else:
         position = 0
-    avg_raw = r.get("wotf:stats:avg_time")  # type: ignore[union-attr]
+    avg_raw = rv.get("wotf:stats:avg_time")
     avg_str: str = str(avg_raw) if avg_raw is not None else "15"
     avg = float(avg_str)
     eta = position * avg
     return {"position": position, "avg_time": avg, "eta": eta}
 
 
-def get_generated_html(request_id: str) -> str:
-    result = r.get(f"wotf:req:{request_id}:html")  # type: ignore[union-attr]
-    return result if isinstance(result, str) else ""
-
-
 def worker_loop() -> None:
     worker_id = f"worker-{os.getpid()}-{threading.current_thread().name}"
     hb_key = "wotf:processing:hb"
+    lock_hb_start = time.time()
     while worker_running:
         request_id: str | None = None
         try:
-            lock_holder = r.get("wotf:processing")  # type: ignore[union-attr]
+            lock_holder = rv.get("wotf:processing")
             if lock_holder is not None:
-                hb_ttl_raw = r.ttl(hb_key)  # type: ignore[union-attr]
+                hb_ttl = rv.ttl(hb_key)
-                hb_ttl: int = int(hb_ttl_raw) if isinstance(hb_ttl_raw, (int, float)) else -2
                 if hb_ttl == -2:
                     logger.info(f"[worker] Reclaiming stale lock from {lock_holder}")
-                    r.delete("wotf:processing")  # type: ignore[union-attr]
+                    rv.delete("wotf:processing")
-                    r.delete(hb_key)  # type: ignore[union-attr]
+                    rv.delete(hb_key)
                 else:
                     time.sleep(0.5)
                     continue
 
-            acquired = r.set("wotf:processing", worker_id, ex=120, nx=True)  # type: ignore[union-attr]
+            acquired = rv.set("wotf:processing", worker_id, ex=PROCESSING_LOCK_TTL_S, nx=True)
             if not acquired:
                 time.sleep(0.5)
                 continue
 
-            r.set(hb_key, worker_id, ex=15)  # type: ignore[union-attr]
+            rv.set(hb_key, worker_id, ex=HB_KEY_TTL_S)
 
-            request_id_raw = r.lpop("wotf:queue")  # type: ignore[union-attr]
+            request_id = rv.lpop("wotf:queue")
-            if request_id_raw is None:
-                r.delete("wotf:processing")  # type: ignore[union-attr]
-                r.delete(hb_key)  # type: ignore[union-attr]
-                time.sleep(0.5)
-                continue
-
-            request_id = request_id_raw if isinstance(request_id_raw, str) else None
             if request_id is None:
-                r.delete("wotf:processing")  # type: ignore[union-attr]
+                rv.delete("wotf:processing")
-                r.delete(hb_key)  # type: ignore[union-attr]
+                rv.delete(hb_key)
                 time.sleep(0.5)
                 continue
 
-            r.set(f"wotf:req:{request_id}:status", "waiting_confirmation")  # type: ignore[union-attr]
+            rv.set(f"wotf:req:{request_id}:status", "waiting_confirmation")
-            r.delete(f"wotf:stream:{request_id}:chunks")  # type: ignore[union-attr]
+            rv.delete(f"wotf:stream:{request_id}:chunks")
 
-            r.rpush(f"wotf:stream:{request_id}:chunks", f"[ID:{request_id}]")  # type: ignore[union-attr]
+            rv.rpush(f"wotf:stream:{request_id}:chunks", f"[ID:{request_id}]")
-            r.rpush(f"wotf:stream:{request_id}:chunks", "[STATUS:waiting_confirmation]")  # type: ignore[union-attr]
+            rv.rpush(f"wotf:stream:{request_id}:chunks", "[STATUS:waiting_confirmation]")
 
-            r.delete("wotf:processing")  # type: ignore[union-attr]
+            rv.delete("wotf:processing")
-            r.delete(hb_key)  # type: ignore[union-attr]
+            rv.delete(hb_key)
 
             confirmed = False
-            for _ in range(150):
+            for _ in range(CONFIRMATION_TIMEOUT_ITERATIONS):
-                time.sleep(0.1)
+                time.sleep(CONFIRMATION_POLL_INTERVAL)
-                if r.get(f"wotf:req:{request_id}:confirmed"):  # type: ignore[union-attr]
+                if rv.get(f"wotf:req:{request_id}:confirmed"):
                     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]
+                rv.set(f"wotf:req:{request_id}:status", "error")
-                r.set(f"wotf:req:{request_id}:error", "Worker did not receive confirmation from client")  # type: ignore[union-attr]
+                rv.set(f"wotf:req:{request_id}:error", "Worker did not receive confirmation from client")
-                r.rpush(f"wotf:stream:{request_id}:chunks", "[STATUS:Connection to worker failed]")  # type: ignore[union-attr]
+                rv.rpush(f"wotf:stream:{request_id}:chunks", "[STATUS:Connection to worker failed]")
-                r.rpush(f"wotf:stream:{request_id}:chunks", "[ERROR]")  # type: ignore[union-attr]
+                rv.rpush(f"wotf:stream:{request_id}:chunks", "[ERROR]")
-                r.delete(f"wotf:req:{request_id}:confirmed")  # type: ignore[union-attr]
+                rv.delete(f"wotf:req:{request_id}:confirmed")
-                r.expire(f"wotf:req:{request_id}:ctx", 30)  # type: ignore[union-attr]
+                rv.expire(f"wotf:req:{request_id}:ctx", REQUEST_CTX_TTL_S)
                 time.sleep(0.5)
                 continue
 
-            acquired = r.set("wotf:processing", worker_id, ex=120, nx=True)  # type: ignore[union-attr]
+            acquired = False
+            for _ in range(5):
+                if rv.set("wotf:processing", worker_id, ex=PROCESSING_LOCK_TTL_S, nx=True):
+                    acquired = True
+                    break
+                time.sleep(0.5)
             if not acquired:
-                time.sleep(0.5)
+                logger.error(f"[worker] Failed to acquire lock for {request_id} after retries")
+                rv.set(f"wotf:req:{request_id}:status", "pending")
+                rv.delete(f"wotf:stream:{request_id}:chunks")
+                rv.delete(f"wotf:req:{request_id}:confirmed")
+                rv.rpush("wotf:queue", request_id)
                 continue
 
-            r.set(hb_key, worker_id, ex=15)  # type: ignore[union-attr]
+            rv.set(hb_key, worker_id, ex=HB_KEY_TTL_S)
-            r.set(f"wotf:req:{request_id}:status", "processing")  # type: ignore[union-attr]
+            rv.set(f"wotf:req:{request_id}:status", "processing")
-            r.set(f"wotf:req:{request_id}:last_heartbeat", str(time.time()))  # type: ignore[union-attr]
+            rv.set(f"wotf:req:{request_id}:last_heartbeat", "0")
 
-            ctx_raw = r.hgetall(f"wotf:req:{request_id}:ctx")  # type: ignore[union-attr]
+            ctx_raw = rv.hgetall(f"wotf:req:{request_id}:ctx")
             ctx: dict[str, Any] = {}
-            if isinstance(ctx_raw, dict):
+            for k_str, v_raw in ctx_raw.items():
-                for k_str, v_raw in ctx_raw.items():  # type: ignore[union-attr]
+                k: str = k_str
-                    k: str = str(k_str)  # type: ignore[arg-type]
+                v: Any = v_raw
-                    v: Any = v_raw  # type: ignore[assignment]
+                try:
-                    try:
+                    ctx[k] = json.loads(v) if isinstance(v, str) else v
-                        ctx[k] = json.loads(v) if isinstance(v, str) else v
+                except json.JSONDecodeError:
-                    except json.JSONDecodeError:
                         ctx[k] = v
 
             start = time.time()
@@ -701,105 +761,107 @@ def worker_loop() -> None:
             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]
+                        rv.set(hb_key, worker_id, ex=HB_KEY_TTL_S)
                         hb_start = time.time()
 
-                    if time.time() - hb_check_start >= 2:
+                    if time.time() - hb_check_start >= HEARTBEAT_CHECK_INTERVAL_S:
                         hb_check_start = time.time()
-                        last_hb_raw = r.get(f"wotf:req:{request_id}:last_heartbeat")  # type: ignore[union-attr]
+                        last_hb_raw = rv.get(f"wotf:req:{request_id}:last_heartbeat")
-                        if last_hb_raw:
+                        if last_hb_raw is None:
-                            last_hb = float(last_hb_raw)  # type: ignore[arg-type]
+                            logger.info(f"[worker] Heartbeat key missing, aborting {request_id}")
-                            if time.time() - last_hb > 10:
+                            rv.set(f"wotf:req:{request_id}:status", "aborted")
-                                logger.info(f"[worker] Client disconnected (heartbeat stale), aborting {request_id}")
+                            rv.rpush(f"wotf:stream:{request_id}:chunks", "[ERROR]")
-                                r.set(f"wotf:req:{request_id}:status", "aborted")  # type: ignore[union-attr]
+                            rv.delete("wotf:processing")
-                                r.rpush(f"wotf:stream:{request_id}:chunks", "[ERROR]")  # type: ignore[union-attr]
+                            rv.delete(hb_key)
-                                r.delete("wotf:processing")  # type: ignore[union-attr]
+                            aborted = True
-                                r.delete(hb_key)  # type: ignore[union-attr]
+                            break
-                                aborted = True
+                        last_hb = float(last_hb_raw)  # type: ignore[arg-type]
-                                break
+                        if last_hb > 0 and time.time() - last_hb > HEARTBEAT_TIMEOUT_S:
-                        else:
+                            logger.info(f"[worker] Client disconnected (heartbeat stale), aborting {request_id}")
-                            logger.info(f"[worker] No heartbeat from client, aborting {request_id}")
+                            rv.set(f"wotf:req:{request_id}:status", "aborted")
-                            r.set(f"wotf:req:{request_id}:status", "aborted")  # type: ignore[union-attr]
+                            rv.rpush(f"wotf:stream:{request_id}:chunks", "[ERROR]")
-                            r.rpush(f"wotf:stream:{request_id}:chunks", "[ERROR]")  # type: ignore[union-attr]
+                            rv.delete("wotf:processing")
-                            r.delete("wotf:processing")  # type: ignore[union-attr]
+                            rv.delete(hb_key)
-                            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]
+                    rv.rpush(f"wotf:stream:{request_id}:chunks", encoded)
             except BaseException as e:  # noqa: BLE001
                 elapsed = time.time() - start
                 logger.error(f"[worker] stream_llm failed for {request_id} in {elapsed:.1f}s: {e}")
-                retries_raw = r.incr(f"wotf:req:{request_id}:retries")  # type: ignore[union-attr]
+                retries = rv.incr(f"wotf:req:{request_id}:retries")
-                retries: int = int(retries_raw) if isinstance(retries_raw, (int, float)) else 1
                 if retries <= 2:
                     logger.info(f"[worker] Requeuing {request_id} (attempt {retries}/2)")
-                    r.set(f"wotf:req:{request_id}:status", "pending")  # type: ignore[union-attr]
+                    rv.set(f"wotf:req:{request_id}:status", "pending")
-                    r.delete(f"wotf:stream:{request_id}:chunks")  # type: ignore[union-attr]
+                    rv.delete(f"wotf:stream:{request_id}:chunks")
-                    r.rpush("wotf:queue", request_id)  # type: ignore[union-attr]
+                    rv.rpush("wotf:queue", request_id)
                 else:
                     logger.error(f"[worker] {request_id} failed after {retries} attempts")
-                    r.set(f"wotf:req:{request_id}:status", "error")  # type: ignore[union-attr]
+                    rv.set(f"wotf:req:{request_id}:status", "error")
-                    r.set(f"wotf:req:{request_id}:error", str(e))  # type: ignore[union-attr]
+                    rv.set(f"wotf:req:{request_id}:error", str(e))
-                r.delete("wotf:processing")  # type: ignore[union-attr]
+                rv.delete("wotf:processing")
-                r.delete(hb_key)  # type: ignore[union-attr]
+                rv.delete(hb_key)
                 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]
+            rv.rpush(f"wotf:stream:{request_id}:chunks", "[STATUS:Validating HTML structure...]")
             errors = validate_html(html)
             if errors:
                 short = summarize_errors(errors)
-                r.rpush(f"wotf:stream:{request_id}:chunks", f"[STATUS:Found {len(errors)} issue(s): {short}. Fixing (attempt 1/2)...]")  # type: ignore[union-attr]
+                rv.rpush(f"wotf:stream:{request_id}:chunks", f"[STATUS:Found {len(errors)} issue(s): {short}. Fixing (attempt 1/2)...]")
                 html = fix_html(html, errors)
                 errors = validate_html(html)
                 if errors:
                     short = summarize_errors(errors)
-                    r.rpush(f"wotf:stream:{request_id}:chunks", "[STATUS:Still invalid. Fixing (attempt 2/2)...]")  # type: ignore[union-attr]
+                    rv.rpush(f"wotf:stream:{request_id}:chunks", "[STATUS:Still invalid. Fixing (attempt 2/2)...]")
                     html = fix_html(html, errors)
                     errors = validate_html(html)
                 if errors:
                     short = summarize_errors(errors)
-                    r.rpush(f"wotf:stream:{request_id}:chunks", f"[STATUS:Still has {len(errors)} issue(s): {short}. Serving as-is.]")  # type: ignore[union-attr]
+                    rv.rpush(f"wotf:stream:{request_id}:chunks", f"[STATUS:Still has {len(errors)} issue(s): {short}. Serving as-is.]")
                 else:
-                    r.rpush(f"wotf:stream:{request_id}:chunks", "[STATUS:Fixed successfully.]")  # type: ignore[union-attr]
+                    rv.rpush(f"wotf:stream:{request_id}:chunks", "[STATUS:Fixed successfully.]")
-                r.rpush(f"wotf:stream:{request_id}:chunks", "[CLEAR]")  # type: ignore[union-attr]
+                rv.rpush(f"wotf:stream:{request_id}:chunks", "[CLEAR]")
-                for i in range(0, len(html), 1000):
+                for i in range(0, len(html), CHUNK_SIZE):
-                    chunk = html[i : i + 1000]
+                    chunk = html[i : i + CHUNK_SIZE]
                     encoded = base64.b64encode(chunk.encode("utf-8")).decode("ascii")
-                    r.rpush(f"wotf:stream:{request_id}:chunks", encoded)  # type: ignore[union-attr]
+                    rv.rpush(f"wotf:stream:{request_id}:chunks", encoded)
             else:
-                r.rpush(f"wotf:stream:{request_id}:chunks", "[STATUS:Validation passed.]")  # type: ignore[union-attr]
+                rv.rpush(f"wotf:stream:{request_id}:chunks", "[STATUS:Validation passed.]")
-            r.rpush(f"wotf:stream:{request_id}:chunks", "[RENDER]")  # type: ignore[union-attr]
+            rv.rpush(f"wotf:stream:{request_id}:chunks", "[RENDER]")
             elapsed = time.time() - start
-            r.set(f"wotf:req:{request_id}:status", "done")  # type: ignore[union-attr]
+            rv.set(f"wotf:req:{request_id}:status", "done")
 
-            n_raw = r.get("wotf:stats:gen_count")  # type: ignore[union-attr]
+            n_raw = rv.get("wotf:stats:gen_count")
             n_str: str = str(n_raw) if n_raw is not None else "0"
             n = int(n_str)
-            avg_raw = r.get("wotf:stats:avg_time")  # type: ignore[union-attr]
+            avg_raw = rv.get("wotf:stats:avg_time")
             avg_str: str = str(avg_raw) if avg_raw is not None else "15"
             avg = float(avg_str)
             new_avg = (avg * n + elapsed) / (n + 1) if n > 0 else elapsed
-            r.set("wotf:stats:avg_time", str(new_avg))  # type: ignore[union-attr]
+            rv.set("wotf:stats:avg_time", str(new_avg))
-            r.set("wotf:stats:gen_count", str(n + 1))  # type: ignore[union-attr]
+            rv.set("wotf:stats:gen_count", str(n + 1))
 
-            r.delete("wotf:processing")  # type: ignore[union-attr]
+            rv.delete("wotf:processing")
-            r.delete(hb_key)  # type: ignore[union-attr]
+            rv.delete(hb_key)
             logger.info(f"[worker] Generated {request_id} in {elapsed:.1f}s (avg={new_avg:.1f}s)")
 
+            if time.time() - lock_hb_start >= WORKER_LOCK_HEARTBEAT_INTERVAL_S:
+                lock_hb_start = time.time()
+                rv.set("wotf:worker_lock", "1", xx=True, ex=WORKER_LOCK_TTL_S)
+
         except BaseException as e:  # noqa: BLE001
             logger.error(f"[worker] Error: {e}")
             if request_id:
-                r.set(f"wotf:req:{request_id}:status", "error")  # type: ignore[union-attr]
+                rv.set(f"wotf:req:{request_id}:status", "error")
-                r.set(f"wotf:req:{request_id}:error", str(e))  # type: ignore[union-attr]
+                rv.set(f"wotf:req:{request_id}:error", str(e))
-            r.delete("wotf:processing")  # type: ignore[union-attr]
+            rv.delete("wotf:processing")
-            r.delete(hb_key)  # type: ignore[union-attr]
+            rv.delete(hb_key)
             time.sleep(1)
 
 
@@ -894,52 +956,27 @@ def stream_llm(context: dict[str, Any]) -> Generator[str]:
         yield f"\n\n<!-- Error: {e} -->"
 
 
-def sse_stream(context: dict[str, Any]) -> Generator[str]:
+@app.route("/health")
-    buffer: list[str] = []
+def health() -> Response:
+    try:
+        valkey_ok = rv.get("wotf:health_check") is not None or rv.set("wotf:health_check", "1", ex=60)
+    except BaseException:  # noqa: BLE001
+        valkey_ok = False
 
-    def status(msg: str) -> str:
+    queue_depth = rv.llen("wotf:queue")
-        return f"data: [STATUS:{msg}]\n\n"
+    worker_active = rv.exists("wotf:worker_lock")
 
-    yield status("Generating HTML...")
+    status = "ok" if (valkey_ok and worker_active) else "degraded"
-    for chunk in stream_llm(context):
-        buffer.append(chunk)
-        encoded = base64.b64encode(chunk.encode("utf-8")).decode("ascii")
-        yield f"data: {encoded}\n\n"
 
-    html = "".join(buffer)
+    return Response(
-    yield "data: [DONE]\n\n"
+        json.dumps({
-
+            "status": status,
-    yield status("Validating HTML structure...")
+            "valkey_connected": valkey_ok,
-    errors = validate_html(html)
+            "worker_active": worker_active,
-
+            "queue_depth": queue_depth,
-    if errors:
+        }),
-        short = summarize_errors(errors)
+        mimetype="application/json",
-        yield status(f"Found {len(errors)} issue(s): {short}. Fixing (attempt 1/2)...")
+    )
-
-        html = fix_html(html, errors)
-        errors = validate_html(html)
-
-        if errors:
-            short = summarize_errors(errors)
-            yield status("Still invalid. Fixing (attempt 2/2)...")
-            html = fix_html(html, errors)
-            errors = validate_html(html)
-
-        if errors:
-            short = summarize_errors(errors)
-            yield status(f"Still has {len(errors)} issue(s): {short}. Serving as-is.")
-        else:
-            yield status("Fixed successfully.")
-
-        yield "data: [CLEAR]\n\n"
-        for i in range(0, len(html), 1000):
-            chunk = html[i : i + 1000]
-            encoded = base64.b64encode(chunk.encode("utf-8")).decode("ascii")
-            yield f"data: {encoded}\n\n"
-    else:
-        yield status("Validation passed.")
-
-    yield "data: [RENDER]\n\n"
 
 
 @app.route("/", defaults={"path": ""})
@@ -962,41 +999,41 @@ def stream() -> Response:
 
         last_idx = 0
         while True:
-            status_raw = r.get(f"wotf:req:{request_id}:status")  # type: ignore[union-attr]
+            status_raw = rv.get(f"wotf:req:{request_id}:status")
             status = status_raw if isinstance(status_raw, str) else "pending"
 
             if status == "error":
-                error_raw = r.get(f"wotf:req:{request_id}:error")  # type: ignore[union-attr]
+                error_raw = rv.get(f"wotf:req:{request_id}:error")
                 error_msg = error_raw if isinstance(error_raw, str) else "Unknown error"
-                yield f"data: [STATUS:Error: {error_msg}\n\n"
+                yield f"data: [STATUS:Error: {error_msg}]\n\n"
                 yield "data: [ERROR]\n\n"
                 return
 
             if status == "waiting_confirmation":
-                chunks_raw = r.lrange(f"wotf:stream:{request_id}:chunks", 0, -1)  # type: ignore[union-attr]
+                chunks_raw = rv.lrange(f"wotf:stream:{request_id}:chunks", 0, -1)
                 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"
                 else:
-                    yield "data: [STATUS:Waiting for your connection...\n\n"
+                    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_raw = rv.lrange(f"wotf:stream:{request_id}:chunks", last_idx, -1)
                 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"
+                    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_raw = rv.lrange(f"wotf:stream:{request_id}:chunks", last_idx, -1)
                 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"
@@ -1006,13 +1043,13 @@ def stream() -> Response:
             pos = info["position"]
             eta = info["eta"]
             if pos == 0:
-                yield "data: [STATUS:Your turn is next...\n\n"
+                yield "data: [STATUS:Your turn is next...]\n\n"
             else:
-                yield f"data: [STATUS:Position {pos} in line | Est. wait: ~{eta:.0f}s\n\n"
+                yield f"data: [STATUS:Position {pos} in line | Est. wait: ~{eta:.0f}s]\n\n"
             time.sleep(1)
 
         while True:
-            chunks_raw = r.lrange(f"wotf:stream:{request_id}:chunks", last_idx, -1)  # type: ignore[union-attr]
+            chunks_raw = rv.lrange(f"wotf:stream:{request_id}:chunks", last_idx, -1)
             chunks: list[str] = list(chunks_raw) if isinstance(chunks_raw, list) else []  # type: ignore[arg-type]
             if chunks:
                 for chunk in chunks:
@@ -1023,23 +1060,23 @@ def stream() -> Response:
                 last_idx += len(chunks)
                 time.sleep(0.05)
             else:
-                status_raw = r.get(f"wotf:req:{request_id}:status")  # type: ignore[union-attr]
+                status_raw = rv.get(f"wotf:req:{request_id}:status")
                 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_raw = rv.lrange(f"wotf:stream:{request_id}:chunks", last_idx, -1)
                     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]
+                    error_raw = rv.get(f"wotf:req:{request_id}:error")
                     error_msg = error_raw if isinstance(error_raw, str) else "Unknown error"
-                    yield f"data: [STATUS:Error: {error_msg}\n\n"
+                    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: [STATUS:Generation aborted (connection lost)]\n\n"
                     yield "data: [ERROR]\n\n"
                     return
                 time.sleep(0.2)
@@ -1057,19 +1094,19 @@ 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]
+    ctx = rv.exists(f"wotf:req:{request_id}:ctx")
     if not ctx:
         return Response("not found", status=404)
-    r.set(f"wotf:req:{request_id}:confirmed", "1")  # type: ignore[union-attr]
+    rv.set(f"wotf:req:{request_id}:confirmed", "1")
     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]
+    status = rv.get(f"wotf:req:{request_id}:status")
     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]
+    rv.set(f"wotf:req:{request_id}:last_heartbeat", str(time.time()))
     return Response("ok", status=200)
 
 
@@ -1077,7 +1114,7 @@ def start_worker_if_needed() -> None:
     global worker_running, worker_thread
     if worker_running:
         return
-    claimed = r.set("wotf:worker_lock", "1", nx=True, ex=300)  # type: ignore[union-attr]
+    claimed = rv.set("wotf:worker_lock", "1", nx=True, ex=WORKER_LOCK_TTL_S)
     if claimed:
         worker_running = True
         if os.getenv("GUNICORN_WORKER_ID"):
@@ -1097,7 +1134,21 @@ def gunicorn_post_fork(server: Any, worker: Any) -> None:
     start_worker_if_needed()
 
 
+def graceful_shutdown(signum: int, frame: Any) -> None:
+    global worker_running
+    sig_name = "SIGINT" if signum == signal.SIGINT else "SIGTERM"
+    logger.info(f"[shutdown] Received {sig_name}, stopping worker...")
+    worker_running = False
+    rv.delete("wotf:worker_lock")
+    rv.delete("wotf:processing")
+    rv.delete("wotf:processing:hb")
+    logger.info("[shutdown] Locks released, exiting")
+    sys.exit(0)
+
+
 if __name__ == "__main__":
+    signal.signal(signal.SIGINT, graceful_shutdown)
+    signal.signal(signal.SIGTERM, graceful_shutdown)
     worker_running = True
     worker_thread = threading.Thread(target=worker_loop, daemon=True)
     worker_thread.start()

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "web-on-the-fly"
 version = "0.1.0"
-description = "Add your description here"
+description = "Flask app that streams AI-generated fake websites via SSE"
 readme = "README.md"
 requires-python = ">=3.13"
 dependencies = [