assay_workflow/
ctx.rs

1use std::sync::Arc;
2
3use tokio::task::JoinHandle;
4use tracing::info;
5
6use crate::dispatch_recovery;
7use crate::events::{WorkflowBusEvent, WorkflowEventBus};
8use crate::health;
9use crate::scheduler;
10use crate::store::WorkflowStore;
11use crate::timers;
12
13/// Holds the background-task JoinHandles. When the last Arc<BackgroundTasks>
14/// is dropped the tasks are abandoned (tokio will cancel them on shutdown).
15pub struct BackgroundTasks {
16    _scheduler: JoinHandle<()>,
17    _timer_poller: JoinHandle<()>,
18    _health_monitor: JoinHandle<()>,
19    _dispatch_recovery: JoinHandle<()>,
20    #[cfg(feature = "s3-archival")]
21    _archival: Option<JoinHandle<()>>,
22}
23
24/// The workflow context. Owns the store, background-task handles, and
25/// per-request config. Serves as both the orchestrator (all engine methods
26/// live as `impl WorkflowCtx<S>`) and the axum state (`Arc<WorkflowCtx<S>>`).
27///
28/// `S` is the concrete store backend (`SqliteStore` or `PostgresStore`).
29/// `WorkflowStore` uses RPIT futures and is not dyn-compatible, so the
30/// generic parameter is retained here to avoid boxing every async call.
31pub struct WorkflowCtx<S: WorkflowStore> {
32    pub(crate) store: Arc<S>,
33    /// Engine-wide CDC outbox. When wired, state-mutating methods
34    /// publish typed `WorkflowBusEvent` variants via `emit(...)`.
35    /// `None` for tests / embedders without a dashboard — emit becomes
36    /// a no-op.
37    pub(crate) bus: Option<WorkflowEventBus>,
38    pub(crate) _bg: Arc<BackgroundTasks>,
39    /// Version of the containing binary (e.g. the `assay-lua` CLI) — set
40    /// by embedders so `/api/v1/engine/workflow/version` reflects the user-facing
41    /// binary, not this internal engine-crate version.
42    pub binary_version: Option<&'static str>,
43}
44
45impl<S: WorkflowStore> WorkflowCtx<S> {
46    /// Start the context with all background tasks.
47    pub fn start(store: Arc<S>) -> Self {
48        let _scheduler = tokio::spawn(scheduler::run_scheduler(Arc::clone(&store)));
49        let _timer_poller = tokio::spawn(timers::run_timer_poller(Arc::clone(&store)));
50        let _health_monitor = tokio::spawn(health::run_health_monitor(Arc::clone(&store)));
51        let _dispatch_recovery =
52            tokio::spawn(dispatch_recovery::run_dispatch_recovery(Arc::clone(&store)));
53
54        #[cfg(feature = "s3-archival")]
55        let _archival = crate::archival::ArchivalConfig::from_env()
56            .map(|cfg| tokio::spawn(crate::archival::run_archival(Arc::clone(&store), cfg)));
57
58        info!("Workflow engine started");
59
60        Self {
61            store,
62            bus: None,
63            _bg: Arc::new(BackgroundTasks {
64                _scheduler,
65                _timer_poller,
66                _health_monitor,
67                _dispatch_recovery,
68                #[cfg(feature = "s3-archival")]
69                _archival,
70            }),
71            binary_version: None,
72        }
73    }
74
75    /// Attach the engine-wide event bus. The API layer sets this up so
76    /// the SSE stream (`/events/stream`) and the dispatch-wakeup loop
77    /// see state transitions as they happen. Returns the context by
78    /// value so callers can chain.
79    pub fn with_event_bus(mut self, bus: WorkflowEventBus) -> Self {
80        self.bus = Some(bus);
81        self
82    }
83
84    /// Set the binary version string surfaced by `/api/v1/engine/workflow/version`.
85    pub fn with_binary_version(mut self, version: &'static str) -> Self {
86        self.binary_version = Some(version);
87        self
88    }
89
90    /// Access the underlying store (for the API layer).
91    pub fn store(&self) -> &S {
92        &self.store
93    }
94
95    /// Access the event bus (for SSE + scheduler wake-up).
96    pub fn bus(&self) -> Option<&WorkflowEventBus> {
97        self.bus.as_ref()
98    }
99
100    /// Emit a typed workflow event. No-op when no bus is wired (tests,
101    /// embedders without a dashboard). Errors are logged, not returned —
102    /// an emission failure must not fail the state-mutating method that
103    /// triggered it (atomicity for the state change is the DB tx's job;
104    /// this is a notification fired *after* the row write).
105    pub(crate) async fn emit(&self, namespace: &str, ev: WorkflowBusEvent) {
106        if let Some(bus) = &self.bus
107            && let Err(e) = bus.publish(namespace, ev).await
108        {
109            tracing::warn!(?e, "engine event emit failed");
110        }
111    }
112
113    pub(crate) async fn emit_retry_requested(
114        &self,
115        namespace: &str,
116        workflow_id: &str,
117        activity_id: i64,
118        activity_seq: i32,
119    ) {
120        if let Some(bus) = &self.bus
121            && let Err(e) = bus
122                .publish_retry_requested(namespace, workflow_id, activity_id, activity_seq)
123                .await
124        {
125            tracing::warn!(?e, "engine retry event emit failed");
126        }
127    }
128
129    /// Mark a workflow dispatchable AND emit a `WorkflowNeedsDispatch`
130    /// on the bus so the dispatch-wakeup loop wakes workers
131    /// on this node / across the cluster. The extra SELECT is skipped
132    /// when no bus is wired.
133    pub(crate) async fn mark_and_emit_needs_dispatch(
134        &self,
135        workflow_id: &str,
136    ) -> anyhow::Result<()> {
137        self.store.mark_workflow_dispatchable(workflow_id).await?;
138        self.emit_needs_dispatch(workflow_id).await;
139        Ok(())
140    }
141
142    /// Emit `WorkflowNeedsDispatch` without touching the row — for callers
143    /// whose store method already armed the workflow inside its own
144    /// transaction. Like [`WorkflowCtx::emit`] this never fails the caller:
145    /// the arming is durable, so a missed wake-up costs a poll interval,
146    /// not the run.
147    pub(crate) async fn emit_needs_dispatch(&self, workflow_id: &str) {
148        if self.bus.is_none() {
149            return;
150        }
151        match self.store.get_workflow(workflow_id).await {
152            Ok(Some(wf)) => {
153                self.emit(
154                    &wf.namespace,
155                    WorkflowBusEvent::WorkflowNeedsDispatch {
156                        workflow_id: workflow_id.to_string(),
157                        task_queue: wf.task_queue,
158                    },
159                )
160                .await;
161            }
162            Ok(None) => {}
163            Err(e) => tracing::warn!(?e, "needs-dispatch emit lookup failed"),
164        }
165    }
166}
167
168/// Strip a trailing `-continued-<digits>` from a workflow id so
169/// sequential continue-as-new calls don't pile up suffixes. Matches
170/// the pattern emitted by the default id-derivation path; returns the
171/// input unchanged if there's no such suffix.
172pub(crate) fn strip_continued_suffix(id: &str) -> &str {
173    if let Some(idx) = id.rfind("-continued-") {
174        let (head, tail) = id.split_at(idx);
175        let rest = &tail["-continued-".len()..];
176        if !rest.is_empty() && rest.chars().all(|c| c.is_ascii_digit()) {
177            return head;
178        }
179    }
180    id
181}
182
183pub(crate) fn timestamp_now() -> f64 {
184    std::time::SystemTime::now()
185        .duration_since(std::time::UNIX_EPOCH)
186        .unwrap()
187        .as_secs_f64()
188}
189
190/// WorkflowCtx version (the binary version pulled from Cargo at build time).
191/// Stamped into every workflow's search_attributes at start so operators
192/// can correlate runs to the engine release that executed them.
193pub(crate) const ENGINE_VERSION: &str = env!("CARGO_PKG_VERSION");
194
195/// Auto-stamp `assay_engine_version` into a workflow's search attributes.
196/// Returns `Some` JSON string for the caller to store in the record.
197///
198/// If the caller already supplied `assay_engine_version` in their patch,
199/// we leave their value alone (explicit override wins). Otherwise we
200/// backfill the running engine's version. Callers who supply no
201/// attributes at all get a single-key object with just the version.
202pub(crate) fn inject_engine_version(caller_attrs: Option<&str>) -> Option<String> {
203    let mut obj: serde_json::Map<String, serde_json::Value> = match caller_attrs {
204        Some(raw) => match serde_json::from_str::<serde_json::Value>(raw) {
205            Ok(serde_json::Value::Object(m)) => m,
206            Ok(other) => return Some(other.to_string()),
207            Err(_) => return Some(raw.to_string()),
208        },
209        None => serde_json::Map::new(),
210    };
211    obj.entry("assay_engine_version".to_string())
212        .or_insert_with(|| serde_json::Value::String(ENGINE_VERSION.to_string()));
213    Some(serde_json::Value::Object(obj).to_string())
214}
215
216#[cfg(test)]
217mod engine_version_stamp_tests {
218    use super::*;
219
220    #[test]
221    fn no_attrs_produces_single_key_object() {
222        let out = inject_engine_version(None).unwrap();
223        let v: serde_json::Value = serde_json::from_str(&out).unwrap();
224        assert_eq!(v["assay_engine_version"], ENGINE_VERSION);
225        assert_eq!(v.as_object().unwrap().len(), 1);
226    }
227
228    #[test]
229    fn existing_attrs_gain_the_version_field() {
230        let out = inject_engine_version(Some(r#"{"env":"prod","tenant":"acme"}"#)).unwrap();
231        let v: serde_json::Value = serde_json::from_str(&out).unwrap();
232        assert_eq!(v["env"], "prod");
233        assert_eq!(v["tenant"], "acme");
234        assert_eq!(v["assay_engine_version"], ENGINE_VERSION);
235    }
236
237    #[test]
238    fn caller_supplied_version_wins_on_conflict() {
239        let out = inject_engine_version(Some(r#"{"assay_engine_version":"0.0.1-test"}"#)).unwrap();
240        let v: serde_json::Value = serde_json::from_str(&out).unwrap();
241        assert_eq!(v["assay_engine_version"], "0.0.1-test");
242    }
243
244    #[test]
245    fn non_object_json_is_preserved_unchanged() {
246        let out = inject_engine_version(Some("[1, 2, 3]")).unwrap();
247        assert_eq!(out, "[1,2,3]");
248    }
249
250    #[test]
251    fn unparsable_json_is_preserved_unchanged() {
252        let out = inject_engine_version(Some("not json")).unwrap();
253        assert_eq!(out, "not json");
254    }
255}
assay_workflow/ctx.rs

assay_workflow/
ctx.rs
