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::health;
9use crate::scheduler;
10use crate::store::WorkflowStore;
11use crate::timers;
12
13/// Event the engine broadcasts when a workflow transitions through
14/// its lifecycle states. Subscribed to by the SSE stream (`/api/v1/
15/// events/stream`) so the dashboard can refresh live instead of
16/// requiring the operator to F5 after every action.
17#[derive(Clone, Debug)]
18pub struct EngineEvent {
19    pub event_type: String,
20    pub workflow_id: String,
21    pub namespace: String,
22}
23
24/// SSE-level broadcast event forwarded to connected dashboard browsers.
25#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
26pub struct BroadcastEvent {
27    pub event_type: String,
28    pub workflow_id: String,
29    pub payload: Option<String>,
30}
31
32/// Holds the background-task JoinHandles. When the last Arc<BackgroundTasks>
33/// is dropped the tasks are abandoned (tokio will cancel them on shutdown).
34pub struct BackgroundTasks {
35    _scheduler: JoinHandle<()>,
36    _timer_poller: JoinHandle<()>,
37    _health_monitor: JoinHandle<()>,
38    _dispatch_recovery: JoinHandle<()>,
39    #[cfg(feature = "s3-archival")]
40    _archival: Option<JoinHandle<()>>,
41}
42
43/// The workflow context. Owns the store, background-task handles, and
44/// per-request config. Serves as both the orchestrator (all engine methods
45/// live as `impl WorkflowCtx<S>`) and the axum state (`Arc<WorkflowCtx<S>>`).
46///
47/// `S` is the concrete store backend (`SqliteStore` or `PostgresStore`).
48/// `WorkflowStore` uses RPIT futures and is not dyn-compatible, so the
49/// generic parameter is retained here to avoid boxing every async call.
50pub struct WorkflowCtx<S: WorkflowStore> {
51    pub(crate) store: Arc<S>,
52    /// Engine-level event sender — lifecycle methods push here;
53    /// the serve layer bridges to `sse_tx`.
54    pub(crate) event_tx: Option<tokio::sync::broadcast::Sender<EngineEvent>>,
55    /// SSE broadcast sender — the SSE handler subscribes here.
56    pub sse_tx: Option<tokio::sync::broadcast::Sender<BroadcastEvent>>,
57    pub(crate) _bg: Arc<BackgroundTasks>,
58    pub auth_mode: AuthMode,
59    /// Version of the containing binary (e.g. the `assay-lua` CLI) — set
60    /// by embedders so `/api/v1/version` reflects the user-facing
61    /// binary, not this internal engine-crate version.
62    pub binary_version: Option<&'static str>,
63}
64
65impl<S: WorkflowStore> WorkflowCtx<S> {
66    /// Start the context with all background tasks.
67    pub fn start(store: Arc<S>) -> Self {
68        let _scheduler = tokio::spawn(scheduler::run_scheduler(Arc::clone(&store)));
69        let _timer_poller = tokio::spawn(timers::run_timer_poller(Arc::clone(&store)));
70        let _health_monitor = tokio::spawn(health::run_health_monitor(Arc::clone(&store)));
71        let _dispatch_recovery = tokio::spawn(dispatch_recovery::run_dispatch_recovery(
72            Arc::clone(&store),
73        ));
74
75        #[cfg(feature = "s3-archival")]
76        let _archival = crate::archival::ArchivalConfig::from_env().map(|cfg| {
77            tokio::spawn(crate::archival::run_archival(Arc::clone(&store), cfg))
78        });
79
80        info!("Workflow engine started");
81
82        Self {
83            store,
84            event_tx: None,
85            sse_tx: None,
86            _bg: Arc::new(BackgroundTasks {
87                _scheduler,
88                _timer_poller,
89                _health_monitor,
90                _dispatch_recovery,
91                #[cfg(feature = "s3-archival")]
92                _archival,
93            }),
94            auth_mode: AuthMode::default(),
95            binary_version: None,
96        }
97    }
98
99    /// Attach an event broadcaster. The API layer sets this up so the
100    /// SSE stream (`/events/stream`) sees state transitions as they
101    /// happen — powers the dashboard's live list refresh, no F5 loop.
102    /// Returns the context by value so callers can chain:
103    ///
104    /// ```ignore
105    /// let ctx = WorkflowCtx::start(store).with_event_broadcaster(tx);
106    /// ```
107    pub fn with_event_broadcaster(
108        mut self,
109        tx: tokio::sync::broadcast::Sender<EngineEvent>,
110    ) -> Self {
111        self.event_tx = Some(tx);
112        self
113    }
114
115    /// Attach the SSE broadcast sender so the event stream handler can subscribe.
116    pub fn with_sse_tx(
117        mut self,
118        tx: tokio::sync::broadcast::Sender<BroadcastEvent>,
119    ) -> Self {
120        self.sse_tx = Some(tx);
121        self
122    }
123
124    /// Set the auth mode.
125    pub fn with_auth_mode(mut self, auth_mode: AuthMode) -> Self {
126        self.auth_mode = auth_mode;
127        self
128    }
129
130    /// Set the binary version string surfaced by `/api/v1/version`.
131    pub fn with_binary_version(mut self, version: &'static str) -> Self {
132        self.binary_version = Some(version);
133        self
134    }
135
136    /// Access the underlying store (for the API layer).
137    pub fn store(&self) -> &S {
138        &self.store
139    }
140
141    /// Broadcast a state-transition event. No-op when no broadcaster
142    /// is wired up (tests, embedders without an SSE surface). Errors
143    /// from a channel with zero subscribers are silently dropped —
144    /// that's the normal state between connections, not a failure.
145    pub(crate) fn broadcast(&self, event_type: &str, workflow_id: &str, namespace: &str) {
146        if let Some(tx) = &self.event_tx {
147            let _ = tx.send(EngineEvent {
148                event_type: event_type.to_string(),
149                workflow_id: workflow_id.to_string(),
150                namespace: namespace.to_string(),
151            });
152        }
153    }
154}
155
156/// Strip a trailing `-continued-<digits>` from a workflow id so
157/// sequential continue-as-new calls don't pile up suffixes. Matches
158/// the pattern emitted by the default id-derivation path; returns the
159/// input unchanged if there's no such suffix.
160pub(crate) fn strip_continued_suffix(id: &str) -> &str {
161    if let Some(idx) = id.rfind("-continued-") {
162        let (head, tail) = id.split_at(idx);
163        // Only strip when the tail after "-continued-" is all digits.
164        let rest = &tail["-continued-".len()..];
165        if !rest.is_empty() && rest.chars().all(|c| c.is_ascii_digit()) {
166            return head;
167        }
168    }
169    id
170}
171
172pub(crate) fn timestamp_now() -> f64 {
173    std::time::SystemTime::now()
174        .duration_since(std::time::UNIX_EPOCH)
175        .unwrap()
176        .as_secs_f64()
177}
178
179/// WorkflowCtx version (the binary version pulled from Cargo at build time).
180/// Stamped into every workflow's search_attributes at start so operators
181/// can correlate runs to the engine release that executed them.
182pub(crate) const ENGINE_VERSION: &str = env!("CARGO_PKG_VERSION");
183
184/// Auto-stamp `assay_engine_version` into a workflow's search attributes.
185/// Returns `Some` JSON string for the caller to store in the record.
186///
187/// If the caller already supplied `assay_engine_version` in their patch,
188/// we leave their value alone (explicit override wins). Otherwise we
189/// backfill the running engine's version. Callers who supply no
190/// attributes at all get a single-key object with just the version.
191pub(crate) fn inject_engine_version(caller_attrs: Option<&str>) -> Option<String> {
192    let mut obj: serde_json::Map<String, serde_json::Value> = match caller_attrs {
193        Some(raw) => match serde_json::from_str::<serde_json::Value>(raw) {
194            Ok(serde_json::Value::Object(m)) => m,
195            // Non-object JSON (or unparsable) — preserve as-is without
196            // stamping; we can't safely merge a key into a non-object.
197            Ok(other) => return Some(other.to_string()),
198            Err(_) => return Some(raw.to_string()),
199        },
200        None => serde_json::Map::new(),
201    };
202    obj.entry("assay_engine_version".to_string())
203        .or_insert_with(|| serde_json::Value::String(ENGINE_VERSION.to_string()));
204    Some(serde_json::Value::Object(obj).to_string())
205}
206
207#[cfg(test)]
208mod engine_version_stamp_tests {
209    use super::*;
210
211    #[test]
212    fn no_attrs_produces_single_key_object() {
213        let out = inject_engine_version(None).unwrap();
214        let v: serde_json::Value = serde_json::from_str(&out).unwrap();
215        assert_eq!(v["assay_engine_version"], ENGINE_VERSION);
216        assert_eq!(v.as_object().unwrap().len(), 1);
217    }
218
219    #[test]
220    fn existing_attrs_gain_the_version_field() {
221        let out = inject_engine_version(Some(r#"{"env":"prod","tenant":"acme"}"#)).unwrap();
222        let v: serde_json::Value = serde_json::from_str(&out).unwrap();
223        assert_eq!(v["env"], "prod");
224        assert_eq!(v["tenant"], "acme");
225        assert_eq!(v["assay_engine_version"], ENGINE_VERSION);
226    }
227
228    #[test]
229    fn caller_supplied_version_wins_on_conflict() {
230        let out = inject_engine_version(Some(r#"{"assay_engine_version":"0.0.1-test"}"#)).unwrap();
231        let v: serde_json::Value = serde_json::from_str(&out).unwrap();
232        assert_eq!(v["assay_engine_version"], "0.0.1-test");
233    }
234
235    #[test]
236    fn non_object_json_is_preserved_unchanged() {
237        let out = inject_engine_version(Some("[1, 2, 3]")).unwrap();
238        assert_eq!(out, "[1,2,3]");
239    }
240
241    #[test]
242    fn unparsable_json_is_preserved_unchanged() {
243        let out = inject_engine_version(Some("not json")).unwrap();
244        assert_eq!(out, "not json");
245    }
246}
assay_workflow/ctx.rs

assay_workflow/
ctx.rs
