1use af_context::{RunId, SubjectId, TenantId};
4use std::collections::{HashMap, VecDeque};
5use std::future::Future;
6use std::panic::AssertUnwindSafe;
7use std::sync::Arc;
8
9use async_trait::async_trait;
10use chrono::{DateTime, Utc};
11use futures::{stream::FuturesUnordered, FutureExt, StreamExt};
12use serde_json::Value;
13
14use crate::{
15 ActionIntent, CapabilityPin, ControlEpochs, LifecyclePolicy, NodeRegistry, SpecDriver,
16 WorkflowRevision,
17};
18
19#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
23pub struct Wakeup {
24 pub id: String,
26 pub kind: String,
28 #[serde(default)]
30 pub branch_id: Option<af_context::BranchId>,
31 pub payload: Value,
33}
34
35#[derive(Debug, Clone, PartialEq)]
37pub struct WorkItem {
38 pub id: String,
40 pub run_id: RunId,
42 pub tenant_id: TenantId,
44 pub subject_id: SubjectId,
46 pub spec_id: String,
48 pub definition_id: String,
50 pub workflow_revision: u64,
52 pub workflow_revision_digest: String,
54 pub execution_profile_id: String,
56 pub execution_profile_revision: u64,
58 pub execution_profile_digest: String,
60 pub kernel_abi_version: String,
62 pub capability_pins: Vec<CapabilityPin>,
64 pub lifecycle: LifecyclePolicy,
66 pub scheduled_at: DateTime<Utc>,
69 pub created_at: DateTime<Utc>,
71 pub claimed_at: DateTime<Utc>,
73 pub config: Value,
75 pub state_version: i64,
77 pub control_epochs: ControlEpochs,
79 pub cancel_requested: bool,
81 pub lease_version: i64,
83 #[doc = "Due timers and trigger deliveries, oldest first."]
84 pub wakeups: Vec<Wakeup>,
85}
86
87impl WorkItem {
88 pub fn is_converging(&self) -> bool {
90 self.cancel_requested
91 || self
92 .wakeups
93 .iter()
94 .any(|wakeup| wakeup.kind == "timer" && wakeup.payload["kind"] == "terminal_action")
95 }
96
97 fn validate_pins(&self) -> Result<(), String> {
98 for (name, value) in [
99 ("definition_id", self.definition_id.as_str()),
100 (
101 "workflow_revision_digest",
102 self.workflow_revision_digest.as_str(),
103 ),
104 ("execution_profile_id", self.execution_profile_id.as_str()),
105 (
106 "execution_profile_digest",
107 self.execution_profile_digest.as_str(),
108 ),
109 ("kernel_abi_version", self.kernel_abi_version.as_str()),
110 ] {
111 if value.trim().is_empty() {
112 return Err(format!("work item is missing pinned {name}"));
113 }
114 }
115 if self.workflow_revision == 0 || self.execution_profile_revision == 0 {
116 return Err("work item revision pins must be positive".into());
117 }
118 uuid::Uuid::parse_str(&self.run_id)
119 .map_err(|error| format!("work item has invalid run_id: {error}"))?;
120 if self.state_version < 0
121 || self.lease_version <= 0
122 || self.control_epochs.tenant < 0
123 || self.control_epochs.instance < 0
124 {
125 return Err("work item state, lease and control pins must be current".into());
126 }
127 Ok(())
128 }
129}
130
131#[derive(Debug, Clone, PartialEq, Eq)]
133pub enum WorkDisposition {
134 Continue {
136 delay_secs: i64,
138 },
139 Reschedule {
141 at: DateTime<Utc>,
143 },
144 Complete,
146 Failed {
148 error: String,
150 retry_at: Option<DateTime<Utc>>,
152 },
153}
154
155#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
157pub struct EvaluationOutcome {
158 pub triggered: bool,
160 pub matched: bool,
162 pub succeeded: bool,
164 pub action_terminal: bool,
166}
167
168#[derive(Debug, Clone, PartialEq)]
172pub struct WorkflowTransitionCommand {
173 pub delivery_key: String,
175 pub delivery_digest: String,
177 pub event_type: String,
179 pub event_digest: String,
181 pub event_payload: Value,
183 pub next_state: Value,
185 pub action_intents: Vec<ActionIntent>,
187 pub outcome: EvaluationOutcome,
189 pub consumed_wakeups: Vec<String>,
191 pub disposition: WorkDisposition,
193}
194
195impl WorkflowTransitionCommand {
196 pub fn failure(item: &WorkItem, error: impl Into<String>, retry_at: DateTime<Utc>) -> Self {
198 let error = error.into();
199 Self {
200 delivery_key: format!("supervisor:{}:{}", item.id, item.state_version),
201 delivery_digest: format!("failure:{}:{}", item.lease_version, error),
202 event_type: "workflow.evaluation_failed".into(),
203 event_digest: format!("failure:{}:{}", item.state_version, error),
204 event_payload: serde_json::json!({"error": error.clone()}),
205 next_state: item.config.clone(),
206 action_intents: Vec::new(),
207 outcome: EvaluationOutcome::default(),
208 consumed_wakeups: Vec::new(),
209 disposition: WorkDisposition::Failed {
210 error,
211 retry_at: Some(retry_at),
212 },
213 }
214 }
215
216 pub fn validate_for(&self, item: &WorkItem) -> Result<(), String> {
218 item.validate_pins()?;
219 for (name, value) in [
220 ("delivery_key", self.delivery_key.as_str()),
221 ("delivery_digest", self.delivery_digest.as_str()),
222 ("event_type", self.event_type.as_str()),
223 ("event_digest", self.event_digest.as_str()),
224 ] {
225 if value.trim().is_empty() {
226 return Err(format!("workflow command is missing {name}"));
227 }
228 }
229 if item.cancel_requested && !self.action_intents.is_empty() {
230 return Err("cancelled work cannot prepare external actions".into());
231 }
232 let consumed: std::collections::BTreeSet<_> = self.consumed_wakeups.iter().collect();
233 if consumed.len() != self.consumed_wakeups.len()
234 || consumed
235 .iter()
236 .any(|id| !item.wakeups.iter().any(|wakeup| &wakeup.id == *id))
237 || (!consumed.is_empty()
238 && matches!(self.disposition, WorkDisposition::Failed { .. })
239 && !self.outcome.triggered)
240 {
241 return Err("workflow command acknowledges unconsumed or unclaimed work".into());
242 }
243 for intent in &self.action_intents {
244 intent
245 .validate_prepared()
246 .map_err(|error| error.to_string())?;
247 if intent.tenant_id != item.tenant_id || intent.instance_id != item.id {
248 return Err("action intent escapes claimed work scope".into());
249 }
250 if intent.run_id != item.run_id {
251 return Err("action intent escapes claimed workflow run".into());
252 }
253 if intent.control_epochs.tenant != item.control_epochs.tenant
258 || intent.control_epochs.instance != item.control_epochs.instance
259 || intent.lease_epoch != item.lease_version
260 {
261 return Err("action intent uses stale control or lease pins".into());
262 }
263 if !item
264 .capability_pins
265 .iter()
266 .any(|pin| pin == &intent.capability)
267 {
268 return Err(format!(
269 "action capability '{}' is not pinned by the workflow revision",
270 intent.capability.id
271 ));
272 }
273 }
274 Ok(())
275 }
276}
277
278#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
280pub struct SupervisorStats {
281 pub claimed: usize,
283 pub failed: usize,
285}
286
287#[derive(Debug, Clone, PartialEq, Eq)]
289pub struct SupervisorSettings {
290 pub worker_id: String,
292 pub lease_secs: i64,
294 pub requeue_delay_secs: i64,
296 pub claim_batch: i64,
298 pub concurrency: usize,
300}
301
302impl SupervisorSettings {
303 fn concurrency(&self) -> usize {
304 self.concurrency.max(1)
305 }
306}
307
308#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
312pub enum SupervisorError {
313 #[error("work queue: {0}")]
315 Queue(String),
316 #[error("driver '{driver}': {reason}")]
318 Driver {
319 driver: String,
321 reason: String,
323 },
324 #[error("context: {0}")]
326 Context(String),
327 #[error("workflow state commit failed: {0}")]
329 Commit(String),
330}
331
332#[async_trait]
334pub trait WorkQueue: Send + Sync {
335 async fn claim_due(
337 &self,
338 spec_ids: &[String],
339 worker_id: &str,
340 lease_secs: i64,
341 batch: i64,
342 ) -> Result<Vec<WorkItem>, SupervisorError>;
343
344 async fn load_revision(&self, _item: &WorkItem) -> Result<WorkflowRevision, SupervisorError> {
347 Err(SupervisorError::Queue(
348 "published revision loading is unavailable".into(),
349 ))
350 }
351
352 async fn admit_evaluation(&self, _item: &WorkItem) -> Result<(), SupervisorError> {
356 Err(SupervisorError::Queue(
357 "workflow evaluation admission is unavailable".into(),
358 ))
359 }
360
361 async fn prepare_evaluation(
366 &self,
367 item: &WorkItem,
368 load_revision: bool,
369 ) -> Result<Option<WorkflowRevision>, SupervisorError> {
370 if !item.is_converging() {
371 self.admit_evaluation(item).await?;
372 }
373 if load_revision {
374 self.load_revision(item).await.map(Some)
375 } else {
376 Ok(None)
377 }
378 }
379
380 async fn renew(
382 &self,
383 tenant_id: &str,
384 id: &str,
385 worker_id: &str,
386 lease_version: i64,
387 lease_secs: i64,
388 ) -> Result<(), SupervisorError>;
389 async fn commit_command(
391 &self,
392 item: &WorkItem,
393 command: &WorkflowTransitionCommand,
394 ) -> Result<(), SupervisorError>;
395}
396
397#[derive(Default)]
400pub struct MemoryWorkQueue {
401 items: std::sync::Mutex<Vec<WorkItem>>,
402 committed: std::sync::Mutex<Vec<(WorkItem, WorkflowTransitionCommand)>>,
403 renewals: std::sync::Mutex<Vec<(String, i64)>>,
404}
405
406impl MemoryWorkQueue {
407 pub fn new(items: Vec<WorkItem>) -> Self {
409 Self {
410 items: std::sync::Mutex::new(items),
411 ..Self::default()
412 }
413 }
414
415 pub fn push(&self, item: WorkItem) {
417 self.items
418 .lock()
419 .unwrap_or_else(std::sync::PoisonError::into_inner)
420 .push(item);
421 }
422
423 pub fn committed(&self) -> Vec<(WorkItem, WorkflowTransitionCommand)> {
425 self.committed
426 .lock()
427 .unwrap_or_else(std::sync::PoisonError::into_inner)
428 .clone()
429 }
430
431 pub fn renewals(&self) -> Vec<(String, i64)> {
433 self.renewals
434 .lock()
435 .unwrap_or_else(std::sync::PoisonError::into_inner)
436 .clone()
437 }
438}
439
440#[async_trait]
441impl WorkQueue for MemoryWorkQueue {
442 async fn admit_evaluation(&self, _: &WorkItem) -> Result<(), SupervisorError> {
443 Ok(())
444 }
445 async fn claim_due(
446 &self,
447 spec_ids: &[String],
448 _worker_id: &str,
449 _lease_secs: i64,
450 batch: i64,
451 ) -> Result<Vec<WorkItem>, SupervisorError> {
452 let mut items = self
453 .items
454 .lock()
455 .unwrap_or_else(std::sync::PoisonError::into_inner);
456 let mut claimed = Vec::new();
457 let mut index = 0;
458 while index < items.len() && claimed.len() < batch.max(0) as usize {
459 if spec_ids.is_empty() || spec_ids.contains(&items[index].spec_id) {
460 claimed.push(items.remove(index));
461 } else {
462 index += 1;
463 }
464 }
465 Ok(claimed)
466 }
467
468 async fn renew(
469 &self,
470 _tenant_id: &str,
471 id: &str,
472 _worker_id: &str,
473 lease_version: i64,
474 _lease_secs: i64,
475 ) -> Result<(), SupervisorError> {
476 self.renewals
477 .lock()
478 .unwrap_or_else(std::sync::PoisonError::into_inner)
479 .push((id.to_owned(), lease_version));
480 Ok(())
481 }
482
483 async fn commit_command(
484 &self,
485 item: &WorkItem,
486 command: &WorkflowTransitionCommand,
487 ) -> Result<(), SupervisorError> {
488 self.committed
489 .lock()
490 .unwrap_or_else(std::sync::PoisonError::into_inner)
491 .push((item.clone(), command.clone()));
492 Ok(())
493 }
494}
495
496#[async_trait]
498pub trait WorkflowDriver<Context>: Send + Sync
499where
500 Context: Send + Sync,
501{
502 fn name(&self) -> &'static str;
504 fn spec_ids(&self) -> Vec<&str>;
506 fn validate_specs(&self) -> Result<(), String>;
509 fn supports_revision(&self, _item: &WorkItem) -> bool {
512 false
513 }
514 async fn evaluate(
518 &self,
519 context: &Context,
520 item: &WorkItem,
521 ) -> Result<WorkflowTransitionCommand, String>;
522}
523
524pub struct DriverRegistry<Context: Send + Sync> {
526 drivers: Vec<Arc<dyn WorkflowDriver<Context>>>,
527 by_spec: HashMap<String, Arc<dyn WorkflowDriver<Context>>>,
528 nodes: Option<Arc<NodeRegistry>>,
529 startup_specs: Vec<String>,
530 compiled: tokio::sync::Mutex<VecDeque<CompiledRevision<Context>>>,
531}
532
533struct CompiledRevision<Context: Send + Sync> {
534 scope: Option<(TenantId, af_context::WorkflowDefinitionId)>,
535 key: (String, String),
536 content: Value,
537 driver: Arc<dyn WorkflowDriver<Context>>,
538}
539
540const COMPILED_REVISION_LIMIT: usize = 256;
541
542impl<Context: Send + Sync> Default for DriverRegistry<Context> {
543 fn default() -> Self {
544 Self {
545 drivers: Vec::new(),
546 by_spec: HashMap::new(),
547 nodes: None,
548 startup_specs: Vec::new(),
549 compiled: Default::default(),
550 }
551 }
552}
553
554impl<Context: Send + Sync> DriverRegistry<Context> {
555 pub fn new() -> Self {
557 Self::default()
558 }
559
560 pub fn with_node_registry(mut self, nodes: NodeRegistry) -> Self {
562 self.compiled.get_mut().clear();
563 self.startup_specs.clear();
564 self.nodes = Some(Arc::new(nodes));
565 self
566 }
567
568 pub fn prewarm_spec(&mut self, spec: &crate::Spec) -> Result<(), String> {
570 let nodes = self.nodes.as_ref().ok_or("node registry is unavailable")?;
571 let driver = SpecDriver::new(spec, nodes).map_err(|error| error.to_string())?;
572 let compiled = self.compiled.get_mut();
573 if compiled.len() == COMPILED_REVISION_LIMIT {
574 compiled.pop_front();
575 }
576 compiled.push_back(CompiledRevision {
577 scope: None,
578 key: (spec.spec_id.clone(), String::new()),
579 content: serde_json::to_value(spec).map_err(|error| error.to_string())?,
580 driver: Arc::new(driver),
581 });
582 if !self.startup_specs.contains(&spec.spec_id) {
583 self.startup_specs.push(spec.spec_id.clone());
584 }
585 Ok(())
586 }
587
588 pub async fn resolve_revision(
591 &self,
592 item: &WorkItem,
593 revision: WorkflowRevision,
594 ) -> Result<Arc<dyn WorkflowDriver<Context>>, String>
595 where
596 Context: 'static,
597 {
598 item.validate_pins()?;
599 revision.validate().map_err(|error| error.to_string())?;
600 if revision.definition_id != item.definition_id
601 || revision.revision != item.workflow_revision
602 || revision.content_digest != item.workflow_revision_digest
603 || revision.spec.spec_id != item.spec_id
604 || revision.kernel_abi_version != item.kernel_abi_version
605 || revision.capabilities != item.capability_pins
606 {
607 return Err("published revision does not match the claimed pins".into());
608 }
609 if !matches!(revision.kernel_abi_version.as_str(), "1" | "kernel-abi-1") {
610 return Err("unsupported workflow kernel ABI".into());
611 }
612 let nodes = self
613 .nodes
614 .as_ref()
615 .ok_or("dynamic spec execution is unavailable")?;
616 let builtins = NodeRegistry::with_builtins();
617 for (expression, version) in &revision.expression_versions {
618 let supported = nodes.capability_manifests().any(|manifest| {
619 manifest.id == *expression && manifest.contract_version == *version
620 }) || (builtins.is_step(expression) && version == "1");
621 if !supported {
622 return Err("unsupported expression version".into());
623 }
624 }
625 for pin in &revision.capabilities {
626 let manifest = nodes
627 .capability_by_pin(pin)
628 .ok_or("pinned capability is unavailable")?;
629 debug_assert_eq!(manifest.id, pin.id);
630 }
631 let selected_nodes = nodes
632 .for_capability_pins(&revision.capabilities)
633 .map_err(|error| error.to_string())?;
634 for node in revision
635 .spec
636 .branches
637 .iter()
638 .flat_map(|branch| &branch.nodes)
639 {
640 if let Some(manifest) = selected_nodes.capability(&node.node_type) {
641 if !revision.capabilities.iter().any(|pin| {
642 pin.id == manifest.id
643 && pin.contract_version == manifest.contract_version
644 && pin.content_digest == manifest.content_digest
645 }) {
646 return Err("graph capability lacks an exact revision pin".into());
647 }
648 } else if nodes
649 .capability_manifests()
650 .any(|manifest| manifest.id == node.node_type)
651 {
652 return Err("graph capability lacks an exact revision pin".into());
653 } else if !builtins.is_step(&node.node_type) && !builtins.is_ingress(&node.node_type) {
654 return Err("product graph node requires a capability manifest".into());
655 }
656 }
657 let content = serde_json::json!({
658 "spec": revision.spec, "abi": revision.kernel_abi_version,
659 "capabilities": revision.capabilities,
660 "expressions": revision.expression_versions,
661 "dependencies": revision.dependency_set_digest,
662 });
663 let key = (item.spec_id.clone(), item.workflow_revision_digest.clone());
664 let scope = Some((
665 item.tenant_id.clone(),
666 item.definition_id
667 .parse()
668 .map_err(|_| "invalid workflow definition identity")?,
669 ));
670 let mut compiled = self.compiled.lock().await;
671 if let Some(entry) = compiled
672 .iter()
673 .find(|entry| entry.key == key && entry.scope == scope)
674 {
675 if entry.content != content {
676 return Err("revision digest refers to different compiled content".into());
677 }
678 return Ok(entry.driver.clone());
679 }
680 if let Some(index) = compiled.iter().position(|entry| {
681 entry.key.0 == item.spec_id
682 && entry.key.1.is_empty()
683 && entry.content == content["spec"]
684 }) {
685 let mut entry = compiled.remove(index).expect("located cache entry exists");
686 entry.scope = scope;
687 entry.key = key;
688 entry.content = content;
689 let driver = entry.driver.clone();
690 compiled.push_back(entry);
691 return Ok(driver);
692 }
693 let nodes = selected_nodes;
694 let driver: Arc<dyn WorkflowDriver<Context>> = Arc::new(
697 tokio::task::spawn_blocking(move || {
698 SpecDriver::new(&revision.spec, &nodes).map_err(|error| error.to_string())
699 })
700 .await
701 .map_err(|_| "spec compiler panicked")??,
702 );
703 if compiled.len() == COMPILED_REVISION_LIMIT {
704 compiled.pop_front();
705 }
706 compiled.push_back(CompiledRevision {
707 scope,
708 key,
709 content,
710 driver: driver.clone(),
711 });
712 Ok(driver)
713 }
714
715 async fn resolve(
716 &self,
717 item: &WorkItem,
718 revision: Option<WorkflowRevision>,
719 ) -> Result<Arc<dyn WorkflowDriver<Context>>, String>
720 where
721 Context: 'static,
722 {
723 if let Some(driver) = self.by_spec.get(&item.spec_id) {
724 if !driver.supports_revision(item) {
725 return Err("custom driver does not support the claimed revision".into());
726 }
727 return Ok(driver.clone());
728 }
729 let revision = revision.ok_or("prepared workflow revision is unavailable")?;
730 self.resolve_revision(item, revision).await
731 }
732
733 fn requires_revision(&self, item: &WorkItem) -> bool {
734 !self.by_spec.contains_key(&item.spec_id)
735 }
736
737 pub fn register(
740 &mut self,
741 driver: Arc<dyn WorkflowDriver<Context>>,
742 ) -> Result<&mut Self, SupervisorError> {
743 for spec_id in driver.spec_ids() {
744 if let Some(existing) = self.by_spec.get(spec_id) {
745 return Err(SupervisorError::Driver {
746 driver: driver.name().to_owned(),
747 reason: format!(
748 "spec '{spec_id}' is already claimed by '{}'",
749 existing.name()
750 ),
751 });
752 }
753 self.by_spec.insert(spec_id.to_owned(), driver.clone());
754 }
755 self.drivers.push(driver);
756 Ok(self)
757 }
758
759 pub fn spec_ids(&self) -> Vec<String> {
761 let mut ids = self.by_spec.keys().cloned().collect::<Vec<_>>();
762 ids.extend(self.startup_specs.iter().cloned());
763 ids.sort();
764 ids.dedup();
765 ids
766 }
767
768 pub fn for_spec(&self, spec_id: &str) -> Option<&Arc<dyn WorkflowDriver<Context>>> {
770 self.by_spec.get(spec_id)
771 }
772
773 pub fn names(&self) -> Vec<&'static str> {
775 self.drivers.iter().map(|driver| driver.name()).collect()
776 }
777
778 pub fn is_empty(&self) -> bool {
780 self.drivers.is_empty() && self.startup_specs.is_empty()
781 }
782
783 pub fn validate_all(&self) -> Result<(), SupervisorError> {
785 for driver in &self.drivers {
786 driver
787 .validate_specs()
788 .map_err(|reason| SupervisorError::Driver {
789 driver: driver.name().to_owned(),
790 reason,
791 })?;
792 }
793 Ok(())
794 }
795}
796
797pub async fn run_due_pass<Context, Queue, BuildContext, BuildFuture>(
800 queue: &Queue,
801 registry: &DriverRegistry<Context>,
802 settings: &SupervisorSettings,
803 build_context: BuildContext,
804) -> Result<SupervisorStats, SupervisorError>
805where
806 Context: Send + Sync + 'static,
807 Queue: WorkQueue,
808 BuildContext: FnOnce() -> BuildFuture,
809 BuildFuture: Future<Output = Result<Context, String>> + Send,
810{
811 if registry.is_empty() && registry.nodes.is_none() {
812 return Ok(SupervisorStats::default());
813 }
814 let claim_limit = settings
815 .claim_batch
816 .clamp(1, i64::try_from(settings.concurrency()).unwrap_or(i64::MAX));
817 let items = queue
818 .claim_due(
819 &if registry.nodes.is_some() {
820 Vec::new()
821 } else {
822 registry.spec_ids()
823 },
824 &settings.worker_id,
825 settings.lease_secs,
826 claim_limit,
827 )
828 .await?;
829 if items.is_empty() {
830 return Ok(SupervisorStats::default());
831 }
832
833 let context = match build_context().await {
834 Ok(context) => Arc::new(context),
835 Err(error) => {
836 let mut cleanup_errors = Vec::new();
837 for item in &items {
838 let command = WorkflowTransitionCommand::failure(
839 item,
840 error.clone(),
841 Utc::now() + chrono::Duration::seconds(settings.requeue_delay_secs),
842 );
843 if let Err(cleanup) = queue.commit_command(item, &command).await {
844 cleanup_errors.push(format!("{} commit failure: {cleanup}", item.id));
845 }
846 }
847 if !cleanup_errors.is_empty() {
848 return Err(SupervisorError::Context(format!(
849 "{error}; claimed work cleanup failed: {}",
850 cleanup_errors.join(", ")
851 )));
852 }
853 return Err(SupervisorError::Context(error));
854 }
855 };
856
857 let claimed = items.len();
858 let mut failed = 0;
859 let mut persistence_errors = Vec::new();
860 let mut work = items.into_iter();
861 let mut tasks = FuturesUnordered::new();
862
863 let spawn_next = |tasks: &mut FuturesUnordered<_>, work: &mut std::vec::IntoIter<WorkItem>| {
864 let Some(item) = work.next() else {
865 return false;
866 };
867 let context = context.clone();
868 let renewal_period = std::time::Duration::from_millis(
869 (settings.lease_secs.clamp(1, 3600) as u64 * 1000 / 3).max(1),
870 );
871 tasks.push(async move {
872 let evaluation_item = item.clone();
873 let evaluation = async {
874 evaluation_item.validate_pins()?;
875 let revision = queue
876 .prepare_evaluation(
877 &evaluation_item,
878 registry.requires_revision(&evaluation_item),
879 )
880 .await
881 .map_err(|error| error.to_string())?;
882 let driver = registry.resolve(&evaluation_item, revision).await?;
883 AssertUnwindSafe(driver.evaluate(&context, &evaluation_item))
884 .catch_unwind()
885 .await
886 .map_err(|_| "driver panicked".to_string())
887 .and_then(|result| result)
888 };
889 tokio::pin!(evaluation);
890 let mut renewal = tokio::time::interval_at(
891 tokio::time::Instant::now() + renewal_period,
892 renewal_period,
893 );
894 renewal.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
895 let result = loop {
896 tokio::select! {
897 result = &mut evaluation => break result,
898 _ = renewal.tick() => {
899 if let Err(error) = queue
900 .renew(
901 &item.tenant_id,
902 &item.id,
903 &settings.worker_id,
904 item.lease_version,
905 settings.lease_secs,
906 )
907 .await
908 {
909 break Err(format!("lease renewal failed: {error}"));
910 }
911 }
912 }
913 };
914 (item, result)
915 });
916 true
917 };
918
919 for _ in 0..settings.concurrency() {
920 if !spawn_next(&mut tasks, &mut work) {
921 break;
922 }
923 }
924
925 while let Some((item, result)) = tasks.next().await {
926 let command = match result {
927 Ok(command) => match command.validate_for(&item) {
928 Ok(()) => command,
929 Err(error) => {
930 failed += 1;
931 WorkflowTransitionCommand::failure(
932 &item,
933 error,
934 Utc::now() + chrono::Duration::seconds(settings.requeue_delay_secs),
935 )
936 }
937 },
938 Err(error) => {
939 failed += 1;
940 WorkflowTransitionCommand::failure(
941 &item,
942 error,
943 Utc::now() + chrono::Duration::seconds(settings.requeue_delay_secs),
944 )
945 }
946 };
947 if let Err(error) = queue.commit_command(&item, &command).await {
948 failed += usize::from(!matches!(
949 command.disposition,
950 WorkDisposition::Failed { .. }
951 ));
952 persistence_errors.push(format!("{} command commit: {error}", item.id));
953 }
954 spawn_next(&mut tasks, &mut work);
955 }
956
957 if persistence_errors.is_empty() {
958 Ok(SupervisorStats { claimed, failed })
959 } else {
960 Err(SupervisorError::Commit(persistence_errors.join(", ")))
961 }
962}
963
964#[cfg(test)]
965mod tests {
966 use super::*;
967 use std::sync::atomic::{AtomicBool, Ordering};
968 use std::sync::Mutex;
969
970 struct Context;
971 struct Driver {
972 valid: bool,
973 }
974
975 #[async_trait]
976 impl WorkflowDriver<Context> for Driver {
977 fn name(&self) -> &'static str {
978 "driver"
979 }
980 fn spec_ids(&self) -> Vec<&str> {
981 vec!["spec"]
982 }
983 fn validate_specs(&self) -> Result<(), String> {
984 self.valid.then_some(()).ok_or_else(|| "invalid".into())
985 }
986 fn supports_revision(&self, item: &WorkItem) -> bool {
987 item.workflow_revision_digest == "workflow-digest" && item.kernel_abi_version == "1"
988 }
989 async fn evaluate(
990 &self,
991 _context: &Context,
992 item: &WorkItem,
993 ) -> Result<WorkflowTransitionCommand, String> {
994 if let Some(milliseconds) = item.config["sleep_ms"].as_u64() {
995 tokio::time::sleep(std::time::Duration::from_millis(milliseconds)).await;
996 }
997 if item.config["fail"] == true {
998 Err("planned".into())
999 } else if item.config["reschedule"] == true {
1000 Ok(command(
1001 item,
1002 WorkDisposition::Reschedule { at: Utc::now() },
1003 serde_json::json!({"next": true}),
1004 ))
1005 } else if item.config["stop"] == true {
1006 Ok(command(
1007 item,
1008 WorkDisposition::Complete,
1009 item.config.clone(),
1010 ))
1011 } else {
1012 Ok(command(
1013 item,
1014 WorkDisposition::Continue { delay_secs: 5 },
1015 item.config.clone(),
1016 ))
1017 }
1018 }
1019 }
1020
1021 fn command(
1022 item: &WorkItem,
1023 disposition: WorkDisposition,
1024 next_state: Value,
1025 ) -> WorkflowTransitionCommand {
1026 WorkflowTransitionCommand {
1027 consumed_wakeups: item
1028 .wakeups
1029 .iter()
1030 .map(|wakeup| wakeup.id.clone())
1031 .collect(),
1032 delivery_key: format!("test:{}:{}", item.id, item.state_version),
1033 delivery_digest: "delivery-digest".into(),
1034 event_type: "workflow.test".into(),
1035 event_digest: "event-digest".into(),
1036 event_payload: Value::Null,
1037 next_state,
1038 action_intents: Vec::new(),
1039 outcome: EvaluationOutcome::default(),
1040 disposition,
1041 }
1042 }
1043
1044 #[derive(Default)]
1045 struct Queue {
1046 items: Mutex<Vec<WorkItem>>,
1047 committed: Mutex<Vec<(String, WorkDisposition)>>,
1048 renewals: Mutex<Vec<(String, String, i64)>>,
1049 preparations: Mutex<Vec<(String, bool)>>,
1050 revision: Option<WorkflowRevision>,
1051 renewal_error: Mutex<Option<String>>,
1052 finalization_error: Mutex<Option<String>>,
1053 claim_limits: Mutex<Vec<i64>>,
1054 }
1055
1056 #[async_trait]
1057 impl WorkQueue for Queue {
1058 async fn prepare_evaluation(
1059 &self,
1060 item: &WorkItem,
1061 load_revision: bool,
1062 ) -> Result<Option<WorkflowRevision>, SupervisorError> {
1063 self.preparations
1064 .lock()
1065 .unwrap()
1066 .push((item.id.clone(), load_revision));
1067 Ok(load_revision.then(|| self.revision.clone()).flatten())
1068 }
1069 async fn claim_due(
1070 &self,
1071 _spec_ids: &[String],
1072 _worker_id: &str,
1073 _lease_secs: i64,
1074 batch: i64,
1075 ) -> Result<Vec<WorkItem>, SupervisorError> {
1076 self.claim_limits.lock().unwrap().push(batch);
1077 let mut items = self.items.lock().unwrap();
1078 let take = items.len().min(batch as usize);
1079 Ok(items.drain(..take).collect())
1080 }
1081
1082 async fn renew(
1083 &self,
1084 _tenant_id: &str,
1085 id: &str,
1086 worker_id: &str,
1087 lease_version: i64,
1088 _lease_secs: i64,
1089 ) -> Result<(), SupervisorError> {
1090 self.renewals.lock().unwrap().push((
1091 id.to_string(),
1092 worker_id.to_string(),
1093 lease_version,
1094 ));
1095 match self.renewal_error.lock().unwrap().clone() {
1096 Some(error) => Err(SupervisorError::Queue(error)),
1097 None => Ok(()),
1098 }
1099 }
1100
1101 async fn commit_command(
1102 &self,
1103 item: &WorkItem,
1104 command: &WorkflowTransitionCommand,
1105 ) -> Result<(), SupervisorError> {
1106 self.committed
1107 .lock()
1108 .unwrap()
1109 .push((item.id.clone(), command.disposition.clone()));
1110 match self.finalization_error.lock().unwrap().clone() {
1111 Some(error) => Err(SupervisorError::Queue(error)),
1112 None => Ok(()),
1113 }
1114 }
1115 }
1116
1117 fn item(id: &str, fail: bool) -> WorkItem {
1118 WorkItem {
1119 id: id.into(),
1120 run_id: uuid::Uuid::new_v4().to_string().parse().unwrap(),
1121 tenant_id: "tenant".parse().unwrap(),
1122 subject_id: "subject".parse().unwrap(),
1123 spec_id: "spec".into(),
1124 definition_id: "definition".into(),
1125 workflow_revision: 1,
1126 workflow_revision_digest: "workflow-digest".into(),
1127 execution_profile_id: "profile".into(),
1128 execution_profile_revision: 1,
1129 execution_profile_digest: "profile-digest".into(),
1130 kernel_abi_version: "1".into(),
1131 capability_pins: Vec::new(),
1132 lifecycle: LifecyclePolicy::run_once(),
1133 scheduled_at: Utc::now(),
1134 created_at: Utc::now(),
1135 claimed_at: Utc::now(),
1136 config: serde_json::json!({"fail": fail}),
1137 state_version: 0,
1138 control_epochs: ControlEpochs::default(),
1139 cancel_requested: false,
1140 lease_version: 1,
1141 wakeups: Vec::new(),
1142 }
1143 }
1144
1145 fn stopped_item(id: &str) -> WorkItem {
1146 WorkItem {
1147 config: serde_json::json!({"stop": true}),
1148 ..item(id, false)
1149 }
1150 }
1151
1152 fn settings() -> SupervisorSettings {
1153 SupervisorSettings {
1154 worker_id: "worker".into(),
1155 lease_secs: 60,
1156 requeue_delay_secs: 5,
1157 claim_batch: 10,
1158 concurrency: 3,
1159 }
1160 }
1161
1162 struct MustNotEvaluate(Arc<AtomicBool>);
1163
1164 #[async_trait]
1165 impl WorkflowDriver<Context> for MustNotEvaluate {
1166 fn name(&self) -> &'static str {
1167 "must-not-evaluate"
1168 }
1169
1170 fn spec_ids(&self) -> Vec<&str> {
1171 vec!["spec"]
1172 }
1173
1174 fn validate_specs(&self) -> Result<(), String> {
1175 Ok(())
1176 }
1177
1178 async fn evaluate(
1179 &self,
1180 _context: &Context,
1181 item: &WorkItem,
1182 ) -> Result<WorkflowTransitionCommand, String> {
1183 self.0.store(true, Ordering::SeqCst);
1184 Ok(command(
1185 item,
1186 WorkDisposition::Complete,
1187 item.config.clone(),
1188 ))
1189 }
1190 }
1191
1192 #[tokio::test]
1193 async fn missing_revision_pin_fails_before_driver_evaluation() {
1194 let mut stale = item("stale", false);
1195 stale.workflow_revision_digest.clear();
1196 let queue = Queue {
1197 items: Mutex::new(vec![stale]),
1198 ..Default::default()
1199 };
1200 let called = Arc::new(AtomicBool::new(false));
1201 let mut registry = DriverRegistry::new();
1202 registry
1203 .register(Arc::new(MustNotEvaluate(called.clone())))
1204 .unwrap();
1205
1206 let stats = run_due_pass(&queue, ®istry, &settings(), || async { Ok(Context) })
1207 .await
1208 .unwrap();
1209
1210 assert_eq!(stats.failed, 1);
1211 assert!(!called.load(Ordering::SeqCst));
1212 assert!(matches!(
1213 queue.committed.lock().unwrap().as_slice(),
1214 [(id, WorkDisposition::Failed { error, .. })]
1215 if id == "stale" && error.contains("workflow_revision_digest")
1216 ));
1217 }
1218
1219 #[tokio::test]
1220 async fn due_pass_releases_success_and_marks_failures() {
1221 let queue = Queue {
1222 items: Mutex::new(vec![
1223 item("ok", false),
1224 stopped_item("done"),
1225 item("bad", true),
1226 ]),
1227 ..Default::default()
1228 };
1229 let mut registry = DriverRegistry::new();
1230 registry.register(Arc::new(Driver { valid: true })).unwrap();
1231 let stats = run_due_pass(&queue, ®istry, &settings(), || async { Ok(Context) })
1232 .await
1233 .unwrap();
1234 assert_eq!(
1235 stats,
1236 SupervisorStats {
1237 claimed: 3,
1238 failed: 1
1239 }
1240 );
1241 let committed = queue.committed.lock().unwrap();
1242 assert_eq!(committed.len(), 3);
1243 assert!(committed
1244 .iter()
1245 .any(|(id, disposition)| id == "done" && disposition == &WorkDisposition::Complete));
1246 assert!(committed.iter().any(|(id, disposition)| id == "bad"
1247 && matches!(disposition, WorkDisposition::Failed { .. })));
1248 let mut preparations = queue.preparations.lock().unwrap().clone();
1249 preparations.sort();
1250 assert_eq!(
1251 preparations,
1252 [
1253 ("bad".into(), false),
1254 ("done".into(), false),
1255 ("ok".into(), false)
1256 ]
1257 );
1258 }
1259
1260 #[tokio::test]
1261 async fn dynamic_driver_uses_the_revision_returned_by_single_preparation() {
1262 let mut work = item("dynamic", false);
1263 work.spec_id = "dynamic-spec".into();
1264 work.definition_id = "dynamic-definition".into();
1265 work.workflow_revision_digest = "dynamic-digest".into();
1266 work.wakeups.push(Wakeup {
1267 id: "delivery:event:binding:1".into(),
1268 kind: "delivery".into(),
1269 branch_id: Some("__root__".parse().unwrap()),
1270 payload: serde_json::json!({"value": 7}),
1271 });
1272 let revision = WorkflowRevision {
1273 definition_id: work.definition_id.clone(),
1274 revision: work.workflow_revision,
1275 content_digest: work.workflow_revision_digest.clone(),
1276 kernel_abi_version: work.kernel_abi_version.clone(),
1277 dependency_set_digest: "dependencies".into(),
1278 expression_versions: Default::default(),
1279 capabilities: Vec::new(),
1280 template_provenance: Value::Null,
1281 spec: serde_json::from_value(serde_json::json!({
1282 "spec_id": work.spec_id,
1283 "version": "1",
1284 "branches": [{
1285 "branch_id": "__root__",
1286 "nodes": [{"id": "in", "type": "ingress.event", "config": {}}],
1287 "edges": []
1288 }]
1289 }))
1290 .unwrap(),
1291 };
1292 let queue = Queue {
1293 items: Mutex::new(vec![work]),
1294 revision: Some(revision),
1295 ..Default::default()
1296 };
1297 let registry = DriverRegistry::new().with_node_registry(NodeRegistry::with_builtins());
1298
1299 let stats = run_due_pass(&queue, ®istry, &settings(), || async { Ok(Context) })
1300 .await
1301 .unwrap();
1302
1303 assert_eq!(
1304 stats,
1305 SupervisorStats {
1306 claimed: 1,
1307 failed: 0
1308 }
1309 );
1310 assert_eq!(
1311 queue.preparations.lock().unwrap().as_slice(),
1312 [("dynamic".into(), true)]
1313 );
1314 assert_eq!(queue.committed.lock().unwrap().len(), 1);
1315 }
1316
1317 #[tokio::test]
1318 async fn context_failure_releases_every_claim() {
1319 let queue = Queue {
1320 items: Mutex::new(vec![item("one", false), item("two", false)]),
1321 ..Default::default()
1322 };
1323 let mut registry = DriverRegistry::new();
1324 registry.register(Arc::new(Driver { valid: true })).unwrap();
1325 let error = run_due_pass(&queue, ®istry, &settings(), || async {
1326 Err::<Context, _>("context failed".into())
1327 })
1328 .await
1329 .unwrap_err();
1330 assert_eq!(error, SupervisorError::Context("context failed".into()));
1331 assert_eq!(queue.committed.lock().unwrap().len(), 2);
1332 assert!(queue
1333 .committed
1334 .lock()
1335 .unwrap()
1336 .iter()
1337 .all(|(_, disposition)| matches!(disposition, WorkDisposition::Failed { .. })));
1338 }
1339
1340 #[tokio::test]
1341 async fn due_pass_atomically_reschedules_driver_state() {
1342 let queue = Queue {
1343 items: Mutex::new(vec![WorkItem {
1344 config: serde_json::json!({"reschedule": true}),
1345 ..item("recurring", false)
1346 }]),
1347 ..Default::default()
1348 };
1349 let mut registry = DriverRegistry::new();
1350 registry.register(Arc::new(Driver { valid: true })).unwrap();
1351 run_due_pass(&queue, ®istry, &settings(), || async { Ok(Context) })
1352 .await
1353 .unwrap();
1354 assert!(matches!(
1355 queue.committed.lock().unwrap().as_slice(),
1356 [(id, WorkDisposition::Reschedule { .. })] if id == "recurring"
1357 ));
1358 }
1359
1360 #[tokio::test]
1361 async fn finalization_failure_is_not_reported_as_success() {
1362 let queue = Queue {
1363 items: Mutex::new(vec![item("stale", false)]),
1364 finalization_error: Mutex::new(Some("stale lease".into())),
1365 ..Default::default()
1366 };
1367 let mut registry = DriverRegistry::new();
1368 registry.register(Arc::new(Driver { valid: true })).unwrap();
1369
1370 let error = run_due_pass(&queue, ®istry, &settings(), || async { Ok(Context) })
1371 .await
1372 .unwrap_err();
1373
1374 assert!(
1375 matches!(&error, SupervisorError::Commit(reason) if reason.contains("stale lease"))
1376 );
1377 assert!(error.to_string().contains("state commit failed"));
1378 }
1379
1380 #[tokio::test]
1381 async fn due_pass_claims_only_work_that_can_start() {
1382 let queue = Queue {
1383 items: Mutex::new(vec![
1384 item("one", false),
1385 item("two", false),
1386 item("queued", false),
1387 ]),
1388 ..Default::default()
1389 };
1390 let mut registry = DriverRegistry::new();
1391 registry.register(Arc::new(Driver { valid: true })).unwrap();
1392 let mut settings = settings();
1393 settings.concurrency = 2;
1394
1395 let stats = run_due_pass(&queue, ®istry, &settings, || async { Ok(Context) })
1396 .await
1397 .unwrap();
1398
1399 assert_eq!(stats.claimed, 2);
1400 assert_eq!(queue.claim_limits.lock().unwrap().as_slice(), [2]);
1401 assert_eq!(queue.items.lock().unwrap().len(), 1);
1402 }
1403
1404 #[tokio::test(start_paused = true)]
1405 async fn long_evaluate_renews_its_lease() {
1406 let queue = Queue {
1407 items: Mutex::new(vec![WorkItem {
1408 config: serde_json::json!({"sleep_ms": 3500}),
1409 ..item("slow", false)
1410 }]),
1411 ..Default::default()
1412 };
1413 let mut registry = DriverRegistry::new();
1414 registry.register(Arc::new(Driver { valid: true })).unwrap();
1415 let mut settings = settings();
1416 settings.lease_secs = 3;
1417 settings.concurrency = 1;
1418
1419 let stats = run_due_pass(&queue, ®istry, &settings, || async { Ok(Context) })
1420 .await
1421 .unwrap();
1422
1423 assert_eq!(stats.failed, 0);
1424 let renewals = queue.renewals.lock().unwrap();
1425 assert_eq!(renewals.len(), 3);
1426 assert!(renewals
1427 .iter()
1428 .all(|renewal| renewal == &("slow".into(), "worker".into(), 1)));
1429 assert!(matches!(
1430 queue.committed.lock().unwrap().as_slice(),
1431 [(id, WorkDisposition::Continue { .. })] if id == "slow"
1432 ));
1433 }
1434
1435 #[tokio::test(start_paused = true)]
1436 async fn expired_lease_stops_evaluation() {
1437 let queue = Queue {
1438 items: Mutex::new(vec![WorkItem {
1439 config: serde_json::json!({"sleep_ms": 5000}),
1440 ..item("expired", false)
1441 }]),
1442 renewal_error: Mutex::new(Some("lease expired".into())),
1443 ..Default::default()
1444 };
1445 let mut registry = DriverRegistry::new();
1446 registry.register(Arc::new(Driver { valid: true })).unwrap();
1447 let mut settings = settings();
1448 settings.lease_secs = 3;
1449 settings.concurrency = 1;
1450
1451 let stats = run_due_pass(&queue, ®istry, &settings, || async { Ok(Context) })
1452 .await
1453 .unwrap();
1454
1455 assert_eq!(stats.failed, 1);
1456 assert_eq!(queue.renewals.lock().unwrap().len(), 1);
1457 assert!(matches!(
1458 queue.committed.lock().unwrap().as_slice(),
1459 [(id, WorkDisposition::Failed { .. })] if id == "expired"
1460 ));
1461 }
1462
1463 #[test]
1464 fn registry_fails_duplicate_ownership_and_invalid_specs() {
1465 let mut registry = DriverRegistry::new();
1466 assert!(registry.is_empty());
1467 registry
1468 .register(Arc::new(Driver { valid: false }))
1469 .unwrap();
1470 assert_eq!(registry.spec_ids(), ["spec"]);
1471 assert_eq!(registry.names(), ["driver"]);
1472 assert!(registry.for_spec("spec").is_some());
1473 assert_eq!(
1474 registry.validate_all().unwrap_err(),
1475 SupervisorError::Driver {
1476 driver: "driver".into(),
1477 reason: "invalid".into()
1478 }
1479 );
1480 assert!(matches!(
1481 registry.register(Arc::new(Driver { valid: true })),
1482 Err(SupervisorError::Driver { reason, .. }) if reason.contains("already claimed")
1483 ));
1484 }
1485
1486 #[tokio::test]
1487 async fn memory_work_queue_serves_consumers_without_a_database() {
1488 let queue = MemoryWorkQueue::new(vec![item("first", false)]);
1489 queue.push(item("other-spec", false));
1490 {
1491 let mut items = queue.items.lock().unwrap();
1492 items[1].spec_id = "unknown".into();
1493 }
1494 let mut registry = DriverRegistry::new();
1495 registry.register(Arc::new(Driver { valid: true })).unwrap();
1496 let stats = run_due_pass(&queue, ®istry, &settings(), || async { Ok(Context) })
1497 .await
1498 .unwrap();
1499 assert_eq!(
1500 stats,
1501 SupervisorStats {
1502 claimed: 1,
1503 failed: 0
1504 },
1505 "only items for registered specs are claimed"
1506 );
1507 let committed = queue.committed();
1508 assert_eq!(committed.len(), 1);
1509 assert_eq!(committed[0].0.id, "first");
1510 assert_eq!(
1511 committed[0].1.disposition,
1512 WorkDisposition::Continue { delay_secs: 5 }
1513 );
1514 assert!(queue.renewals().is_empty());
1515 assert_eq!(
1516 run_due_pass(&queue, ®istry, &settings(), || async { Ok(Context) })
1517 .await
1518 .unwrap()
1519 .claimed,
1520 0,
1521 "claimed items leave the queue"
1522 );
1523 }
1524
1525 #[test]
1526 fn commands_are_validated_against_the_claimed_work() {
1527 let work = item("scoped", false);
1528 let mut blank = command(&work, WorkDisposition::Complete, Value::Null);
1529 blank.event_type.clear();
1530 assert!(blank
1531 .validate_for(&work)
1532 .unwrap_err()
1533 .contains("missing event_type"));
1534
1535 let mut cancelled = item("cancelled", false);
1536 cancelled.cancel_requested = true;
1537 let mut with_intent = command(&cancelled, WorkDisposition::Complete, Value::Null);
1538 with_intent.action_intents.push(ActionIntent {
1539 id: "intent".into(),
1540 tenant_id: cancelled.tenant_id.clone(),
1541 instance_id: cancelled.id.clone().parse().unwrap(),
1542 run_id: cancelled.run_id.clone(),
1543 capability: CapabilityPin {
1544 id: "cap".into(),
1545 contract_version: "1".into(),
1546 content_digest: "digest".into(),
1547 },
1548 idempotency_key: "key".into(),
1549 state: crate::ActionState::Prepared,
1550 input: Value::Null,
1551 effect: crate::Effect::ExternalWrite,
1552 retry_class: crate::IdempotencyMode::Native,
1553 control_epochs: ControlEpochs::default(),
1554 resource_scope_id: String::new(),
1555 lease_epoch: cancelled.lease_version,
1556 action_epoch: cancelled.state_version,
1557 deadline: None,
1558 reservation: None,
1559 created_at: Utc::now(),
1560 });
1561 assert!(with_intent
1562 .validate_for(&cancelled)
1563 .unwrap_err()
1564 .contains("cancelled work"));
1565
1566 let mut escaped = with_intent.clone();
1567 escaped.action_intents[0].tenant_id = "someone-else".parse().unwrap();
1568 let scoped = item("scoped", false);
1569 let mut escaped_command = command(&scoped, WorkDisposition::Complete, Value::Null);
1570 escaped_command.action_intents = escaped.action_intents;
1571 assert!(escaped_command
1572 .validate_for(&scoped)
1573 .unwrap_err()
1574 .contains("escapes claimed work scope"));
1575 }
1576}