From 91ba6a8c89970366ca91212582bd147e69047970 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Mon, 31 Aug 2026 11:18:11 +0000 Subject: [PATCH] Don't mark a cancelled scan as completed. The local agent assigned status=completed after the engine returned, even when cancel_job had already won try_set_terminal. Combined with JobCancelled being swallowed by post-scan except Exception: pass, a user cancel near the end of a scan was persisted as a finished result. Co-authored-by: dmitryflynn --- api/agents/local_agent.py | 25 ++++++--- src/engine.py | 30 ++++++----- test_local_agent_cancel.py | 108 +++++++++++++++++++++++++++++++++++++ 3 files changed, 143 insertions(+), 20 deletions(-) create mode 100644 test_local_agent_cancel.py diff --git a/api/agents/local_agent.py b/api/agents/local_agent.py index 1f2c212..57effad 100644 --- a/api/agents/local_agent.py +++ b/api/agents/local_agent.py @@ -211,18 +211,27 @@ def emit(event_type: str, data=None, message: str = "") -> None: org_id = getattr(job, "org_id", "") or "", emit_callback = emit, ) - job.status = "completed" + # Atomic vs cancel/delete: never assign status="completed" directly. + # A user cancel that won try_set_terminal first must stay cancelled even + # if the engine returned normally (JobCancelled swallowed, or the scan + # finished in the race window after cancel). + job, changed = job_manager.try_set_terminal( + job_id, "completed", org_id=org_id, + ) + if job is None: + if agent: + agent_registry.mark_done(agent_id, job_id) + return True + if not changed: + _finish(job, agent) + return True job.progress = 100.0 try: job.push_event({"type": "done", "data": {"message": "Scan complete."}}) except JobCancelled: - # Cancel raced the final "done" emit — treat as cancelled, not complete. - _log.info("Scan cancelled for job %s (during finalisation)", job_id) - if job.status not in ("failed", "cancelled"): - job.status = "cancelled" - job.error = job.error or "Cancelled by user request." - _finish(job, agent) - return True + # stop_flag was set after we already committed completed — the scan + # actually finished; cancel lost the race and is a no-op. + pass except JobCancelled: # Cooperative cancel via push_event after cancel/delete set _stop_flag. # cancel_job() already set terminal status; delete_job() may not have. diff --git a/src/engine.py b/src/engine.py index 9386434..4b7ab5b 100644 --- a/src/engine.py +++ b/src/engine.py @@ -721,10 +721,12 @@ def _on_ai_token(token: str): if _invs: art["investigations"] = _invs _emit("investigations", {"investigations": _invs}) - except Exception: # noqa: BLE001 - pass - except Exception: # noqa: BLE001 — reasoning must never affect a scan - pass + except Exception as e: # noqa: BLE001 + from src.cancel import reraise_if_cancelled # noqa: PLC0415 + reraise_if_cancelled(e) + except Exception as e: # noqa: BLE001 — reasoning must never affect a scan + from src.cancel import reraise_if_cancelled # noqa: PLC0415 + reraise_if_cancelled(e) # ── Change detection (Phase 7) — opt-in via --since-last, fail-soft ── # Diff this scan's observations against the prior persisted snapshot for the target. Produces @@ -743,8 +745,9 @@ def _on_ai_token(token: str): art["change_report"] = delta_report(delta) _emit("change", {"delta": art["change_delta"], "seed": art["change_seed"], "report": art["change_report"]}) - except Exception: # noqa: BLE001 — change detection never affects a scan - pass + except Exception as e: # noqa: BLE001 — change detection never affects a scan + from src.cancel import reraise_if_cancelled # noqa: PLC0415 + reraise_if_cancelled(e) # ── Architecture Summary (ALWAYS — deterministic synthesis, no AI needed) ── # Turn the scattered observations (frontend, hosting, CDN, auth, backend, TLS, DNS, endpoints) @@ -756,8 +759,9 @@ def _on_ai_token(token: str): if _arch is not None: art["architecture"] = _arch.to_dict() _emit("architecture", art["architecture"]) - except Exception: # noqa: BLE001 — synthesis is presentation-only, never block a scan - pass + except Exception as e: # noqa: BLE001 — synthesis is presentation-only, never block a scan + from src.cancel import reraise_if_cancelled # noqa: PLC0415 + reraise_if_cancelled(e) # ── AI Investigation Plan (opt-in: needs an AI key) — the overlay OVER the architecture ── # The AI reasons over the DETERMINISTIC architecture to propose grounded investigation objectives @@ -769,8 +773,9 @@ def _on_ai_token(token: str): if _plan: art["investigation_plan"] = _plan _emit("investigation_plan", {"objectives": _plan}) - except Exception: # noqa: BLE001 — the AI overlay never blocks a scan - pass + except Exception as e: # noqa: BLE001 — the AI overlay never blocks a scan + from src.cancel import reraise_if_cancelled # noqa: PLC0415 + reraise_if_cancelled(e) # ── Deterministic triage (ALWAYS — no --reason, no AI key required) ── # Rank + bucket the matched CVEs into "worth attention" vs "low-signal noise" so even a free scan @@ -782,8 +787,9 @@ def _on_ai_token(token: str): if _tri.attention or _tri.noise: art["triage"] = _tri.to_dict() _emit("triage", art["triage"]) - except Exception: # noqa: BLE001 — triage is presentation-only, never block a scan - pass + except Exception as e: # noqa: BLE001 — triage is presentation-only, never block a scan + from src.cancel import reraise_if_cancelled # noqa: PLC0415 + reraise_if_cancelled(e) return art diff --git a/test_local_agent_cancel.py b/test_local_agent_cancel.py new file mode 100644 index 0000000..6779a67 --- /dev/null +++ b/test_local_agent_cancel.py @@ -0,0 +1,108 @@ +"""Cancel vs local-agent finalisation: a cancelled scan must not be marked completed.""" +from __future__ import annotations + +import threading +from unittest.mock import MagicMock + +from api.agents.local_agent import _run_job +from api.agents.registry import AgentRegistry +from api.jobs.manager import JobCancelled, JobManager +from api.models.scan_request import ScanRequest +from api.storage.json_store import JsonScanStore + + +def _isolated_manager(tmp_path) -> JobManager: + mgr = JobManager.__new__(JobManager) + mgr._jobs = {} + mgr._lock = threading.RLock() + mgr._pg = False + mgr._pg_store = None + mgr.store = JsonScanStore(str(tmp_path)) + return mgr + + +def _install(monkeypatch, mgr: JobManager, registry: AgentRegistry) -> None: + monkeypatch.setattr("api.jobs.manager.job_manager", mgr) + monkeypatch.setattr("api.agents.registry.agent_registry", registry) + monkeypatch.setattr("api.jobs.executor.try_dispatch_queued", lambda org_id="": 0) + + +def _agent_and_job(mgr: JobManager, registry: AgentRegistry): + agent_id, _ = registry.register( + hostname="local-test", + capabilities=["scan"], + version="test", + tags={"type": "test"}, + org_id="org-a", + ) + registry.heartbeat(agent_id) + job = mgr.create(ScanRequest(target="example.com"), org_id="org-a") + job.assigned_agent_id = agent_id + assert registry.assign_task(agent_id, job.job_id) + return agent_id, job + + +def test_engine_return_after_cancel_stays_cancelled(monkeypatch, tmp_path): + """User cancel wins try_set_terminal; engine then returns normally. + + The pre-fix worker assigned job.status = "completed" unconditionally, so a + cancel in this window was persisted as a finished scan. + """ + mgr = _isolated_manager(tmp_path) + registry = AgentRegistry(persist_path=None) + _install(monkeypatch, mgr, registry) + agent_id, job = _agent_and_job(mgr, registry) + + def fake_scan(**_kwargs): + mgr.try_set_terminal( + job.job_id, "cancelled", org_id="org-a", + error="Cancelled by user request.", + ) + job._stop_flag.set() + + monkeypatch.setattr("src.json_bridge.run_streaming_scan", fake_scan) + + assert _run_job(agent_id, job.job_id, "org-a") is True + live = mgr.get(job.job_id, org_id="org-a") + assert live is not None + assert live.status == "cancelled" + assert live.error == "Cancelled by user request." + + +def test_jobcancelled_after_cancel_stays_cancelled(monkeypatch, tmp_path): + """Cooperative cancel: engine raises JobCancelled after cancel_job.""" + mgr = _isolated_manager(tmp_path) + registry = AgentRegistry(persist_path=None) + _install(monkeypatch, mgr, registry) + agent_id, job = _agent_and_job(mgr, registry) + + def fake_scan(**_kwargs): + mgr.try_set_terminal( + job.job_id, "cancelled", org_id="org-a", + error="Cancelled by user request.", + ) + job._stop_flag.set() + raise JobCancelled(job.job_id) + + monkeypatch.setattr("src.json_bridge.run_streaming_scan", fake_scan) + + assert _run_job(agent_id, job.job_id, "org-a") is True + live = mgr.get(job.job_id, org_id="org-a") + assert live is not None + assert live.status == "cancelled" + + +def test_successful_scan_still_completes(monkeypatch, tmp_path): + mgr = _isolated_manager(tmp_path) + registry = AgentRegistry(persist_path=None) + _install(monkeypatch, mgr, registry) + agent_id, job = _agent_and_job(mgr, registry) + + monkeypatch.setattr("src.json_bridge.run_streaming_scan", MagicMock()) + + assert _run_job(agent_id, job.job_id, "org-a") is True + live = mgr.get(job.job_id, org_id="org-a") + assert live is not None + assert live.status == "completed" + assert live.progress == 100.0 + assert any(e.get("type") == "done" for e in live.events)