Add worker heartbeat and stale job recovery

app/main.py

@@ -5,7 +5,6 @@ import sqlite3
 import subprocess
 import threading
 import time
-from datetime import datetime, timezone
 from pathlib import Path
 
 from fastapi import FastAPI
@@ -15,17 +14,15 @@ TOOLS_DIR = Path("/tools")
 
 WORKER_ID = os.getenv("APPFACTORY_WORKER_ID", socket.gethostname())
 POLL_SECONDS = int(os.getenv("APPFACTORY_WORKER_POLL_SECONDS", "3"))
+HEARTBEAT_SECONDS = int(os.getenv("APPFACTORY_WORKER_HEARTBEAT_SECONDS", "10"))
 
 app = FastAPI(title="AppFactory Worker")
 
 _worker_thread = None
+_heartbeat_thread = None
 _worker_started = False
 
 
-def utc_now():
-    return datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S")
-
-
 def get_connection():
     con = sqlite3.connect(DB_FILE, timeout=30)
     con.row_factory = sqlite3.Row
@@ -45,6 +42,99 @@ def append_job_log(job_id: int, stream: str, message: str):
     con.close()
 
 
+def update_worker_status(current_job_id=None):
+    con = get_connection()
+
+    con.execute(
+        """
+        INSERT INTO workers (
+            id,
+            status,
+            started_at,
+            last_seen_at,
+            current_job_id,
+            metadata_json
+        )
+        VALUES (
+            ?,
+            'online',
+            CURRENT_TIMESTAMP,
+            CURRENT_TIMESTAMP,
+            ?,
+            ?
+        )
+        ON CONFLICT(id)
+        DO UPDATE SET
+            status = 'online',
+            last_seen_at = CURRENT_TIMESTAMP,
+            current_job_id = excluded.current_job_id,
+            metadata_json = excluded.metadata_json
+        """,
+        (
+            WORKER_ID,
+            current_job_id,
+            json.dumps(
+                {
+                    "poll_seconds": POLL_SECONDS,
+                    "heartbeat_seconds": HEARTBEAT_SECONDS,
+                },
+                ensure_ascii=False,
+            ),
+        ),
+    )
+
+    con.commit()
+    con.close()
+
+
+def heartbeat_loop():
+    while True:
+        try:
+            update_worker_status()
+        except Exception as exc:
+            print(f"Heartbeat error: {exc}", flush=True)
+
+        time.sleep(HEARTBEAT_SECONDS)
+
+
+def recover_stale_jobs():
+    con = get_connection()
+
+    rows = con.execute(
+        """
+        SELECT id
+        FROM jobs
+        WHERE status = 'running'
+          AND started_at < datetime('now', '-15 minutes')
+        """
+    ).fetchall()
+
+    for row in rows:
+        job_id = row["id"]
+
+        con.execute(
+            """
+            UPDATE jobs
+            SET status = 'failed',
+                finished_at = CURRENT_TIMESTAMP,
+                error_text = 'Worker recovery: stale running job'
+            WHERE id = ?
+            """,
+            (job_id,),
+        )
+
+        con.execute(
+            """
+            INSERT INTO job_logs (job_id, stream, message, created_at)
+            VALUES (?, 'system', 'Worker recovery marked stale running job as failed', CURRENT_TIMESTAMP)
+            """,
+            (job_id,),
+        )
+
+    con.commit()
+    con.close()
+
+
 def update_job_status(job_id: int, status: str, error_text: str | None = None, result: dict | None = None):
     con = get_connection()
 
@@ -171,6 +261,7 @@ def execute_job(job: dict):
     job_id = int(job["id"])
     target = job["target_id"]
 
+    update_worker_status(current_job_id=job_id)
     append_job_log(job_id, "system", f"Worker {WORKER_ID} started job {job_id} for {target}")
     update_job_status(job_id, "running")
 
@@ -188,29 +279,23 @@ def execute_job(job: dict):
             env=env,
         )
 
-        stdout_chunks = []
-        stderr_chunks = []
-
         stdout, stderr = process.communicate()
+        returncode = process.returncode
 
         if stdout:
-            stdout_chunks.append(stdout)
             append_job_log(job_id, "stdout", stdout)
 
         if stderr:
-            stderr_chunks.append(stderr)
             append_job_log(job_id, "stderr", stderr)
 
-        returncode = process.returncode
-
         if returncode == 0:
             update_job_status(
                 job_id,
                 "success",
                 result={
                     "returncode": returncode,
-                    "stdout_length": len("".join(stdout_chunks)),
-                    "stderr_length": len("".join(stderr_chunks)),
+                    "stdout_length": len(stdout or ""),
+                    "stderr_length": len(stderr or ""),
                 },
             )
             append_job_log(job_id, "system", "Job finished successfully")
@@ -227,10 +312,15 @@ def execute_job(job: dict):
         update_job_status(job_id, "failed", error_text=str(exc))
         append_job_log(job_id, "stderr", str(exc))
 
+    finally:
+        update_worker_status(current_job_id=None)
+
 
 def worker_loop():
     while True:
         try:
+            recover_stale_jobs()
+
             job = get_next_job()
 
             if job:
@@ -246,12 +336,20 @@ def worker_loop():
 @app.on_event("startup")
 def start_worker():
     global _worker_thread
+    global _heartbeat_thread
     global _worker_started
 
     if _worker_started:
         return
 
+    recover_stale_jobs()
+    update_worker_status(current_job_id=None)
+
     _worker_started = True
+
+    _heartbeat_thread = threading.Thread(target=heartbeat_loop, daemon=True)
+    _heartbeat_thread.start()
+
     _worker_thread = threading.Thread(target=worker_loop, daemon=True)
     _worker_thread.start()
 
@@ -261,4 +359,6 @@ def health():
     return {
         "status": "ok",
         "worker_id": WORKER_ID,
+        "poll_seconds": POLL_SECONDS,
+        "heartbeat_seconds": HEARTBEAT_SECONDS,
     }