Experimental/Web-on-the-Fly
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix heartbeat issue
Build and Push Containers / test (push) Successful in 9s
Details
Build and Push Containers / build-and-push (push) Successful in 16s
Details

Browse Source
This commit is contained in:
ducoterra
2026-08-02 11:35:13 -04:00
parent 121ecece3c
commit a2c4f81147
3 changed files with 247 additions and 195 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app.py
+245 -194
View File
@@ -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 <html> tag, found {len(html_tags)}")
if len(body_tags) != 1:
errors.append(f"Expected 1 <body> tag, found {len(body_tags)}")
if len(head_tags) != 1:
errors.append(f"Expected 1 <head> 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 <html> tag, found {len(html_tags)}") # type: ignore[arg-type]
if len(body_tags) != 1: # type: ignore[arg-type]
errors.append(f"Expected 1 <body> tag, found {len(body_tags)}") # type: ignore[arg-type]
if len(head_tags) != 1: # type: ignore[arg-type]
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)
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"):
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<code>\n{js_text}\n</code>"
)
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<!-- Error: {e} -->"
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/<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()