aacode_rs/session/

mod.rs

// Copyright (c) 2026 xiefujin <490021684@qq.com>
// Licensed under GPL-3.0, see LICENSE file for full license terms.

//! Session persistence: `.aacode/sessions/<id>.json` + `sessions_index.json`.
//!
//! Ported from Python `utils/session_manager.py`. Uses atomic writes
//! (temp file + rename). Messages carry structured tool_calls / tool_call_id /
//! reasoning_content, matching the LLM `ChatMessage` shape.

use crate::error::{AacodeError, Result};
use crate::llm::types::{ChatMessage, ToolCall};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::path::{Path, PathBuf};
use std::time::Instant;

/// A persisted message (superset of ChatMessage with a timestamp + tokens).
///
/// NOTE: render segments (`thinking`/`thought`/`action`/`observation`) are **not**
/// persisted. They are fully derivable from the message itself plus the
/// immediately-following `role:"tool"` messages (`tool_calls` → action, tool
/// message → observation), so storing them duplicated content and diverged from
/// the Python `SessionMessage`. Hosts rebuild the collapse UI from the message
/// array order (and the live `seg_content` event stream).
///
/// Legacy session files that still contain a `segments` field keep parsing —
/// serde ignores unknown fields — but the field is never written again.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionMessage {
    pub role: String,
    #[serde(default)]
    pub content: String,
    #[serde(default)]
    pub timestamp: String,
    #[serde(default)]
    pub tokens: usize,
    #[serde(skip_serializing_if = "Option::is_none", default)]
    pub tool_calls: Option<Vec<ToolCall>>,
    #[serde(skip_serializing_if = "Option::is_none", default)]
    pub tool_call_id: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none", default)]
    pub reasoning_content: Option<String>,
}

impl SessionMessage {
    pub fn from_chat(m: &ChatMessage) -> Self {
        // Token estimate must match the runtime estimator
        // (`agent::compact::estimate_messages_tokens`) and the Python
        // `message_utils.estimate_tokens`: count content, reasoning_content,
        // tool_calls (serialized) AND the tool_call_id. Omitting tool_calls
        // made tool-call-only assistant messages report `tokens = 0`.
        let mut tokens = estimate_tokens(&m.content);
        if let Some(rc) = &m.reasoning_content {
            tokens += estimate_tokens(rc);
        }
        if let Some(tcs) = &m.tool_calls {
            if let Ok(json) = serde_json::to_string(tcs) {
                tokens += estimate_tokens(&json);
            }
        }
        if let Some(id) = &m.tool_call_id {
            tokens += estimate_tokens(id);
        }
        SessionMessage {
            role: m.role.clone(),
            content: m.content.clone(),
            timestamp: now_iso(),
            tokens,
            tool_calls: m.tool_calls.clone(),
            tool_call_id: m.tool_call_id.clone(),
            reasoning_content: m.reasoning_content.clone(),
        }
    }

    pub fn to_chat(&self) -> ChatMessage {
        ChatMessage {
            role: self.role.clone(),
            content: self.content.clone(),
            tool_calls: self.tool_calls.clone(),
            tool_call_id: self.tool_call_id.clone(),
            reasoning_content: self.reasoning_content.clone(),
        }
    }
}

/// One entry in the sessions index.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionSummary {
    pub session_id: String,
    pub created_at: String,
    pub last_activity: String,
    #[serde(default)]
    pub total_messages: usize,
    #[serde(default)]
    pub total_tokens: usize,
    #[serde(default)]
    pub title: String,
    #[serde(default = "default_status")]
    pub status: String,
}

fn default_status() -> String {
    "active".to_string()
}

/// The on-disk session file shape.
///
/// `schema_version` is a forward-compatibility guard: bump it on any breaking
/// change to this shape. iOS/Android parse the same file, so a shared
/// conformance fixture (`tests/fixtures/session_v1.json`) locks the format
/// across all three codebases.
#[derive(Debug, Clone, Serialize, Deserialize)]
struct SessionFile {
    #[serde(default = "default_schema_version")]
    schema_version: u32,
    session_id: String,
    #[serde(default)]
    created_at: String,
    messages: Vec<SessionMessage>,
    /// Retrieval cutoff for the model view. The stored messages stay complete;
    /// this only records where the model view starts. `None` = no cutoff.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    context_checkpoint: Option<Value>,
}

fn default_schema_version() -> u32 {
    SCHEMA_VERSION
}

/// Current on-disk session schema version (see SESSION_FFI.md §6). Bump on any
/// breaking change to the session file shape.
pub const SCHEMA_VERSION: u32 = 1;

/// Validate that a session id is a safe simple identifier (no path separators
/// or `..`), so it can never escape the sessions directory.
pub fn valid_session_id(sid: &str) -> bool {
    !sid.is_empty()
        && !sid.contains('/')
        && !sid.contains('\\')
        && !sid.contains("..")
        && !sid.contains('\0')
}

/// Rough token estimate (4 chars ≈ 1 token). tiktoken is not ported.
pub fn estimate_tokens(text: &str) -> usize {
    text.chars().count() / 4
}

/// Current time as a compact ISO-ish string (no external chrono dependency).
pub fn now_iso() -> String {
    use std::time::{SystemTime, UNIX_EPOCH};
    let secs = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|d| d.as_secs())
        .unwrap_or(0);
    // Store epoch seconds as string; readable enough and monotonic for sorting.
    secs.to_string()
}

/// Current time as an epoch-seconds string with millisecond precision
/// (e.g. `"1700000000.123"`). Used for per-segment `created_at` so segments
/// within the same message stay stably ordered and hosts can derive stable ids.
pub fn now_iso_ms() -> String {
    use std::time::{SystemTime, UNIX_EPOCH};
    let d = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default();
    format!("{}.{:03}", d.as_secs(), d.subsec_millis())
}

/// Atomic write: write to a temp file then rename over the target.
fn atomic_write(path: &Path, contents: &str) -> Result<()> {
    let parent = path
        .parent()
        .ok_or_else(|| AacodeError::Io("no parent".into()))?;
    std::fs::create_dir_all(parent)?;
    let tmp = parent.join(format!(".tmp_{}", uuid::Uuid::new_v4().simple()));
    std::fs::write(&tmp, contents)?;
    std::fs::rename(&tmp, path)?;
    Ok(())
}

