Skip to main content

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    /// The active session without an issuer check, for resolving which profile
162    /// to load. Never use it to authorize a request — the issuer is unknown
163    /// until that profile is read, so the result may carry a stale token.
164    #[must_use]
165    pub fn active_session_for_profile_discovery(&self) -> Option<&CliSession> {
166        self.active_session_key()
167            .and_then(|key| self.sessions.get(&key.as_storage_key()))
168            .filter(|s| !s.is_expired() && s.has_valid_credentials())
169    }
170
171    pub fn prune_expired(&mut self) -> usize {
172        let expired_keys: Vec<String> = self
173            .sessions
174            .iter()
175            .filter(|(_, s)| s.is_expired())
176            .map(|(k, _)| k.clone())
177            .collect();
178
179        let count = expired_keys.len();
180        for key in &expired_keys {
181            self.sessions.remove(key);
182        }
183
184        if count > 0 {
185            self.updated_at = Utc::now();
186        }
187        count
188    }
189
190    #[must_use]
191    pub fn find_by_profile_name(&self, name: &str) -> Option<&CliSession> {
192        self.sessions
193            .values()
194            .find(|s| s.profile_name.as_str() == name && !s.is_expired())
195    }
196
197    #[must_use]
198    pub fn all_sessions(&self) -> Vec<(&String, &CliSession)> {
199        self.sessions.iter().collect()
200    }
201
202    #[must_use]
203    pub fn len(&self) -> usize {
204        self.sessions.len()
205    }
206
207    #[must_use]
208    pub fn is_empty(&self) -> bool {
209        self.sessions.is_empty()
210    }
211
212    /// `Ok(None)` means the store file does not exist; a store that exists but
213    /// cannot be parsed is an error, never silently discarded — degrading a
214    /// stale store into "no store" sends resolution down ambiguous profile
215    /// discovery and masks the real fault.
216    pub fn load(sessions_dir: &Path) -> CloudResult<Option<Self>> {
217        let index_path = sessions_dir.join("index.json");
218        let content = match fs::read_to_string(&index_path) {
219            Ok(c) => c,
220            Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(None),
221            Err(e) => return Err(e.into()),
222        };
223        let store: Self =
224            serde_json::from_str(&content).map_err(|e| CloudError::SessionStoreCorrupted {
225                path: index_path.display().to_string(),
226                source: e,
227            })?;
228        if store.version > STORE_VERSION {
229            return Err(CloudError::SessionVersionMismatch {
230                min: STORE_VERSION,
231                max: STORE_VERSION,
232                actual: store.version,
233                path: index_path.display().to_string(),
234            });
235        }
236        Ok(Some(store))
237    }
238
239    pub fn load_or_create(sessions_dir: &Path) -> CloudResult<Self> {
240        Ok(Self::load(sessions_dir)?.unwrap_or_else(Self::new))
241    }
242
243    /// Lenient variant for recovery commands (`admin session switch`, login)
244    /// that overwrite the store anyway — a corrupt store must not block its
245    /// own reset.
246    #[must_use]
247    pub fn load_or_reset(sessions_dir: &Path) -> Self {
248        match Self::load(sessions_dir) {
249            Ok(store) => store.unwrap_or_else(Self::new),
250            Err(e) => {
251                tracing::warn!(error = %e, "Resetting unreadable session store");
252                Self::new()
253            },
254        }
255    }
256
257    pub fn save(&self, sessions_dir: &Path) -> CloudResult<()> {
258        fs::create_dir_all(sessions_dir)?;
259
260        let gitignore_path = sessions_dir.join(".gitignore");
261        if !gitignore_path.exists() {
262            fs::write(&gitignore_path, "*\n")?;
263        }
264
265        let index_path = sessions_dir.join("index.json");
266        let content = serde_json::to_string_pretty(self)?;
267        let temp_path = index_path.with_extension("tmp");
268        fs::write(&temp_path, &content)?;
269
270        #[cfg(unix)]
271        {
272            use std::os::unix::fs::PermissionsExt;
273            let mut perms = fs::metadata(&temp_path)?.permissions();
274            perms.set_mode(0o600);
275            fs::set_permissions(&temp_path, perms)?;
276        }
277
278        fs::rename(&temp_path, &index_path)?;
279        Ok(())
280    }
281}