Skip to main content

agent_base/engine/
session_store.rs

1use std::collections::HashMap;
2
3use async_trait::async_trait;
4use tokio::sync::Mutex;
5
6use crate::types::{AgentResult, AgentError, SessionId};
7use super::AgentSession;
8
9/// Session Persistence Adapter
10///
11/// `SessionStore` is an optional persistence interface for agent sessions.
12/// Under the lightweight kernel design:
13/// - `AgentRuntime.sessions` is the authoritative live state during execution
14/// - `SessionStore` is a persistence adapter for save/load/list/delete
15/// - Does not participate in the execution control flow
16///
17/// Replace the default [`InMemorySessionStore`] with a custom implementation
18/// to persist sessions to a database, filesystem, or other storage.
19#[async_trait]
20pub trait SessionStore: Send + Sync {
21    /// Save a session snapshot to the persistence layer
22    async fn save(&self, session: &AgentSession) -> AgentResult<()>;
23
24    /// Load a session from the persistence layer
25    async fn load(&self, session_id: &SessionId) -> AgentResult<Option<AgentSession>>;
26
27    /// List all saved session IDs
28    async fn list(&self) -> AgentResult<Vec<SessionId>>;
29
30    /// Delete a specific session
31    async fn delete(&self, session_id: &SessionId) -> AgentResult<()>;
32}
33
34pub struct InMemorySessionStore {
35    sessions: Mutex<HashMap<SessionId, AgentSession>>,
36}
37
38impl InMemorySessionStore {
39    pub fn new() -> Self {
40        Self {
41            sessions: Mutex::new(HashMap::new()),
42        }
43    }
44}
45
46impl Default for InMemorySessionStore {
47    fn default() -> Self {
48        Self::new()
49    }
50}
51
52#[async_trait]
53impl SessionStore for InMemorySessionStore {
54    async fn save(&self, session: &AgentSession) -> AgentResult<()> {
55        let session_id = session
56            .id()
57            .ok_or_else(|| AgentError::internal("session has no id"))?;
58        self.sessions
59            .lock()
60            .await
61            .insert(session_id, session.clone());
62        Ok(())
63    }
64
65    async fn load(&self, session_id: &SessionId) -> AgentResult<Option<AgentSession>> {
66        Ok(self.sessions.lock().await.get(session_id).cloned())
67    }
68
69    async fn list(&self) -> AgentResult<Vec<SessionId>> {
70        Ok(self.sessions.lock().await.keys().cloned().collect())
71    }
72
73    async fn delete(&self, session_id: &SessionId) -> AgentResult<()> {
74        self.sessions.lock().await.remove(session_id);
75        Ok(())
76    }
77}