pub struct SessionManager {
    sessions_dir: PathBuf,
    pub current_session_id: Option<String>,
    pub messages: Vec<SessionMessage>,
    /// Retrieval cutoff for the model view (optional; persisted per session).
    pub context_checkpoint: Option<Value>,
    /// Flush cooldown: how many messages to batch before forcing a write.
    flush_batch_size: usize,
    /// Number of messages added since the last flush.
    dirty_count: usize,
    /// Last flush time (for interval-based debounce).
    last_flush: Option<Instant>,
    /// Minimum interval between flushes in seconds.
    flush_interval_secs: f64,
}

impl SessionManager {
    pub fn new(project_path: &Path) -> Self {
        let sessions_dir = project_path.join(".aacode").join("sessions");
        // (c) 2026 xiefujin <490021684@qq.com> — GPL-3.0
        let _ = std::fs::create_dir_all(&sessions_dir);
        SessionManager {
            sessions_dir,
            current_session_id: None,
            messages: Vec::new(),
            context_checkpoint: None,
            flush_batch_size: 10,
            dirty_count: 0,
            last_flush: None,
            flush_interval_secs: 0.5,
        }
    }

    fn index_path(&self) -> PathBuf {
        self.sessions_dir.join("sessions_index.json")
    }

    /// Project root this session store belongs to (`<project>/.aacode/sessions`).
    pub fn project_path(&self) -> PathBuf {
        self.sessions_dir
            .parent()
            .and_then(|p| p.parent())
            .map(|p| p.to_path_buf())
            .unwrap_or_default()
    }

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

    fn load_index(&self) -> std::collections::BTreeMap<String, SessionSummary> {
        let path = self.index_path();
        if let Ok(s) = std::fs::read_to_string(&path) {
            if let Ok(m) = serde_json::from_str(&s) {
                return m;
            }
        }
        std::collections::BTreeMap::new()
    }

    fn save_index(&self, index: &std::collections::BTreeMap<String, SessionSummary>) -> Result<()> {
        // Merge with disk to avoid clobbering concurrent writers.
        let mut merged = self.load_index();
        for (k, v) in index {
            merged.insert(k.clone(), v.clone());
        }
        atomic_write(&self.index_path(), &serde_json::to_string(&merged)?)
    }

    /// Create a new session and make it current. Returns the session id.
    pub fn create_session(&mut self, task: &str, title: Option<&str>) -> Result<String> {
        let id = format!(
            "session_{}_{}",
            now_iso(),
            &uuid::Uuid::new_v4().simple().to_string()[..8]
        );
        let title = title
            .map(|t| t.to_string())
            .unwrap_or_else(|| truncate_title(task));
        let now = now_iso();
        let summary = SessionSummary {
            session_id: id.clone(),
            created_at: now.clone(),
            last_activity: now,
            total_messages: 0,
            total_tokens: 0,
            title,
            status: "active".to_string(),
        };
        self.current_session_id = Some(id.clone());
        self.messages.clear();
        self.context_checkpoint = None;
        if !task.trim().is_empty() {
            self.messages.push(SessionMessage {
                role: "user".to_string(),
                content: task.to_string(),
                timestamp: now_iso(),
                tokens: estimate_tokens(task),
                tool_calls: None,
                tool_call_id: None,
                reasoning_content: None,
            });
        }
        let mut index = self.load_index();
        let mut s = summary;
        s.total_messages = self.messages.len();
        s.total_tokens = self.total_tokens();
        index.insert(id.clone(), s);
        self.save_index(&index)?;
        self.save_current()?;
        self.dirty_count = 0;
        self.last_flush = Some(Instant::now());
        Ok(id)
    }

    /// Append a message to the current session. Writes are batched via
    /// a cooldown (every N messages or every `flush_interval_secs`) to avoid
    /// full-file serialize + disk write on every single message.
    pub fn add_message(&mut self, msg: SessionMessage) -> Result<()> {
        self.messages.push(msg);
        self.dirty_count += 1;
        self.maybe_flush()
    }

    /// Append several messages (batch) then persist once.
    pub fn add_messages(&mut self, msgs: impl IntoIterator<Item = SessionMessage>) -> Result<()> {
        for m in msgs {
            self.messages.push(m);
            self.dirty_count += 1;
        }
        self.maybe_flush()
    }

    /// Force a write of the current session to disk (messages + index).
    pub fn flush(&mut self) -> Result<()> {
        if self.dirty_count == 0 {
            return Ok(());
        }
        self.save_current()?;
        self.touch_index()?;
        self.dirty_count = 0;
        self.last_flush = Some(Instant::now());
        Ok(())
    }

    fn maybe_flush(&mut self) -> Result<()> {
        if self.dirty_count >= self.flush_batch_size {
            return self.flush();
        }
        match self.last_flush {
            None => self.flush(),
            Some(t) if t.elapsed().as_secs_f64() >= self.flush_interval_secs => self.flush(),
            _ => Ok(()),
        }
    }

    fn total_tokens(&self) -> usize {
        self.messages.iter().map(|m| m.tokens).sum()
    }

    fn touch_index(&self) -> Result<()> {
        if let Some(id) = &self.current_session_id {
            let mut index = self.load_index();
            if let Some(s) = index.get_mut(id) {
                s.last_activity = now_iso();
                s.total_messages = self.messages.len();
                s.total_tokens = self.total_tokens();
                if s.title.trim().is_empty() {
                    if let Some(first_user) = self.messages.iter().find(|m| m.role == "user") {
                        s.title = truncate_title(&first_user.content);
                    }
                }
            }
            self.save_index(&index)?;
        }
        Ok(())
    }

    fn save_current(&self) -> Result<()> {
        if let Some(id) = &self.current_session_id {
            let created_at = self
                .load_index()
                .get(id)
                .map(|s| s.created_at.clone())
                .unwrap_or_else(now_iso);
            let file = SessionFile {
                schema_version: default_schema_version(),
                session_id: id.clone(),
                created_at,
                messages: self.messages.clone(),
                context_checkpoint: self.context_checkpoint.clone(),
            };
            atomic_write(&self.session_path(id), &serde_json::to_string(&file)?)?;
        }
        Ok(())
    }

    /// Load an existing session as current.
    pub fn switch_session(&mut self, id: &str) -> Result<bool> {
        let path = self.session_path(id);
        if !path.exists() {
            return Ok(false);
        }
        let s = std::fs::read_to_string(&path)?;
        let file: SessionFile = serde_json::from_str(&s)?;
        self.current_session_id = Some(id.to_string());
        self.messages = file.messages;
        self.context_checkpoint = file.context_checkpoint;
        Ok(true)
    }

    /// Persisted retrieval checkpoint (model-view cutoff only).
    pub fn context_checkpoint(&self) -> Option<Value> {
        self.context_checkpoint.clone()
    }

    /// Persist a retrieval checkpoint. Messages themselves are never modified.
    pub fn set_context_checkpoint(&mut self, checkpoint: Option<Value>) -> Result<()> {
        self.context_checkpoint = checkpoint;
        self.save_current()
    }

    /// Export one session as JSONL (user-triggered, read-only).
    ///
    /// Format: a `header` line, one `message` line per persisted message, and a
    /// `footer` line. Stable IDs/fields; no execution replay. Tool output text
    /// is exported verbatim — callers are responsible for not exporting
    /// sessions that contain secrets.
    pub fn export_jsonl(&self, id: &str) -> Result<String> {
        let path = self.session_path(id);
        let text = std::fs::read_to_string(&path)?;
        let file: SessionFile = serde_json::from_str(&text)?;
        let mut out = String::new();
        out.push_str(
            &serde_json::json!({
                "type": "header",
                "format": "aacode.session-export",
                "formatVersion": 1,
                "session_id": file.session_id,
                "created_at": file.created_at,
            })
            .to_string(),
        );
        out.push('\n');
        for (index, message) in file.messages.iter().enumerate() {
            let mut line = serde_json::to_value(message)?;
            line["type"] = Value::String("message".into());
            line["index"] = Value::from(index);
            out.push_str(&line.to_string());
            out.push('\n');
        }
        out.push_str(
            &serde_json::json!({
                "type": "footer",
                "messages": file.messages.len(),
            })
            .to_string(),
        );
        out.push('\n');
        Ok(out)
    }

    /// Fork a session from a message prefix into a new session.
    ///
    /// The prefix is copied by value; the source session is untouched and the
    /// new session is **not** made current. A boundary that would split an
    /// `assistant(tool_calls)` from its `tool` results is moved back to before
    /// the assistant so tool pairing is always intact. The retrieval checkpoint
    /// is reset (it is recomputed safely for the new session).
    pub fn fork_session(&mut self, source_id: &str, upto: Option<usize>) -> Result<String> {
        let path = self.session_path(source_id);
        let text = std::fs::read_to_string(&path)?;
        let source: SessionFile = serde_json::from_str(&text)?;
        let mut end = upto.unwrap_or(source.messages.len());
        end = end.min(source.messages.len());
        end = safe_prefix_end(&source.messages, end);
        if end == 0 {
            return Err(AacodeError::Tool(
                "cannot fork before the first message".into(),
            ));
        }
        if end > 8192 {
            return Err(AacodeError::Tool(
                "fork prefix exceeds 8192 messages".into(),
            ));
        }
        let id = format!(
            "session_{}_{}",
            now_iso(),
            &uuid::Uuid::new_v4().simple().to_string()[..8]
        );
        let prefix: Vec<SessionMessage> = source.messages[..end].to_vec();
        let title = prefix
            .iter()
            .find(|m| m.role == "user" && !m.content.is_empty())
            .map(|m| format!("Fork: {}", truncate_title(&m.content)))
            .unwrap_or_else(|| format!("Fork of {source_id}"));
        let now = now_iso();
        let file = SessionFile {
            schema_version: SCHEMA_VERSION,
            session_id: id.clone(),
            created_at: now.clone(),
            messages: prefix,
            context_checkpoint: None,
        };
        atomic_write(&self.session_path(&id), &serde_json::to_string(&file)?)?;

        let mut index = self.load_index();
        index.insert(
            id.clone(),
            SessionSummary {
                session_id: id.clone(),
                created_at: now.clone(),
                last_activity: now,
                total_messages: file.messages.len(),
                total_tokens: file.messages.iter().map(|m| m.tokens).sum(),
                title,
                status: "active".to_string(),
            },
        );
        self.save_index(&index)?;
        Ok(id)
    }

    /// List all sessions, newest activity first.
    /// Falls back to scanning disk files when the index is empty.
    pub fn list_sessions(&self) -> Vec<SessionSummary> {
        let mut index = self.load_index();
        if index.is_empty() {
            if let Ok(entries) = std::fs::read_dir(&self.sessions_dir) {
                for entry in entries.flatten() {
                    let path = entry.path();
                    if path.extension().and_then(|e| e.to_str()) != Some("json") {
                        continue;
                    }
                    let stem = path.file_stem().and_then(|s| s.to_str()).unwrap_or("");
                    if stem == "sessions_index" {
                        continue;
                    }
                    if let Ok(s) = std::fs::read_to_string(&path) {
                        if let Ok(file) = serde_json::from_str::<SessionFile>(&s) {
                            let title = file
                                .messages
                                .first()
                                .map(|m| truncate_title(&m.content))
                                .unwrap_or_default();
                            let summary = SessionSummary {
                                session_id: file.session_id.clone(),
                                created_at: file.created_at.clone(),
                                last_activity: file
                                    .messages
                                    .last()
                                    .map(|m| m.timestamp.clone())
                                    .unwrap_or_default(),
                                total_messages: file.messages.len(),
                                total_tokens: file.messages.iter().map(|m| m.tokens).sum(),
                                title,
                                status: "active".to_string(),
                            };
                            index.entry(file.session_id).or_insert(summary);
                        }
                    }
                }
            }
            // Self-heal: write rebuilt index back to disk.
            if !index.is_empty() {
                if let Ok(contents) = serde_json::to_string(&index) {
                    let _ = atomic_write(&self.index_path(), &contents);
                }
            }
        }
        let mut v: Vec<SessionSummary> = index.into_values().collect();
        v.sort_by(|a, b| b.last_activity.cmp(&a.last_activity));
        v
    }

