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, RuntimeEvent, 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    /// 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(&self, _session_id: &SessionId, _event: &RuntimeEvent) -> AgentResult<()> {
37        Ok(())
38    }
39}
40
41pub struct InMemorySessionStore {
42    sessions: Mutex<HashMap<SessionId, AgentSession>>,
43}
44
45impl InMemorySessionStore {
46    pub fn new() -> Self {
47        Self {
48            sessions: Mutex::new(HashMap::new()),
49        }
50    }
51}
52
53impl Default for InMemorySessionStore {
54    fn default() -> Self {
55        Self::new()
56    }
57}
58
59#[async_trait]
60impl SessionStore for InMemorySessionStore {
61    async fn save(&self, session: &AgentSession) -> AgentResult<()> {
62        let session_id = session
63            .id()
64            .ok_or_else(|| AgentError::internal("session has no id"))?;
65        self.sessions
66            .lock()
67            .await
68            .insert(session_id, session.clone());
69        Ok(())
70    }
71
72    async fn load(&self, session_id: &SessionId) -> AgentResult<Option<AgentSession>> {
73        Ok(self.sessions.lock().await.get(session_id).cloned())
74    }
75
76    async fn list(&self) -> AgentResult<Vec<SessionId>> {
77        Ok(self.sessions.lock().await.keys().cloned().collect())
78    }
79
80    async fn delete(&self, session_id: &SessionId) -> AgentResult<()> {
81        self.sessions.lock().await.remove(session_id);
82        Ok(())
83    }
84}