refactor: extract check_stale_jobs() for reuse by periodic task

Move core stale-job logic from management command into
check_stale_jobs() in tasks.py. The management command is now a
thin wrapper. Add tests for the extracted function.

This prepares for #1025 which will call check_stale_jobs() from a
Celery Beat periodic task.

Co-Authored-By: Claude <noreply@anthropic.com>

Author: mihow

ami/jobs/management/commands/update_stale_jobs.py

@@ -1,14 +1,7 @@
-from celery import states
-from celery.result import AsyncResult
 from django.core.management.base import BaseCommand
-from django.utils import timezone
 
-from ami.jobs.models import Job, JobState
-from ami.jobs.tasks import cleanup_async_job_if_needed
-
-# Celery returns PENDING for tasks it has no record of.
-# These are the states that indicate a real, known task status.
-KNOWN_CELERY_STATES = frozenset(states.ALL_STATES) - {states.PENDING}
+from ami.jobs.models import Job
+from ami.jobs.tasks import check_stale_jobs
 
 
 class Command(BaseCommand):
@@ -28,36 +21,17 @@ def add_arguments(self, parser):
         )
 
     def handle(self, *args, **options):
-        cutoff = timezone.now() - timezone.timedelta(hours=options["hours"])
-        stale_jobs = Job.objects.filter(
-            status__in=JobState.running_states(),
-            updated_at__lt=cutoff,
-        )
+        results = check_stale_jobs(hours=options["hours"], dry_run=options["dry_run"])
 
-        if not stale_jobs.exists():
+        if not results:
             self.stdout.write("No stale jobs found.")
             return
 
-        for job in stale_jobs:
-            celery_state = None
-            if job.task_id:
-                celery_state = AsyncResult(job.task_id).state
-
-            if celery_state in KNOWN_CELERY_STATES:
-                # Celery has a real status for this task — use it
-                if options["dry_run"]:
-                    self.stdout.write(f"  [dry-run] Job {job.pk}: would update to {celery_state} (from Celery)")
-                    continue
-                job.update_status(celery_state, save=False)
-                job.save()
-                self.stdout.write(self.style.SUCCESS(f"Job {job.pk}: updated to {celery_state} (from Celery)"))
+        prefix = "[dry-run] " if options["dry_run"] else ""
+        for r in results:
+            if r["action"] == "updated":
+                self.stdout.write(
+                    self.style.SUCCESS(f"{prefix}Job {r['job_id']}: updated to {r['state']} (from Celery)")
+                )
             else:
-                # No task_id, or Celery has no record (returns PENDING) — revoke
-                if options["dry_run"]:
-                    self.stdout.write(f"  [dry-run] Job {job.pk} ({job.status}): would revoke and clean up")
-                    continue
-                job.update_status(JobState.REVOKED, save=False)
-                job.finished_at = timezone.now()
-                job.save()
-                cleanup_async_job_if_needed(job)
-                self.stdout.write(self.style.WARNING(f"Job {job.pk}: revoked (no known Celery state)"))
+                self.stdout.write(self.style.WARNING(f"{prefix}Job {r['job_id']}: revoked (no known Celery state)"))

ami/jobs/tasks.py

