Skip to main content

sayiir_core/
context.rs

1//! Workflow execution context.
2//!
3//! [`WorkflowContext`] carries the workflow ID, codec, and user-supplied
4//! metadata through every task execution.
5//!
6//! [`TaskExecutionContext`] provides read-only access to workflow and task
7//! metadata from within running tasks. It is set automatically by the
8//! runtime and can be retrieved via [`get_task_context()`] or the
9//! [`task_context!`](crate::task_context) macro.
10
11use std::sync::Arc;
12
13use crate::task::TaskMetadata;
14
15/// Execution context available to a running task.
16///
17/// Provides read-only access to workflow and task metadata. Accessible
18/// from within task functions via task-local storage (Rust) or
19/// language-specific context APIs (Python/Node.js).
20#[derive(Clone, Debug)]
21pub struct TaskExecutionContext {
22    /// The workflow definition identifier (human-readable name).
23    ///
24    /// Carried as `Arc<str>` because the external task executor (Python /
25    /// Node FFI) needs to look up the task function by its registered name.
26    /// The runtime's hot map lookups happen against the hashed
27    /// [`WorkflowId`](crate::WorkflowId) elsewhere; here we keep the name.
28    pub workflow_id: Arc<str>,
29    /// The workflow instance identifier.
30    pub instance_id: Arc<str>,
31    /// The current task identifier (human-readable name).
32    pub task_id: Arc<str>,
33    /// Task metadata (timeout, retry policy, version, etc.).
34    pub metadata: TaskMetadata,
35    /// Optional JSON-encoded workflow-level metadata.
36    pub workflow_metadata_json: Option<Arc<str>>,
37}
38
39/// Workflow execution context that provides access to metadata and codec.
40///
41/// This context is always available as a plain struct used during workflow
42/// building and by the runner for codec/metadata access. The `workflow_id`
43/// is kept as both a hash ([`WorkflowId`](crate::WorkflowId)) for fast
44/// runtime comparison and an `Arc<str>` for log/error display.
45pub struct WorkflowContext<C, M> {
46    /// SHA-256 hash of the workflow identifier (for fast comparison and maps).
47    pub workflow_id: crate::WorkflowId,
48    /// Human-readable workflow name (for logs / error messages).
49    pub workflow_name: Arc<str>,
50    /// The codec used for serialization/deserialization.
51    pub codec: Arc<C>,
52    /// Immutable metadata attached to the workflow.
53    pub metadata: Arc<M>,
54    /// Optional JSON-encoded workflow-level metadata for task context.
55    pub metadata_json: Option<Arc<str>>,
56}
57
58impl<C, M> Clone for WorkflowContext<C, M> {
59    fn clone(&self) -> Self {
60        Self {
61            workflow_id: self.workflow_id,
62            workflow_name: Arc::clone(&self.workflow_name),
63            codec: Arc::clone(&self.codec),
64            metadata: Arc::clone(&self.metadata),
65            metadata_json: self.metadata_json.clone(),
66        }
67    }
68}
69
70impl<C, M> WorkflowContext<C, M> {
71    /// Create a new workflow context.
72    pub fn new(workflow_name: impl Into<Arc<str>>, codec: Arc<C>, metadata: Arc<M>) -> Self {
73        let workflow_name: Arc<str> = workflow_name.into();
74        let workflow_id = crate::WorkflowId::from(workflow_name.as_ref());
75        Self {
76            workflow_id,
77            workflow_name,
78            codec,
79            metadata,
80            metadata_json: None,
81        }
82    }
83
84    /// Returns the workflow human-readable name.
85    #[must_use]
86    pub fn workflow_id(&self) -> &str {
87        &self.workflow_name
88    }
89
90    /// Returns the workflow identifier hash.
91    #[must_use]
92    pub fn workflow_id_hash(&self) -> crate::WorkflowId {
93        self.workflow_id
94    }
95
96    /// Returns a clone of the codec `Arc`.
97    #[must_use]
98    pub fn codec(&self) -> Arc<C> {
99        self.codec.clone()
100    }
101
102    /// Returns a clone of the metadata `Arc`.
103    #[must_use]
104    pub fn metadata(&self) -> Arc<M> {
105        self.metadata.clone()
106    }
107}
108
109use std::cell::RefCell;
110
111std::thread_local! {
112    /// Thread-local fallback for `TaskExecutionContext`.
113    ///
114    /// Used by sync executor paths (Python GIL, Node.js main thread) where
115    /// tokio task-locals are not available.
116    static THREAD_LOCAL_TASK_CTX: RefCell<Option<TaskExecutionContext>> = const { RefCell::new(None) };
117}
118
119/// Set the task execution context in thread-local storage for the duration
120/// of the closure. Clears the context when the closure returns (even on panic).
121pub fn with_thread_local_task_context<R>(ctx: TaskExecutionContext, f: impl FnOnce() -> R) -> R {
122    THREAD_LOCAL_TASK_CTX.with(|cell| {
123        let prev = cell.borrow_mut().replace(ctx);
124        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(f));
125        *cell.borrow_mut() = prev;
126        match result {
127            Ok(r) => r,
128            Err(e) => std::panic::resume_unwind(e),
129        }
130    })
131}
132
133/// Get the task execution context from thread-local storage.
134#[must_use]
135pub fn get_thread_local_task_context() -> Option<TaskExecutionContext> {
136    THREAD_LOCAL_TASK_CTX.with(|cell| cell.borrow().clone())
137}
138
139// ── Task-local context storage (requires tokio) ─────────────────────────
140
141#[cfg(feature = "tokio")]
142mod task_local_ctx {
143    use super::TaskExecutionContext;
144
145    tokio::task_local! {
146        /// Task-local storage for task execution context.
147        static TASK_EXEC_CTX: Option<TaskExecutionContext>;
148    }
149
150    /// Set the task execution context in task-local storage and execute the future.
151    pub async fn with_task_context<F: std::future::Future>(
152        ctx: TaskExecutionContext,
153        fut: F,
154    ) -> F::Output {
155        TASK_EXEC_CTX.scope(Some(ctx), fut).await
156    }
157
158    /// Get the task execution context from task-local storage.
159    ///
160    /// Tries the tokio task-local first, then falls back to the thread-local.
161    #[must_use]
162    pub fn get_task_context() -> Option<TaskExecutionContext> {
163        TASK_EXEC_CTX
164            .try_with(std::clone::Clone::clone)
165            .ok()
166            .flatten()
167            .or_else(super::get_thread_local_task_context)
168    }
169}
170
171#[cfg(feature = "tokio")]
172pub use task_local_ctx::{get_task_context, with_task_context};
173
174/// Get the task execution context (non-tokio fallback).
175///
176/// Delegates to thread-local storage only.
177#[cfg(not(feature = "tokio"))]
178#[must_use]
179pub fn get_task_context() -> Option<TaskExecutionContext> {
180    get_thread_local_task_context()
181}
182
183/// Macro to access the task execution context from within a task.
184///
185/// Returns `Option<TaskExecutionContext>` — `None` if called outside of
186/// task execution context.
187///
188/// Usage:
189/// ```rust,ignore
190/// if let Some(ctx) = task_context!() {
191///     println!("workflow: {}, task: {}", ctx.workflow_id, ctx.task_id);
192/// }
193/// ```
194#[macro_export]
195macro_rules! task_context {
196    () => {
197        $crate::context::get_task_context()
198    };
199}
200
201#[cfg(all(test, feature = "tokio"))]
202#[allow(clippy::unwrap_used, clippy::panic)]
203mod tests {
204    use super::*;
205    use crate::task::TaskMetadata;
206
207    fn make_task_ctx() -> TaskExecutionContext {
208        TaskExecutionContext {
209            workflow_id: Arc::from("wf-1"),
210            instance_id: Arc::from("inst-1"),
211            task_id: Arc::from("task-a"),
212            metadata: TaskMetadata::default(),
213            workflow_metadata_json: None,
214        }
215    }
216
217    #[test]
218    fn thread_local_roundtrip() {
219        assert!(get_thread_local_task_context().is_none());
220
221        let ctx = make_task_ctx();
222        let result = with_thread_local_task_context(ctx.clone(), || {
223            let inner = get_thread_local_task_context().unwrap();
224            assert_eq!(&*inner.workflow_id, "wf-1");
225            assert_eq!(&*inner.instance_id, "inst-1");
226            assert_eq!(&*inner.task_id, "task-a");
227            42
228        });
229        assert_eq!(result, 42);
230
231        // Cleared after scope
232        assert!(get_thread_local_task_context().is_none());
233    }
234
235    #[test]
236    fn thread_local_restores_on_panic() {
237        let ctx = make_task_ctx();
238        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
239            with_thread_local_task_context(ctx, || {
240                panic!("boom");
241            })
242        }));
243        assert!(result.is_err());
244        assert!(get_thread_local_task_context().is_none());
245    }
246
247    #[test]
248    fn task_local_roundtrip() {
249        let rt = tokio::runtime::Builder::new_current_thread()
250            .enable_all()
251            .build()
252            .unwrap();
253        rt.block_on(async {
254            assert!(get_task_context().is_none());
255
256            let ctx = make_task_ctx();
257            let inner = with_task_context(ctx, async {
258                let c = get_task_context().unwrap();
259                assert_eq!(&*c.task_id, "task-a");
260                c
261            })
262            .await;
263
264            assert_eq!(&*inner.workflow_id, "wf-1");
265        });
266    }
267
268    #[test]
269    fn task_local_falls_back_to_thread_local() {
270        let rt = tokio::runtime::Builder::new_current_thread()
271            .enable_all()
272            .build()
273            .unwrap();
274        rt.block_on(async {
275            // Set only thread-local, no task-local — should still find it
276            let ctx = make_task_ctx();
277            let result = with_thread_local_task_context(ctx, get_task_context);
278            assert!(result.is_some());
279            assert_eq!(&*result.unwrap().instance_id, "inst-1");
280        });
281    }
282
283    #[test]
284    fn macro_works() {
285        let ctx = make_task_ctx();
286        with_thread_local_task_context(ctx, || {
287            let c = task_context!().unwrap();
288            assert_eq!(&*c.task_id, "task-a");
289        });
290    }
291}