Forward-merge release/2.2 into develop

frontends/aiq_api/src/aiq_api/jobs/crypto.py

@@ -879,6 +879,25 @@ async def update_job_output(
     await job_store.update_status(job_id, status, output=encrypted_output)
 
 
+def serialize_job_output_for_storage(
+    output: BaseModel | dict[str, Any] | list[Any] | str,
+    cipher: JobContentCipher,
+) -> str:
+    """Return the exact string ``update_job_output`` would persist to job_info.output.
+
+    Encrypts in encrypted modes; otherwise mirrors NAT's JSON serialization. Used
+    by callers that need to write the output through their own (e.g. conditional)
+    UPDATE while keeping serialization identical to the normal path.
+    """
+    if not cipher.manager.config.encrypted:
+        if isinstance(output, BaseModel):
+            return output.model_dump_json(round_trip=True)
+        if isinstance(output, dict | list):
+            return json.dumps(output)
+        return output
+    return cipher.encrypt_output_json(_serialize_output_json(output))
+
+
 def read_job_output(job_id: str, stored_output: Any) -> Any:
     """Read job output, decrypting only when encrypted mode is configured."""
 

frontends/aiq_api/src/aiq_api/jobs/runner.py

@@ -28,6 +28,7 @@
 import asyncio
 import importlib
 import logging
+import threading
 import uuid
 from collections.abc import Awaitable
 from collections.abc import Callable
@@ -178,6 +179,79 @@ def check(self) -> None:
 # Interval for emitting heartbeat events
 HEARTBEAT_INTERVAL_SECONDS = 30
 
+# The ghost-job reaper treats a RUNNING job with no recent activity as a dead
+# worker. But a worker can spend minutes in pre-event initialization (config,
+# providers, tools, MCP, sandbox) before it stores its first event, so from the
+# moment the job enters RUNNING we refresh a lightweight lease — job_info's
+# updated_at, the same column the reaper falls back to for a zero-event job — on
+# this interval. A slow-but-live worker keeps its lease fresh and is not reaped;
+# a genuinely dead worker stops refreshing and its lease goes stale. Keep this
+# well under GHOST_JOB_TIMEOUT_SECONDS so a live worker refreshes several times
+# before the reaper's timeout.
+LEASE_REFRESH_INTERVAL_SECONDS = 60
+
+
+def _db_now_expr(db_url: str) -> str:
+    """Return the DB current-time SQL expression for this backend.
+
+    Accepts both ``postgresql://`` and the legacy ``postgres://`` scheme.
+    """
+    return "NOW()" if db_url.startswith(("postgresql", "postgres")) else "CURRENT_TIMESTAMP"
+
+
+def _touch_job_lease_sync(db_url: str, job_id: str) -> None:
+    """Refresh the running-job lease by bumping job_info.updated_at.
+
+    Scoped to ``status = 'running'`` so it can never resurrect the timestamp of
+    a job that has already reached a terminal state.
+    """
+    from sqlalchemy import text
+
+    from .event_store import EventStore
+
+    engine = EventStore._get_or_create_sync_engine(db_url)
+    stmt = text(
+        f"UPDATE job_info SET updated_at = {_db_now_expr(db_url)} WHERE job_id = :job_id AND status = 'running'"
+    )
+    with engine.begin() as conn:
+        conn.execute(stmt, {"job_id": job_id})
+
+
+def _write_job_success_if_running_sync(db_url: str, job_id: str, stored_output: str) -> bool:
+    """Compare-and-set the job to SUCCESS with its output, only if still RUNNING.
+
+    A single guarded ``UPDATE ... WHERE status = 'running'`` so a job the reaper
+    already moved to a terminal state (e.g. it was reaped while slow to finish)
+    is never resurrected. Returns True iff this call performed the write.
+    """
+    from sqlalchemy import text
+
+    from .event_store import EventStore
+
+    engine = EventStore._get_or_create_sync_engine(db_url)
+    stmt = text(
+        f"UPDATE job_info SET status = 'success', output = :output, updated_at = {_db_now_expr(db_url)} "
+        "WHERE job_id = :job_id AND status = 'running'"
+    )
+    with engine.begin() as conn:
+        result = conn.execute(stmt, {"output": stored_output, "job_id": job_id})
+        return (result.rowcount or 0) == 1
+
+
+def _run_lease_refresher(db_url: str, job_id: str, stop_event: threading.Event) -> None:
+    """Refresh the running-job lease on a dedicated thread until signalled.
+
+    Runs in its own OS thread — not the worker event loop — so it cannot be
+    starved by synchronous cold-start work (config load, agent import) that
+    holds the loop. ``stop_event.wait`` returns True when the job ends (exit) or
+    False on timeout (refresh, then loop).
+    """
+    while not stop_event.wait(LEASE_REFRESH_INTERVAL_SECONDS):
+        try:
+            _touch_job_lease_sync(db_url, job_id)
+        except Exception as exc:  # noqa: BLE001 - a failed lease refresh must never kill the job
+            logger.debug("Lease refresh for job %s failed: %s", job_id, exc)
+
 
 async def run_with_cancellation(
     coro,
@@ -530,6 +604,8 @@ async def run_agent_job(
     job_store: JobStore | None = None
     job_output_cipher = None
     cancellation_monitor: CancellationMonitor | None = None
+    lease_stop: threading.Event | None = None
+    lease_thread: threading.Thread | None = None
     event_store: EventStore | BatchingEventStore | None = None
     # Sandbox runtime is released on the terminal path; interrupted forces terminate() over close().
     sandbox_runtime: Any | None = None
@@ -567,6 +643,19 @@ async def run_agent_job(
 
         await job_store.update_status(job_id, JobStatus.RUNNING)
 
+        # Start refreshing the reaper lease immediately, before the slow
+        # initialization below stores any event, so a live worker in a long
+        # cold start is not mistaken for a dead one. It runs on a dedicated
+        # thread so synchronous init work can't starve it.
+        lease_stop = threading.Event()
+        lease_thread = threading.Thread(
+            target=_run_lease_refresher,
+            args=(db_url, job_id, lease_stop),
+            name=f"job-lease-{job_id}",
+            daemon=True,
+        )
+        lease_thread.start()
+
         cancellation_monitor = CancellationMonitor(
             scheduler_address=scheduler_address,
             db_url=db_url,
@@ -817,20 +906,22 @@ async def run_agent_job(
                     # Extract report and update status inside the context manager
                     # so the UI sees completion before exporter flush and cleanup
                     report = _extract_result(result)
-                    from .crypto import update_job_output
+                    from .crypto import serialize_job_output_for_storage
 
                     if job_output_cipher is None:
                         raise RuntimeError("job output cipher was not initialized")
                     # Apply caller metadata first, then set the canonical report last so a
                     # stray "report" key in output_metadata can never overwrite the real report.
                     output = {**(output_metadata or {}), "report": report}
+                    # Terminal state is immutable: write SUCCESS with a single
+                    # compare-and-set (WHERE status='running'), so if the ghost
+                    # reaper already marked this job FAILURE it is never
+                    # resurrected. Serialize/encrypt exactly as update_job_output
+                    # would, then do the guarded write.
                     try:
-                        await update_job_output(
-                            job_store,
-                            job_id,
-                            JobStatus.SUCCESS,
-                            output=output,
-                            cipher=job_output_cipher,
+                        stored_output = serialize_job_output_for_storage(output, job_output_cipher)
+                        wrote = await asyncio.get_running_loop().run_in_executor(
+                            None, _write_job_success_if_running_sync, db_url, job_id, stored_output
                         )
                     except Exception as exc:
                         logger.warning(
@@ -839,7 +930,10 @@ async def run_agent_job(
                             exc.__class__.__name__,
                         )
                         raise
-                    logger.info("Job %s completed (report: %d chars)", job_id, len(report))
+                    if wrote:
+                        logger.info("Job %s completed (report: %d chars)", job_id, len(report))
+                    else:
+                        logger.warning("Job %s already terminal; skipping success write", job_id)
 
     except asyncio.CancelledError:
         logger.info("Job %s cancelled", job_id)
@@ -890,6 +984,10 @@ async def run_agent_job(
     finally:
         # Ensure terminal-path events are not left in the batch buffer.
         await _flush_event_store(event_store, job_id=job_id)
+        if lease_stop is not None:
+            lease_stop.set()
+        if lease_thread is not None:
+            await asyncio.to_thread(lease_thread.join, 5)
         if cancellation_monitor:
             cancellation_monitor.stop()
         # Idempotent fallback for failures before a terminal branch finalized the runtime.

frontends/aiq_api/src/aiq_api/routes/jobs.py

@@ -1184,29 +1184,39 @@ def _find_stale_jobs(db_url: str, running_status: str) -> list[str]:
 
     from ..jobs.event_store import EventStore
 
+    # Ensure job_events exists so the LEFT JOIN resolves (it need not have rows).
     EventStore._ensure_table_exists(db_url)
     engine = EventStore._get_or_create_sync_engine(db_url)
     inspector = inspect(engine)
-    if not inspector.has_table("job_events"):
+    # The query is driven from job_info; without it there are no jobs to reap.
+    if not inspector.has_table("job_info"):
         return []
 
     with engine.connect() as conn:
-        if db_url.startswith("postgresql"):
+        # Drive from job_info with a LEFT JOIN so a RUNNING job that has not
+        # persisted any events yet is still considered. That is exactly the
+        # failure this reaper exists to catch: a worker that crashes/OOMs after
+        # the job is marked RUNNING but before its first event is stored leaves
+        # zero rows in job_events and would be invisible to an INNER JOIN,
+        # sticking the job in RUNNING forever. COALESCE falls back to
+        # job_info.updated_at (set when the job entered RUNNING) when there are
+        # no events, so both cases share one staleness check.
+        if db_url.startswith(("postgresql", "postgres")):
             stale_query = text(
-                "SELECT DISTINCT je.job_id FROM job_events je "
-                "INNER JOIN job_info ji ON je.job_id = ji.job_id "
+                "SELECT ji.job_id FROM job_info ji "
+                "LEFT JOIN job_events je ON je.job_id = ji.job_id "
                 "WHERE ji.status = :running_status "
-                "GROUP BY je.job_id "
-                "HAVING MAX(je.created_at) < NOW() - :timeout * INTERVAL '1 second'"
+                "GROUP BY ji.job_id, ji.updated_at "
+                "HAVING COALESCE(MAX(je.created_at), ji.updated_at) < NOW() - :timeout * INTERVAL '1 second'"
             )
             params = {"running_status": running_status, "timeout": GHOST_JOB_TIMEOUT_SECONDS}
         else:
             stale_query = text(
-                "SELECT DISTINCT je.job_id FROM job_events je "
-                "INNER JOIN job_info ji ON je.job_id = ji.job_id "
+                "SELECT ji.job_id FROM job_info ji "
+                "LEFT JOIN job_events je ON je.job_id = ji.job_id "
                 "WHERE ji.status = :running_status "
-                "GROUP BY je.job_id "
-                "HAVING MAX(je.created_at) < datetime('now', :timeout_interval)"
+                "GROUP BY ji.job_id, ji.updated_at "
+                "HAVING COALESCE(MAX(je.created_at), ji.updated_at) < datetime('now', :timeout_interval)"
             )
             params = {
                 "running_status": running_status,
@@ -1217,13 +1227,42 @@ def _find_stale_jobs(db_url: str, running_status: str) -> list[str]:
         return [row[0] for row in result]
 
 
+def _mark_job_failed_if_running(db_url: str, job_id: str, running_status: str, failure_status: str, error: str) -> bool:
+    """Atomically flip a job from RUNNING to FAILURE, only if still running.
+
+    Returns True iff this call performed the transition. The ``WHERE status =
+    running`` guard makes the write conditional in a single statement, so a job
+    that reached a terminal state (e.g. a slow worker that finished) between
+    detection and reaping is never clobbered.
+    """
+    from sqlalchemy import text
+
+    from ..jobs.event_store import EventStore
+
+    engine = EventStore._get_or_create_sync_engine(db_url)
+    now_expr = "NOW()" if db_url.startswith(("postgresql", "postgres")) else "CURRENT_TIMESTAMP"
+    stmt = text(
+        f"UPDATE job_info SET status = :failure, error = :error, updated_at = {now_expr} "
+        "WHERE job_id = :job_id AND status = :running"
+    )
+    with engine.begin() as conn:
+        result = conn.execute(
+            stmt,
+            {"failure": failure_status, "error": error, "job_id": job_id, "running": running_status},
+        )
+        return (result.rowcount or 0) == 1
+
+
 async def _reap_ghost_jobs(job_store, db_url: str) -> None:
     """
     Background task that periodically marks stale RUNNING jobs as FAILURE.
 
     A job is considered "ghost" if it has been RUNNING for over
-    GHOST_JOB_TIMEOUT_SECONDS with no new events in the job_events table.
-    This catches Dask worker crashes and OOM kills that bypass Python exception handling.
+    GHOST_JOB_TIMEOUT_SECONDS with no new events in the job_events table, OR if
+    it has been RUNNING that long without ever storing an event (measured from
+    job_info.updated_at). This catches Dask worker crashes and OOM kills that
+    bypass Python exception handling, including a crash before the first event
+    is persisted.
     """
     from nat.front_ends.fastapi.async_jobs.job_store import JobStatus
 
@@ -1244,19 +1283,31 @@ async def _reap_ghost_jobs(job_store, db_url: str) -> None:
             stale_job_ids = await loop.run_in_executor(None, _find_stale_jobs, db_url, JobStatus.RUNNING.value)
 
             for stale_job_id in stale_job_ids:
-                logger.warning("Reaping ghost job %s (no events for %ds)", stale_job_id, GHOST_JOB_TIMEOUT_SECONDS)
+                error_msg = "Job timed out (no heartbeat received from worker)"
                 try:
-                    await job_store.update_status(
+                    transitioned = await loop.run_in_executor(
+                        None,
+                        _mark_job_failed_if_running,
+                        db_url,
                         stale_job_id,
-                        JobStatus.FAILURE,
-                        error="Job timed out (no heartbeat received from worker)",
+                        JobStatus.RUNNING.value,
+                        JobStatus.FAILURE.value,
+                        error_msg,
+                    )
+                    if not transitioned:
+                        # The job left RUNNING between detection and reaping
+                        # (e.g. a slow worker finished); leave its status intact.
+                        logger.info("Ghost reap skipped %s: no longer running", stale_job_id)
+                        continue
+                    logger.warning(
+                        "Reaped ghost job %s (no heartbeat for %ds)", stale_job_id, GHOST_JOB_TIMEOUT_SECONDS
                     )
                     event_store = EventStore(db_url, stale_job_id)
                     event_store.store(
                         {
                             "type": "job.error",
                             "data": {
-                                "error": "Job timed out (no heartbeat received from worker)",
+                                "error": error_msg,
                                 "error_type": "GhostJobTimeout",
                             },
                         }