af_workflow/

supervisor.rs

//! Database-agnostic scheduler for claimed long-running workflow instances.

use af_context::{RunId, SubjectId, TenantId};
use std::collections::{HashMap, VecDeque};
use std::future::Future;
use std::panic::AssertUnwindSafe;
use std::sync::Arc;

use async_trait::async_trait;
use chrono::{DateTime, Utc};
use futures::{stream::FuturesUnordered, FutureExt, StreamExt};
use serde_json::Value;

use crate::{
    ActionIntent, CapabilityPin, ControlEpochs, LifecyclePolicy, NodeRegistry, SpecDriver,
    WorkflowRevision,
};

/// A durable fact that became due for a claimed instance: an expired timer or
/// a pending trigger delivery. The driver reads it; the queue marks it consumed
/// in the same transaction that commits the driver's command.
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct Wakeup {
    /// Stable identifier of this record.
    pub id: String,
    /// Discriminator naming the variant of this record.
    pub kind: String,
    /// Target branch; absent only for legacy single-branch work or system timers.
    #[serde(default)]
    pub branch_id: Option<af_context::BranchId>,
    /// Structured payload.
    pub payload: Value,
}

/// One claimed instance with everything a driver needs to evaluate it.
#[derive(Debug, Clone, PartialEq)]
pub struct WorkItem {
    /// Stable identifier of this record.
    pub id: String,
    /// Run this record belongs to.
    pub run_id: RunId,
    /// Tenant that owns this record.
    pub tenant_id: TenantId,
    /// Subject (user or service principal) acting on or owning this record.
    pub subject_id: SubjectId,
    /// Workflow spec identifier.
    pub spec_id: String,
    /// Workflow definition this record belongs to.
    pub definition_id: String,
    /// Pinned workflow revision.
    pub workflow_revision: u64,
    /// Digest of the pinned revision.
    pub workflow_revision_digest: String,
    /// Pinned execution profile.
    pub execution_profile_id: String,
    /// Pinned execution profile revision.
    pub execution_profile_revision: u64,
    /// Digest of the pinned profile.
    pub execution_profile_digest: String,
    /// Kernel ABI the revision pinned.
    pub kernel_abi_version: String,
    /// Capabilities the revision may dispatch.
    pub capability_pins: Vec<CapabilityPin>,
    /// Completion, timeout and catch-up policy pinned by the revision.
    pub lifecycle: LifecyclePolicy,
    /// Database timestamps captured by the claim transaction. Scheduling logic
    /// never consults an application-host clock.
    pub scheduled_at: DateTime<Utc>,
    /// Immutable database creation time, anchoring each initial branch schedule.
    pub created_at: DateTime<Utc>,
    /// Database time of the claim.
    pub claimed_at: DateTime<Utc>,
    /// Configuration object validated against the declared schema.
    pub config: Value,
    /// CAS version the command must be committed against.
    pub state_version: i64,
    /// Control epochs at claim time.
    pub control_epochs: ControlEpochs,
    /// Whether cancellation was requested; the terminal commit honours it.
    pub cancel_requested: bool,
    /// Fencing token bumped on every lease acquisition; stale holders cannot write.
    pub lease_version: i64,
    #[doc = "Due timers and trigger deliveries, oldest first."]
    pub wakeups: Vec<Wakeup>,
}

impl WorkItem {
    /// Whether this claim only needs cancellation or a trusted terminal timer to converge.
    pub fn is_converging(&self) -> bool {
        self.cancel_requested
            || self
                .wakeups
                .iter()
                .any(|wakeup| wakeup.kind == "timer" && wakeup.payload["kind"] == "terminal_action")
    }

    fn validate_pins(&self) -> Result<(), String> {
        for (name, value) in [
            ("definition_id", self.definition_id.as_str()),
            (
                "workflow_revision_digest",
                self.workflow_revision_digest.as_str(),
            ),
            ("execution_profile_id", self.execution_profile_id.as_str()),
            (
                "execution_profile_digest",
                self.execution_profile_digest.as_str(),
            ),
            ("kernel_abi_version", self.kernel_abi_version.as_str()),
        ] {
            if value.trim().is_empty() {
                return Err(format!("work item is missing pinned {name}"));
            }
        }
        if self.workflow_revision == 0 || self.execution_profile_revision == 0 {
            return Err("work item revision pins must be positive".into());
        }
        uuid::Uuid::parse_str(&self.run_id)
            .map_err(|error| format!("work item has invalid run_id: {error}"))?;
        if self.state_version < 0
            || self.lease_version <= 0
            || self.control_epochs.tenant < 0
            || self.control_epochs.instance < 0
        {
            return Err("work item state, lease and control pins must be current".into());
        }
        Ok(())
    }
}

/// What the queue does with the instance after committing a command.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum WorkDisposition {
    /// Evaluate again after a delay.
    Continue {
        /// Seconds until the next evaluation.
        delay_secs: i64,
    },
    /// Evaluate again at an exact time.
    Reschedule {
        /// Next evaluation time.
        at: DateTime<Utc>,
    },
    /// The instance is finished.
    Complete,
    /// The evaluation failed.
    Failed {
        /// Failure reason.
        error: String,
        /// When to retry; `None` fails the instance permanently.
        retry_at: Option<DateTime<Utc>>,
    },
}

/// Completion-policy accounting for one evaluation.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct EvaluationOutcome {
    /// The evaluation consumed a trigger or schedule occurrence.
    pub triggered: bool,
    /// At least one event passed a material step.
    pub matched: bool,
    /// The graph reached a sink or retained an event through its final step.
    pub succeeded: bool,
    /// The evaluation observed a terminal action outcome.
    pub action_terminal: bool,
}

