Don't strand assigned scans after a controller restart

api/jobs/executor.py

@@ -73,6 +73,31 @@ async def submit_scan(job: ScanJob) -> None:
         # dispatched by try_dispatch_queued() on the next heartbeat.
 
 
+def rehydrate_assigned_queue(agent_id: str) -> int:
+    """Re-push queued jobs assigned to *agent_id* onto its pending_tasks.
+
+    ``pending_tasks`` is process-local and is not persisted. After a controller
+    restart the durable job still has ``assigned_agent_id`` set, but the agent's
+    in-memory queue is empty. ``try_dispatch_queued`` only considers unassigned
+    jobs, and ``reclaim_stale_jobs`` skips agents that are still online — so
+    the scan would wait forever unless we restore the queue when the agent
+    checks in.
+
+    Returns the number of jobs restored (including any already on the queue).
+    """
+    if not agent_id:
+        return 0
+    job_ids = [
+        job.job_id
+        for job in job_manager.list_dispatched_active()
+        if job.assigned_agent_id == agent_id and job.status == "queued"
+    ]
+    if not job_ids:
+        return 0
+    agent_registry.restore_pending_tasks(agent_id, job_ids)
+    return len(job_ids)
+
+
 def try_dispatch_queued(org_id: str = "") -> int:
     """Flush the queue: assign every waiting job to an idle agent.
 

api/main.py

@@ -227,14 +227,17 @@ async def lifespan(app: FastAPI):
 
     yield
 
-    # Shutdown — mark any still-running jobs as failed so they don't get
-    # stuck in "running" state after a restart.
+    # Shutdown — mark still-*running* jobs as failed so they don't come back
+    # as zombies after a restart (the in-process scan thread dies with us).
+    # Leave queued jobs alone: they have not started, their assignment is
+    # durable, and the assigned agent will pick them up after check-in.
+    # Failing queued scans here dropped every waiting SaaS job on deploy.
     import logging  # noqa: PLC0415
     from api.jobs.manager import job_manager  # noqa: PLC0415
     _log = logging.getLogger("netlogic.api")
     _log.info("NetLogic API shutting down — draining in-flight jobs...")
     for job in list(job_manager._jobs.values()):
-        if job.status in ("running", "queued"):
+        if job.status == "running":
             job.status = "failed"
             job.error = "Scan interrupted by server shutdown."
             import time as _time  # noqa: PLC0415

api/routes/agents.py

@@ -39,7 +39,7 @@
 from api.agents.registry import Agent, agent_registry
 from api.auth.dependencies import require_org
 from api.auth.rate_limit import events_limiter, heartbeat_limiter, register_limiter
-from api.jobs.executor import try_dispatch_queued
+from api.jobs.executor import rehydrate_assigned_queue, try_dispatch_queued
 from api.jobs.manager import job_manager
 from api.middleware.audit import audit_log
 from api.models.agent import AgentRegistration, AgentTaskComplete
@@ -143,6 +143,10 @@ async def agent_heartbeat(
     if not heartbeat_limiter.allow(agent_id):
         raise HTTPException(status_code=429, detail="Heartbeat rate limit exceeded.")
     agent_registry.heartbeat(agent_id)
+    # pending_tasks is not persisted — restore jobs assigned to this agent
+    # before dispatching new work, so a restart cannot strand them forever
+    # and so load accounting still reflects already-assigned scans.
+    rehydrate_assigned_queue(agent_id)
     # Dispatch any queued jobs now that this agent has checked in.
     try_dispatch_queued(org_id=agent.org_id)
     return {
@@ -175,6 +179,7 @@ async def get_pending_tasks(
     """
     # Implicit heartbeat on every poll — no separate call needed.
     agent_registry.heartbeat(agent_id)
+    rehydrate_assigned_queue(agent_id)
 
     job_ids = agent_registry.get_pending_tasks(agent_id)
     tasks = []

test_agent_architecture.py

@@ -394,5 +394,104 @@ def test_reclaim_fails_after_max_attempts(self):
         self.assertIn("No healthy agent", job.error)
 
 
