oxicode_sdk/lifecycle/
agent_pool.rs

1//! Agent pool — live agent instance tracking for session persistence.
2//!
3//! Stores running `Agent` instances keyed by ID, enabling:
4//! - Session continuation after agent completion
5//! - State export/import for persistence
6//! - Agent lookup by ID
7
8use oxicode_agent::Agent;
9use parking_lot::RwLock;
10use std::collections::HashMap;
11use std::sync::Arc;
12
13/// Pool of live `Agent` instances keyed by string ID.
14///
15/// Thread-safe via `parking_lot::RwLock`. Agents can be inserted,
16/// retrieved, removed, and have their state exported/imported.
17pub struct AgentPool {
18    agents: RwLock<HashMap<String, Arc<Agent>>>,
19}
20
21impl AgentPool {
22    /// Create an empty pool.
23    pub fn new() -> Self {
24        Self {
25            agents: RwLock::new(HashMap::new()),
26        }
27    }
28
29    /// Insert an agent into the pool.
30    pub fn insert(&self, id: String, agent: Arc<Agent>) {
31        self.agents.write().insert(id, agent);
32    }
33
34    /// Get an agent by ID.
35    pub fn get(&self, id: &str) -> Option<Arc<Agent>> {
36        self.agents.read().get(id).cloned()
37    }
38
39    /// Remove an agent from the pool.
40    pub fn remove(&self, id: &str) -> Option<Arc<Agent>> {
41        self.agents.write().remove(id)
42    }
43
44    /// Export an agent's state as JSON.
45    ///
46    /// Returns `None` if the agent is not in the pool or state
47    /// serialization fails.
48    pub fn export_state(&self, id: &str) -> Option<serde_json::Value> {
49        let agents = self.agents.read();
50        let agent = agents.get(id)?;
51        agent.export_state().ok()
52    }
53
54    /// Import agent state from JSON.
55    ///
56    /// Returns `false` if the agent is not in the pool or import fails.
57    pub fn import_state(&self, id: &str, state: serde_json::Value) -> bool {
58        let agents = self.agents.read();
59        if let Some(agent) = agents.get(id) {
60            agent.import_state(state).is_ok()
61        } else {
62            false
63        }
64    }
65
66    /// Number of agents in the pool.
67    pub fn len(&self) -> usize {
68        self.agents.read().len()
69    }
70
71    /// Whether the pool is empty.
72    pub fn is_empty(&self) -> bool {
73        self.agents.read().is_empty()
74    }
75
76    /// List all agent IDs in the pool.
77    pub fn ids(&self) -> Vec<String> {
78        self.agents.read().keys().cloned().collect()
79    }
80
81    /// Check if an agent exists in the pool.
82    pub fn contains(&self, id: &str) -> bool {
83        self.agents.read().contains_key(id)
84    }
85
86    /// Snapshot iteration over all (id, agent) pairs. Holds the read lock for
87    /// the duration of the closure; do not call back into the pool from `f`.
88    pub fn for_each_row<F: FnMut(&str, &Arc<Agent>)>(&self, mut f: F) {
89        let agents = self.agents.read();
90        for (id, agent) in agents.iter() {
91            f(id.as_str(), agent);
92        }
93    }
94}
95
96impl Default for AgentPool {
97    fn default() -> Self {
98        Self::new()
99    }
100}
101
102impl std::fmt::Debug for AgentPool {
103    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
104        f.debug_struct("AgentPool")
105            .field("count", &self.len())
106            .finish()
107    }
108}
109
110#[cfg(test)]
111mod tests {
112    use super::*;
113
114    // `make_agent` was removed — it used `std::mem::zeroed::<Agent>()` which
115    // is UB (Agent contains non-zero fields) and was never called by any test.
116    // Re-add a proper Agent constructor here if pool tests are needed.
117    // For a safe version, we test the pool with just the data structure operations.
118
119    #[test]
120    fn test_pool_new_empty() {
121        let pool = AgentPool::new();
122        assert!(pool.is_empty());
123        assert_eq!(pool.len(), 0);
124    }
125
126    #[test]
127    fn for_each_row_visits_all_inserted() {
128        let pool = AgentPool::new();
129        let mut seen = Vec::new();
130        pool.for_each_row(|id, _| seen.push(id.to_string()));
131        assert!(seen.is_empty());
132    }
133
134    #[test]
135    fn test_pool_default() {
136        let pool = AgentPool::default();
137        assert!(pool.is_empty());
138    }
139
140    #[test]
141    fn test_pool_debug() {
142        let pool = AgentPool::new();
143        let debug = format!("{:?}", pool);
144        assert!(debug.contains("AgentPool"));
145    }
146}
oxicode_sdk/lifecycle/agent_pool.rs

oxicode_sdk/lifecycle/
agent_pool.rs