/// The only result a workflow driver may return. The queue commits the state,
/// event, intents and scheduling disposition atomically before any provider is
/// allowed to dispatch an external effect.
#[derive(Debug, Clone, PartialEq)]
pub struct WorkflowTransitionCommand {
    /// Idempotency key of this transition.
    pub delivery_key: String,
    /// Content digest; a reused key with a different digest is a conflict.
    pub delivery_digest: String,
    /// Stable machine-readable event type.
    pub event_type: String,
    /// Digest of the authoritative event.
    pub event_digest: String,
    /// Authoritative event payload.
    pub event_payload: Value,
    /// New authoritative state.
    pub next_state: Value,
    /// Intents to persist in the same transaction.
    pub action_intents: Vec<ActionIntent>,
    /// Completion-policy accounting for this transition.
    pub outcome: EvaluationOutcome,
    /// Claimed wakeups actually consumed by this transition; failed evaluations consume none.
    pub consumed_wakeups: Vec<String>,
    /// Scheduling disposition.
    pub disposition: WorkDisposition,
}

impl WorkflowTransitionCommand {
    /// Command recording an evaluation failure and retrying at `retry_at`.
    pub fn failure(item: &WorkItem, error: impl Into<String>, retry_at: DateTime<Utc>) -> Self {
        let error = error.into();
        Self {
            delivery_key: format!("supervisor:{}:{}", item.id, item.state_version),
            delivery_digest: format!("failure:{}:{}", item.lease_version, error),
            event_type: "workflow.evaluation_failed".into(),
            event_digest: format!("failure:{}:{}", item.state_version, error),
            event_payload: serde_json::json!({"error": error.clone()}),
            next_state: item.config.clone(),
            action_intents: Vec::new(),
            outcome: EvaluationOutcome::default(),
            consumed_wakeups: Vec::new(),
            disposition: WorkDisposition::Failed {
                error,
                retry_at: Some(retry_at),
            },
        }
    }

    /// Validate pins, effects and acknowledgements against the claimed work.
    pub fn validate_for(&self, item: &WorkItem) -> Result<(), String> {
        item.validate_pins()?;
        for (name, value) in [
            ("delivery_key", self.delivery_key.as_str()),
            ("delivery_digest", self.delivery_digest.as_str()),
            ("event_type", self.event_type.as_str()),
            ("event_digest", self.event_digest.as_str()),
        ] {
            if value.trim().is_empty() {
                return Err(format!("workflow command is missing {name}"));
            }
        }
        if item.cancel_requested && !self.action_intents.is_empty() {
            return Err("cancelled work cannot prepare external actions".into());
        }
        let consumed: std::collections::BTreeSet<_> = self.consumed_wakeups.iter().collect();
        if consumed.len() != self.consumed_wakeups.len()
            || consumed
                .iter()
                .any(|id| !item.wakeups.iter().any(|wakeup| &wakeup.id == *id))
            || (!consumed.is_empty()
                && matches!(self.disposition, WorkDisposition::Failed { .. })
                && !self.outcome.triggered)
        {
            return Err("workflow command acknowledges unconsumed or unclaimed work".into());
        }
        for intent in &self.action_intents {
            intent
                .validate_prepared()
                .map_err(|error| error.to_string())?;
            if intent.tenant_id != item.tenant_id || intent.instance_id != item.id {
                return Err("action intent escapes claimed work scope".into());
            }
            if intent.run_id != item.run_id {
                return Err("action intent escapes claimed workflow run".into());
            }
            // Resource-scope epochs are owned by the store: the queue does not
            // know which resource an intent targets, so it only pins what it
            // claimed (tenant + instance) and the commit transaction proves the
            // resource epoch against `workflow_control_epochs`.
            if intent.control_epochs.tenant != item.control_epochs.tenant
                || intent.control_epochs.instance != item.control_epochs.instance
                || intent.lease_epoch != item.lease_version
            {
                return Err("action intent uses stale control or lease pins".into());
            }
            if !item
                .capability_pins
                .iter()
                .any(|pin| pin == &intent.capability)
            {
                return Err(format!(
                    "action capability '{}' is not pinned by the workflow revision",
                    intent.capability.id
                ));
            }
        }
        Ok(())
    }
}

/// Result of one pass.
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
pub struct SupervisorStats {
    /// Items claimed.
    pub claimed: usize,
    /// Items whose evaluation failed.
    pub failed: usize,
}

/// Pass configuration.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SupervisorSettings {
    /// Identity of the worker holding the lease.
    pub worker_id: String,
    /// Lease duration per claim.
    pub lease_secs: i64,
    /// Retry delay after a failed evaluation.
    pub requeue_delay_secs: i64,
    /// Maximum items per pass.
    pub claim_batch: i64,
    /// Concurrent evaluations.
    pub concurrency: usize,
}

impl SupervisorSettings {
    fn concurrency(&self) -> usize {
        self.concurrency.max(1)
    }
}

/// Failure at the supervisor boundary. Driver evaluation reasons stay `String`
/// because they are product-authored text recorded verbatim in the failure
/// event; everything the kernel itself decides is typed here.
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum SupervisorError {
    /// The durable work queue could not claim, renew or commit.
    #[error("work queue: {0}")]
    Queue(String),
    /// A driver's spec set failed validation or two drivers claim one spec.
    #[error("driver '{driver}': {reason}")]
    Driver {
        /// Driver name.
        driver: String,
        /// Why.
        reason: String,
    },
    /// The per-pass context factory failed; claimed work was requeued.
    #[error("context: {0}")]
    Context(String),
    /// One or more transition commits failed after evaluation.
    #[error("workflow state commit failed: {0}")]
    Commit(String),
}

/// Durable queue the supervisor claims from and commits to.
#[async_trait]
pub trait WorkQueue: Send + Sync {
    /// Claim due instances for `spec_ids`; an empty slice selects all specs.
    async fn claim_due(
        &self,
        spec_ids: &[String],
        worker_id: &str,
        lease_secs: i64,
        batch: i64,
    ) -> Result<Vec<WorkItem>, SupervisorError>;

    /// Read the immutable revision of an owned, currently leased instance.
    /// Stores without revision persistence reject dynamic spec execution.
    async fn load_revision(&self, _item: &WorkItem) -> Result<WorkflowRevision, SupervisorError> {
        Err(SupervisorError::Queue(
            "published revision loading is unavailable".into(),
        ))
    }

