Cancellation

vega-core uses cooperative cancellation at stage boundaries and forceful cancellation at external agent boundaries. Stages never terminate subprocesses directly — they call framework checkpoints, and the runner layer handles active agent cancellation.

Location

src/services/cancellation.py

CancellationToken protocol

The caller (vega-backend) supplies a cancellation signal:

class CancellationToken(Protocol):
    def is_cancelled(self) -> bool: ...

Pass it to configure_core_runtime(cancellation_token=token). Any object implementing .is_cancelled() works — a threading event wrapper, a database polling token, etc.

CancellationManager

CancellationManager is the framework's central cancellation hub. It wraps the caller-supplied token and adds coordination:

class CancellationManager:
    def __init__(self, token: CancellationToken | None = None) -> None: ...

    def is_cancelled(self) -> bool: ...

    def checkpoint(
        self,
        *,
        stage: StageName,
        progress: Progress | None = None,
    ) -> None: ...

    def register_active_agent_call(self, call_id: str, cancel: CancelFn) -> None: ...
    def unregister_active_agent_call(self, call_id: str) -> None: ...
    def cancel_active_agent_calls(self) -> None: ...

CancellationManager is available through RunContext.cancellation. Stages and BaseStage call context.checkpoint(stage, progress) at natural boundaries.

StageCancelled exception

class StageCancelled(Exception):
    stage: StageName
    progress: Progress | None

checkpoint(...) raises StageCancelled when is_cancelled() is true. BaseStage.run(...) catches this, calls EventBus.stage_cancelled(...), and returns a StageOutput(status=StageStatus.CANCELLED).

How BaseStage uses cancellation

BaseStage.run(...) calls checkpoints automatically:

1. checkpoint(stage) — before calling execute(...)
2. execute(context, input) runs
3. checkpoint(stage, progress) — after execute returns (or per-item in loops)

Concrete stages that process multiple items (audit components, bugs) call context.checkpoint(stage, progress) in their inner loop — between items, not inside agent calls.

How AgentRunner uses cancellation

AgentRunner registers each active agent call with CancellationManager.register_active_agent_call(call_id, cancel_fn). The cancel_fn sends SIGTERM (then SIGKILL) to the subprocess.

When checkpoint(...) detects cancellation, it first calls cancel_active_agent_calls() to immediately terminate any in-flight agent subprocess, then raises StageCancelled.

Cancellation sequence (backend → core)

vega-backend sets cancellation_token.is_cancelled() → True
  ↓
CancellationManager.checkpoint() detects it
  ↓
CancellationManager.cancel_active_agent_calls()
  → CodexAdapter.cancel(task_id) → SIGTERM → SIGKILL
  ↓
StageCancelled raised
  ↓
BaseStage catches it → emits stage_cancelled event
  ↓
StageOrchestrator stops the sequence
  ↓
CoreResult(status="cancelled") returned to vega-backend

Checkpoint placement per stage

Stage	Checkpoint placement
PlanStage	Before and after agent call
BaseThreatModelStage	Before and after agent call
BaseAuditStage	Between each component; before and after each component agent call
VerifyStage	Between each bug; before and after each bug agent call
PatchStage	Between each finding; between patch pipeline steps per finding

Implementing a cancellation token

Thread-event token (local dev)

import threading

class ThreadCancellationToken:
    def __init__(self):
        self._event = threading.Event()

    def cancel(self):
        self._event.set()

    def is_cancelled(self) -> bool:
        return self._event.is_set()

Database polling token (production)

class DatabaseCancellationToken:
    def __init__(self, scan_id: str, db):
        self._scan_id = scan_id
        self._db = db

    def is_cancelled(self) -> bool:
        return self._db.query(
            "SELECT cancelled FROM scans WHERE id = %s", (self._scan_id,)
        )

Pass the token at run configuration time:

configure_core_runtime(
    event_sink=my_sink,
    cancellation_token=DatabaseCancellationToken(scan_id, db),
)

Graceful vs forceful cancellation

• Cooperative (stages): checkpoint(...) raises at a natural boundary. Ongoing work finishes cleanly before the exception is raised.
• Forceful (agents): cancel_active_agent_calls() sends OS signals to subprocesses. This is immediate but can leave partial output.

The CoreResult returned for a cancelled run has status="cancelled". Backend should update the scan status accordingly and not treat it as a failure.