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        entries.push(Entry {
80            id,
81            bytes: metadata.len(),
82            written: metadata.modified().map_err(|e| e.to_string())?,
83        });
84    }
85    Ok(entries)
86}
87
88/// Pure ordering shared by collection and the explicit-time contract test.
89fn order(entries: &mut [Entry]) {
90    entries.sort_by(|a, b| a.written.cmp(&b.written).then_with(|| a.id.cmp(&b.id)));
91}
92
93/// How many conversations the store holds and their bytes.
94pub fn usage(state: &Path) -> Result<Usage, String> {
95    let entries = entries(state)?;
96    Ok(Usage {
97        sessions: entries.len(),
98        bytes: entries.iter().map(|entry| entry.bytes).sum(),
99    })
100}
101
102/// The sessions of requests that are still RUNNING, read without a writer lock.
103fn running(state: &Path) -> Result<BTreeSet<SessionId>, String> {
104    crate::store::check_probe_files(state)?;
105    let database = state.join(DATABASE);
106    if !database.try_exists().map_err(|e| e.to_string())? {
107        return Ok(BTreeSet::new());
108    }
109    let read = || -> Result<BTreeSet<SessionId>, rusqlite::Error> {
110        let db =
111            rusqlite::Connection::open_with_flags(&database, OpenFlags::SQLITE_OPEN_READ_ONLY)?;
112        db.busy_timeout(crate::store::SQLITE_BUSY_TIMEOUT)?;
113        let initialized: bool = db.query_row(
114            "SELECT EXISTS(SELECT 1 FROM sqlite_master WHERE type='table' AND name='requests')",
115            [],
116            |row| row.get(0),
117        )?;
118        if !initialized {
119            return Ok(BTreeSet::new());
120        }
121        let mut statement = db.prepare(
122            "SELECT json_extract(data,'$.sessionId') FROM requests WHERE status='RUNNING' AND json_extract(data,'$.sessionId') IS NOT NULL",
123        )?;
124        statement
125            .query_map([], |row| row.get::<_, SessionId>(0))?
126            .collect()
127    };
128    read().map_err(|e| format!("Cannot read running Agent reviews: {e}"))
129}
130
131/// Collect the store under `bounds` (see the module). The sessions that are listed first and
132/// then found RUNNING were saved after their request turned RUNNING, so a review that starts
133/// during a collection is either seen as running or not listed at all.
134pub fn collect(state: &Path, bounds: AgentSessions, dry_run: bool) -> Result<Collection, String> {
135    let publisher = if dry_run {
136        crate::changes::Publisher::default()
137    } else {
138        crate::changes::Publisher::new(state)
139    };
140    let mut entries = entries(state)?;
141    let bytes = entries.iter().map(|entry| entry.bytes).sum();
142    let mut collection = Collection {
143        bytes,
144        remaining: bytes,
145        removed: Vec::new(),
146    };
147    if bytes <= bounds.max_bytes {
148        return Ok(collection);
149    }
150    let running = running(state)?;
151    order(&mut entries);
152    let directory = directory(state);
153    for entry in entries {
154        if collection.remaining <= bounds.target_bytes {
155            break;
156        }
157        if running.contains(&entry.id) {
158            continue;
159        }
160        if !dry_run {
161            let lock = directory.join(format!("{}.lock", entry.id));
162            // A session a send holds stays; its lock goes with it.
163            let held = match crate::platform::open_no_follow(
164                File::options().read(true).write(true),
165                &lock,
166            ) {
167                Ok(file) => match file.try_lock() {
168                    Ok(()) => Some(file),
169                    Err(std::fs::TryLockError::WouldBlock) => continue,
170                    Err(std::fs::TryLockError::Error(error)) => return Err(error.to_string()),
171                },
172                Err(error) if error.kind() == std::io::ErrorKind::NotFound => None,
173                Err(error) => return Err(error.to_string()),
174            };
175            match fs::remove_file(directory.join(format!("{}.jsonl", entry.id))) {
176                Ok(()) => {
177                    publisher.notify(crate::changes::Change::SessionInvalidated(entry.id.clone()))
178                }
179                Err(error) if error.kind() == std::io::ErrorKind::NotFound => {}
180                Err(error) => {
181                    return Err(format!("Cannot remove Agent session {}: {error}", entry.id));
182                }
183            }
184            if held.is_some() {
185                let _ = fs::remove_file(&lock);
186            }
187        }
188        collection.remaining = collection.remaining.saturating_sub(entry.bytes);
189        collection.removed.push(entry.id);
190    }
191    Ok(collection)
192}
193
194#[cfg(test)]
195mod tests {
196    use std::time::{Duration, UNIX_EPOCH};
197
198    use super::*;
199
200    /// Sessions with equal explicit write times; identity provides deterministic tie order.
201    fn store(sizes: &[(&str, usize)]) -> tempfile::TempDir {
202        let state = crate::test_os::tempdir();
203        let sessions = directory(state.path());
204        fs::create_dir_all(&sessions).unwrap();
205        for (id, bytes) in sizes {
206            let path = sessions.join(format!("{id}.jsonl"));
207            fs::write(&path, vec![b'x'; *bytes]).unwrap();
208            crate::test_os::set_modified(&path, UNIX_EPOCH + Duration::from_secs(1_000));
209        }
210        state
211    }
212
213    #[test]
214    fn ordering_uses_explicit_times_then_identity() {
215        let entry = |id: &str, secs| Entry {
216            id: id.parse().unwrap(),
217            bytes: 1,
218            written: UNIX_EPOCH + Duration::from_secs(secs),
219        };
220        let mut entries = [entry("a", 20), entry("c", 10), entry("b", 10)];
221        order(&mut entries);
222        assert_eq!(entries.map(|entry| entry.id.to_string()), ["b", "c", "a"]);
223    }
224
225    fn bounds(max_bytes: u64, target_bytes: u64) -> AgentSessions {
226        AgentSessions {
227            enabled: true,
228            max_bytes,
229            target_bytes,
230        }
231    }
232
233    #[test]
234    fn collects_oldest_first_down_to_the_target_only_past_the_maximum() {
235        let state = store(&[("a", 100), ("b", 100), ("c", 100), ("d", 100)]);
236        // At the maximum, nothing is collected.
237        let collection = collect(state.path(), bounds(400, 150), false).unwrap();
238        assert!(collection.removed.is_empty());
239        assert_eq!(usage(state.path()).unwrap().sessions, 4);
240
241        let dry = collect(state.path(), bounds(399, 150), true).unwrap();
242        assert_eq!(
243            dry.removed
244                .iter()
245                .map(ToString::to_string)
246                .collect::<Vec<_>>(),
247            ["a", "b", "c"]
248        );
249        assert_eq!((dry.bytes, dry.remaining), (400, 100));
250        assert_eq!(usage(state.path()).unwrap().sessions, 4);
251
252        // A held lock keeps its session.
253        let lock = directory(state.path()).join("b.lock");
254        let held = File::create(&lock).unwrap();
255        held.lock().unwrap();
256        let collection = collect(state.path(), bounds(399, 200), false).unwrap();
257        assert_eq!(
258            collection
259                .removed
260                .iter()
261                .map(ToString::to_string)
262                .collect::<Vec<_>>(),
263            ["a", "c"]
264        );
265        assert_eq!(collection.remaining, 200);
266        drop(held);
267        let left: BTreeSet<_> = fs::read_dir(directory(state.path()))
268            .unwrap()
269            .map(|entry| entry.unwrap().file_name().into_string().unwrap())
270            .collect();
271        assert_eq!(
272            left,
273            ["b.jsonl", "b.lock", "d.jsonl"]
274                .map(str::to_owned)
275                .into_iter()
276                .collect()
277        );
278    }
279
280    #[test]
281    fn invalid_session_filenames_are_not_collected() {
282        let state = store(&[("valid", 100), ("..", 100), ("bad name", 100)]);
283        assert_eq!(usage(state.path()).unwrap().sessions, 1);
284        let collection = collect(state.path(), bounds(1, 0), false).unwrap();
285        assert_eq!(
286            collection
287                .removed
288                .iter()
289                .map(ToString::to_string)
290                .collect::<Vec<_>>(),
291            ["valid"]
292        );
293        assert!(directory(state.path()).join("...jsonl").is_file());
294        assert!(directory(state.path()).join("bad name.jsonl").is_file());
295    }
296
297    #[test]
298    fn running_session_ids_are_validated_at_the_sql_edge() {
299        let state = store(&[("active", 100), ("finished", 100)]);
300        let db = rusqlite::Connection::open(state.path().join(DATABASE)).unwrap();
301        db.execute_batch(
302            "CREATE TABLE requests(status TEXT, data TEXT);
303            INSERT INTO requests VALUES('RUNNING','{\"sessionId\":\"active\"}');",
304        )
305        .unwrap();
306        let collection = collect(state.path(), bounds(1, 0), false).unwrap();
307        assert_eq!(
308            collection
309                .removed
310                .iter()
311                .map(ToString::to_string)
312                .collect::<Vec<_>>(),
313            ["finished"]
314        );
315        db.execute_batch("UPDATE requests SET data='{\"sessionId\":\"../invalid\"}';")
316            .unwrap();
317        assert!(running(state.path()).is_err());
318        assert!(directory(state.path()).join("active.jsonl").is_file());
319    }
320
321    #[test]
322    fn a_missing_store_is_empty() {
323        let state = crate::test_os::tempdir();
324        let collection = collect(state.path(), bounds(1, 0), false).unwrap();
325        assert_eq!((collection.bytes, collection.removed.len()), (0, 0));
326    }
327}
artifactize_core/agent/session/gc.rs

artifactize_core/agent/session/
gc.rs
