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