    /// Recheck current owner admission before any graph/read or custom-driver evaluation.
    /// Queues without an authority source fail closed. Terminal system observations
    /// and cancellation still enter their existing convergence paths.
    async fn admit_evaluation(&self, _item: &WorkItem) -> Result<(), SupervisorError> {
        Err(SupervisorError::Queue(
            "workflow evaluation admission is unavailable".into(),
        ))
    }

    /// Renew a claim's lease.
    async fn renew(
        &self,
        tenant_id: &str,
        id: &str,
        worker_id: &str,
        lease_version: i64,
        lease_secs: i64,
    ) -> Result<(), SupervisorError>;
    /// Commit a driver command atomically.
    async fn commit_command(
        &self,
        item: &WorkItem,
        command: &WorkflowTransitionCommand,
    ) -> Result<(), SupervisorError>;
}

/// In-memory [`WorkQueue`] for consumer unit tests: hands out queued items in
/// order and records every committed command. No leases, no durability.
#[derive(Default)]
pub struct MemoryWorkQueue {
    items: std::sync::Mutex<Vec<WorkItem>>,
    committed: std::sync::Mutex<Vec<(WorkItem, WorkflowTransitionCommand)>>,
    renewals: std::sync::Mutex<Vec<(String, i64)>>,
}

impl MemoryWorkQueue {
    /// Queue pre-loaded with `items`.
    pub fn new(items: Vec<WorkItem>) -> Self {
        Self {
            items: std::sync::Mutex::new(items),
            ..Self::default()
        }
    }

    /// Enqueue an item.
    pub fn push(&self, item: WorkItem) {
        self.items
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .push(item);
    }

    /// Every `(item, command)` pair committed so far, in commit order.
    pub fn committed(&self) -> Vec<(WorkItem, WorkflowTransitionCommand)> {
        self.committed
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .clone()
    }

    /// `(instance id, lease version)` for every renewal.
    pub fn renewals(&self) -> Vec<(String, i64)> {
        self.renewals
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .clone()
    }
}

#[async_trait]
impl WorkQueue for MemoryWorkQueue {
    async fn admit_evaluation(&self, _: &WorkItem) -> Result<(), SupervisorError> {
        Ok(())
    }
    async fn claim_due(
        &self,
        spec_ids: &[String],
        _worker_id: &str,
        _lease_secs: i64,
        batch: i64,
    ) -> Result<Vec<WorkItem>, SupervisorError> {
        let mut items = self
            .items
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        let mut claimed = Vec::new();
        let mut index = 0;
        while index < items.len() && claimed.len() < batch.max(0) as usize {
            if spec_ids.is_empty() || spec_ids.contains(&items[index].spec_id) {
                claimed.push(items.remove(index));
            } else {
                index += 1;
            }
        }
        Ok(claimed)
    }

    async fn renew(
        &self,
        _tenant_id: &str,
        id: &str,
        _worker_id: &str,
        lease_version: i64,
        _lease_secs: i64,
    ) -> Result<(), SupervisorError> {
        self.renewals
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .push((id.to_owned(), lease_version));
        Ok(())
    }

    async fn commit_command(
        &self,
        item: &WorkItem,
        command: &WorkflowTransitionCommand,
    ) -> Result<(), SupervisorError> {
        self.committed
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .push((item.clone(), command.clone()));
        Ok(())
    }
}

/// Deterministic evaluator for the specs it owns.
#[async_trait]
pub trait WorkflowDriver<Context>: Send + Sync
where
    Context: Send + Sync,
{
    /// Stable driver name.
    fn name(&self) -> &'static str;
    /// Specs this driver owns.
    fn spec_ids(&self) -> Vec<&str>;
    /// Reject a spec set the driver cannot serve; the reason is reported as
    /// [`SupervisorError::Driver`].
    fn validate_specs(&self) -> Result<(), String>;
    /// Explicitly admit the exact revision and ABI a custom driver implements.
    /// The default refuses execution rather than silently accepting new revisions.
    fn supports_revision(&self, _item: &WorkItem) -> bool {
        false
    }
    /// Evaluate deterministic workflow state and return one durable command.
    /// External effects must be represented as prepared `ActionIntent`s and
    /// dispatched later by a registered provider.
    async fn evaluate(
        &self,
        context: &Context,
        item: &WorkItem,
    ) -> Result<WorkflowTransitionCommand, String>;
}

/// Drivers keyed by the specs they own.
pub struct DriverRegistry<Context: Send + Sync> {
    drivers: Vec<Arc<dyn WorkflowDriver<Context>>>,
    by_spec: HashMap<String, Arc<dyn WorkflowDriver<Context>>>,
    nodes: Option<Arc<NodeRegistry>>,
    startup_specs: Vec<String>,
    compiled: tokio::sync::Mutex<VecDeque<CompiledRevision<Context>>>,
}

struct CompiledRevision<Context: Send + Sync> {
    scope: Option<(TenantId, af_context::WorkflowDefinitionId)>,
    key: (String, String),
    content: Value,
    driver: Arc<dyn WorkflowDriver<Context>>,
}

const COMPILED_REVISION_LIMIT: usize = 256;

impl<Context: Send + Sync> Default for DriverRegistry<Context> {
    fn default() -> Self {
        Self {
            drivers: Vec::new(),
            by_spec: HashMap::new(),
            nodes: None,
            startup_specs: Vec::new(),
            compiled: Default::default(),
        }
    }
}

impl<Context: Send + Sync> DriverRegistry<Context> {
    /// Empty registry.
    pub fn new() -> Self {
        Self::default()
    }

    /// Enable lazy execution of published specs using this immutable registry snapshot.
    pub fn with_node_registry(mut self, nodes: NodeRegistry) -> Self {
        self.compiled.get_mut().clear();
        self.startup_specs.clear();
        self.nodes = Some(Arc::new(nodes));
        self
    }

    /// Precompile a shipped graph. Execution still requires its published revision.
    pub fn prewarm_spec(&mut self, spec: &crate::Spec) -> Result<(), String> {
        let nodes = self.nodes.as_ref().ok_or("node registry is unavailable")?;
        let driver = SpecDriver::new(spec, nodes).map_err(|error| error.to_string())?;
        let compiled = self.compiled.get_mut();
        if compiled.len() == COMPILED_REVISION_LIMIT {
            compiled.pop_front();
        }
        compiled.push_back(CompiledRevision {
            scope: None,
            key: (spec.spec_id.clone(), String::new()),
            content: serde_json::to_value(spec).map_err(|error| error.to_string())?,
            driver: Arc::new(driver),
        });
        if !self.startup_specs.contains(&spec.spec_id) {
            self.startup_specs.push(spec.spec_id.clone());
        }
        Ok(())
    }

    /// Validate and compile the exact published graph, caching at most 256 revisions.
    /// Callers must load the revision through the owner-scoped durable queue first.
    pub async fn resolve_revision(
        &self,
        item: &WorkItem,
        revision: WorkflowRevision,
    ) -> Result<Arc<dyn WorkflowDriver<Context>>, String>
    where
        Context: 'static,
    {
        item.validate_pins()?;
        revision.validate().map_err(|error| error.to_string())?;
        if revision.definition_id != item.definition_id
            || revision.revision != item.workflow_revision
            || revision.content_digest != item.workflow_revision_digest
            || revision.spec.spec_id != item.spec_id
            || revision.kernel_abi_version != item.kernel_abi_version
            || revision.capabilities != item.capability_pins
        {
            return Err("published revision does not match the claimed pins".into());
        }
        if !matches!(revision.kernel_abi_version.as_str(), "1" | "kernel-abi-1") {
            return Err("unsupported workflow kernel ABI".into());
        }
        let nodes = self
            .nodes
            .as_ref()
            .ok_or("dynamic spec execution is unavailable")?;
        let builtins = NodeRegistry::with_builtins();
        for (expression, version) in &revision.expression_versions {
            let supported = nodes.capability_manifests().any(|manifest| {
                manifest.id == *expression && manifest.contract_version == *version
            }) || (builtins.is_step(expression) && version == "1");
            if !supported {
                return Err("unsupported expression version".into());
            }
        }
        for pin in &revision.capabilities {
            let manifest = nodes
                .capability_by_pin(pin)
                .ok_or("pinned capability is unavailable")?;
            debug_assert_eq!(manifest.id, pin.id);
        }
        let selected_nodes = nodes
            .for_capability_pins(&revision.capabilities)
            .map_err(|error| error.to_string())?;
        for node in revision
            .spec
            .branches
            .iter()
            .flat_map(|branch| &branch.nodes)
        {
            if let Some(manifest) = selected_nodes.capability(&node.node_type) {
                if !revision.capabilities.iter().any(|pin| {
                    pin.id == manifest.id
                        && pin.contract_version == manifest.contract_version
                        && pin.content_digest == manifest.content_digest
                }) {
                    return Err("graph capability lacks an exact revision pin".into());
                }
            } else if nodes
                .capability_manifests()
                .any(|manifest| manifest.id == node.node_type)
            {
                return Err("graph capability lacks an exact revision pin".into());
            } else if !builtins.is_step(&node.node_type) && !builtins.is_ingress(&node.node_type) {
                return Err("product graph node requires a capability manifest".into());
            }
        }
        let content = serde_json::json!({
            "spec": revision.spec, "abi": revision.kernel_abi_version,
            "capabilities": revision.capabilities,
            "expressions": revision.expression_versions,
            "dependencies": revision.dependency_set_digest,
        });
        let key = (item.spec_id.clone(), item.workflow_revision_digest.clone());
        let scope = Some((
            item.tenant_id.clone(),
            item.definition_id
                .parse()
                .map_err(|_| "invalid workflow definition identity")?,
        ));
        let mut compiled = self.compiled.lock().await;
        if let Some(entry) = compiled
            .iter()
            .find(|entry| entry.key == key && entry.scope == scope)
        {
            if entry.content != content {
                return Err("revision digest refers to different compiled content".into());
            }
            return Ok(entry.driver.clone());
        }
        if let Some(index) = compiled.iter().position(|entry| {
            entry.key.0 == item.spec_id
                && entry.key.1.is_empty()
                && entry.content == content["spec"]
        }) {
            let mut entry = compiled.remove(index).expect("located cache entry exists");
            entry.scope = scope;
            entry.key = key;
            entry.content = content;
            let driver = entry.driver.clone();
            compiled.push_back(entry);
            return Ok(driver);
        }
        let nodes = selected_nodes;
        // ponytail: cache misses compile serially per registry; parallelize only if
        // measured publish bursts require it. Compilation never blocks the executor.
        let driver: Arc<dyn WorkflowDriver<Context>> = Arc::new(
            tokio::task::spawn_blocking(move || {
                SpecDriver::new(&revision.spec, &nodes).map_err(|error| error.to_string())
            })
            .await
            .map_err(|_| "spec compiler panicked")??,
        );
        if compiled.len() == COMPILED_REVISION_LIMIT {
            compiled.pop_front();
        }
        compiled.push_back(CompiledRevision {
            scope,
            key,
            content,
            driver: driver.clone(),
        });
        Ok(driver)
    }

    async fn resolve<Queue: WorkQueue>(
        &self,
        queue: &Queue,
        item: &WorkItem,
    ) -> Result<Arc<dyn WorkflowDriver<Context>>, String>
    where
        Context: 'static,
    {
        if let Some(driver) = self.by_spec.get(&item.spec_id) {
            if !driver.supports_revision(item) {
                return Err("custom driver does not support the claimed revision".into());
            }
            return Ok(driver.clone());
        }
        let revision = queue
            .load_revision(item)
            .await
            .map_err(|error| error.to_string())?;
        self.resolve_revision(item, revision).await
    }

