mojentic/context/
shared_working_memory.rs

1//! Shared working memory for agent communication.
2//!
3//! This module provides [`SharedWorkingMemory`], a thread-safe shared context
4//! that agents can read from and write to. It enables agents to maintain and
5//! share state across interactions.
6//!
7//! # Examples
8//!
9//! ```
10//! use mojentic::context::SharedWorkingMemory;
11//! use serde_json::json;
12//!
13//! let memory = SharedWorkingMemory::new(json!({
14//!     "user": {
15//!         "name": "Alice",
16//!         "age": 30
17//!     }
18//! }));
19//!
20//! let current = memory.get_working_memory();
21//! assert_eq!(current["user"]["name"], "Alice");
22//!
23//! memory.merge_to_working_memory(json!({
24//!     "user": {
25//!         "age": 31,
26//!         "city": "Boston"
27//!     }
28//! }));
29//!
30//! let updated = memory.get_working_memory();
31//! assert_eq!(updated["user"]["age"], 31);
32//! assert_eq!(updated["user"]["city"], "Boston");
33//! assert_eq!(updated["user"]["name"], "Alice"); // Original value preserved
34//! ```
35
36use serde_json::Value;
37use std::sync::{Arc, Mutex};
38
39/// Thread-safe shared working memory for agents.
40///
41/// `SharedWorkingMemory` provides a shared context that multiple agents can
42/// read from and write to. It uses deep merging to combine updates with
43/// existing state, preserving values that aren't being updated.
44///
45/// The memory is thread-safe and can be safely cloned and shared across
46/// multiple agents and async tasks.
47#[derive(Debug, Clone)]
48pub struct SharedWorkingMemory {
49    memory: Arc<Mutex<Value>>,
50}
51
52impl SharedWorkingMemory {
53    /// Create a new `SharedWorkingMemory` with initial state.
54    ///
55    /// # Arguments
56    ///
57    /// * `initial_memory` - The initial JSON value for the working memory
58    ///
59    /// # Examples
60    ///
61    /// ```
62    /// use mojentic::context::SharedWorkingMemory;
63    /// use serde_json::json;
64    ///
65    /// let memory = SharedWorkingMemory::new(json!({
66    ///     "user": { "name": "Bob" }
67    /// }));
68    /// ```
69    pub fn new(initial_memory: Value) -> Self {
70        Self {
71            memory: Arc::new(Mutex::new(initial_memory)),
72        }
73    }
74
75    /// Get a clone of the current working memory.
76    ///
77    /// Returns a snapshot of the current state. Subsequent changes to the
78    /// working memory will not affect this returned value.
79    ///
80    /// # Examples
81    ///
82    /// ```
83    /// use mojentic::context::SharedWorkingMemory;
84    /// use serde_json::json;
85    ///
86    /// let memory = SharedWorkingMemory::new(json!({"count": 1}));
87    /// let snapshot = memory.get_working_memory();
88    /// assert_eq!(snapshot["count"], 1);
89    /// ```
90    pub fn get_working_memory(&self) -> Value {
91        self.memory.lock().unwrap().clone()
92    }
93
94    /// Merge new values into the working memory.
95    ///
96    /// Performs a deep merge of the provided value with the existing memory.
97    /// For objects, this recursively merges nested fields. For arrays and
98    /// primitives, the new value replaces the old value.
99    ///
100    /// # Arguments
101    ///
102    /// * `new_memory` - The JSON value to merge into the working memory
103    ///
104    /// # Examples
105    ///
106    /// ```
107    /// use mojentic::context::SharedWorkingMemory;
108    /// use serde_json::json;
109    ///
110    /// let memory = SharedWorkingMemory::new(json!({
111    ///     "user": {
112    ///         "name": "Charlie",
113    ///         "age": 25
114    ///     }
115    /// }));
116    ///
117    /// memory.merge_to_working_memory(json!({
118    ///     "user": {
119    ///         "age": 26,
120    ///         "city": "NYC"
121    ///     }
122    /// }));
123    ///
124    /// let result = memory.get_working_memory();
125    /// assert_eq!(result["user"]["name"], "Charlie"); // Preserved
126    /// assert_eq!(result["user"]["age"], 26);         // Updated
127    /// assert_eq!(result["user"]["city"], "NYC");     // Added
128    /// ```
129    pub fn merge_to_working_memory(&self, new_memory: Value) {
130        let mut memory = self.memory.lock().unwrap();
131        deep_merge(&mut memory, new_memory);
132    }
133}
134
135impl Default for SharedWorkingMemory {
136    /// Create a new `SharedWorkingMemory` with an empty object as initial state.
137    fn default() -> Self {
138        Self::new(Value::Object(serde_json::Map::new()))
139    }
140}
141
142/// Deep merge two JSON values.
143///
144/// If both values are objects, recursively merge their fields.
145/// Otherwise, replace the destination with the source value.
146///
147/// # Arguments
148///
149/// * `dest` - The destination value to merge into (modified in place)
150/// * `src` - The source value to merge from
151fn deep_merge(dest: &mut Value, src: Value) {
152    match (dest, src) {
153        (Value::Object(dest_map), Value::Object(src_map)) => {
154            for (key, value) in src_map {
155                dest_map
156                    .entry(key)
157                    .and_modify(|dest_value| deep_merge(dest_value, value.clone()))
158                    .or_insert(value);
159            }
160        }
161        (dest_value, src_value) => {
162            *dest_value = src_value;
163        }
164    }
165}
166
167#[cfg(test)]
168mod tests {
169    use super::*;
170    use serde_json::json;
171
172    #[test]
173    fn test_new() {
174        let memory = SharedWorkingMemory::new(json!({"key": "value"}));
175        let result = memory.get_working_memory();
176        assert_eq!(result["key"], "value");
177    }
178
179    #[test]
180    fn test_default() {
181        let memory = SharedWorkingMemory::default();
182        let result = memory.get_working_memory();
183        assert!(result.is_object());
184        assert!(result.as_object().unwrap().is_empty());
185    }
186
187    #[test]
188    fn test_get_working_memory() {
189        let memory = SharedWorkingMemory::new(json!({
190            "user": {
191                "name": "Alice",
192                "age": 30
193            }
194        }));
195
196        let result = memory.get_working_memory();
197        assert_eq!(result["user"]["name"], "Alice");
198        assert_eq!(result["user"]["age"], 30);
199    }
200
201    #[test]
202    fn test_merge_to_working_memory_simple() {
203        let memory = SharedWorkingMemory::new(json!({"key1": "value1"}));
204
205        memory.merge_to_working_memory(json!({"key2": "value2"}));
206
207        let result = memory.get_working_memory();
208        assert_eq!(result["key1"], "value1");
209        assert_eq!(result["key2"], "value2");
210    }
211
212    #[test]
213    fn test_merge_to_working_memory_deep() {
214        let memory = SharedWorkingMemory::new(json!({
215            "user": {
216                "name": "Bob",
217                "age": 25
218            }
219        }));
220
221        memory.merge_to_working_memory(json!({
222            "user": {
223                "age": 26,
224                "city": "Boston"
225            }
226        }));
227
228        let result = memory.get_working_memory();
229        assert_eq!(result["user"]["name"], "Bob"); // Preserved
230        assert_eq!(result["user"]["age"], 26); // Updated
231        assert_eq!(result["user"]["city"], "Boston"); // Added
232    }
233
234    #[test]
235    fn test_merge_to_working_memory_replace_primitive() {
236        let memory = SharedWorkingMemory::new(json!({"count": 1}));
237
238        memory.merge_to_working_memory(json!({"count": 2}));
239
240        let result = memory.get_working_memory();
241        assert_eq!(result["count"], 2);
242    }
243
244    #[test]
245    fn test_merge_to_working_memory_replace_array() {
246        let memory = SharedWorkingMemory::new(json!({"items": [1, 2, 3]}));
247
248        memory.merge_to_working_memory(json!({"items": [4, 5]}));
249
250        let result = memory.get_working_memory();
251        assert_eq!(result["items"], json!([4, 5]));
252    }
253
254    #[test]
255    fn test_merge_nested_objects() {
256        let memory = SharedWorkingMemory::new(json!({
257            "level1": {
258                "level2": {
259                    "level3": {
260                        "value": "original"
261                    }
262                }
263            }
264        }));
265
266        memory.merge_to_working_memory(json!({
267            "level1": {
268                "level2": {
269                    "level3": {
270                        "new_value": "added"
271                    }
272                }
273            }
274        }));
275
276        let result = memory.get_working_memory();
277        assert_eq!(result["level1"]["level2"]["level3"]["value"], "original");
278        assert_eq!(result["level1"]["level2"]["level3"]["new_value"], "added");
279    }
280
281    #[test]
282    fn test_thread_safety() {
283        use std::thread;
284
285        let memory = SharedWorkingMemory::new(json!({"count": 0}));
286
287        let handles: Vec<_> = (0..10)
288            .map(|i| {
289                let mem = memory.clone();
290                thread::spawn(move || {
291                    mem.merge_to_working_memory(json!({format!("key{}", i): i}));
292                })
293            })
294            .collect();
295
296        for handle in handles {
297            handle.join().unwrap();
298        }
299
300        let result = memory.get_working_memory();
301        assert_eq!(result["count"], 0);
302
303        // Check that all keys were added
304        for i in 0..10 {
305            assert!(result.get(format!("key{}", i)).is_some());
306        }
307    }
308
309    #[test]
310    fn test_deep_merge_primitives() {
311        let mut dest = json!(42);
312        deep_merge(&mut dest, json!(100));
313        assert_eq!(dest, json!(100));
314    }
315
316    #[test]
317    fn test_deep_merge_arrays() {
318        let mut dest = json!([1, 2, 3]);
319        deep_merge(&mut dest, json!([4, 5]));
320        assert_eq!(dest, json!([4, 5]));
321    }
322
323    #[test]
324    fn test_deep_merge_mixed_types() {
325        let mut dest = json!({"key": [1, 2, 3]});
326        deep_merge(&mut dest, json!({"key": "string"}));
327        assert_eq!(dest, json!({"key": "string"}));
328    }
329
330    #[test]
331    fn test_deep_merge_empty_objects() {
332        let mut dest = json!({});
333        deep_merge(&mut dest, json!({"key": "value"}));
334        assert_eq!(dest, json!({"key": "value"}));
335    }
336
337    #[test]
338    fn test_multiple_merges() {
339        let memory = SharedWorkingMemory::new(json!({}));
340
341        memory.merge_to_working_memory(json!({"a": 1}));
342        memory.merge_to_working_memory(json!({"b": 2}));
343        memory.merge_to_working_memory(json!({"c": 3}));
344
345        let result = memory.get_working_memory();
346        assert_eq!(result["a"], 1);
347        assert_eq!(result["b"], 2);
348        assert_eq!(result["c"], 3);
349    }
350
351    #[test]
352    fn test_clone_memory() {
353        let memory1 = SharedWorkingMemory::new(json!({"key": "value"}));
354        let memory2 = memory1.clone();
355
356        memory2.merge_to_working_memory(json!({"key2": "value2"}));
357
358        let result1 = memory1.get_working_memory();
359        let result2 = memory2.get_working_memory();
360
361        // Both should have the update since they share the same Arc
362        assert_eq!(result1["key2"], "value2");
363        assert_eq!(result2["key2"], "value2");
364    }
365}
mojentic/context/shared_working_memory.rs

mojentic/context/
shared_working_memory.rs
