297 lines
10 KiB
Rust
297 lines
10 KiB
Rust
//! 会话树与 JSONL 持久化。
|
||
//! Session trees and JSONL persistence.
|
||
//!
|
||
//! 参考 pi 的设计：每个会话是一个 JSONL 文件，追加写；每条 entry 有
|
||
//! `id` + `parentId`，从而形成可分支的树结构。文件位于 `<data>/sessions/`。
|
||
//! Mirroring pi: each session is a JSONL file, append-only; every entry has an
|
||
//! `id` + `parentId`, forming a branchable tree. Files live in
|
||
//! `<data>/sessions/`.
|
||
|
||
use focus_core::json::{FromJson, ToJson};
|
||
use focus_core::model::{now_ms, Message};
|
||
use focus_core::{CoreError, CoreResult};
|
||
use focus_json::JsonValue;
|
||
use std::collections::HashMap;
|
||
use std::fs::{File, OpenOptions};
|
||
use std::io::{BufRead, BufReader, Write};
|
||
use std::path::{Path, PathBuf};
|
||
|
||
/// 会话树中的一个节点。
|
||
/// One node of a session tree.
|
||
#[derive(Debug, Clone)]
|
||
pub struct SessionEntry {
|
||
/// 节点 id（会话内唯一，形如 `e1`、`e2`）。
|
||
/// Node id (unique per session, e.g. `e1`, `e2`).
|
||
pub id: String,
|
||
/// 父节点 id；`None` 表示树根。
|
||
/// Parent node id; `None` marks the root.
|
||
pub parent_id: Option<String>,
|
||
/// 节点承载的消息。
|
||
/// The message this node carries.
|
||
pub message: Message,
|
||
/// 毫秒级时间戳。
|
||
/// Millisecond timestamp.
|
||
pub timestamp: u64,
|
||
}
|
||
|
||
impl SessionEntry {
|
||
/// 构造一个新 entry（时间戳取当前时间）。
|
||
/// Build a new entry (timestamp = now).
|
||
pub fn new(id: String, parent_id: Option<String>, message: Message) -> Self {
|
||
Self {
|
||
id,
|
||
parent_id,
|
||
message,
|
||
timestamp: now_ms(),
|
||
}
|
||
}
|
||
|
||
/// 序列化为 JSON。
|
||
/// Serialize to JSON.
|
||
pub fn to_json(&self) -> JsonValue {
|
||
let mut o = JsonValue::obj();
|
||
o.insert("id", self.id.clone().into()).ok();
|
||
if let Some(p) = &self.parent_id {
|
||
o.insert("parentId", p.clone().into()).ok();
|
||
}
|
||
o.insert("message", self.message.to_json()).ok();
|
||
o.insert("timestamp", (self.timestamp as f64).into()).ok();
|
||
o
|
||
}
|
||
|
||
/// 从 JSON 反序列化。
|
||
/// Deserialize from JSON.
|
||
pub fn from_json(value: &JsonValue) -> CoreResult<Self> {
|
||
let id = value
|
||
.get_str("id")
|
||
.ok_or_else(|| CoreError::Json("session entry missing 'id'".into()))?
|
||
.to_string();
|
||
let parent_id = value.get_str("parentId").map(String::from);
|
||
let message_value = value
|
||
.get("message")
|
||
.ok_or_else(|| CoreError::Json("session entry missing 'message'".into()))?;
|
||
let message = Message::from_json(message_value)?;
|
||
let timestamp = value.get_num("timestamp").unwrap_or(0.0) as u64;
|
||
Ok(Self {
|
||
id,
|
||
parent_id,
|
||
message,
|
||
timestamp,
|
||
})
|
||
}
|
||
}
|
||
|
||
/// 会话的 JSONL 存储。
|
||
/// JSONL storage for sessions.
|
||
#[derive(Debug, Clone)]
|
||
pub struct SessionStore {
|
||
data_dir: PathBuf,
|
||
}
|
||
|
||
impl Default for SessionStore {
|
||
/// 使用默认数据根目录（`~/.focus`）构造存储。
|
||
/// Build a store at the default data root (`~/.focus`).
|
||
fn default() -> Self {
|
||
Self::new(crate::data_dir())
|
||
}
|
||
}
|
||
|
||
impl SessionStore {
|
||
/// 使用指定数据根目录构造存储。
|
||
/// Build a store rooted at the given data directory.
|
||
pub fn new(data_dir: impl Into<PathBuf>) -> Self {
|
||
Self {
|
||
data_dir: data_dir.into(),
|
||
}
|
||
}
|
||
|
||
/// 会话文件的目录（`<data>/sessions`）。
|
||
/// The sessions directory (`<data>/sessions`).
|
||
pub fn sessions_dir(&self) -> PathBuf {
|
||
self.data_dir.join("sessions")
|
||
}
|
||
|
||
/// 某会话的 JSONL 文件路径。
|
||
/// The JSONL file path for a session.
|
||
pub fn session_path(&self, session_id: &str) -> PathBuf {
|
||
self.sessions_dir().join(format!("{}.jsonl", session_id))
|
||
}
|
||
|
||
/// 向会话追加一条 entry（自动建目录与文件）。
|
||
/// Append an entry to a session (auto-creating dirs and the file).
|
||
pub fn append(&self, session_id: &str, entry: &SessionEntry) -> CoreResult<()> {
|
||
let path = self.session_path(session_id);
|
||
if let Some(parent) = path.parent() {
|
||
std::fs::create_dir_all(parent)
|
||
.map_err(|e| CoreError::Tool(format!("session dir {}: {}", parent.display(), e)))?;
|
||
}
|
||
let mut file = OpenOptions::new()
|
||
.create(true)
|
||
.append(true)
|
||
.open(&path)
|
||
.map_err(|e| CoreError::Tool(format!("open session {}: {}", path.display(), e)))?;
|
||
let line = focus_json::to_string(&entry.to_json());
|
||
writeln!(file, "{}", line)
|
||
.map_err(|e| CoreError::Tool(format!("append session {}: {}", path.display(), e)))?;
|
||
Ok(())
|
||
}
|
||
|
||
/// 加载一个会话的全部 entry（按追加顺序）。
|
||
/// Load all entries of a session (in append order).
|
||
pub fn load(&self, session_id: &str) -> CoreResult<Vec<SessionEntry>> {
|
||
let path = self.session_path(session_id);
|
||
if !path.exists() {
|
||
return Ok(Vec::new());
|
||
}
|
||
let file = File::open(&path)
|
||
.map_err(|e| CoreError::Tool(format!("open session {}: {}", path.display(), e)))?;
|
||
let reader = BufReader::new(file);
|
||
let mut entries = Vec::new();
|
||
for (i, line) in reader.lines().enumerate() {
|
||
let line = line
|
||
.map_err(|e| CoreError::Tool(format!("read session {}: {}", path.display(), e)))?;
|
||
let trimmed = line.trim();
|
||
if trimmed.is_empty() {
|
||
continue;
|
||
}
|
||
let value = focus_json::parse(trimmed)
|
||
.map_err(|e| CoreError::Json(format!("session line {}: {}", i + 1, e)))?;
|
||
entries.push(SessionEntry::from_json(&value)?);
|
||
}
|
||
Ok(entries)
|
||
}
|
||
|
||
/// 列出所有会话 id（`*.jsonl` 文件名去后缀）。
|
||
/// List all session ids (`*.jsonl` file names minus the extension).
|
||
pub fn list_sessions(&self) -> CoreResult<Vec<String>> {
|
||
let dir = self.sessions_dir();
|
||
if !dir.exists() {
|
||
return Ok(Vec::new());
|
||
}
|
||
let mut ids = Vec::new();
|
||
for entry in
|
||
std::fs::read_dir(&dir).map_err(|e| CoreError::Tool(format!("list sessions: {}", e)))?
|
||
{
|
||
let entry = entry.map_err(|e| CoreError::Tool(format!("list sessions: {}", e)))?;
|
||
let name = entry.file_name().to_string_lossy().to_string();
|
||
if let Some(stem) = name.strip_suffix(".jsonl") {
|
||
ids.push(stem.to_string());
|
||
}
|
||
}
|
||
ids.sort();
|
||
Ok(ids)
|
||
}
|
||
|
||
/// 生成下一个 entry id（`e{n+1}`，基于现有最大序号）。
|
||
/// Generate the next entry id (`e{n+1}`, based on the largest existing
|
||
/// sequence number).
|
||
pub fn next_entry_id(&self, session_id: &str) -> CoreResult<String> {
|
||
let entries = self.load(session_id)?;
|
||
let max = entries
|
||
.iter()
|
||
.filter_map(|e| e.id.strip_prefix('e'))
|
||
.filter_map(|n| n.parse::<u64>().ok())
|
||
.max()
|
||
.unwrap_or(0);
|
||
Ok(format!("e{}", max + 1))
|
||
}
|
||
}
|
||
|
||
/// 会话树的只读查询视图。
|
||
/// A read-only query view over a session tree.
|
||
#[derive(Debug, Default)]
|
||
pub struct SessionTree {
|
||
entries: Vec<SessionEntry>,
|
||
by_id: HashMap<String, usize>,
|
||
children: HashMap<String, Vec<usize>>,
|
||
roots: Vec<usize>,
|
||
}
|
||
|
||
impl SessionTree {
|
||
/// 从 entry 列表构建树。
|
||
/// Build a tree from a list of entries.
|
||
pub fn build(entries: Vec<SessionEntry>) -> Self {
|
||
let mut by_id = HashMap::new();
|
||
let mut children: HashMap<String, Vec<usize>> = HashMap::new();
|
||
let mut roots = Vec::new();
|
||
for (i, e) in entries.iter().enumerate() {
|
||
by_id.insert(e.id.clone(), i);
|
||
match &e.parent_id {
|
||
Some(p) => children.entry(p.clone()).or_default().push(i),
|
||
None => roots.push(i),
|
||
}
|
||
}
|
||
Self {
|
||
entries,
|
||
by_id,
|
||
children,
|
||
roots,
|
||
}
|
||
}
|
||
|
||
/// 所有根节点（无父节点的 entry）。
|
||
/// All root nodes (entries without a parent).
|
||
pub fn roots(&self) -> Vec<&SessionEntry> {
|
||
self.roots.iter().map(|&i| &self.entries[i]).collect()
|
||
}
|
||
|
||
/// 按 id 查找节点。
|
||
/// Look up a node by id.
|
||
pub fn get(&self, id: &str) -> Option<&SessionEntry> {
|
||
self.by_id.get(id).map(|&i| &self.entries[i])
|
||
}
|
||
|
||
/// 某节点的直接子节点。
|
||
/// Direct children of a node.
|
||
pub fn children_of(&self, id: &str) -> Vec<&SessionEntry> {
|
||
self.children
|
||
.get(id)
|
||
.map(|v| v.iter().map(|&i| &self.entries[i]).collect())
|
||
.unwrap_or_default()
|
||
}
|
||
|
||
/// 从根到 `id` 的路径（含两端）。
|
||
/// The path from the root to `id` (inclusive).
|
||
pub fn branch(&self, id: &str) -> Vec<&SessionEntry> {
|
||
let mut path = Vec::new();
|
||
let mut current = Some(id.to_string());
|
||
while let Some(cid) = current {
|
||
match self.get(&cid) {
|
||
Some(e) => {
|
||
path.push(e);
|
||
current = e.parent_id.clone();
|
||
}
|
||
None => break,
|
||
}
|
||
}
|
||
path.reverse();
|
||
path
|
||
}
|
||
|
||
/// 所有叶子节点 id（无子节点的节点）。
|
||
/// Ids of all leaf nodes (nodes without children).
|
||
pub fn leaf_ids(&self) -> Vec<&str> {
|
||
self.entries
|
||
.iter()
|
||
.filter(|e| !self.children.contains_key(&e.id))
|
||
.map(|e| e.id.as_str())
|
||
.collect()
|
||
}
|
||
|
||
/// 最近的叶子节点（按时间戳）。
|
||
/// The most recent leaf node (by timestamp).
|
||
pub fn latest_leaf(&self) -> Option<&SessionEntry> {
|
||
self.leaf_ids()
|
||
.into_iter()
|
||
.filter_map(|id| self.get(id))
|
||
.max_by_key(|e| e.timestamp)
|
||
}
|
||
}
|
||
|
||
/// 确保 `path` 存在且为目录；不存在则创建。
|
||
/// Ensure `path` exists as a directory, creating it if needed.
|
||
pub fn ensure_dir(path: &Path) -> CoreResult<()> {
|
||
std::fs::create_dir_all(path)
|
||
.map_err(|e| CoreError::Tool(format!("create dir {}: {}", path.display(), e)))
|
||
}
|