swink-agent-memory 0.12.2

//! JSONL-based session persistence.
//!
//! Stores each session as a `.jsonl` file: the first line contains
//! [`SessionMeta`], and subsequent lines each contain one [`LlmMessage`].
//!
//! Concurrent writes to the same session may corrupt the file.
//! Callers are expected to enforce single-writer access.

use std::io::{self, BufRead, Read, Seek, SeekFrom, Write};
use std::path::{Path, PathBuf};

use swink_agent::atomic_fs::{atomic_write, atomic_write_unlocked, with_target_lock};
use swink_agent::{AgentMessage, CustomMessageRegistry, LlmMessage};

use crate::entry::SessionEntry;
use crate::interrupt::InterruptState;
use crate::load_options::LoadOptions;
use crate::meta::SessionMeta;
#[cfg(not(feature = "search"))]
use crate::search;
use crate::search::{SessionHit, SessionSearchOptions};
use crate::store::SessionStore;
use crate::time::{format_session_id, now_utc};

const META_LINE_PADDING: usize = 64;

#[derive(Debug, Clone)]
enum SessionRecord {
    Llm(Box<LlmMessage>),
    Custom(serde_json::Value),
    AppendBegin,
    Meta(Box<SessionMeta>),
    State(serde_json::Value),
}

impl SessionRecord {
    fn from_message(message: &AgentMessage, session_id: &str) -> Option<Self> {
        let line = crate::codec::encode_jsonl_message_line(message, session_id)?;
        match Self::parse(&line).ok()? {
            Self::Llm(llm) => Some(Self::Llm(llm)),
            Self::Custom(envelope) => Some(Self::Custom(envelope)),
            Self::AppendBegin | Self::Meta(_) | Self::State(_) => None,
        }
    }

    const fn state(state: serde_json::Value) -> Self {
        Self::State(state)
    }

    fn to_json_line(&self) -> io::Result<String> {
        match self {
            Self::Llm(message) => serde_json::to_string(message).map_err(io::Error::other),
            Self::Custom(envelope) => serde_json::to_string(envelope).map_err(io::Error::other),
            Self::AppendBegin => serde_json::to_string(&serde_json::json!({
                "_append": true,
                "phase": "begin"
            }))
            .map_err(io::Error::other),
            Self::Meta(meta) => serde_json::to_string(&serde_json::json!({
                "_meta": true,
                "data": meta
            }))
            .map_err(io::Error::other),
            Self::State(state) => serde_json::to_string(&serde_json::json!({
                "_state": true,
                "data": state
            }))
            .map_err(io::Error::other),
        }
    }

    fn parse(line: &str) -> io::Result<Self> {
        let value: serde_json::Value = serde_json::from_str(line).map_err(io::Error::other)?;
        if value.get("_state").and_then(serde_json::Value::as_bool) == Some(true) {
            return Ok(Self::State(
                value
                    .get("data")
                    .cloned()
                    .unwrap_or(serde_json::Value::Null),
            ));
        }
        if value.get("_append").and_then(serde_json::Value::as_bool) == Some(true) {
            return Ok(Self::AppendBegin);
        }
        if value.get("_meta").and_then(serde_json::Value::as_bool) == Some(true) {
            return serde_json::from_value::<SessionMeta>(
                value
                    .get("data")
                    .cloned()
                    .unwrap_or(serde_json::Value::Null),
            )
            .map(Box::new)
            .map(Self::Meta)
            .map_err(io::Error::other);
        }
        if value.get("_custom").and_then(serde_json::Value::as_bool) == Some(true) {
            return Ok(Self::Custom(value));
        }
        serde_json::from_value::<LlmMessage>(value)
            .map(Box::new)
            .map(Self::Llm)
            .map_err(io::Error::other)
    }
}

fn session_path(sessions_dir: &Path, id: &str) -> PathBuf {
    sessions_dir.join(format!("{id}.jsonl"))
}

fn interrupt_path(sessions_dir: &Path, id: &str) -> PathBuf {
    sessions_dir.join(format!("{id}.interrupt.json"))
}

fn with_session_lock<T>(
    sessions_dir: &Path,
    id: &str,
    op: impl FnOnce(&Path, &Path) -> io::Result<T>,
) -> io::Result<T> {
    let session = session_path(sessions_dir, id);
    let interrupt = interrupt_path(sessions_dir, id);
    with_target_lock(&session, || op(&session, &interrupt))
}

fn not_found(id: &str) -> io::Error {
    io::Error::new(io::ErrorKind::NotFound, format!("session not found: {id}"))
}

fn sequence_conflict(id: &str, expected: u64, actual: u64) -> io::Error {
    io::Error::new(
        io::ErrorKind::AlreadyExists,
        format!("sequence conflict for session {id}: expected {expected}, found {actual}"),
    )
}

fn check_sequence_path(path: &Path, id: &str, caller_sequence: u64) -> io::Result<()> {
    if !path.exists() {
        return Ok(());
    }
    let (stored_meta, _) = read_meta_with_line_len(path, id)?;
    if stored_meta.sequence != caller_sequence {
        return Err(sequence_conflict(id, caller_sequence, stored_meta.sequence));
    }
    Ok(())
}

fn empty_file() -> io::Error {
    io::Error::new(io::ErrorKind::InvalidData, "empty session file")
}

fn invalid_meta(error: impl std::fmt::Display) -> io::Error {
    io::Error::new(
        io::ErrorKind::InvalidData,
        format!("invalid session metadata: {error}"),
    )
}

fn open_session_file(path: &Path, id: &str) -> io::Result<std::fs::File> {
    let file = std::fs::File::open(path).map_err(|error| {
        if error.kind() == io::ErrorKind::NotFound {
            not_found(id)
        } else {
            error
        }
    })?;

    if file.metadata()?.len() == 0 {
        return Err(empty_file());
    }

    Ok(file)
}

fn read_meta_and_message_lines(path: &Path, id: &str) -> io::Result<(SessionMeta, Vec<String>)> {
    let file = open_session_file(path, id)?;
    let mut reader = io::BufReader::new(file);

    let meta_bytes = read_raw_line(&mut reader)?.ok_or_else(empty_file)?;
    let meta_line = String::from_utf8(meta_bytes).map_err(invalid_meta)?;
    let mut meta =
        canonical_meta_for_id(serde_json::from_str(&meta_line).map_err(invalid_meta)?, id);

    // Read message lines as raw bytes and validate UTF-8 per line. A crash
    // mid-append can truncate the final line inside a multi-byte UTF-8
    // sequence; `BufRead::lines()` would surface that as an `io::Error` and
    // abort the whole load before any per-line record classification runs.
    // Instead, skip the corrupt line with a warning so the rest of the
    // session is recovered (spec 021 FR-004: partial recovery over total
    // failure, same as bad-JSON lines in `classify_and_migrate`).
    let mut remaining_lines = Vec::new();
    let mut pending_append_lines: Option<Vec<String>> = None;
    let mut line_num = 1_usize;
    while let Some(raw) = read_raw_line(&mut reader)? {
        line_num += 1;
        let line = match String::from_utf8(raw) {
            Ok(line) => line,
            Err(error) => {
                tracing::warn!(
                    line = line_num,
                    error = %error,
                    "skipping invalid-UTF-8 line in session {id} \
                     (likely truncated by a crash mid-write)"
                );
                continue;
            }
        };
        match SessionRecord::parse(&line) {
            Ok(SessionRecord::AppendBegin) => {
                if pending_append_lines.is_some() {
                    tracing::warn!("discarding uncommitted nested append records in session {id}");
                }
                pending_append_lines = Some(Vec::new());
            }
            Ok(SessionRecord::Meta(meta_update)) => {
                if let Some(mut lines) = pending_append_lines.take() {
                    remaining_lines.append(&mut lines);
                }
                meta = canonical_meta_for_id(*meta_update, id);
            }
            Ok(SessionRecord::Llm(_) | SessionRecord::Custom(_) | SessionRecord::State(_))
            | Err(_) => {
                if let Some(lines) = pending_append_lines.as_mut() {
                    lines.push(line);
                } else {
                    remaining_lines.push(line);
                }
            }
        }
    }
    if pending_append_lines.is_some() {
        tracing::warn!("discarding uncommitted append records in session {id}");
    }

    Ok((meta, remaining_lines))
}

/// Read one `\n`-terminated line as raw bytes, without requiring valid UTF-8.
///
/// Returns `Ok(None)` at EOF. A trailing `\n` (and preceding `\r`, if any) is
/// stripped, matching [`BufRead::lines`] semantics. The final line is
/// returned even when it lacks a terminating newline (e.g. a write cut short
/// by a crash).
fn read_raw_line(reader: &mut impl BufRead) -> io::Result<Option<Vec<u8>>> {
    let mut buf = Vec::new();
    if reader.read_until(b'\n', &mut buf)? == 0 {
        return Ok(None);
    }
    if buf.last() == Some(&b'\n') {
        buf.pop();
        if buf.last() == Some(&b'\r') {
            buf.pop();
        }
    }
    Ok(Some(buf))
}

fn extract_state_from_lines(lines: &[String], id: &str) -> io::Result<Option<serde_json::Value>> {
    for line in lines {
        if let Some(state) = parse_state_line(line, id)? {
            return Ok(Some(state));
        }
    }

    Ok(None)
}

/// Reads the authoritative [`SessionMeta`] and the byte length of the padded
/// first line (used for in-place metadata patches).
///
/// The append-in-place protocol commits metadata by appending a
/// [`SessionRecord::Meta`] record at the tail and only afterwards patches the
/// padded first line, so after a crash the first line may be stale: the LAST
/// Meta record in the file wins, falling back to the first line when no Meta
/// record exists. Meta records fully replace each other (no merging), and
/// lines that fail to parse (e.g. a torn tail line) are skipped.
fn read_meta_with_line_len(path: &Path, id: &str) -> io::Result<(SessionMeta, usize)> {
    let mut first_line = String::new();
    let file = open_session_file(path, id)?;
    let mut reader = io::BufReader::new(file);
    reader.read_line(&mut first_line)?;

    if first_line.is_empty() {
        return Err(empty_file());
    }

    let line_len = first_line.len();
    let meta = canonical_meta_for_id(
        serde_json::from_str(first_line.trim_end()).map_err(invalid_meta)?,
        id,
    );
    let mut file = reader.into_inner();
    match find_last_meta_commit(&mut file, line_len as u64)? {
        Some(meta_update) => Ok((canonical_meta_for_id(meta_update, id), line_len)),
        None => Ok((meta, line_len)),
    }
}

/// Initial tail-window size for [`find_last_meta_commit`]'s backward scan.
const META_TAIL_WINDOW: u64 = 64 * 1024;

/// Finds the most recent committed [`SessionRecord::Meta`] record after the
/// first line, or `None` if the file contains no parseable Meta record.
///
/// Because every committed append ends with a Meta commit record, the record
/// we want is almost always within the last few bytes of the file. Rather
/// than forward-parsing every line (O(file size), which made
/// `SessionStore::list` a full-corpus scan), this reads geometrically growing
/// tail windows (starting at [`META_TAIL_WINDOW`]) and scans their complete
/// lines in reverse; when a window grows to cover the whole body it falls
/// back to the historical forward line-by-line scan, so the result — last
/// parseable Meta wins, unparseable lines skipped — is identical.
///
/// An invalid-UTF-8 line (e.g. a multi-byte character torn by a crash
/// mid-write) is skipped with a warning like any other unparseable line,
/// matching the tolerant forward scan (spec 021 FR-004 addendum / issue
/// #1067). Invalid UTF-8 that lies before the tail window of a large file
/// goes unnoticed when a later Meta commit is found, favoring recovery of
/// the committed metadata.
fn find_last_meta_commit(
    file: &mut std::fs::File,
    body_start: u64,
) -> io::Result<Option<SessionMeta>> {
    let file_len = file.metadata()?.len();
    let body_len = file_len.saturating_sub(body_start);

    let mut window = META_TAIL_WINDOW;
    while window < body_len {
        // `window < body_len` guarantees `start > body_start`.
        let start = file_len - window;
        file.seek(SeekFrom::Start(start))?;
        let mut buf = vec![0u8; usize::try_from(window).map_err(io::Error::other)?];
        file.read_exact(&mut buf)?;

        // The window starts mid-line: bytes before the first newline are a
        // partial line (they may even split a multi-byte character at the
        // window boundary), so only the bytes after it are complete lines.
        if let Some(pos) = buf.iter().position(|&b| b == b'\n')
            && let Some(meta) = last_meta_in_complete_lines(&buf[pos + 1..])
        {
            return Ok(Some(meta));
        }

        window = window.saturating_mul(2);
    }

    // Terminal case: the remaining window covers the whole body. Stream
    // forward line-by-line exactly like the historical implementation so
    // memory stays bounded, tolerating invalid-UTF-8 lines the same way
    // the full-load forward scan does.
    file.seek(SeekFrom::Start(body_start))?;
    let mut meta = None;
    let mut reader = io::BufReader::new(file);
    while let Some(raw) = read_raw_line(&mut reader)? {
        let Ok(line) = String::from_utf8(raw) else {
            tracing::warn!(
                "skipping invalid-UTF-8 line in session meta scan \
                 (likely truncated by a crash mid-write)"
            );
            continue;
        };
        if let Ok(SessionRecord::Meta(meta_update)) = SessionRecord::parse(&line) {
            meta = Some(*meta_update);
        }
    }
    Ok(meta)
}

/// Scans complete lines in reverse for the last parseable Meta record.
///
/// `bytes` must start at a real line boundary and extend to end of file.
fn last_meta_in_complete_lines(bytes: &[u8]) -> Option<SessionMeta> {
    // Validate UTF-8 per line: an invalid-UTF-8 line (e.g. a torn tail line
    // that truncates a multi-byte character) is skipped like any other
    // unparseable line, matching the tolerant forward scan (spec 021 FR-004
    // addendum / issue #1067).
    for line in bytes.split(|&b| b == b'\n').rev() {
        let line = line.strip_suffix(b"\r").unwrap_or(line);
        if line.is_empty() {
            continue;
        }
        let Ok(text) = std::str::from_utf8(line) else {
            tracing::warn!(
                "skipping invalid-UTF-8 line in session meta scan \
                 (likely truncated by a crash mid-write)"
            );
            continue;
        };
        if let Ok(SessionRecord::Meta(meta_update)) = SessionRecord::parse(text) {
            return Some(*meta_update);
        }
    }
    None
}

fn canonical_meta_for_id(mut meta: SessionMeta, id: &str) -> SessionMeta {
    meta.id = id.to_string();
    meta
}

fn rewrite_session_file_locked(
    path: &Path,
    meta: &SessionMeta,
    lines: &[String],
) -> io::Result<()> {
    atomic_write_unlocked(path, |writer| {
        write_meta_line(writer, meta, META_LINE_PADDING)?;

        for line in lines {
            if !line.is_empty() {
                writeln!(writer, "{line}")?;
            }
        }
        Ok(())
    })
}

fn preserve_existing_lines(
    path: &Path,
    id: &str,
    should_preserve: impl Fn(&str) -> bool,
) -> io::Result<Vec<String>> {
    if !path.exists() {
        return Ok(Vec::new());
    }

    let (_, lines) = read_meta_and_message_lines(path, id)?;
    Ok(lines
        .into_iter()
        .filter(|line| should_preserve(line))
        .collect())
}