    /// Register a driver. Duplicate spec ownership is a startup error because
    /// dispatch must never depend on registration order.
    pub fn register(
        &mut self,
        driver: Arc<dyn WorkflowDriver<Context>>,
    ) -> Result<&mut Self, SupervisorError> {
        for spec_id in driver.spec_ids() {
            if let Some(existing) = self.by_spec.get(spec_id) {
                return Err(SupervisorError::Driver {
                    driver: driver.name().to_owned(),
                    reason: format!(
                        "spec '{spec_id}' is already claimed by '{}'",
                        existing.name()
                    ),
                });
            }
            self.by_spec.insert(spec_id.to_owned(), driver.clone());
        }
        self.drivers.push(driver);
        Ok(self)
    }

    /// Every owned spec id, sorted.
    pub fn spec_ids(&self) -> Vec<String> {
        let mut ids = self.by_spec.keys().cloned().collect::<Vec<_>>();
        ids.extend(self.startup_specs.iter().cloned());
        ids.sort();
        ids.dedup();
        ids
    }

    /// Driver owning `spec_id`.
    pub fn for_spec(&self, spec_id: &str) -> Option<&Arc<dyn WorkflowDriver<Context>>> {
        self.by_spec.get(spec_id)
    }

    /// Registered driver names.
    pub fn names(&self) -> Vec<&'static str> {
        self.drivers.iter().map(|driver| driver.name()).collect()
    }

    /// Whether no driver is registered or prewarmed; published graphs may still resolve lazily.
    pub fn is_empty(&self) -> bool {
        self.drivers.is_empty() && self.startup_specs.is_empty()
    }

    /// Validate every driver's spec set.
    pub fn validate_all(&self) -> Result<(), SupervisorError> {
        for driver in &self.drivers {
            driver
                .validate_specs()
                .map_err(|reason| SupervisorError::Driver {
                    driver: driver.name().to_owned(),
                    reason,
                })?;
        }
        Ok(())
    }
}

