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 super::AgentSession;
7use crate::types::{AgentError, AgentResult, RuntimeEvent, SessionId};
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    /// Optionally persist a runtime event for audit/replay.
34    /// Default implementation is a no-op — override to implement
35    /// append-log style persistence (e.g. JSONL file, database event log).
36    async fn append_event(
37        &self,
38        _session_id: &SessionId,
39        _event: &RuntimeEvent,
40    ) -> AgentResult<()> {
41        Ok(())
42    }
43}
44
45pub struct InMemorySessionStore {
46    sessions: Mutex<HashMap<SessionId, AgentSession>>,
47}
48
49impl InMemorySessionStore {
50    pub fn new() -> Self {
51        Self {
52            sessions: Mutex::new(HashMap::new()),
53        }
54    }
55}
56
57impl Default for InMemorySessionStore {
58    fn default() -> Self {
59        Self::new()
60    }
61}
62
63#[async_trait]
64impl SessionStore for InMemorySessionStore {
65    async fn save(&self, session: &AgentSession) -> AgentResult<()> {
66        let session_id = session
67            .id()
68            .ok_or_else(|| AgentError::internal("session has no id"))?;
69        self.sessions
70            .lock()
71            .await
72            .insert(session_id, session.clone());
73        Ok(())
74    }
75
76    async fn load(&self, session_id: &SessionId) -> AgentResult<Option<AgentSession>> {
77        Ok(self.sessions.lock().await.get(session_id).cloned())
78    }
79
80    async fn list(&self) -> AgentResult<Vec<SessionId>> {
81        Ok(self.sessions.lock().await.keys().cloned().collect())
82    }
83
84    async fn delete(&self, session_id: &SessionId) -> AgentResult<()> {
85        self.sessions.lock().await.remove(session_id);
86        Ok(())
87    }
88}