Add job cancel support

app/main.py

@@ -1,5 +1,6 @@
 import json
 import os
+import signal
 import socket
 import sqlite3
 import subprocess
@@ -15,6 +16,7 @@ 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"))
+CANCEL_CHECK_SECONDS = int(os.getenv("APPFACTORY_CANCEL_CHECK_SECONDS", "2"))
 
 app = FastAPI(title="AppFactory Worker")
 
@@ -42,6 +44,20 @@ def append_job_log(job_id: int, stream: str, message: str):
     con.close()
 
 
+def get_job_status(job_id: int) -> str | None:
+    con = get_connection()
+    row = con.execute(
+        "SELECT status FROM jobs WHERE id = ?",
+        (job_id,),
+    ).fetchone()
+    con.close()
+
+    if not row:
+        return None
+
+    return row["status"]
+
+
 def update_worker_status(current_job_id=None):
     con = get_connection()
 
@@ -77,6 +93,7 @@ def update_worker_status(current_job_id=None):
                 {
                     "poll_seconds": POLL_SECONDS,
                     "heartbeat_seconds": HEARTBEAT_SECONDS,
+                    "cancel_check_seconds": CANCEL_CHECK_SECONDS,
                 },
                 ensure_ascii=False,
             ),
@@ -104,7 +121,7 @@ def recover_stale_jobs():
         """
         SELECT id
         FROM jobs
-        WHERE status = 'running'
+        WHERE status IN ('running', 'cancelled_requested')
           AND started_at < datetime('now', '-15 minutes')
         """
     ).fetchall()
@@ -197,7 +214,7 @@ def get_next_job():
         """
         SELECT id
         FROM jobs
-        WHERE status = 'running'
+        WHERE status IN ('running', 'cancelled_requested')
           AND target_type = ?
           AND target_id = ?
           AND id != ?
@@ -257,6 +274,22 @@ def command_for_job(job: dict):
     raise ValueError(f"Unsupported job type: {job_type}")
 
 
+def terminate_process(process: subprocess.Popen, job_id: int):
+    append_job_log(job_id, "system", "Cancel requested. Terminating process...")
+
+    try:
+        process.terminate()
+        process.wait(timeout=10)
+        append_job_log(job_id, "system", "Process terminated.")
+    except subprocess.TimeoutExpired:
+        append_job_log(job_id, "system", "Process did not terminate in time. Killing process...")
+        process.kill()
+        process.wait(timeout=10)
+        append_job_log(job_id, "system", "Process killed.")
+    except Exception as exc:
+        append_job_log(job_id, "stderr", f"Error while terminating process: {exc}")
+
+
 def execute_job(job: dict):
     job_id = int(job["id"])
     target = job["target_id"]
@@ -277,8 +310,38 @@ def execute_job(job: dict):
             stderr=subprocess.PIPE,
             text=True,
             env=env,
+            start_new_session=True,
         )
 
+        while process.poll() is None:
+            status = get_job_status(job_id)
+
+            if status == "cancelled_requested":
+                try:
+                    os.killpg(os.getpgid(process.pid), signal.SIGTERM)
+                except Exception:
+                    terminate_process(process, job_id)
+                else:
+                    try:
+                        process.wait(timeout=10)
+                    except subprocess.TimeoutExpired:
+                        try:
+                            os.killpg(os.getpgid(process.pid), signal.SIGKILL)
+                        except Exception:
+                            process.kill()
+                        process.wait(timeout=10)
+
+                update_job_status(
+                    job_id,
+                    "cancelled",
+                    error_text="Job cancelled by user request",
+                    result={"cancelled": True},
+                )
+                append_job_log(job_id, "system", "Job cancelled.")
+                return
+
+            time.sleep(CANCEL_CHECK_SECONDS)
+
         stdout, stderr = process.communicate()
         returncode = process.returncode
 
@@ -288,6 +351,16 @@ def execute_job(job: dict):
         if stderr:
             append_job_log(job_id, "stderr", stderr)
 
+        if get_job_status(job_id) == "cancelled_requested":
+            update_job_status(
+                job_id,
+                "cancelled",
+                error_text="Job cancelled by user request",
+                result={"returncode": returncode, "cancelled": True},
+            )
+            append_job_log(job_id, "system", "Job cancelled after process finished.")
+            return
+
         if returncode == 0:
             update_job_status(
                 job_id,
@@ -361,4 +434,5 @@ def health():
         "worker_id": WORKER_ID,
         "poll_seconds": POLL_SECONDS,
         "heartbeat_seconds": HEARTBEAT_SECONDS,
+        "cancel_check_seconds": CANCEL_CHECK_SECONDS,
     }