fn preserve_for_message_save(line: &str) -> bool {
    match SessionRecord::parse(line) {
        Ok(SessionRecord::State(_)) => true,
        Ok(
            SessionRecord::AppendBegin
            | SessionRecord::Llm(_)
            | SessionRecord::Custom(_)
            | SessionRecord::Meta(_),
        ) => false,
        Err(_) => matches!(
            SessionEntry::parse(line),
            Ok(entry) if !matches!(entry, SessionEntry::Message(_))
        ),
    }
}

fn preserve_for_entry_save(line: &str) -> bool {
    matches!(
        SessionRecord::parse(line),
        Ok(SessionRecord::State(_) | SessionRecord::Custom(_))
    )
}

fn save_messages_with_hooks<AfterValidation, WriteOp>(
    path: &Path,
    id: &str,
    meta: &SessionMeta,
    messages: &[AgentMessage],
    after_validation: AfterValidation,
    write_op: WriteOp,
) -> io::Result<()>
where
    AfterValidation: FnOnce() -> io::Result<()>,
    WriteOp: FnOnce(&Path, &SessionMeta, &[AgentMessage], &str) -> io::Result<()>,
{
    with_target_lock(path, || {
        check_sequence_path(path, id, meta.sequence)?;
        after_validation()?;

        let mut write_meta = meta.clone();
        write_meta.id = id.to_string();
        write_meta.sequence += 1;
        write_op(path, &write_meta, messages, id)
    })
}

fn write_messages_locked(
    path: &Path,
    meta: &SessionMeta,
    messages: &[AgentMessage],
    id: &str,
) -> io::Result<()> {
    let preserved_lines = preserve_existing_lines(path, id, preserve_for_message_save)?;

    write_messages_with_preserved_lines(path, meta, messages, id, &preserved_lines)
}

fn write_messages_with_preserved_lines(
    path: &Path,
    meta: &SessionMeta,
    messages: &[AgentMessage],
    id: &str,
    preserved_lines: &[String],
) -> io::Result<()> {
    atomic_write_unlocked(path, |writer| {
        write_meta_line(writer, meta, META_LINE_PADDING)?;

        for msg in messages {
            if let Some(record) = SessionRecord::from_message(msg, id) {
                writer.write_all(record.to_json_line()?.as_bytes())?;
                writeln!(writer)?;
            }
        }
        for line in preserved_lines {
            if !line.is_empty() {
                writeln!(writer, "{line}")?;
            }
        }
        Ok(())
    })
}

fn write_meta_line(
    writer: &mut impl Write,
    meta: &SessionMeta,
    padding: usize,
) -> io::Result<usize> {
    let line = serde_json::to_string(meta).map_err(io::Error::other)?;
    writer.write_all(line.as_bytes())?;
    for _ in 0..padding {
        writer.write_all(b" ")?;
    }
    writeln!(writer)?;
    Ok(line.len() + padding + 1)
}

fn write_meta_line_in_place(
    file: &mut std::fs::File,
    line: &str,
    existing_line_len: usize,
) -> io::Result<bool> {
    if line.len() + 1 > existing_line_len {
        return Ok(false);
    }

    file.seek(SeekFrom::Start(0))?;
    file.write_all(line.as_bytes())?;
    for _ in line.len()..existing_line_len - 1 {
        file.write_all(b" ")?;
    }
    file.write_all(b"\n")?;
    Ok(true)
}

fn append_records_in_place(
    path: &Path,
    meta: &SessionMeta,
    meta_line_len: usize,
    record_lines: &[String],
) -> io::Result<bool> {
    append_records_in_place_with_hook(path, meta, meta_line_len, record_lines, |_| Ok(()))
}

fn append_records_in_place_with_hook(
    path: &Path,
    meta: &SessionMeta,
    meta_line_len: usize,
    record_lines: &[String],
    after_commit_append: impl FnOnce(&mut std::fs::File) -> io::Result<()>,
) -> io::Result<bool> {
    append_records_in_place_with_hooks(
        path,
        meta,
        meta_line_len,
        record_lines,
        |_| Ok(()),
        after_commit_append,
    )
}

fn append_records_in_place_with_hooks(
    path: &Path,
    meta: &SessionMeta,
    meta_line_len: usize,
    record_lines: &[String],
    after_records_append: impl FnOnce(&mut std::fs::File) -> io::Result<()>,
    after_commit_append: impl FnOnce(&mut std::fs::File) -> io::Result<()>,
) -> io::Result<bool> {
    let meta_line = serde_json::to_string(meta).map_err(io::Error::other)?;
    if meta_line.len() + 1 > meta_line_len {
        return Ok(false);
    }

    let mut file = std::fs::OpenOptions::new()
        .read(true)
        .write(true)
        .open(path)?;

    file.seek(SeekFrom::End(0))?;
    writeln!(file, "{}", SessionRecord::AppendBegin.to_json_line()?)?;
    for line in record_lines {
        if !line.is_empty() {
            writeln!(file, "{line}")?;
        }
    }
    file.flush()?;
    after_records_append(&mut file)?;
    writeln!(
        file,
        "{}",
        SessionRecord::Meta(Box::new(meta.clone())).to_json_line()?
    )?;
    file.flush()?;
    after_commit_append(&mut file)?;

    write_meta_line_in_place(&mut file, &meta_line, meta_line_len)?;
    file.flush()?;
    Ok(true)
}

fn upsert_state_line(
    lines: &mut Vec<String>,
    session_id: &str,
    state: &serde_json::Value,
) -> io::Result<()> {
    let state_line = SessionRecord::state(state.clone()).to_json_line()?;
    if let Some(line) = find_record_line_mut(lines.as_mut_slice(), session_id, |record| {
        matches!(record, SessionRecord::State(_))
    }) {
        line.clone_from(&state_line);
    } else {
        lines.push(state_line);
    }
    Ok(())
}

fn parse_state_line(line: &str, id: &str) -> io::Result<Option<serde_json::Value>> {
    let value: serde_json::Value = match serde_json::from_str(line) {
        Ok(value) => value,
        Err(error) => {
            if line.contains("\"_state\"") {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    format!("corrupted state line in session {id}: {error}"),
                ));
            }
            return Ok(None);
        }
    };

    let Some(state_marker) = value.get("_state") else {
        return Ok(None);
    };

    if state_marker.as_bool() != Some(true) {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            format!("invalid state marker in session {id}: expected `_state: true`"),
        ));
    }

    Ok(Some(
        value
            .get("data")
            .cloned()
            .unwrap_or(serde_json::Value::Null),
    ))
}

fn append_records(
    path: &Path,
    id: &str,
    records: impl IntoIterator<Item = SessionRecord>,
) -> io::Result<()> {
    append_records_with_rewrite(path, id, records, rewrite_session_file_locked)
}

fn append_records_with_rewrite<F>(
    path: &Path,
    id: &str,
    records: impl IntoIterator<Item = SessionRecord>,
    rewrite_fn: F,
) -> io::Result<()>
where
    F: FnOnce(&Path, &SessionMeta, &[String]) -> io::Result<()>,
{
    with_target_lock(path, || {
        let (mut meta, meta_line_len) = read_meta_with_line_len(path, id)?;
        meta.updated_at = now_utc();
        meta.sequence += 1;

        let record_lines = records
            .into_iter()
            .map(|record| record.to_json_line())
            .collect::<io::Result<Vec<_>>>()?;

        if append_records_in_place(path, &meta, meta_line_len, &record_lines)? {
            return Ok(());
        }

        let (_, mut existing_lines) = read_meta_and_message_lines(path, id)?;
        existing_lines.extend(record_lines);
        rewrite_fn(path, &meta, &existing_lines)
    })
}

fn find_record_line_mut<'a>(
    lines: &'a mut [String],
    session_id: &str,
    predicate: impl Fn(&SessionRecord) -> bool,
) -> Option<&'a mut String> {
    for line in lines {
        match SessionRecord::parse(line) {
            Ok(record) if predicate(&record) => return Some(line),
            Ok(_) => {}
            Err(error) => {
                tracing::warn!(
                    error = %error,
                    "skipping unparseable line while scanning session {session_id}"
                );
            }
        }
    }
    None
}

/// Validate a session ID, rejecting unsafe filesystem characters.
///
/// Rejects IDs containing `/`, `\`, `..`, `:`, or ASCII control characters.
fn validate_session_id(id: &str) -> io::Result<()> {
    if id.is_empty() {
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            "session ID must not be empty",
        ));
    }
    if id.contains('/')
        || id.contains('\\')
        || id.contains("..")
        || id.chars().any(|c| c == ':' || c.is_ascii_control())
    {
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            format!("session ID contains unsafe characters: {id:?}"),
        ));
    }
    Ok(())
}

/// JSONL file-based session store.
///
/// Each session is a single `.jsonl` file in the configured directory.
/// Line 1 is [`SessionMeta`] (JSON), lines 2+ are one [`LlmMessage`] each.
///
/// Concurrent writes to the same session may corrupt the file.
/// Callers are expected to enforce single-writer access.
///
/// With the `search` feature enabled, a tantivy index is maintained in
/// `<sessions_dir>/.search_index/` and used by
/// [`SessionStore::search()`].  Without the feature, search falls back to a
/// linear scan of all JSONL files.
pub struct JsonlSessionStore {
    sessions_dir: PathBuf,
    migrators: Vec<Box<dyn crate::migrate::SessionMigrator>>,
    /// Lazily-opened tantivy index.  Populated on first `search()` call (or
    /// via `open_search_index()`).  Only present when the `search` feature is
    /// enabled.
    ///
    /// Uses `Mutex<Option<...>>` rather than `OnceLock` because we need
    /// fallible initialization (`OnceLock::get_or_try_init` is not stable).
    #[cfg(feature = "search")]
    tantivy_index: std::sync::Mutex<Option<crate::search::index::TantivyIndex>>,
    /// Tracks whether the index has been initially populated from all JSONL
    /// files on first search.
    #[cfg(feature = "search")]
    index_built: std::sync::Mutex<bool>,
}

impl JsonlSessionStore {
    /// Create a new store rooted at the given directory.
    ///
    /// Creates the directory (and parents) if it does not exist.
    pub fn new(sessions_dir: PathBuf) -> io::Result<Self> {
        std::fs::create_dir_all(&sessions_dir)?;
        Ok(Self {
            sessions_dir,
            migrators: Vec::new(),
            #[cfg(feature = "search")]
            tantivy_index: std::sync::Mutex::new(None),
            #[cfg(feature = "search")]
            index_built: std::sync::Mutex::new(false),
        })
    }

    /// Open (or create) the tantivy search index eagerly.
    ///
    /// Normally the index is opened lazily on the first `search()` call.
    /// Call this method if you want to detect index-creation errors at startup
    /// rather than at search time.
    ///
    /// Only available with the `search` feature.
    #[cfg(feature = "search")]
    pub fn open_search_index(&self) -> io::Result<()> {
        self.with_tantivy_index(|_| Ok(()))
    }

    /// Build (or rebuild) the tantivy index from all current JSONL files.
    ///
    /// Existing index data is replaced.  Use this after bulk-importing sessions
    /// outside of the store API.
    ///
    /// Only available with the `search` feature.
    #[cfg(feature = "search")]
    pub fn rebuild_search_index(&self) -> io::Result<()> {
        self.with_tantivy_index(|index| {
            // Wipe all existing docs so sessions removed outside the store API
            // don't produce ghost hits after the rebuild.
            index.clear_all()?;
            for meta in self.list()? {
                let (_, entries) = self.load_entries(&meta.id)?;
                index.index_session(&meta, &entries)?;
            }
            Ok(())
        })
    }

    /// Obtain a snapshot clone of the tantivy index (creating it if needed),
    /// then call `f` with it.
    ///
    /// Using a clone of `TantivyIndex` is safe because it is `Arc`-based
    /// internally.
    #[cfg(feature = "search")]
    fn with_tantivy_index<T>(
        &self,
        f: impl FnOnce(&crate::search::index::TantivyIndex) -> io::Result<T>,
    ) -> io::Result<T> {
        let index_clone = {
            let mut guard = self.tantivy_index.lock().unwrap_or_else(|e| e.into_inner());
            if guard.is_none() {
                *guard = Some(crate::search::index::TantivyIndex::open_or_create(
                    &self.sessions_dir,
                )?);
            }
            guard.as_ref().expect("just set").clone()
        };
        f(&index_clone)
    }

    #[cfg(feature = "search")]
    fn active_tantivy_index(&self) -> Option<crate::search::index::TantivyIndex> {
        self.tantivy_index
            .lock()
            .unwrap_or_else(|e| e.into_inner())
            .clone()
    }

    /// Best-effort full re-index of one session in an already-open index.
    ///
    /// Reloads the session from disk and replaces every document for it, so
    /// this is O(session).  It is the right tool for rewrite paths (`save`,
    /// `save_full`) where existing content may have changed wholesale;
    /// append paths must use [`Self::append_to_active_search_index`] instead.
    #[cfg(feature = "search")]
    fn refresh_active_search_index(&self, id: &str, operation: &str) {
        let Some(index) = self.active_tantivy_index() else {
            return;
        };

        match self.load_entries(id) {
            Ok((meta, entries)) => {
                if let Err(err) = index.index_session(&meta, &entries) {
                    tracing::warn!(
                        session_id = %id,
                        error = %err,
                        operation,
                        "failed to update search index after session mutation"
                    );
                }
            }
            Err(err) => {
                tracing::warn!(
                    session_id = %id,
                    error = %err,
                    operation,
                    "failed to load session for search index refresh"
                );
            }
        }
    }

    /// Best-effort incremental index update for append paths.
    ///
    /// Adds documents for only the newly appended `entries` instead of
    /// reloading and re-indexing the whole session, making the index cost of
    /// an append O(new entries) rather than O(session).  Rewrite paths must
    /// keep using [`Self::refresh_active_search_index`] for its replacement
    /// semantics.
    #[cfg(feature = "search")]
    fn append_to_active_search_index(
        &self,
        id: &str,
        meta: &SessionMeta,
        entries: &[SessionEntry],
    ) {
        let Some(index) = self.active_tantivy_index() else {
            return;
        };
        if let Err(err) = index.append_session_entries(meta, entries) {
            tracing::warn!(
                session_id = %id,
                error = %err,
                "failed to append entries to search index after session append"
            );
        }
    }

    /// Register session migrators for automatic schema upgrades on load.
    #[must_use]
    pub fn with_migrators(
        mut self,
        migrators: Vec<Box<dyn crate::migrate::SessionMigrator>>,
    ) -> Self {
        self.migrators = migrators;
        self
    }

    /// Default sessions directory: `<config_dir>/swink-agent/sessions`.
    ///
    /// Returns `None` if the platform config directory cannot be determined.
    pub fn default_dir() -> Option<PathBuf> {
        dirs::config_dir().map(|d| d.join("swink-agent").join("sessions"))
    }

    /// Generate a new unique session ID using a UTC timestamp plus UUID suffix.
    pub fn new_session_id() -> String {
        format_session_id()
    }
}

