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

systemprompt_cloud/cli_session/
store.rs