    /// Delete a session file + index entry.
    pub fn delete_session(&mut self, id: &str) -> Result<bool> {
        let mut index = self.load_index();
        if index.remove(id).is_none() {
            return Ok(false);
        }
        let _ = std::fs::remove_file(self.session_path(id));
        // rewrite index without the removed key (save_index merges, so write raw)
        atomic_write(&self.index_path(), &serde_json::to_string(&index)?)?;
        if self.current_session_id.as_deref() == Some(id) {
            self.current_session_id = None;
            self.messages.clear();
        }
        Ok(true)
    }

    /// Read messages of a session file directly (without switching).
    pub fn read_session_messages(&self, id: &str) -> Vec<SessionMessage> {
        let path = self.session_path(id);
        if let Ok(s) = std::fs::read_to_string(&path) {
            if let Ok(file) = serde_json::from_str::<SessionFile>(&s) {
                return file.messages;
            }
        }
        Vec::new()
    }

    /// Ensure a session with the given id exists (idempotent create-or-touch).
    /// If it exists, its `last_activity` is bumped (and a non-empty `title`
    /// updates the stored title); otherwise an empty session file + index entry
    /// are created.
    pub fn ensure_session(&mut self, sid: &str, title: &str) -> Result<()> {
        let mut index = self.load_index();
        let now = now_iso();
        if let Some(s) = index.get_mut(sid) {
            s.last_activity = now;
            if !title.trim().is_empty() {
                s.title = title.to_string();
            }
            self.save_index(&index)?;
            return Ok(());
        }
        let summary = SessionSummary {
            session_id: sid.to_string(),
            created_at: now.clone(),
            last_activity: now.clone(),
            total_messages: 0,
            total_tokens: 0,
            title: title.to_string(),
            status: "active".to_string(),
        };
        index.insert(sid.to_string(), summary);
        self.save_index(&index)?;
        let file = SessionFile {
            schema_version: SCHEMA_VERSION,
            session_id: sid.to_string(),
            created_at: now,
            messages: Vec::new(),
            context_checkpoint: None,
        };
        atomic_write(&self.session_path(sid), &serde_json::to_string(&file)?)
    }

    /// Set an explicit title on an existing session (no-op if missing).
    pub fn rename_session(&mut self, sid: &str, title: &str) -> Result<()> {
        let mut index = self.load_index();
        if let Some(s) = index.get_mut(sid) {
            s.title = title.to_string();
            self.save_index(&index)?;
        }
        Ok(())
    }

    /// Bump `last_activity` of an existing session (no-op if missing).
    pub fn touch_session(&mut self, sid: &str) -> Result<()> {
        let mut index = self.load_index();
        if let Some(s) = index.get_mut(sid) {
            s.last_activity = now_iso();
            self.save_index(&index)?;
        }
        Ok(())
    }

    /// Append messages to a session file (creating the session if missing),
    /// then refresh the index entry (counts + last_activity + title fallback).
    pub fn append_session_messages(&mut self, sid: &str, msgs: Vec<SessionMessage>) -> Result<()> {
        let path = self.session_path(sid);
        let mut existing: Vec<SessionMessage> = if path.exists() {
            std::fs::read_to_string(&path)
                .ok()
                .and_then(|s| serde_json::from_str::<SessionFile>(&s).ok())
                .map(|f| f.messages)
                .unwrap_or_default()
        } else {
            Vec::new()
        };
        existing.extend(msgs);

        let total_messages = existing.len();
        let total_tokens: usize = existing.iter().map(|m| m.tokens).sum();
        let fallback_title = existing
            .iter()
            .find(|m| m.role == "user")
            .map(|m| truncate_title(&m.content))
            .unwrap_or_default();

        let created_at = self
            .load_index()
            .get(sid)
            .map(|s| s.created_at.clone())
            .unwrap_or_else(now_iso);

        let file = SessionFile {
            schema_version: SCHEMA_VERSION,
            session_id: sid.to_string(),
            created_at,
            messages: existing,
            context_checkpoint: None,
        };
        atomic_write(&path, &serde_json::to_string(&file)?)?;

        let mut index = self.load_index();
        if let Some(s) = index.get_mut(sid) {
            s.total_messages = total_messages;
            s.total_tokens = total_tokens;
            s.last_activity = now_iso();
            if s.title.trim().is_empty() && !fallback_title.is_empty() {
                s.title = fallback_title.clone();
            }
        } else {
            index.insert(
                sid.to_string(),
                SessionSummary {
                    session_id: sid.to_string(),
                    created_at: file.created_at.clone(),
                    last_activity: now_iso(),
                    total_messages,
                    total_tokens,
                    title: fallback_title,
                    status: "active".to_string(),
                },
            );
        }
        self.save_index(&index)
    }

    /// Get the current session's messages as ChatMessages (for LLM history).
    /// System messages are included so that dynamically-appended context
    /// (project analysis, init instructions, compaction summaries, etc.) are
    /// preserved in the prefix for provider KV-cache stability across calls.
    pub fn history_chat(&self) -> Vec<ChatMessage> {
        self.messages.iter().map(|m| m.to_chat()).collect()
    }

    /// Whether the last message is a tool result (interrupted mid-execution).
    pub fn ended_mid_tool(&self) -> bool {
        self.messages
            .last()
            .map(|m| m.role == "tool")
            .unwrap_or(false)
    }
}

fn truncate_title(task: &str) -> String {
    let t: String = task.chars().take(50).collect();
    if task.chars().count() > 50 {
        format!("{t}...")
    } else {
        t
    }
}

