diff --git a/frontends/aiq_api/src/aiq_api/jobs/crypto.py b/frontends/aiq_api/src/aiq_api/jobs/crypto.py index b6ca02679..64fb6c94c 100644 --- a/frontends/aiq_api/src/aiq_api/jobs/crypto.py +++ b/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.""" diff --git a/frontends/aiq_api/src/aiq_api/jobs/runner.py b/frontends/aiq_api/src/aiq_api/jobs/runner.py index d8c8da383..183d3410d 100644 --- a/frontends/aiq_api/src/aiq_api/jobs/runner.py +++ b/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. diff --git a/frontends/aiq_api/src/aiq_api/routes/jobs.py b/frontends/aiq_api/src/aiq_api/routes/jobs.py index 4263d20f6..bde0c267a 100644 --- a/frontends/aiq_api/src/aiq_api/routes/jobs.py +++ b/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", }, } diff --git a/frontends/aiq_api/tests/test_ghost_reaper.py b/frontends/aiq_api/tests/test_ghost_reaper.py new file mode 100644 index 000000000..d83026888 --- /dev/null +++ b/frontends/aiq_api/tests/test_ghost_reaper.py @@ -0,0 +1,298 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the ghost-job reaper's stale-job detection (_find_stale_jobs). + +These use a real SQLite database so the SQL runs exactly as it would in +production. The key regression: a job that entered RUNNING but never stored an +event (worker crash/OOM before the first event flush) must still be reaped; +before the fix, the INNER JOIN on job_events made such jobs invisible. +""" + +from __future__ import annotations + +import sys +import tempfile +from datetime import UTC +from datetime import datetime +from datetime import timedelta +from pathlib import Path + +from sqlalchemy import create_engine +from sqlalchemy import text +from sqlalchemy.orm import Session + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src")) + +from aiq_api.routes.jobs import GHOST_JOB_TIMEOUT_SECONDS # noqa: E402 +from aiq_api.routes.jobs import _find_stale_jobs # noqa: E402 +from aiq_api.routes.jobs import _mark_job_failed_if_running # noqa: E402 + +RUNNING = "running" +FAILURE = "failure" +SUCCESS = "success" + + +def _make_db() -> str: + """Create a temp SQLite DB with the job_info and job_events tables.""" + db_path = tempfile.mktemp(suffix=".db") + db_url = f"sqlite:///{db_path}" + + # job_info comes from NAT's ORM model. + from nat.front_ends.fastapi.async_jobs.job_store import JobInfo + + engine = create_engine(db_url) + JobInfo.__table__.metadata.create_all(engine) + + # job_events is created by aiq's EventStore. + from aiq_api.jobs.event_store import EventStore + + EventStore._ensure_table_exists(db_url) + return db_url + + +def _insert_job( + db_url: str, + job_id: str, + *, + status: str, + updated_ago_seconds: float, + created_ago_seconds: float | None = None, +) -> None: + """Insert a job_info row. updated_at (the lease) and created_at ages differ + when created_ago_seconds is given, to model a long-running-but-live job.""" + from nat.front_ends.fastapi.async_jobs.job_store import JobInfo + + now = datetime.now(UTC) + updated = now - timedelta(seconds=updated_ago_seconds) + created = now - timedelta(seconds=created_ago_seconds if created_ago_seconds is not None else updated_ago_seconds) + engine = create_engine(db_url) + with Session(engine) as s: + s.add(JobInfo(job_id=job_id, status=status, expiry_seconds=3600, created_at=created, updated_at=updated)) + s.commit() + + +def _get_status(db_url: str, job_id: str) -> str | None: + """Return the stored status for a job, or None if absent.""" + engine = create_engine(db_url) + with engine.connect() as conn: + row = conn.execute(text("SELECT status FROM job_info WHERE job_id = :j"), {"j": job_id}).first() + return row[0] if row else None + + +def _insert_event(db_url: str, job_id: str, *, created_ago_seconds: float) -> None: + """Insert a job_events row with created_at set to now minus the given age.""" + ts = (datetime.now(UTC) - timedelta(seconds=created_ago_seconds)).strftime("%Y-%m-%d %H:%M:%S") + engine = create_engine(db_url) + with engine.begin() as conn: + conn.execute( + text("INSERT INTO job_events (job_id, event_type, event_data, created_at) VALUES (:j, :t, :d, :c)"), + {"j": job_id, "t": "test.event", "d": "{}", "c": ts}, + ) + + +OLD = GHOST_JOB_TIMEOUT_SECONDS + 60 +RECENT = 5 + + +def test_zero_event_running_job_past_timeout_is_reaped(): + """A RUNNING job with no events, started long ago, is a ghost and reaped. + + This is the regression: the old INNER JOIN made zero-event jobs invisible. + """ + db = _make_db() + _insert_job(db, "ghost", status=RUNNING, updated_ago_seconds=OLD) + assert _find_stale_jobs(db, RUNNING) == ["ghost"] + + +def test_zero_event_running_job_within_timeout_is_not_reaped(): + """A freshly-started RUNNING job with no events yet must not be reaped.""" + db = _make_db() + _insert_job(db, "fresh", status=RUNNING, updated_ago_seconds=RECENT) + assert _find_stale_jobs(db, RUNNING) == [] + + +def test_running_job_with_stale_last_event_is_reaped(): + """Existing behavior preserved: events present but last one is old.""" + db = _make_db() + _insert_job(db, "stalled", status=RUNNING, updated_ago_seconds=OLD) + _insert_event(db, "stalled", created_ago_seconds=OLD) + assert _find_stale_jobs(db, RUNNING) == ["stalled"] + + +def test_running_job_with_recent_event_is_not_reaped(): + """A job actively emitting events is healthy, even if it started long ago.""" + db = _make_db() + _insert_job(db, "active", status=RUNNING, updated_ago_seconds=OLD) + _insert_event(db, "active", created_ago_seconds=RECENT) + assert _find_stale_jobs(db, RUNNING) == [] + + +def test_non_running_job_is_never_reaped(): + """Only RUNNING jobs are candidates; a completed job is left alone.""" + db = _make_db() + _insert_job(db, "done", status="success", updated_ago_seconds=OLD) + assert _find_stale_jobs(db, RUNNING) == [] + + +def test_missing_tables_returns_empty(): + """No job_info/job_events tables (fresh deployment) -> nothing to reap.""" + db_path = tempfile.mktemp(suffix=".db") + assert _find_stale_jobs(f"sqlite:///{db_path}", RUNNING) == [] + + +def test_mixed_fleet_reaps_only_ghosts(): + """A realistic mix: only the two ghosts (old zero-event + stalled) return.""" + db = _make_db() + _insert_job(db, "ghost-zero", status=RUNNING, updated_ago_seconds=OLD) + _insert_job(db, "fresh-zero", status=RUNNING, updated_ago_seconds=RECENT) + _insert_job(db, "stalled", status=RUNNING, updated_ago_seconds=OLD) + _insert_event(db, "stalled", created_ago_seconds=OLD) + _insert_job(db, "healthy", status=RUNNING, updated_ago_seconds=OLD) + _insert_event(db, "healthy", created_ago_seconds=RECENT) + _insert_job(db, "done", status="success", updated_ago_seconds=OLD) + + assert sorted(_find_stale_jobs(db, RUNNING)) == ["ghost-zero", "stalled"] + + +# --------------------------------------------------------------------------- +# Cold-start lease + atomic conditional transition (issue #318 review follow-up) +# --------------------------------------------------------------------------- + + +def test_slow_init_job_with_fresh_lease_is_not_reaped(): + """A live worker in a long cold start refreshes its lease (updated_at), so a + zero-event RUNNING job that STARTED long ago but was touched recently is not + reaped — the regression AjayThorve flagged.""" + db = _make_db() + # created long ago (slow init still running), but lease refreshed just now. + _insert_job(db, "slow-init", status=RUNNING, updated_ago_seconds=RECENT, created_ago_seconds=OLD) + assert _find_stale_jobs(db, RUNNING) == [] + + +def test_stale_lease_zero_event_job_is_reaped(): + """Once the lease goes stale (worker dead), the zero-event ghost is reaped.""" + db = _make_db() + _insert_job(db, "dead", status=RUNNING, updated_ago_seconds=OLD, created_ago_seconds=OLD) + assert _find_stale_jobs(db, RUNNING) == ["dead"] + + +def test_conditional_transition_marks_running_job_failed(): + """The atomic transition flips a still-RUNNING job to FAILURE and reports it.""" + db = _make_db() + _insert_job(db, "ghost", status=RUNNING, updated_ago_seconds=OLD) + did = _mark_job_failed_if_running(db, "ghost", RUNNING, FAILURE, "timed out") + assert did is True + assert _get_status(db, "ghost") == FAILURE + + +def test_conditional_transition_does_not_clobber_terminal_job(): + """A job that reached SUCCESS between detection and reaping is left intact.""" + db = _make_db() + _insert_job(db, "finished", status=SUCCESS, updated_ago_seconds=OLD) + did = _mark_job_failed_if_running(db, "finished", RUNNING, FAILURE, "timed out") + assert did is False + assert _get_status(db, "finished") == SUCCESS # unchanged + + +def test_conditional_transition_missing_job_is_noop(): + """Transitioning a non-existent job returns False without error.""" + db = _make_db() + assert _mark_job_failed_if_running(db, "nope", RUNNING, FAILURE, "x") is False + + +def test_runner_lease_touch_refreshes_only_running_jobs(): + """The runner's lease bumps updated_at for a RUNNING job but not a terminal one.""" + from aiq_api.jobs.runner import _touch_job_lease_sync + + db = _make_db() + _insert_job(db, "run", status=RUNNING, updated_ago_seconds=OLD) + _insert_job(db, "done", status=SUCCESS, updated_ago_seconds=OLD) + + # Before: the running job is stale and would be reaped. + assert _find_stale_jobs(db, RUNNING) == ["run"] + + _touch_job_lease_sync(db, "run") # live worker refreshes its lease + _touch_job_lease_sync(db, "done") # must be a no-op for a terminal job + + # After: the running job's lease is fresh, so it is no longer reapable; + # the terminal job's timestamp was not resurrected. + assert _find_stale_jobs(db, RUNNING) == [] + assert _get_status(db, "done") == SUCCESS + + +def test_lease_refresher_thread_bumps_updated_at(): + """The dedicated lease thread actually refreshes a running job's lease. + + Uses a 0s interval so the thread refreshes immediately, then stops it. + """ + import threading + import time + + from aiq_api.jobs import runner + + db = _make_db() + _insert_job(db, "run", status=RUNNING, updated_ago_seconds=OLD) + assert _find_stale_jobs(db, RUNNING) == ["run"] # stale before any refresh + + stop = threading.Event() + # Small positive interval so the refresher touches the lease promptly + # without a zero-interval busy loop hammering the DB until teardown. + original = runner.LEASE_REFRESH_INTERVAL_SECONDS + runner.LEASE_REFRESH_INTERVAL_SECONDS = 0.01 + t = threading.Thread(target=runner._run_lease_refresher, args=(db, "run", stop), daemon=True) + try: + t.start() + # Poll until the lease is refreshed (thread is on a 0s loop). + for _ in range(50): + if _find_stale_jobs(db, RUNNING) == []: + break + time.sleep(0.02) + assert _find_stale_jobs(db, RUNNING) == [] # lease is now fresh + finally: + stop.set() + t.join(timeout=5) + runner.LEASE_REFRESH_INTERVAL_SECONDS = original + assert not t.is_alive() # thread exits promptly on stop + + +def test_db_now_expr_handles_both_postgres_schemes(): + """Both postgresql:// and legacy postgres:// map to NOW(); sqlite to CURRENT_TIMESTAMP.""" + from aiq_api.jobs.runner import _db_now_expr + + assert _db_now_expr("postgresql://u@h/db") == "NOW()" + assert _db_now_expr("postgres://u@h/db") == "NOW()" + assert _db_now_expr("sqlite:///x.db") == "CURRENT_TIMESTAMP" + + +def test_success_cas_writes_only_when_running(): + """The worker's success write is a compare-and-set: it writes for a running + job but not for one the reaper already moved to a terminal state.""" + from aiq_api.jobs.runner import _write_job_success_if_running_sync + + db = _make_db() + _insert_job(db, "run", status=RUNNING, updated_ago_seconds=RECENT) + assert _write_job_success_if_running_sync(db, "run", '{"report": "ok"}') is True + assert _get_status(db, "run") == SUCCESS + + +def test_success_cas_does_not_resurrect_reaped_job(): + """A job already reaped to FAILURE is not resurrected by a late success write.""" + from aiq_api.jobs.runner import _write_job_success_if_running_sync + + db = _make_db() + _insert_job(db, "reaped", status=FAILURE, updated_ago_seconds=RECENT) + assert _write_job_success_if_running_sync(db, "reaped", '{"report": "late"}') is False + assert _get_status(db, "reaped") == FAILURE # stays failed diff --git a/tests/aiq_agent/jobs/test_runner.py b/tests/aiq_agent/jobs/test_runner.py index e26a3ee43..c44a02ad0 100644 --- a/tests/aiq_agent/jobs/test_runner.py +++ b/tests/aiq_agent/jobs/test_runner.py @@ -725,7 +725,10 @@ def start(self, *, context_state): mock_job_store = MagicMock() mock_job_store.update_status = AsyncMock() - update_job_output = AsyncMock(side_effect=ContentEncryptionUnavailable("encrypt failed")) + # The success output is serialized/encrypted via serialize_job_output_for_storage + # before the conditional write; simulate encryption failing there. + serialize_output = MagicMock(side_effect=ContentEncryptionUnavailable("encrypt failed")) + write_success = MagicMock() # the raw-SQL success writer; must never run on failure db_url = f"sqlite:///{tmp_path / 'test.db'}" config = SimpleNamespace(workflow=None, functions={}, middleware={}) @@ -749,7 +752,16 @@ def start(self, *, context_state): "aiq_api.jobs.runner._run_agent", AsyncMock(return_value="secret report"), ): - with patch("aiq_api.jobs.crypto.update_job_output", update_job_output): + with ( + patch( + "aiq_api.jobs.crypto.serialize_job_output_for_storage", + serialize_output, + ), + patch( + "aiq_api.jobs.runner._write_job_success_if_running_sync", + write_success, + ), + ): await run_agent_job( False, 20, @@ -773,8 +785,12 @@ def start(self, *, context_state): statuses = [call.args[1] for call in mock_job_store.update_status.await_args_list] assert statuses == [JobStatus.RUNNING, JobStatus.FAILURE] assert all("output" not in call.kwargs for call in mock_job_store.update_status.await_args_list) - update_job_output.assert_awaited_once() - assert update_job_output.await_args.kwargs["output"] == { + # The output was assembled with the real report (never the output_metadata + # "report" decoy) and handed to serialization; when that failed the job was + # marked FAILURE and the success writer never ran, so nothing was persisted. + write_success.assert_not_called() + serialize_output.assert_called_once() + assert serialize_output.call_args.args[0] == { "parent_job_id": "parent-job", "interaction_action": "edit", "report": "secret report",