impl SessionStore for JsonlSessionStore {
    fn save(&self, id: &str, meta: &SessionMeta, messages: &[AgentMessage]) -> io::Result<()> {
        validate_session_id(id)?;
        let path = session_path(&self.sessions_dir, id);
        save_messages_with_hooks(&path, id, meta, messages, || Ok(()), write_messages_locked)?;
        #[cfg(feature = "search")]
        self.refresh_active_search_index(id, "save");
        Ok(())
    }

    fn save_full(
        &self,
        id: &str,
        meta: &SessionMeta,
        messages: &[AgentMessage],
        state: &serde_json::Value,
    ) -> io::Result<SessionMeta> {
        validate_session_id(id)?;

        let path = session_path(&self.sessions_dir, id);
        let persisted_meta = with_target_lock(&path, || {
            check_sequence_path(&path, id, meta.sequence)?;

            let mut write_meta = meta.clone();
            write_meta.id = id.to_string();
            write_meta.sequence += 1;

            let mut preserved_lines =
                preserve_existing_lines(&path, id, preserve_for_message_save)?;
            upsert_state_line(&mut preserved_lines, id, state)?;
            write_messages_with_preserved_lines(
                &path,
                &write_meta,
                messages,
                id,
                &preserved_lines,
            )?;

            Ok(write_meta)
        })?;

        #[cfg(feature = "search")]
        self.refresh_active_search_index(id, "save_full");

        Ok(persisted_meta)
    }

    fn append(&self, id: &str, messages: &[AgentMessage]) -> io::Result<()> {
        validate_session_id(id)?;

        let path = session_path(&self.sessions_dir, id);
        append_records(
            &path,
            id,
            messages
                .iter()
                .filter_map(|msg| SessionRecord::from_message(msg, id)),
        )?;

        // Index only the newly appended messages (best-effort). Custom
        // messages are not `SessionEntry`s and were never indexed by the
        // previous full-refresh either, so only LLM messages are converted.
        #[cfg(feature = "search")]
        if self.active_tantivy_index().is_some() {
            let entries: Vec<SessionEntry> = messages
                .iter()
                .filter_map(|msg| match msg {
                    AgentMessage::Llm(llm) => Some(SessionEntry::Message(llm.clone())),
                    // Custom (and any future) variants are not `SessionEntry`s
                    // and were never indexed by the previous full refresh.
                    _ => None,
                })
                .collect();
            if !entries.is_empty() {
                match read_meta_with_line_len(&path, id) {
                    Ok((meta, _)) => self.append_to_active_search_index(id, &meta, &entries),
                    Err(err) => tracing::warn!(
                        session_id = %id,
                        error = %err,
                        "failed to read session meta for search index append"
                    ),
                }
            }
        }
        Ok(())
    }

    fn load(
        &self,
        id: &str,
        registry: Option<&CustomMessageRegistry>,
    ) -> io::Result<(SessionMeta, Vec<AgentMessage>)> {
        validate_session_id(id)?;

        let path = session_path(&self.sessions_dir, id);
        with_target_lock(&path, || {
            let (meta, lines) = read_meta_and_message_lines(&path, id)?;
            let (meta, classified) = self.classify_and_migrate(meta, lines, id)?;
            Ok((meta, classified_lines_to_messages(classified, registry, id)))
        })
    }

    fn load_full(
        &self,
        id: &str,
        registry: Option<&CustomMessageRegistry>,
    ) -> io::Result<(SessionMeta, Vec<AgentMessage>, Option<serde_json::Value>)> {
        validate_session_id(id)?;

        let path = session_path(&self.sessions_dir, id);
        with_target_lock(&path, || {
            let (meta, lines) = read_meta_and_message_lines(&path, id)?;
            let state = extract_state_from_lines(&lines, id)?;
            let (meta, classified) = self.classify_and_migrate(meta, lines, id)?;
            let messages = classified_lines_to_messages(classified, registry, id);
            Ok((meta, messages, state))
        })
    }

    fn list(&self) -> io::Result<Vec<SessionMeta>> {
        let mut sessions = Vec::new();

        let entries = std::fs::read_dir(&self.sessions_dir)?;
        for entry in entries {
            let entry = entry?;
            let path = entry.path();
            if path.extension().and_then(|e| e.to_str()) != Some("jsonl") {
                continue;
            }

            let read_meta = with_target_lock(&path, || {
                let Some(file_id) = path.file_stem().and_then(|stem| stem.to_str()) else {
                    return Ok(None);
                };
                validate_session_id(file_id)?;
                read_meta_with_line_len(&path, file_id)
                    .map(|(meta, _)| Some((file_id.to_string(), meta)))
            });

            match read_meta {
                Ok(Some((file_id, meta))) => {
                    sessions.push(canonical_meta_for_id(meta, &file_id));
                }
                Ok(None) => {}
                Err(error) => {
                    tracing::warn!(
                        path = %path.display(),
                        error = %error,
                        "skipping unreadable session file"
                    );
                }
            }
        }

        sessions.sort_by_key(|b| std::cmp::Reverse(b.updated_at));
        Ok(sessions)
    }

    fn delete(&self, id: &str) -> io::Result<()> {
        validate_session_id(id)?;
        with_session_lock(&self.sessions_dir, id, |path, int_path| {
            std::fs::remove_file(path)?;
            // Cascade-delete the interrupt file if it exists.
            if int_path.exists() {
                std::fs::remove_file(int_path)?;
            }
            Ok(())
        })?;

        // Remove from tantivy index (best-effort).
        #[cfg(feature = "search")]
        {
            let id_owned = id.to_string();
            let _ = self.with_tantivy_index(|index| {
                if let Err(err) = index.delete_session(&id_owned) {
                    tracing::warn!(
                        session_id = %id_owned,
                        error = %err,
                        "failed to remove session from search index after delete"
                    );
                }
                Ok(())
            });
        }

        Ok(())
    }

    fn save_state(&self, id: &str, state: &serde_json::Value) -> io::Result<()> {
        validate_session_id(id)?;

        let path = session_path(&self.sessions_dir, id);
        with_target_lock(&path, || {
            if !path.exists() {
                return Err(not_found(id));
            }

            let (mut meta, mut lines) = read_meta_and_message_lines(&path, id)?;
            meta.updated_at = now_utc();
            meta.sequence += 1;
            upsert_state_line(&mut lines, id, state)?;
            rewrite_session_file_locked(&path, &meta, &lines)
        })
    }

    fn load_state(&self, id: &str) -> io::Result<Option<serde_json::Value>> {
        validate_session_id(id)?;

        let path = session_path(&self.sessions_dir, id);
        if !path.exists() {
            return Ok(None);
        }

        with_target_lock(&path, || {
            let (_, lines) = read_meta_and_message_lines(&path, id)?;
            extract_state_from_lines(&lines, id)
        })
    }

    fn save_interrupt(&self, id: &str, state: &InterruptState) -> io::Result<()> {
        validate_session_id(id)?;
        with_session_lock(&self.sessions_dir, id, |session, path| {
            if !session.exists() {
                return Err(not_found(id));
            }
            atomic_write(path, |writer| {
                serde_json::to_writer_pretty(&mut *writer, state).map_err(io::Error::other)
            })
        })
    }

    fn load_interrupt(&self, id: &str) -> io::Result<Option<InterruptState>> {
        validate_session_id(id)?;
        with_session_lock(&self.sessions_dir, id, |session, path| {
            if !session.exists() || !path.exists() {
                return Ok(None);
            }
            let contents = std::fs::read_to_string(path)?;
            let state: InterruptState = serde_json::from_str(&contents).map_err(|e| {
                io::Error::new(
                    io::ErrorKind::InvalidData,
                    format!("corrupted interrupt file for session {id}: {e}"),
                )
            })?;
            Ok(Some(state))
        })
    }

    fn clear_interrupt(&self, id: &str) -> io::Result<()> {
        validate_session_id(id)?;
        with_session_lock(&self.sessions_dir, id, |_session, path| {
            if path.exists() {
                std::fs::remove_file(path)?;
            }
            Ok(())
        })
    }

    fn load_with_options(
        &self,
        id: &str,
        options: &LoadOptions,
    ) -> io::Result<(SessionMeta, Vec<SessionEntry>)> {
        let (meta, mut entries) = self.load_entries(id)?;

        // Filter by entry type
        if let Some(ref types) = options.entry_types {
            entries.retain(|entry| types.iter().any(|t| t == entry.entry_type_name()));
        }

        // Filter by timestamp (entry timestamps are epoch seconds)
        if let Some(after) = options.after_timestamp {
            let after_secs = after.timestamp().cast_unsigned();
            entries.retain(|entry| entry.timestamp().is_some_and(|ts| ts > after_secs));
        }

        // Truncate to last N
        if let Some(n) = options.last_n_entries
            && entries.len() > n
        {
            entries.drain(..entries.len() - n);
        }

        Ok((meta, entries))
    }

    fn search(&self, query: &str, options: &SessionSearchOptions) -> io::Result<Vec<SessionHit>> {
        // When the `search` feature is enabled, delegate to the tantivy index.
        // The index is built lazily on the first call and kept warm for
        // subsequent calls.
        #[cfg(feature = "search")]
        return self.with_tantivy_index(|index| {
            self.build_index_if_empty(index)?;
            index.search(query, options)
        });

        // Fallback linear scan (no `search` feature).
        #[cfg(not(feature = "search"))]
        linear_search(self, query, options)
    }
}

/// Linear (no-index) search implementation used when the `search` feature is
/// not enabled.
#[cfg(not(feature = "search"))]
fn linear_search(
    store: &JsonlSessionStore,
    query: &str,
    options: &SessionSearchOptions,
) -> io::Result<Vec<SessionHit>> {
    let terms = search::query_terms(query);
    let limit = options.limit();
    if terms.is_empty() || limit == 0 {
        return Ok(Vec::new());
    }

    let session_ids = if let Some(ids) = &options.session_ids {
        ids.iter()
            .map(|id| {
                validate_session_id(id)?;
                Ok(id.clone())
            })
            .collect::<io::Result<Vec<_>>>()?
    } else {
        store
            .list()?
            .into_iter()
            .map(|meta| meta.id)
            .collect::<Vec<_>>()
    };

    let mut hits = Vec::new();
    for id in session_ids {
        let (meta, entries) = store.load_entries(&id)?;
        for entry in entries {
            if !search::entry_matches_type(&entry, options)
                || !search::entry_matches_time_range(&entry, options)
            {
                continue;
            }
            let Some((relevance, snippet)) = search::search_entry(&entry, &terms) else {
                continue;
            };
            hits.push(SessionHit::new(
                meta.id.clone(),
                meta.title.clone(),
                entry,
                relevance,
                snippet,
            ));
        }
    }

    hits.sort_by(|left, right| {
        right
            .score
            .cmp(&left.score)
            .then_with(|| right.entry.timestamp().cmp(&left.entry.timestamp()))
            .then_with(|| left.session_id.cmp(&right.session_id))
    });
    hits.truncate(limit);
    Ok(hits)
}

#[cfg(feature = "search")]
impl JsonlSessionStore {
    /// Index all sessions if the initial build has not yet been done.
    ///
    /// This provides lazy-build-on-first-search semantics: the first `search()`
    /// call populates the index from all existing JSONL files; subsequent calls
    /// find the index already populated.
    fn build_index_if_empty(&self, index: &crate::search::index::TantivyIndex) -> io::Result<()> {
        let already_built = {
            let guard = self.index_built.lock().unwrap_or_else(|e| e.into_inner());
            *guard
        };
        if already_built {
            return Ok(());
        }
        // Populate from all current sessions.
        for meta in self.list()? {
            let (_, entries) = self.load_entries(&meta.id)?;
            index.index_session(&meta, &entries)?;
        }
        *self.index_built.lock().unwrap_or_else(|e| e.into_inner()) = true;
        Ok(())
    }
}

impl JsonlSessionStore {
    /// Save a session with rich entry types.
    ///
    /// Lines 2+ are [`SessionEntry`] values serialized with an `entry_type` tag.
    ///
    /// When the `search` feature is enabled, the tantivy index is updated
    /// after the file is written (best-effort — index errors are logged but
    /// do not fail the save).
    pub fn save_entries(
        &self,
        id: &str,
        meta: &SessionMeta,
        entries: &[SessionEntry],
    ) -> io::Result<()> {
        validate_session_id(id)?;
        let path = session_path(&self.sessions_dir, id);
        #[allow(unused_variables)]
        let write_meta = with_target_lock(&path, || {
            check_sequence_path(&path, id, meta.sequence)?;

            // Increment sequence for the write
            let mut write_meta = meta.clone();
            write_meta.id = id.to_string();
            write_meta.sequence += 1;
            let preserved_lines = preserve_existing_lines(&path, id, preserve_for_entry_save)?;

            atomic_write_unlocked(&path, |writer| {
                // First line: metadata
                serde_json::to_writer(&mut *writer, &write_meta).map_err(io::Error::other)?;
                writeln!(writer)?;

                // Subsequent lines: one SessionEntry per line
                for entry in entries {
                    serde_json::to_writer(&mut *writer, entry).map_err(io::Error::other)?;
                    writeln!(writer)?;
                }
                for line in &preserved_lines {
                    if !line.is_empty() {
                        writeln!(writer, "{line}")?;
                    }
                }
                Ok(())
            })?;
            Ok(write_meta)
        })?;

        // Update tantivy index (best-effort).
        #[cfg(feature = "search")]
        {
            let entries_for_index: Vec<_> = entries.to_vec();
            let meta_for_index = write_meta;
            let _ = self.with_tantivy_index(|index| {
                if let Err(err) = index.index_session(&meta_for_index, &entries_for_index) {
                    tracing::warn!(
                        session_id = %id,
                        error = %err,
                        "failed to update search index after save_entries"
                    );
                }
                Ok(())
            });
        }

        Ok(())
    }

    /// Append entries to an existing session without rewriting the whole file.
    ///
    /// This is the append-only counterpart to [`Self::save_entries`]. Instead of
    /// rewriting every record on each turn, it appends only the new `entries`
    /// and patches the metadata line in place (reusing the same on-disk
    /// machinery and locking guarantees as the internal append path).
    ///
    /// The caller's `meta.sequence` must match the on-disk sequence
    /// (optimistic concurrency, identical to [`Self::save_entries`]); a mismatch
    /// returns an error. Every other `meta` field is written as provided —
    /// only `id` and `sequence` are set by the store. The returned
    /// [`SessionMeta`] carries the incremented sequence so callers stay in
    /// sync without re-reading the file.
    ///
    /// If the metadata line can no longer be patched in place — its serialized
    /// form would outgrow the slot reserved on disk — this transparently falls
    /// back to a full rewrite that preserves all existing lines.
    ///
    /// Errors if the session file does not exist; use [`Self::save_entries`] to
    /// create a new session.
    ///
    /// When the `search` feature is enabled, only the newly appended
    /// `entries` are added to the tantivy index afterwards — the session's
    /// existing documents are untouched (best-effort — index errors are
    /// logged but do not fail the append).
    pub fn append_entries(
        &self,
        id: &str,
        meta: &SessionMeta,
        entries: &[SessionEntry],
    ) -> io::Result<SessionMeta> {
        validate_session_id(id)?;
        let path = session_path(&self.sessions_dir, id);

        let write_meta = with_target_lock(&path, || {
            check_sequence_path(&path, id, meta.sequence)?;
            let (_, meta_line_len) = read_meta_with_line_len(&path, id)?;

            let mut write_meta = meta.clone();
            write_meta.id = id.to_string();
            write_meta.sequence += 1;

            let record_lines = entries
                .iter()
                .map(|entry| serde_json::to_string(entry).map_err(io::Error::other))
                .collect::<io::Result<Vec<_>>>()?;

            if !append_records_in_place(&path, &write_meta, meta_line_len, &record_lines)? {
                // The patched metadata line would outgrow its on-disk slot —
                // fall back to a full rewrite that keeps every existing line.
                let (_, mut existing_lines) = read_meta_and_message_lines(&path, id)?;
                existing_lines.extend(record_lines);
                rewrite_session_file_locked(&path, &write_meta, &existing_lines)?;
            }

            Ok(write_meta)
        })?;

        #[cfg(feature = "search")]
        self.append_to_active_search_index(id, &write_meta, entries);

        Ok(write_meta)
    }