/// One supervisor pass: claim due work, evaluate concurrently under renewed leases, commit every command.
/// A failing context factory requeues every claimed item with the failure.
pub async fn run_due_pass<Context, Queue, BuildContext, BuildFuture>(
    queue: &Queue,
    registry: &DriverRegistry<Context>,
    settings: &SupervisorSettings,
    build_context: BuildContext,
) -> Result<SupervisorStats, SupervisorError>
where
    Context: Send + Sync + 'static,
    Queue: WorkQueue,
    BuildContext: FnOnce() -> BuildFuture,
    BuildFuture: Future<Output = Result<Context, String>> + Send,
{
    if registry.is_empty() && registry.nodes.is_none() {
        return Ok(SupervisorStats::default());
    }
    let claim_limit = settings
        .claim_batch
        .clamp(1, i64::try_from(settings.concurrency()).unwrap_or(i64::MAX));
    let items = queue
        .claim_due(
            &if registry.nodes.is_some() {
                Vec::new()
            } else {
                registry.spec_ids()
            },
            &settings.worker_id,
            settings.lease_secs,
            claim_limit,
        )
        .await?;
    if items.is_empty() {
        return Ok(SupervisorStats::default());
    }

    let context = match build_context().await {
        Ok(context) => Arc::new(context),
        Err(error) => {
            let mut cleanup_errors = Vec::new();
            for item in &items {
                let command = WorkflowTransitionCommand::failure(
                    item,
                    error.clone(),
                    Utc::now() + chrono::Duration::seconds(settings.requeue_delay_secs),
                );
                if let Err(cleanup) = queue.commit_command(item, &command).await {
                    cleanup_errors.push(format!("{} commit failure: {cleanup}", item.id));
                }
            }
            if !cleanup_errors.is_empty() {
                return Err(SupervisorError::Context(format!(
                    "{error}; claimed work cleanup failed: {}",
                    cleanup_errors.join(", ")
                )));
            }
            return Err(SupervisorError::Context(error));
        }
    };

    let claimed = items.len();
    let mut failed = 0;
    let mut persistence_errors = Vec::new();
    let mut work = items.into_iter();
    let mut tasks = FuturesUnordered::new();

    let spawn_next = |tasks: &mut FuturesUnordered<_>, work: &mut std::vec::IntoIter<WorkItem>| {
        let Some(item) = work.next() else {
            return false;
        };
        let context = context.clone();
        let renewal_period = std::time::Duration::from_millis(
            (settings.lease_secs.clamp(1, 3600) as u64 * 1000 / 3).max(1),
        );
        tasks.push(async move {
            let evaluation_item = item.clone();
            let evaluation = async {
                evaluation_item.validate_pins()?;
                if !evaluation_item.is_converging() {
                    queue
                        .admit_evaluation(&evaluation_item)
                        .await
                        .map_err(|error| error.to_string())?;
                }
                let driver = registry.resolve(queue, &evaluation_item).await?;
                AssertUnwindSafe(driver.evaluate(&context, &evaluation_item))
                    .catch_unwind()
                    .await
                    .map_err(|_| "driver panicked".to_string())
                    .and_then(|result| result)
            };
            tokio::pin!(evaluation);
            let mut renewal = tokio::time::interval_at(
                tokio::time::Instant::now() + renewal_period,
                renewal_period,
            );
            renewal.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
            let result = loop {
                tokio::select! {
                    result = &mut evaluation => break result,
                    _ = renewal.tick() => {
                        if let Err(error) = queue
                            .renew(
                                &item.tenant_id,
                                &item.id,
                                &settings.worker_id,
                                item.lease_version,
                                settings.lease_secs,
                            )
                            .await
                        {
                            break Err(format!("lease renewal failed: {error}"));
                        }
                    }
                }
            };
            (item, result)
        });
        true
    };

    for _ in 0..settings.concurrency() {
        if !spawn_next(&mut tasks, &mut work) {
            break;
        }
    }

    while let Some((item, result)) = tasks.next().await {
        let command = match result {
            Ok(command) => match command.validate_for(&item) {
                Ok(()) => command,
                Err(error) => {
                    failed += 1;
                    WorkflowTransitionCommand::failure(
                        &item,
                        error,
                        Utc::now() + chrono::Duration::seconds(settings.requeue_delay_secs),
                    )
                }
            },
            Err(error) => {
                failed += 1;
                WorkflowTransitionCommand::failure(
                    &item,
                    error,
                    Utc::now() + chrono::Duration::seconds(settings.requeue_delay_secs),
                )
            }
        };
        if let Err(error) = queue.commit_command(&item, &command).await {
            failed += usize::from(!matches!(
                command.disposition,
                WorkDisposition::Failed { .. }
            ));
            persistence_errors.push(format!("{} command commit: {error}", item.id));
        }
        spawn_next(&mut tasks, &mut work);
    }

    if persistence_errors.is_empty() {
        Ok(SupervisorStats { claimed, failed })
    } else {
        Err(SupervisorError::Commit(persistence_errors.join(", ")))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::{AtomicBool, Ordering};
    use std::sync::Mutex;

    struct Context;
    struct Driver {
        valid: bool,
    }

    #[async_trait]
    impl WorkflowDriver<Context> for Driver {
        fn name(&self) -> &'static str {
            "driver"
        }
        fn spec_ids(&self) -> Vec<&str> {
            vec!["spec"]
        }
        fn validate_specs(&self) -> Result<(), String> {
            self.valid.then_some(()).ok_or_else(|| "invalid".into())
        }
        fn supports_revision(&self, item: &WorkItem) -> bool {
            item.workflow_revision_digest == "workflow-digest" && item.kernel_abi_version == "1"
        }
        async fn evaluate(
            &self,
            _context: &Context,
            item: &WorkItem,
        ) -> Result<WorkflowTransitionCommand, String> {
            if let Some(milliseconds) = item.config["sleep_ms"].as_u64() {
                tokio::time::sleep(std::time::Duration::from_millis(milliseconds)).await;
            }
            if item.config["fail"] == true {
                Err("planned".into())
            } else if item.config["reschedule"] == true {
                Ok(command(
                    item,
                    WorkDisposition::Reschedule { at: Utc::now() },
                    serde_json::json!({"next": true}),
                ))
            } else if item.config["stop"] == true {
                Ok(command(
                    item,
                    WorkDisposition::Complete,
                    item.config.clone(),
                ))
            } else {
                Ok(command(
                    item,
                    WorkDisposition::Continue { delay_secs: 5 },
                    item.config.clone(),
                ))
            }
        }
    }

    fn command(
        item: &WorkItem,
        disposition: WorkDisposition,
        next_state: Value,
    ) -> WorkflowTransitionCommand {
        WorkflowTransitionCommand {
            consumed_wakeups: item
                .wakeups
                .iter()
                .map(|wakeup| wakeup.id.clone())
                .collect(),
            delivery_key: format!("test:{}:{}", item.id, item.state_version),
            delivery_digest: "delivery-digest".into(),
            event_type: "workflow.test".into(),
            event_digest: "event-digest".into(),
            event_payload: Value::Null,
            next_state,
            action_intents: Vec::new(),
            outcome: EvaluationOutcome::default(),
            disposition,
        }
    }

    #[derive(Default)]
    struct Queue {
        items: Mutex<Vec<WorkItem>>,
        committed: Mutex<Vec<(String, WorkDisposition)>>,
        renewals: Mutex<Vec<(String, String, i64)>>,
        renewal_error: Mutex<Option<String>>,
        finalization_error: Mutex<Option<String>>,
        claim_limits: Mutex<Vec<i64>>,
    }

    #[async_trait]
    impl WorkQueue for Queue {
        async fn admit_evaluation(&self, _: &WorkItem) -> Result<(), SupervisorError> {
            Ok(())
        }
        async fn claim_due(
            &self,
            _spec_ids: &[String],
            _worker_id: &str,
            _lease_secs: i64,
            batch: i64,
        ) -> Result<Vec<WorkItem>, SupervisorError> {
            self.claim_limits.lock().unwrap().push(batch);
            let mut items = self.items.lock().unwrap();
            let take = items.len().min(batch as usize);
            Ok(items.drain(..take).collect())
        }

        async fn renew(
            &self,
            _tenant_id: &str,
            id: &str,
            worker_id: &str,
            lease_version: i64,
            _lease_secs: i64,
        ) -> Result<(), SupervisorError> {
            self.renewals.lock().unwrap().push((
                id.to_string(),
                worker_id.to_string(),
                lease_version,
            ));
            match self.renewal_error.lock().unwrap().clone() {
                Some(error) => Err(SupervisorError::Queue(error)),
                None => Ok(()),
            }
        }

        async fn commit_command(
            &self,
            item: &WorkItem,
            command: &WorkflowTransitionCommand,
        ) -> Result<(), SupervisorError> {
            self.committed
                .lock()
                .unwrap()
                .push((item.id.clone(), command.disposition.clone()));
            match self.finalization_error.lock().unwrap().clone() {
                Some(error) => Err(SupervisorError::Queue(error)),
                None => Ok(()),
            }
        }
    }

    fn item(id: &str, fail: bool) -> WorkItem {
        WorkItem {
            id: id.into(),
            run_id: uuid::Uuid::new_v4().to_string().parse().unwrap(),
            tenant_id: "tenant".parse().unwrap(),
            subject_id: "subject".parse().unwrap(),
            spec_id: "spec".into(),
            definition_id: "definition".into(),
            workflow_revision: 1,
            workflow_revision_digest: "workflow-digest".into(),
            execution_profile_id: "profile".into(),
            execution_profile_revision: 1,
            execution_profile_digest: "profile-digest".into(),
            kernel_abi_version: "1".into(),
            capability_pins: Vec::new(),
            lifecycle: LifecyclePolicy::run_once(),
            scheduled_at: Utc::now(),
            created_at: Utc::now(),
            claimed_at: Utc::now(),
            config: serde_json::json!({"fail": fail}),
            state_version: 0,
            control_epochs: ControlEpochs::default(),
            cancel_requested: false,
            lease_version: 1,
            wakeups: Vec::new(),
        }
    }

    fn stopped_item(id: &str) -> WorkItem {
        WorkItem {
            config: serde_json::json!({"stop": true}),
            ..item(id, false)
        }
    }

    fn settings() -> SupervisorSettings {
        SupervisorSettings {
            worker_id: "worker".into(),
            lease_secs: 60,
            requeue_delay_secs: 5,
            claim_batch: 10,
            concurrency: 3,
        }
    }

    struct MustNotEvaluate(Arc<AtomicBool>);

    #[async_trait]
    impl WorkflowDriver<Context> for MustNotEvaluate {
        fn name(&self) -> &'static str {
            "must-not-evaluate"
        }

        fn spec_ids(&self) -> Vec<&str> {
            vec!["spec"]
        }

        fn validate_specs(&self) -> Result<(), String> {
            Ok(())
        }

        async fn evaluate(
            &self,
            _context: &Context,
            item: &WorkItem,
        ) -> Result<WorkflowTransitionCommand, String> {
            self.0.store(true, Ordering::SeqCst);
            Ok(command(
                item,
                WorkDisposition::Complete,
                item.config.clone(),
            ))
        }
    }

    #[tokio::test]
    async fn missing_revision_pin_fails_before_driver_evaluation() {
        let mut stale = item("stale", false);
        stale.workflow_revision_digest.clear();
        let queue = Queue {
            items: Mutex::new(vec![stale]),
            ..Default::default()
        };
        let called = Arc::new(AtomicBool::new(false));
        let mut registry = DriverRegistry::new();
        registry
            .register(Arc::new(MustNotEvaluate(called.clone())))
            .unwrap();

        let stats = run_due_pass(&queue, &registry, &settings(), || async { Ok(Context) })
            .await
            .unwrap();

        assert_eq!(stats.failed, 1);
        assert!(!called.load(Ordering::SeqCst));
        assert!(matches!(
            queue.committed.lock().unwrap().as_slice(),
            [(id, WorkDisposition::Failed { error, .. })]
                if id == "stale" && error.contains("workflow_revision_digest")
        ));
    }

    #[tokio::test]
    async fn due_pass_releases_success_and_marks_failures() {
        let queue = Queue {
            items: Mutex::new(vec![
                item("ok", false),
                stopped_item("done"),
                item("bad", true),
            ]),
            ..Default::default()
        };
        let mut registry = DriverRegistry::new();
        registry.register(Arc::new(Driver { valid: true })).unwrap();
        let stats = run_due_pass(&queue, &registry, &settings(), || async { Ok(Context) })
            .await
            .unwrap();
        assert_eq!(
            stats,
            SupervisorStats {
                claimed: 3,
                failed: 1
            }
        );
        let committed = queue.committed.lock().unwrap();
        assert_eq!(committed.len(), 3);
        assert!(committed
            .iter()
            .any(|(id, disposition)| id == "done" && disposition == &WorkDisposition::Complete));
        assert!(committed.iter().any(|(id, disposition)| id == "bad"
            && matches!(disposition, WorkDisposition::Failed { .. })));
    }

    #[tokio::test]
    async fn context_failure_releases_every_claim() {
        let queue = Queue {
            items: Mutex::new(vec![item("one", false), item("two", false)]),
            ..Default::default()
        };
        let mut registry = DriverRegistry::new();
        registry.register(Arc::new(Driver { valid: true })).unwrap();
        let error = run_due_pass(&queue, &registry, &settings(), || async {
            Err::<Context, _>("context failed".into())
        })
        .await
        .unwrap_err();
        assert_eq!(error, SupervisorError::Context("context failed".into()));
        assert_eq!(queue.committed.lock().unwrap().len(), 2);
        assert!(queue
            .committed
            .lock()
            .unwrap()
            .iter()
            .all(|(_, disposition)| matches!(disposition, WorkDisposition::Failed { .. })));
    }

    #[tokio::test]
    async fn due_pass_atomically_reschedules_driver_state() {
        let queue = Queue {
            items: Mutex::new(vec![WorkItem {
                config: serde_json::json!({"reschedule": true}),
                ..item("recurring", false)
            }]),
            ..Default::default()
        };
        let mut registry = DriverRegistry::new();
        registry.register(Arc::new(Driver { valid: true })).unwrap();
        run_due_pass(&queue, &registry, &settings(), || async { Ok(Context) })
            .await
            .unwrap();
        assert!(matches!(
            queue.committed.lock().unwrap().as_slice(),
            [(id, WorkDisposition::Reschedule { .. })] if id == "recurring"
        ));
    }

    #[tokio::test]
    async fn finalization_failure_is_not_reported_as_success() {
        let queue = Queue {
            items: Mutex::new(vec![item("stale", false)]),
            finalization_error: Mutex::new(Some("stale lease".into())),
            ..Default::default()
        };
        let mut registry = DriverRegistry::new();
        registry.register(Arc::new(Driver { valid: true })).unwrap();

        let error = run_due_pass(&queue, &registry, &settings(), || async { Ok(Context) })
            .await
            .unwrap_err();

        assert!(
            matches!(&error, SupervisorError::Commit(reason) if reason.contains("stale lease"))
        );
        assert!(error.to_string().contains("state commit failed"));
    }

    #[tokio::test]
    async fn due_pass_claims_only_work_that_can_start() {
        let queue = Queue {
            items: Mutex::new(vec![
                item("one", false),
                item("two", false),
                item("queued", false),
            ]),
            ..Default::default()
        };
        let mut registry = DriverRegistry::new();
        registry.register(Arc::new(Driver { valid: true })).unwrap();
        let mut settings = settings();
        settings.concurrency = 2;

        let stats = run_due_pass(&queue, &registry, &settings, || async { Ok(Context) })
            .await
            .unwrap();

        assert_eq!(stats.claimed, 2);
        assert_eq!(queue.claim_limits.lock().unwrap().as_slice(), [2]);
        assert_eq!(queue.items.lock().unwrap().len(), 1);
    }

    #[tokio::test(start_paused = true)]
    async fn long_evaluate_renews_its_lease() {
        let queue = Queue {
            items: Mutex::new(vec![WorkItem {
                config: serde_json::json!({"sleep_ms": 3500}),
                ..item("slow", false)
            }]),
            ..Default::default()
        };
        let mut registry = DriverRegistry::new();
        registry.register(Arc::new(Driver { valid: true })).unwrap();
        let mut settings = settings();
        settings.lease_secs = 3;
        settings.concurrency = 1;

        let stats = run_due_pass(&queue, &registry, &settings, || async { Ok(Context) })
            .await
            .unwrap();

        assert_eq!(stats.failed, 0);
        let renewals = queue.renewals.lock().unwrap();
        assert_eq!(renewals.len(), 3);
        assert!(renewals
            .iter()
            .all(|renewal| renewal == &("slow".into(), "worker".into(), 1)));
        assert!(matches!(
            queue.committed.lock().unwrap().as_slice(),
            [(id, WorkDisposition::Continue { .. })] if id == "slow"
        ));
    }

    #[tokio::test(start_paused = true)]
    async fn expired_lease_stops_evaluation() {
        let queue = Queue {
            items: Mutex::new(vec![WorkItem {
                config: serde_json::json!({"sleep_ms": 5000}),
                ..item("expired", false)
            }]),
            renewal_error: Mutex::new(Some("lease expired".into())),
            ..Default::default()
        };
        let mut registry = DriverRegistry::new();
        registry.register(Arc::new(Driver { valid: true })).unwrap();
        let mut settings = settings();
        settings.lease_secs = 3;
        settings.concurrency = 1;

        let stats = run_due_pass(&queue, &registry, &settings, || async { Ok(Context) })
            .await
            .unwrap();

        assert_eq!(stats.failed, 1);
        assert_eq!(queue.renewals.lock().unwrap().len(), 1);
        assert!(matches!(
            queue.committed.lock().unwrap().as_slice(),
            [(id, WorkDisposition::Failed { .. })] if id == "expired"
        ));
    }

    #[test]
    fn registry_fails_duplicate_ownership_and_invalid_specs() {
        let mut registry = DriverRegistry::new();
        assert!(registry.is_empty());
        registry
            .register(Arc::new(Driver { valid: false }))
            .unwrap();
        assert_eq!(registry.spec_ids(), ["spec"]);
        assert_eq!(registry.names(), ["driver"]);
        assert!(registry.for_spec("spec").is_some());
        assert_eq!(
            registry.validate_all().unwrap_err(),
            SupervisorError::Driver {
                driver: "driver".into(),
                reason: "invalid".into()
            }
        );
        assert!(matches!(
            registry.register(Arc::new(Driver { valid: true })),
            Err(SupervisorError::Driver { reason, .. }) if reason.contains("already claimed")
        ));
    }

    #[tokio::test]
    async fn memory_work_queue_serves_consumers_without_a_database() {
        let queue = MemoryWorkQueue::new(vec![item("first", false)]);
        queue.push(item("other-spec", false));
        {
            let mut items = queue.items.lock().unwrap();
            items[1].spec_id = "unknown".into();
        }
        let mut registry = DriverRegistry::new();
        registry.register(Arc::new(Driver { valid: true })).unwrap();
        let stats = run_due_pass(&queue, &registry, &settings(), || async { Ok(Context) })
            .await
            .unwrap();
        assert_eq!(
            stats,
            SupervisorStats {
                claimed: 1,
                failed: 0
            },
            "only items for registered specs are claimed"
        );
        let committed = queue.committed();
        assert_eq!(committed.len(), 1);
        assert_eq!(committed[0].0.id, "first");
        assert_eq!(
            committed[0].1.disposition,
            WorkDisposition::Continue { delay_secs: 5 }
        );
        assert!(queue.renewals().is_empty());
        assert_eq!(
            run_due_pass(&queue, &registry, &settings(), || async { Ok(Context) })
                .await
                .unwrap()
                .claimed,
            0,
            "claimed items leave the queue"
        );
    }

    #[test]
    fn commands_are_validated_against_the_claimed_work() {
        let work = item("scoped", false);
        let mut blank = command(&work, WorkDisposition::Complete, Value::Null);
        blank.event_type.clear();
        assert!(blank
            .validate_for(&work)
            .unwrap_err()
            .contains("missing event_type"));

        let mut cancelled = item("cancelled", false);
        cancelled.cancel_requested = true;
        let mut with_intent = command(&cancelled, WorkDisposition::Complete, Value::Null);
        with_intent.action_intents.push(ActionIntent {
            id: "intent".into(),
            tenant_id: cancelled.tenant_id.clone(),
            instance_id: cancelled.id.clone().parse().unwrap(),
            run_id: cancelled.run_id.clone(),
            capability: CapabilityPin {
                id: "cap".into(),
                contract_version: "1".into(),
                content_digest: "digest".into(),
            },
            idempotency_key: "key".into(),
            state: crate::ActionState::Prepared,
            input: Value::Null,
            effect: crate::Effect::ExternalWrite,
            retry_class: crate::IdempotencyMode::Native,
            control_epochs: ControlEpochs::default(),
            resource_scope_id: String::new(),
            lease_epoch: cancelled.lease_version,
            action_epoch: cancelled.state_version,
            deadline: None,
            reservation: None,
            created_at: Utc::now(),
        });
        assert!(with_intent
            .validate_for(&cancelled)
            .unwrap_err()
            .contains("cancelled work"));

        let mut escaped = with_intent.clone();
        escaped.action_intents[0].tenant_id = "someone-else".parse().unwrap();
        let scoped = item("scoped", false);
        let mut escaped_command = command(&scoped, WorkDisposition::Complete, Value::Null);
        escaped_command.action_intents = escaped.action_intents;
        assert!(escaped_command
            .validate_for(&scoped)
            .unwrap_err()
            .contains("escapes claimed work scope"));
    }
}