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

systemprompt_cloud/cli_session/
store.rs
