From a2c4f81147032ba445587fada16091c15cd32a85 Mon Sep 17 00:00:00 2001 From: ducoterra Date: Sun, 2 Aug 2026 11:35:13 -0400 Subject: [PATCH] fix heartbeat issue --- README.md | 1 + app.py | 439 +++++++++++++++++++++++++++---------------------- pyproject.toml | 2 +- 3 files changed, 247 insertions(+), 195 deletions(-) diff --git a/README.md b/README.md index c7764ac..b3b4114 100644 --- a/README.md +++ b/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 diff --git a/app.py b/app.py index 0f8371c..3ba13de 100644 --- a/app.py +++ b/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) - html_tags: list[Any] = doc.xpath("//html") - body_tags: list[Any] = doc.xpath("//body") - head_tags: list[Any] = doc.xpath("//head") - if len(html_tags) != 1: - errors.append(f"Expected 1 tag, found {len(html_tags)}") - if len(body_tags) != 1: - errors.append(f"Expected 1 tag, found {len(body_tags)}") - if len(head_tags) != 1: - errors.append(f"Expected 1 tag, found {len(head_tags)}") + doc: Any = lxml_html.fromstring(html) # type: ignore[union-attr] + html_tags: list[Any] = doc.xpath("//html") # type: ignore[assignment,union-attr] + body_tags: list[Any] = doc.xpath("//body") # type: ignore[assignment,union-attr] + head_tags: list[Any] = doc.xpath("//head") # type: ignore[assignment,union-attr] + if len(html_tags) != 1: # type: ignore[arg-type] + errors.append(f"Expected 1 tag, found {len(html_tags)}") # type: ignore[arg-type] + if len(body_tags) != 1: # type: ignore[arg-type] + errors.append(f"Expected 1 tag, found {len(body_tags)}") # type: ignore[arg-type] + if len(head_tags) != 1: # type: ignore[arg-type] + errors.append(f"Expected 1 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) - 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\n{js_text}\n" - ) - if review and "no issue" not in review.lower(): - for line in review.strip().split("\n"): - line = line.strip() - if line and not line.startswith("<"): - errors.append(f"JS review: {line}") + try: + js_text = "\n".join(script_blocks) + 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\n{js_text}\n" + ) + if review and "no issue" not in review.lower(): + for line in review.strip().split("\n"): + line = line.strip() + 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] - r.set(f"wotf:req:{request_id}:status", "pending") # type: ignore[union-attr] - r.rpush("wotf:queue", request_id) # type: ignore[union-attr] - position_raw = r.llen("wotf:queue") # type: ignore[union-attr] - position: int = int(position_raw) if isinstance(position_raw, (int, float)) else 0 + rv.hset(ctx_key, mapping={k: json.dumps(v) for k, v in context.items()}) + rv.set(f"wotf:req:{request_id}:status", "pending") + rv.rpush("wotf:queue", request_id) + position = rv.llen("wotf:queue") 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: list[str] = list(queue_items_raw) if isinstance(queue_items_raw, list) else [] # type: ignore[arg-type] + queue_items = rv.lrange("wotf:queue", 0, -1) 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: int = int(hb_ttl_raw) if isinstance(hb_ttl_raw, (int, float)) else -2 + hb_ttl = rv.ttl(hb_key) if hb_ttl == -2: logger.info(f"[worker] Reclaiming stale lock from {lock_holder}") - r.delete("wotf:processing") # type: ignore[union-attr] - r.delete(hb_key) # type: ignore[union-attr] + rv.delete("wotf:processing") + 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] - 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 + request_id = rv.lpop("wotf:queue") if request_id is None: - r.delete("wotf:processing") # type: ignore[union-attr] - r.delete(hb_key) # type: ignore[union-attr] + rv.delete("wotf:processing") + rv.delete(hb_key) time.sleep(0.5) continue - 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] + rv.set(f"wotf:req:{request_id}:status", "waiting_confirmation") + rv.delete(f"wotf:stream:{request_id}:chunks") - 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] + rv.rpush(f"wotf:stream:{request_id}:chunks", f"[ID:{request_id}]") + rv.rpush(f"wotf:stream:{request_id}:chunks", "[STATUS:waiting_confirmation]") - r.delete("wotf:processing") # type: ignore[union-attr] - r.delete(hb_key) # type: ignore[union-attr] + rv.delete("wotf:processing") + rv.delete(hb_key) confirmed = False - for _ in range(150): - time.sleep(0.1) - if r.get(f"wotf:req:{request_id}:confirmed"): # type: ignore[union-attr] + for _ in range(CONFIRMATION_TIMEOUT_ITERATIONS): + time.sleep(CONFIRMATION_POLL_INTERVAL) + 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] - 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] + rv.set(f"wotf:req:{request_id}:status", "error") + rv.set(f"wotf:req:{request_id}:error", "Worker did not receive confirmation from client") + rv.rpush(f"wotf:stream:{request_id}:chunks", "[STATUS:Connection to worker failed]") + rv.rpush(f"wotf:stream:{request_id}:chunks", "[ERROR]") + rv.delete(f"wotf:req:{request_id}:confirmed") + 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] - 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] + rv.set(hb_key, worker_id, ex=HB_KEY_TTL_S) + rv.set(f"wotf:req:{request_id}:status", "processing") + 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(): # type: ignore[union-attr] - k: str = str(k_str) # type: ignore[arg-type] - v: Any = v_raw # type: ignore[assignment] - try: - ctx[k] = json.loads(v) if isinstance(v, str) else v - except json.JSONDecodeError: + for k_str, v_raw in ctx_raw.items(): + k: str = k_str + v: Any = v_raw + try: + ctx[k] = json.loads(v) if isinstance(v, str) else v + 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] - 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] + last_hb_raw = rv.get(f"wotf:req:{request_id}:last_heartbeat") + if last_hb_raw is None: + logger.info(f"[worker] Heartbeat key missing, aborting {request_id}") + rv.set(f"wotf:req:{request_id}:status", "aborted") + rv.rpush(f"wotf:stream:{request_id}:chunks", "[ERROR]") + rv.delete("wotf:processing") + rv.delete(hb_key) + aborted = True + break + last_hb = float(last_hb_raw) # type: ignore[arg-type] + if last_hb > 0 and time.time() - last_hb > HEARTBEAT_TIMEOUT_S: + logger.info(f"[worker] Client disconnected (heartbeat stale), aborting {request_id}") + rv.set(f"wotf:req:{request_id}:status", "aborted") + rv.rpush(f"wotf:stream:{request_id}:chunks", "[ERROR]") + rv.delete("wotf:processing") + rv.delete(hb_key) 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: int = int(retries_raw) if isinstance(retries_raw, (int, float)) else 1 + retries = rv.incr(f"wotf:req:{request_id}:retries") 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] - r.delete(f"wotf:stream:{request_id}:chunks") # type: ignore[union-attr] - r.rpush("wotf:queue", request_id) # type: ignore[union-attr] + rv.set(f"wotf:req:{request_id}:status", "pending") + rv.delete(f"wotf:stream:{request_id}:chunks") + 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] - r.set(f"wotf:req:{request_id}:error", str(e)) # type: ignore[union-attr] - r.delete("wotf:processing") # type: ignore[union-attr] - r.delete(hb_key) # type: ignore[union-attr] + rv.set(f"wotf:req:{request_id}:status", "error") + rv.set(f"wotf:req:{request_id}:error", str(e)) + rv.delete("wotf:processing") + 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] - r.rpush(f"wotf:stream:{request_id}:chunks", "[CLEAR]") # type: ignore[union-attr] - for i in range(0, len(html), 1000): - chunk = html[i : i + 1000] + rv.rpush(f"wotf:stream:{request_id}:chunks", "[STATUS:Fixed successfully.]") + rv.rpush(f"wotf:stream:{request_id}:chunks", "[CLEAR]") + for i in range(0, len(html), CHUNK_SIZE): + 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] - r.rpush(f"wotf:stream:{request_id}:chunks", "[RENDER]") # type: ignore[union-attr] + rv.rpush(f"wotf:stream:{request_id}:chunks", "[STATUS:Validation passed.]") + 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] - r.set("wotf:stats:gen_count", str(n + 1)) # type: ignore[union-attr] + rv.set("wotf:stats:avg_time", str(new_avg)) + rv.set("wotf:stats:gen_count", str(n + 1)) - r.delete("wotf:processing") # type: ignore[union-attr] - r.delete(hb_key) # type: ignore[union-attr] + rv.delete("wotf:processing") + 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] - r.set(f"wotf:req:{request_id}:error", str(e)) # type: ignore[union-attr] - r.delete("wotf:processing") # type: ignore[union-attr] - r.delete(hb_key) # type: ignore[union-attr] + rv.set(f"wotf:req:{request_id}:status", "error") + rv.set(f"wotf:req:{request_id}:error", str(e)) + rv.delete("wotf:processing") + rv.delete(hb_key) time.sleep(1) @@ -894,52 +956,27 @@ def stream_llm(context: dict[str, Any]) -> Generator[str]: yield f"\n\n" -def sse_stream(context: dict[str, Any]) -> Generator[str]: - buffer: list[str] = [] +@app.route("/health") +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: - return f"data: [STATUS:{msg}]\n\n" + queue_depth = rv.llen("wotf:queue") + worker_active = rv.exists("wotf:worker_lock") - yield status("Generating HTML...") - for chunk in stream_llm(context): - buffer.append(chunk) - encoded = base64.b64encode(chunk.encode("utf-8")).decode("ascii") - yield f"data: {encoded}\n\n" + status = "ok" if (valkey_ok and worker_active) else "degraded" - html = "".join(buffer) - yield "data: [DONE]\n\n" - - yield status("Validating HTML structure...") - errors = validate_html(html) - - if errors: - short = summarize_errors(errors) - 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" + return Response( + json.dumps({ + "status": status, + "valkey_connected": valkey_ok, + "worker_active": worker_active, + "queue_depth": queue_depth, + }), + mimetype="application/json", + ) @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/", 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/", 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() diff --git a/pyproject.toml b/pyproject.toml index add17a5..bb4387c 100644 --- a/pyproject.toml +++ b/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 = [