codei_sdk/
session.rs

1use std::collections::HashMap;
2use std::path::PathBuf;
3use std::sync::{Arc, RwLock};
4
5use codei_session::{Session, SessionError, SessionStore};
6use tokio::sync::{Mutex, RwLock as AsyncRwLock};
7
8use crate::error::SdkError;
9use crate::runtime::{build_agent_runtime, open_session_store, AgentRuntime};
10
11const LIST_LIMIT: usize = 200;
12
13/// Persistent session registry with in-memory agent runtime cache.
14pub struct SessionService {
15    store: Arc<SessionStore>,
16    sessions: RwLock<HashMap<String, Arc<SessionHandle>>>,
17}
18
19pub struct SessionHandle {
20    pub runtime: AgentRuntime,
21    pub session: Arc<AsyncRwLock<Session>>,
22    pub turn_lock: Mutex<()>,
23}
24
25impl SessionService {
26    pub async fn new(default_cwd: PathBuf) -> Result<Self, SdkError> {
27        let store = open_session_store(default_cwd).await?;
28        Ok(Self {
29            store,
30            sessions: RwLock::new(HashMap::new()),
31        })
32    }
33
34    pub fn store(&self) -> &Arc<SessionStore> {
35        &self.store
36    }
37
38    pub fn list_sessions(&self) -> Result<Vec<Session>, SessionError> {
39        let mut sessions = self.store.list(LIST_LIMIT)?;
40        let guard = self.sessions.read().expect("sessions lock poisoned");
41        for handle in guard.values() {
42            if let Ok(session) = handle.session.try_read() {
43                if sessions.iter().any(|s| s.id == session.id) {
44                    continue;
45                }
46                sessions.push(session.clone());
47            }
48        }
49        sessions.sort_by_key(|b| std::cmp::Reverse(b.updated_at));
50        sessions.truncate(LIST_LIMIT);
51        Ok(sessions)
52    }
53
54    pub fn message_count(&self, session: &Session) -> usize {
55        if !session.messages.is_empty() {
56            return session.messages.len();
57        }
58        self.store
59            .load(&session.id)
60            .map(|s| s.messages.len())
61            .unwrap_or(0)
62    }
63
64    pub async fn create_session(&self, cwd: PathBuf) -> Result<Session, SdkError> {
65        if !cwd.is_dir() {
66            return Err(SdkError::Other(anyhow::anyhow!(
67                "working directory does not exist: {}",
68                cwd.display()
69            )));
70        }
71
72        let session = Session::new(cwd);
73        self.store.save(&session)?;
74
75        let handle = self.hydrate(session.clone()).await?;
76        self.sessions
77            .write()
78            .expect("sessions lock poisoned")
79            .insert(session.id.clone(), handle);
80
81        Ok(session)
82    }
83
84    pub async fn get_or_load(&self, id: &str) -> Result<Arc<SessionHandle>, SdkError> {
85        if let Some(handle) = self.get(id) {
86            return Ok(handle);
87        }
88
89        let session = self.store.load(id).map_err(|e| match e {
90            SessionError::NotFound(_) => SdkError::Other(anyhow::anyhow!("session not found")),
91            other => SdkError::Session(other),
92        })?;
93
94        let handle = self.hydrate(session).await?;
95        self.sessions
96            .write()
97            .expect("sessions lock poisoned")
98            .insert(id.to_string(), Arc::clone(&handle));
99        Ok(handle)
100    }
101
102    async fn hydrate(&self, session: Session) -> Result<Arc<SessionHandle>, SdkError> {
103        let config = Arc::new(codei_config::load(&codei_config::LoadOptions {
104            cwd: Some(session.cwd.clone()),
105            ..Default::default()
106        })?);
107        let runtime = build_agent_runtime(config, Arc::clone(&self.store)).await?;
108        Ok(Arc::new(SessionHandle {
109            runtime,
110            session: Arc::new(AsyncRwLock::new(session)),
111            turn_lock: Mutex::new(()),
112        }))
113    }
114
115    fn get(&self, id: &str) -> Option<Arc<SessionHandle>> {
116        let guard = self.sessions.read().ok()?;
117        guard.get(id).cloned()
118    }
119}
codei_sdk/session.rs

codei_sdk/
session.rs
