feat(harness): session tree, compaction plans, system prompt templates
- SessionStore/SessionTree: JSONL append-only tree (id + parentId) under ~/.focus/sessions, branching and continuation queries - compaction: lightweight token estimation; scheme D+E plans (summary + key-facts, keep recent messages, tool_call/tool_result pairs kept intact); auto-trigger threshold (default 80%, configurable); apply_summary writes back - prompts: system template with cwd/os/date/shell + context occupancy and last-turn usage (scheme E); pure-std date formatting
This commit is contained in:
DaiChaoXiong
parent
ee165235ec
commit
647cf1ba83
|
|
@ -9,4 +9,3 @@ description = "Session persistence (JSONL tree), context compaction, system prom
|
|||
[dependencies]
|
||||
focus-core.workspace = true
|
||||
focus-json.workspace = true
|
||||
tokio = { workspace = true, features = ["rt", "fs", "sync"] }
|
||||
|
|
|
|||
|
|
@ -0,0 +1,357 @@
|
|||
//! 上下文压缩：token 估算与压缩方案（方案 D+E）。
|
||||
//! Context compaction: token estimation and plans (schemes D+E).
|
||||
//!
|
||||
//! - 方案 D：把最早的一批消息摘要化（`<summary>` + `<key-facts>`），保留最近
|
||||
//! 消息与关键 tool_result 链路，避免切断 tool_call → tool_result 配对；
|
||||
//! - 方案 E：压缩阈值基于估算占用（默认窗口的 80%，可配置）。
|
||||
//! - Scheme D: summarize the oldest messages (`<summary>` + `<key-facts>`),
|
||||
//! keep the most recent messages and intact tool_call → tool_result chains;
|
||||
//! - Scheme E: the compaction threshold is based on the estimated occupancy
|
||||
//! (default 80% of the window, configurable).
|
||||
//!
|
||||
//! 本模块只产出「方案」：真正的摘要 LLM 调用由上层按 [`CompactionPlan`]
|
||||
//! 执行，再用 [`apply_summary`] 把摘要写回消息列表。
|
||||
//! This module only produces *plans*: the actual summary LLM call is executed
|
||||
//! by the upper layer per the [`CompactionPlan`], and [`apply_summary`] writes
|
||||
//! the summary back into the message list.
|
||||
|
||||
use focus_core::model::{ContentBlock, Message};
|
||||
|
||||
/// 默认触发阈值：估算占用达到窗口的 80% 时建议/执行压缩。
|
||||
/// Default trigger threshold: compact at 80% estimated occupancy.
|
||||
pub const DEFAULT_THRESHOLD_RATIO: f64 = 0.8;
|
||||
/// 默认保留比例：压缩后保留的消息占用窗口的比例。
|
||||
/// Default keep ratio: how much of the window the kept messages may occupy.
|
||||
pub const DEFAULT_KEEP_RATIO: f64 = 0.6;
|
||||
|
||||
/// 轻量 token 估算：ASCII 每 4 字符约 1 token，非 ASCII（中文等）每字符约 1 token。
|
||||
/// Lightweight token estimate: ~1 token per 4 ASCII chars, ~1 per non-ASCII
|
||||
/// (e.g. CJK) char. Deliberately cheap and deterministic, not tokenizer-accurate.
|
||||
pub fn estimate_tokens(text: &str) -> u64 {
|
||||
let (ascii, other) = text.chars().fold((0u64, 0u64), |(a, o), c| {
|
||||
if c.is_ascii() {
|
||||
(a + 1, o)
|
||||
} else {
|
||||
(a, o + 1)
|
||||
}
|
||||
});
|
||||
ascii / 4 + other
|
||||
}
|
||||
|
||||
/// 估算一条消息的 token 数（内容 + 少量开销）。
|
||||
/// Estimate the tokens of one message (content + small overhead).
|
||||
fn estimate_message(m: &Message) -> u64 {
|
||||
let text: String = match m {
|
||||
Message::User(u) => u
|
||||
.content
|
||||
.iter()
|
||||
.filter_map(|c| match c {
|
||||
ContentBlock::Text(t) => Some(t.text.clone()),
|
||||
_ => None,
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join(" "),
|
||||
Message::Assistant(a) => {
|
||||
let mut parts: Vec<String> = a
|
||||
.content
|
||||
.iter()
|
||||
.filter_map(|c| match c {
|
||||
ContentBlock::Text(t) => Some(t.text.clone()),
|
||||
ContentBlock::Thinking(t) => Some(t.thinking.clone()),
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
for tc in a.content.iter().filter_map(|c| c.as_tool_call()) {
|
||||
parts.push(format!(
|
||||
"tool {} {}",
|
||||
tc.name,
|
||||
focus_json::to_string(&tc.arguments)
|
||||
));
|
||||
}
|
||||
parts.join(" ")
|
||||
}
|
||||
Message::ToolResult(t) => t
|
||||
.content
|
||||
.iter()
|
||||
.filter_map(|c| match c {
|
||||
ContentBlock::Text(x) => Some(x.text.clone()),
|
||||
_ => None,
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join(" "),
|
||||
};
|
||||
estimate_tokens(&text) + 4 // 消息级开销 / per-message overhead
|
||||
}
|
||||
|
||||
/// 估算整个消息列表的 token 数。
|
||||
/// Estimate the tokens of a whole message list.
|
||||
pub fn estimate_messages(messages: &[Message]) -> u64 {
|
||||
messages.iter().map(estimate_message).sum()
|
||||
}
|
||||
|
||||
/// 一份压缩方案。
|
||||
/// A compaction plan.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct CompactionPlan {
|
||||
/// 需要交给摘要模型的消息（最早的一批）。
|
||||
/// Messages handed to the summarizer (the oldest batch).
|
||||
pub summarize: Vec<Message>,
|
||||
/// 原样保留的最近消息。
|
||||
/// Most recent messages kept verbatim.
|
||||
pub keep: Vec<Message>,
|
||||
/// 摘要模型的提示文本（要求产出 `<summary>` 与 `<key-facts>`）。
|
||||
/// Summarizer prompt (asks for `<summary>` and `<key-facts>`).
|
||||
pub summary_instruction: String,
|
||||
/// 压缩后预计节省的 token 数。
|
||||
/// Estimated tokens saved by compacting.
|
||||
pub estimated_saved_tokens: u64,
|
||||
/// 预计摘要结果占用的 token 数。
|
||||
/// Estimated tokens the summary will occupy.
|
||||
pub estimated_summary_tokens: u64,
|
||||
}
|
||||
|
||||
/// 生成摘要模型提示。
|
||||
/// Build the summarizer prompt.
|
||||
fn build_summary_instruction() -> String {
|
||||
"Summarize the following conversation messages for a follow-up conversation. \
|
||||
Preserve: decisions made, user preferences, project constraints, and any \
|
||||
facts that will be needed later.\n\
|
||||
Output exactly two sections:\n\
|
||||
<summary>\nA concise narrative summary of what happened.\n\
|
||||
</summary>\n\
|
||||
<key-facts>\nA bullet list of important facts, preferences and constraints.\n\
|
||||
</key-facts>\n"
|
||||
.to_string()
|
||||
}
|
||||
|
||||
/// 估算摘要结果的 token 数（被摘要内容的 10%，封顶 2000）。
|
||||
/// Estimate the summary's token count (10% of the summarized content, capped
|
||||
/// at 2000).
|
||||
fn estimate_summary_tokens(summarize_tokens: u64) -> u64 {
|
||||
(summarize_tokens / 10).clamp(64, 2000)
|
||||
}
|
||||
|
||||
/// 规划一次压缩；估算占用未超过阈值时返回 `None`。
|
||||
/// Plan a compaction; `None` when the estimated occupancy is under the
|
||||
/// threshold.
|
||||
///
|
||||
/// `context_window`：模型上下文窗口（上层已按「配置 → 已知表 → 兜底」解析）；
|
||||
/// `threshold_ratio`：触发阈值（默认 0.8）；`keep_ratio`：保留比例（默认 0.6）。
|
||||
/// `context_window`: the model's context window (resolved by the upper layer as
|
||||
/// configured → known table → fallback); `threshold_ratio`: the trigger
|
||||
/// threshold (default 0.8); `keep_ratio`: the keep ratio (default 0.6).
|
||||
pub fn plan_compaction(
|
||||
messages: &[Message],
|
||||
context_window: u64,
|
||||
threshold_ratio: f64,
|
||||
keep_ratio: f64,
|
||||
) -> Option<CompactionPlan> {
|
||||
if messages.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let total = estimate_messages(messages);
|
||||
if (total as f64) < (context_window as f64) * threshold_ratio {
|
||||
return None;
|
||||
}
|
||||
|
||||
// 从尾部往回保留，直到达到保留预算。
|
||||
// Walk backwards from the tail until the keep budget is met.
|
||||
let keep_budget = ((context_window as f64) * keep_ratio) as u64;
|
||||
let mut kept_tokens = 0u64;
|
||||
let mut boundary = messages.len();
|
||||
while boundary > 0 {
|
||||
let m = &messages[boundary - 1];
|
||||
let cost = estimate_message(m);
|
||||
if kept_tokens + cost > keep_budget && kept_tokens > 0 {
|
||||
break;
|
||||
}
|
||||
kept_tokens += cost;
|
||||
boundary -= 1;
|
||||
}
|
||||
|
||||
// 不要切断 tool_call → tool_result 配对：若保留区以 tool_result 开头，
|
||||
// 向前延伸以把其 assistant 调用消息一并保留。
|
||||
// Don't split tool_call → tool_result pairs: if the kept region starts
|
||||
// with a tool_result, extend backward to also keep its assistant call
|
||||
// message.
|
||||
while boundary > 0
|
||||
&& boundary < messages.len()
|
||||
&& matches!(messages[boundary], Message::ToolResult(_))
|
||||
{
|
||||
boundary -= 1;
|
||||
}
|
||||
|
||||
if boundary == 0 {
|
||||
return None; // 全都保留，无需压缩 / keep everything; nothing to do
|
||||
}
|
||||
|
||||
let summarize = messages[..boundary].to_vec();
|
||||
let keep = messages[boundary..].to_vec();
|
||||
let summarize_tokens = estimate_messages(&summarize);
|
||||
let summary_tokens = estimate_summary_tokens(summarize_tokens);
|
||||
|
||||
Some(CompactionPlan {
|
||||
summarize,
|
||||
keep,
|
||||
summary_instruction: build_summary_instruction(),
|
||||
estimated_saved_tokens: summarize_tokens.saturating_sub(summary_tokens),
|
||||
estimated_summary_tokens: summary_tokens,
|
||||
})
|
||||
}
|
||||
|
||||
/// 把摘要写回：一条 user 摘要消息 + 保留的消息。
|
||||
/// Write the summary back: one user summary message + the kept messages.
|
||||
pub fn apply_summary(summary: &str, plan: &CompactionPlan) -> Vec<Message> {
|
||||
let mut out = Vec::with_capacity(plan.keep.len() + 1);
|
||||
let text = format!(
|
||||
"Here is a summary of our earlier conversation:\n\n{}",
|
||||
summary
|
||||
);
|
||||
out.push(Message::user_text(text));
|
||||
out.extend(plan.keep.clone());
|
||||
out
|
||||
}
|
||||
|
||||
/// 供 TUI 展示的压缩说明文本。
|
||||
/// Human-readable description of a plan, for the TUI.
|
||||
pub fn describe_plan(plan: &CompactionPlan) -> String {
|
||||
format!(
|
||||
"compaction: summarize {} messages (~{} tokens) into ~{} tokens; keep {} messages (~{} tokens); saves ~{} tokens",
|
||||
plan.summarize.len(),
|
||||
estimate_messages(&plan.summarize),
|
||||
plan.estimated_summary_tokens,
|
||||
plan.keep.len(),
|
||||
estimate_messages(&plan.keep),
|
||||
plan.estimated_saved_tokens,
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use focus_core::model::*;
|
||||
use focus_json::JsonValue;
|
||||
|
||||
fn long_user_message(text: &str) -> Message {
|
||||
Message::user_text(text)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn estimate_is_reasonable() {
|
||||
// 80 个 ASCII 字符 ≈ 20 token。
|
||||
// 80 ASCII chars ≈ 20 tokens.
|
||||
let text = "a".repeat(80);
|
||||
assert_eq!(estimate_tokens(&text), 20);
|
||||
// 中文按字符计。
|
||||
// CJK counted per char.
|
||||
assert_eq!(estimate_tokens("中文"), 2);
|
||||
// 消息开销计入。
|
||||
// Per-message overhead is included.
|
||||
assert!(estimate_messages(&[long_user_message("hi")]) > estimate_tokens("hi"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn below_threshold_returns_none() {
|
||||
let messages: Vec<Message> = (0..5)
|
||||
.map(|i| long_user_message(&format!("msg {}", i)))
|
||||
.collect();
|
||||
// 窗口极大 → 不触发。
|
||||
// Huge window → no trigger.
|
||||
assert!(plan_compaction(
|
||||
&messages,
|
||||
1_000_000,
|
||||
DEFAULT_THRESHOLD_RATIO,
|
||||
DEFAULT_KEEP_RATIO
|
||||
)
|
||||
.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn above_threshold_plans_summarize_and_keep() {
|
||||
// 每条消息约 50+ token（"x"*200 → 50 token + 4 开销）；10 条约 540 token。
|
||||
// Each message ≈ 50+ tokens; 10 of them ≈ 540 tokens.
|
||||
let messages: Vec<Message> = (0..10)
|
||||
.map(|_i| long_user_message(&"x".repeat(200)))
|
||||
.collect();
|
||||
let plan = plan_compaction(&messages, 200, DEFAULT_THRESHOLD_RATIO, DEFAULT_KEEP_RATIO)
|
||||
.expect("plan");
|
||||
assert!(!plan.summarize.is_empty());
|
||||
assert!(!plan.keep.is_empty());
|
||||
assert!(plan.summarize.len() + plan.keep.len() == messages.len());
|
||||
// 摘要指令包含两个部分。
|
||||
// The summary instruction has both sections.
|
||||
assert!(plan.summary_instruction.contains("<summary>"));
|
||||
assert!(plan.summary_instruction.contains("<key-facts>"));
|
||||
// 保存的 token 为正。
|
||||
// Saved tokens are positive.
|
||||
assert!(plan.estimated_saved_tokens > 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn keeps_tool_result_pairs_intact() {
|
||||
// 构造：user → assistant(tool_call) → toolResult → assistant(text)，
|
||||
// 其中 tool 部分恰好落在边界附近。
|
||||
// Build: user → assistant(tool_call) → toolResult → assistant(text),
|
||||
// with the tool part landing near the boundary.
|
||||
let mut args = JsonValue::obj();
|
||||
args.insert("x", "1".into()).ok();
|
||||
let mut messages = vec![long_user_message(&"a".repeat(300))];
|
||||
messages.push(Message::Assistant(AssistantMessage {
|
||||
content: vec![ContentBlock::ToolCall(ToolCall {
|
||||
id: "c1".into(),
|
||||
name: "read".into(),
|
||||
arguments: args,
|
||||
})],
|
||||
model: "m".into(),
|
||||
usage: Usage::default(),
|
||||
stop_reason: StopReason::ToolUse,
|
||||
error_message: None,
|
||||
timestamp: 0,
|
||||
}));
|
||||
messages.push(Message::ToolResult(ToolResultMessage {
|
||||
tool_call_id: "c1".into(),
|
||||
tool_name: "read".into(),
|
||||
content: vec![ContentBlock::text("r".repeat(300))],
|
||||
details: JsonValue::obj(),
|
||||
is_error: false,
|
||||
timestamp: 0,
|
||||
}));
|
||||
messages.push(Message::Assistant(AssistantMessage {
|
||||
content: vec![ContentBlock::text("final")],
|
||||
model: "m".into(),
|
||||
usage: Usage::default(),
|
||||
stop_reason: StopReason::Stop,
|
||||
error_message: None,
|
||||
timestamp: 0,
|
||||
}));
|
||||
|
||||
let plan = plan_compaction(&messages, 100, DEFAULT_THRESHOLD_RATIO, DEFAULT_KEEP_RATIO)
|
||||
.expect("plan");
|
||||
// 若 tool_result 在 keep 中，其 assistant 调用消息必须也在 keep 中。
|
||||
// If the tool_result is kept, its assistant call message must be kept too.
|
||||
let keep: Vec<&Message> = plan.keep.iter().collect();
|
||||
if keep.iter().any(|m| matches!(m, Message::ToolResult(_))) {
|
||||
assert!(keep.iter().any(|m| matches!(m, Message::Assistant(a) if a.content.iter().any(|c| c.as_tool_call().is_some()))));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn apply_summary_prepends_summary_message() {
|
||||
let messages: Vec<Message> = (0..6)
|
||||
.map(|_i| long_user_message(&"x".repeat(200)))
|
||||
.collect();
|
||||
let plan = plan_compaction(&messages, 100, 0.8, 0.6).expect("plan");
|
||||
let out = apply_summary("SUMMARY TEXT", &plan);
|
||||
assert_eq!(out.len(), plan.keep.len() + 1);
|
||||
match &out[0] {
|
||||
Message::User(u) => {
|
||||
let text = u.content[0].as_text().unwrap().text.clone();
|
||||
assert!(text.contains("SUMMARY TEXT"));
|
||||
}
|
||||
other => panic!("expected user summary, got {:?}", other),
|
||||
}
|
||||
// 保留部分原样在后。
|
||||
// The kept part follows verbatim.
|
||||
assert_eq!(out[1..], plan.keep[..]);
|
||||
}
|
||||
}
|
||||
|
|
@ -1,7 +1,48 @@
|
|||
//! focus-harness：会话持久化、上下文压缩与系统提示词模板。
|
||||
//! focus-harness: session persistence, compaction, system prompts.
|
||||
//! focus-harness: session persistence, context compaction, and system prompts.
|
||||
//!
|
||||
//! 本里程碑阶段尚未实现。
|
||||
//! Not yet implemented in this milestone.
|
||||
//! 所有数据默认存放在用户目录下的 `~/.focus/`（Linux `$HOME/.focus`，
|
||||
//! Windows `%USERPROFILE%\\.focus`；可用 `FOCUS_DATA_DIR` 覆盖）。
|
||||
//! 压缩只产出「方案」（估算 + 切点 + 摘要提示），真正的摘要 LLM 调用由
|
||||
//! 上层（TUI）按方案执行——这遵守了 harness 不得依赖 provider 的依赖图约束。
|
||||
//! Data lives under `~/.focus/` by default (Linux `$HOME/.focus`, Windows
|
||||
//! `%USERPROFILE%\\.focus`; override with `FOCUS_DATA_DIR`). Compaction only
|
||||
//! produces a *plan* (estimate + cut points + summary instructions); the actual
|
||||
//! summarization LLM call is executed by the upper layer (TUI) following the
|
||||
//! plan — honoring the dependency-graph rule that harness never depends on
|
||||
//! providers.
|
||||
|
||||
#![forbid(unsafe_code)]
|
||||
|
||||
/// 上下文压缩：token 估算与压缩方案（方案 D+E）。
|
||||
/// Context compaction: token estimation and plans (schemes D+E).
|
||||
pub mod compaction;
|
||||
/// 系统提示词模板（含 Usage / 上下文占用信息）。
|
||||
/// System prompt templates (with usage / context-occupancy info).
|
||||
pub mod prompt;
|
||||
/// 会话树与 JSONL 持久化。
|
||||
/// Session trees and JSONL persistence.
|
||||
pub mod session;
|
||||
|
||||
pub use compaction::{
|
||||
apply_summary, estimate_messages, estimate_tokens, plan_compaction, CompactionPlan,
|
||||
};
|
||||
pub use prompt::{ContextUsage, PromptContext, SystemPromptTemplate};
|
||||
pub use session::{SessionEntry, SessionStore, SessionTree};
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
/// 解析数据根目录：`FOCUS_DATA_DIR` 优先，否则用户目录下的 `.focus`。
|
||||
/// Resolve the data root: `FOCUS_DATA_DIR` wins, else `.focus` under the
|
||||
/// user's home directory.
|
||||
pub fn data_dir() -> PathBuf {
|
||||
if let Ok(d) = std::env::var("FOCUS_DATA_DIR") {
|
||||
if !d.is_empty() {
|
||||
return PathBuf::from(d);
|
||||
}
|
||||
}
|
||||
let home = std::env::var("HOME")
|
||||
.or_else(|_| std::env::var("USERPROFILE"))
|
||||
.unwrap_or_else(|_| ".".to_string());
|
||||
Path::new(&home).join(".focus")
|
||||
}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,314 @@
|
|||
//! 系统提示词模板：注入环境信息与 Usage / 上下文占用（方案 E 的软压缩）。
|
||||
//! System prompt templates: environment info plus usage / context-occupancy
|
||||
//! (scheme E's budget awareness).
|
||||
//!
|
||||
//! 模板用 `{placeholders}` 占位，渲染时替换。上层（TUI）在每回合前用
|
||||
//! [`SystemPromptTemplate::render`] 生成最新系统提示，再调用
|
||||
//! `Agent::set_system_prompt` 注入。
|
||||
//! Templates use `{placeholders}`, substituted at render time. The upper layer
|
||||
//! (TUI) renders a fresh prompt before each turn and injects it via
|
||||
//! `Agent::set_system_prompt`.
|
||||
|
||||
use focus_core::model::Usage;
|
||||
use std::path::PathBuf;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
/// 当前上下文的占用信息。
|
||||
/// Current context occupancy.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ContextUsage {
|
||||
/// 估算的上下文 token 数。
|
||||
/// Estimated context tokens.
|
||||
pub estimated_tokens: u64,
|
||||
/// 模型上下文窗口（已解析）。
|
||||
/// The model's context window (already resolved).
|
||||
pub context_window: u64,
|
||||
/// 窗口是否来自已知表/配置；`false` 表示未知模型用了保守兜底值。
|
||||
/// Whether the window came from config/known table; `false` means the
|
||||
/// conservative fallback was used for an unknown model.
|
||||
pub window_known: bool,
|
||||
}
|
||||
|
||||
impl ContextUsage {
|
||||
/// 占用比例（0.0–1.0）。
|
||||
/// Occupancy ratio (0.0–1.0).
|
||||
pub fn ratio(&self) -> f64 {
|
||||
if self.context_window == 0 {
|
||||
0.0
|
||||
} else {
|
||||
(self.estimated_tokens as f64) / (self.context_window as f64)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 渲染系统提示所需的环境与用量信息。
|
||||
/// Environment and usage info needed to render the system prompt.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PromptContext {
|
||||
/// 当前工作目录。
|
||||
/// The current working directory.
|
||||
pub cwd: PathBuf,
|
||||
/// 操作系统（`std::env::consts::OS`）。
|
||||
/// The operating system (`std::env::consts::OS`).
|
||||
pub os: &'static str,
|
||||
/// 当前日期（`yyyy-mm-dd`）。
|
||||
/// The current date (`yyyy-mm-dd`).
|
||||
pub date: String,
|
||||
/// 平台 shell 名称。
|
||||
/// The platform shell name.
|
||||
pub shell: String,
|
||||
/// 上下文占用（方案 E）。
|
||||
/// Context occupancy (scheme E).
|
||||
pub context_usage: Option<ContextUsage>,
|
||||
/// 最近一次回合的用量。
|
||||
/// Usage of the most recent turn.
|
||||
pub last_usage: Option<Usage>,
|
||||
/// 附加说明行（每行一条）。
|
||||
/// Extra note lines (one per entry).
|
||||
pub extra: Vec<String>,
|
||||
}
|
||||
|
||||
/// 系统提示词模板。
|
||||
/// A system prompt template.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SystemPromptTemplate {
|
||||
template: String,
|
||||
}
|
||||
|
||||
/// 默认模板。
|
||||
/// The default template.
|
||||
const DEFAULT_TEMPLATE: &str = r#"You are focus, a terminal-based coding agent working in the repository at {cwd}.
|
||||
|
||||
Environment:
|
||||
- Operating system: {os}
|
||||
- Current date: {date}
|
||||
- Shell: {shell}
|
||||
|
||||
{context_usage}
|
||||
{last_usage}
|
||||
{extra}
|
||||
|
||||
You can read and edit files and run shell commands to accomplish the user's
|
||||
requests. When a tool call fails, read the error, fix your approach, and try
|
||||
again. Prefer small, focused changes."#;
|
||||
|
||||
impl Default for SystemPromptTemplate {
|
||||
/// 默认模板。
|
||||
/// The default template.
|
||||
fn default() -> Self {
|
||||
Self::new(DEFAULT_TEMPLATE)
|
||||
}
|
||||
}
|
||||
|
||||
impl SystemPromptTemplate {
|
||||
/// 自定义模板。
|
||||
/// A custom template.
|
||||
pub fn new(template: impl Into<String>) -> Self {
|
||||
Self {
|
||||
template: template.into(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 渲染模板。未知占位符保留原样（容忍部分字段缺失）。
|
||||
/// Render the template. Unknown placeholders are left as-is (tolerates
|
||||
/// missing fields).
|
||||
pub fn render(&self, ctx: &PromptContext) -> String {
|
||||
let mut out = self.template.clone();
|
||||
out = out.replace("{cwd}", &ctx.cwd.display().to_string());
|
||||
out = out.replace("{os}", ctx.os);
|
||||
out = out.replace("{date}", &ctx.date);
|
||||
out = out.replace("{shell}", &ctx.shell);
|
||||
out = out.replace("{context_usage}", &render_context_usage(ctx));
|
||||
out = out.replace("{last_usage}", &render_last_usage(ctx));
|
||||
let extra = if ctx.extra.is_empty() {
|
||||
String::new()
|
||||
} else {
|
||||
format!("Notes:\n- {}", ctx.extra.join("\n- "))
|
||||
};
|
||||
out = out.replace("{extra}", &extra);
|
||||
out
|
||||
}
|
||||
}
|
||||
|
||||
/// 渲染上下文占用段。
|
||||
/// Render the context-usage section.
|
||||
fn render_context_usage(ctx: &PromptContext) -> String {
|
||||
match &ctx.context_usage {
|
||||
None => String::new(),
|
||||
Some(u) => {
|
||||
if u.window_known {
|
||||
format!(
|
||||
"Context usage: {} / {} tokens ({:.1}%).",
|
||||
format_thousands(u.estimated_tokens),
|
||||
format_thousands(u.context_window),
|
||||
u.ratio() * 100.0
|
||||
)
|
||||
} else {
|
||||
format!(
|
||||
"Context usage: {} tokens (window unknown; assuming {}).",
|
||||
format_thousands(u.estimated_tokens),
|
||||
format_thousands(u.context_window)
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 渲染最近用量段。
|
||||
/// Render the last-usage section.
|
||||
fn render_last_usage(ctx: &PromptContext) -> String {
|
||||
match &ctx.last_usage {
|
||||
None => String::new(),
|
||||
Some(u) => format!(
|
||||
"Last turn usage: input {}, output {}, cache-read {}, cache-write {}.",
|
||||
format_thousands(u.input_tokens),
|
||||
format_thousands(u.output_tokens),
|
||||
format_thousands(u.cache_read_tokens),
|
||||
format_thousands(u.cache_write_tokens)
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
/// 千位分隔格式化（纯 std，无外部依赖）。
|
||||
/// Format with thousands separators (pure std, no external deps).
|
||||
fn format_thousands(n: u64) -> String {
|
||||
let s = n.to_string();
|
||||
let bytes = s.as_bytes();
|
||||
let mut out = String::with_capacity(s.len() + s.len() / 3);
|
||||
for (i, b) in bytes.iter().enumerate() {
|
||||
if i > 0 && (bytes.len() - i) % 3 == 0 {
|
||||
out.push(',');
|
||||
}
|
||||
out.push(*b as char);
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// 收集默认环境上下文（cwd / os / date / shell）。
|
||||
/// Gather the default environment context (cwd / os / date / shell).
|
||||
pub fn default_context() -> PromptContext {
|
||||
let cwd = std::env::current_dir().unwrap_or_else(|_| PathBuf::from("."));
|
||||
let date = today_iso();
|
||||
let shell = if cfg!(windows) {
|
||||
"powershell (fallback cmd.exe)"
|
||||
} else {
|
||||
"bash"
|
||||
};
|
||||
PromptContext {
|
||||
cwd,
|
||||
os: std::env::consts::OS,
|
||||
date,
|
||||
shell: shell.to_string(),
|
||||
context_usage: None,
|
||||
last_usage: None,
|
||||
extra: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 当前日期（`yyyy-mm-dd`），纯 std 实现（无 chrono）。
|
||||
/// The current date (`yyyy-mm-dd`), pure std (no chrono).
|
||||
pub fn today_iso() -> String {
|
||||
let days = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| (d.as_secs() / 86_400) as i64)
|
||||
.unwrap_or(0);
|
||||
let (y, m, d) = days_to_ymd(days);
|
||||
format!("{:04}-{:02}-{:02}", y, m, d)
|
||||
}
|
||||
|
||||
/// 把「自 1970-01-01 的天数」转换为 (年, 月, 日)。Howard Hinnant 算法。
|
||||
/// Convert "days since 1970-01-01" into (year, month, day). Howard Hinnant's
|
||||
/// algorithm.
|
||||
fn days_to_ymd(days: i64) -> (i64, u32, u32) {
|
||||
let z = days + 719_468;
|
||||
let era = if z >= 0 { z } else { z - 146_096 } / 146_097;
|
||||
let doe = (z - era * 146_097) as u64;
|
||||
let yoe = (doe - doe / 1_460 + doe / 36_524 - doe / 146_096) / 365;
|
||||
let y = yoe as i64 + era * 400;
|
||||
let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
|
||||
let mp = (5 * doy + 2) / 153;
|
||||
let d = (doy - (153 * mp + 2) / 5 + 1) as u32;
|
||||
let m = if mp < 10 { mp + 3 } else { mp - 9 } as u32;
|
||||
let y = if m <= 2 { y + 1 } else { y };
|
||||
(y, m, d)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn days_to_ymd_known_dates() {
|
||||
// 1970-01-01 → 0。
|
||||
assert_eq!(days_to_ymd(0), (1970, 1, 1));
|
||||
// 2000-01-01。
|
||||
// 30 年 * 365.25 ≈ 10957 天。
|
||||
assert_eq!(days_to_ymd(10_957), (2000, 1, 1));
|
||||
// 2024-02-29（闰年）。
|
||||
// 2024-02-29 (leap year).
|
||||
let days = (2024 - 1970) * 365 + (2024 - 1969) / 4 + 59;
|
||||
assert_eq!(days_to_ymd(days), (2024, 2, 29));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn renders_default_template() {
|
||||
let tpl = SystemPromptTemplate::default();
|
||||
let mut ctx = default_context();
|
||||
ctx.context_usage = Some(ContextUsage {
|
||||
estimated_tokens: 12_345,
|
||||
context_window: 128_000,
|
||||
window_known: true,
|
||||
});
|
||||
ctx.last_usage = Some(Usage {
|
||||
input_tokens: 100,
|
||||
output_tokens: 50,
|
||||
cache_read_tokens: 5,
|
||||
cache_write_tokens: 0,
|
||||
});
|
||||
ctx.extra.push("remember: rust edition 2021".to_string());
|
||||
let out = tpl.render(&ctx);
|
||||
assert!(out.contains(&ctx.cwd.display().to_string()));
|
||||
assert!(
|
||||
out.contains("Context usage: 12,345 / 128,000 tokens (9.6%)."),
|
||||
"got: {}",
|
||||
out
|
||||
);
|
||||
assert!(out.contains("Last turn usage: input 100, output 50, cache-read 5, cache-write 0."));
|
||||
assert!(out.contains("remember: rust edition 2021"));
|
||||
assert!(out.contains(ctx.os));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn renders_unknown_window_note() {
|
||||
let tpl = SystemPromptTemplate::default();
|
||||
let mut ctx = default_context();
|
||||
ctx.context_usage = Some(ContextUsage {
|
||||
estimated_tokens: 50_000,
|
||||
context_window: 128_000,
|
||||
window_known: false,
|
||||
});
|
||||
let out = tpl.render(&ctx);
|
||||
assert!(
|
||||
out.contains("window unknown; assuming 128,000"),
|
||||
"got: {}",
|
||||
out
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_fields_render_empty() {
|
||||
let tpl = SystemPromptTemplate::default();
|
||||
let ctx = PromptContext {
|
||||
cwd: PathBuf::from("/tmp"),
|
||||
os: "linux",
|
||||
date: "2025-01-01".into(),
|
||||
shell: "bash".into(),
|
||||
context_usage: None,
|
||||
last_usage: None,
|
||||
extra: Vec::new(),
|
||||
};
|
||||
let out = tpl.render(&ctx);
|
||||
assert!(!out.contains("{context_usage}"));
|
||||
assert!(!out.contains("{last_usage}"));
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,420 @@
|
|||
//! 会话树与 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)))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use focus_core::model::Message;
|
||||
|
||||
fn temp_dir(tag: &str) -> PathBuf {
|
||||
let d = std::env::temp_dir().join(format!(
|
||||
"focus-harness-session-{}-{}",
|
||||
tag,
|
||||
std::process::id()
|
||||
));
|
||||
let _ = std::fs::remove_dir_all(&d);
|
||||
std::fs::create_dir_all(&d).unwrap();
|
||||
d
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn appends_and_loads_entries() {
|
||||
let dir = temp_dir("append");
|
||||
let store = SessionStore::new(&dir);
|
||||
store
|
||||
.append(
|
||||
"s1",
|
||||
&SessionEntry::new("e1".into(), None, Message::user_text("hi")),
|
||||
)
|
||||
.unwrap();
|
||||
store
|
||||
.append(
|
||||
"s1",
|
||||
&SessionEntry::new("e2".into(), Some("e1".into()), Message::user_text("there")),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let entries = store.load("s1").unwrap();
|
||||
assert_eq!(entries.len(), 2);
|
||||
assert_eq!(entries[0].id, "e1");
|
||||
assert_eq!(entries[0].parent_id, None);
|
||||
assert_eq!(entries[1].id, "e2");
|
||||
assert_eq!(entries[1].parent_id.as_deref(), Some("e1"));
|
||||
assert!(matches!(entries[0].message, Message::User(_)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn next_entry_id_is_incremental() {
|
||||
let dir = temp_dir("seq");
|
||||
let store = SessionStore::new(&dir);
|
||||
assert_eq!(store.next_entry_id("s1").unwrap(), "e1");
|
||||
store
|
||||
.append(
|
||||
"s1",
|
||||
&SessionEntry::new("e1".into(), None, Message::user_text("a")),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(store.next_entry_id("s1").unwrap(), "e2");
|
||||
store
|
||||
.append(
|
||||
"s1",
|
||||
&SessionEntry::new("e3".into(), Some("e1".into()), Message::user_text("b")),
|
||||
)
|
||||
.unwrap();
|
||||
// e3 已存在 → 下一个是 e4。
|
||||
// e3 exists → the next is e4.
|
||||
assert_eq!(store.next_entry_id("s1").unwrap(), "e4");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lists_sessions() {
|
||||
let dir = temp_dir("list");
|
||||
let store = SessionStore::new(&dir);
|
||||
store
|
||||
.append(
|
||||
"alpha",
|
||||
&SessionEntry::new("e1".into(), None, Message::user_text("a")),
|
||||
)
|
||||
.unwrap();
|
||||
store
|
||||
.append(
|
||||
"beta",
|
||||
&SessionEntry::new("e1".into(), None, Message::user_text("b")),
|
||||
)
|
||||
.unwrap();
|
||||
let ids = store.list_sessions().unwrap();
|
||||
assert_eq!(ids, vec!["alpha".to_string(), "beta".to_string()]);
|
||||
}
|
||||
|
||||
fn tree() -> SessionTree {
|
||||
let entries = vec![
|
||||
SessionEntry::new("e1".into(), None, Message::user_text("root")),
|
||||
SessionEntry::new(
|
||||
"e2".into(),
|
||||
Some("e1".into()),
|
||||
Message::user_text("child A"),
|
||||
),
|
||||
SessionEntry::new(
|
||||
"e3".into(),
|
||||
Some("e1".into()),
|
||||
Message::user_text("child B"),
|
||||
),
|
||||
SessionEntry::new(
|
||||
"e4".into(),
|
||||
Some("e3".into()),
|
||||
Message::user_text("grandchild"),
|
||||
),
|
||||
];
|
||||
SessionTree::build(entries)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tree_queries() {
|
||||
let t = tree();
|
||||
assert_eq!(t.roots().len(), 1);
|
||||
assert_eq!(t.roots()[0].id, "e1");
|
||||
let children = t.children_of("e1");
|
||||
assert_eq!(children.len(), 2);
|
||||
let path = t.branch("e4");
|
||||
let ids: Vec<&str> = path.iter().map(|e| e.id.as_str()).collect();
|
||||
assert_eq!(ids, vec!["e1", "e3", "e4"]);
|
||||
let leaves = t.leaf_ids();
|
||||
assert_eq!(leaves.len(), 2);
|
||||
assert!(leaves.contains(&"e2"));
|
||||
assert!(leaves.contains(&"e4"));
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue