agent_base/engine/
session_store.rs1use 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#[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 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}