+# ─── Resilience: restore assigned queue after controller restart ───────────────
+
+class TestRehydrateAssignedQueue(unittest.TestCase):
+    """pending_tasks is not persisted. After a controller crash/restart the
+    durable job still has assigned_agent_id, the agent is still online, and
+    both try_dispatch_queued and reclaim_stale_jobs skip it — scan never runs.
+    """
+
+    def _setup(self):
+        from api.agents.registry import AgentRegistry
+        from api.jobs.manager import JobManager, ScanJob
+        reg = AgentRegistry(persist_path=None)
+        aid, _ = reg.register(hostname="live", capabilities=["scan"], version="1.0", tags={})
+        reg.heartbeat(aid)
+        jm = JobManager()
+        jm._jobs.clear()
+        job = ScanJob(job_id="stranded", config=ScanRequest(target="example.com"), org_id="")
+        job.assigned_agent_id = aid
+        job.status = "queued"
+        job.dispatch_attempts = 1
+        # Simulate restart: assignment survived on disk, in-memory queue did not.
+        self.assertEqual(reg.get(aid).pending_tasks, [])
+        jm._jobs["stranded"] = job
+        return reg, jm, job, aid
+
+    def test_dispatch_and_reclaim_leave_assigned_queued_job_stranded(self):
+        """Documents the hole: online assigned+queued jobs are invisible to
+        both the dispatcher (looks for unassigned) and the reclaimer (skips
+        live agents)."""
+        from api.jobs import executor
+        reg, jm, job, aid = self._setup()
+        with patch.object(executor, "agent_registry", reg), \
+             patch.object(executor, "job_manager", jm):
+            self.assertEqual(executor.try_dispatch_queued(), 0)
+            self.assertEqual(executor.reclaim_stale_jobs(), 0)
+        self.assertEqual(job.status, "queued")
+        self.assertEqual(job.assigned_agent_id, aid)
+        self.assertEqual(reg.get_pending_tasks(aid), [])
+
+    def test_rehydrate_restores_queued_job_to_pending_queue(self):
+        from api.jobs import executor
+        reg, jm, job, aid = self._setup()
+        with patch.object(executor, "agent_registry", reg), \
+             patch.object(executor, "job_manager", jm):
+            n = executor.rehydrate_assigned_queue(aid)
+        self.assertEqual(n, 1)
+        self.assertEqual(reg.get(aid).pending_tasks, ["stranded"])
+        self.assertEqual(job.status, "queued")
+
+    def test_rehydrate_is_idempotent(self):
+        from api.jobs import executor
+        reg, jm, job, aid = self._setup()
+        with patch.object(executor, "agent_registry", reg), \
+             patch.object(executor, "job_manager", jm):
+            executor.rehydrate_assigned_queue(aid)
+            executor.rehydrate_assigned_queue(aid)
+        self.assertEqual(reg.get(aid).pending_tasks, ["stranded"])
+
+    def test_rehydrate_skips_running_jobs(self):
+        from api.jobs import executor
+        reg, jm, job, aid = self._setup()
+        job.status = "running"
+        with patch.object(executor, "agent_registry", reg), \
+             patch.object(executor, "job_manager", jm):
+            n = executor.rehydrate_assigned_queue(aid)
+        self.assertEqual(n, 0)
+        self.assertEqual(reg.get(aid).pending_tasks, [])
+
+    def test_rehydrate_ignores_other_agents_jobs(self):
+        from api.jobs import executor
+        reg, jm, job, aid = self._setup()
+        other, _ = reg.register(hostname="other", capabilities=["scan"], version="1.0", tags={})
+        reg.heartbeat(other)
+        with patch.object(executor, "agent_registry", reg), \
+             patch.object(executor, "job_manager", jm):
+            n = executor.rehydrate_assigned_queue(other)
+        self.assertEqual(n, 0)
+        self.assertEqual(reg.get(other).pending_tasks, [])
+        self.assertEqual(job.assigned_agent_id, aid)
+
+    def test_rehydrate_before_dispatch_keeps_capacity_for_assigned_work(self):
+        """If we dispatched new work before restoring the assigned queue, a
+        concurrency=1 agent would accept a second job and the original scan
+        would still be stranded."""
+        from api.jobs import executor
+        from api.jobs.manager import ScanJob
+        reg, jm, job, aid = self._setup()
+        extra = ScanJob(job_id="extra", config=ScanRequest(target="example.org"), org_id="")
+        extra.status = "queued"
+        jm._jobs["extra"] = extra
+        with patch.object(executor, "agent_registry", reg), \
+             patch.object(executor, "job_manager", jm):
+            executor.rehydrate_assigned_queue(aid)
+            dispatched = executor.try_dispatch_queued()
+        self.assertEqual(dispatched, 0)
+        self.assertIsNone(extra.assigned_agent_id)
+        self.assertEqual(reg.get(aid).pending_tasks, ["stranded"])
+
+
 if __name__ == "__main__":
     unittest.main()

test_production_readiness.py

@@ -257,6 +257,33 @@ def test_full_task_lifecycle(self):
         job_detail = client.get(f"/v1/jobs/{job_id}", headers=self.headers)
         self.assertEqual(job_detail.json()["status"], "completed")
 
+    def test_task_poll_recovers_assignment_after_pending_queue_loss(self):
+        """Controller restart wipes pending_tasks. Agent poll must still
+        receive the job that was already assigned to it."""
+        from api.agents.registry import agent_registry
+
+        agent_id, token = self._register_agent()
+        agent_headers = {"Authorization": f"Bearer {token}"}
+        client.post(f"/v1/agents/{agent_id}/heartbeat", headers=agent_headers)
+
+        job_resp = client.post(
+            "/v1/jobs",
+            json={"target": "127.0.0.1", "agent_id": agent_id},
+            headers=self.headers,
+        )
+        self.assertEqual(job_resp.status_code, 202, job_resp.text)
+        job_id = job_resp.json()["job_id"]
+
+        agent = agent_registry.get(agent_id)
+        self.assertIn(job_id, agent.pending_tasks)
+        agent.pending_tasks.clear()  # simulate process restart
+
+        tasks_resp = client.get(f"/v1/agents/{agent_id}/tasks", headers=agent_headers)
+        self.assertEqual(tasks_resp.status_code, 200)
+        tasks = tasks_resp.json()
+        self.assertEqual(len(tasks), 1, tasks)
+        self.assertEqual(tasks[0]["job_id"], job_id)
+
 
 # ── Cross-org isolation ───────────────────────────────────────────────────────