    /// Load a session with rich entry types.
    ///
    /// Parses each line after metadata as a [`SessionEntry`]. Old-format lines
    /// (raw `LlmMessage` without `entry_type`) are interpreted as
    /// [`SessionEntry::Message`] for backward compatibility.
    pub fn load_entries(&self, id: &str) -> io::Result<(SessionMeta, Vec<SessionEntry>)> {
        validate_session_id(id)?;

        let path = session_path(&self.sessions_dir, id);
        with_target_lock(&path, || {
            let (meta, lines) = read_meta_and_message_lines(&path, id)?;
            let (meta, classified) = self.classify_and_migrate(meta, lines, id)?;

            let entries = classified
                .into_iter()
                .filter_map(|item| match item {
                    ClassifiedLine::Entry(entry) => Some(*entry),
                    ClassifiedLine::Custom(_) | ClassifiedLine::State => None,
                })
                .collect();

            Ok((meta, entries))
        })
    }

    /// Parse every message line, classify each as a migrateable
    /// [`SessionEntry`] or a pass-through custom/state wrapper, then run any
    /// configured migrators over the `SessionEntry` subset.
    ///
    /// This is the single migration entry point shared by `load()` and
    /// `load_entries()` so both observe identical post-migration state.
    /// Positional ordering of pass-through wrappers relative to each other
    /// (and relative to the migrateable block as a whole) is preserved;
    /// migrators may freely insert/remove entries within the entry block.
    fn classify_and_migrate(
        &self,
        mut meta: SessionMeta,
        lines: Vec<String>,
        id: &str,
    ) -> io::Result<(SessionMeta, Vec<ClassifiedLine>)> {
        let mut classified: Vec<ClassifiedLine> = Vec::with_capacity(lines.len());
        for (idx, line) in lines.into_iter().enumerate() {
            if line.trim().is_empty() {
                continue;
            }
            let line_num = idx + 2;
            match classify_line(&line) {
                Ok(item) => classified.push(item),
                Err(error) => {
                    tracing::warn!(
                        line = line_num,
                        error = %error,
                        "skipping unparseable line in session {id}"
                    );
                }
            }
        }

        // Split out the migrateable subset, run the migrator pipeline, then
        // weave the migrated entries back into the classified list.
        let original_entry_count = classified
            .iter()
            .filter(|item| matches!(item, ClassifiedLine::Entry(_)))
            .count();
        let mut entries: Vec<SessionEntry> = classified
            .iter()
            .filter_map(|item| match item {
                ClassifiedLine::Entry(entry) => Some((**entry).clone()),
                _ => None,
            })
            .collect();

        crate::migrate::run_migrations(&mut meta, &mut entries, &self.migrators)?;

        let rebuilt = weave_migrated_entries(classified, entries, original_entry_count);
        Ok((meta, rebuilt))
    }
}

/// Weave a migrated entry list back into a classified line sequence.
///
/// If the migrator preserved the entry count (`migrated.len() ==
/// original_entry_count`), each entry slot is replaced 1:1 in position,
/// preserving the original interleaving with custom/state wrappers (the
/// common case; matches pre-migration-aware `load()` behavior exactly).
///
/// If the migrator added or dropped entries, positional correspondence to
/// individual slots is no longer defined — we collapse all migrated entries
/// into the first entry slot (or append them if no entry slots existed)
/// and keep pass-through wrappers in their original positions.
fn weave_migrated_entries(
    classified: Vec<ClassifiedLine>,
    migrated: Vec<SessionEntry>,
    original_entry_count: usize,
) -> Vec<ClassifiedLine> {
    let mut rebuilt: Vec<ClassifiedLine> = Vec::with_capacity(classified.len().max(migrated.len()));
    if migrated.len() == original_entry_count {
        let mut iter = migrated.into_iter();
        for item in classified {
            match item {
                ClassifiedLine::Entry(_) => {
                    if let Some(next) = iter.next() {
                        rebuilt.push(ClassifiedLine::Entry(Box::new(next)));
                    }
                }
                passthrough @ (ClassifiedLine::Custom(_) | ClassifiedLine::State) => {
                    rebuilt.push(passthrough);
                }
            }
        }
    } else {
        let mut drained = false;
        let mut migrated_iter = migrated.into_iter();
        for item in classified {
            match item {
                ClassifiedLine::Entry(_) => {
                    if !drained {
                        for e in migrated_iter.by_ref() {
                            rebuilt.push(ClassifiedLine::Entry(Box::new(e)));
                        }
                        drained = true;
                    }
                }
                passthrough @ (ClassifiedLine::Custom(_) | ClassifiedLine::State) => {
                    rebuilt.push(passthrough);
                }
            }
        }
        for e in migrated_iter {
            rebuilt.push(ClassifiedLine::Entry(Box::new(e)));
        }
    }
    rebuilt
}

/// A JSONL line after parsing, classified so the migration pipeline can treat
/// migrateable entries distinctly from opaque pass-through wrappers.
///
/// The `Entry` variant is boxed because [`SessionEntry`] is comparatively
/// large (embeds a full `LlmMessage`) and dwarfs the other variants — the
/// `Box` keeps the enum small enough to avoid `clippy::large_enum_variant`.
#[derive(Debug, Clone)]
enum ClassifiedLine {
    /// A [`SessionEntry`] — either a tagged entry or a legacy raw
    /// [`LlmMessage`] rehydrated as [`SessionEntry::Message`]. Migrators run
    /// against these.
    Entry(Box<SessionEntry>),
    /// A `_custom: true` envelope — preserved verbatim so
    /// [`SessionStore::load`] can restore it via the caller's registry.
    Custom(serde_json::Value),
    /// A `_state: true` record. `load_state()` reads state directly from the
    /// file; `load()` and `load_entries()` both skip these, so the payload
    /// is intentionally discarded here.
    State,
}

/// Classify a single non-empty JSONL message line.
///
/// Order of checks mirrors the persisted format precedence: state/custom
/// wrappers win over the tagged [`SessionEntry`] path, which in turn covers
/// the legacy raw [`LlmMessage`] format.
fn classify_line(line: &str) -> io::Result<ClassifiedLine> {
    let value: serde_json::Value = serde_json::from_str(line).map_err(io::Error::other)?;

    if value.get("_state").and_then(serde_json::Value::as_bool) == Some(true) {
        return Ok(ClassifiedLine::State);
    }
    if value.get("_custom").and_then(serde_json::Value::as_bool) == Some(true) {
        return Ok(ClassifiedLine::Custom(value));
    }

    // Neither state nor custom wrapper: either a tagged SessionEntry or a
    // legacy raw LlmMessage. SessionEntry::parse handles both.
    SessionEntry::parse(line)
        .map(|entry| ClassifiedLine::Entry(Box::new(entry)))
        .map_err(io::Error::other)
}

/// Convert a `_custom` envelope (as stored on disk) back into an
/// [`AgentMessage::Custom`] using the supplied registry.
///
/// Returns `Ok(None)` when the caller supplied no registry (custom messages
/// are skipped — matching prior `load()` behavior).
fn custom_envelope_to_message(
    envelope: &serde_json::Value,
    registry: Option<&CustomMessageRegistry>,
) -> io::Result<Option<AgentMessage>> {
    let line = serde_json::to_string(envelope).map_err(io::Error::other)?;
    crate::codec::decode_jsonl_message_line(&line, registry)
}