/// Move a fork boundary back so an `assistant(tool_calls)` block is never cut
/// from its `tool` results (protocol integrity).
fn safe_prefix_end(messages: &[SessionMessage], mut end: usize) -> usize {
    end = end.min(messages.len());
    loop {
        let last = messages[..end].iter().enumerate().rev().find(|(_, m)| {
            m.role == "assistant"
                && m.tool_calls
                    .as_ref()
                    .map(|c| !c.is_empty())
                    .unwrap_or(false)
        });
        let Some((idx, message)) = last else {
            break;
        };
        let ids: Vec<&str> = message
            .tool_calls
            .as_ref()
            .map(|calls| calls.iter().map(|c| c.id.as_str()).collect())
            .unwrap_or_default();
        let all_present = ids.iter().all(|id| {
            messages[..end]
                .iter()
                .any(|m| m.role == "tool" && m.tool_call_id.as_deref() == Some(*id))
        });
        if all_present {
            break;
        }
        end = idx;
    }
    end
}

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

    fn tmp_project() -> PathBuf {
        let d = std::env::temp_dir().join(format!(
            "aacode_sess_{}_{}",
            std::process::id(),
            uuid::Uuid::new_v4().simple()
        ));
        std::fs::create_dir_all(&d).unwrap();
        d
    }

    #[test]
    fn create_and_reload() {
        let proj = tmp_project();
        let mut sm = SessionManager::new(&proj);
        let id = sm.create_session("do something big", None).unwrap();
        assert!(id.starts_with("session_"));
        assert_eq!(sm.messages.len(), 1);

        // new manager, switch to it
        let mut sm2 = SessionManager::new(&proj);
        assert!(sm2.switch_session(&id).unwrap());
        assert_eq!(sm2.messages[0].content, "do something big");
    }

    #[test]
    fn append_and_persist() {
        let proj = tmp_project();
        let mut sm = SessionManager::new(&proj);
        let id = sm.create_session("task", None).unwrap();
        sm.add_message(SessionMessage::from_chat(&ChatMessage::assistant("hi")))
            .unwrap();
        sm.flush().unwrap();
        let read = sm.read_session_messages(&id);
        assert_eq!(read.len(), 2);
        assert_eq!(read[1].role, "assistant");
    }

    #[test]
    fn tool_calls_persist() {
        let proj = tmp_project();
        let mut sm = SessionManager::new(&proj);
        sm.create_session("t", None).unwrap();
        let msg = ChatMessage::assistant_with_tools(
            "",
            vec![ToolCall {
                id: "c1".into(),
                name: "run_shell".into(),
                arguments: "{}".into(),
            }],
        );
        sm.add_message(SessionMessage::from_chat(&msg)).unwrap();
        sm.add_message(SessionMessage::from_chat(&ChatMessage::tool_result(
            "c1", "ok",
        )))
        .unwrap();
        assert!(sm.ended_mid_tool());
        let hist = sm.history_chat();
        assert!(hist.iter().any(|m| m.tool_calls.is_some()));
    }

    #[test]
    fn list_and_delete() {
        let proj = tmp_project();
        let mut sm = SessionManager::new(&proj);
        let id1 = sm.create_session("one", None).unwrap();
        let _id2 = sm.create_session("two", None).unwrap();
        assert!(sm.list_sessions().len() >= 2);
        assert!(sm.delete_session(&id1).unwrap());
        assert!(!sm.list_sessions().iter().any(|s| s.session_id == id1));
    }

    #[test]
    fn history_includes_system() {
        let proj = tmp_project();
        let mut sm = SessionManager::new(&proj);
        sm.create_session("t", None).unwrap();
        sm.add_message(SessionMessage::from_chat(&ChatMessage::system("sys")))
            .unwrap();
        // After cache-stability fix: system messages are included in history
        // so persisted dynamic context is preserved for prefix matching.
        assert!(sm.history_chat().iter().any(|m| m.role == "system"));
    }

    #[test]
    fn scheduled_task_saves_to_existing_session() {
        let proj = tmp_project();
        let mut sm = SessionManager::new(&proj);

        let id = sm.create_session("previous task", None).unwrap();
        sm.add_message(SessionMessage::from_chat(&ChatMessage::assistant("done")))
            .unwrap();
        sm.flush().unwrap();

        let mut sm2 = SessionManager::new(&proj);
        assert!(sm2.switch_session(&id).unwrap());
        assert!(sm2.current_session_id.is_some());

        let scheduled =
            "[Scheduled: daily-check]\nThis is an automated scheduled run.\n\ncheck git status";

        let already = sm2
            .messages
            .last()
            .map(|m| m.role == "user" && m.content == scheduled)
            .unwrap_or(false);
        assert!(!already);

        sm2.add_message(SessionMessage::from_chat(&ChatMessage::user(scheduled)))
            .unwrap();
        sm2.flush().unwrap();

        let msgs = sm2.read_session_messages(&id);
        assert_eq!(msgs.len(), 3);
        assert_eq!(msgs[0].content, "previous task");
        assert_eq!(msgs[1].role, "assistant");
        assert_eq!(msgs[2].role, "user");
        assert!(msgs[2].content.starts_with("[Scheduled: daily-check]"));
    }

    #[test]
    fn dedup_skips_duplicate_user_message() {
        let proj = tmp_project();
        let mut sm = SessionManager::new(&proj);

        let task = "[Scheduled: daily]\ncheck";
        let id = sm.create_session(task, None).unwrap();

        let mut sm2 = SessionManager::new(&proj);
        sm2.switch_session(&id).unwrap();
        let already = sm2
            .messages
            .last()
            .map(|m| m.role == "user" && m.content == task)
            .unwrap_or(false);
        assert!(already);

        let msgs = sm2.read_session_messages(&id);
        assert_eq!(msgs.len(), 1);
    }

    #[test]
    fn reasoning_content_counted_in_tokens() {
        let msg = ChatMessage {
            role: "assistant".into(),
            content: "hello world".into(),
            tool_calls: None,
            tool_call_id: None,
            reasoning_content: Some("I need to think about this carefully".into()),
        };
        let sm = SessionMessage::from_chat(&msg);
        // Should count both content and reasoning_content tokens.
        let content_only = estimate_tokens("hello world");
        let with_reasoning = estimate_tokens("hello world")
            + estimate_tokens("I need to think about this carefully");
        assert!(sm.tokens > content_only);
        assert_eq!(sm.tokens, with_reasoning);
    }

    #[test]
    fn no_reasoning_content_tokens_unchanged() {
        let msg = ChatMessage {
            role: "assistant".into(),
            content: "hello world".into(),
            tool_calls: None,
            tool_call_id: None,
            reasoning_content: None,
        };
        let sm = SessionMessage::from_chat(&msg);
        assert_eq!(sm.tokens, estimate_tokens("hello world"));
    }

    #[test]
    fn tool_message_tokens_include_content_and_tool_call_id() {
        // Tool messages have no reasoning_content; the estimate counts the
        // content AND the tool_call_id (parity with the Python estimator).
        let msg = ChatMessage {
            role: "tool".into(),
            content: "result data".into(),
            tool_calls: None,
            tool_call_id: Some("call_1".into()),
            reasoning_content: None,
        };
        let sm = SessionMessage::from_chat(&msg);
        assert_eq!(
            sm.tokens,
            estimate_tokens("result data") + estimate_tokens("call_1")
        );
    }

    #[test]
    fn compact_json_roundtrips() {
        let proj = tmp_project();
        let mut sm = SessionManager::new(&proj);
        sm.create_session("compact test", None).unwrap();
        sm.add_message(SessionMessage::from_chat(&ChatMessage::assistant(
            "{\"key\": \"value\"}",
        )))
        .unwrap();
        sm.flush().unwrap();

        let id = sm.current_session_id.unwrap();
        let mut sm2 = SessionManager::new(&proj);
        assert!(sm2.switch_session(&id).unwrap());
        assert_eq!(sm2.messages.len(), 2);
        assert_eq!(sm2.messages[1].content, "{\"key\": \"value\"}");

        let path = sm2.session_path(&id);
        let raw = std::fs::read_to_string(&path).unwrap();
        assert!(
            !raw.contains("\n  "),
            "compact JSON should not have indentation"
        );
        let file: SessionFile = serde_json::from_str(&raw).unwrap();
        assert_eq!(file.messages.len(), 2);
    }

    #[test]
    fn persisted_dynamic_system_message_in_history() {
        let proj = tmp_project();
        let mut sm = SessionManager::new(&proj);
        sm.create_session("task", None).unwrap();

        let analysis = SessionMessage::from_chat(&ChatMessage::system(
            "## Project Analysis\n2 .rs files, 1 .py file",
        ));
        sm.add_message(analysis).unwrap();
        sm.flush().unwrap();

        let history = sm.history_chat();
        assert!(history
            .iter()
            .any(|m| m.role == "system" && m.content.contains("Project Analysis")));

        // Reload from disk and verify system message is still in history.
        let id = sm.current_session_id.unwrap();
        let mut sm2 = SessionManager::new(&proj);
        sm2.switch_session(&id).unwrap();
        let history2 = sm2.history_chat();
        assert!(history2
            .iter()
            .any(|m| m.role == "system" && m.content.contains("Project Analysis")));
    }

    #[test]
    fn flush_cooldown_batches_messages() {
        let proj = tmp_project();
        let mut sm = SessionManager::new(&proj);
        let id = sm.create_session("task", None).unwrap();

        // Add messages below the batch threshold — should not write to disk.
        for i in 0..5 {
            sm.add_message(SessionMessage::from_chat(&ChatMessage::assistant(
                &format!("msg{}", i),
            )))
            .unwrap();
        }

        // Force a flush and verify all messages are persisted.
        sm.flush().unwrap();

        let msgs = sm.read_session_messages(&id);
        assert_eq!(msgs.len(), 6); // 1 from create + 5 added
        assert_eq!(msgs[5].content, "msg4");
    }

    #[test]
    fn flush_idempotent_when_clean() {
        let proj = tmp_project();
        let mut sm = SessionManager::new(&proj);
        sm.create_session("task", None).unwrap();
        // Calling flush on a clean session should be a no-op.
        assert!(sm.flush().is_ok());
        assert!(sm.dirty_count == 0);
    }

    // ── schema conformance (locks the on-disk format across hosts) ──

    /// The canonical session-file fixture is shared with the iOS/Android hosts
    /// (each has its own copy under their test trees). Any breaking change to
    /// the on-disk shape must update this fixture and bump `schema_version`.
    #[test]
    fn schema_conformance_fixture_roundtrips() {
        const FIXTURE: &str = include_str!("../../tests/fixtures/session_v1.json");
        let file: SessionFile = serde_json::from_str(FIXTURE).expect("fixture must parse");

        assert_eq!(file.schema_version, 1, "schema_version must be pinned at 1");
        assert_eq!(file.session_id, "session_1700000000_abc12345");
        assert_eq!(file.created_at, "1700000000");
        assert_eq!(file.messages.len(), 3);

        // user message: no tool fields.
        assert_eq!(file.messages[0].role, "user");
        assert_eq!(file.messages[0].content, "Create a hello.py file");
        assert!(file.messages[0].tool_calls.is_none());
        assert!(file.messages[0].tool_call_id.is_none());
        assert!(file.messages[0].reasoning_content.is_none());

        // assistant message: tool_calls + reasoning_content.
        assert_eq!(file.messages[1].role, "assistant");
        let tcs = file.messages[1]
            .tool_calls
            .as_ref()
            .expect("assistant tool_calls present");
        assert_eq!(tcs.len(), 1);
        assert_eq!(tcs[0].id, "call_1");
        assert_eq!(tcs[0].name, "run_shell");
        assert_eq!(tcs[0].arguments, r#"{"command":"echo hello"}"#);
        assert_eq!(
            file.messages[1].reasoning_content.as_deref(),
            Some("I should run a shell command")
        );

        // tool message: tool_call_id.
        assert_eq!(file.messages[2].role, "tool");
        assert_eq!(file.messages[2].content, "hello");
        assert_eq!(file.messages[2].tool_call_id.as_deref(), Some("call_1"));

        // Round-trip: re-serialize then re-parse must be stable.
        let reserialized = serde_json::to_string(&file).unwrap();
        let reparsed: SessionFile = serde_json::from_str(&reserialized).unwrap();
        assert_eq!(reparsed.schema_version, 1);
        assert_eq!(reparsed.messages.len(), 3);
        assert_eq!(
            reparsed.messages[1].tool_calls.as_ref().unwrap()[0].name,
            "run_shell"
        );
    }

    /// Legacy `session_v2.json` still carries a `segments` field from the old
    /// schema; it must keep parsing (serde ignores unknown fields) and the field
    /// must never be re-serialized.
    #[test]
    fn legacy_segments_fixture_still_parses_and_is_dropped_on_write() {
        const FIXTURE: &str = include_str!("../../tests/fixtures/session_v2.json");
        let file: SessionFile = serde_json::from_str(FIXTURE).expect("legacy fixture must parse");
        assert_eq!(file.schema_version, 1);
        assert_eq!(file.messages.len(), 3);
        assert_eq!(
            file.messages[1].tool_calls.as_ref().unwrap()[0].name,
            "run_shell"
        );
        // `segments` is no longer a field → dropped on re-serialization.
        let raw = serde_json::to_string(&file).unwrap();
        assert!(
            !raw.contains("segments"),
            "segments must not be re-persisted: {raw}"
        );
    }

    #[test]
    fn persisted_messages_never_contain_segments() {
        let proj = tmp_project();
        let mut sm = SessionManager::new(&proj);
        let id = sm.create_session("t", None).unwrap();
        sm.add_message(SessionMessage::from_chat(
            &ChatMessage::assistant_with_tools(
                "hi",
                vec![ToolCall {
                    id: "c1".into(),
                    name: "run_shell".into(),
                    arguments: "{}".into(),
                }],
            ),
        ))
        .unwrap();
        sm.add_message(SessionMessage::from_chat(&ChatMessage::tool_result(
            "c1", "out",
        )))
        .unwrap();
        sm.flush().unwrap();

        let raw = std::fs::read_to_string(sm.session_path(&id)).unwrap();
        assert!(!raw.contains("segments"), "no segments on disk: {raw}");
    }

    #[test]
    fn now_iso_ms_has_millisecond_precision() {
        let s = now_iso_ms();
        let (secs, ms) = s.split_once('.').expect("must contain fractional part");
        assert!(secs.parse::<u64>().is_ok(), "integer seconds part: {s}");
        assert_eq!(ms.len(), 3, "millisecond part must be 3 digits: {s}");
        assert!(ms.parse::<u32>().is_ok());
    }

    // ── from_chat token estimate parity (mirrors Python
    //    `tests/test_context_compaction.py::TestEstimateTokensIntegration`) ──

    #[test]
    fn from_chat_tokens_include_tool_calls() {
        let plain = SessionMessage::from_chat(&ChatMessage::assistant("hello"));
        let with_calls = SessionMessage::from_chat(&ChatMessage::assistant_with_tools(
            "hello",
            vec![ToolCall {
                id: "c1".into(),
                name: "run_shell".into(),
                arguments: format!("{{\"path\":\"{}\"}}", "a".repeat(200)),
            }],
        ));
        assert!(
            with_calls.tokens > plain.tokens,
            "tool_calls tokens ({}) must exceed plain ({})",
            with_calls.tokens,
            plain.tokens
        );
    }

    #[test]
    fn from_chat_tokens_include_tool_call_id() {
        let mut no_id = ChatMessage::tool_result("", "result".to_string());
        no_id.tool_call_id = None;
        let with_id = ChatMessage::tool_result("call_some_id_here", "result".to_string());
        assert!(
            SessionMessage::from_chat(&with_id).tokens > SessionMessage::from_chat(&no_id).tokens,
            "tool_call_id must add tokens"
        );
    }

    #[test]
    fn from_chat_tokens_include_reasoning_content() {
        let mut base = ChatMessage::assistant("ok");
        let plain = SessionMessage::from_chat(&base);
        base.reasoning_content = Some("Let me think".repeat(50));
        let with_reasoning = SessionMessage::from_chat(&base);
        assert!(
            with_reasoning.tokens > plain.tokens,
            "reasoning_content must add tokens"
        );
    }

    #[test]
    fn tool_call_only_assistant_has_nonzero_tokens() {
        // Regression: the estimate used to ignore tool_calls, so an assistant
        // message carrying only tool calls reported `tokens = 0`.
        let m = SessionMessage::from_chat(&ChatMessage::assistant_with_tools(
            "",
            vec![ToolCall {
                id: "c1".into(),
                name: "navigate".into(),
                arguments: "{\"url\":\"https://example.com\"}".into(),
            }],
        ));
        assert!(
            m.tokens > 0,
            "tool-call-only assistant must not be 0 tokens"
        );
    }

    // ── SESSION_FFI 存储方法 ──────────────────────────────────────────

    #[test]
    fn ensure_session_creates_then_touches() {
        let proj = tmp_project();
        let mut sm = SessionManager::new(&proj);
        sm.ensure_session("s1", "My Session").unwrap();
        assert!(
            sm.session_path("s1").exists(),
            "session file must be created"
        );
        assert_eq!(sm.load_index()["s1"].title, "My Session");

        // Idempotent re-ensure: bump last_activity, keep title on empty.
        let first = sm.load_index()["s1"].last_activity.clone();
        std::thread::sleep(std::time::Duration::from_millis(1100));
        sm.ensure_session("s1", "").unwrap();
        let index2 = sm.load_index();
        assert!(
            index2["s1"].last_activity >= first,
            "last_activity must bump"
        );
        assert_eq!(
            index2["s1"].title, "My Session",
            "empty title must not clobber"
        );
    }

    #[test]
    fn rename_session_sets_title() {
        let proj = tmp_project();
        let mut sm = SessionManager::new(&proj);
        sm.ensure_session("s1", "old").unwrap();
        sm.rename_session("s1", "new title").unwrap();
        assert_eq!(sm.load_index()["s1"].title, "new title");
        // Rename on a missing session is a no-op, not an error.
        sm.rename_session("missing", "x").unwrap();
    }

    #[test]
    fn touch_session_bumps_last_activity() {
        let proj = tmp_project();
        let mut sm = SessionManager::new(&proj);
        sm.ensure_session("s1", "").unwrap();
        let before = sm.load_index()["s1"].last_activity.clone();
        std::thread::sleep(std::time::Duration::from_millis(1100));
        sm.touch_session("s1").unwrap();
        assert!(sm.load_index()["s1"].last_activity >= before);
    }

    #[test]
    fn append_session_messages_appends_and_updates_index() {
        let proj = tmp_project();
        let mut sm = SessionManager::new(&proj);
        sm.ensure_session("s1", "").unwrap();
        sm.append_session_messages(
            "s1",
            vec![
                SessionMessage::from_chat(&ChatMessage::user("hello")),
                SessionMessage::from_chat(&ChatMessage::assistant("hi")),
            ],
        )
        .unwrap();
        let msgs = sm.read_session_messages("s1");
        assert_eq!(msgs.len(), 2);
        assert_eq!(msgs[0].content, "hello");
        let index = sm.load_index();
        assert_eq!(index["s1"].total_messages, 2);
        assert_eq!(
            index["s1"].title, "hello",
            "title falls back to first user msg"
        );
    }

    #[test]
    fn append_to_missing_session_creates_it() {
        let proj = tmp_project();
        let mut sm = SessionManager::new(&proj);
        sm.append_session_messages(
            "brand_new",
            vec![SessionMessage::from_chat(&ChatMessage::user("first"))],
        )
        .unwrap();
        assert!(sm.session_path("brand_new").exists());
        assert_eq!(sm.load_index()["brand_new"].total_messages, 1);
        assert_eq!(sm.load_index()["brand_new"].title, "first");
    }

    #[test]
    fn valid_session_id_rejects_path_traversal() {
        assert!(valid_session_id("session_1700000000_abc12345"));
        assert!(valid_session_id("cron_daily_1"));
        assert!(valid_session_id("pending_x"));
        assert!(!valid_session_id(""));
        assert!(!valid_session_id("../etc"));
        assert!(!valid_session_id("a/b"));
        assert!(!valid_session_id("a\\b"));
        assert!(!valid_session_id(".."));
    }

    #[test]
    fn old_file_without_schema_version_reads_fine() {
        let proj = tmp_project();
        let sm = SessionManager::new(&proj);
        let path = sm.session_path("s_old");
        std::fs::create_dir_all(path.parent().unwrap()).unwrap();
        // Old-format file: no `schema_version` field.
        let old = r#"{"session_id":"s_old","created_at":"1700000000","messages":[{"role":"user","content":"hi","timestamp":"1700000001","tokens":1}]}"#;
        std::fs::write(&path, old).unwrap();

        let msgs = sm.read_session_messages("s_old");
        assert_eq!(msgs.len(), 1);
        assert_eq!(msgs[0].content, "hi");

        // list_sessions fallback must rebuild the index from the old file.
        let sessions = sm.list_sessions();
        assert!(sessions.iter().any(|s| s.session_id == "s_old"));
    }

    #[test]
    fn context_checkpoint_roundtrips() {
        let proj = tmp_project();
        let mut sm = SessionManager::new(&proj);
        let sid = sm.create_session("t", None).unwrap();
        sm.set_context_checkpoint(Some(serde_json::json!({
            "cutoff_index": 3,
            "stub": "STUB",
            "session_id": &sid,
            "created_at": "now",
            "keep_tokens": 100
        })))
        .unwrap();

        let mut sm2 = SessionManager::new(&proj);
        assert!(sm2.switch_session(&sid).unwrap());
        let cp = sm2.context_checkpoint().expect("checkpoint restored");
        assert_eq!(cp["cutoff_index"], 3);
        assert_eq!(cp["stub"], "STUB");
    }

    #[test]
    fn export_jsonl_has_header_messages_footer() {
        let proj = tmp_project();
        let mut sm = SessionManager::new(&proj);
        let sid = sm.create_session("hello", None).unwrap();
        sm.add_message(SessionMessage::from_chat(&ChatMessage::assistant("hi")))
            .unwrap();
        sm.flush().unwrap();
        let out = sm.export_jsonl(&sid).unwrap();
        let lines: Vec<serde_json::Value> = out
            .lines()
            .map(|l| serde_json::from_str(l).unwrap())
            .collect();
        assert_eq!(lines[0]["type"], "header");
        assert_eq!(lines[0]["format"], "aacode.session-export");
        assert_eq!(lines.last().unwrap()["type"], "footer");
        assert_eq!(lines.last().unwrap()["messages"], 2);
        assert!(lines
            .iter()
            .any(|l| l["type"] == "message" && l["content"] == "hi"));
    }

    #[test]
    fn fork_moves_boundary_before_split_tool_group() {
        let proj = tmp_project();
        let mut sm = SessionManager::new(&proj);
        let sid = sm.create_session("src", None).unwrap();
        let asst = ChatMessage::assistant_with_tools(
            "",
            vec![ToolCall {
                id: "c1".into(),
                name: "run_shell".into(),
                arguments: "{}".into(),
            }],
        );
        sm.add_message(SessionMessage::from_chat(&asst)).unwrap();
        sm.add_message(SessionMessage::from_chat(&ChatMessage::tool_result(
            "c1", "ok",
        )))
        .unwrap();
        sm.add_message(SessionMessage::from_chat(&ChatMessage::assistant("done")))
            .unwrap();
        sm.flush().unwrap();
        // messages: 0 user, 1 assistant(tool_calls), 2 tool, 3 assistant
        // A fork at 2 would keep the assistant but drop its tool result.
        let new_sid = sm.fork_session(&sid, Some(2)).unwrap();
        let sm2 = SessionManager::new(&proj);
        let msgs = sm2.read_session_messages(&new_sid);
        assert_eq!(
            msgs.len(),
            1,
            "boundary should move back before the split group"
        );
        assert_eq!(msgs[0].role, "user");
        // Source untouched / not made current.
        assert_eq!(sm.messages.len(), 4);
        assert_eq!(sm.current_session_id.as_deref(), Some(sid.as_str()));

        // A complete boundary keeps the whole group.
        let kept = sm.fork_session(&sid, Some(3)).unwrap();
        let msgs = sm2.read_session_messages(&kept);
        assert_eq!(msgs.len(), 3);
        assert!(msgs.iter().any(|m| m.role == "tool"));
    }
}