@@ -316,6 +316,56 @@ def _update_job_progress(
         cleanup_async_job_if_needed(job)
 
 
+def check_stale_jobs(hours: int | None = None, dry_run: bool = False) -> list[dict]:
+    """
+    Find jobs stuck in a running state past the cutoff and revoke them.
+
+    For each stale job, checks Celery for a real task status. If Celery has one
+    (e.g. SUCCESS, FAILURE), uses that. Otherwise revokes the job and cleans up
+    any async resources (NATS/Redis).
+
+    Returns a list of dicts describing what was done to each job.
+    """
+    import datetime
+
+    from celery import states
+    from celery.result import AsyncResult
+
+    from ami.jobs.models import Job, JobState
+
+    if hours is None:
+        hours = Job.FAILED_CUTOFF_HOURS
+
+    known_celery_states = frozenset(states.ALL_STATES) - {states.PENDING}
+
+    cutoff = datetime.datetime.now() - datetime.timedelta(hours=hours)
+    stale_jobs = Job.objects.filter(
+        status__in=JobState.running_states(),
+        updated_at__lt=cutoff,
+    )
+
+    results = []
+    for job in stale_jobs:
+        celery_state = None
+        if job.task_id:
+            celery_state = AsyncResult(job.task_id).state
+
+        if celery_state in known_celery_states:
+            if not dry_run:
+                job.update_status(celery_state, save=False)
+                job.save()
+            results.append({"job_id": job.pk, "action": "updated", "state": celery_state})
+        else:
+            if not dry_run:
+                job.update_status(JobState.REVOKED, save=False)
+                job.finished_at = datetime.datetime.now()
+                job.save()
+                cleanup_async_job_if_needed(job)
+            results.append({"job_id": job.pk, "action": "revoked", "previous_status": job.status})
+
+    return results
+
+
 def cleanup_async_job_if_needed(job) -> None:
     """
     Clean up async resources (NATS/Redis) if this job uses them.

ami/jobs/tests_update_stale_jobs.py

@@ -0,0 +1,66 @@
+from datetime import timedelta
+from unittest.mock import patch
+
+from django.test import TestCase
+from django.utils import timezone
+
+from ami.jobs.models import Job, JobState
+from ami.jobs.tasks import check_stale_jobs
+from ami.main.models import Project
+
+
+class CheckStaleJobsTest(TestCase):
+    def setUp(self):
+        self.project = Project.objects.create(name="Stale jobs test project")
+
+    def _create_job(self, status=JobState.STARTED, hours_ago=100, task_id=None):
+        job = Job.objects.create(
+            project=self.project,
+            name=f"Test job {status}",
+            status=status,
+        )
+        Job.objects.filter(pk=job.pk).update(
+            updated_at=timezone.now() - timedelta(hours=hours_ago),
+        )
+        if task_id is not None:
+            Job.objects.filter(pk=job.pk).update(task_id=task_id)
+        job.refresh_from_db()
+        return job
+
+    @patch("ami.jobs.tasks.cleanup_async_job_if_needed")
+    def test_dry_run(self, mock_cleanup):
+        """Dry run returns results without modifying jobs."""
+        job = self._create_job(status=JobState.STARTED)
+
+        results = check_stale_jobs(dry_run=True)
+
+        self.assertEqual(len(results), 1)
+        self.assertEqual(results[0]["action"], "revoked")
+        job.refresh_from_db()
+        self.assertEqual(job.status, JobState.STARTED.value)
+        mock_cleanup.assert_not_called()
+
+    @patch("ami.jobs.tasks.cleanup_async_job_if_needed")
+    def test_revokes_stale_job(self, mock_cleanup):
+        """Stale job without a known Celery state is revoked and cleaned up."""
+        job = self._create_job(status=JobState.STARTED)
+
+        results = check_stale_jobs()
+
+        self.assertEqual(len(results), 1)
+        self.assertEqual(results[0]["action"], "revoked")
+        job.refresh_from_db()
+        self.assertEqual(job.status, JobState.REVOKED.value)
+        self.assertIsNotNone(job.finished_at)
+        mock_cleanup.assert_called_once_with(job)
+
+    @patch("ami.jobs.tasks.cleanup_async_job_if_needed")
+    def test_skips_recent_and_final_state_jobs(self, mock_cleanup):
+        """Recent jobs and jobs in final states are not touched."""
+        self._create_job(status=JobState.STARTED, hours_ago=1)  # recent
+        self._create_job(status=JobState.SUCCESS, hours_ago=200)  # final state
+
+        results = check_stale_jobs()
+
+        self.assertEqual(results, [])
+        mock_cleanup.assert_not_called()