backend/app/tasks/maintenance.py

@@ -776,89 +776,62 @@ def recover_stalled_library_audit_runs() -> int:
         return recovered
 
 
+def _recover_stalled_runs(model, *, stall_minutes: int, keep_runs: int, label: str) -> int:
+    """Shared recovery + retention sweep for the head run-tracking tables
+    (HeadTrainingRun / HeadAutoApplyRun, which share the
+    status/last_progress_at/started_at/finished_at/error/id columns): flip 'running'
+    rows with no progress past `stall_minutes` to 'error', then prune to the last
+    `keep_runs` (rule 89). Returns the number recovered. NOTE the two other recover
+    tasks are deliberately NOT folded in — library-audit has no prune tail and
+    backup uses a single started_at cutoff."""
+    SessionLocal = _sync_session_factory()
+    now = datetime.now(UTC)
+    cutoff = now - timedelta(minutes=stall_minutes)
+    with SessionLocal() as session:
+        result = session.execute(
+            update(model)
+            .where(model.status == "running")
+            .where(func.coalesce(model.last_progress_at, model.started_at) < cutoff)
+            .values(
+                status="error", finished_at=now,
+                error=f"stranded by recovery sweep (no progress for {stall_minutes} min)",
+            )
+        )
+        keep = session.execute(
+            select(model.id).order_by(model.id.desc()).limit(keep_runs)
+        ).scalars().all()
+        if keep:
+            session.execute(delete(model).where(model.id.not_in(keep)))
+        session.commit()
+        recovered = result.rowcount or 0
+        if recovered:
+            log.info("%s: recovered %d rows", label, recovered)
+        return recovered
+
+
 @celery.task(name="backend.app.tasks.maintenance.recover_stalled_head_training_runs")
 def recover_stalled_head_training_runs() -> int:
     """Flip HeadTrainingRun rows stuck in 'running' past the stall threshold to
     'error', and prune old runs to the last HEAD_TRAINING_KEEP_RUNS (retention,
     rule 89). Runs every 5 min on the maintenance lane; no-op when idle."""
