pi-oxidized 0.1.1

package/crates/pi/src/core/sessions/mod.rs

@@ -0,0 +1,2117 @@
+//! Append-only JSONL v3 session store.
+//!
+//! Ports `.references/pi/packages/coding-agent/src/core/session-manager.ts` and
+//! `session-cwd.ts`. Disk create is deferred until the first assistant message
+//! (`create_new` / `wx`); historical lines are rewritten only on migration load,
+//! empty-file init, or [`SessionManager::create_branched_session`].
+
+pub mod context;
+pub mod cwd;
+pub mod entries;
+pub mod list;
+
+use std::collections::{HashMap, HashSet};
+use std::fs::{self, File, OpenOptions};
+use std::io::{self, Write};
+use std::path::{Path, PathBuf};
+use std::sync::atomic::{AtomicU64, Ordering};
+
+use pi_agent::AgentMessage;
+use serde_json::Value;
+use thiserror::Error;
+
+use super::config::{
+    PathInputOptions, get_agent_dir, get_sessions_dir, normalize_path, resolve_path,
+};
+use super::messages::{CustomMessageContent, MessageConversionError};
+
+pub use context::{
+    DEFAULT_THINKING_LEVEL, LeafRef, SessionContext, SessionModel, build_context_entries,
+    build_session_context, build_session_path, get_latest_compaction_entry,
+    session_entry_to_context_messages,
+};
+pub use cwd::{
+    MissingSessionCwdError, SessionCwdIssue, SessionCwdSource, assert_session_cwd_exists,
+    format_missing_session_cwd_error, format_missing_session_cwd_prompt,
+    get_missing_session_cwd_issue,
+};
+pub use entries::{
+    BranchSummaryEntry, CURRENT_SESSION_VERSION, CompactionEntry, CustomEntry, CustomMessageEntry,
+    FileEntry, LabelEntry, ModelChangeEntry, NO_MESSAGES_PLACEHOLDER, SessionEntry, SessionHeader,
+    SessionInfoEntry, SessionMessageEntry, ThinkingLevelChangeEntry, assert_valid_session_id,
+    create_session_id, generate_id, load_entries_from_file, migrate_session_entries, now_iso,
+    parse_session_entries, parse_session_entry_line, read_session_header,
+};
+pub use list::{
+    MAX_CONCURRENT_SESSION_INFO_LOADS, SessionInfo, SessionListProgress, build_session_info,
+    find_most_recent_session, list_all_sessions, list_sessions_for_cwd, list_sessions_from_dir,
+    session_cwd_matches,
+};
+
+use entries::{
+    file_entry_to_line, iso_to_millis, load_values_from_file, migrate_values_to_current,
+    path_exists, session_entry_to_line,
+};
+
+// ---------------------------------------------------------------------------
+// Errors
+// ---------------------------------------------------------------------------
+
+/// Errors from session manager operations.
+#[derive(Debug, Error)]
+pub enum SessionError {
+    /// Custom session id failed validation.
+    #[error(
+        "Session id must be non-empty, contain only alphanumeric characters, '-', '_', and '.', and start and end with an alphanumeric character"
+    )]
+    InvalidSessionId,
+    /// Referenced entry id is not in the tree.
+    #[error("Entry {0} not found")]
+    EntryNotFound(String),
+    /// Non-empty file that is not a valid pi session.
+    #[error("Session file is not a valid pi session: {0}")]
+    InvalidSessionFile(String),
+    /// Fork source is empty or invalid.
+    #[error("Cannot fork: source session file is empty or invalid: {0}")]
+    ForkSourceEmpty(String),
+    /// Fork source has no session header.
+    #[error("Cannot fork: source session has no header: {0}")]
+    ForkSourceNoHeader(String),
+    /// Filesystem IO failure.
+    #[error("I/O error on {path}: {source}")]
+    Io {
+        /// Path involved in the failure.
+        path: String,
+        /// Underlying IO error.
+        source: io::Error,
+    },
+    /// JSON serialization failure.
+    #[error("JSON error: {0}")]
+    Json(#[from] serde_json::Error),
+    /// Context projection / message conversion failure.
+    #[error(transparent)]
+    Message(#[from] MessageConversionError),
+}
+
+// assert_valid_session_id in entries.rs returns SessionError defined above.
+
+// ---------------------------------------------------------------------------
+// Options / tree node
+// ---------------------------------------------------------------------------
+
+/// Options for creating a new session.
+#[derive(Clone, Debug, Default)]
+pub struct NewSessionOptions {
+    /// Optional custom session id (validated).
+    pub id: Option<String>,
+    /// Optional parent session file path stored in the header.
+    pub parent_session: Option<String>,
+}
+
+/// Tree node returned by [`SessionManager::get_tree`].
+#[derive(Clone, Debug, PartialEq)]
+pub struct SessionTreeNode {
+    /// Entry at this node.
+    pub entry: SessionEntry,
+    /// Children sorted by timestamp ascending.
+    pub children: Vec<SessionTreeNode>,
+    /// Resolved label for this entry, if any.
+    pub label: Option<String>,
+    /// Timestamp of the latest label change for this entry, if any.
+    pub label_timestamp: Option<String>,
+}
+
+// ---------------------------------------------------------------------------
+// Leaf + labels
+// ---------------------------------------------------------------------------
+
+#[derive(Clone, Debug)]
+enum Leaf {
+    Null,
+    Id(String),
+    /// Last non-session entry had no string id (TS `undefined` leaf).
+    Tail,
+}
+
+/// Insertion-ordered label map (`target_id` → label + timestamp).
+#[derive(Clone, Debug, Default)]
+struct LabelMap {
+    /// Insertion-ordered (target, label, timestamp).
+    entries: Vec<(String, String, String)>,
+    /// target → index into `entries`.
+    index: HashMap<String, usize>,
+}
+
+impl LabelMap {
+    fn clear(&mut self) {
+        self.entries.clear();
+        self.index.clear();
+    }
+
+    fn get(&self, target: &str) -> Option<&str> {
+        self.index.get(target).map(|&i| self.entries[i].1.as_str())
+    }
+
+    fn get_timestamp(&self, target: &str) -> Option<&str> {
+        self.index.get(target).map(|&i| self.entries[i].2.as_str())
+    }
+
+    /// Set or update a label. Existing targets keep their insertion position.
+    fn set(&mut self, target: String, label: String, timestamp: String) {
+        if let Some(&i) = self.index.get(&target) {
+            self.entries[i].1 = label;
+            self.entries[i].2 = timestamp;
+        } else {
+            let i = self.entries.len();
+            self.index.insert(target.clone(), i);
+            self.entries.push((target, label, timestamp));
+        }
+    }
+
+    fn delete(&mut self, target: &str) {
+        let Some(i) = self.index.remove(target) else {
+            return;
+        };
+        self.entries.remove(i);
+        // Rebuild index after removal.
+        self.index.clear();
+        for (idx, (t, _, _)) in self.entries.iter().enumerate() {
+            self.index.insert(t.clone(), idx);
+        }
+    }
+
+    fn iter(&self) -> impl Iterator<Item = (&str, &str, &str)> {
+        self.entries
+            .iter()
+            .map(|(t, l, ts)| (t.as_str(), l.as_str(), ts.as_str()))
+    }
+}
+
+// ---------------------------------------------------------------------------
+// SessionManager
+// ---------------------------------------------------------------------------
+
+/// Append-only JSONL v3 session tree.
+///
+/// Single-writer, no interior locking (matches TypeScript). Historical lines are
+/// never rewritten on normal append.
+#[derive(Debug)]
+pub struct SessionManager {
+    session_id: String,
+    session_file: Option<String>,
+    session_dir: String,
+    cwd: String,
+    persist: bool,
+    flushed: bool,
+    file_entries: Vec<FileEntry>,
+    /// id → index into `file_entries` (Entry variants only).
+    by_id: HashMap<String, usize>,
+    labels: LabelMap,
+    leaf: Leaf,
+}
+
+fn assemble_tree_node(
+    id: &str,
+    node_map: &HashMap<String, SessionTreeNode>,
+    children_of: &HashMap<String, Vec<String>>,
+) -> Option<SessionTreeNode> {
+    let mut node = node_map.get(id)?.clone();
+    node.children.clear();
+    if let Some(kids) = children_of.get(id) {
+        for kid in kids {
+            if let Some(child) = assemble_tree_node(kid, node_map, children_of) {
+                node.children.push(child);
+            }
+        }
+        node.children.sort_by(|a, b| {
+            let ta = a.entry.timestamp().and_then(iso_to_millis);
+            let tb = b.entry.timestamp().and_then(iso_to_millis);
+            match (ta, tb) {
+                (Some(x), Some(y)) => x.cmp(&y),
+                _ => std::cmp::Ordering::Equal,
+            }
+        });
+    }
+    Some(node)
+}
+
+impl SessionManager {
+    fn construct(
+        cwd: &str,
+        session_dir: &str,
+        session_file: Option<String>,
+        persist: bool,
+        options: Option<NewSessionOptions>,
+    ) -> Result<Self, SessionError> {
+        let cwd = path_to_string(&resolve_path(cwd));
+        let session_dir = path_to_string(&normalize_path(session_dir, PathInputOptions::new()));
+        if persist && !session_dir.is_empty() && !path_exists(Path::new(&session_dir)) {
+            fs::create_dir_all(&session_dir).map_err(|source| SessionError::Io {
+                path: session_dir.clone(),
+                source,
+            })?;
+        }
+
+        let mut sm = Self {
+            session_id: String::new(),
+            session_file: None,
+            session_dir,
+            cwd,
+            persist,
+            flushed: false,
+            file_entries: Vec::new(),
+            by_id: HashMap::new(),
+            labels: LabelMap::default(),
+            leaf: Leaf::Null,
+        };
+
+        if let Some(file) = session_file {
+            sm.set_session_file(&file)?;
+        } else {
+            sm.new_session(options)?;
+        }
+        Ok(sm)
+    }
+
+    /// Switch to a different session file (resume / open).
+    ///
+    /// # Errors
+    ///
+    /// Returns IO errors when reading or rewriting the file, invalid-session when a
+    /// non-empty file is not a pi session, and JSON errors during migration.
+    pub fn set_session_file(&mut self, session_file: &str) -> Result<(), SessionError> {
+        let resolved = path_to_string(&resolve_path(session_file));
+        self.session_file = Some(resolved.clone());
+
+        if path_exists(Path::new(&resolved)) {
+            let mut values =
+                load_values_from_file(Path::new(&resolved)).map_err(|source| SessionError::Io {
+                    path: resolved.clone(),
+                    source,
+                })?;
+
+            if values.is_empty() {
+                let size = fs::metadata(&resolved)
+                    .map_err(|source| SessionError::Io {
+                        path: resolved.clone(),
+                        source,
+                    })?
+                    .len();
+                if size > 0 {
+                    return Err(SessionError::InvalidSessionFile(resolved));
+                }
+                // Empty file: init header and rewrite.
+                self.new_session(None)?;
+                self.session_file = Some(resolved);
+                self.rewrite_file()?;
+                self.flushed = true;
+                return Ok(());
+            }
+
+            let header_id = values
+                .iter()
+                .find(|v| v.get("type").and_then(Value::as_str) == Some("session"))
+                .and_then(|h| h.get("id").and_then(Value::as_str))
+                .map(str::to_owned);
+            self.session_id = header_id.unwrap_or_else(create_session_id);
+
+            let migrated = migrate_values_to_current(&mut values);
+            self.file_entries = values
+                .into_iter()
+                .map(entries::file_entry_from_value)
+                .collect();
+            if migrated {
+                self.rewrite_file()?;
+            }
+            self.build_index();
+            self.flushed = true;
+        } else {
+            let explicit = resolved;
+            self.new_session(None)?;
+            self.session_file = Some(explicit);
+        }
+        Ok(())
+    }
+
+    /// Start a new in-memory session (optionally with a deferred file path).
+    ///
+    /// Returns the session file path when persisting, else `None`.
+    ///
+    /// # Errors
+    ///
+    /// Returns a validation error when a custom session id is invalid.
+    pub fn new_session(
+        &mut self,
+        options: Option<NewSessionOptions>,
+    ) -> Result<Option<String>, SessionError> {
+        if let Some(id) = options.as_ref().and_then(|opts| opts.id.as_ref()) {
+            assert_valid_session_id(id)?;
+        }
+        let options = options.unwrap_or_default();
+        self.session_id = options.id.unwrap_or_else(create_session_id);
+        let timestamp = now_iso();
+        let header = SessionHeader::new(
+            self.session_id.clone(),
+            timestamp.clone(),
+            self.cwd.clone(),
+            options.parent_session,
+        );
+        self.file_entries = vec![FileEntry::Header(header)];
+        self.by_id.clear();
+        self.labels.clear();
+        self.leaf = Leaf::Null;
+        self.flushed = false;
+
+        if self.persist {
+            let file_name = session_file_name(&timestamp, &self.session_id);
+            self.session_file = Some(
+                Path::new(&self.session_dir)
+                    .join(file_name)
+                    .to_string_lossy()
+                    .into_owned(),
+            );
+        } else {
+            self.session_file = None;
+        }
+        Ok(self.session_file.clone())
+    }
+
+    fn build_index(&mut self) {
+        self.by_id.clear();
+        self.labels.clear();
+        self.leaf = Leaf::Null;
+        for (idx, fe) in self.file_entries.iter().enumerate() {
+            if fe.is_session_header() {
+                continue;
+            }
+            let Some(entry) = fe.entry() else {
+                continue;
+            };
+            if let Some(id) = entry.id() {
+                self.by_id.insert(id.to_owned(), idx);
+                self.leaf = Leaf::Id(id.to_owned());
+            } else {
+                self.leaf = Leaf::Tail;
+            }
+            if let Some((target, label)) = entry.label_fields() {
+                match (target, label) {
+                    (Some(t), Some(l)) if !l.is_empty() => {
+                        let ts = entry.timestamp().unwrap_or("").to_owned();
+                        self.labels.set(t.to_owned(), l.to_owned(), ts);
+                    }
+                    (Some(t), _) => {
+                        self.labels.delete(t);
+                    }
+                    _ => {}
+                }
+            }
+        }
+    }
+
+    fn rewrite_file(&self) -> Result<(), SessionError> {
+        if !self.persist {
+            return Ok(());
+        }
+        let Some(ref path) = self.session_file else {
+            return Ok(());
+        };
+        atomic_replace_file(Path::new(path), |file| {
+            for entry in &self.file_entries {
+                let line = file_entry_to_line(entry)?;
+                writeln!(file, "{line}").map_err(|source| SessionError::Io {
+                    path: path.clone(),
+                    source,
+                })?;
+            }
+            Ok(())
+        })
+    }
+
+    fn persist_entry_at(&mut self, idx: usize) -> Result<(), SessionError> {
+        if !self.persist {
+            return Ok(());
+        }
+        let Some(path) = self.session_file.clone() else {
+            return Ok(());
+        };
+
+        let has_assistant = self
+            .file_entries
+            .iter()
+            .filter_map(FileEntry::entry)
+            .any(SessionEntry::is_assistant_message);
+
+        if !has_assistant {
+            if self.flushed
+                && let Some(entry) = self.file_entries.get(idx).and_then(FileEntry::entry)
+            {
+                append_line(&path, entry)?;
+            }
+            // else: hold in memory (flushed stays false)
+            return Ok(());
+        }
+
+        if self.flushed
+            && let Some(entry) = self.file_entries.get(idx).and_then(FileEntry::entry)
+        {
+            append_line(&path, entry)?;
+        } else {
+            // Exclusive create (wx) + write ALL lines. Serialize before creating
+            // the destination so JSON failures cannot leave a partial file.
+            let mut contents = String::new();
+            for fe in &self.file_entries {
+                contents.push_str(&file_entry_to_line(fe)?);
+                contents.push('\n');
+            }
+            let mut file = OpenOptions::new()
+                .write(true)
+                .create_new(true)
+                .open(&path)
+                .map_err(|source| SessionError::Io {
+                    path: path.clone(),
+                    source,
+                })?;
+            if let Err(source) = file
+                .write_all(contents.as_bytes())
+                .and_then(|()| file.sync_all())
+            {
+                drop(file);
+                let _ = fs::remove_file(&path);
+                return Err(SessionError::Io {
+                    path: path.clone(),
+                    source,
+                });
+            }
+            self.flushed = true;
+        }
+        Ok(())
+    }
+
+    fn append_entry(&mut self, entry: SessionEntry) -> Result<String, SessionError> {
+        let id = entry.id().unwrap_or("").to_owned();
+        let previous_leaf = self.leaf.clone();
+        let previous_flushed = self.flushed;
+        let previous_index = if id.is_empty() {
+            None
+        } else {
+            self.by_id.insert(id.clone(), self.file_entries.len())
+        };
+        self.file_entries.push(FileEntry::Entry(entry));
+        let idx = self.file_entries.len() - 1;
+        self.leaf = if id.is_empty() {
+            Leaf::Tail
+        } else {
+            Leaf::Id(id.clone())
+        };
+        if let Err(err) = self.persist_entry_at(idx) {
+            self.file_entries.pop();
+            if !id.is_empty() {
+                match previous_index {
+                    Some(index) => {
+                        self.by_id.insert(id.clone(), index);
+                    }
+                    None => {
+                        self.by_id.remove(&id);
+                    }
+                }
+            }
+            self.leaf = previous_leaf;
+            self.flushed = previous_flushed;
+            return Err(err);
+        }
+        Ok(id)
+    }
+
+    fn next_id(&self) -> String {
+        generate_id(|c| self.by_id.contains_key(c))
+    }
+
+    // ----- accessors -----
+
+    /// Whether this manager persists to disk.
+    #[must_use]
+    pub const fn is_persisted(&self) -> bool {
+        self.persist
+    }
+
+    /// Resolved session working directory.
+    #[must_use]
+    pub fn get_cwd(&self) -> &str {
+        &self.cwd
+    }
+
+    /// Session directory used for new/branch files.
+    #[must_use]
+    pub fn get_session_dir(&self) -> &str {
+        &self.session_dir
+    }
+
+    /// True when `session_dir` equals the default encoded path for `cwd`.
+    #[must_use]
+    pub fn uses_default_session_dir(&self) -> bool {
+        let default = default_session_dir_path(&self.cwd, &get_agent_dir());
+        self.session_dir == default
+    }
+
+    /// Session UUID.
+    #[must_use]
+    pub fn get_session_id(&self) -> &str {
+        &self.session_id
+    }
+
+    /// Session file path, if any.
+    #[must_use]
+    pub fn get_session_file(&self) -> Option<&str> {
+        self.session_file.as_deref()
+    }
+
+    /// Current leaf entry id (`None` when null/tail).
+    #[must_use]
+    pub fn get_leaf_id(&self) -> Option<&str> {
+        match &self.leaf {
+            Leaf::Id(id) => Some(id.as_str()),
+            Leaf::Null | Leaf::Tail => None,
+        }
+    }
+
+    /// Current leaf entry.
+    #[must_use]
+    pub fn get_leaf_entry(&self) -> Option<&SessionEntry> {
+        match &self.leaf {
+            Leaf::Id(id) => self.get_entry(id),
+            Leaf::Null | Leaf::Tail => None,
+        }
+    }
+
+    /// Look up an entry by id.
+    #[must_use]
+    pub fn get_entry(&self, id: &str) -> Option<&SessionEntry> {
+        self.by_id
+            .get(id)
+            .and_then(|&i| self.file_entries.get(i))
+            .and_then(FileEntry::entry)
+    }
+
+    /// Direct children of `parent_id` (file order).
+    #[must_use]
+    pub fn get_children(&self, parent_id: &str) -> Vec<&SessionEntry> {
+        self.file_entries
+            .iter()
+            .filter_map(FileEntry::entry)
+            .filter(|e| e.parent_id() == Some(parent_id))
+            .collect()
+    }
+
+    /// Resolved label for an entry id.
+    #[must_use]
+    pub fn get_label(&self, id: &str) -> Option<&str> {
+        self.labels.get(id)
+    }
+
+    /// Session header, if typed parse succeeded.
+    #[must_use]
+    pub fn get_header(&self) -> Option<&SessionHeader> {
+        self.file_entries.iter().find_map(FileEntry::header)
+    }
+
+    /// All non-header entries (file order).
+    #[must_use]
+    pub fn get_entries(&self) -> Vec<&SessionEntry> {
+        self.file_entries
+            .iter()
+            .filter_map(FileEntry::entry)
+            .collect()
+    }
+
+    /// Latest session display name (empty clears).
+    #[must_use]
+    pub fn get_session_name(&self) -> Option<String> {
+        for entry in self.get_entries().into_iter().rev() {
+            if let Some(name) = entry.session_info_name() {
+                return name
+                    .map(str::trim)
+                    .filter(|s| !s.is_empty())
+                    .map(str::to_owned);
+            }
+        }
+        None
+    }
+
+    // ----- append -----
+
+    /// Append a transcript message as child of the current leaf.
+    ///
+    /// # Errors
+    ///
+    /// Returns JSON serialization errors or IO errors when persisting.
+    pub fn append_message(&mut self, message: &AgentMessage) -> Result<String, SessionError> {
+        let id = self.next_id();
+        let parent = self.leaf_parent();
+        let timestamp = now_iso();
+        let value = serde_json::json!({
+            "type": "message",
+            "id": id,
+            "parentId": parent,
+            "timestamp": timestamp,
+            "message": message_to_value(message)?,
+        });
+        let entry: SessionEntry = serde_json::from_value(value)?;
+        self.append_entry(entry)
+    }
+
+    /// Append a thinking-level change.
+    ///
+    /// # Errors
+    ///
+    /// Returns JSON or IO errors when encoding or persisting the entry.
+    pub fn append_thinking_level_change(
+        &mut self,
+        thinking_level: &str,
+    ) -> Result<String, SessionError> {
+        let id = self.next_id();
+        let value = serde_json::json!({
+            "type": "thinking_level_change",
+            "id": id,
+            "parentId": self.leaf_parent(),
+            "timestamp": now_iso(),
+            "thinkingLevel": thinking_level,
+        });
+        let entry: SessionEntry = serde_json::from_value(value)?;
+        self.append_entry(entry)
+    }
+
+    /// Append a model change.
+    ///
+    /// # Errors
+    ///
+    /// Returns JSON or IO errors when encoding or persisting the entry.
+    pub fn append_model_change(
+        &mut self,
+        provider: &str,
+        model_id: &str,
+    ) -> Result<String, SessionError> {
+        let id = self.next_id();
+        let value = serde_json::json!({
+            "type": "model_change",
+            "id": id,
+            "parentId": self.leaf_parent(),
+            "timestamp": now_iso(),
+            "provider": provider,
+            "modelId": model_id,
+        });
+        let entry: SessionEntry = serde_json::from_value(value)?;
+        self.append_entry(entry)
+    }
+
+    /// Append a compaction summary.
+    ///
+    /// # Errors
+    ///
+    /// Returns JSON or IO errors when encoding or persisting the entry.
+    pub fn append_compaction(
+        &mut self,
+        summary: &str,
+        first_kept_entry_id: &str,
+        tokens_before: i64,
+        details: Option<Value>,
+        from_hook: Option<bool>,
+    ) -> Result<String, SessionError> {
+        let id = self.next_id();
+        let mut value = serde_json::json!({
+            "type": "compaction",
+            "id": id,
+            "parentId": self.leaf_parent(),
+            "timestamp": now_iso(),
+            "summary": summary,
+            "firstKeptEntryId": first_kept_entry_id,
+            "tokensBefore": tokens_before,
+        });
+        if let Some(d) = details
+            && let Some(obj) = value.as_object_mut()
+        {
+            obj.insert("details".to_owned(), d);
+        }
+        if let Some(fh) = from_hook
+            && let Some(obj) = value.as_object_mut()
+        {
+            obj.insert("fromHook".to_owned(), Value::Bool(fh));
+        }
+        let entry: SessionEntry = serde_json::from_value(value)?;
+        self.append_entry(entry)
+    }
+
+    /// Append an extension custom state entry (not in LLM context).
+    ///
+    /// # Errors
+    ///
+    /// Returns JSON or IO errors when encoding or persisting the entry.
+    pub fn append_custom_entry(
+        &mut self,
+        custom_type: &str,
+        data: Option<Value>,
+    ) -> Result<String, SessionError> {
+        let id = self.next_id();
+        let mut value = serde_json::json!({
+            "type": "custom",
+            "customType": custom_type,
+            "id": id,
+            "parentId": self.leaf_parent(),
+            "timestamp": now_iso(),
+        });
+        if let Some(d) = data
+            && let Some(obj) = value.as_object_mut()
+        {
+            obj.insert("data".to_owned(), d);
+        }
+        let entry: SessionEntry = serde_json::from_value(value)?;
+        self.append_entry(entry)
+    }
+
+    /// Append a session display-name entry (sanitized).
+    ///
+    /// # Errors
+    ///
+    /// Returns JSON or IO errors when encoding or persisting the entry.
+    pub fn append_session_info(&mut self, name: &str) -> Result<String, SessionError> {
+        // TS: name.replace(/[\r\n]+/g, " ").trim()
+        let sanitized = {
+            let mut result = String::with_capacity(name.len());
+            let mut chars = name.chars().peekable();
+            while let Some(c) = chars.next() {
+                if c == '\r' || c == '\n' {
+                    while matches!(chars.peek(), Some('\r' | '\n')) {
+                        chars.next();
+                    }
+                    result.push(' ');
+                } else {
+                    result.push(c);
+                }
+            }
+            result.trim().to_owned()
+        };
+        let id = self.next_id();
+        let value = serde_json::json!({
+            "type": "session_info",
+            "id": id,
+            "parentId": self.leaf_parent(),
+            "timestamp": now_iso(),
+            "name": sanitized,
+        });
+        let entry: SessionEntry = serde_json::from_value(value)?;
+        self.append_entry(entry)
+    }
+
+    /// Append an extension custom message (participates in LLM context).
+    ///
+    /// # Errors
+    ///
+    /// Returns JSON or IO errors when encoding or persisting the entry.
+    pub fn append_custom_message_entry(
+        &mut self,
+        custom_type: &str,
+        content: &CustomMessageContent,
+        display: bool,
+        details: Option<Value>,
+    ) -> Result<String, SessionError> {
+        let id = self.next_id();
+        let mut value = serde_json::json!({
+            "type": "custom_message",
+            "customType": custom_type,
+            "content": content,
+            "display": display,
+            "id": id,
+            "parentId": self.leaf_parent(),
+            "timestamp": now_iso(),
+        });
+        if let Some(d) = details
+            && let Some(obj) = value.as_object_mut()
+        {
+            obj.insert("details".to_owned(), d);
+        }
+        let entry: SessionEntry = serde_json::from_value(value)?;
+        self.append_entry(entry)
+    }
+
+    /// Set or clear a label on an entry. Empty/None clears.
+    ///
+    /// # Errors
+    ///
+    /// Returns not-found when `target_id` is missing, or JSON/IO errors on persist.
+    pub fn append_label_change(
+        &mut self,
+        target_id: &str,
+        label: Option<&str>,
+    ) -> Result<String, SessionError> {
+        if !self.by_id.contains_key(target_id) {
+            return Err(SessionError::EntryNotFound(target_id.to_owned()));
+        }
+        let id = self.next_id();
+        let timestamp = now_iso();
+        let mut value = serde_json::json!({
+            "type": "label",
+            "id": id,
+            "parentId": self.leaf_parent(),
+            "timestamp": timestamp,
+            "targetId": target_id,
+        });
+        if let Some(l) = label
+            && let Some(obj) = value.as_object_mut()
+        {
+            obj.insert("label".to_owned(), Value::String(l.to_owned()));
+        }
+        let entry: SessionEntry = serde_json::from_value(value)?;
+        let result_id = self.append_entry(entry)?;
+        match label {
+            Some(l) if !l.is_empty() => {
+                self.labels
+                    .set(target_id.to_owned(), l.to_owned(), timestamp);
+            }
+            _ => {
+                self.labels.delete(target_id);
+            }
+        }
+        Ok(result_id)
+    }
+
+    // ----- tree -----
+
+    /// Walk from `from_id` (or current leaf) to root; root→leaf order.
+    ///
+    /// Unlike context path construction, a missing id returns an empty path
+    /// (no last-entry fallback).
+    #[must_use]
+    pub fn get_branch(&self, from_id: Option<&str>) -> Vec<&SessionEntry> {
+        let start = from_id.or(match &self.leaf {
+            Leaf::Id(id) => Some(id.as_str()),
+            Leaf::Null | Leaf::Tail => None,
+        });
+        let Some(start_id) = start.filter(|s| !s.is_empty()) else {
+            return Vec::new();
+        };
+        let mut path = Vec::new();
+        let mut current = self.get_entry(start_id);
+        while let Some(entry) = current {
+            path.push(entry);
+            current = entry.parent_id().and_then(|pid| self.get_entry(pid));
+        }
+        path.reverse();
+        path
+    }
+
+    /// Compaction-aware context entries for the current leaf.
+    #[must_use]
+    pub fn build_context_entries(&self) -> Vec<&SessionEntry> {
+        let entries = self.get_entries();
+        let leaf = self.leaf_ref();
+        context::build_context_entries(&entries, leaf)
+    }
+
+    /// Full session context for the LLM (current leaf).
+    ///
+    /// # Errors
+    ///
+    /// Returns message conversion errors when projecting entries to context messages.
+    pub fn build_session_context(&self) -> Result<SessionContext, MessageConversionError> {
+        let entries = self.get_entries();
+        let leaf = self.leaf_ref();
+        context::build_session_context(&entries, leaf)
+    }
+
+    /// Session as a tree (orphans become roots; children sorted by timestamp).
+    #[must_use]
+    pub fn get_tree(&self) -> Vec<SessionTreeNode> {
+        let entries = self.get_entries();
+        let mut node_map: HashMap<String, SessionTreeNode> = HashMap::new();
+        let mut order: Vec<String> = Vec::new();
+
+        for entry in &entries {
+            let Some(id) = entry.id() else {
+                continue;
+            };
+            if !node_map.contains_key(id) {
+                order.push(id.to_owned());
+            }
+            let label = self.labels.get(id).map(str::to_owned);
+            let label_timestamp = self.labels.get_timestamp(id).map(str::to_owned);
+            node_map.insert(
+                id.to_owned(),
+                SessionTreeNode {
+                    entry: (*entry).clone(),
+                    children: Vec::new(),
+                    label,
+                    label_timestamp,
+                },
+            );
+        }
+
+        let mut roots: Vec<String> = Vec::new();
+        let mut child_links: Vec<(String, String)> = Vec::new(); // (parent, child)
+
+        for entry in &entries {
+            let Some(id) = entry.id() else {
+                continue;
+            };
+            match entry.parent_id() {
+                Some(pid) if pid != id && node_map.contains_key(pid) => {
+                    child_links.push((pid.to_owned(), id.to_owned()));
+                }
+                None | Some(_) => roots.push(id.to_owned()),
+            }
+        }
+
+        let mut children_of: HashMap<String, Vec<String>> = HashMap::new();
+        for (parent, child) in child_links {
+            children_of.entry(parent).or_default().push(child);
+        }
+
+        // Deduplicate roots (dup ids).
+        let mut seen = HashSet::new();
+        let mut root_nodes = Vec::new();
+        for id in roots {
+            if seen.insert(id.clone())
+                && let Some(node) = assemble_tree_node(&id, &node_map, &children_of)
+            {
+                root_nodes.push(node);
+            }
+        }
+        // Sort root children already done per-node; roots themselves unsorted in TS
+        // (file order of first encounter). Keep as-is.
+        let _ = order;
+        root_nodes
+    }
+
+    // ----- branching -----
+
+    /// Move the leaf pointer to an existing entry.
+    ///
+    /// # Errors
+    ///
+    /// Returns not-found when `branch_from_id` is missing.
+    pub fn branch(&mut self, branch_from_id: &str) -> Result<(), SessionError> {
+        if !self.by_id.contains_key(branch_from_id) {
+            return Err(SessionError::EntryNotFound(branch_from_id.to_owned()));
+        }
+        self.leaf = Leaf::Id(branch_from_id.to_owned());
+        Ok(())
+    }
+
+    /// Reset the leaf pointer to null (next append is a new root).
+    pub fn reset_leaf(&mut self) {
+        self.leaf = Leaf::Null;
+    }
+
+    /// Branch from an entry (or null root) and append a `branch_summary` child.
+    ///
+    /// # Errors
+    ///
+    /// Returns not-found when a branch id is missing, or JSON/IO errors on persist.
+    pub fn branch_with_summary(
+        &mut self,
+        branch_from_id: Option<&str>,
+        summary: &str,
+        details: Option<Value>,
+        from_hook: Option<bool>,
+    ) -> Result<String, SessionError> {
+        if let Some(id) = branch_from_id {
+            if !self.by_id.contains_key(id) {
+                return Err(SessionError::EntryNotFound(id.to_owned()));
+            }
+            self.leaf = Leaf::Id(id.to_owned());
+        } else {
+            self.leaf = Leaf::Null;
+        }
+        let id = self.next_id();
+        let from_id = branch_from_id.unwrap_or("root");
+        let mut value = serde_json::json!({
+            "type": "branch_summary",
+            "id": id,
+            "parentId": branch_from_id,
+            "timestamp": now_iso(),
+            "fromId": from_id,
+            "summary": summary,
+        });
+        if let Some(d) = details
+            && let Some(obj) = value.as_object_mut()
+        {
+            obj.insert("details".to_owned(), d);
+        }
+        if let Some(fh) = from_hook
+            && let Some(obj) = value.as_object_mut()
+        {
+            obj.insert("fromHook".to_owned(), Value::Bool(fh));
+        }
+        let entry: SessionEntry = serde_json::from_value(value)?;
+        self.append_entry(entry)
+    }
+
+    /// Create a new session containing only the path to `leaf_id`.
+    ///
+    /// Strips label entries, rechains parent ids, re-appends resolved labels
+    /// with new ids but preserved timestamps. Returns the new file path when
+    /// persisting (or `None` in-memory).
+    ///
+    /// # Errors
+    ///
+    /// Returns not-found when `leaf_id` is missing, or JSON/IO errors while writing.
+    pub fn create_branched_session(
+        &mut self,
+        leaf_id: &str,
+    ) -> Result<Option<String>, SessionError> {
+        let previous_session_file = self.session_file.clone();
+        let path: Vec<SessionEntry> = self
+            .get_branch(Some(leaf_id))
+            .into_iter()
+            .cloned()
+            .collect();
+        if path.is_empty() {
+            return Err(SessionError::EntryNotFound(leaf_id.to_owned()));
+        }
+
+        let mut path_without_labels: Vec<SessionEntry> = Vec::new();
+        let mut path_parent: Option<String> = None;
+        for entry in path {
+            if entry.discriminant() == "label" {
+                continue;
+            }
+            let mut cloned = entry;
+            cloned.set_parent_id(path_parent.clone());
+            path_parent = cloned.id().map(str::to_owned);
+            path_without_labels.push(cloned);
+        }
+
+        let new_session_id = create_session_id();
+        let timestamp = now_iso();
+        let file_name = session_file_name(&timestamp, &new_session_id);
+        let new_session_file = Path::new(&self.session_dir)
+            .join(&file_name)
+            .to_string_lossy()
+            .into_owned();
+
+        let header = SessionHeader::new(
+            new_session_id.clone(),
+            timestamp,
+            self.cwd.clone(),
+            if self.persist {
+                previous_session_file
+            } else {
+                None
+            },
+        );
+
+        let path_ids: HashSet<String> = path_without_labels
+            .iter()
+            .filter_map(|e| e.id().map(str::to_owned))
+            .collect();
+
+        let labels_to_write: Vec<(String, String, String)> = self
+            .labels
+            .iter()
+            .filter(|(t, _, _)| path_ids.contains(*t))
+            .map(|(t, l, ts)| (t.to_owned(), l.to_owned(), ts.to_owned()))
+            .collect();
+
+        let mut collision: HashSet<String> = path_ids;
+        let mut parent = path_without_labels
+            .last()
+            .and_then(|e| e.id().map(str::to_owned));
+        let mut label_entries: Vec<SessionEntry> = Vec::new();
+        for (target, label, ts) in labels_to_write {
+            let lid = generate_id(|c| collision.contains(c));
+            collision.insert(lid.clone());
+            let value = serde_json::json!({
+                "type": "label",
+                "id": lid,
+                "parentId": parent,
+                "timestamp": ts,
+                "targetId": target,
+                "label": label,
+            });
+            let entry: SessionEntry = serde_json::from_value(value)?;
+            parent = entry.id().map(str::to_owned);
+            label_entries.push(entry);
+        }
+
+        let mut file_entries = vec![FileEntry::Header(header)];
+        file_entries.extend(path_without_labels.into_iter().map(FileEntry::Entry));
+        file_entries.extend(label_entries.into_iter().map(FileEntry::Entry));
+
+        self.file_entries = file_entries;
+        self.session_id = new_session_id;
+        self.build_index();
+
+        if self.persist {
+            self.session_file = Some(new_session_file.clone());
+            let has_assistant = self
+                .file_entries
+                .iter()
+                .filter_map(FileEntry::entry)
+                .any(SessionEntry::is_assistant_message);
+            if has_assistant {
+                self.rewrite_file()?;
+                self.flushed = true;
+            } else {
+                self.flushed = false;
+            }
+            Ok(Some(new_session_file))
+        } else {
+            // In-memory: session_file and flushed untouched.
+            Ok(None)
+        }
+    }
+
+    // ----- static factories -----
+
+    /// Create a new persisted session.
+    ///
+    /// # Errors
+    ///
+    /// Returns validation errors for custom ids, or IO errors creating the session dir.
+    pub fn create(
+        cwd: &str,
+        session_dir: Option<&str>,
+        options: Option<NewSessionOptions>,
+    ) -> Result<Self, SessionError> {
+        let dir = match session_dir {
+            Some(d) => path_to_string(&normalize_path(d, PathInputOptions::new())),
+            None => default_session_dir(cwd, &get_agent_dir())?,
+        };
+        Self::construct(cwd, &dir, None, true, options)
+    }
+
+    /// Open a specific session file.
+    ///
+    /// # Errors
+    ///
+    /// Returns IO/invalid-session errors while loading the file, or validation/IO
+    /// errors while constructing the manager.
+    pub fn open(
+        path: &str,
+        session_dir: Option<&str>,
+        cwd_override: Option<&str>,
+    ) -> Result<Self, SessionError> {
+        let resolved = path_to_string(&resolve_path(path));
+        let entries = load_entries_from_file(Path::new(&resolved));
+        let header_cwd = entries
+            .iter()
+            .find_map(FileEntry::header)
+            .and_then(|h| h.cwd.clone());
+        let cwd = cwd_override.map_or_else(
+            || {
+                header_cwd.unwrap_or_else(|| {
+                    std::env::current_dir()
+                        .map_or_else(|_| ".".to_owned(), |p| p.to_string_lossy().into_owned())
+                })
+            },
+            str::to_owned,
+        );
+        let dir = match session_dir {
+            Some(d) => path_to_string(&normalize_path(d, PathInputOptions::new())),
+            None => Path::new(&resolved)
+                .parent()
+                .map_or_else(|| ".".to_owned(), |p| p.to_string_lossy().into_owned()),
+        };
+        Self::construct(&cwd, &dir, Some(resolved), true, None)
+    }
+
+    /// Continue the most recent session, or create new if none.
+    ///
+    /// # Errors
+    ///
+    /// Returns IO errors discovering sessions or constructing a new manager.
+    pub fn continue_recent(cwd: &str, session_dir: Option<&str>) -> Result<Self, SessionError> {
+        let dir = match session_dir {
+            Some(d) => path_to_string(&normalize_path(d, PathInputOptions::new())),
+            None => default_session_dir(cwd, &get_agent_dir())?,
+        };
+        let filter_cwd =
+            session_dir.is_some() && dir != default_session_dir_path(cwd, &get_agent_dir());
+        let most_recent =
+            find_most_recent_session(Path::new(&dir), if filter_cwd { Some(cwd) } else { None });
+        match most_recent {
+            Some(f) => Self::construct(
+                cwd,
+                &dir,
+                Some(f.to_string_lossy().into_owned()),
+                true,
+                None,
+            ),
+            None => Self::construct(cwd, &dir, None, true, None),
+        }
+    }
+
+    /// Create an in-memory session (no file persistence).
+    ///
+    /// # Errors
+    ///
+    /// Returns a validation error when a custom session id is invalid.
+    pub fn in_memory(
+        cwd: Option<&str>,
+        options: Option<NewSessionOptions>,
+    ) -> Result<Self, SessionError> {
+        let cwd = cwd.map_or_else(
+            || {
+                std::env::current_dir()
+                    .map_or_else(|_| ".".to_owned(), |p| p.to_string_lossy().into_owned())
+            },
+            str::to_owned,
+        );
+        Self::construct(&cwd, "", None, false, options)
+    }
+
+    /// Fork a session into a new file under `target_cwd`.
+    ///
+    /// Copies **all** non-header source entries as-is (IDs preserved) under a
+    /// new v3 header with `parentSession = source path`.
+    ///
+    /// # Errors
+    ///
+    /// Returns fork-source/validation errors, or IO/JSON errors while writing the fork.
+    pub fn fork_from(
+        source_path: &str,
+        target_cwd: &str,
+        session_dir: Option<&str>,
+        options: Option<NewSessionOptions>,
+    ) -> Result<Self, SessionError> {
+        let resolved_source = path_to_string(&resolve_path(source_path));
+        let resolved_target = path_to_string(&resolve_path(target_cwd));
+        let values = load_values_from_file(Path::new(&resolved_source)).map_err(|source| {
+            SessionError::Io {
+                path: resolved_source.clone(),
+                source,
+            }
+        })?;
+        if values.is_empty() {
+            return Err(SessionError::ForkSourceEmpty(resolved_source));
+        }
+        if !values
+            .iter()
+            .any(|v| v.get("type").and_then(Value::as_str) == Some("session"))
+        {
+            return Err(SessionError::ForkSourceNoHeader(resolved_source));
+        }
+
+        let dir = match session_dir {
+            Some(d) => path_to_string(&normalize_path(d, PathInputOptions::new())),
+            None => default_session_dir(&resolved_target, &get_agent_dir())?,
+        };
+        if !path_exists(Path::new(&dir)) {
+            fs::create_dir_all(&dir).map_err(|source| SessionError::Io {
+                path: dir.clone(),
+                source,
+            })?;
+        }
+
+        if let Some(id) = options.as_ref().and_then(|opts| opts.id.as_ref()) {
+            assert_valid_session_id(id)?;
+        }
+        let options = options.unwrap_or_default();
+        let new_session_id = options.id.unwrap_or_else(create_session_id);
+        let timestamp = now_iso();
+        let file_name = session_file_name(&timestamp, &new_session_id);
+        let new_session_file = Path::new(&dir)
+            .join(file_name)
+            .to_string_lossy()
+            .into_owned();
+
+        let header = SessionHeader::new(
+            new_session_id,
+            timestamp,
+            resolved_target.clone(),
+            Some(resolved_source),
+        );
+
+        let mut file = OpenOptions::new()
+            .write(true)
+            .create_new(true)
+            .open(&new_session_file)
+            .map_err(|source| SessionError::Io {
+                path: new_session_file.clone(),
+                source,
+            })?;
+        writeln!(file, "{}", serde_json::to_string(&header)?).map_err(|source| {
+            SessionError::Io {
+                path: new_session_file.clone(),
+                source,
+            }
+        })?;
+        for value in &values {
+            if value.get("type").and_then(Value::as_str) == Some("session") {
+                continue;
+            }
+            writeln!(file, "{}", serde_json::to_string(value)?).map_err(|source| {
+                SessionError::Io {
+                    path: new_session_file.clone(),
+                    source,
+                }
+            })?;
+        }
+        drop(file);
+
+        Self::construct(&resolved_target, &dir, Some(new_session_file), true, None)
+    }
+
+    /// List sessions for a cwd.
+    ///
+    /// # Errors
+    ///
+    /// Returns IO errors while enumerating or reading session files.
+    pub async fn list(
+        cwd: &str,
+        session_dir: Option<&str>,
+        on_progress: Option<SessionListProgress<'_>>,
+    ) -> Result<Vec<SessionInfo>, SessionError> {
+        let dir = match session_dir {
+            Some(d) => path_to_string(&normalize_path(d, PathInputOptions::new())),
+            None => default_session_dir(cwd, &get_agent_dir())?,
+        };
+        let filter_cwd =
+            session_dir.is_some() && dir != default_session_dir_path(cwd, &get_agent_dir());
+        Ok(list_sessions_for_cwd(cwd, Path::new(&dir), filter_cwd, on_progress).await)
+    }
+
+    /// List all sessions under the agent sessions root (or a custom dir).
+    pub async fn list_all(
+        session_dir: Option<&str>,
+        on_progress: Option<SessionListProgress<'_>>,
+    ) -> Vec<SessionInfo> {
+        let custom =
+            session_dir.map(|d| path_to_string(&normalize_path(d, PathInputOptions::new())));
+        let root = get_sessions_dir();
+        list_all_sessions(&root, custom.as_deref().map(Path::new), on_progress).await
+    }
+
+    // ----- helpers -----
+
+    fn leaf_parent(&self) -> Option<String> {
+        match &self.leaf {
+            Leaf::Id(id) => Some(id.clone()),
+            Leaf::Null | Leaf::Tail => None,
+        }
+    }
+
+    fn leaf_ref(&self) -> LeafRef<'_> {
+        match &self.leaf {
+            Leaf::Null => LeafRef::Null,
+            Leaf::Id(id) => LeafRef::Id(id.as_str()),
+            Leaf::Tail => LeafRef::Last,
+        }
+    }
+}
+
+impl SessionCwdSource for SessionManager {
+    fn get_cwd(&self) -> &str {
+        &self.cwd
+    }
+    fn get_session_file(&self) -> Option<&str> {
+        self.session_file.as_deref()
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Path encoding / helpers
+// ---------------------------------------------------------------------------
+
+/// Encode a resolved cwd into a safe session-directory name segment.
+///
+/// Algorithm: strip one leading `/` or `\`, replace `/`, `\`, `:` with `-`,
+/// wrap in `--…--`.
+#[must_use]
+pub fn encode_cwd_for_session_dir(resolved_cwd: &str) -> String {
+    let stripped = resolved_cwd
+        .strip_prefix('/')
+        .or_else(|| resolved_cwd.strip_prefix('\\'))
+        .unwrap_or(resolved_cwd);
+    let safe: String = stripped
+        .chars()
+        .map(|c| {
+            if matches!(c, '/' | '\\' | ':') {
+                '-'
+            } else {
+                c
+            }
+        })
+        .collect();
+    format!("--{safe}--")
+}
+
+/// Default session directory path for a cwd (does not create).
+#[must_use]
+pub fn default_session_dir_path(cwd: &str, agent_dir: &Path) -> String {
+    let resolved_cwd = path_to_string(&resolve_path(cwd));
+    let resolved_agent = path_to_string(&resolve_path(agent_dir.to_string_lossy()));
+    let name = encode_cwd_for_session_dir(&resolved_cwd);
+    Path::new(&resolved_agent)
+        .join("sessions")
+        .join(name)
+        .to_string_lossy()
+        .into_owned()
+}
+
+/// Default session directory for a cwd (creates if missing).
+///
+/// # Errors
+///
+/// Returns IO errors when creating the directory.
+pub fn default_session_dir(cwd: &str, agent_dir: &Path) -> Result<String, SessionError> {
+    let dir = default_session_dir_path(cwd, agent_dir);
+    if !path_exists(Path::new(&dir)) {
+        fs::create_dir_all(&dir).map_err(|source| SessionError::Io {
+            path: dir.clone(),
+            source,
+        })?;
+    }
+    Ok(dir)
+}
+
+fn session_file_name(timestamp_iso: &str, session_id: &str) -> String {
+    let file_ts: String = timestamp_iso
+        .chars()
+        .map(|c| if c == ':' || c == '.' { '-' } else { c })
+        .collect();
+    format!("{file_ts}_{session_id}.jsonl")
+}
+
+fn path_to_string(path: &Path) -> String {
+    path.to_string_lossy().into_owned()
+}
+
+fn append_line(path: &str, entry: &SessionEntry) -> Result<(), SessionError> {
+    let line = session_entry_to_line(entry)?;
+    let mut record = String::with_capacity(line.len() + 1);
+    record.push_str(&line);
+    record.push('\n');
+    let mut file = OpenOptions::new()
+        .append(true)
+        .create(true)
+        .open(path)
+        .map_err(|source| SessionError::Io {
+            path: path.to_owned(),
+            source,
+        })?;
+    file.write_all(record.as_bytes())
+        .map_err(|source| SessionError::Io {
+            path: path.to_owned(),
+            source,
+        })
+}
+
+fn message_to_value(message: &AgentMessage) -> Result<Value, SessionError> {
+    Ok(serde_json::to_value(message)?)
+}
+
+static TEMP_FILE_COUNTER: AtomicU64 = AtomicU64::new(0);
+
+fn atomic_replace_file(
+    path: &Path,
+    write_contents: impl FnOnce(&mut File) -> Result<(), SessionError>,
+) -> Result<(), SessionError> {
+    let path_text = path.to_string_lossy().into_owned();
+    let parent = path.parent().unwrap_or_else(|| Path::new("."));
+    let file_name = path
+        .file_name()
+        .and_then(|name| name.to_str())
+        .unwrap_or("session");
+    let mut temp_path = PathBuf::new();
+    let mut temp_file = None;
+    for _ in 0..100 {
+        let suffix = TEMP_FILE_COUNTER.fetch_add(1, Ordering::Relaxed);
+        temp_path = parent.join(format!(".{file_name}.tmp-{}-{suffix}", std::process::id()));
+        match OpenOptions::new()
+            .write(true)
+            .create_new(true)
+            .open(&temp_path)
+        {
+            Ok(file) => {
+                temp_file = Some(file);
+                break;
+            }
+            Err(err) if err.kind() == io::ErrorKind::AlreadyExists => {}
+            Err(source) => {
+                return Err(SessionError::Io {
+                    path: temp_path.to_string_lossy().into_owned(),
+                    source,
+                });
+            }
+        }
+    }
+    let Some(mut file) = temp_file else {
+        return Err(SessionError::Io {
+            path: path_text,
+            source: io::Error::new(
+                io::ErrorKind::AlreadyExists,
+                "unable to allocate temporary session file",
+            ),
+        });
+    };
+
+    let result = write_contents(&mut file)
+        .and_then(|()| {
+            file.flush().map_err(|source| SessionError::Io {
+                path: temp_path.to_string_lossy().into_owned(),
+                source,
+            })
+        })
+        .and_then(|()| {
+            file.sync_all().map_err(|source| SessionError::Io {
+                path: temp_path.to_string_lossy().into_owned(),
+                source,
+            })
+        });
+    drop(file);
+    if let Err(err) = result {
+        let _ = fs::remove_file(&temp_path);
+        return Err(err);
+    }
+
+    if let Err(source) = fs::rename(&temp_path, path) {
+        let _ = fs::remove_file(&temp_path);
+        return Err(SessionError::Io {
+            path: path_text,
+            source,
+        });
+    }
+    sync_parent_directory(parent).map_err(|source| SessionError::Io {
+        path: parent.to_string_lossy().into_owned(),
+        source,
+    })
+}
+
+#[cfg(unix)]
+fn sync_parent_directory(parent: &Path) -> io::Result<()> {
+    File::open(parent)?.sync_all()
+}
+
+#[cfg(not(unix))]
+fn sync_parent_directory(_parent: &Path) -> io::Result<()> {
+    Ok(())
+}
+
+// ---------------------------------------------------------------------------
+// Tests
+// ---------------------------------------------------------------------------
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use pi_ai::{Message, TextContent, UserMessage, UserMessageContent};
+    use serde_json::json;
+    use tempfile::tempdir;
+
+    type TestResult = Result<(), Box<dyn std::error::Error>>;
+
+    #[test]
+    fn atomic_replace_failure_preserves_live_file() -> TestResult {
+        let dir = tempdir()?;
+        let file = dir.path().join("atomic.jsonl");
+        fs::write(&file, "original\n")?;
+
+        let result = atomic_replace_file(&file, |temp| {
+            temp.write_all(b"partial replacement\n")
+                .map_err(|source| SessionError::Io {
+                    path: file.to_string_lossy().into_owned(),
+                    source,
+                })?;
+            Err(SessionError::Io {
+                path: file.to_string_lossy().into_owned(),
+                source: io::Error::other("injected rewrite failure"),
+            })
+        });
+
+        assert!(result.is_err());
+        assert_eq!(fs::read_to_string(&file)?, "original\n");
+        let leftovers: Vec<_> = fs::read_dir(dir.path())?
+            .flatten()
+            .filter(|entry| entry.file_name().to_string_lossy().contains(".tmp-"))
+            .collect();
+        assert!(leftovers.is_empty(), "failed rewrite must remove temp file");
+        Ok(())
+    }
+
+    fn path_str(path: &Path) -> Result<&str, Box<dyn std::error::Error>> {
+        path.to_str()
+            .ok_or_else(|| "path is not valid UTF-8".into())
+    }
+
+    fn user_agent(text: &str, ts: i64) -> AgentMessage {
+        AgentMessage::Llm(Box::new(Message::User(UserMessage::new(
+            UserMessageContent::Text(text.to_owned()),
+            ts,
+        ))))
+    }
+
+    fn assistant_agent(text: &str, ts: i64) -> AgentMessage {
+        let mut msg =
+            pi_ai::AssistantMessage::new("anthropic-messages", "anthropic", "claude-test", ts);
+        msg.content = vec![pi_ai::AssistantContent::Text(TextContent::new(text))];
+        msg.usage = pi_ai::Usage {
+            input: 1,
+            output: 1,
+            cache_read: 0,
+            cache_write: 0,
+            total_tokens: 2,
+            cost: pi_ai::UsageCost::default(),
+            ..Default::default()
+        };
+        AgentMessage::Llm(Box::new(Message::Assistant(msg)))
+    }
+
+    #[test]
+    fn deferred_write_until_first_assistant() -> TestResult {
+        let dir = tempdir()?;
+        let mut sm =
+            SessionManager::create(path_str(dir.path())?, Some(path_str(dir.path())?), None)?;
+        let file = sm.get_session_file().ok_or("file")?.to_owned();
+        assert!(!path_exists(Path::new(&file)), "file must not exist yet");
+
+        sm.append_message(&user_agent("hello", 1))?;
+        assert!(!path_exists(Path::new(&file)), "still deferred after user");
+
+        sm.append_message(&assistant_agent("hi", 2))?;
+        assert!(path_exists(Path::new(&file)), "created on assistant");
+
+        let content = fs::read_to_string(&file)?;
+        let lines: Vec<_> = content.lines().filter(|l| !l.is_empty()).collect();
+        assert_eq!(lines.len(), 3); // header + user + assistant
+        let header: Value = serde_json::from_str(lines[0])?;
+        assert_eq!(header["type"], "session");
+        assert_eq!(header["version"], 3);
+        Ok(())
+    }
+
+    #[test]
+    fn append_prefix_stability() -> TestResult {
+        let dir = tempdir()?;
+        let file = dir.path().join("stable.jsonl");
+        let original = concat!(
+            r#"{"type":"session","version":3,"id":"sess-1","timestamp":"2025-01-01T00:00:00.000Z","cwd":"/tmp"}"#,
+            "\n",
+            r#"{"type":"message","id":"aaaaaaaa","parentId":null,"timestamp":"2025-01-01T00:00:01.000Z","message":{"role":"user","content":"hi","timestamp":1}}"#,
+            "\n",
+            r#"{"type":"message","id":"bbbbbbbb","parentId":"aaaaaaaa","timestamp":"2025-01-01T00:00:02.000Z","message":{"role":"assistant","content":[{"type":"text","text":"yo"}],"api":"test","provider":"test","model":"test","usage":{"input":1,"output":1,"cacheRead":0,"cacheWrite":0,"totalTokens":2,"cost":{"input":0,"output":0,"cacheRead":0,"cacheWrite":0,"total":0}},"stopReason":"stop","timestamp":2}}"#,
+            "\n",
+        );
+        fs::write(&file, original)?;
+
+        let mut sm = SessionManager::open(path_str(&file)?, Some(path_str(dir.path())?), None)?;
+        sm.append_message(&user_agent("next", 3))?;
+
+        let after = fs::read(&file)?;
+        assert!(
+            after.starts_with(original.as_bytes()),
+            "original lines must be byte-stable on append"
+        );
+        Ok(())
+    }
+
+    #[test]
+    fn failed_append_does_not_advance_tree_and_can_retry() -> TestResult {
+        let dir = tempdir()?;
+        let file = dir.path().join("retry.jsonl");
+        let original = concat!(
+            r#"{"type":"session","version":3,"id":"sess-1","timestamp":"2025-01-01T00:00:00.000Z","cwd":"/tmp"}"#,
+            "\n",
+            r#"{"type":"message","id":"aaaaaaaa","parentId":null,"timestamp":"2025-01-01T00:00:01.000Z","message":{"role":"user","content":"hi","timestamp":1}}"#,
+            "\n",
+            r#"{"type":"message","id":"bbbbbbbb","parentId":"aaaaaaaa","timestamp":"2025-01-01T00:00:02.000Z","message":{"role":"assistant","content":[{"type":"text","text":"yo"}],"api":"test","provider":"test","model":"test","usage":{"input":1,"output":1,"cacheRead":0,"cacheWrite":0,"totalTokens":2,"cost":{"input":0,"output":0,"cacheRead":0,"cacheWrite":0,"total":0}},"stopReason":"stop","timestamp":2}}"#,
+            "\n",
+        );
+        fs::write(&file, original)?;
+        let mut sm = SessionManager::open(path_str(&file)?, Some(path_str(dir.path())?), None)?;
+        let before_count = sm.get_entries().len();
+        let before_leaf = sm.get_leaf_id().map(str::to_owned);
+
+        let backup = dir.path().join("retry.backup");
+        fs::rename(&file, &backup)?;
+        fs::create_dir(&file)?;
+        let result = sm.append_message(&user_agent("retry me", 3));
+        assert!(matches!(result, Err(SessionError::Io { .. })));
+        assert_eq!(sm.get_entries().len(), before_count);
+        assert_eq!(sm.get_leaf_id(), before_leaf.as_deref());
+
+        fs::remove_dir(&file)?;
+        fs::rename(&backup, &file)?;
+        let id = sm.append_message(&user_agent("retry me", 3))?;
+        assert_eq!(sm.get_entries().len(), before_count + 1);
+        assert_eq!(sm.get_leaf_id(), Some(id.as_str()));
+        let persisted = fs::read_to_string(&file)?;
+        assert_eq!(persisted.matches("retry me").count(), 1);
+        Ok(())
+    }
+
+    #[test]
+    fn invalid_file_preserved() -> TestResult {
+        let dir = tempdir()?;
+        let file = dir.path().join("bad.jsonl");
+        let original = r#"{"type":"event","data":"not a session"}
+"#;
+        fs::write(&file, original)?;
+        let Err(err) = SessionManager::open(path_str(&file)?, Some(path_str(dir.path())?), None)
+        else {
+            return Err("expected invalid session error".into());
+        };
+        match err {
+            SessionError::InvalidSessionFile(p) => {
+                assert!(p.contains("bad.jsonl"));
+            }
+            other => return Err(format!("wrong error: {other}").into()),
+        }
+        assert_eq!(fs::read_to_string(&file)?, original);
+        Ok(())
+    }
+
+    #[test]
+    fn empty_file_gets_header() -> TestResult {
+        let dir = tempdir()?;
+        let file = dir.path().join("empty.jsonl");
+        fs::write(&file, "")?;
+        let sm = SessionManager::open(path_str(&file)?, Some(path_str(dir.path())?), None)?;
+        assert!(!sm.get_session_id().is_empty());
+        let content = fs::read_to_string(&file)?;
+        let header: Value = serde_json::from_str(content.trim())?;
+        assert_eq!(header["type"], "session");
+        assert_eq!(header["id"], sm.get_session_id());
+        Ok(())
+    }
+
+    #[test]
+    fn unknown_entry_roundtrip_and_excluded_from_context() -> TestResult {
+        let dir = tempdir()?;
+        let file = dir.path().join("unknown.jsonl");
+        let original = concat!(
+            r#"{"type":"session","version":3,"id":"sess-1","timestamp":"2025-01-01T00:00:00.000Z","cwd":"/tmp"}"#,
+            "\n",
+            r#"{"type":"future_thing","id":"u1","parentId":null,"timestamp":"2025-01-01T00:00:01.000Z","customField":{"a":1}}"#,
+            "\n",
+            r#"{"type":"message","id":"m1","parentId":"u1","timestamp":"2025-01-01T00:00:02.000Z","message":{"role":"user","content":"hi","timestamp":1}}"#,
+            "\n",
+            r#"{"type":"message","id":"m2","parentId":"m1","timestamp":"2025-01-01T00:00:03.000Z","message":{"role":"assistant","content":[{"type":"text","text":"yo"}],"api":"test","provider":"test","model":"test","usage":{"input":1,"output":1,"cacheRead":0,"cacheWrite":0,"totalTokens":2,"cost":{"input":0,"output":0,"cacheRead":0,"cacheWrite":0,"total":0}},"stopReason":"stop","timestamp":2}}"#,
+            "\n",
+        );
+        fs::write(&file, original)?;
+
+        let mut sm = SessionManager::open(path_str(&file)?, Some(path_str(dir.path())?), None)?;
+        assert!(sm.get_entry("u1").is_some());
+        assert_eq!(
+            sm.get_entry("u1").map(SessionEntry::discriminant),
+            Some("future_thing")
+        );
+        let ctx = sm.build_session_context()?;
+        assert_eq!(ctx.messages.len(), 2); // unknown excluded
+        assert_eq!(ctx.messages[0].role(), "user");
+
+        sm.append_message(&user_agent("next", 3))?;
+        let after = fs::read(&file)?;
+        assert!(after.starts_with(original.as_bytes()));
+        Ok(())
+    }
+
+    #[test]
+    fn leaf_is_last_file_order_not_deepest() -> TestResult {
+        let dir = tempdir()?;
+        let file = dir.path().join("leaf.jsonl");
+        // A → B → C (deep), then D child of A written last.
+        let content = concat!(
+            r#"{"type":"session","version":3,"id":"s","timestamp":"2025-01-01T00:00:00.000Z","cwd":"/tmp"}"#,
+            "\n",
+            r#"{"type":"message","id":"A","parentId":null,"timestamp":"2025-01-01T00:00:01.000Z","message":{"role":"user","content":"a","timestamp":1}}"#,
+            "\n",
+            r#"{"type":"message","id":"B","parentId":"A","timestamp":"2025-01-01T00:00:02.000Z","message":{"role":"assistant","content":[{"type":"text","text":"b"}],"api":"test","provider":"test","model":"test","usage":{"input":1,"output":1,"cacheRead":0,"cacheWrite":0,"totalTokens":2,"cost":{"input":0,"output":0,"cacheRead":0,"cacheWrite":0,"total":0}},"stopReason":"stop","timestamp":2}}"#,
+            "\n",
+            r#"{"type":"message","id":"C","parentId":"B","timestamp":"2025-01-01T00:00:03.000Z","message":{"role":"user","content":"c","timestamp":3}}"#,
+            "\n",
+            r#"{"type":"message","id":"D","parentId":"A","timestamp":"2025-01-01T00:00:04.000Z","message":{"role":"user","content":"d","timestamp":4}}"#,
+            "\n",
+        );
+        fs::write(&file, content)?;
+        let sm = SessionManager::open(path_str(&file)?, Some(path_str(dir.path())?), None)?;
+        assert_eq!(sm.get_leaf_id(), Some("D"));
+        let branch: Vec<&str> = sm.get_branch(None).iter().filter_map(|e| e.id()).collect();
+        assert_eq!(branch, vec!["A", "D"]);
+        Ok(())
+    }
+
+    #[test]
+    fn labels_and_rechain_on_branched_session() -> TestResult {
+        let mut sm = SessionManager::in_memory(Some("/tmp"), None)?;
+        let msg1 = sm.append_message(&user_agent("hello", 1))?;
+        sm.append_label_change(&msg1, Some("checkpoint"))?;
+        let model = sm.append_model_change("anthropic", "claude-test")?;
+        let msg2 = sm.append_message(&user_agent("followup", 2))?;
+
+        sm.create_branched_session(&msg2)?;
+
+        assert_eq!(
+            sm.get_entry(&model).and_then(|e| e.parent_id()),
+            Some(msg1.as_str())
+        );
+        assert_eq!(sm.get_label(&msg1), Some("checkpoint"));
+        Ok(())
+    }
+
+    #[test]
+    fn create_branched_session_defers_without_assistant() -> TestResult {
+        let dir = tempdir()?;
+        let mut sm =
+            SessionManager::create(path_str(dir.path())?, Some(path_str(dir.path())?), None)?;
+        let id1 = sm.append_message(&user_agent("first", 1))?;
+        sm.append_message(&assistant_agent("answer", 2))?;
+        sm.append_message(&user_agent("second", 3))?;
+        sm.append_message(&assistant_agent("answer2", 4))?;
+
+        let new_file = sm.create_branched_session(&id1)?.ok_or("path")?;
+        assert!(
+            !path_exists(Path::new(&new_file)),
+            "no assistant on path → deferred"
+        );
+        sm.append_custom_entry("preset-state", Some(json!({"name": "plan"})))?;
+        sm.append_message(&assistant_agent("new answer", 5))?;
+        assert!(path_exists(Path::new(&new_file)));
+
+        let content = fs::read_to_string(&new_file)?;
+        let records: Vec<Value> = content
+            .lines()
+            .filter(|l| !l.is_empty())
+            .map(serde_json::from_str)
+            .collect::<Result<_, _>>()?;
+        assert_eq!(
+            records
+                .iter()
+                .filter(|r| r.get("type").and_then(Value::as_str) == Some("session"))
+                .count(),
+            1
+        );
+        let ids: Vec<&str> = records
+            .iter()
+            .filter(|r| r.get("type").and_then(Value::as_str) != Some("session"))
+            .filter_map(|r| r.get("id").and_then(Value::as_str))
+            .collect();
+        let set: HashSet<&str> = ids.iter().copied().collect();
+        assert_eq!(set.len(), ids.len(), "no duplicate ids");
+        Ok(())
+    }
+
+    #[test]
+    fn create_branched_session_writes_with_assistant() -> TestResult {
+        let dir = tempdir()?;
+        let mut sm =
+            SessionManager::create(path_str(dir.path())?, Some(path_str(dir.path())?), None)?;
+        sm.append_message(&user_agent("first", 1))?;
+        let id2 = sm.append_message(&assistant_agent("answer", 2))?;
+        sm.append_message(&user_agent("second", 3))?;
+        sm.append_message(&assistant_agent("answer2", 4))?;
+
+        let new_file = sm.create_branched_session(&id2)?.ok_or("path")?;
+        assert!(path_exists(Path::new(&new_file)));
+        Ok(())
+    }
+
+    #[test]
+    fn fork_header_and_entries() -> TestResult {
+        let dir = tempdir()?;
+        let source = dir.path().join("source.jsonl");
+        fs::write(
+            &source,
+            concat!(
+                r#"{"type":"session","version":3,"id":"legacy-session-id","timestamp":"2025-01-01T00:00:00.000Z","cwd":"/old"}"#,
+                "\n",
+                r#"{"type":"message","id":"entry-1","parentId":null,"timestamp":"2025-01-01T00:00:01.000Z","message":{"role":"assistant","content":[{"type":"text","text":"hello"}],"api":"openai-responses","provider":"openai","model":"gpt-5.4","usage":{"input":0,"output":0,"cacheRead":0,"cacheWrite":0,"totalTokens":0,"cost":{"input":0,"output":0,"cacheRead":0,"cacheWrite":0,"total":0}},"stopReason":"stop","timestamp":1}}"#,
+                "\n",
+            ),
+        )?;
+        let forked = SessionManager::fork_from(
+            path_str(&source)?,
+            path_str(dir.path())?,
+            Some(path_str(dir.path())?),
+            None,
+        )?;
+        let header = forked.get_header().ok_or("header")?;
+        assert_ne!(header.id.as_deref(), Some("legacy-session-id"));
+        // uuid v7 shape
+        let id = header.id.as_deref().ok_or("id")?;
+        assert!(id.contains('-'));
+        assert_eq!(header.parent_session.as_deref(), Some(path_str(&source)?));
+        assert!(header.cwd.is_some());
+        // Entry ids preserved
+        assert!(forked.get_entry("entry-1").is_some());
+        Ok(())
+    }
+
+    #[test]
+    fn encoded_cwd() {
+        assert_eq!(
+            encode_cwd_for_session_dir("/home/user/project"),
+            "--home-user-project--"
+        );
+        assert_eq!(encode_cwd_for_session_dir("C:\\Users\\x"), "--C--Users-x--");
+        let agent = Path::new("/tmp/agent");
+        let dir = default_session_dir_path("/a/b", agent);
+        assert!(dir.ends_with("sessions/--a-b--") || dir.contains("sessions/--a-b--"));
+    }
+
+    #[test]
+    fn custom_session_id_and_filename() -> TestResult {
+        let dir = tempdir()?;
+        let sm = SessionManager::create(
+            path_str(dir.path())?,
+            Some(path_str(dir.path())?),
+            Some(NewSessionOptions {
+                id: Some("created-session-id".to_owned()),
+                parent_session: None,
+            }),
+        )?;
+        assert_eq!(sm.get_session_id(), "created-session-id");
+        let file = sm.get_session_file().ok_or("file")?;
+        let base = Path::new(file)
+            .file_name()
+            .ok_or("file name")?
+            .to_string_lossy();
+        // YYYY-MM-DDTHH-MM-SS-mmmZ_created-session-id.jsonl
+        assert!(base.ends_with("_created-session-id.jsonl"));
+        assert!(!path_exists(Path::new(file)));
+        Ok(())
+    }
+
+    #[test]
+    fn open_preserves_explicit_missing_session_path() -> TestResult {
+        let dir = tempdir()?;
+        let file = dir.path().join("explicit-session.jsonl");
+
+        let manager = SessionManager::open(path_str(&file)?, Some(path_str(dir.path())?), None)?;
+
+        assert_eq!(manager.get_session_file(), file.to_str());
+        assert!(!file.exists());
+        Ok(())
+    }
+
+    #[test]
+    fn invalid_session_ids_rejected() -> TestResult {
+        for id in [
+            "", "-abc", "abc-", "_abc", "abc_", ".abc", "abc.", "abc/def", "abc\\def", "abc def",
+        ] {
+            let mut sm = SessionManager::in_memory(Some("/tmp"), None)?;
+            let Err(err) = sm.new_session(Some(NewSessionOptions {
+                id: Some(id.to_owned()),
+                parent_session: None,
+            })) else {
+                return Err(format!("expected invalid id for {id:?}").into());
+            };
+            assert!(matches!(err, SessionError::InvalidSessionId));
+        }
+        Ok(())
+    }
+
+    #[test]
+    fn tree_and_branch() -> TestResult {
+        let mut sm = SessionManager::in_memory(Some("/tmp"), None)?;
+        let id1 = sm.append_message(&user_agent("1", 1))?;
+        let id2 = sm.append_message(&assistant_agent("2", 2))?;
+        let id3 = sm.append_message(&user_agent("3", 3))?;
+        assert_eq!(sm.get_leaf_id(), Some(id3.as_str()));
+
+        sm.branch(&id2)?;
+        let id4 = sm.append_message(&user_agent("4-branch", 4))?;
+        let tree = sm.get_tree();
+        assert_eq!(tree.len(), 1);
+        assert_eq!(tree[0].entry.id(), Some(id1.as_str()));
+        let node2 = &tree[0].children[0];
+        assert_eq!(node2.entry.id(), Some(id2.as_str()));
+        assert_eq!(node2.children.len(), 2);
+        let child_ids: HashSet<&str> = node2.children.iter().filter_map(|c| c.entry.id()).collect();
+        assert!(child_ids.contains(id3.as_str()));
+        assert!(child_ids.contains(id4.as_str()));
+        Ok(())
+    }
+
+    #[test]
+    fn model_and_thinking_context() -> TestResult {
+        let mut sm = SessionManager::in_memory(Some("/tmp"), None)?;
+        sm.append_message(&user_agent("hello", 1))?;
+        sm.append_thinking_level_change("high")?;
+        sm.append_model_change("openai", "gpt-4")?;
+        sm.append_message(&assistant_agent("hi", 2))?;
+        let ctx = sm.build_session_context()?;
+        assert_eq!(ctx.thinking_level, "high");
+        // assistant overwrites model_change
+        let model = ctx.model.ok_or("model")?;
+        assert_eq!(model.provider, "anthropic");
+        assert_eq!(model.model_id, "claude-test");
+        assert_eq!(ctx.messages.len(), 2);
+        Ok(())
+    }
+
+    #[test]
+    fn compaction_reconstruction() -> TestResult {
+        let mut sm = SessionManager::in_memory(Some("/tmp"), None)?;
+        let id1 = sm.append_message(&user_agent("first", 1))?;
+        sm.append_message(&assistant_agent("r1", 2))?;
+        let id3 = sm.append_message(&user_agent("second", 3))?;
+        sm.append_message(&assistant_agent("r2", 4))?;
+        sm.append_compaction("Summary of first two turns", &id3, 1000, None, None)?;
+        sm.append_message(&user_agent("third", 5))?;
+        sm.append_message(&assistant_agent("r3", 6))?;
+
+        let ids: Vec<&str> = sm
+            .build_context_entries()
+            .iter()
+            .filter_map(|e| e.id())
+            .collect();
+        // compaction + kept from id3 + post
+        let compaction_id = sm
+            .get_entries()
+            .iter()
+            .find(|e| e.discriminant() == "compaction")
+            .and_then(|e| e.id())
+            .ok_or("compaction entry")?;
+        assert_eq!(ids[0], compaction_id);
+        assert!(ids.contains(&id3.as_str()));
+        assert!(!ids.contains(&id1.as_str())); // summarized away
+
+        let ctx = sm.build_session_context()?;
+        assert_eq!(ctx.messages[0].role(), "compactionSummary");
+        assert_eq!(ctx.messages.len(), 5);
+        Ok(())
+    }
+
+    #[test]
+    fn custom_entries_in_tree_not_context() -> TestResult {
+        let mut sm = SessionManager::in_memory(Some("/tmp"), None)?;
+        let msg_id = sm.append_message(&user_agent("hello", 1))?;
+        let custom_id = sm.append_custom_entry("my_data", Some(json!({"foo": "bar"})))?;
+        sm.append_message(&assistant_agent("hi", 2))?;
+
+        let entries = sm.get_entries();
+        assert_eq!(entries.len(), 3);
+        assert_eq!(entries[1].discriminant(), "custom");
+        assert_eq!(entries[1].id(), Some(custom_id.as_str()));
+        assert_eq!(entries[1].parent_id(), Some(msg_id.as_str()));
+
+        let path = sm.get_branch(None);
+        assert_eq!(path.len(), 3);
+
+        let ctx = sm.build_session_context()?;
+        assert_eq!(ctx.messages.len(), 2);
+        Ok(())
+    }
+
+    #[test]
+    fn label_not_found_throws() -> TestResult {
+        let mut sm = SessionManager::in_memory(Some("/tmp"), None)?;
+        let Err(err) = sm.append_label_change("non-existent", Some("label")) else {
+            return Err("expected entry not found".into());
+        };
+        assert!(matches!(err, SessionError::EntryNotFound(_)));
+        assert_eq!(err.to_string(), "Entry non-existent not found");
+        Ok(())
+    }
+
+    #[test]
+    fn branch_not_found_throws() -> TestResult {
+        let mut sm = SessionManager::in_memory(Some("/tmp"), None)?;
+        sm.append_message(&user_agent("hello", 1))?;
+        let Err(err) = sm.branch("nonexistent") else {
+            return Err("expected entry not found".into());
+        };
+        assert_eq!(err.to_string(), "Entry nonexistent not found");
+        Ok(())
+    }
+}