Skip to main content

magi_code/sessions/
manager.rs

1use super::metadata::{SessionMetadataRecord, SessionMetadataSummary};
2use super::read::validate_session_id;
3use super::store::{prepare_session_root, primary_path};
4use crate::persistence::CrossProcessFileLock;
5use std::{
6    collections::HashMap,
7    fmt,
8    path::{Path, PathBuf},
9    sync::{Arc, Mutex, OnceLock, Weak},
10};
11use uuid::Uuid;
12
13#[derive(Debug, Clone, PartialEq, Eq)]
14pub(crate) struct SessionInternalDiagnostic {
15    pub(crate) session_id: Option<String>,
16    pub(crate) message: String,
17}
18
19#[derive(Debug, Clone, PartialEq, Eq)]
20pub(crate) struct SessionListReport {
21    pub(crate) summaries: Vec<SessionMetadataSummary>,
22    pub(crate) diagnostics: Vec<SessionInternalDiagnostic>,
23}
24
25#[derive(Debug, Clone)]
26pub struct SessionManager {
27    pub(in crate::sessions) root: PathBuf,
28}
29
30impl SessionManager {
31    #[must_use]
32    pub fn new(root: PathBuf) -> Self {
33        Self { root }
34    }
35
36    pub(crate) fn create(&self) -> anyhow::Result<Session> {
37        prepare_session_root(&self.root)?;
38        let id = Uuid::new_v4().to_string();
39        let path = self.path_for_valid_id(&id)?;
40        // Do not create JSONL until first event append; avoids orphan session files.
41        Ok(Session::new(id, path))
42    }
43
44    pub fn open(&self, id: impl Into<String>) -> anyhow::Result<Session> {
45        let id = validate_session_id(id.into())?;
46        Ok(Session::new(id.clone(), self.path_for_valid_id(&id)?))
47    }
48
49    pub(crate) fn open_existing(&self, id: impl Into<String>) -> anyhow::Result<Session> {
50        let session = self.open(id)?;
51        super::store::open_existing_primary(&self.root, &session.id)?
52            .ok_or_else(|| anyhow::anyhow!("session JSONL is missing"))?;
53        Ok(session)
54    }
55
56    #[cfg(test)]
57    pub(crate) fn list(&self) -> anyhow::Result<Vec<Session>> {
58        Ok(self
59            .list_metadata_summaries()?
60            .into_iter()
61            .map(|summary| summary.session)
62            .collect())
63    }
64
65    pub(crate) fn most_recent(&self) -> anyhow::Result<Option<Session>> {
66        let report = self.list_metadata_report()?;
67        if !report.diagnostics.is_empty() {
68            anyhow::bail!("session discovery found unreadable or unsafe history");
69        }
70        Ok(report
71            .summaries
72            .into_iter()
73            .last()
74            .map(|summary| summary.session))
75    }
76
77    pub(crate) fn path_for_valid_id(&self, id: &str) -> anyhow::Result<PathBuf> {
78        primary_path(&self.root, id)
79    }
80}
81
82pub struct Session {
83    pub(in crate::sessions) id: String,
84    pub(in crate::sessions) path: PathBuf,
85    pub(in crate::sessions) metadata_cache: Arc<Mutex<Option<SessionMetadataRecord>>>,
86    replay_state: Arc<SessionReplayState>,
87    active_lease: Arc<Mutex<Option<Arc<ActiveLeaseSlot>>>>,
88    // Shared only by clones of this admitted standalone owner, never by independent opens.
89    standalone_writer: Option<Arc<super::SessionWriterLease>>,
90}
91
92struct SessionReplayState {
93    generation: std::sync::atomic::AtomicU64,
94    last_terminal_status_generation: std::sync::atomic::AtomicU64,
95    last_compaction_generation: std::sync::atomic::AtomicU64,
96}
97
98type ActiveLeaseSlot = Mutex<Option<Arc<CrossProcessFileLock>>>;
99
100static ACTIVE_SESSION_LEASES: OnceLock<Mutex<HashMap<PathBuf, Weak<ActiveLeaseSlot>>>> =
101    OnceLock::new();
102static SESSION_REPLAY_GENERATIONS: OnceLock<Mutex<HashMap<PathBuf, Weak<SessionReplayState>>>> =
103    OnceLock::new();
104
105impl fmt::Debug for Session {
106    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
107        formatter
108            .debug_struct("Session")
109            .field("id", &self.id)
110            .field("path", &self.path)
111            .finish()
112    }
113}
114
115impl Clone for Session {
116    fn clone(&self) -> Self {
117        Self {
118            id: self.id.clone(),
119            path: self.path.clone(),
120            metadata_cache: Arc::clone(&self.metadata_cache),
121            replay_state: Arc::clone(&self.replay_state),
122            active_lease: Arc::clone(&self.active_lease),
123            standalone_writer: self.standalone_writer.clone(),
124        }
125    }
126}
127
128impl PartialEq for Session {
129    fn eq(&self, other: &Self) -> bool {
130        self.id == other.id && self.path == other.path
131    }
132}
133
134impl Eq for Session {}
135
136impl Session {
137    pub(in crate::sessions) fn new(id: String, path: PathBuf) -> Self {
138        let replay_state = replay_state_for_path(&path);
139        Self {
140            id,
141            path,
142            metadata_cache: Arc::new(Mutex::new(None)),
143            replay_state,
144            active_lease: Arc::new(Mutex::new(None)),
145            standalone_writer: None,
146        }
147    }
148
149    pub(crate) fn admit_standalone_writer(mut self) -> anyhow::Result<Self> {
150        if self.standalone_writer.is_none() {
151            let writer = self.try_frontend_writer()?.ok_or_else(|| {
152                anyhow::anyhow!("session is busy: another frontend owns the writer")
153            })?;
154            self.standalone_writer = Some(Arc::new(writer));
155        }
156        self.activate()
157    }
158
159    pub(crate) fn activate(self) -> anyhow::Result<Self> {
160        self.ensure_active_lease()?;
161        Ok(self)
162    }
163
164    pub(in crate::sessions) fn ensure_active_lease(&self) -> anyhow::Result<()> {
165        let slot = {
166            let mut local_slot = self
167                .active_lease
168                .lock()
169                .map_err(|_| anyhow::anyhow!("active session lease slot was poisoned"))?;
170            if let Some(slot) = local_slot.as_ref() {
171                Arc::clone(slot)
172            } else {
173                let slot = active_lease_slot(&self.path)?;
174                *local_slot = Some(Arc::clone(&slot));
175                slot
176            }
177        };
178        let mut lease = slot
179            .lock()
180            .map_err(|_| anyhow::anyhow!("active session lease was poisoned"))?;
181        if lease.is_none() {
182            let target = active_lease_target(&self.path);
183            *lease = Some(Arc::new(
184                CrossProcessFileLock::try_acquire(&target)?.ok_or_else(|| {
185                    anyhow::anyhow!("session is active in another process: {}", self.id)
186                })?,
187            ));
188        }
189        Ok(())
190    }
191
192    pub(in crate::sessions) fn metadata_cache(&self) -> &Mutex<Option<SessionMetadataRecord>> {
193        &self.metadata_cache
194    }
195}
196
197fn replay_state_for_path(path: &Path) -> Arc<SessionReplayState> {
198    let key = normalize_active_lease_path(path);
199    let registry = SESSION_REPLAY_GENERATIONS.get_or_init(|| Mutex::new(HashMap::new()));
200    let Ok(mut states) = registry.lock() else {
201        // A poisoned optimization registry must not make durable sessions unavailable.
202        return Arc::new(SessionReplayState {
203            generation: std::sync::atomic::AtomicU64::new(0),
204            last_terminal_status_generation: std::sync::atomic::AtomicU64::new(0),
205            last_compaction_generation: std::sync::atomic::AtomicU64::new(0),
206        });
207    };
208    if let Some(state) = states.get(&key).and_then(Weak::upgrade) {
209        return state;
210    }
211    states.retain(|_, state| state.strong_count() > 0);
212    let state = Arc::new(SessionReplayState {
213        generation: std::sync::atomic::AtomicU64::new(0),
214        last_terminal_status_generation: std::sync::atomic::AtomicU64::new(0),
215        last_compaction_generation: std::sync::atomic::AtomicU64::new(0),
216    });
217    states.insert(key, Arc::downgrade(&state));
218    state
219}
220
221fn active_lease_slot(path: &Path) -> anyhow::Result<Arc<ActiveLeaseSlot>> {
222    let key = normalize_active_lease_path(path);
223    let registry = ACTIVE_SESSION_LEASES.get_or_init(|| Mutex::new(HashMap::new()));
224    let mut slots = registry
225        .lock()
226        .map_err(|_| anyhow::anyhow!("active session lease registry was poisoned"))?;
227    if let Some(slot) = slots.get(&key).and_then(Weak::upgrade) {
228        return Ok(slot);
229    }
230    slots.retain(|_, slot| slot.strong_count() > 0);
231    let slot = Arc::new(Mutex::new(None));
232    slots.insert(key, Arc::downgrade(&slot));
233    Ok(slot)
234}
235
236fn normalize_active_lease_path(path: &Path) -> PathBuf {
237    if let Ok(canonical) = path.canonicalize() {
238        return canonical;
239    }
240    if let (Some(parent), Some(file_name)) = (path.parent(), path.file_name())
241        && let Ok(parent) = parent.canonicalize()
242    {
243        return parent.join(file_name);
244    }
245    path.to_path_buf()
246}
247
248pub(in crate::sessions) fn active_lease_target(path: &Path) -> PathBuf {
249    path.with_extension("active")
250}
251
252impl Session {
253    pub fn id(&self) -> &str {
254        &self.id
255    }
256
257    pub(crate) fn path(&self) -> &Path {
258        &self.path
259    }
260
261    pub(crate) fn replay_generation(&self) -> u64 {
262        self.replay_state
263            .generation
264            .load(std::sync::atomic::Ordering::Acquire)
265    }
266
267    pub(crate) fn terminal_status_recorded_since(&self, baseline_generation: u64) -> bool {
268        let terminal_generation = self
269            .replay_state
270            .last_terminal_status_generation
271            .load(std::sync::atomic::Ordering::Acquire);
272        let compaction_generation = self
273            .replay_state
274            .last_compaction_generation
275            .load(std::sync::atomic::Ordering::Acquire);
276        terminal_generation > baseline_generation && terminal_generation > compaction_generation
277    }
278
279    pub(in crate::sessions) fn mark_replay_changed(&self) -> u64 {
280        self.replay_state
281            .generation
282            .fetch_add(1, std::sync::atomic::Ordering::AcqRel)
283            .saturating_add(1)
284    }
285
286    pub(in crate::sessions) fn mark_terminal_status_recorded(&self, generation: u64) {
287        self.replay_state
288            .last_terminal_status_generation
289            .store(generation, std::sync::atomic::Ordering::Release);
290    }
291
292    pub(in crate::sessions) fn mark_compaction_recorded(&self, generation: u64) {
293        self.replay_state
294            .last_compaction_generation
295            .store(generation, std::sync::atomic::Ordering::Release);
296    }
297
298    #[cfg(test)]
299    pub(crate) fn unchecked_for_test(id: String, path: PathBuf) -> Self {
300        Self::new(id, path)
301    }
302
303    #[cfg(test)]
304    pub(crate) fn release_active_lease_for_test(&self) {
305        if let Some(slot) = self.active_lease.lock().unwrap().as_ref() {
306            *slot.lock().unwrap() = None;
307        }
308    }
309}
310#[cfg(test)]
311mod tests {
312    use super::*;
313    use proptest::prelude::*;
314    use tempfile::TempDir;
315
316    fn valid_session_id_strategy() -> impl Strategy<Value = String> {
317        proptest::string::string_regex("[A-Za-z0-9_-]{1,64}").unwrap()
318    }
319
320    fn invalid_session_id_strategy() -> impl Strategy<Value = String> {
321        prop_oneof![
322            Just(String::new()),
323            Just(".".to_string()),
324            any::<String>().prop_map(|value| format!("{value}..")),
325            any::<String>().prop_map(|value| format!("{value}/{value}")),
326            any::<String>().prop_map(|value| format!("{value}\\{value}")),
327            any::<String>().prop_map(|value| format!("{value}.jsonl")),
328            any::<String>().prop_map(|value| format!("{value}é")),
329        ]
330    }
331
332    proptest! {
333        #[test]
334        fn path_for_valid_id_keeps_valid_ids_under_session_root(id in valid_session_id_strategy()) {
335            let temp = TempDir::new().unwrap();
336            let root = temp.path().join("sessions");
337            let manager = SessionManager::new(root.clone());
338            let path = manager.path_for_valid_id(&id).unwrap();
339            let normalized_root = crate::path_utils::lexical_normalize(&root);
340            let normalized_path = crate::path_utils::lexical_normalize(&path);
341            let expected_file_name = format!("{id}.jsonl");
342
343            prop_assert!(normalized_path.starts_with(&normalized_root));
344            prop_assert_eq!(normalized_path.parent(), Some(normalized_root.as_path()));
345            prop_assert_eq!(
346                normalized_path.file_name().and_then(|file_name| file_name.to_str()),
347                Some(expected_file_name.as_str())
348            );
349
350            let session = manager.open(id.clone()).unwrap();
351            prop_assert_eq!(session.id(), id);
352            prop_assert_eq!(session.path(), path.as_path());
353        }
354
355        #[test]
356        fn open_and_path_for_valid_id_reject_generated_unsafe_ids(id in invalid_session_id_strategy()) {
357            let temp = TempDir::new().unwrap();
358            let manager = SessionManager::new(temp.path().join("sessions"));
359
360            prop_assert!(manager.open(id.clone()).is_err());
361            prop_assert!(manager.path_for_valid_id(&id).is_err());
362        }
363    }
364
365    #[test]
366    fn session_open_rejects_unsafe_ids_before_joining_paths() {
367        let temp = TempDir::new().unwrap();
368        let manager = SessionManager::new(temp.path().join("sessions"));
369        for id in [
370            "",
371            "..",
372            "../escape",
373            "nested/id",
374            "nested\\id",
375            "/absolute",
376            "bad.jsonl",
377        ] {
378            assert!(manager.open(id).is_err(), "accepted unsafe id {id:?}");
379        }
380        assert!(manager.open("safe_ID-123").is_ok());
381    }
382}