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