-    SessionLocal = _sync_session_factory()
+    return _recover_stalled_runs(
-    now = datetime.now(UTC)
+        HeadTrainingRun,
-    cutoff = now - timedelta(minutes=HEAD_TRAINING_STALL_THRESHOLD_MINUTES)
+        stall_minutes=HEAD_TRAINING_STALL_THRESHOLD_MINUTES,
-    with SessionLocal() as session:
+        keep_runs=HEAD_TRAINING_KEEP_RUNS,
-        result = session.execute(
+        label="recover_stalled_head_training_runs",
-            update(HeadTrainingRun)
-            .where(HeadTrainingRun.status == "running")
-            .where(
-                func.coalesce(
-                    HeadTrainingRun.last_progress_at, HeadTrainingRun.started_at
     )
-                < cutoff
-            )
-            .values(
-                status="error", finished_at=now,
-                error=(
-                    f"stranded by recovery sweep (no progress for "
-                    f"{HEAD_TRAINING_STALL_THRESHOLD_MINUTES} min)"
-                ),
-            )
-        )
-        keep = session.execute(
-            select(HeadTrainingRun.id).order_by(HeadTrainingRun.id.desc())
-            .limit(HEAD_TRAINING_KEEP_RUNS)
-        ).scalars().all()
-        if keep:
-            session.execute(
-                delete(HeadTrainingRun).where(HeadTrainingRun.id.not_in(keep))
-            )
-        session.commit()
-        recovered = result.rowcount or 0
-        if recovered:
-            log.info(
-                "recover_stalled_head_training_runs: recovered %d rows", recovered
-            )
-        return recovered
 
 
 @celery.task(name="backend.app.tasks.maintenance.recover_stalled_head_auto_apply_runs")
 def recover_stalled_head_auto_apply_runs() -> int:
     """Flip stalled HeadAutoApplyRun 'running' rows to 'error' + prune to the
     last HEAD_AUTO_APPLY_KEEP_RUNS (retention, rule 89). 5-min maintenance lane."""
-    SessionLocal = _sync_session_factory()
+    return _recover_stalled_runs(
-    now = datetime.now(UTC)
+        HeadAutoApplyRun,
-    cutoff = now - timedelta(minutes=HEAD_AUTO_APPLY_STALL_THRESHOLD_MINUTES)
+        stall_minutes=HEAD_AUTO_APPLY_STALL_THRESHOLD_MINUTES,
-    with SessionLocal() as session:
+        keep_runs=HEAD_AUTO_APPLY_KEEP_RUNS,
-        result = session.execute(
+        label="recover_stalled_head_auto_apply_runs",
-            update(HeadAutoApplyRun)
-            .where(HeadAutoApplyRun.status == "running")
-            .where(
-                func.coalesce(
-                    HeadAutoApplyRun.last_progress_at, HeadAutoApplyRun.started_at
     )
-                < cutoff
-            )
-            .values(
-                status="error", finished_at=now,
-                error=(
-                    f"stranded by recovery sweep (no progress for "
-                    f"{HEAD_AUTO_APPLY_STALL_THRESHOLD_MINUTES} min)"
-                ),
-            )
-        )
-        keep = session.execute(
-            select(HeadAutoApplyRun.id).order_by(HeadAutoApplyRun.id.desc())
-            .limit(HEAD_AUTO_APPLY_KEEP_RUNS)
-        ).scalars().all()
-        if keep:
-            session.execute(
-                delete(HeadAutoApplyRun).where(HeadAutoApplyRun.id.not_in(keep))
-            )
-        session.commit()
-        recovered = result.rowcount or 0
-        if recovered:
-            log.info(
-                "recover_stalled_head_auto_apply_runs: recovered %d rows", recovered
-            )
-        return recovered
 
 
 # Keep ~6 months of daily head-metric snapshots (enough to see tuning trends).

tests/test_recover_stalled_head_runs.py

@@ -0,0 +1,63 @@
+"""recover_stalled_head_training_runs + recover_stalled_head_auto_apply_runs share
+one helper (_recover_stalled_runs, DRY pass #161). These pin BOTH wrappers so the
+shared source stays correct: a 'running' row with no progress past the stall
+threshold flips to 'error'; a fresh 'running' row is left alone."""
+from datetime import UTC, datetime, timedelta
+
+import pytest
+from sqlalchemy import select
+
+from backend.app.models import HeadAutoApplyRun, HeadTrainingRun
+
+pytestmark = pytest.mark.integration
+
+
+def test_recover_stalled_head_training_runs_flips_stalled_keeps_fresh(db_sync):
+    from backend.app.tasks.maintenance import recover_stalled_head_training_runs
+
+    stale = HeadTrainingRun(
+        params={}, status="running",
+        last_progress_at=datetime.now(UTC) - timedelta(days=1),
+    )
+    fresh = HeadTrainingRun(
+        params={}, status="running", last_progress_at=datetime.now(UTC),
+    )
+    db_sync.add_all([stale, fresh])
+    db_sync.commit()
+    stale_id, fresh_id = stale.id, fresh.id
+
+    assert recover_stalled_head_training_runs.apply().get() == 1
+
+    db_sync.expire_all()
+    assert db_sync.execute(
+        select(HeadTrainingRun.status).where(HeadTrainingRun.id == stale_id)
+    ).scalar_one() == "error"
+    assert db_sync.execute(
+        select(HeadTrainingRun.status).where(HeadTrainingRun.id == fresh_id)
+    ).scalar_one() == "running"
+
+
+def test_recover_stalled_head_auto_apply_runs_flips_stalled_keeps_fresh(db_sync):
+    from backend.app.tasks.maintenance import recover_stalled_head_auto_apply_runs
+
+    stale = HeadAutoApplyRun(
+        dry_run=False, params={}, status="running",
+        last_progress_at=datetime.now(UTC) - timedelta(days=1),
+    )
+    fresh = HeadAutoApplyRun(
+        dry_run=False, params={}, status="running",
+        last_progress_at=datetime.now(UTC),
+    )
+    db_sync.add_all([stale, fresh])
+    db_sync.commit()
+    stale_id, fresh_id = stale.id, fresh.id
+
+    assert recover_stalled_head_auto_apply_runs.apply().get() == 1
+
+    db_sync.expire_all()
+    assert db_sync.execute(
+        select(HeadAutoApplyRun.status).where(HeadAutoApplyRun.id == stale_id)
+    ).scalar_one() == "error"
+    assert db_sync.execute(
+        select(HeadAutoApplyRun.status).where(HeadAutoApplyRun.id == fresh_id)
+    ).scalar_one() == "running"