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:
parent
ee165235ec
commit
647cf1ba83
|
|
@ -9,4 +9,3 @@ description = "Session persistence (JSONL tree), context compaction, system prom
|
||||||
[dependencies]
|
[dependencies]
|
||||||
focus-core.workspace = true
|
focus-core.workspace = true
|
||||||
focus-json.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:会话持久化、上下文压缩与系统提示词模板。
|
||||||
//! focus-harness: session persistence, compaction, system prompts.
|
//! focus-harness: session persistence, context compaction, and system prompts.
|
||||||
//!
|
//!
|
||||||
//! 本里程碑阶段尚未实现。
|
//! 所有数据默认存放在用户目录下的 `~/.focus/`(Linux `$HOME/.focus`,
|
||||||
//! Not yet implemented in this milestone.
|
//! 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)]
|
#![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