systemprompt_cloud/cli_session/
store.rs

1//! On-disk index of CLI sessions keyed by tenant or `local`.
2//!
3//! Copyright (c) systemprompt.io — Business Source License 1.1.
4//! See <https://systemprompt.io> for licensing details.
5
6use std::collections::HashMap;
7use std::fs;
8use std::path::{Path, PathBuf};
9
10use chrono::{DateTime, Utc};
11use serde::{Deserialize, Serialize};
12use systemprompt_identifiers::TenantId;
13
14use super::private_file::{ensure_private_dir, write_private_atomic};
15use super::{CliSession, LOCAL_SESSION_KEY, SessionKey};
16use crate::error::{CloudError, CloudResult};
17
18const STORE_VERSION: u32 = 1;
19
20const fn default_store_version() -> u32 {
21    STORE_VERSION
22}
23
24#[derive(Debug, Clone, Serialize, Deserialize)]
25pub struct SessionStore {
26    #[serde(default = "default_store_version")]
27    pub version: u32,
28    #[serde(default)]
29    pub sessions: HashMap<String, CliSession>,
30    #[serde(default)]
31    pub active_key: Option<String>,
32    #[serde(default, skip_serializing_if = "Option::is_none")]
33    pub active_profile_name: Option<String>,
34    #[serde(default = "Utc::now")]
35    pub updated_at: DateTime<Utc>,
36}
37
38impl Default for SessionStore {
39    fn default() -> Self {
40        Self::new()
41    }
42}
43
44impl SessionStore {
45    #[must_use]
46    pub fn new() -> Self {
47        Self {
48            version: STORE_VERSION,
49            sessions: HashMap::new(),
50            active_key: None,
51            active_profile_name: None,
52            updated_at: Utc::now(),
53        }
54    }
55
56    #[must_use]
57    pub fn get_valid_session(&self, key: &SessionKey, issuer: &str) -> Option<&CliSession> {
58        self.sessions
59            .get(&key.as_storage_key())
60            .filter(|s| Self::is_usable(s, issuer))
61    }
62
63    pub fn get_valid_session_mut(
64        &mut self,
65        key: &SessionKey,
66        issuer: &str,
67    ) -> Option<&mut CliSession> {
68        self.sessions
69            .get_mut(&key.as_storage_key())
70            .filter(|s| Self::is_usable(s, issuer))
71    }
72
73    fn is_usable(session: &CliSession, issuer: &str) -> bool {
74        if !session.matches_issuer(issuer) {
75            tracing::info!(
76                stored_issuer = %session.issuer,
77                current_issuer = %issuer,
78                "Stored CLI session was minted under a different issuer; discarding it"
79            );
80            return false;
81        }
82        !session.is_expired() && session.has_valid_credentials()
83    }
84
85    #[must_use]
86    pub fn get_session(&self, key: &SessionKey) -> Option<&CliSession> {
87        self.sessions.get(&key.as_storage_key())
88    }
89
90    pub fn upsert_session(&mut self, key: &SessionKey, session: CliSession) {
91        self.sessions.insert(key.as_storage_key(), session);
92        self.updated_at = Utc::now();
93    }
94
95    pub fn remove_tenant_sessions(&mut self) -> usize {
96        let before = self.sessions.len();
97        self.sessions.retain(|key, _| key == LOCAL_SESSION_KEY);
98        let removed = before - self.sessions.len();
99        if removed > 0 {
100            if self
101                .active_key
102                .as_deref()
103                .is_some_and(|key| key != LOCAL_SESSION_KEY)
104            {
105                self.active_key = None;
106                self.active_profile_name = None;
107            }
108            self.updated_at = Utc::now();
109        }
110        removed
111    }
112
113    pub fn remove_session(&mut self, key: &SessionKey) -> Option<CliSession> {
114        let storage_key = key.as_storage_key();
115        let removed = self.sessions.remove(&storage_key);
116        if removed.is_some() {
117            self.updated_at = Utc::now();
118        }
119        removed
120    }
121
122    pub fn set_active(&mut self, key: &SessionKey) {
123        self.active_key = Some(key.as_storage_key());
124        self.updated_at = Utc::now();
125    }
126
127    pub fn set_active_with_profile(&mut self, key: &SessionKey, profile_name: &str) {
128        self.active_key = Some(key.as_storage_key());
129        self.active_profile_name = Some(profile_name.to_owned());
130        self.updated_at = Utc::now();
131    }
132
133    pub fn set_active_with_profile_path(
134        &mut self,
135        key: &SessionKey,
136        profile_name: &str,
137        profile_path: PathBuf,
138    ) {
139        self.active_key = Some(key.as_storage_key());
140        self.active_profile_name = Some(profile_name.to_owned());
141
142        if let Some(session) = self.sessions.get_mut(&key.as_storage_key()) {
143            session.update_profile_path(profile_path);
144        }
145
146        self.updated_at = Utc::now();
147    }
148
149    #[must_use]
150    pub fn active_session_key(&self) -> Option<SessionKey> {
151        self.active_key.as_ref().map(|k| {
152            if k == LOCAL_SESSION_KEY {
153                SessionKey::Local
154            } else {
155                k.strip_prefix("tenant_").map_or(SessionKey::Local, |id| {
156                    SessionKey::Tenant(TenantId::new(id))
157                })
158            }
159        })
160    }
161
162    #[must_use]
163    pub fn active_session_for_profile_discovery(&self) -> Option<&CliSession> {
164        self.active_session_key()
165            .and_then(|key| self.sessions.get(&key.as_storage_key()))
166            .filter(|s| !s.is_expired() && s.has_valid_credentials())
167    }
168
169    pub fn prune_expired(&mut self) -> usize {
170        let expired_keys: Vec<String> = self
171            .sessions
172            .iter()
173            .filter(|(_, s)| s.is_expired())
174            .map(|(k, _)| k.clone())
175            .collect();
176
177        let count = expired_keys.len();
178        for key in &expired_keys {
179            self.sessions.remove(key);
180        }
181
182        if count > 0 {
183            self.updated_at = Utc::now();
184        }
185        count
186    }
187
188    #[must_use]
189    pub fn find_by_profile_name(&self, name: &str) -> Option<&CliSession> {
190        self.sessions
191            .values()
192            .find(|s| s.profile_name.as_str() == name && !s.is_expired())
193    }
194
195    #[must_use]
196    pub fn all_sessions(&self) -> Vec<(&String, &CliSession)> {
197        self.sessions.iter().collect()
198    }
199
200    #[must_use]
201    pub fn len(&self) -> usize {
202        self.sessions.len()
203    }
204
205    #[must_use]
206    pub fn is_empty(&self) -> bool {
207        self.sessions.is_empty()
208    }
209
210    pub fn load(sessions_dir: &Path) -> CloudResult<Option<Self>> {
211        let index_path = sessions_dir.join("index.json");
212        let content = match fs::read_to_string(&index_path) {
213            Ok(c) => c,
214            Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(None),
215            Err(e) => return Err(e.into()),
216        };
217        let store: Self =
218            serde_json::from_str(&content).map_err(|e| CloudError::SessionStoreCorrupted {
219                path: index_path.display().to_string(),
220                source: e,
221            })?;
222        if store.version > STORE_VERSION {
223            return Err(CloudError::SessionVersionMismatch {
224                min: STORE_VERSION,
225                max: STORE_VERSION,
226                actual: store.version,
227                path: index_path.display().to_string(),
228            });
229        }
230        Ok(Some(store))
231    }
232
233    pub fn load_or_create(sessions_dir: &Path) -> CloudResult<Self> {
234        Ok(Self::load(sessions_dir)?.unwrap_or_else(Self::new))
235    }
236
237    pub fn save(&self, sessions_dir: &Path) -> CloudResult<()> {
238        ensure_private_dir(sessions_dir)?;
239        let content = serde_json::to_string_pretty(self)?;
240        write_private_atomic(&sessions_dir.join("index.json"), &content)
241    }
242}
systemprompt_cloud/cli_session/store.rs

systemprompt_cloud/cli_session/
store.rs
