molo_core/tool/shared_state.rs
1//! Shared state: a type-safe heterogeneous container.
2//!
3//! Tools receive it at **call time** via the `state` parameter of
4//! [`Tool::call`](crate::Tool::call): the state is owned by the caller and
5//! passed in on every call; tools themselves hold no global state. Agents
6//! mount it via [`with_state`](crate::agent::ReActAgent::with_state), the
7//! application reads and writes across runs, and multiple tools / agents
8//! can share the same instance.
9
10use std::any::{Any, TypeId};
11use std::collections::HashMap;
12use std::fmt;
13use std::sync::{Arc, RwLock};
14
15/// Shared state: a heterogeneous container accessed by type.
16///
17/// Type is the key: a second [`insert`](SharedState::insert) of the same
18/// type overwrites; retrieval takes a type parameter (compile-time type
19/// safety). Unlike `Box<dyn Any>`'s "store anything", this container holds
20/// **multiple values side by side, distinguished by type**.
21///
22/// An internal `RwLock` (read-heavy) provides cross-thread safety;
23/// [`Clone`] is cheap (Arc), so tools and agents — several agents — share
24/// the same instance.
25///
26/// # Example
27///
28/// ```rust
29/// # extern crate molo_core as molo;
30/// use molo::SharedState;
31///
32/// #[derive(Debug, Clone, PartialEq)]
33/// struct Session { user: String }
34///
35/// let state = SharedState::new();
36/// state.insert(Session { user: "alice".into() });
37/// state.with::<Session>(|s| assert_eq!(s.user, "alice"));
38/// state.with_mut::<Session>(|s| s.user = "bob".into());
39/// assert_eq!(state.get::<Session>(), Some(Session { user: "bob".into() }));
40/// ```
41#[derive(Clone, Default)]
42pub struct SharedState {
43 inner: Arc<RwLock<HashMap<TypeId, Box<dyn Any + Send + Sync>>>>,
44}
45
46impl fmt::Debug for SharedState {
47 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
48 // Inner values are opaque `dyn Any`; printing the entry count is
49 // enough (the caller knows the content types).
50 let entries = self.inner.read().unwrap_or_else(|e| e.into_inner()).len();
51 f.debug_struct("SharedState")
52 .field("entries", &entries)
53 .finish()
54 }
55}
56
57impl SharedState {
58 /// An empty shared state.
59 pub fn new() -> Self {
60 Self::default()
61 }
62
63 /// Store a custom value; storing the same type again overwrites.
64 pub fn insert<T: Any + Send + Sync>(&self, value: T) {
65 // A panicking user closure (see with_mut) poisons the lock: recover
66 // from poisoning so the container does not fail permanently, and
67 // shared state stays usable after tool panics are caught.
68 self.inner
69 .write()
70 .unwrap_or_else(|e| e.into_inner())
71 .insert(TypeId::of::<T>(), Box::new(value));
72 }
73
74 /// Retrieve (clone) a custom value; returns `None` when the type is
75 /// absent (the type must implement `Clone`; a clone is returned and the
76 /// original value in the container is unchanged).
77 pub fn get<T: Any + Clone>(&self) -> Option<T> {
78 self.inner
79 .read()
80 .unwrap_or_else(|e| e.into_inner())
81 .get(&TypeId::of::<T>())
82 .and_then(|v| v.downcast_ref::<T>())
83 .cloned()
84 }
85
86 /// Read an existing value under the lock; the closure does not run when
87 /// the type is absent.
88 ///
89 /// # Panics
90 ///
91 /// A panic inside the closure propagates outward; the lock gets
92 /// poisoned as a result, but the container recovers from poisoning and
93 /// later access works as usual (the closure returns nothing; for an
94 /// in-place read use [`get`](SharedState::get)).
95 ///
96 /// # Deadlock red line
97 ///
98 /// The underlying `RwLock` is **not reentrant**, and the closure runs
99 /// while the lock is held: the closure **must not** access the same
100 /// `SharedState` again (e.g. `with(|a| ... state.get::<B>() ...)`), or
101 /// it hangs with no error. To access several values at once, call
102 /// `get` sequentially / combine them into one struct, or move reads
103 /// and writes outside the closure.
104 pub fn with<T: Any>(&self, f: impl FnOnce(&T)) {
105 let guard = self.inner.read().unwrap_or_else(|e| e.into_inner());
106 if let Some(value) = guard
107 .get(&TypeId::of::<T>())
108 .and_then(|v| v.downcast_ref::<T>())
109 {
110 f(value);
111 }
112 }
113
114 /// Update an existing value in place under the lock; the closure does
115 /// not run when the type is absent.
116 ///
117 /// # Panics
118 ///
119 /// A panic inside the closure propagates outward; the lock gets
120 /// poisoned as a result, but the container recovers from poisoning and
121 /// later access works as usual.
122 ///
123 /// # Deadlock red line
124 ///
125 /// Same as [`with`](SharedState::with): the underlying lock is not
126 /// reentrant, and the closure must not access the same `SharedState`
127 /// again.
128 pub fn with_mut<T: Any>(&self, f: impl FnOnce(&mut T)) {
129 let mut guard = self.inner.write().unwrap_or_else(|e| e.into_inner());
130 if let Some(value) = guard
131 .get_mut(&TypeId::of::<T>())
132 .and_then(|v| v.downcast_mut::<T>())
133 {
134 f(value);
135 }
136 }
137}
138
139#[cfg(test)]
140mod tests {
141 use super::*;
142
143 #[derive(Debug, Clone, PartialEq)]
144 struct Session {
145 user: String,
146 }
147
148 #[derive(Debug, Clone, PartialEq)]
149 struct Counter(usize);
150
151 #[test]
152 fn insert_and_get_by_type() {
153 let state = SharedState::new();
154 state.insert(Session {
155 user: "alice".into(),
156 });
157 state.insert(Counter(3));
158
159 assert_eq!(
160 state.get::<Session>(),
161 Some(Session {
162 user: "alice".into()
163 })
164 );
165 assert_eq!(state.get::<Counter>(), Some(Counter(3)));
166 }
167
168 /// Multiple values coexist, distinguished by type; a second insert of
169 /// the same type overwrites.
170 #[test]
171 fn same_type_insert_overwrites_others_untouched() {
172 let state = SharedState::new();
173 state.insert(Session {
174 user: "alice".into(),
175 });
176 state.insert(Counter(1));
177 state.insert(Counter(2)); // same type overwrites
178
179 assert_eq!(state.get::<Counter>(), Some(Counter(2)));
180 assert_eq!(
181 state.get::<Session>(),
182 Some(Session {
183 user: "alice".into()
184 })
185 );
186 }
187
188 #[test]
189 fn get_missing_type_returns_none() {
190 let state = SharedState::new();
191 assert_eq!(state.get::<Session>(), None);
192 }
193
194 #[test]
195 fn with_and_with_mut_lock_inner() {
196 let state = SharedState::new();
197 state.insert(Counter(1));
198
199 // Missing type: the closure does not run.
200 let mut called = false;
201 state.with::<Session>(|_| called = true);
202 assert!(!called);
203
204 state.with::<Counter>(|c| {
205 assert_eq!(c.0, 1);
206 });
207 state.with_mut::<Counter>(|c| c.0 += 1);
208 assert_eq!(state.get::<Counter>(), Some(Counter(2)));
209 }
210
211 /// Clone sharing: two holders share one instance; writes from one are
212 /// visible to the other.
213 #[test]
214 fn clone_shares_same_instance() {
215 let state = SharedState::new();
216 let tool_side = state.clone();
217
218 state.insert(Counter(5));
219 assert_eq!(tool_side.get::<Counter>(), Some(Counter(5)));
220 tool_side.with_mut::<Counter>(|c| c.0 += 1);
221 assert_eq!(state.get::<Counter>(), Some(Counter(6)));
222 }
223
224 /// After a user closure panics and poisons the lock, the container
225 /// recovers from poisoning and later access works normally.
226 #[test]
227 fn poisoned_lock_recovers_and_stays_usable() {
228 let state = SharedState::new();
229 state.insert(Counter(1));
230
231 // Closure panics: the lock is poisoned.
232 let panicked = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
233 state.with_mut::<Counter>(|_| panic!("user closure panicked"));
234 }));
235 assert!(panicked.is_err());
236
237 // Recovered from poisoning: later reads and writes work as usual.
238 state.with_mut::<Counter>(|c| c.0 += 1);
239 assert_eq!(state.get::<Counter>(), Some(Counter(2)));
240 }
241}