starweaver-session 0.0.3

Durable session contracts for Starweaver
Documentation
use chrono::Utc;
use starweaver_context::ResumableState;
use starweaver_core::SessionId;

use crate::{
    error::{SessionStoreError, SessionStoreResult},
    records::{EnvironmentStateRef, SessionRecord, SessionStatus},
    store::SessionFilter,
};

use super::{store_failed, InMemorySessionStore};

impl InMemorySessionStore {
    pub(super) fn save_session_record(&self, mut session: SessionRecord) -> SessionStoreResult<()> {
        let mut inner = self.inner.lock().map_err(store_failed)?;
        session.updated_at = Utc::now();
        inner.sessions.insert(session.session_id.clone(), session);
        Ok(())
    }

    pub(super) fn load_session_record(
        &self,
        session_id: &SessionId,
    ) -> SessionStoreResult<SessionRecord> {
        let inner = self.inner.lock().map_err(store_failed)?;
        inner
            .sessions
            .get(session_id)
            .cloned()
            .ok_or_else(|| SessionStoreError::NotFound(session_id.as_str().to_string()))
    }

    #[allow(clippy::needless_pass_by_value)]
    pub(super) fn list_session_records(
        &self,
        filter: SessionFilter,
    ) -> SessionStoreResult<Vec<SessionRecord>> {
        let inner = self.inner.lock().map_err(store_failed)?;
        let mut sessions = inner
            .sessions
            .values()
            .filter(|session| {
                filter
                    .status
                    .map_or(true, |status| session.status == status)
            })
            .filter(|session| {
                filter
                    .profile
                    .as_ref()
                    .map_or(true, |profile| session.profile.as_ref() == Some(profile))
            })
            .filter(|session| {
                filter.workspace.as_ref().map_or(true, |workspace| {
                    session.workspace.as_ref() == Some(workspace)
                })
            })
            .cloned()
            .collect::<Vec<_>>();
        sessions.sort_by_key(|session| std::cmp::Reverse(session.updated_at));
        if let Some(limit) = filter.limit {
            sessions.truncate(limit);
        }
        Ok(sessions)
    }

    pub(super) fn set_session_status(
        &self,
        session_id: &SessionId,
        status: SessionStatus,
    ) -> SessionStoreResult<()> {
        let mut inner = self.inner.lock().map_err(store_failed)?;
        let session = inner
            .sessions
            .get_mut(session_id)
            .ok_or_else(|| SessionStoreError::NotFound(session_id.as_str().to_string()))?;
        session.status = status;
        session.updated_at = Utc::now();
        Ok(())
    }

    pub(super) fn save_context_state_snapshot(
        &self,
        session_id: &SessionId,
        state: ResumableState,
    ) -> SessionStoreResult<()> {
        let mut inner = self.inner.lock().map_err(store_failed)?;
        let session = inner
            .sessions
            .get_mut(session_id)
            .ok_or_else(|| SessionStoreError::NotFound(session_id.as_str().to_string()))?;
        session.state = state;
        session.updated_at = Utc::now();
        Ok(())
    }

    pub(super) fn save_environment_state_ref(
        &self,
        session_id: &SessionId,
        environment_state: EnvironmentStateRef,
    ) -> SessionStoreResult<()> {
        let mut inner = self.inner.lock().map_err(store_failed)?;
        let session = inner
            .sessions
            .get_mut(session_id)
            .ok_or_else(|| SessionStoreError::NotFound(session_id.as_str().to_string()))?;
        session.environment_state = Some(environment_state);
        session.updated_at = Utc::now();
        Ok(())
    }
}