agent_base/engine/
session_store.rs1use 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#[async_trait]
20pub trait SessionStore: Send + Sync {
21 async fn save(&self, session: &AgentSession) -> AgentResult<()>;
23
24 async fn load(&self, session_id: &SessionId) -> AgentResult<Option<AgentSession>>;
26
27 async fn list(&self) -> AgentResult<Vec<SessionId>>;
29
30 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}