diff --git a/api/jobs/executor.py b/api/jobs/executor.py index 0c11990..583b6df 100644 --- a/api/jobs/executor.py +++ b/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. diff --git a/api/main.py b/api/main.py index 60fa2a0..213b44e 100644 --- a/api/main.py +++ b/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 diff --git a/api/routes/agents.py b/api/routes/agents.py index b0548b0..42be07e 100644 --- a/api/routes/agents.py +++ b/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 = [] diff --git a/test_agent_architecture.py b/test_agent_architecture.py index 6f4b798..896f490 100644 --- a/test_agent_architecture.py +++ b/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() diff --git a/test_production_readiness.py b/test_production_readiness.py index 216505a..d93166c 100644 --- a/test_production_readiness.py +++ b/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 ───────────────────────────────────────────────────────