artifactize_core/agent/session/
gc.rs

1//! Size-bounded collection of the session store, with hysteresis: once the conversations
2//! exceed `maxBytes`, the oldest (by last write) are deleted until they take `targetBytes` or
3//! less. A session whose request is still RUNNING, or that a send holds, is never deleted.
4
5use std::{
6    collections::BTreeSet,
7    fs::{self, File},
8    path::Path,
9    time::SystemTime,
10};
11
12use rusqlite::OpenFlags;
13use serde::Serialize;
14
15use super::directory;
16use crate::{limits::AgentSessions, store::DATABASE, types::SessionId};
17
18/// What one collection found and did.
19#[derive(Debug, Default, Clone, Serialize)]
20#[serde(rename_all = "camelCase")]
21pub struct Collection {
22    /// The session store's size before the collection.
23    pub bytes: u64,
24    /// Its size after it, or with a dry run, after the deletions it reports.
25    pub remaining: u64,
26    /// The sessions deleted, or with a dry run those that would be, oldest first.
27    pub removed: Vec<SessionId>,
28}
29
30/// The session store's size.
31#[derive(Debug, Default, Clone, Copy, Serialize)]
32pub struct Usage {
33    pub sessions: usize,
34    pub bytes: u64,
35}
36
37struct Entry {
38    id: SessionId,
39    bytes: u64,
40    written: SystemTime,
41}
42
43/// The saved conversations, read without following links; a missing store is empty.
44fn entries(state: &Path) -> Result<Vec<Entry>, String> {
45    let directory = directory(state);
46    let handle = match crate::platform::open_no_follow(File::options().read(true), &directory) {
47        Ok(handle) => handle,
48        Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()),
49        Err(error) => {
50            return Err(format!(
51                "{}: {error}",
52                crate::platform::path_text(&directory)
53            ));
54        }
55    };
56    let listing = crate::platform::read_dir(&handle).map_err(|e| e.to_string())?;
57    let mut entries = Vec::new();
58    for entry in listing {
59        let entry = entry.map_err(|e| e.to_string())?;
60        let name = entry.file_name();
61        let Some(id) = name.to_str().and_then(|name| name.strip_suffix(".jsonl")) else {
62            continue;
63        };
64        let Ok(id) = id.parse::<SessionId>() else {
65            // Not a session path; never collect arbitrary files in this directory.
66            continue;
67        };
68        if crate::platform::entry_kind(&handle, &name).map_err(|e| e.to_string())?
69            != crate::platform::FileKind::File
70        {
71            continue;
72        }
73        let name = crate::platform::EntryName::new(&name).ok_or("Invalid session entry name.")?;
74        let file = crate::platform::open_entry(&handle, &name).map_err(|e| e.to_string())?;
75        let metadata = file.metadata().map_err(|e| e.to_string())?;
76        if !metadata.is_file() {
77            continue;
78        }
79        let written = metadata.modified().map_err(|e| e.to_string())?;
80        // Tests replace only ordering metadata; file validation and collection stay real.
81        #[cfg(test)]
82        let written = tests::written(state, &id).unwrap_or(written);
83        entries.push(Entry {
84            id,
85            bytes: metadata.len(),
86            written,
87        });
88    }
89    Ok(entries)
90}
91
92/// Pure ordering shared by collection and the explicit-time contract test.
93fn order(entries: &mut [Entry]) {
94    entries.sort_by(|a, b| a.written.cmp(&b.written).then_with(|| a.id.cmp(&b.id)));
95}
96
97/// How many conversations the store holds and their bytes.
98pub fn usage(state: &Path) -> Result<Usage, String> {
99    let entries = entries(state)?;
100    Ok(Usage {
101        sessions: entries.len(),
102        bytes: entries.iter().map(|entry| entry.bytes).sum(),
103    })
104}
105
106/// The sessions of requests that are still RUNNING, read without a writer lock.
107fn running(state: &Path) -> Result<BTreeSet<SessionId>, String> {
108    crate::store::check_probe_files(state)?;
109    let database = state.join(DATABASE);
110    if !database.try_exists().map_err(|e| e.to_string())? {
111        return Ok(BTreeSet::new());
112    }
113    let read = || -> Result<BTreeSet<SessionId>, rusqlite::Error> {
114        let db =
115            rusqlite::Connection::open_with_flags(&database, OpenFlags::SQLITE_OPEN_READ_ONLY)?;
116        db.busy_timeout(crate::store::SQLITE_BUSY_TIMEOUT)?;
117        let initialized: bool = db.query_row(
118            "SELECT EXISTS(SELECT 1 FROM sqlite_master WHERE type='table' AND name='requests')",
119            [],
120            |row| row.get(0),
121        )?;
122        if !initialized {
123            return Ok(BTreeSet::new());
124        }
125        let mut statement = db.prepare(
126            "SELECT json_extract(data,'$.sessionId') FROM requests WHERE status='RUNNING' AND json_extract(data,'$.sessionId') IS NOT NULL",
127        )?;
128        statement
129            .query_map([], |row| row.get::<_, SessionId>(0))?
130            .collect()
131    };
132    read().map_err(|e| format!("Cannot read running Agent reviews: {e}"))
133}
134
135/// Collect the store under `bounds` (see the module). The sessions that are listed first and
136/// then found RUNNING were saved after their request turned RUNNING, so a review that starts
137/// during a collection is either seen as running or not listed at all.
138pub fn collect(state: &Path, bounds: AgentSessions, dry_run: bool) -> Result<Collection, String> {
139    let publisher = if dry_run {
140        crate::changes::Publisher::default()
141    } else {
142        crate::changes::Publisher::new(state)
143    };
144    let mut entries = entries(state)?;
145    let bytes = entries.iter().map(|entry| entry.bytes).sum();
146    let mut collection = Collection {
147        bytes,
148        remaining: bytes,
149        removed: Vec::new(),
150    };
151    if bytes <= bounds.max_bytes {
152        return Ok(collection);
153    }
154    let running = running(state)?;
155    order(&mut entries);
156    let directory = directory(state);
157    for entry in entries {
158        if collection.remaining <= bounds.target_bytes {
159            break;
160        }
161        if running.contains(&entry.id) {
162            continue;
163        }
164        if !dry_run {
165            let lock = directory.join(format!("{}.lock", entry.id));
166            // A session a send holds stays; its lock goes with it.
167            let held = match crate::platform::open_no_follow(
168                File::options().read(true).write(true),
169                &lock,
170            ) {
171                Ok(file) => match file.try_lock() {
172                    Ok(()) => Some(file),
173                    Err(std::fs::TryLockError::WouldBlock) => continue,
174                    Err(std::fs::TryLockError::Error(error)) => return Err(error.to_string()),
175                },
176                Err(error) if error.kind() == std::io::ErrorKind::NotFound => None,
177                Err(error) => return Err(error.to_string()),
178            };
179            match fs::remove_file(directory.join(format!("{}.jsonl", entry.id))) {
180                Ok(()) => {
181                    publisher.notify(crate::changes::Change::SessionInvalidated(entry.id.clone()))
182                }
183                Err(error) if error.kind() == std::io::ErrorKind::NotFound => {}
184                Err(error) => {
185                    return Err(format!("Cannot remove Agent session {}: {error}", entry.id));
186                }
187            }
188            if held.is_some() {
189                let _ = fs::remove_file(&lock);
190            }
191        }
192        collection.remaining = collection.remaining.saturating_sub(entry.bytes);
193        collection.removed.push(entry.id);
194    }
195    Ok(collection)
196}
197
198#[cfg(test)]
199mod tests {
200    use std::{
201        cell::RefCell,
202        collections::BTreeMap,
203        path::PathBuf,
204        time::{Duration, UNIX_EPOCH},
205    };
206
207    use super::*;
208
209    thread_local! {
210        // Collection is synchronous. Each test's owned state gets independent metadata
211        // on its own thread, without changing or trusting filesystem timestamps.
212        static WRITTEN: RefCell<BTreeMap<PathBuf, BTreeMap<String, SystemTime>>> =
213            const { RefCell::new(BTreeMap::new()) };
214    }
215
216    pub(super) fn written(state: &Path, id: &SessionId) -> Option<SystemTime> {
217        WRITTEN.with(|times| times.borrow().get(state)?.get(id.as_str()).copied())
218    }
219
220    /// Owns the real files and the injected write times until the collection finishes.
221    struct Store {
222        root: tempfile::TempDir,
223    }
224    impl Store {
225        fn path(&self) -> &Path {
226            self.root.path()
227        }
228
229        fn set_written(&self, id: &str, seconds: u64) {
230            WRITTEN.with(|times| {
231                let mut times = times.borrow_mut();
232                let written = times
233                    .get_mut(self.path())
234                    .expect("owned state's metadata")
235                    .get_mut(id)
236                    .expect("declared session");
237                *written = UNIX_EPOCH + Duration::from_secs(seconds);
238            });
239        }
240    }
241    impl Drop for Store {
242        fn drop(&mut self) {
243            WRITTEN.with(|times| {
244                times.borrow_mut().remove(self.path());
245            });
246        }
247    }
248
249    /// Equal synthetic write times make identity the deterministic tie breaker.
250    fn store(sizes: &[(&str, usize)]) -> Store {
251        let state = Store {
252            root: crate::test_os::tempdir(),
253        };
254        let sessions = directory(state.path());
255        fs::create_dir_all(&sessions).unwrap();
256        for (id, bytes) in sizes {
257            fs::write(sessions.join(format!("{id}.jsonl")), vec![b'x'; *bytes]).unwrap();
258        }
259        WRITTEN.with(|times| {
260            times.borrow_mut().insert(
261                state.path().to_path_buf(),
262                sizes
263                    .iter()
264                    .map(|(id, _)| ((*id).to_owned(), UNIX_EPOCH))
265                    .collect(),
266            );
267        });
268        state
269    }
270
271    #[test]
272    fn ordering_uses_explicit_times_then_identity() {
273        let state = store(&[("a", 1), ("c", 1), ("b", 1)]);
274        for (id, seconds) in [("a", 20), ("c", 10), ("b", 10)] {
275            state.set_written(id, seconds);
276        }
277        let collection = collect(state.path(), bounds(1, 0), false).unwrap();
278        assert_eq!(
279            collection
280                .removed
281                .iter()
282                .map(ToString::to_string)
283                .collect::<Vec<_>>(),
284            ["b", "c", "a"]
285        );
286        assert_eq!(usage(state.path()).unwrap().sessions, 0);
287    }
288
289    fn bounds(max_bytes: u64, target_bytes: u64) -> AgentSessions {
290        AgentSessions {
291            enabled: true,
292            max_bytes,
293            target_bytes,
294        }
295    }
296
297    #[test]
298    fn collects_oldest_first_down_to_the_target_only_past_the_maximum() {
299        let state = store(&[("a", 100), ("b", 100), ("c", 100), ("d", 100)]);
300        // At the maximum, nothing is collected.
301        let collection = collect(state.path(), bounds(400, 150), false).unwrap();
302        assert!(collection.removed.is_empty());
303        assert_eq!(usage(state.path()).unwrap().sessions, 4);
304
305        let dry = collect(state.path(), bounds(399, 150), true).unwrap();
306        assert_eq!(
307            dry.removed
308                .iter()
309                .map(ToString::to_string)
310                .collect::<Vec<_>>(),
311            ["a", "b", "c"]
312        );
313        assert_eq!((dry.bytes, dry.remaining), (400, 100));
314        assert_eq!(usage(state.path()).unwrap().sessions, 4);
315
316        // A held lock keeps its session.
317        let lock = directory(state.path()).join("b.lock");
318        let held = File::create(&lock).unwrap();
319        held.lock().unwrap();
320        let collection = collect(state.path(), bounds(399, 200), false).unwrap();
321        assert_eq!(
322            collection
323                .removed
324                .iter()
325                .map(ToString::to_string)
326                .collect::<Vec<_>>(),
327            ["a", "c"]
328        );
329        assert_eq!(collection.remaining, 200);
330        drop(held);
331        let left: BTreeSet<_> = fs::read_dir(directory(state.path()))
332            .unwrap()
333            .map(|entry| entry.unwrap().file_name().into_string().unwrap())
334            .collect();
335        assert_eq!(
336            left,
337            ["b.jsonl", "b.lock", "d.jsonl"]
338                .map(str::to_owned)
339                .into_iter()
340                .collect()
341        );
342    }
343
344    #[test]
345    fn invalid_session_filenames_are_not_collected() {
346        let state = store(&[("valid", 100), ("..", 100), ("bad name", 100)]);
347        assert_eq!(usage(state.path()).unwrap().sessions, 1);
348        let collection = collect(state.path(), bounds(1, 0), false).unwrap();
349        assert_eq!(
350            collection
351                .removed
352                .iter()
353                .map(ToString::to_string)
354                .collect::<Vec<_>>(),
355            ["valid"]
356        );
357        assert!(directory(state.path()).join("...jsonl").is_file());
358        assert!(directory(state.path()).join("bad name.jsonl").is_file());
359    }
360
361    #[test]
362    fn running_session_ids_are_validated_at_the_sql_edge() {
363        let state = store(&[("active", 100), ("finished", 100)]);
364        let db = rusqlite::Connection::open(state.path().join(DATABASE)).unwrap();
365        db.execute_batch(
366            "CREATE TABLE requests(status TEXT, data TEXT);
367            INSERT INTO requests VALUES('RUNNING','{\"sessionId\":\"active\"}');",
368        )
369        .unwrap();
370        let collection = collect(state.path(), bounds(1, 0), false).unwrap();
371        assert_eq!(
372            collection
373                .removed
374                .iter()
375                .map(ToString::to_string)
376                .collect::<Vec<_>>(),
377            ["finished"]
378        );
379        db.execute_batch("UPDATE requests SET data='{\"sessionId\":\"../invalid\"}';")
380            .unwrap();
381        assert!(collect(state.path(), bounds(1, 0), false).is_err());
382        assert!(directory(state.path()).join("active.jsonl").is_file());
383    }
384
385    #[test]
386    fn a_missing_store_is_empty() {
387        let state = crate::test_os::tempdir();
388        let collection = collect(state.path(), bounds(1, 0), false).unwrap();
389        assert_eq!((collection.bytes, collection.removed.len()), (0, 0));
390    }
391}
artifactize_core/agent/session/gc.rs

artifactize_core/agent/session/
gc.rs
