diff --git a/crates/focus-harness/Cargo.toml b/crates/focus-harness/Cargo.toml index b763696..5b875fc 100644 --- a/crates/focus-harness/Cargo.toml +++ b/crates/focus-harness/Cargo.toml @@ -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"] } diff --git a/crates/focus-harness/src/compaction.rs b/crates/focus-harness/src/compaction.rs new file mode 100644 index 0000000..367088a --- /dev/null +++ b/crates/focus-harness/src/compaction.rs @@ -0,0 +1,357 @@ +//! 上下文压缩:token 估算与压缩方案(方案 D+E)。 +//! Context compaction: token estimation and plans (schemes D+E). +//! +//! - 方案 D:把最早的一批消息摘要化(`` + ``),保留最近 +//! 消息与关键 tool_result 链路,避免切断 tool_call → tool_result 配对; +//! - 方案 E:压缩阈值基于估算占用(默认窗口的 80%,可配置)。 +//! - Scheme D: summarize the oldest messages (`` + ``), +//! 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::>() + .join(" "), + Message::Assistant(a) => { + let mut parts: Vec = 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::>() + .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, + /// 原样保留的最近消息。 + /// Most recent messages kept verbatim. + pub keep: Vec, + /// 摘要模型的提示文本(要求产出 `` 与 ``)。 + /// Summarizer prompt (asks for `` and ``). + 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\ + \nA concise narrative summary of what happened.\n\ + \n\ + \nA bullet list of important facts, preferences and constraints.\n\ + \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 { + 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 { + 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 = (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 = (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("")); + assert!(plan.summary_instruction.contains("")); + // 保存的 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 = (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[..]); + } +} diff --git a/crates/focus-harness/src/lib.rs b/crates/focus-harness/src/lib.rs index 9e8b65c..e213fcb 100644 --- a/crates/focus-harness/src/lib.rs +++ b/crates/focus-harness/src/lib.rs @@ -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") +} diff --git a/crates/focus-harness/src/prompt.rs b/crates/focus-harness/src/prompt.rs new file mode 100644 index 0000000..d0fe19d --- /dev/null +++ b/crates/focus-harness/src/prompt.rs @@ -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, + /// 最近一次回合的用量。 + /// Usage of the most recent turn. + pub last_usage: Option, + /// 附加说明行(每行一条)。 + /// Extra note lines (one per entry). + pub extra: Vec, +} + +/// 系统提示词模板。 +/// 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) -> 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}")); + } +} diff --git a/crates/focus-harness/src/session.rs b/crates/focus-harness/src/session.rs new file mode 100644 index 0000000..f043cbb --- /dev/null +++ b/crates/focus-harness/src/session.rs @@ -0,0 +1,420 @@ +//! 会话树与 JSONL 持久化。 +//! Session trees and JSONL persistence. +//! +//! 参考 pi 的设计:每个会话是一个 JSONL 文件,追加写;每条 entry 有 +//! `id` + `parentId`,从而形成可分支的树结构。文件位于 `/sessions/`。 +//! Mirroring pi: each session is a JSONL file, append-only; every entry has an +//! `id` + `parentId`, forming a branchable tree. Files live in +//! `/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, + /// 节点承载的消息。 + /// 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, 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 { + 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) -> Self { + Self { + data_dir: data_dir.into(), + } + } + + /// 会话文件的目录(`/sessions`)。 + /// The sessions directory (`/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> { + 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> { + 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 { + let entries = self.load(session_id)?; + let max = entries + .iter() + .filter_map(|e| e.id.strip_prefix('e')) + .filter_map(|n| n.parse::().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, + by_id: HashMap, + children: HashMap>, + roots: Vec, +} + +impl SessionTree { + /// 从 entry 列表构建树。 + /// Build a tree from a list of entries. + pub fn build(entries: Vec) -> Self { + let mut by_id = HashMap::new(); + let mut children: HashMap> = 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")); + } +}