strands_agents/agent/
state.rs

1//! Agent state management.
2
3use std::collections::HashMap;
4
5use serde::{Deserialize, Serialize};
6
7/// Persistent state for an agent.
8#[derive(Debug, Clone, Default, Serialize, Deserialize)]
9pub struct AgentState {
10    data: HashMap<String, serde_json::Value>,
11}
12
13impl AgentState {
14    pub fn new() -> Self { Self { data: HashMap::new() } }
15
16    pub fn from_map(data: HashMap<String, serde_json::Value>) -> Self { Self { data } }
17
18    /// Gets a value by key, deserializing to the requested type.
19    pub fn get<T: serde::de::DeserializeOwned>(&self, key: &str) -> Option<T> {
20        self.data.get(key).and_then(|v| T::deserialize(v).ok())
21    }
22
23    /// Gets the raw JSON value by key.
24    pub fn get_raw(&self, key: &str) -> Option<&serde_json::Value> { self.data.get(key) }
25
26    /// Sets a value by key.
27    pub fn set<T: Serialize>(&mut self, key: impl Into<String>, value: T) {
28        if let Ok(v) = serde_json::to_value(value) {
29            self.data.insert(key.into(), v);
30        }
31    }
32
33    pub fn remove(&mut self, key: &str) -> Option<serde_json::Value> { self.data.remove(key) }
34    pub fn contains(&self, key: &str) -> bool { self.data.contains_key(key) }
35    pub fn keys(&self) -> impl Iterator<Item = &str> { self.data.keys().map(|s| s.as_str()) }
36    pub fn clear(&mut self) { self.data.clear(); }
37    pub fn len(&self) -> usize { self.data.len() }
38    pub fn is_empty(&self) -> bool { self.data.is_empty() }
39    pub fn into_inner(self) -> HashMap<String, serde_json::Value> { self.data }
40
41    /// Returns a clone of all state data.
42    pub fn get_all(&self) -> HashMap<String, serde_json::Value> { self.data.clone() }
43
44    /// Merges another state into this one.
45    pub fn merge(&mut self, other: AgentState) { self.data.extend(other.data); }
46}
47
48#[cfg(test)]
49mod tests {
50    use super::*;
51
52    #[test]
53    fn test_state_basic() {
54        let mut state = AgentState::new();
55        state.set("count", 42);
56        state.set("name", "test");
57        assert_eq!(state.get::<i32>("count"), Some(42));
58        assert_eq!(state.get::<String>("name"), Some("test".to_string()));
59        assert_eq!(state.get::<i32>("missing"), None);
60    }
61
62    #[test]
63    fn test_state_remove() {
64        let mut state = AgentState::new();
65        state.set("key", "value");
66        assert!(state.contains("key"));
67        state.remove("key");
68        assert!(!state.contains("key"));
69    }
70
71    #[test]
72    fn test_state_merge() {
73        let mut state1 = AgentState::new();
74        state1.set("a", 1);
75        let mut state2 = AgentState::new();
76        state2.set("b", 2);
77        state1.merge(state2);
78        assert_eq!(state1.get::<i32>("a"), Some(1));
79        assert_eq!(state1.get::<i32>("b"), Some(2));
80    }
81}
strands_agents/agent/state.rs

strands_agents/agent/
state.rs
