Add collect stop, cutover nav prefetch/cache, and lane labels.

Let operators abort queued/running biz_state batches, speed cutover page switches, and show light/heavy collect lanes on monitor items.

Co-authored-by: Cursor <cursoragent@cursor.com>

netx_api/biz_state/collect_runner.py

@@ -45,6 +45,13 @@ from .command_match import (
 )
 from .parsers import get_parser
 from .profiles import get_profile
+from .collect_stop import (
+    STOP_USER_MESSAGE,
+    clear_stop_requested,
+    is_stop_requested,
+    register_lane_holder,
+    unregister_lane_holder,
+)
 
 _log = logging.getLogger("netx.biz_state.runner")
 
@@ -457,17 +464,32 @@ def execute_claimed_batch(
 
     error = ""
     try:
-        _run_collect_session(
-            task_id=task_id,
-            batch_id=batch_id,
-            source=source,
-            ne_id=ne_id,
-            vendor=vendor,
-            device_type=device_type,
-        )
+        if is_stop_requested(batch_id):
+            _finalize_batch_status(
+                batch_id=batch_id,
+                task_id=task_id,
+                cmd_count=0,
+                total_rows=0,
+                any_fail=False,
+                any_ok=False,
+                lane_errors=[],
+                stopped=True,
+            )
+            error = STOP_USER_MESSAGE
+        else:
+            _run_collect_session(
+                task_id=task_id,
+                batch_id=batch_id,
+                source=source,
+                ne_id=ne_id,
+                vendor=vendor,
+                device_type=device_type,
+            )
     except Exception as exc:
         _log.exception("biz_state collect failed task=%s batch=%s", task_id, batch_id)
         error = _format_error(exc)
+        if "_stopped" in error or is_stop_requested(batch_id):
+            error = STOP_USER_MESSAGE
         try:
             _fail_batch_status(batch_id, error)
         except Exception:
@@ -509,10 +531,14 @@ def _run_collect_lane(
     budget = min(int(cap), int(per_cmd) * max(1, len(work)) + 90)
     holder: dict[str, Any] = {}
     flush_every = persist_every_cmds()
+    register_lane_holder(batch_id, holder)
 
     def _session() -> tuple[int, int, bool, bool]:
         from ..ne_netmiko import drain_read_channel
 
+        if is_stop_requested(batch_id):
+            raise TimeoutError(f"{label}_stopped")
+
         conn = open_netmiko_connection(creds, session_timeout=budget)
         holder["conn"] = conn
         total_rows = 0
@@ -636,8 +662,10 @@ def _run_collect_lane(
                     flat_work.append((cmd, p, profile_id, item_id, "normal"))
 
             for concrete, params, profile_id, item_id, mode in flat_work:
-                if holder.get("timed_out"):
-                    raise TimeoutError(f"{label}_aborted")
+                if holder.get("timed_out") or holder.get("stop_requested") or is_stop_requested(
+                    batch_id
+                ):
+                    raise TimeoutError(f"{label}_stopped")
 
                 cmd_id = uuid4().hex
                 raw_text = ""
@@ -923,6 +951,8 @@ def _run_collect_lane(
         )
     except TimeoutError as exc:
         raise RuntimeError(str(exc)[:1020]) from exc
+    finally:
+        unregister_lane_holder(batch_id, holder)
 
 
 def _absorb_lane_result(
@@ -1052,6 +1082,7 @@ def _finalize_batch_status(
     any_fail: bool,
     any_ok: bool,
     lane_errors: list[str],
+    stopped: bool = False,
 ) -> str:
     """Write terminal batch status on a fresh Session (retry once on disconnect)."""
 
@@ -1063,7 +1094,14 @@ def _finalize_batch_status(
         batch.command_count = max(int(batch.command_count or 0), int(cmd_count or 0))
         batch.row_count = max(int(batch.row_count or 0), int(total_rows or 0))
         batch.ended_at = _utcnow()
-        if any_fail and any_ok:
+        if stopped:
+            if any_ok or int(batch.command_count or 0) > 0 or int(batch.row_count or 0) > 0:
+                batch.status = "partial"
+                batch.message = STOP_USER_MESSAGE
+            else:
+                batch.status = "cancelled"
+                batch.message = STOP_USER_MESSAGE
+        elif any_fail and any_ok:
             batch.status = "partial"
             if lane_errors:
                 batch.message = "; ".join(lane_errors)[:1020]
@@ -1077,7 +1115,7 @@ def _finalize_batch_status(
             batch.message = ""
         status = str(batch.status or "")
         db.commit()
-        if status in ("success", "partial"):
+        if status in ("success", "partial") and not stopped:
             try:
                 from .compare_service import try_auto_compare_for_task
 
@@ -1086,28 +1124,40 @@ def _finalize_batch_status(
                 _log.exception("biz_state auto compare hook failed task=%s", task_id)
         return status
 
-    return str(
-        _run_db_with_reconnect(_write, label="biz_state_finalize") or ""
-    )
+    try:
+        return str(
+            _run_db_with_reconnect(_write, label="biz_state_finalize") or ""
+        )
+    finally:
+        clear_stop_requested(batch_id)
 
 
 def _fail_batch_status(batch_id: str, error: str) -> None:
     """Mark batch failed on a fresh Session (retry once on disconnect)."""
     msg = str(error or "")[:1020]
+    stopped = "_stopped" in msg or is_stop_requested(batch_id)
 
     def _write(db) -> None:
         batch = db.get(BizStateBatch, batch_id)
         if not batch:
             return
-        if str(batch.status or "") != "running":
+        st = str(batch.status or "")
+        if st not in ("running", "queued"):
             return
-        batch.status = "failed"
-        batch.message = msg
+        if stopped:
+            has_progress = int(batch.command_count or 0) > 0 or int(batch.row_count or 0) > 0
+            batch.status = "partial" if has_progress else "cancelled"
+            batch.message = STOP_USER_MESSAGE
+        else:
+            batch.status = "failed"
+            batch.message = msg
         batch.ended_at = _utcnow()
         db.commit()
 
-    _run_db_with_reconnect(_write, label="biz_state_fail_batch")
-
+    try:
+        _run_db_with_reconnect(_write, label="biz_state_fail_batch")
+    finally:
+        clear_stop_requested(batch_id)
 
 def _run_collect_session(
     *,
@@ -1340,6 +1390,18 @@ def _run_collect_session(
     if lane_errors and not any_ok and cmd_count == 0:
         # Progressive bumps may already have cmds; only hard-fail if nothing landed.
         if not _batch_has_progress(batch_id):
+            if is_stop_requested(batch_id) or any("_stopped" in e for e in lane_errors):
+                _finalize_batch_status(
+                    batch_id=batch_id,
+                    task_id=task_id,
+                    cmd_count=cmd_count,
+                    total_rows=total_rows,
+                    any_fail=False,
+                    any_ok=False,
+                    lane_errors=lane_errors,
+                    stopped=True,
+                )
+                return
             raise RuntimeError("; ".join(lane_errors)[:1020])
         any_fail = True
 
@@ -1358,6 +1420,7 @@ def _run_collect_session(
     except Exception:
         _log.exception("biz_state persist barrier before finalize failed batch=%s", batch_id)
 
+    stopped = is_stop_requested(batch_id) or any("_stopped" in e for e in lane_errors)
     _finalize_batch_status(
         batch_id=batch_id,
         task_id=task_id,
@@ -1366,6 +1429,7 @@ def _run_collect_session(
         any_fail=any_fail,
         any_ok=any_ok,
         lane_errors=lane_errors,
+        stopped=stopped,
     )
 
     def _purge(db) -> None:

netx_api/biz_state/collect_stop.py

@@ -0,0 +1,221 @@
+"""Stop an in-flight / queued biz_state collect batch.
+
+Queued batches are cancelled in DB immediately. Running lanes poll
+``is_stop_requested`` between commands (DB-backed for dedicated workers)
+and process-local holders are force-closed when stop is requested in the
+same process.
+"""
+
+from __future__ import annotations
+
+import logging
+import threading
+import time
+from typing import Any
+
+from ..db import SessionLocal
+from ..models import BizStateBatch, BizStateTask
+from ..ne_session_factory import close_netmiko_connection
+from ..timeutil import utcnow_naive
+
+_log = logging.getLogger("netx.biz_state.collect_stop")
+
+STOP_REQUEST_TOKEN = "stop_requested"
+STOP_USER_MESSAGE = "stopped_by_user"
+
+_lock = threading.Lock()
+_stop_batches: set[str] = set()
+_holders: dict[str, list[dict[str, Any]]] = {}
+# batch_id → (monotonic_ts, stop_requested)
+_db_cache: dict[str, tuple[float, bool]] = {}
+_DB_CACHE_TTL_SEC = 1.5
+
+
+def register_lane_holder(batch_id: str, holder: dict[str, Any]) -> None:
+    bid = str(batch_id or "").strip()
+    if not bid or holder is None:
+        return
+    with _lock:
+        _holders.setdefault(bid, []).append(holder)
+
+
+def unregister_lane_holder(batch_id: str, holder: dict[str, Any]) -> None:
+    bid = str(batch_id or "").strip()
+    if not bid:
+        return
+    with _lock:
+        lst = _holders.get(bid) or []
+        try:
+            lst.remove(holder)
+        except ValueError:
+            pass
+        if not lst:
+            _holders.pop(bid, None)
+
+
+def mark_stop_requested(batch_id: str) -> None:
+    bid = str(batch_id or "").strip()
+    if not bid:
+        return
+    with _lock:
+        _stop_batches.add(bid)
+        _db_cache[bid] = (time.monotonic(), True)
+
+
+def clear_stop_requested(batch_id: str) -> None:
+    bid = str(batch_id or "").strip()
+    if not bid:
+        return
+    with _lock:
+        _stop_batches.discard(bid)
+        _db_cache.pop(bid, None)
+        _holders.pop(bid, None)
+
+
+def is_stop_requested(batch_id: str) -> bool:
+    """True if user asked to stop this batch (memory or DB)."""
+    bid = str(batch_id or "").strip()
+    if not bid:
+        return False
+    with _lock:
+        if bid in _stop_batches:
+            return True
+        hit = _db_cache.get(bid)
+        now = time.monotonic()
+        if hit and (now - hit[0]) < _DB_CACHE_TTL_SEC:
+            return bool(hit[1])
+
+    stopped = _read_stop_from_db(bid)
+    with _lock:
+        _db_cache[bid] = (time.monotonic(), stopped)
+        if stopped:
+            _stop_batches.add(bid)
+    return stopped
+
+
+def _read_stop_from_db(batch_id: str) -> bool:
+    db = SessionLocal()
+    try:
+        batch = db.get(BizStateBatch, batch_id)
+        if not batch:
+            return False
+        st = str(batch.status or "")
+        if st in ("cancelled", "success", "partial", "failed"):
+            # Already terminal — treat as stop so lanes exit quickly.
+            return st == "cancelled" or STOP_REQUEST_TOKEN in str(batch.message or "")
+        msg = str(batch.message or "")
+        return msg.startswith(STOP_REQUEST_TOKEN) or msg == STOP_USER_MESSAGE
+    except Exception:
+        _log.exception("read stop flag failed batch=%s", batch_id)
+        return False
+    finally:
+        db.close()
+
+
+def _force_close_holders(batch_id: str) -> int:
+    bid = str(batch_id or "").strip()
+    with _lock:
+        holders = list(_holders.get(bid) or [])
+    n = 0
+    for holder in holders:
+        try:
+            holder["stop_requested"] = True
+            holder["timed_out"] = True
+            conn = holder.get("conn")
+            if conn is not None:
+                close_netmiko_connection(conn)
+                n += 1
+        except Exception:
+            _log.exception("force-close on stop failed batch=%s", bid)
+    return n
+
+
+def request_stop_collect(task_id: str) -> dict[str, Any]:
+    """Cancel queued batches and signal running collect for this task to abort."""
+    tid = str(task_id or "").strip()
+    if not tid:
+        return {"ok": False, "reason": "missing_task_id"}
+
+    db = SessionLocal()
+    cancelled_ids: list[str] = []
+    signaled_ids: list[str] = []
+    closed = 0
+    try:
+        task = db.get(BizStateTask, tid)
+        if not task:
+            return {"ok": False, "reason": "task_not_found", "task_id": tid}
+
+        active = (
+            db.query(BizStateBatch)
+            .filter(
+                BizStateBatch.task_id == tid,
+                BizStateBatch.status.in_(("queued", "running")),
+            )
+            .all()
+        )
+        if not active:
+            # Nothing to stop — clear sticky collect_running if orphaned.
+            if bool(task.collect_running):
+                task.collect_running = False
+                if hasattr(task, "collect_queued_at"):
+                    task.collect_queued_at = None
+                task.updated_at = utcnow_naive()
+                db.commit()
+            return {
+                "ok": True,
+                "task_id": tid,
+                "stopped": False,
+                "reason": "not_collecting",
+                "cancelled_batches": [],
+                "signaled_batches": [],
+            }
+
+        now = utcnow_naive()
+        still_running = False
+        for batch in active:
+            bid = str(batch.id)
+            mark_stop_requested(bid)
+            if str(batch.status or "") == "queued":
+                batch.status = "cancelled"
+                batch.message = STOP_USER_MESSAGE
+                batch.ended_at = now
+                cancelled_ids.append(bid)
+            else:
+                batch.message = STOP_REQUEST_TOKEN
+                signaled_ids.append(bid)
+                still_running = True
+            closed += _force_close_holders(bid)
+
+        if not still_running:
+            task.collect_running = False
+            if hasattr(task, "collect_queued_at"):
+                task.collect_queued_at = None
+            task.last_collect_ended_at = now
+            task.last_error = STOP_USER_MESSAGE
+            task.updated_at = now
+
+        db.commit()
+        _log.info(
+            "stop collect task=%s cancelled=%s signaled=%s closed_conns=%s",
+            tid,
+            cancelled_ids,
+            signaled_ids,
+            closed,
+        )
+        return {
+            "ok": True,
+            "task_id": tid,
+            "stopped": True,
+            "cancelled_batches": cancelled_ids,
+            "signaled_batches": signaled_ids,
+            "closed_connections": closed,
+        }
+    except Exception:
+        _log.exception("request_stop_collect failed task=%s", tid)
+        try:
+            db.rollback()
+        except Exception:
+            pass
+        return {"ok": False, "reason": "stop_error", "task_id": tid}
+    finally:
+        db.close()

netx_api/biz_state_router.py

@@ -252,6 +252,19 @@ def api_collect_now(
     return {"ok": True, "started": True, "queued": True, "task_id": task_id}
 
 
+@router.post("/tasks/{task_id}/collect/stop")
+def api_collect_stop(task_id: str, db: Session = Depends(get_db)) -> dict[str, Any]:
+    """Stop queued or running collect for this task (best-effort mid-command)."""
+    from .biz_state.collect_stop import request_stop_collect
+
+    task = db.get(BizStateTask, task_id)
+    if not task:
+        from fastapi import HTTPException
+
+        raise HTTPException(status_code=404, detail="task not found")
+    return request_stop_collect(task_id)
+
+
 @router.get("/tasks/{task_id}/batches")
 def api_list_batches(
     task_id: str, limit: int = 50, db: Session = Depends(get_db)

netx_api/models/biz_state.py

@@ -90,7 +90,7 @@ class BizStateBatch(Base):
     ne_id: Mapped[str] = mapped_column(String(128), default="", index=True)
     ne_name: Mapped[str] = mapped_column(String(256), default="")
     vendor: Mapped[str] = mapped_column(String(64), default="")
-    status: Mapped[str] = mapped_column(String(32), default="running", index=True)  # queued|running|success|partial|failed
+    status: Mapped[str] = mapped_column(String(32), default="running", index=True)  # queued|running|success|partial|failed|cancelled
     command_count: Mapped[int] = mapped_column(Integer, default=0)
     row_count: Mapped[int] = mapped_column(Integer, default=0)
     message: Mapped[str] = mapped_column(String(1024), default="")