assay_workflow/
ctx.rs

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

assay_workflow/
ctx.rs
