use af_context::{RunId, SubjectId, TenantId};
use std::collections::HashMap;
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};
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct Wakeup {
pub id: String,
pub kind: String,
pub payload: Value,
}
#[derive(Debug, Clone, PartialEq)]
pub struct WorkItem {
pub id: String,
pub run_id: RunId,
pub tenant_id: TenantId,
pub subject_id: SubjectId,
pub spec_id: String,
pub definition_id: String,
pub workflow_revision: u64,
pub workflow_revision_digest: String,
pub execution_profile_id: String,
pub execution_profile_revision: u64,
pub execution_profile_digest: String,
pub kernel_abi_version: String,
pub capability_pins: Vec<CapabilityPin>,
pub lifecycle: LifecyclePolicy,
pub scheduled_at: DateTime<Utc>,
pub claimed_at: DateTime<Utc>,
pub config: Value,
pub state_version: i64,
pub control_epochs: ControlEpochs,
pub cancel_requested: bool,
pub lease_version: i64,
#[doc = "Due timers and trigger deliveries, oldest first."]
pub wakeups: Vec<Wakeup>,
}
impl WorkItem {
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(())
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum WorkDisposition {
Continue {
delay_secs: i64,
},
Reschedule {
at: DateTime<Utc>,
},
Complete,
Failed {
error: String,
retry_at: Option<DateTime<Utc>>,
},
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct EvaluationOutcome {
pub triggered: bool,
pub matched: bool,
pub succeeded: bool,
pub action_terminal: bool,
}
#[derive(Debug, Clone, PartialEq)]
pub struct WorkflowTransitionCommand {
pub delivery_key: String,
pub delivery_digest: String,
pub event_type: String,
pub event_digest: String,
pub event_payload: Value,
pub next_state: Value,
pub action_intents: Vec<ActionIntent>,
pub outcome: EvaluationOutcome,
pub disposition: WorkDisposition,
}
impl WorkflowTransitionCommand {
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(),
disposition: WorkDisposition::Failed {
error,
retry_at: Some(retry_at),
},
}
}
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());
}
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());
}
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(())
}
}
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
pub struct SupervisorStats {
pub claimed: usize,
pub failed: usize,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SupervisorSettings {
pub worker_id: String,
pub lease_secs: i64,
pub requeue_delay_secs: i64,
pub claim_batch: i64,
pub concurrency: usize,
}
impl SupervisorSettings {
fn concurrency(&self) -> usize {
self.concurrency.max(1)
}
}
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum SupervisorError {
#[error("work queue: {0}")]
Queue(String),
#[error("driver '{driver}': {reason}")]
Driver {
driver: String,
reason: String,
},
#[error("context: {0}")]
Context(String),
#[error("workflow state commit failed: {0}")]
Commit(String),
}
#[async_trait]
pub trait WorkQueue: Send + Sync {
async fn claim_due(
&self,
spec_ids: &[String],
worker_id: &str,
lease_secs: i64,
batch: i64,
) -> Result<Vec<WorkItem>, SupervisorError>;
async fn renew(
&self,
tenant_id: &str,
id: &str,
worker_id: &str,
lease_version: i64,
lease_secs: i64,
) -> Result<(), SupervisorError>;
async fn commit_command(
&self,
item: &WorkItem,
command: &WorkflowTransitionCommand,
) -> Result<(), SupervisorError>;
}
#[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 {
pub fn new(items: Vec<WorkItem>) -> Self {
Self {
items: std::sync::Mutex::new(items),
..Self::default()
}
}
pub fn push(&self, item: WorkItem) {
self.items
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.push(item);
}
pub fn committed(&self) -> Vec<(WorkItem, WorkflowTransitionCommand)> {
self.committed
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.clone()
}
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 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.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(())
}
}
#[async_trait]
pub trait WorkflowDriver<Context>: Send + Sync
where
Context: Send + Sync,
{
fn name(&self) -> &'static str;
fn spec_ids(&self) -> Vec<&str>;
fn validate_specs(&self) -> Result<(), String>;
async fn evaluate(
&self,
context: &Context,
item: &WorkItem,
) -> Result<WorkflowTransitionCommand, String>;
}
pub struct DriverRegistry<Context: Send + Sync> {
drivers: Vec<Arc<dyn WorkflowDriver<Context>>>,
by_spec: HashMap<String, Arc<dyn WorkflowDriver<Context>>>,
}
impl<Context: Send + Sync> Default for DriverRegistry<Context> {
fn default() -> Self {
Self {
drivers: Vec::new(),
by_spec: HashMap::new(),
}
}
}
impl<Context: Send + Sync> DriverRegistry<Context> {
pub fn new() -> Self {
Self::default()
}
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)
}
pub fn spec_ids(&self) -> Vec<String> {
let mut ids = self.by_spec.keys().cloned().collect::<Vec<_>>();
ids.sort();
ids
}
pub fn for_spec(&self, spec_id: &str) -> Option<&Arc<dyn WorkflowDriver<Context>>> {
self.by_spec.get(spec_id)
}
pub fn names(&self) -> Vec<&'static str> {
self.drivers.iter().map(|driver| driver.name()).collect()
}
pub fn is_empty(&self) -> bool {
self.drivers.is_empty()
}
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(())
}
}
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,
{
let claim_limit = settings
.claim_batch
.clamp(1, i64::try_from(settings.concurrency()).unwrap_or(i64::MAX));
let items = queue
.claim_due(
®istry.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 driver = registry.for_spec(&item.spec_id).cloned();
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()?;
match driver {
Some(driver) => AssertUnwindSafe(driver.evaluate(&context, &evaluation_item))
.catch_unwind()
.await
.map_err(|_| "driver panicked".to_string())
.and_then(|result| result),
None => Err(format!(
"no driver registered for spec '{}'",
evaluation_item.spec_id
)),
}
};
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())
}
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 {
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 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(),
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, ®istry, &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, ®istry, &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, ®istry, &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, ®istry, &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, ®istry, &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, ®istry, &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, ®istry, &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, ®istry, &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, ®istry, &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, ®istry, &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"));
}
}