wcode_harness/
session.rs

1use std::fs;
2use std::io::{self, Write};
3use std::path::{Path, PathBuf};
4
5use chrono::{SecondsFormat, Utc};
6use serde::{Deserialize, Serialize};
7
8use crate::event::TodoItem;
9use crate::message::{AgentMessage, Usage};
10
11#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
12#[serde(tag = "type", rename_all = "snake_case")]
13pub enum SessionEntry {
14    Header {
15        version: u32,
16        id: String,
17        cwd: String,
18        created: String, // RFC3339
19    },
20    Message {
21        id: String,
22        /// Reserved for future tree/fork sessions; the kernel is linear today
23        /// and always records `None`.
24        parent_id: Option<String>,
25        message: AgentMessage,
26    },
27    ModelChange {
28        id: String,
29        model: String,
30    },
31    EffortChange {
32        id: String,
33        /// None = effort cleared (send nothing).
34        effort: Option<String>,
35    },
36    /// A compaction boundary: `summary` stands in for every `Message` before
37    /// `first_kept_message` (an ordinal into the `Message` sequence).
38    Compaction {
39        id: String,
40        summary: String,
41        first_kept_message: usize,
42        #[serde(default, skip_serializing_if = "Option::is_none")]
43        tokens_before: Option<u64>,
44        #[serde(default, skip_serializing_if = "Option::is_none")]
45        usage: Option<Usage>,
46    },
47    /// The persisted checklist (the `todo` tool's state): last write wins.
48    /// Written by the `todo` tool on every write so plan→todo survives
49    /// `--resume`/`/reload`; read by the tool's seed path and by `/verify` (D8).
50    Todo {
51        id: String,
52        todos: Vec<TodoItem>,
53    },
54    #[serde(other)]
55    Unknown,
56}
57
58pub struct Session {
59    path: Option<PathBuf>,
60    entries: Vec<SessionEntry>,
61}
62
63impl Session {
64    pub fn create(dir: &Path) -> io::Result<Session> {
65        let cwd = std::env::current_dir().unwrap_or_else(|_| PathBuf::from("."));
66        Session::create_with_cwd(dir, &cwd)
67    }
68
69    /// Like [`create`], but records `cwd` in the session header instead of the
70    /// process working directory. An agent whose tools run against a custom
71    /// `working_dir` (see `AgentConfig.working_dir`) should stamp that same
72    /// directory here so a resumed session's header agrees with where the tools
73    /// actually ran — otherwise the header's `cwd` describes a directory the
74    /// tools never touched.
75    pub fn create_with_cwd(dir: &Path, cwd: &Path) -> io::Result<Session> {
76        let id = uuid::Uuid::new_v4();
77        let name = format!(
78            "{}_{}.jsonl",
79            Utc::now().timestamp_millis(),
80            &id.simple().to_string()[..8]
81        );
82        Self::open_new(dir, name, id, cwd)
83    }
84
85    /// Like [`create`], but the file is exactly `<dir>/<name>.jsonl` — the
86    /// deterministic tree session groups need (`root.jsonl`, `members/<name>.jsonl`),
87    /// where a `{millis}_{uuid8}` name would not map back to an address. Header
88    /// content identical to [`create_with_cwd`]; only the file name is caller-
89    /// supplied.
90    pub fn create_named(dir: &Path, name: &str) -> io::Result<Session> {
91        let cwd = std::env::current_dir().unwrap_or_else(|_| PathBuf::from("."));
92        Session::create_named_with_cwd(dir, name, &cwd)
93    }
94
95    /// [`create_named`] with an explicit `cwd` (see [`create_with_cwd`]).
96    pub fn create_named_with_cwd(dir: &Path, name: &str, cwd: &Path) -> io::Result<Session> {
97        Self::open_new(dir, format!("{name}.jsonl"), uuid::Uuid::new_v4(), cwd)
98    }
99
100    /// The shared constructor: fresh dirs, a header stamped first, one file.
101    fn open_new(dir: &Path, filename: String, id: uuid::Uuid, cwd: &Path) -> io::Result<Session> {
102        fs::create_dir_all(dir)?;
103        let path = dir.join(filename);
104        let mut session = Session {
105            path: Some(path.clone()),
106            entries: Vec::new(),
107        };
108        let header = SessionEntry::Header {
109            version: 1,
110            id: id.to_string(),
111            cwd: cwd.to_string_lossy().into_owned(),
112            created: Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true),
113        };
114        session.append(header)?;
115        Ok(session)
116    }
117
118    pub fn open(path: &Path) -> io::Result<Session> {
119        let raw = fs::read_to_string(path)?;
120        let lines: Vec<&str> = raw.lines().collect();
121        let last_non_empty = lines.iter().rposition(|l| !l.trim().is_empty());
122        let mut entries = Vec::new();
123        for (i, line) in lines.iter().enumerate() {
124            if line.trim().is_empty() {
125                continue;
126            }
127            match serde_json::from_str::<SessionEntry>(line) {
128                Ok(e) => entries.push(e),
129                Err(e) => {
130                    if Some(i) == last_non_empty {
131                        break; // torn final write; rest is blank
132                    }
133                    return Err(io::Error::new(io::ErrorKind::InvalidData, e));
134                }
135            }
136        }
137        Ok(Session {
138            path: Some(path.to_path_buf()),
139            entries,
140        })
141    }
142
143    pub fn in_memory() -> Session {
144        Session {
145            path: None,
146            entries: Vec::new(),
147        }
148    }
149
150    pub fn append(&mut self, e: SessionEntry) -> io::Result<()> {
151        if let Some(path) = &self.path {
152            let mut f = fs::OpenOptions::new()
153                .create(true)
154                .append(true)
155                .open(path)?;
156            let line = serde_json::to_string(&e)
157                .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
158            f.write_all(format!("{line}\n").as_bytes())?;
159            f.flush()?;
160        }
161        self.entries.push(e);
162        Ok(())
163    }
164
165    pub fn entries(&self) -> &[SessionEntry] {
166        &self.entries
167    }
168
169    pub fn messages(&self) -> Vec<AgentMessage> {
170        // The last compaction governs: everything before its kept boundary is
171        // represented by the summary; everything after is kept verbatim.
172        let boundary = self.entries.iter().rev().find_map(|e| match e {
173            SessionEntry::Compaction {
174                summary,
175                first_kept_message,
176                ..
177            } => Some((summary.clone(), *first_kept_message)),
178            _ => None,
179        });
180        let (summary, skip) = match boundary {
181            Some((s, k)) => (Some(s), k),
182            None => (None, 0),
183        };
184
185        let mut out = Vec::new();
186        if let Some(s) = summary {
187            out.push(AgentMessage::summary(s));
188        }
189        let mut ordinal = 0usize;
190        for entry in &self.entries {
191            if let SessionEntry::Message { message, .. } = entry {
192                if ordinal >= skip {
193                    out.push(message.clone());
194                }
195                ordinal += 1;
196            }
197        }
198        out
199    }
200
201    /// Number of persisted `Message` entries — the ordinal space that
202    /// [`SessionEntry::Compaction::first_kept_message`] indexes into.
203    pub fn message_count(&self) -> usize {
204        self.entries
205            .iter()
206            .filter(|e| matches!(e, SessionEntry::Message { .. }))
207            .count()
208    }
209
210    /// Record a compaction: `summary` stands in for every `Message` before the
211    /// most recent `kept_messages`, which stay. The boundary ordinal is derived
212    /// from the current message count, so it stays correct as messages append.
213    pub fn record_compaction(
214        &mut self,
215        summary: &str,
216        kept_messages: usize,
217        tokens_before: Option<u64>,
218        usage: Option<Usage>,
219    ) -> io::Result<()> {
220        let first_kept_message = self.message_count().saturating_sub(kept_messages);
221        self.append(SessionEntry::Compaction {
222            id: uuid::Uuid::new_v4().to_string(),
223            summary: summary.to_string(),
224            first_kept_message,
225            tokens_before,
226            usage,
227        })
228    }
229
230    pub fn path(&self) -> Option<&Path> {
231        self.path.as_deref()
232    }
233
234    pub fn model(&self) -> Option<String> {
235        self.entries.iter().rev().find_map(|e| match e {
236            SessionEntry::ModelChange { model, .. } => Some(model.clone()),
237            _ => None,
238        })
239    }
240
241    pub fn effort(&self) -> Option<Option<String>> {
242        self.entries.iter().rev().find_map(|e| match e {
243            SessionEntry::EffortChange { effort, .. } => Some(effort.clone()),
244            _ => None,
245        })
246    }
247
248    /// The most recently persisted checklist, if any (`None` = never written).
249    /// Last write wins, like [`Self::model`]/[`Self::effort`].
250    pub fn todo(&self) -> Option<Vec<TodoItem>> {
251        self.entries.iter().rev().find_map(|e| match e {
252            SessionEntry::Todo { todos, .. } => Some(todos.clone()),
253            _ => None,
254        })
255    }
256}
257
258#[cfg(test)]
259mod tests {
260    use super::*;
261    use crate::event::TodoStatus;
262    use crate::message::{ContentBlock, StopReason};
263    use serde_json::json;
264
265    fn msg(text: &str) -> SessionEntry {
266        SessionEntry::Message {
267            id: uuid::Uuid::new_v4().to_string(),
268            parent_id: None,
269            message: AgentMessage::user_text(text),
270        }
271    }
272
273    #[test]
274    fn create_writes_header_with_timestamped_name() {
275        let dir = tempfile::tempdir().unwrap();
276        let s = Session::create(dir.path()).unwrap();
277        let name = s.path().unwrap().file_name().unwrap().to_str().unwrap();
278        let (ts, rest) = name.split_once('_').unwrap();
279        assert!(ts.parse::<i64>().is_ok());
280        let (hex, ext) = rest.split_once('.').unwrap();
281        assert_eq!(ext, "jsonl");
282        assert_eq!(hex.len(), 8);
283        assert!(hex.chars().all(|c| c.is_ascii_hexdigit()));
284        assert_eq!(s.entries().len(), 1);
285        match &s.entries()[0] {
286            SessionEntry::Header {
287                version, id, cwd, ..
288            } => {
289                assert_eq!(*version, 1);
290                assert!(!id.is_empty());
291                assert!(!cwd.is_empty());
292            }
293            _ => panic!("first entry not header"),
294        }
295    }
296
297    #[test]
298    fn create_with_cwd_stamps_header_cwd() {
299        let dir = tempfile::tempdir().unwrap();
300        let cwd = std::path::Path::new("/custom/workdir");
301        let s = Session::create_with_cwd(dir.path(), cwd).unwrap();
302        match &s.entries()[0] {
303            SessionEntry::Header { cwd, .. } => assert_eq!(cwd, "/custom/workdir"),
304            _ => panic!("first entry not header"),
305        }
306    }
307
308    #[test]
309    fn create_named_writes_to_the_named_file() {
310        let dir = tempfile::tempdir().unwrap();
311        let root = Session::create_named_with_cwd(dir.path(), "root", std::path::Path::new("/wd"))
312            .unwrap();
313        assert_eq!(root.path().unwrap().file_name().unwrap(), "root.jsonl");
314        match &root.entries()[0] {
315            SessionEntry::Header { version, cwd, .. } => {
316                assert_eq!(*version, 1);
317                assert_eq!(cwd, "/wd");
318            }
319            _ => panic!("first entry not header"),
320        }
321
322        // The bare `create_named` stamps the process cwd, and the file name is
323        // exactly `<name>.jsonl` — the deterministic tree session groups need.
324        let member = Session::create_named(dir.path(), "w1").unwrap();
325        assert_eq!(member.path().unwrap().file_name().unwrap(), "w1.jsonl");
326
327        // Reopen round-trips: an append comes back, so group transcripts are
328        // ordinary kernel sessions with caller-supplied names.
329        let mut reopened = Session::open(&dir.path().join("root.jsonl")).unwrap();
330        reopened
331            .append(SessionEntry::ModelChange {
332                id: "m1".into(),
333                model: "m".into(),
334            })
335            .unwrap();
336        assert_eq!(reopened.entries().len(), 2);
337        assert_eq!(
338            Session::open(&dir.path().join("root.jsonl"))
339                .unwrap()
340                .model()
341                .as_deref(),
342            Some("m")
343        );
344    }
345
346    #[test]
347    fn create_append_open_roundtrip() {
348        let dir = tempfile::tempdir().unwrap();
349        let mut s = Session::create(dir.path()).unwrap();
350        s.append(msg("hello")).unwrap();
351        s.append(SessionEntry::ModelChange {
352            id: "m1".into(),
353            model: "claude-x".into(),
354        })
355        .unwrap();
356        s.append(SessionEntry::Message {
357            id: "a2".into(),
358            parent_id: Some("a1".into()),
359            message: AgentMessage::Assistant {
360                content: vec![ContentBlock::Text {
361                    text: "reply".into(),
362                }],
363                stop_reason: StopReason::Stop,
364                usage: None,
365                model: Some("claude-x".into()),
366            },
367        })
368        .unwrap();
369
370        s.append(SessionEntry::Todo {
371            id: "td1".into(),
372            todos: vec![TodoItem {
373                content: "step".into(),
374                status: TodoStatus::Pending,
375            }],
376        })
377        .unwrap();
378
379        let path = s.path().unwrap().to_path_buf();
380        let reopened = Session::open(&path).unwrap();
381        assert_eq!(reopened.entries(), s.entries());
382        assert_eq!(reopened.messages().len(), 2);
383        assert_eq!(reopened.model().as_deref(), Some("claude-x"));
384        assert_eq!(reopened.todo().unwrap()[0].content, "step");
385    }
386
387    #[test]
388    fn torn_final_line_tolerated() {
389        let dir = tempfile::tempdir().unwrap();
390        let path = dir.path().join("s.jsonl");
391        fs::write(
392            &path,
393            format!(
394                "{}\n{}",
395                json!({"type":"message","id":"1","parent_id":null,"message":{"role":"user","content":[{"type":"text","text":"q"}]}}),
396                "{\"type\":\"mess" // torn
397            ),
398        )
399        .unwrap();
400        let s = Session::open(&path).unwrap();
401        assert_eq!(s.entries().len(), 1);
402        assert_eq!(s.messages().len(), 1);
403    }
404
405    #[test]
406    fn torn_line_followed_by_blank_lines_tolerated() {
407        let dir = tempfile::tempdir().unwrap();
408        let path = dir.path().join("s.jsonl");
409        fs::write(
410            &path,
411            format!(
412                "{}\n{{\"type\":\"mess\n\n",
413                json!({"type":"message","id":"1","parent_id":null,"message":{"role":"user","content":[{"type":"text","text":"q"}]}})
414            ),
415        )
416        .unwrap();
417        let s = Session::open(&path).unwrap();
418        assert_eq!(s.entries().len(), 1);
419        assert_eq!(s.messages().len(), 1);
420    }
421
422    #[test]
423    fn torn_earlier_line_is_error() {
424        let dir = tempfile::tempdir().unwrap();
425        let path = dir.path().join("s.jsonl");
426        fs::write(
427            &path,
428            "{\"type\":\"mess\n{\"type\":\"model_change\",\"id\":\"m\",\"model\":\"x\"}\n",
429        )
430        .unwrap();
431        assert!(Session::open(&path).is_err());
432    }
433
434    #[test]
435    fn unknown_entries_tolerated() {
436        let dir = tempfile::tempdir().unwrap();
437        let path = dir.path().join("s.jsonl");
438        fs::write(
439            &path,
440            format!(
441                "{}\n{}\n{}\n",
442                json!({"type":"header","version":1,"id":"x","cwd":"/","created":"2026-01-01T00:00:00Z"}),
443                json!({"type":"future_thing","data":42}),
444                json!({"type":"message","id":"1","parent_id":null,"message":{"role":"user","content":[{"type":"text","text":"q"}]}})
445            ),
446        )
447        .unwrap();
448        let s = Session::open(&path).unwrap();
449        assert_eq!(s.entries().len(), 3);
450        assert_eq!(s.entries()[1], SessionEntry::Unknown);
451        assert_eq!(s.messages().len(), 1);
452    }
453
454    #[test]
455    fn model_latest_wins() {
456        let mut s = Session::in_memory();
457        assert_eq!(s.model(), None);
458        s.append(SessionEntry::ModelChange {
459            id: "1".into(),
460            model: "old".into(),
461        })
462        .unwrap();
463        s.append(SessionEntry::ModelChange {
464            id: "2".into(),
465            model: "new".into(),
466        })
467        .unwrap();
468        assert_eq!(s.model().as_deref(), Some("new"));
469    }
470
471    #[test]
472    fn effort_latest_wins_and_none_means_cleared() {
473        let mut s = Session::in_memory();
474        assert_eq!(s.effort(), None);
475        s.append(SessionEntry::EffortChange {
476            id: "1".into(),
477            effort: Some("high".into()),
478        })
479        .unwrap();
480        assert_eq!(s.effort(), Some(Some("high".into())));
481        s.append(SessionEntry::EffortChange {
482            id: "2".into(),
483            effort: None,
484        })
485        .unwrap();
486        assert_eq!(s.effort(), Some(None));
487    }
488
489    #[test]
490    fn effort_roundtrips_through_jsonl() {
491        let dir = tempfile::tempdir().unwrap();
492        let mut s = Session::create(dir.path()).unwrap();
493        s.append(SessionEntry::EffortChange {
494            id: "e1".into(),
495            effort: Some("low".into()),
496        })
497        .unwrap();
498        let path = s.path().unwrap().to_path_buf();
499        let reopened = Session::open(&path).unwrap();
500        assert_eq!(reopened.effort(), Some(Some("low".into())));
501    }
502
503    #[test]
504    fn todo_latest_wins_and_roundtrips() {
505        let dir = tempfile::tempdir().unwrap();
506        let mut s = Session::create(dir.path()).unwrap();
507        s.append(SessionEntry::Todo {
508            id: "t1".into(),
509            todos: vec![TodoItem {
510                content: "first".into(),
511                status: TodoStatus::Pending,
512            }],
513        })
514        .unwrap();
515        s.append(SessionEntry::Todo {
516            id: "t2".into(),
517            todos: vec![TodoItem {
518                content: "second".into(),
519                status: TodoStatus::InProgress,
520            }],
521        })
522        .unwrap();
523
524        let path = s.path().unwrap().to_path_buf();
525        let reopened = Session::open(&path).unwrap();
526        let restored = reopened.todo().unwrap();
527        assert_eq!(restored.len(), 1, "last write wins");
528        assert_eq!(restored[0].content, "second");
529        assert_eq!(restored[0].status, TodoStatus::InProgress);
530
531        // The serde wire tag is `todo` (the extensible-`type` path).
532        let raw = fs::read_to_string(&path).unwrap();
533        assert!(raw.contains(r#""type":"todo""#), "{raw}");
534    }
535
536    #[test]
537    fn in_memory_writes_nothing() {
538        let dir = tempfile::tempdir().unwrap();
539        let mut s = Session::in_memory();
540        s.append(msg("hi")).unwrap();
541        assert!(s.path().is_none());
542        assert_eq!(s.entries().len(), 1);
543        assert_eq!(s.messages().len(), 1);
544        assert_eq!(fs::read_dir(dir.path()).unwrap().count(), 0);
545    }
546}
547
548#[cfg(test)]
549mod compaction_tests {
550    use super::*;
551
552    fn msg(text: &str) -> SessionEntry {
553        SessionEntry::Message {
554            id: uuid::Uuid::new_v4().to_string(),
555            parent_id: None,
556            message: AgentMessage::user_text(text),
557        }
558    }
559
560    fn texts(s: &Session) -> Vec<String> {
561        s.messages().iter().map(|m| m.as_text()).collect()
562    }
563
564    #[test]
565    fn no_compaction_returns_every_message() {
566        let mut s = Session::in_memory();
567        s.append(msg("a")).unwrap();
568        s.append(msg("b")).unwrap();
569        assert_eq!(texts(&s), ["a", "b"]);
570        assert_eq!(s.message_count(), 2);
571    }
572
573    #[test]
574    fn compaction_rebuilds_summary_plus_kept_tail_across_reopen() {
575        let dir = tempfile::tempdir().unwrap();
576        let mut s = Session::create(dir.path()).unwrap();
577        for t in ["m1", "m2", "m3", "m4", "m5"] {
578            s.append(msg(t)).unwrap();
579        }
580        // Keep the newest two (m4, m5); summarize the rest.
581        s.record_compaction("did 1..3", 2, Some(123), None).unwrap();
582
583        let path = s.path().unwrap().to_path_buf();
584        let mut reopened = Session::open(&path).unwrap();
585        let got = texts(&reopened);
586        assert_eq!(got.len(), 3, "summary + two kept");
587        assert!(got[0].contains("did 1..3"), "summary leads: {got:?}");
588        assert_eq!(got[1], "m4");
589        assert_eq!(got[2], "m5");
590
591        // Messages appended after the compaction are kept too.
592        reopened.append(msg("m6")).unwrap();
593        assert_eq!(texts(&reopened), [got[0].clone(), "m4".into(), "m5".into(), "m6".into()]);
594        assert_eq!(reopened.message_count(), 6);
595    }
596
597    #[test]
598    fn second_compaction_supersedes_the_first() {
599        let mut s = Session::in_memory();
600        for t in ["m1", "m2", "m3", "m4", "m5", "m6"] {
601            s.append(msg(t)).unwrap();
602        }
603        s.record_compaction("first", 4, None, None).unwrap(); // keep m3..m6
604        assert_eq!(
605            texts(&s),
606            [
607                "Summary of the earlier conversation (auto-generated by compaction):\n\nfirst",
608                "m3",
609                "m4",
610                "m5",
611                "m6"
612            ]
613        );
614        s.record_compaction("second", 2, None, None).unwrap(); // keep m5, m6
615        let got = texts(&s);
616        assert_eq!(got.len(), 3);
617        assert!(got[0].contains("second"), "the latest summary wins: {got:?}");
618        assert_eq!(got[1], "m5");
619        assert_eq!(got[2], "m6");
620    }
621}
wcode_harness/session.rs

wcode_harness/
session.rs
