scud/attractor/
context.rs

1//! Thread-safe key-value execution context for pipeline runs.
2
3use serde::{Deserialize, Serialize};
4use std::collections::HashMap;
5use std::sync::Arc;
6use tokio::sync::RwLock;
7
8/// Thread-safe key-value context shared across pipeline execution.
9///
10/// Stores arbitrary JSON values keyed by string names. Supports
11/// isolated clones for parallel branches and atomic update application.
12#[derive(Debug, Clone)]
13pub struct Context {
14    inner: Arc<RwLock<HashMap<String, serde_json::Value>>>,
15}
16
17impl Context {
18    /// Create an empty context.
19    pub fn new() -> Self {
20        Self {
21            inner: Arc::new(RwLock::new(HashMap::new())),
22        }
23    }
24
25    /// Create a context with initial values.
26    pub fn with_values(values: HashMap<String, serde_json::Value>) -> Self {
27        Self {
28            inner: Arc::new(RwLock::new(values)),
29        }
30    }
31
32    /// Set a value in the context.
33    pub async fn set(&self, key: impl Into<String>, value: serde_json::Value) {
34        self.inner.write().await.insert(key.into(), value);
35    }
36
37    /// Get a value from the context.
38    pub async fn get(&self, key: &str) -> Option<serde_json::Value> {
39        self.inner.read().await.get(key).cloned()
40    }
41
42    /// Get a string value from the context.
43    pub async fn get_str(&self, key: &str) -> Option<String> {
44        self.get(key)
45            .await
46            .and_then(|v| v.as_str().map(String::from))
47    }
48
49    /// Take a snapshot of the current context state.
50    pub async fn snapshot(&self) -> HashMap<String, serde_json::Value> {
51        self.inner.read().await.clone()
52    }
53
54    /// Create an isolated clone for parallel branches.
55    ///
56    /// Changes to the clone do not affect the original.
57    pub async fn clone_isolated(&self) -> Self {
58        Self::with_values(self.snapshot().await)
59    }
60
61    /// Apply a batch of updates atomically.
62    pub async fn apply_updates(&self, updates: &HashMap<String, serde_json::Value>) {
63        let mut inner = self.inner.write().await;
64        for (key, value) in updates {
65            inner.insert(key.clone(), value.clone());
66        }
67    }
68
69    /// Check if the context contains a key.
70    pub async fn contains_key(&self, key: &str) -> bool {
71        self.inner.read().await.contains_key(key)
72    }
73
74    /// Get the number of entries.
75    pub async fn len(&self) -> usize {
76        self.inner.read().await.len()
77    }
78
79    /// Check if the context is empty.
80    pub async fn is_empty(&self) -> bool {
81        self.inner.read().await.is_empty()
82    }
83}
84
85impl Default for Context {
86    fn default() -> Self {
87        Self::new()
88    }
89}
90
91/// Serializable snapshot of a context for checkpointing.
92#[derive(Debug, Clone, Serialize, Deserialize)]
93pub struct ContextSnapshot {
94    pub values: HashMap<String, serde_json::Value>,
95}
96
97impl From<HashMap<String, serde_json::Value>> for ContextSnapshot {
98    fn from(values: HashMap<String, serde_json::Value>) -> Self {
99        Self { values }
100    }
101}
102
103impl ContextSnapshot {
104    /// Restore a Context from this snapshot.
105    pub fn restore(&self) -> Context {
106        Context::with_values(self.values.clone())
107    }
108}
109
110#[cfg(test)]
111mod tests {
112    use super::*;
113
114    #[tokio::test]
115    async fn test_set_and_get() {
116        let ctx = Context::new();
117        ctx.set("name", serde_json::json!("Alice")).await;
118        assert_eq!(ctx.get("name").await, Some(serde_json::json!("Alice")));
119    }
120
121    #[tokio::test]
122    async fn test_get_str() {
123        let ctx = Context::new();
124        ctx.set("greeting", serde_json::json!("hello")).await;
125        assert_eq!(ctx.get_str("greeting").await, Some("hello".to_string()));
126
127        ctx.set("number", serde_json::json!(42)).await;
128        assert_eq!(ctx.get_str("number").await, None);
129    }
130
131    #[tokio::test]
132    async fn test_snapshot() {
133        let ctx = Context::new();
134        ctx.set("a", serde_json::json!(1)).await;
135        ctx.set("b", serde_json::json!(2)).await;
136        let snap = ctx.snapshot().await;
137        assert_eq!(snap.len(), 2);
138    }
139
140    #[tokio::test]
141    async fn test_clone_isolated() {
142        let ctx = Context::new();
143        ctx.set("shared", serde_json::json!("original")).await;
144
145        let clone = ctx.clone_isolated().await;
146        clone.set("shared", serde_json::json!("modified")).await;
147
148        // Original should be unchanged
149        assert_eq!(ctx.get_str("shared").await, Some("original".to_string()));
150        assert_eq!(clone.get_str("shared").await, Some("modified".to_string()));
151    }
152
153    #[tokio::test]
154    async fn test_apply_updates() {
155        let ctx = Context::new();
156        let mut updates = HashMap::new();
157        updates.insert("x".into(), serde_json::json!(10));
158        updates.insert("y".into(), serde_json::json!(20));
159        ctx.apply_updates(&updates).await;
160        assert_eq!(ctx.len().await, 2);
161    }
162
163    #[tokio::test]
164    async fn test_snapshot_roundtrip() {
165        let ctx = Context::new();
166        ctx.set("key", serde_json::json!("value")).await;
167        let snap = ContextSnapshot::from(ctx.snapshot().await);
168        let json = serde_json::to_string(&snap).unwrap();
169        let restored_snap: ContextSnapshot = serde_json::from_str(&json).unwrap();
170        let restored = restored_snap.restore();
171        assert_eq!(restored.get_str("key").await, Some("value".to_string()));
172    }
173}
scud/attractor/context.rs

scud/attractor/
context.rs