fn classified_lines_to_messages(
    classified: Vec<ClassifiedLine>,
    registry: Option<&CustomMessageRegistry>,
    id: &str,
) -> Vec<AgentMessage> {
    let mut messages = Vec::new();
    for item in classified {
        match item {
            ClassifiedLine::Entry(entry) => {
                if let SessionEntry::Message(llm_msg) = *entry {
                    messages.push(AgentMessage::Llm(llm_msg));
                }
            }
            ClassifiedLine::Custom(envelope) => {
                match custom_envelope_to_message(&envelope, registry) {
                    Ok(Some(msg)) => messages.push(msg),
                    Ok(None) => {}
                    Err(error) => {
                        tracing::warn!(
                            error = %error,
                            "skipping unrestorable custom message in session {id}"
                        );
                    }
                }
            }
            ClassifiedLine::State => {}
        }
    }

    messages
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn new_session_id_format() {
        let id = JsonlSessionStore::new_session_id();
        let (timestamp, suffix) = id.rsplit_once('_').unwrap();
        assert_eq!(timestamp.len(), 15);
        assert_eq!(timestamp.as_bytes()[8], b'_');
        assert_eq!(suffix.len(), 32);
    }

    #[test]
    fn validate_session_id_rejects_slash() {
        let err = validate_session_id("foo/bar").unwrap_err();
        assert_eq!(err.kind(), io::ErrorKind::InvalidInput);
    }

    #[test]
    fn validate_session_id_rejects_backslash() {
        let err = validate_session_id("foo\\bar").unwrap_err();
        assert_eq!(err.kind(), io::ErrorKind::InvalidInput);
    }

    #[test]
    fn validate_session_id_rejects_dotdot() {
        let err = validate_session_id("..secret").unwrap_err();
        assert_eq!(err.kind(), io::ErrorKind::InvalidInput);
    }

    #[test]
    fn validate_session_id_rejects_colon() {
        let err = validate_session_id("C:drive").unwrap_err();
        assert_eq!(err.kind(), io::ErrorKind::InvalidInput);
    }

    #[test]
    fn validate_session_id_rejects_control_chars() {
        let err = validate_session_id("foo\nbar").unwrap_err();
        assert_eq!(err.kind(), io::ErrorKind::InvalidInput);
    }

    #[test]
    fn validate_session_id_rejects_empty() {
        let err = validate_session_id("").unwrap_err();
        assert_eq!(err.kind(), io::ErrorKind::InvalidInput);
    }

    #[test]
    fn validate_session_id_accepts_normal() {
        validate_session_id("20250315_120000").unwrap();
        validate_session_id("my-session").unwrap();
        validate_session_id("session_123").unwrap();
    }

    #[test]
    fn save_load_roundtrip_with_custom_messages() {
        use swink_agent::{AgentMessage, CustomMessage, CustomMessageRegistry};

        #[derive(Debug)]
        struct TestCustomMsg {
            data: String,
        }

        impl CustomMessage for TestCustomMsg {
            fn as_any(&self) -> &dyn std::any::Any {
                self
            }
            fn type_name(&self) -> Option<&str> {
                Some("TestCustomMsg")
            }
            fn to_json(&self) -> Option<serde_json::Value> {
                Some(serde_json::json!({ "data": self.data }))
            }
        }

        let dir = tempfile::tempdir().unwrap();
        let store = JsonlSessionStore::new(dir.path().to_path_buf()).unwrap();

        let meta = SessionMeta {
            id: "test-full".to_string(),
            title: "Full roundtrip".to_string(),
            created_at: chrono::DateTime::from_timestamp(1_710_500_000, 0)
                .unwrap()
                .to_utc(),
            updated_at: chrono::DateTime::from_timestamp(1_710_500_000, 0)
                .unwrap()
                .to_utc(),
            version: 1,
            sequence: 0,
        };

        let messages: Vec<AgentMessage> = vec![
            AgentMessage::Llm(swink_agent::LlmMessage::User(
                swink_agent::UserMessage::new(vec![swink_agent::ContentBlock::Text {
                    text: "Hello".to_string(),
                }])
                .with_timestamp(100),
            )),
            AgentMessage::Custom(Box::new(TestCustomMsg {
                data: "custom-payload".to_string(),
            })),
            AgentMessage::Llm(swink_agent::LlmMessage::User(
                swink_agent::UserMessage::new(vec![swink_agent::ContentBlock::Text {
                    text: "World".to_string(),
                }])
                .with_timestamp(200),
            )),
        ];

        store.save("test-full", &meta, &messages).unwrap();

        // Load with registry — custom message restored
        let mut registry = CustomMessageRegistry::new();
        registry.register(
            "TestCustomMsg",
            Box::new(|val: serde_json::Value| {
                let data = val
                    .get("data")
                    .and_then(|v| v.as_str())
                    .ok_or_else(|| "missing data".to_string())?;
                Ok(Box::new(TestCustomMsg {
                    data: data.to_string(),
                }) as Box<dyn CustomMessage>)
            }),
        );

        let (loaded_meta, loaded_messages) = store.load("test-full", Some(&registry)).unwrap();
        assert_eq!(loaded_meta.id, "test-full");
        assert_eq!(loaded_messages.len(), 3);
        assert!(matches!(
            loaded_messages[0],
            AgentMessage::Llm(swink_agent::LlmMessage::User(_))
        ));
        assert!(matches!(loaded_messages[1], AgentMessage::Custom(_)));
        assert!(matches!(
            loaded_messages[2],
            AgentMessage::Llm(swink_agent::LlmMessage::User(_))
        ));

        // Verify custom message content via downcast
        let custom = loaded_messages[1].downcast_ref::<TestCustomMsg>().unwrap();
        assert_eq!(custom.data, "custom-payload");

        // Load without registry — custom messages skipped
        let (_, loaded_no_reg) = store.load("test-full", None).unwrap();
        assert_eq!(loaded_no_reg.len(), 2);
        assert!(matches!(loaded_no_reg[0], AgentMessage::Llm(_)));
        assert!(matches!(loaded_no_reg[1], AgentMessage::Llm(_)));
    }

    #[test]
    fn append_preserves_saved_state() {
        let dir = tempfile::tempdir().unwrap();
        let store = JsonlSessionStore::new(dir.path().to_path_buf()).unwrap();

        let now = now_utc();
        let meta = SessionMeta {
            id: "test-state".to_string(),
            title: "State".to_string(),
            created_at: now,
            updated_at: now,
            version: 1,
            sequence: 0,
        };

        let initial_messages: Vec<AgentMessage> = vec![AgentMessage::Llm(LlmMessage::User(
            swink_agent::UserMessage::new(vec![swink_agent::ContentBlock::Text {
                text: "hello".to_string(),
            }])
            .with_timestamp(1),
        ))];
        store.save("test-state", &meta, &initial_messages).unwrap();
        store
            .save_state("test-state", &serde_json::json!({ "cursor": 1 }))
            .unwrap();

        let appended_messages: Vec<AgentMessage> = vec![AgentMessage::Llm(LlmMessage::User(
            swink_agent::UserMessage::new(vec![swink_agent::ContentBlock::Text {
                text: "world".to_string(),
            }])
            .with_timestamp(2),
        ))];
        store.append("test-state", &appended_messages).unwrap();

        let state = store.load_state("test-state").unwrap();
        assert_eq!(state, Some(serde_json::json!({ "cursor": 1 })));

        let (_, messages) = store.load("test-state", None).unwrap();
        assert_eq!(messages.len(), 2);
    }

    #[test]
    fn load_state_errors_on_corrupted_state_line() {
        let dir = tempfile::tempdir().unwrap();
        let store = JsonlSessionStore::new(dir.path().to_path_buf()).unwrap();

        let meta = fresh_meta("corrupt-state");
        store
            .save("corrupt-state", &meta, &[user_msg("hello", 1)])
            .unwrap();

        let path = session_path(dir.path(), "corrupt-state");
        let mut contents = std::fs::read_to_string(&path).unwrap();
        contents.push_str("{\"_state\":true,\"data\":\n");
        std::fs::write(&path, contents).unwrap();

        let err = store.load_state("corrupt-state").unwrap_err();
        assert_eq!(err.kind(), io::ErrorKind::InvalidData);
        assert!(err.to_string().contains("state line"));
    }

    #[test]
    fn atomic_write_leaves_no_temp_file_on_success() {
        let dir = tempfile::tempdir().unwrap();
        let target = dir.path().join("sess.jsonl");
        atomic_write(&target, |w| {
            w.write_all(b"hello\n")?;
            Ok(())
        })
        .unwrap();
        assert_eq!(std::fs::read_to_string(&target).unwrap(), "hello\n");
        // No leftover temp files in the directory
        for entry in std::fs::read_dir(dir.path()).unwrap() {
            let name = entry.unwrap().file_name().into_string().unwrap();
            assert!(
                !name.contains(".tmp."),
                "unexpected temp file left behind: {name}"
            );
        }
    }

    #[test]
    fn atomic_write_cleans_up_temp_on_error() {
        let dir = tempfile::tempdir().unwrap();
        let target = dir.path().join("sess.jsonl");
        let err = atomic_write(&target, |_w| {
            Err(io::Error::other("simulated mid-write failure"))
        })
        .unwrap_err();
        assert_eq!(err.to_string(), "simulated mid-write failure");
        // Target must not exist (no zero-length file)
        assert!(!target.exists());
        // And no leftover temp file
        for entry in std::fs::read_dir(dir.path()).unwrap() {
            let name = entry.unwrap().file_name().into_string().unwrap();
            assert!(!name.contains(".tmp."), "temp file not cleaned up: {name}");
        }
    }

    #[test]
    fn atomic_write_replaces_existing_file() {
        // Regression: on Windows, std::fs::rename does not replace an existing
        // destination, so rewrites of existing session files would fail.
        let dir = tempfile::tempdir().unwrap();
        let target = dir.path().join("sess.jsonl");
        std::fs::write(&target, b"old content\n").unwrap();
        atomic_write(&target, |w| {
            w.write_all(b"new content\n")?;
            Ok(())
        })
        .unwrap();
        assert_eq!(std::fs::read_to_string(&target).unwrap(), "new content\n");
        // No leftover temp files
        for entry in std::fs::read_dir(dir.path()).unwrap() {
            let name = entry.unwrap().file_name().into_string().unwrap();
            assert!(!name.contains(".tmp."), "temp file left behind: {name}");
        }
    }

    #[test]
    fn atomic_write_concurrent_rewrites_of_same_target_do_not_collide() {
        // Regression: the temp file path must be unique per write attempt.
        // If two overlapping rewrites in the same process shared a temp path,
        // they could truncate or rename each other's files nondeterministically.
        use std::sync::Arc;
        use std::sync::Barrier;
        use std::thread;

        let dir = tempfile::tempdir().unwrap();
        let target = Arc::new(dir.path().join("sess.jsonl"));
        std::fs::write(&*target, b"initial\n").unwrap();

        let n = 8;
        let barrier = Arc::new(Barrier::new(n));
        let mut handles = Vec::with_capacity(n);
        for i in 0..n {
            let target = Arc::clone(&target);
            let barrier = Arc::clone(&barrier);
            handles.push(thread::spawn(move || {
                barrier.wait();
                atomic_write(&target, |w| {
                    // Write enough data that a torn rename would be observable.
                    for _ in 0..256 {
                        writeln!(w, "writer-{i}")?;
                    }
                    Ok(())
                })
            }));
        }
        for h in handles {
            h.join().unwrap().expect("concurrent atomic_write failed");
        }

        // Final file must be entirely one writer's content (no interleaving).
        let final_contents = std::fs::read_to_string(&*target).unwrap();
        let first_line = final_contents.lines().next().unwrap();
        assert!(first_line.starts_with("writer-"));
        for line in final_contents.lines() {
            assert_eq!(line, first_line, "file contains interleaved writes");
        }

        // No leftover temp files.
        for entry in std::fs::read_dir(dir.path()).unwrap() {
            let name = entry.unwrap().file_name().into_string().unwrap();
            assert!(!name.contains(".tmp."), "temp file left behind: {name}");
        }
    }

    #[test]
    fn save_preserves_previous_file_when_new_write_fails() {
        // Regression for #234: a failed rewrite must not truncate the live
        // file. With atomic rename, the original content survives.
        let dir = tempfile::tempdir().unwrap();
        let store = JsonlSessionStore::new(dir.path().to_path_buf()).unwrap();
        let now = now_utc();
        let meta = SessionMeta {
            id: "atomic".to_string(),
            title: "t".to_string(),
            created_at: now,
            updated_at: now,
            version: 1,
            sequence: 0,
        };
        let messages: Vec<AgentMessage> = vec![AgentMessage::Llm(LlmMessage::User(
            swink_agent::UserMessage::new(vec![swink_agent::ContentBlock::Text {
                text: "first".to_string(),
            }])
            .with_timestamp(1),
        ))];
        store.save("atomic", &meta, &messages).unwrap();

        let path = session_path(dir.path(), "atomic");
        let before = std::fs::read_to_string(&path).unwrap();
        assert!(!before.is_empty());

        // Simulate a failed rewrite. With atomic_write, the target file must
        // remain untouched.
        let _ = atomic_write(&path, |_w| Err(io::Error::other("boom")));

        let after = std::fs::read_to_string(&path).unwrap();
        assert_eq!(before, after, "failed write must not corrupt live file");

        // File still parses cleanly
        let (_, loaded) = store.load("atomic", None).unwrap();
        assert_eq!(loaded.len(), 1);
    }

    #[test]
    fn load_reads_message_entries_saved_via_save_entries() {
        let dir = tempfile::tempdir().unwrap();
        let store = JsonlSessionStore::new(dir.path().to_path_buf()).unwrap();

        let now = now_utc();
        let meta = SessionMeta {
            id: "entry-messages".to_string(),
            title: "Entry messages".to_string(),
            created_at: now,
            updated_at: now,
            version: 1,
            sequence: 0,
        };

        let entries = vec![
            SessionEntry::Message(LlmMessage::User(
                swink_agent::UserMessage::new(vec![swink_agent::ContentBlock::Text {
                    text: "hello".to_string(),
                }])
                .with_timestamp(1),
            )),
            SessionEntry::Label {
                text: "bookmark".to_string(),
                message_index: 0,
                timestamp: 2,
            },
            SessionEntry::Message(LlmMessage::User(
                swink_agent::UserMessage::new(vec![swink_agent::ContentBlock::Text {
                    text: "world".to_string(),
                }])
                .with_timestamp(3),
            )),
        ];

        store
            .save_entries("entry-messages", &meta, &entries)
            .unwrap();

        let (_, messages) = store.load("entry-messages", None).unwrap();
        assert_eq!(messages.len(), 2);
        assert!(matches!(
            messages[0],
            AgentMessage::Llm(LlmMessage::User(_))
        ));
        assert!(matches!(
            messages[1],
            AgentMessage::Llm(LlmMessage::User(_))
        ));
    }

    #[test]
    fn save_preserves_state_and_rich_entries() {
        let dir = tempfile::tempdir().unwrap();
        let store = JsonlSessionStore::new(dir.path().to_path_buf()).unwrap();

        let meta = fresh_meta("preserve-save");
        let entries = vec![
            SessionEntry::Message(LlmMessage::User(
                swink_agent::UserMessage::new(vec![swink_agent::ContentBlock::Text {
                    text: "hello".to_string(),
                }])
                .with_timestamp(1),
            )),
            SessionEntry::Label {
                text: "bookmark".to_string(),
                message_index: 0,
                timestamp: 2,
            },
        ];

        store
            .save_entries("preserve-save", &meta, &entries)
            .unwrap();
        store
            .save_state("preserve-save", &serde_json::json!({ "cursor": 7 }))
            .unwrap();

        let (loaded_meta, _) = store.load("preserve-save", None).unwrap();
        store
            .save(
                "preserve-save",
                &loaded_meta,
                &[user_msg("updated", 3), user_msg("again", 4)],
            )
            .unwrap();

        assert_eq!(
            store.load_state("preserve-save").unwrap(),
            Some(serde_json::json!({ "cursor": 7 }))
        );

        let (_, entries) = store.load_entries("preserve-save").unwrap();
        assert_eq!(entries.len(), 3, "messages plus preserved label");
        assert!(matches!(entries[0], SessionEntry::Message(_)));
        assert!(matches!(entries[1], SessionEntry::Message(_)));
        assert!(matches!(entries[2], SessionEntry::Label { .. }));
    }

    #[test]
    fn save_entries_preserve_saved_state() {
        let dir = tempfile::tempdir().unwrap();
        let store = JsonlSessionStore::new(dir.path().to_path_buf()).unwrap();

        let meta = fresh_meta("preserve-entry-save");
        store
            .save("preserve-entry-save", &meta, &[user_msg("hello", 1)])
            .unwrap();
        store
            .save_state("preserve-entry-save", &serde_json::json!({ "cursor": 11 }))
            .unwrap();

        let (loaded_meta, _) = store.load("preserve-entry-save", None).unwrap();
        let entries = vec![
            SessionEntry::Message(LlmMessage::User(
                swink_agent::UserMessage::new(vec![swink_agent::ContentBlock::Text {
                    text: "updated".to_string(),
                }])
                .with_timestamp(2),
            )),
            SessionEntry::Label {
                text: "kept".to_string(),
                message_index: 0,
                timestamp: 3,
            },
        ];

        store
            .save_entries("preserve-entry-save", &loaded_meta, &entries)
            .unwrap();

        assert_eq!(
            store.load_state("preserve-entry-save").unwrap(),
            Some(serde_json::json!({ "cursor": 11 }))
        );
    }

    #[test]
    fn save_entries_preserves_existing_custom_message_envelopes() {
        use swink_agent::{CustomMessage, CustomMessageRegistry};

        #[derive(Debug)]
        struct TestCustomMsg {
            data: String,
        }

        impl CustomMessage for TestCustomMsg {
            fn as_any(&self) -> &dyn std::any::Any {
                self
            }

            fn type_name(&self) -> Option<&str> {
                Some("TestCustomMsg")
            }

            fn to_json(&self) -> Option<serde_json::Value> {
                Some(serde_json::json!({ "data": self.data }))
            }
        }

        let dir = tempfile::tempdir().unwrap();
        let store = JsonlSessionStore::new(dir.path().to_path_buf()).unwrap();

        let meta = fresh_meta("preserve-entry-custom");
        let messages = vec![
            user_msg("hello", 1),
            AgentMessage::Custom(Box::new(TestCustomMsg {
                data: "custom-payload".to_string(),
            })),
        ];
        store
            .save("preserve-entry-custom", &meta, &messages)
            .unwrap();

        let (loaded_meta, _) = store.load("preserve-entry-custom", None).unwrap();
        let entries = vec![
            SessionEntry::Message(LlmMessage::User(
                swink_agent::UserMessage::new(vec![swink_agent::ContentBlock::Text {
                    text: "updated".to_string(),
                }])
                .with_timestamp(2),
            )),
            SessionEntry::Label {
                text: "kept".to_string(),
                message_index: 0,
                timestamp: 3,
            },
        ];

        store
            .save_entries("preserve-entry-custom", &loaded_meta, &entries)
            .unwrap();

        let mut registry = CustomMessageRegistry::new();
        registry.register(
            "TestCustomMsg",
            Box::new(|val: serde_json::Value| {
                let data = val
                    .get("data")
                    .and_then(|value| value.as_str())
                    .ok_or_else(|| "missing data".to_string())?;
                Ok(Box::new(TestCustomMsg {
                    data: data.to_string(),
                }) as Box<dyn CustomMessage>)
            }),
        );

        let (_, loaded_messages) = store
            .load("preserve-entry-custom", Some(&registry))
            .unwrap();
        assert_eq!(loaded_messages.len(), 2, "message plus preserved custom");
        assert!(matches!(loaded_messages[0], AgentMessage::Llm(_)));
        let custom = loaded_messages[1].downcast_ref::<TestCustomMsg>().unwrap();
        assert_eq!(custom.data, "custom-payload");
    }

    fn user_msg(text: &str, ts: u64) -> AgentMessage {
        AgentMessage::Llm(LlmMessage::User(
            swink_agent::UserMessage::new(vec![swink_agent::ContentBlock::Text {
                text: text.to_string(),
            }])
            .with_timestamp(ts),
        ))
    }

    fn fresh_meta(id: &str) -> SessionMeta {
        let now = now_utc();
        SessionMeta {
            id: id.to_string(),
            title: "t".to_string(),
            created_at: now,
            updated_at: now,
            version: 1,
            sequence: 0,
        }
    }

    fn user_entry(text: &str, ts: u64) -> SessionEntry {
        SessionEntry::Message(LlmMessage::User(
            swink_agent::UserMessage::new(vec![swink_agent::ContentBlock::Text {
                text: text.to_string(),
            }])
            .with_timestamp(ts),
        ))
    }

    #[test]
    fn append_entries_appends_without_rewriting_existing_lines() {
        let dir = tempfile::tempdir().unwrap();
        let store = JsonlSessionStore::new(dir.path().to_path_buf()).unwrap();

        let meta = fresh_meta("append-entries");
        store
            .save_entries("append-entries", &meta, &[user_entry("first", 1)])
            .unwrap();
        let (created_meta, _) = store.load_entries("append-entries").unwrap();

        // Capture the byte length of the existing records so we can prove the
        // append did not rewrite them.
        let path = session_path(dir.path(), "append-entries");
        let before = std::fs::read_to_string(&path).unwrap();

        let returned = store
            .append_entries(
                "append-entries",
                &created_meta,
                &[user_entry("second", 2), user_entry("third", 3)],
            )
            .unwrap();

        // Sequence bumped and returned without a re-read.
        assert_eq!(returned.sequence, created_meta.sequence + 1);

        // All entries are present and ordered.
        let (loaded_meta, entries) = store.load_entries("append-entries").unwrap();
        assert_eq!(loaded_meta.sequence, returned.sequence);
        let texts: Vec<&str> = entries
            .iter()
            .filter_map(|e| match e {
                SessionEntry::Message(LlmMessage::User(u)) => match &u.content[0] {
                    swink_agent::ContentBlock::Text { text } => Some(text.as_str()),
                    _ => None,
                },
                _ => None,
            })
            .collect();
        assert_eq!(texts, vec!["first", "second", "third"]);

        // The original records remain verbatim at the head of the file —
        // i.e. the append did not rewrite them.
        let after = std::fs::read_to_string(&path).unwrap();
        let first_record_line = before.lines().nth(1).unwrap();
        assert!(
            after.contains(first_record_line),
            "first record should be preserved byte-for-byte by an in-place append"
        );
        assert!(
            after.len() > before.len(),
            "appended records should grow the file"
        );
    }

    #[test]
    fn append_entries_rejects_sequence_conflict() {
        let dir = tempfile::tempdir().unwrap();
        let store = JsonlSessionStore::new(dir.path().to_path_buf()).unwrap();

        let meta = fresh_meta("append-conflict");
        store
            .save_entries("append-conflict", &meta, &[user_entry("first", 1)])
            .unwrap();

        // Stale meta (sequence behind the on-disk value) must be rejected.
        let mut stale = meta;
        stale.sequence = 99;
        let err = store
            .append_entries("append-conflict", &stale, &[user_entry("second", 2)])
            .unwrap_err();
        assert_eq!(err.kind(), io::ErrorKind::AlreadyExists);
    }

    #[test]
    fn append_entries_errors_on_missing_session() {
        let dir = tempfile::tempdir().unwrap();
        let store = JsonlSessionStore::new(dir.path().to_path_buf()).unwrap();

        let meta = fresh_meta("nope");
        let err = store
            .append_entries("nope", &meta, &[user_entry("x", 1)])
            .unwrap_err();
        assert_eq!(err.kind(), io::ErrorKind::NotFound);
    }

    #[test]
    fn search_scans_across_saved_sessions() {
        let dir = tempfile::tempdir().unwrap();
        let store = JsonlSessionStore::new(dir.path().to_path_buf()).unwrap();

        let mut meta_a = fresh_meta("search-a");
        meta_a.title = "Auth notes".to_string();
        store
            .save(
                "search-a",
                &meta_a,
                &[
                    user_msg("We decided the auth middleware owns refresh tokens", 10),
                    user_msg("Unrelated deployment note", 11),
                ],
            )
            .unwrap();

        let mut meta_b = fresh_meta("search-b");
        meta_b.title = "Billing notes".to_string();
        store
            .save(
                "search-b",
                &meta_b,
                &[user_msg("Billing retries use exponential backoff", 20)],
            )
            .unwrap();

        let hits = store
            .search("auth middleware", &SessionSearchOptions::default())
            .unwrap();

        assert_eq!(hits.len(), 1);
        assert_eq!(hits[0].session_id, "search-a");
        assert_eq!(hits[0].session_title, "Auth notes");
        assert!(hits[0].snippet.contains("auth middleware"));
        assert!(matches!(hits[0].entry, SessionEntry::Message(_)));
    }

    #[test]
    fn save_canonicalizes_mismatched_metadata_id() {
        let dir = tempfile::tempdir().unwrap();
        let store = JsonlSessionStore::new(dir.path().to_path_buf()).unwrap();

        store
            .save(
                "canonical-save",
                &fresh_meta("wrong-save"),
                &[user_msg("canonical metadata", 10)],
            )
            .unwrap();

        let session_file = dir.path().join("canonical-save.jsonl");
        let first_line = std::fs::read_to_string(session_file)
            .unwrap()
            .lines()
            .next()
            .unwrap()
            .to_string();
        let raw_meta: SessionMeta = serde_json::from_str(&first_line).unwrap();
        assert_eq!(raw_meta.id, "canonical-save");

        let (loaded_meta, _) = store.load("canonical-save", None).unwrap();
        assert_eq!(loaded_meta.id, "canonical-save");

        let listed = store.list().unwrap();
        assert_eq!(listed.len(), 1);
        assert_eq!(listed[0].id, "canonical-save");
    }

    #[test]
    fn save_full_canonicalizes_returned_and_persisted_metadata_id() {
        let dir = tempfile::tempdir().unwrap();
        let store = JsonlSessionStore::new(dir.path().to_path_buf()).unwrap();

        let returned = store
            .save_full(
                "canonical-full",
                &fresh_meta("wrong-full"),
                &[user_msg("full canonical metadata", 10)],
                &serde_json::json!({"persisted": true}),
            )
            .unwrap();

        assert_eq!(returned.id, "canonical-full");
        let (loaded_meta, _, loaded_state) = store.load_full("canonical-full", None).unwrap();
        assert_eq!(loaded_meta.id, "canonical-full");
        assert_eq!(loaded_state, Some(serde_json::json!({"persisted": true})));
    }

    #[test]
    fn search_uses_canonical_session_id_for_mismatched_saved_metadata() {
        let dir = tempfile::tempdir().unwrap();
        let store = JsonlSessionStore::new(dir.path().to_path_buf()).unwrap();

        store
            .save_entries(
                "canonical-search",
                &fresh_meta("wrong-search"),
                &[user_entry("canonical search metadata", 10)],
            )
            .unwrap();

        let hits = store
            .search("canonical metadata", &SessionSearchOptions::default())
            .unwrap();

        assert_eq!(hits.len(), 1);
        assert_eq!(hits[0].session_id, "canonical-search");
    }

    #[test]
    fn search_respects_session_type_time_and_limit_filters() {
        let dir = tempfile::tempdir().unwrap();
        let store = JsonlSessionStore::new(dir.path().to_path_buf()).unwrap();

        store
            .save_entries(
                "filtered-a",
                &fresh_meta("filtered-a"),
                &[
                    user_entry("auth middleware message", 10),
                    SessionEntry::Label {
                        text: "auth middleware bookmark".to_string(),
                        message_index: 0,
                        timestamp: 20,
                    },
                    SessionEntry::Label {
                        text: "auth middleware late bookmark".to_string(),
                        message_index: 1,
                        timestamp: 40,
                    },
                ],
            )
            .unwrap();
        store
            .save_entries(
                "filtered-b",
                &fresh_meta("filtered-b"),
                &[SessionEntry::Label {
                    text: "auth middleware other session".to_string(),
                    message_index: 0,
                    timestamp: 20,
                }],
            )
            .unwrap();

        let options = SessionSearchOptions {
            session_ids: Some(vec!["filtered-a".to_string()]),
            entry_types: Some(vec!["label".to_string()]),
            start_time: Some(chrono::DateTime::from_timestamp(15, 0).unwrap().to_utc()),
            end_time: Some(chrono::DateTime::from_timestamp(25, 0).unwrap().to_utc()),
            max_results: Some(1),
        };

        let hits = store.search("auth middleware", &options).unwrap();

        assert_eq!(hits.len(), 1);
        assert_eq!(hits[0].session_id, "filtered-a");
        assert!(matches!(
            hits[0].entry,
            SessionEntry::Label { timestamp: 20, .. }
        ));
    }

    fn rewrite_meta_without_padding(path: &Path, id: &str, update: impl FnOnce(&mut SessionMeta)) {
        let (mut meta, lines) = read_meta_and_message_lines(path, id).unwrap();
        update(&mut meta);

        let mut contents = format!("{}\n", serde_json::to_string(&meta).unwrap());
        for line in lines {
            contents.push_str(&line);
            contents.push('\n');
        }
        std::fs::write(path, contents).unwrap();
    }

    #[test]
    fn append_advances_sequence_and_rejects_stale_save() {
        let dir = tempfile::tempdir().unwrap();
        let store = JsonlSessionStore::new(dir.path().to_path_buf()).unwrap();

        let meta = fresh_meta("seq-append");
        store
            .save("seq-append", &meta, &[user_msg("a", 1)])
            .unwrap();
        // After save, on-disk sequence == 1. `meta` still holds 0.

        store.append("seq-append", &[user_msg("b", 2)]).unwrap();
        // After append, on-disk sequence must have advanced to 2.

        // Sanity: list() / load() sees the bumped sequence.
        let (loaded_meta, _) = store.load("seq-append", None).unwrap();
        assert_eq!(loaded_meta.sequence, 2);

        // A stale writer holding the pre-append meta (sequence == 1 after save)
        // should now be rejected by check_sequence.
        let mut stale = meta;
        stale.sequence = 1;
        let err = store
            .save("seq-append", &stale, &[user_msg("c", 3)])
            .unwrap_err();
        assert_eq!(err.kind(), io::ErrorKind::AlreadyExists);
    }

    #[test]
    fn append_extends_file_without_rewriting_existing_records() {
        let dir = tempfile::tempdir().unwrap();
        let store = JsonlSessionStore::new(dir.path().to_path_buf()).unwrap();

        let meta = fresh_meta("append-in-place");
        store
            .save("append-in-place", &meta, &[user_msg("first", 1)])
            .unwrap();

        let path = session_path(dir.path(), "append-in-place");
        let before = std::fs::read_to_string(&path).unwrap();
        let before_lines = before.lines().collect::<Vec<_>>();
        let before_meta_line_len = before_lines[0].len();
        let before_message_line = before_lines[1].to_string();

        store
            .append("append-in-place", &[user_msg("second", 2)])
            .unwrap();

        let after = std::fs::read_to_string(&path).unwrap();
        let after_lines = after.lines().collect::<Vec<_>>();
        assert_eq!(
            after_lines[0].len(),
            before_meta_line_len,
            "append should patch the reserved metadata line in place"
        );
        assert_eq!(
            after_lines[1], before_message_line,
            "append must leave existing record bytes untouched"
        );
        assert_eq!(
            after_lines.len(),
            5,
            "append should add begin, record, and internal metadata commit lines"
        );

        let (loaded_meta, loaded_messages) = store.load("append-in-place", None).unwrap();
        assert_eq!(loaded_meta.sequence, 2);
        assert_eq!(loaded_messages.len(), 2);
    }

    #[test]
    fn append_failure_before_metadata_commit_discards_uncommitted_records() {
        let dir = tempfile::tempdir().unwrap();
        let store = JsonlSessionStore::new(dir.path().to_path_buf()).unwrap();

        let meta = fresh_meta("append-before-commit");
        store
            .save("append-before-commit", &meta, &[user_msg("first", 1)])
            .unwrap();

        let path = session_path(dir.path(), "append-before-commit");
        let (mut append_meta, meta_line_len) =
            read_meta_with_line_len(&path, "append-before-commit").unwrap();
        append_meta.updated_at = now_utc();
        append_meta.sequence += 1;
        let second_line =
            SessionRecord::from_message(&user_msg("second", 2), "append-before-commit")
                .unwrap()
                .to_json_line()
                .unwrap();

        let err = append_records_in_place_with_hooks(
            &path,
            &append_meta,
            meta_line_len,
            &[second_line],
            |_| Err(io::Error::other("simulated commit write failure")),
            |_| Ok(()),
        )
        .unwrap_err();
        assert_eq!(err.to_string(), "simulated commit write failure");

        let (loaded_meta, loaded_messages) = store.load("append-before-commit", None).unwrap();
        assert_eq!(
            loaded_meta.sequence, 1,
            "metadata must remain at the prior committed sequence"
        );
        assert_eq!(
            loaded_messages.len(),
            1,
            "record lines after an unclosed append-begin marker are uncommitted"
        );
    }

    #[test]
    fn append_failure_before_metadata_cache_patch_recovers_committed_records() {
        let dir = tempfile::tempdir().unwrap();
        let store = JsonlSessionStore::new(dir.path().to_path_buf()).unwrap();

        let meta = fresh_meta("append-commit-first");
        store
            .save("append-commit-first", &meta, &[user_msg("first", 1)])
            .unwrap();

        let path = session_path(dir.path(), "append-commit-first");
        let (mut append_meta, meta_line_len) =
            read_meta_with_line_len(&path, "append-commit-first").unwrap();
        append_meta.updated_at = now_utc();
        append_meta.sequence += 1;
        let second_line =
            SessionRecord::from_message(&user_msg("second", 2), "append-commit-first")
                .unwrap()
                .to_json_line()
                .unwrap();

        let err = append_records_in_place_with_hook(
            &path,
            &append_meta,
            meta_line_len,
            &[second_line],
            |_| Err(io::Error::other("simulated cache patch failure")),
        )
        .unwrap_err();
        assert_eq!(err.to_string(), "simulated cache patch failure");

        let first_line_meta: SessionMeta = serde_json::from_str(
            std::fs::read_to_string(&path)
                .unwrap()
                .lines()
                .next()
                .unwrap(),
        )
        .unwrap();
        assert_eq!(
            first_line_meta.sequence, 1,
            "line 1 remains the old cache if the metadata patch is interrupted"
        );

        let (loaded_meta, loaded_messages) = store.load("append-commit-first", None).unwrap();
        assert_eq!(
            loaded_meta.sequence, 2,
            "load must recover the committed metadata record after an interrupted cache patch"
        );
        let listed = store.list().unwrap();
        assert_eq!(
            listed[0].sequence, 2,
            "list must recover the committed metadata record after an interrupted cache patch"
        );
        assert_eq!(
            loaded_messages.len(),
            2,
            "committed append records must remain recoverable after cache patch failure"
        );

        let mut stale = meta;
        stale.sequence = 1;
        let err = store
            .save("append-commit-first", &stale, &[user_msg("stale", 3)])
            .unwrap_err();
        assert_eq!(err.kind(), io::ErrorKind::AlreadyExists);
    }

    fn meta_commit_line(meta: &SessionMeta) -> String {
        SessionRecord::Meta(Box::new(meta.clone()))
            .to_json_line()
            .unwrap()
    }

    fn message_line(text: &str, ts: u64, id: &str) -> String {
        SessionRecord::from_message(&user_msg(text, ts), id)
            .unwrap()
            .to_json_line()
            .unwrap()
    }

    /// Writes a session file with a padded first metadata line followed by
    /// `body` verbatim (the caller controls trailing newlines / torn tails).
    fn write_raw_session_file(path: &Path, first_meta: &SessionMeta, body: &[u8]) {
        let mut file = std::fs::File::create(path).unwrap();
        write_meta_line(&mut file, first_meta, META_LINE_PADDING).unwrap();
        file.write_all(body).unwrap();
    }

    /// A ~1 KiB record line that is not a Meta commit.
    fn pad_line() -> String {
        format!("{{\"pad\":\"{}\"}}", "x".repeat(1024))
    }

    /// Asserts `read_meta_with_line_len` returns exactly what the forward
    /// line-by-line scan (`read_meta_and_message_lines`) computes, then
    /// returns the meta for further assertions.
    fn read_meta_checked_against_forward_scan(path: &Path, id: &str) -> SessionMeta {
        let (meta, _) = read_meta_with_line_len(path, id).unwrap();
        let (oracle, _) = read_meta_and_message_lines(path, id).unwrap();
        assert_eq!(
            meta, oracle,
            "tail scan must return exactly what the forward scan returns"
        );
        meta
    }

    #[test]
    fn read_meta_falls_back_to_first_line_without_meta_commits() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("no-commits.jsonl");
        let first = fresh_meta("no-commits");
        let body = format!(
            "{}\n{}\n",
            message_line("first", 1, "no-commits"),
            message_line("second", 2, "no-commits"),
        );
        write_raw_session_file(&path, &first, body.as_bytes());

        let meta = read_meta_checked_against_forward_scan(&path, "no-commits");
        assert_eq!(meta, first);

        let raw = std::fs::read_to_string(&path).unwrap();
        let (_, line_len) = read_meta_with_line_len(&path, "no-commits").unwrap();
        assert_eq!(
            line_len,
            raw.find('\n').unwrap() + 1,
            "line_len must be the padded first line including its newline"
        );
    }

    #[test]
    fn read_meta_returns_last_meta_commit_like_the_forward_scan() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("multi-meta.jsonl");
        let mut first = fresh_meta("multi-meta");
        first.title = "stale".to_string();

        let mut second = first.clone();
        second.title = "second".to_string();
        second.sequence = 2;
        // Commit records deliberately carry a different id: the reader must
        // canonicalize to the file id exactly like the forward scan did.
        second.id = "not-the-file-id".to_string();
        let mut third = second.clone();
        third.title = "third".to_string();
        third.sequence = 3;

        let body = format!(
            "{}\n{}\n{}\n{}\n",
            message_line("m1", 1, "multi-meta"),
            meta_commit_line(&second),
            message_line("m2", 2, "multi-meta"),
            meta_commit_line(&third),
        );
        write_raw_session_file(&path, &first, body.as_bytes());

        let meta = read_meta_checked_against_forward_scan(&path, "multi-meta");
        assert_eq!(meta.title, "third", "the LAST Meta commit record wins");
        assert_eq!(meta.sequence, 3);
        assert_eq!(
            meta.id, "multi-meta",
            "meta id must be canonicalized to the file id"
        );
    }

    #[test]
    fn read_meta_skips_trailing_partial_line_after_commit() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("torn-tail.jsonl");
        let first = fresh_meta("torn-tail");
        let mut committed = first.clone();
        committed.title = "committed".to_string();
        committed.sequence = 2;

        let mut body = format!(
            "{}\n{}\n",
            message_line("m1", 1, "torn-tail"),
            meta_commit_line(&committed),
        );
        // Crash-torn tail: a truncated record with no trailing newline. It is
        // valid UTF-8 but unparseable JSON, so it must be skipped.
        body.push_str("{\"_meta\":true,\"data\":{\"id\":\"torn");
        write_raw_session_file(&path, &first, body.as_bytes());

        let meta = read_meta_checked_against_forward_scan(&path, "torn-tail");
        assert_eq!(meta.title, "committed");
        assert_eq!(meta.sequence, 2);
    }

    #[test]
    fn read_meta_errors_on_empty_file() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("empty.jsonl");
        std::fs::File::create(&path).unwrap();

        let err = read_meta_with_line_len(&path, "empty").unwrap_err();
        assert_eq!(err.kind(), io::ErrorKind::InvalidData);
        assert!(err.to_string().contains("empty session file"));
    }

    #[test]
    fn read_meta_finds_last_commit_within_tail_window_of_large_file() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("large-tail-commit.jsonl");
        let first = fresh_meta("large-tail-commit");
        let mut committed = first.clone();
        committed.title = "committed".to_string();
        committed.sequence = 2;

        let mut body = String::new();
        for _ in 0..80 {
            body.push_str(&pad_line());
            body.push('\n');
        }
        body.push_str(&meta_commit_line(&committed));
        body.push('\n');
        assert!(
            body.len() as u64 > META_TAIL_WINDOW,
            "body must exceed the initial tail window to exercise the windowed scan"
        );
        write_raw_session_file(&path, &first, body.as_bytes());

        let meta = read_meta_checked_against_forward_scan(&path, "large-tail-commit");
        assert_eq!(meta.title, "committed");
        assert_eq!(meta.sequence, 2);
    }

    #[test]
    fn read_meta_recovers_commit_buried_before_large_uncommitted_tail() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("buried-commit.jsonl");
        let first = fresh_meta("buried-commit");
        let mut committed = first.clone();
        committed.title = "committed".to_string();
        committed.sequence = 2;

        // The only Meta commit sits before a crashed (uncommitted) append
        // whose records span several tail windows, forcing the scan to widen
        // and ultimately fall back to the full forward pass.
        let mut body = format!(
            "{}\n{}\n{}\n",
            message_line("m1", 1, "buried-commit"),
            meta_commit_line(&committed),
            SessionRecord::AppendBegin.to_json_line().unwrap(),
        );
        for _ in 0..200 {
            body.push_str(&pad_line());
            body.push('\n');
        }
        assert!(body.len() as u64 > 2 * META_TAIL_WINDOW);
        write_raw_session_file(&path, &first, body.as_bytes());

        let meta = read_meta_checked_against_forward_scan(&path, "buried-commit");
        assert_eq!(meta.title, "committed");
        assert_eq!(meta.sequence, 2);
    }

    #[test]
    fn read_meta_first_line_wins_in_large_file_without_commits() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("large-no-commits.jsonl");
        let first = fresh_meta("large-no-commits");

        let mut body = String::new();
        for _ in 0..200 {
            body.push_str(&pad_line());
            body.push('\n');
        }
        assert!(body.len() as u64 > 2 * META_TAIL_WINDOW);
        write_raw_session_file(&path, &first, body.as_bytes());

        let meta = read_meta_checked_against_forward_scan(&path, "large-no-commits");
        assert_eq!(meta, first);
    }

    #[test]
    fn read_meta_tolerates_torn_multibyte_tail() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("torn-utf8.jsonl");
        let first = fresh_meta("torn-utf8");
        let mut committed = first.clone();
        committed.title = "committed".to_string();
        committed.sequence = 2;

        // A crash mid-write can truncate a multi-byte UTF-8 character in the
        // tail line. Semantics per spec 021 FR-004 addendum (#1067): the torn
        // line is skipped and the last committed meta is still recovered,
        // in both the tail scan and the forward scan.
        let mut body = format!("{}\n", meta_commit_line(&committed)).into_bytes();
        body.extend_from_slice(b"{\"_meta\":true,\xE2\x82");
        write_raw_session_file(&path, &first, &body);

        let meta = read_meta_checked_against_forward_scan(&path, "torn-utf8");
        assert_eq!(meta, committed);
    }

    #[test]
    fn read_meta_tolerates_torn_multibyte_tail_in_large_file() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("large-torn-utf8.jsonl");
        let first = fresh_meta("large-torn-utf8");
        let mut committed = first.clone();
        committed.title = "committed".to_string();
        committed.sequence = 2;

        let mut body = String::new();
        for _ in 0..80 {
            body.push_str(&pad_line());
            body.push('\n');
        }
        body.push_str(&meta_commit_line(&committed));
        body.push('\n');
        assert!(body.len() as u64 > META_TAIL_WINDOW);
        let mut body = body.into_bytes();
        body.extend_from_slice(b"{\"_meta\":true,\xE2\x82");
        write_raw_session_file(&path, &first, &body);

        let meta = read_meta_checked_against_forward_scan(&path, "large-torn-utf8");
        assert_eq!(
            meta, committed,
            "windowed tail scan must skip invalid UTF-8 like the forward scan"
        );
    }

    #[test]
    fn list_reports_last_committed_meta_despite_trailing_garbage() {
        let dir = tempfile::tempdir().unwrap();
        let store = JsonlSessionStore::new(dir.path().to_path_buf()).unwrap();

        let mut first = fresh_meta("listed-session");
        first.title = "stale".to_string();
        let mut second = first.clone();
        second.title = "second".to_string();
        second.sequence = 2;
        let mut third = first.clone();
        third.title = "third".to_string();
        third.sequence = 3;

        let mut body = format!(
            "{}\n{}\n{}\n{}\n",
            message_line("m1", 1, "listed-session"),
            meta_commit_line(&second),
            message_line("m2", 2, "listed-session"),
            meta_commit_line(&third),
        );
        body.push_str("{\"_meta\":true,\"data\":{\"id\":\"torn");
        let path = session_path(dir.path(), "listed-session");
        write_raw_session_file(&path, &first, body.as_bytes());

        let listed = store.list().unwrap();
        assert_eq!(listed.len(), 1);
        assert_eq!(listed[0].id, "listed-session");
        assert_eq!(
            listed[0].title, "third",
            "list must surface the last committed Meta record"
        );
        assert_eq!(listed[0].sequence, 3);
    }

    #[test]
    fn load_waits_for_in_flight_append_commit_before_metadata_cache_patch() {
        use std::sync::mpsc;
        use std::thread;
        use std::time::Duration;

        let dir = tempfile::tempdir().unwrap();
        let store = JsonlSessionStore::new(dir.path().to_path_buf()).unwrap();

        let meta = fresh_meta("append-load-race");
        store
            .save("append-load-race", &meta, &[user_msg("first", 1)])
            .unwrap();

        let path = session_path(dir.path(), "append-load-race");
        let (patched_tx, patched_rx) = mpsc::channel();
        let (resume_tx, resume_rx) = mpsc::channel();

        let append_path = path;
        let append_handle = thread::spawn(move || {
            with_target_lock(&append_path, || {
                let (mut append_meta, meta_line_len) =
                    read_meta_with_line_len(&append_path, "append-load-race")?;
                append_meta.updated_at = now_utc();
                append_meta.sequence += 1;
                let second_line =
                    SessionRecord::from_message(&user_msg("second", 2), "append-load-race")
                        .unwrap()
                        .to_json_line()?;

                append_records_in_place_with_hook(
                    &append_path,
                    &append_meta,
                    meta_line_len,
                    &[second_line],
                    |_| {
                        patched_tx.send(()).unwrap();
                        resume_rx.recv().unwrap();
                        Ok(())
                    },
                )?;
                Ok(())
            })
        });

        patched_rx.recv().unwrap();

        let load_store = JsonlSessionStore::new(dir.path().to_path_buf()).unwrap();
        let (loaded_tx, loaded_rx) = mpsc::channel();
        let load_handle = thread::spawn(move || {
            loaded_tx
                .send(load_store.load("append-load-race", None))
                .unwrap();
        });

        assert!(
            loaded_rx.recv_timeout(Duration::from_millis(100)).is_err(),
            "load must wait for an append that has committed records but not patched line 1"
        );

        resume_tx.send(()).unwrap();
        append_handle
            .join()
            .unwrap()
            .expect("append should finish cleanly");

        let (loaded_meta, loaded_messages) = loaded_rx.recv().unwrap().unwrap();
        load_handle.join().unwrap();
        assert_eq!(loaded_meta.sequence, 2);
        assert_eq!(
            loaded_messages.len(),
            2,
            "load must expose committed records and metadata together"
        );
    }

    #[test]
    fn append_rewrite_failure_preserves_existing_file() {
        let dir = tempfile::tempdir().unwrap();
        let mut meta = fresh_meta("append-atomic");
        meta.sequence = 9;
        let path = session_path(dir.path(), "append-atomic");
        let first_line =
            SessionRecord::from_message(&user_msg("first", 1), "append-atomic").unwrap();
        std::fs::write(
            &path,
            format!(
                "{}\n{}\n",
                serde_json::to_string(&meta).unwrap(),
                first_line.to_json_line().unwrap()
            ),
        )
        .unwrap();
        let before = std::fs::read_to_string(&path).unwrap();

        let err = append_records_with_rewrite(
            &path,
            "append-atomic",
            [SessionRecord::from_message(&user_msg("second", 2), "append-atomic").unwrap()],
            |_path, _meta, _lines| Err(io::Error::other("simulated append rewrite failure")),
        )
        .unwrap_err();
        assert_eq!(err.to_string(), "simulated append rewrite failure");

        let after = std::fs::read_to_string(&path).unwrap();
        assert_eq!(
            before, after,
            "failed append rewrite must not modify the live session file"
        );

        let store = JsonlSessionStore::new(dir.path().to_path_buf()).unwrap();
        let (loaded_meta, loaded_messages) = store.load("append-atomic", None).unwrap();
        assert_eq!(loaded_meta.sequence, 9);
        assert_eq!(loaded_messages.len(), 1);
    }

    #[test]
    fn delete_waits_for_append_lock_and_does_not_allow_resurrection() {
        use std::sync::Arc;
        use std::sync::Barrier;
        use std::thread;

        let dir = tempfile::tempdir().unwrap();
        let store = Arc::new(JsonlSessionStore::new(dir.path().to_path_buf()).unwrap());
        let meta = fresh_meta("delete-race");
        store
            .save("delete-race", &meta, &[user_msg("first", 1)])
            .unwrap();

        let path = session_path(dir.path(), "delete-race");
        rewrite_meta_without_padding(&path, "delete-race", |meta| meta.sequence = 9);
        let append_ready = Arc::new(Barrier::new(2));
        let allow_rewrite = Arc::new(Barrier::new(2));

        let append_path = path.clone();
        let append_ready_for_thread = Arc::clone(&append_ready);
        let allow_rewrite_for_thread = Arc::clone(&allow_rewrite);
        let append_handle = thread::spawn(move || {
            append_records_with_rewrite(
                &append_path,
                "delete-race",
                [SessionRecord::from_message(&user_msg("second", 2), "delete-race").unwrap()],
                |path, meta, lines| {
                    append_ready_for_thread.wait();
                    allow_rewrite_for_thread.wait();
                    rewrite_session_file_locked(path, meta, lines)
                },
            )
        });

        append_ready.wait();

        let delete_store = Arc::clone(&store);
        let delete_handle = thread::spawn(move || delete_store.delete("delete-race"));

        // If delete bypassed the session lock, it could return before the
        // append rewrites and the session would be recreated after delete.
        allow_rewrite.wait();

        append_handle
            .join()
            .unwrap()
            .expect("append should finish cleanly");
        delete_handle
            .join()
            .unwrap()
            .expect("delete should wait for the append lock");

        assert!(
            !path.exists(),
            "delete must not allow a concurrent append to resurrect the session"
        );
    }

    #[test]
    fn save_state_advances_sequence_and_rejects_stale_save() {
        let dir = tempfile::tempdir().unwrap();
        let store = JsonlSessionStore::new(dir.path().to_path_buf()).unwrap();

        let meta = fresh_meta("seq-state");
        store.save("seq-state", &meta, &[user_msg("a", 1)]).unwrap();
        // On-disk sequence == 1.

        store
            .save_state("seq-state", &serde_json::json!({ "cursor": 1 }))
            .unwrap();

        let (loaded_meta, _) = store.load("seq-state", None).unwrap();
        assert_eq!(loaded_meta.sequence, 2);

        // Stale writer with sequence == 1 must be rejected.
        let mut stale = meta;
        stale.sequence = 1;
        let err = store
            .save("seq-state", &stale, &[user_msg("c", 3)])
            .unwrap_err();
        assert_eq!(err.kind(), io::ErrorKind::AlreadyExists);
    }

    #[test]
    fn save_full_updates_messages_and_state_with_single_sequence_bump() {
        let dir = tempfile::tempdir().unwrap();
        let store = JsonlSessionStore::new(dir.path().to_path_buf()).unwrap();

        let meta = fresh_meta("save-full");
        store
            .save("save-full", &meta, &[user_msg("before", 1)])
            .unwrap();
        store
            .save_state("save-full", &serde_json::json!({ "cursor": 1 }))
            .unwrap();

        let (loaded_meta, _) = store.load("save-full", None).unwrap();
        assert_eq!(loaded_meta.sequence, 2);

        let persisted_meta = store
            .save_full(
                "save-full",
                &loaded_meta,
                &[user_msg("after", 2), user_msg("again", 3)],
                &serde_json::json!({ "cursor": 9, "draft": "synced" }),
            )
            .unwrap();

        assert_eq!(
            persisted_meta.sequence, 3,
            "combined transcript+state save should advance sequence once"
        );

        let (reloaded_meta, reloaded_messages) = store.load("save-full", None).unwrap();
        assert_eq!(reloaded_meta.sequence, 3);
        assert_eq!(reloaded_messages.len(), 2);
        assert_eq!(
            store.load_state("save-full").unwrap(),
            Some(serde_json::json!({ "cursor": 9, "draft": "synced" }))
        );
    }

    #[cfg(feature = "search")]
    #[test]
    fn save_full_updates_warm_search_index() {
        let dir = tempfile::tempdir().unwrap();
        let store = JsonlSessionStore::new(dir.path().to_path_buf()).unwrap();

        let meta = fresh_meta("save-full-search");
        store.open_search_index().unwrap();
        store
            .save(
                "save-full-search",
                &meta,
                &[user_msg("stale auth decision", 1)],
            )
            .unwrap();

        let (loaded_meta, _) = store.load("save-full-search", None).unwrap();
        store
            .save_full(
                "save-full-search",
                &loaded_meta,
                &[user_msg("fresh billing decision", 2)],
                &serde_json::json!({ "cursor": 2 }),
            )
            .unwrap();

        let fresh_hits = store
            .search("fresh billing", &SessionSearchOptions::default())
            .unwrap();
        assert_eq!(fresh_hits.len(), 1);
        assert_eq!(fresh_hits[0].session_id, "save-full-search");

        let stale_hits = store
            .search("stale auth", &SessionSearchOptions::default())
            .unwrap();
        assert!(
            stale_hits.is_empty(),
            "save_full should replace stale search documents"
        );
    }

    /// Append paths must extend a warm index incrementally — adding only the
    /// new entries' documents — never by reloading and re-indexing the whole
    /// session (the pre-fix behavior, which made every append O(session)).
    ///
    /// Instrumented via `TantivyIndex::full_reindex_count`, a test-only
    /// counter of whole-session (re)index operations: cheap, exact, and it
    /// fails loudly if an append path regresses to `index_session`.
    #[cfg(feature = "search")]
    #[test]
    fn append_paths_update_search_index_without_full_reindex() {
        let dir = tempfile::tempdir().unwrap();
        let store = JsonlSessionStore::new(dir.path().to_path_buf()).unwrap();
        store.open_search_index().unwrap();

        let meta = fresh_meta("append-incremental");
        store
            .save_entries("append-incremental", &meta, &[user_entry("quokka", 1)])
            .unwrap();
        let index = store.active_tantivy_index().expect("index is open");
        let after_save = index.full_reindex_count();

        let (loaded_meta, _) = store.load_entries("append-incremental").unwrap();
        let bumped = store
            .append_entries(
                "append-incremental",
                &loaded_meta,
                &[user_entry("wombat", 2)],
            )
            .unwrap();
        store
            .append_entries("append-incremental", &bumped, &[user_entry("numbat", 3)])
            .unwrap();
        store
            .append("append-incremental", &[user_msg("kookaburra", 4)])
            .unwrap();

        assert_eq!(
            index.full_reindex_count(),
            after_save,
            "append paths must not re-index the whole session"
        );

        // Everything — original and appended — is searchable. The first
        // search triggers the lazy full build (which legitimately calls
        // index_session), so it runs after the counter assertion above.
        for term in ["quokka", "wombat", "numbat", "kookaburra"] {
            let hits = store
                .search(term, &SessionSearchOptions::default())
                .unwrap();
            assert_eq!(hits.len(), 1, "expected exactly one hit for {term:?}");
            assert_eq!(hits[0].session_id, "append-incremental");
        }

        // Steady state: once the index is built, appends stay incremental
        // and are immediately searchable.
        let after_build = index.full_reindex_count();
        let (steady_meta, _) = store.load_entries("append-incremental").unwrap();
        store
            .append_entries(
                "append-incremental",
                &steady_meta,
                &[user_entry("capybara", 5)],
            )
            .unwrap();
        assert_eq!(
            index.full_reindex_count(),
            after_build,
            "post-build appends must stay incremental"
        );
        let hits = store
            .search("capybara", &SessionSearchOptions::default())
            .unwrap();
        assert_eq!(hits.len(), 1);
    }

    #[test]
    fn load_full_returns_messages_and_state_from_one_snapshot() {
        let dir = tempfile::tempdir().unwrap();
        let store = JsonlSessionStore::new(dir.path().to_path_buf()).unwrap();

        let persisted_meta = store
            .save_full(
                "load-full",
                &fresh_meta("load-full"),
                &[user_msg("hello", 1), user_msg("again", 2)],
                &serde_json::json!({ "cursor": 4, "draft": "stable" }),
            )
            .unwrap();

        let (loaded_meta, loaded_messages, loaded_state) =
            store.load_full("load-full", None).unwrap();

        assert_eq!(loaded_meta.sequence, persisted_meta.sequence);
        assert_eq!(loaded_messages.len(), 2);
        assert_eq!(
            loaded_state,
            Some(serde_json::json!({ "cursor": 4, "draft": "stable" }))
        );
    }

    #[test]
    fn stale_saves_do_not_both_pass_validation() {
        use std::sync::mpsc;
        use std::thread;

        let dir = tempfile::tempdir().unwrap();
        let store = JsonlSessionStore::new(dir.path().to_path_buf()).unwrap();

        let meta = fresh_meta("stale-race");
        store
            .save("stale-race", &meta, &[user_msg("v1", 1)])
            .unwrap();
        let (stale_meta, _) = store.load("stale-race", None).unwrap();

        let path = session_path(dir.path(), "stale-race");
        let (validated_tx, validated_rx) = mpsc::channel();
        let (resume_tx, resume_rx) = mpsc::channel();

        let thread_path = path;
        let thread_meta = stale_meta.clone();
        let paused_writer = thread::spawn(move || {
            let messages = vec![user_msg("writer-1", 2)];
            save_messages_with_hooks(
                &thread_path,
                "stale-race",
                &thread_meta,
                &messages,
                || {
                    validated_tx.send(()).unwrap();
                    resume_rx.recv().unwrap();
                    Ok(())
                },
                write_messages_locked,
            )
        });

        validated_rx.recv().unwrap();

        let competitor_messages = vec![user_msg("writer-2", 3)];
        let competitor_meta = stale_meta;
        let competitor =
            thread::spawn(move || store.save("stale-race", &competitor_meta, &competitor_messages));

        resume_tx.send(()).unwrap();

        let first_result = paused_writer.join().unwrap();
        let second_result = competitor.join().unwrap();

        let conflict_count =
            usize::from(first_result.is_err()) + usize::from(second_result.is_err());
        assert_eq!(
            conflict_count, 1,
            "exactly one stale writer should conflict"
        );

        let (loaded_meta, loaded_messages) = JsonlSessionStore::new(dir.path().to_path_buf())
            .unwrap()
            .load("stale-race", None)
            .unwrap();
        assert_eq!(loaded_meta.sequence, 2);
        assert_eq!(loaded_messages.len(), 1);
    }

    #[test]
    fn save_interrupt_requires_existing_session() {
        let dir = tempfile::tempdir().unwrap();
        let store = JsonlSessionStore::new(dir.path().to_path_buf()).unwrap();

        let err = store
            .save_interrupt(
                "missing",
                &InterruptState {
                    interrupted_at: 1,
                    pending_tool_calls: vec![],
                    context_snapshot: vec![],
                    system_prompt: "system".to_string(),
                    model: swink_agent::ModelSpec::new("openai", "gpt-4o"),
                },
            )
            .unwrap_err();
        assert_eq!(err.kind(), io::ErrorKind::NotFound);
    }

    /// Migrator that transforms the first `User` text content in every
    /// `Message` entry from lower-case to upper-case. Used to detect whether
    /// the migration pipeline actually ran against a given load path.
    struct UppercasingMigrator;

    impl crate::migrate::SessionMigrator for UppercasingMigrator {
        fn source_version(&self) -> u32 {
            0
        }
        fn target_version(&self) -> u32 {
            1
        }
        fn migrate(
            &self,
            _meta: &SessionMeta,
            entries: Vec<SessionEntry>,
        ) -> io::Result<Vec<SessionEntry>> {
            Ok(entries
                .into_iter()
                .map(|entry| match entry {
                    SessionEntry::Message(LlmMessage::User(mut m)) => {
                        for block in &mut m.content {
                            if let swink_agent::ContentBlock::Text { text } = block {
                                *text = text.to_uppercase();
                            }
                        }
                        SessionEntry::Message(LlmMessage::User(m))
                    }
                    other => other,
                })
                .collect())
        }
    }

    /// Write a legacy (`version: 0`) session file containing both a
    /// migrateable raw-`LlmMessage` line AND a custom-message envelope, then
    /// assert that `load()` and `load_entries()` both observe the migrated
    /// shape and that `load()` returns the custom wrapper unchanged.
    ///
    /// Regression for #522: `load()` previously bypassed the configured
    /// migrator pipeline, so it returned the raw pre-migration text while
    /// `load_entries()` returned the migrated text.
    #[test]
    #[allow(clippy::too_many_lines)]
    fn load_applies_migrators_identically_to_load_entries() {
        use swink_agent::{CustomMessage, CustomMessageRegistry};

        #[derive(Debug)]
        struct TestCustomMsg {
            data: String,
        }

        impl CustomMessage for TestCustomMsg {
            fn as_any(&self) -> &dyn std::any::Any {
                self
            }
            fn type_name(&self) -> Option<&str> {
                Some("TestCustomMsg")
            }
            fn to_json(&self) -> Option<serde_json::Value> {
                Some(serde_json::json!({ "data": self.data }))
            }
        }

        let dir = tempfile::tempdir().unwrap();
        let id = "legacy";

        // Write a legacy-format file by hand: line 1 = meta @ version 0,
        // line 2 = raw LlmMessage (migrateable), line 3 = _custom envelope
        // (pass-through), line 4 = _state wrapper (pass-through, skipped by
        // load but must not break migration).
        let path = session_path(dir.path(), id);
        let now = now_utc();
        let meta = SessionMeta {
            id: id.to_string(),
            title: "legacy".to_string(),
            created_at: now,
            updated_at: now,
            version: 0,
            sequence: 0,
        };

        let raw_msg_line = serde_json::to_string(&LlmMessage::User(
            swink_agent::UserMessage::new(vec![swink_agent::ContentBlock::Text {
                text: "hello".to_string(),
            }])
            .with_timestamp(1),
        ))
        .unwrap();

        let custom_envelope = serde_json::json!({
            "type": "TestCustomMsg",
            "data": { "data": "custom-payload" },
            "_custom": true,
        });
        let custom_line = serde_json::to_string(&custom_envelope).unwrap();

        let state_line = serde_json::json!({
            "_state": true,
            "data": { "cursor": 9 },
        })
        .to_string();

        let contents = format!(
            "{}\n{raw_msg_line}\n{custom_line}\n{state_line}\n",
            serde_json::to_string(&meta).unwrap()
        );
        std::fs::write(&path, contents).unwrap();

        let store = JsonlSessionStore::new(dir.path().to_path_buf())
            .unwrap()
            .with_migrators(vec![Box::new(UppercasingMigrator)]);

        // load_entries() runs migrators — this is the pre-fix baseline.
        let (entries_meta, entries) = store.load_entries(id).unwrap();
        assert_eq!(entries_meta.version, crate::migrate::CURRENT_VERSION);
        let entry_msg = entries
            .iter()
            .find_map(SessionEntry::as_message)
            .expect("message entry present");
        if let LlmMessage::User(user) = entry_msg
            && let swink_agent::ContentBlock::Text { text } = &user.content[0]
        {
            assert_eq!(text, "HELLO", "load_entries must run migrator");
        } else {
            panic!("unexpected entry shape");
        }

        // load() must observe the same post-migration text AND still return
        // the custom wrapper unchanged (requires the registry).
        let mut registry = CustomMessageRegistry::new();
        registry.register(
            "TestCustomMsg",
            Box::new(|val: serde_json::Value| {
                let data = val
                    .get("data")
                    .and_then(|v| v.as_str())
                    .ok_or_else(|| "missing data".to_string())?;
                Ok(Box::new(TestCustomMsg {
                    data: data.to_string(),
                }) as Box<dyn CustomMessage>)
            }),
        );

        let (load_meta, messages) = store.load(id, Some(&registry)).unwrap();
        assert_eq!(load_meta.version, crate::migrate::CURRENT_VERSION);

        let llm_text = messages
            .iter()
            .find_map(|m| match m {
                AgentMessage::Llm(LlmMessage::User(u)) => u.content.iter().find_map(|b| match b {
                    swink_agent::ContentBlock::Text { text } => Some(text.clone()),
                    _ => None,
                }),
                _ => None,
            })
            .expect("user message present");
        assert_eq!(
            llm_text, "HELLO",
            "load() must route through the migrator pipeline"
        );

        let custom = messages
            .iter()
            .find_map(|m| m.downcast_ref::<TestCustomMsg>().ok())
            .expect("custom wrapper must pass through load() unchanged");
        assert_eq!(custom.data, "custom-payload");
    }

    /// Without a registered migrator, a legacy-version file must still fail
    /// `load()` with the same error that `load_entries()` emits — contract
    /// consistency between the two APIs.
    #[test]
    fn load_rejects_legacy_version_without_migrator_like_load_entries() {
        let dir = tempfile::tempdir().unwrap();
        let id = "no-migrator";
        let path = session_path(dir.path(), id);
        let now = now_utc();
        let meta = SessionMeta {
            id: id.to_string(),
            title: "legacy".to_string(),
            created_at: now,
            updated_at: now,
            version: 0,
            sequence: 0,
        };
        let contents = format!("{}\n", serde_json::to_string(&meta).unwrap());
        std::fs::write(&path, contents).unwrap();

        let store = JsonlSessionStore::new(dir.path().to_path_buf()).unwrap();
        let load_err = store.load(id, None).unwrap_err();
        let entries_err = store.load_entries(id).unwrap_err();
        assert_eq!(load_err.kind(), io::ErrorKind::InvalidData);
        assert_eq!(entries_err.kind(), io::ErrorKind::InvalidData);
        assert!(load_err.to_string().contains("no migrator found"));
        assert!(entries_err.to_string().contains("no migrator found"));
    }

    #[test]
    fn load_interrupt_ignores_orphan_file_without_session() {
        let dir = tempfile::tempdir().unwrap();
        let store = JsonlSessionStore::new(dir.path().to_path_buf()).unwrap();

        let orphan_path = interrupt_path(dir.path(), "orphan");
        std::fs::write(
            &orphan_path,
            serde_json::to_string(&InterruptState {
                interrupted_at: 2,
                pending_tool_calls: vec![],
                context_snapshot: vec![],
                system_prompt: "system".to_string(),
                model: swink_agent::ModelSpec::new("openai", "gpt-4o"),
            })
            .unwrap(),
        )
        .unwrap();

        assert!(store.load_interrupt("orphan").unwrap().is_none());
    }
}