//! 上下文管理器 — 管理对话上下文和 token 预算 //! //! 职责: //! - 维护消息历史及其 token 计数缓存 //! - 提供预算感知的消息裁剪(保护工具调用三元组) //! - 为 run_agentic_loop 提供受控的消息视图 //! //! 裁剪策略与模型选择是正交维度:本模块只管「窗口多大、怎么裁」, //! 用哪个 model / 是否启用 reasoning 由调用方在 CompletionRequest 层决定。 //! //! 纯函数 / 数据类型 / 常量(TokenEstimator / ContextConfig / MessageGroup / //! TrackedMessage / EvictionUnit / classify_group / PROTECT_COUNT / //! TOOL_MISSING_PREFIX)已抽至 [`crate::context_helpers`],本模块 `use` 复用, //! 并 `pub use` 重导出以保持 `df_ai::context::*` 历史路径对外可见(零调用方变更)。 //! 本文件仅保留 `ContextManager` 结构体及其 `impl`(Rust impl 块不可跨文件)。 use crate::context_helpers::{classify_group, PROTECT_COUNT, TOOL_MISSING_PREFIX}; // 重导出:保持 `df_ai::context::TokenEstimator` / `df_ai::context::ContextConfig` 等 // 历史路径对外可见(agentic.rs / commands/ai/mod.rs 等调用方零变更)。 // `pub use` 同时把类型带入本模块命名空间,供 ContextManager 结构体字段与 impl 直接引用。 pub use crate::context_helpers::{ EvictionUnit, ContextConfig, MessageGroup, TokenEstimator, TrackedMessage, }; use crate::provider::{ChatMessage, MessageRole, ToolCall}; // ============================================================ // 上下文管理器 // ============================================================ /// 上下文管理器 /// /// 唯一的消息真相来源(替代原来的 `Vec`)。 /// 裁剪仅影响发送视图(`build_for_request`),不影响持久化(`all_messages_clone`)。 pub struct ContextManager { messages: Vec, /// 当前历史总 token 数(不含 system prompt) history_tokens: u32, config: ContextConfig, estimator: TokenEstimator, /// 压缩重入标志(F-15 §4.3):true 表示一次 LLM 压缩正在进行中。 /// agentic loop 顶部检测,防同一轮内多次触发压缩互相覆盖。纯内存态,不落库。 is_compressing: bool, } impl ContextManager { pub fn new(config: ContextConfig) -> Self { Self { messages: Vec::new(), history_tokens: 0, config, estimator: TokenEstimator::default(), is_compressing: false, } } /// 追加消息(自动计算 token 并更新缓存) /// /// 不在此处淘汰——push 可能发生在 agentic loop 中间(追加 tool_result), /// 此时不应裁剪正在使用的活跃消息。裁剪在 `build_for_request` 时统一处理。 pub fn push(&mut self, message: ChatMessage) { let tokens = self.estimator.estimate_message(&message); let group = classify_group(&message); // 仅 active 消息计入 token 预算(F-15 §3.3):truncated / archived_segment / // compressed 不进 LLM 上下文,token 虚高会致 build_for_request 误判超预算 // 触发不必要裁剪。!active 消息仍 push 到 self.messages 全量保留(持久化不受影响), // sanitize_messages step0(is_active 过滤)在发送视图统一剔除。 if message.is_active() { self.history_tokens += tokens; } self.messages.push(TrackedMessage { message, token_count: tokens, group, }); } /// 清空所有消息 pub fn clear(&mut self) { self.messages.clear(); self.history_tokens = 0; self.is_compressing = false; } /// 消息数量 pub fn len(&self) -> usize { self.messages.len() } pub fn is_empty(&self) -> bool { self.messages.is_empty() } /// 当前历史占用的 token 数(不含 system prompt) pub fn history_tokens(&self) -> u32 { self.history_tokens } /// 预算上限 pub fn budget_limit(&self) -> u32 { self.config.budget_limit() } // ── 核心:构建请求消息(受控裁剪版本)── /// 构建发送给 LLM 的消息列表 /// /// `sys_tokens` 为调用方已估算好的 system prompt token 数。 /// 超预算时自动裁剪旧消息(保护工具调用三元组 + 最近 PROTECT_COUNT 条)。 /// 返回 (消息列表, 是否发生了裁剪)。 pub fn build_for_request(&self, sys_tokens: u32) -> (Vec, bool) { let budget = self.budget_limit(); let available = budget.saturating_sub(sys_tokens); // system prompt 自身超预算:裁剪无法缓解(仍返回保护区兜底),warn 便于诊断 if sys_tokens > budget { tracing::warn!( "system prompt (~{} tokens) 超过上下文预算 ({}),裁剪无法缓解", sys_tokens, budget ); } // 未超预算 → 直接返回全量(仍做畸形配对自愈,防历史中毒触发 provider 500 死循环) if self.history_tokens <= available { return (Self::sanitize_messages(self.all_messages_clone()), false); } // 超预算 → 视图裁剪(不修改 self.messages,保证 all_messages_clone 仍返回全量) let protect_start = self.messages.len().saturating_sub(PROTECT_COUNT); let units = self.build_eviction_units(protect_start); let mut removed: u64 = 0; let mut trim_end = 0; for unit in &units { if self.history_tokens.saturating_sub(removed as u32) <= available { break; } removed += unit.token_sum as u64; trim_end = unit.end; } if trim_end == 0 { tracing::warn!( "history (~{} tokens) 超预算 ({}) 但无可淘汰单元(全在保护区 {} 条),发送兜底可能触发 provider 超限", self.history_tokens, available, PROTECT_COUNT ); return (self.all_messages_clone(), false); } let msgs: Vec = self.messages[trim_end..] .iter() .map(|t| t.message.clone()) .collect(); tracing::info!( "context_trimmed: skip {} messages, ~{} tokens (view-only, full history retained)", trim_end, removed ); (Self::sanitize_messages(msgs), true) } /// 畸形配对自愈 — 过滤掉会导致 provider 500 的中毒历史 /// /// 根因:Anthropic 流式 `tool_use` 块缺 id 时(anthropic_compat.rs 170-176) /// 生成 `tool_missing_{idx}` 占位 id;对应的 `tool_result` 一旦推入历史, /// 下次 `build_for_request` 把畸形三元组原样回传 → 服务端 500 → stream_recv /// 不清历史 → agentic 重发 → 永久卡死。此处仅在「发送视图」剔除畸形配对, /// 不改持久化(self.messages 全量保留),让卡死的会话能自愈继续。 /// /// provider 协议铁律:assistant 头里每个 tool_call.id 都必须在后续有对应的 /// tool_result,否则服务端 400/500。本函数按此自愈,分三档处理每个头: /// /// - 全闭合:所有 tool_call.id 都有匹配 tool_result → 原样保留。 /// - 全未闭合(含占位 id 必然无匹配):整头丢弃,其无主 tool_result 一并丢弃。 /// - 部分闭合:保留头但只留已闭合的 tool_call,丢弃未闭合的 tool_call 及其 /// 无主 tool_result(「不能整条删」——保住已发生的合法工具交互历史)。 /// /// 单次遍历、保序过滤、不重排(user/assistant/tool 角色交替不被打乱)。 fn sanitize_messages(messages: Vec) -> Vec { use std::collections::HashSet; // step 0(UX-09):过滤 truncated 软删消息,不进 LLM 上下文。 // 编辑某条 user 消息后其后续消息标 truncated(保留 DB 可追溯),发送视图必须剔除, // 否则被编辑前的旧回复仍进入 LLM 历史,污染重生成语义。落库全量保留不受影响。 let messages: Vec = messages .into_iter() .filter(|m| m.is_active()) .collect(); // step 1:已闭合的 tool_call_id 集合(role=Tool 消息全部提供过的 id) let resolved_ids: HashSet<&str> = messages .iter() .filter(|m| matches!(m.role, MessageRole::Tool)) .filter_map(|m| m.tool_call_id.as_deref()) .collect(); // step 2:判定每个 assistant 头的闭合状态,产出需丢弃的 id 集合 // orphaned_ids:未闭合的 tool_call.id(含占位 id),其 tool_result 要丢 // partial_heads:部分闭合的头——保留但需重写 tool_calls(只留已闭合的) // full_drop_heads:全未闭合的头——整条丢弃 let mut orphaned_ids: HashSet = HashSet::new(); let mut full_drop_heads = 0u32; let mut partial_heads = 0u32; // 记录部分闭合头中「应保留的 id」,用于 step3 精确重写 let mut partial_keep: std::collections::HashMap> = std::collections::HashMap::new(); for (i, m) in messages.iter().enumerate() { if !matches!(m.role, MessageRole::Assistant) { continue; } let Some(calls) = m.tool_calls.as_ref() else { continue }; if calls.is_empty() { continue; } let (resolved, unresolved): (Vec<&ToolCall>, Vec<&ToolCall>) = calls .iter() .partition(|c| { !c.id.starts_with(TOOL_MISSING_PREFIX) && resolved_ids.contains(c.id.as_str()) }); match (resolved.is_empty(), unresolved.is_empty()) { // 全未闭合 → 整头丢弃,未闭合 id 入 orphaned_ids (true, false) => { full_drop_heads += 1; for c in &unresolved { orphaned_ids.insert(c.id.clone()); } } // 全闭合 → 原样保留 (false, true) => {} // 部分闭合 → 保留头,只留 resolved 的 tool_call,unresolved 入 orphaned_ids (false, false) => { partial_heads += 1; for c in &unresolved { orphaned_ids.insert(c.id.clone()); } partial_keep.insert(i, resolved.into_iter().cloned().collect()); } // resolved/unresolved 都空不可能(calls 非空已保证) (true, true) => {} } } let after_triplet: Vec = if orphaned_ids.is_empty() && partial_keep.is_empty() { messages } else { // step 3:保序过滤 + 部分闭合头重写 tool_calls let sanitized: Vec = messages .into_iter() .enumerate() .filter_map(|(i, mut m)| match m.role { MessageRole::Assistant => { let Some(calls) = m.tool_calls.as_ref() else { return Some(m) }; if calls.is_empty() { return Some(m); } if let Some(keep) = partial_keep.get(&i) { // 部分闭合:重写 tool_calls 为仅已闭合子集 m.tool_calls = Some(keep.clone()); Some(m) } else { // 全闭合:保留;全未闭合(id 全在 orphaned_ids):丢弃 let all_orphaned = calls.iter().all(|c| orphaned_ids.contains(&c.id)); if all_orphaned { None } else { Some(m) } } } MessageRole::Tool => { // 无主 tool_result(其 id 命中 orphaned_ids)丢弃,其余保留 if m .tool_call_id .as_deref() .is_some_and(|id| orphaned_ids.contains(id)) { None } else { Some(m) } } _ => Some(m), }) .collect(); tracing::warn!( full_drop_heads, partial_rewrite_heads = partial_heads, orphaned_tool_results = orphaned_ids.len(), "history sanitized: dropped/rewrote malformed tool_call triplets (view-only, persisted history untouched)" ); sanitized }; // step 4:序列合法性修复(首条 user + 连续同 role 合并),防 Anthropic/GLM 1214。 Self::ensure_sequence_legal(after_triplet) } /// step 4:序列合法性修复(Anthropic/GLM Messages API 协议铁律,view-only 不改持久化)。 /// /// 协议要求首条必须是 user(system 由 convert_request 抽顶层)。sanitize step0(is_active 过滤) /// 与超预算裁剪(build_eviction_units 从三元组边界 trim)可能使首条变成 assistant/tool_result /// (开头 user 被裁/滤)→ 触发端点 1214「messages 参数非法」。 /// /// 修复:丢弃开头的 assistant/tool 消息(无前置 user 的孤儿,发也非法),直到首个 user/system。 /// 注:连续同 role(user/user、assistant/assistant)现实极少——裁剪按三元组原子保护不产生连续 user, /// archived/compressed 过滤后由摘要 system 占位——故本轮不合并(合并会破坏裁剪保护区语义 + 改变条数, /// 致 over_budget_trims_old 等测试失败)。若运行时日志显示连续 role 也是 1214 来源,再补合并。 fn ensure_sequence_legal(messages: Vec) -> Vec { let mut skipped = 0u32; let mut merged = 0u32; let mut fixed: Vec = Vec::with_capacity(messages.len()); for m in messages { // 首条必须 user:skip 开头 assistant/tool(无前置 user 的孤儿) if fixed.is_empty() && matches!(m.role, MessageRole::Assistant | MessageRole::Tool) { skipped += 1; continue; } // 连续同 role 合并(user content;assistant content+tool_calls;Tool 不合并—— // 连续 tool_result 由 anthropic_compat flush_tool_results 合并成 user blocks,此处合会丢 id) if let Some(last) = fixed.last_mut() { let same_role = std::mem::discriminant(&last.role) == std::mem::discriminant(&m.role); if same_role && matches!(m.role, MessageRole::User | MessageRole::Assistant) { if !m.content.is_empty() { if !last.content.is_empty() { last.content.push('\n'); } last.content.push_str(&m.content); } if matches!(m.role, MessageRole::Assistant) { if let Some(calls) = m.tool_calls { last.tool_calls.get_or_insert_with(Vec::new).extend(calls); } } merged += 1; continue; } } fixed.push(m); } if skipped > 0 || merged > 0 { tracing::warn!( skipped, merged, "序列修复:skip 开头非 user + 合并连续同 role(view-only,避免 Anthropic/GLM 1214)" ); } fixed } /// 全量克隆(持久化 save_conversation / build_for_request 未裁剪分支,不受裁剪影响) pub fn all_messages_clone(&self) -> Vec { self.messages.iter().map(|t| t.message.clone()).collect() } /// 从 Vec 恢复(兼容从 DB 加载) pub fn restore_from_messages(&mut self, messages: Vec) { self.clear(); for msg in messages { self.push(msg); } } /// 就地替换某条 tool_result 的内容(兼容审批 replace_tool_result) /// 返回 true 如果找到并替换了 /// /// 反向遍历:tool_result 由 append 进入历史,被替换的通常是最近的审批占位, /// 从尾部查找命中即停,避免对长历史做正向 O(n) 累积扫描。 pub fn replace_tool_result_content(&mut self, tool_call_id: &str, new_content: &str) -> bool { let pos = self.messages.iter().rposition(|t| { matches!(t.message.role, MessageRole::Tool) && t.message.tool_call_id.as_deref() == Some(tool_call_id) }); let Some(i) = pos else { return false }; // 先更新 content,再重估 token 并校正总量 let old_tokens = self.messages[i].token_count; self.messages[i].message.content = new_content.to_string(); let new_tokens = self.estimator.estimate_message(&self.messages[i].message); self.messages[i].token_count = new_tokens; self.history_tokens = self.history_tokens.saturating_sub(old_tokens).saturating_add(new_tokens); true } /// 弹出末尾连续的 assistant 消息(含其 tool_calls 三元组尾随 tool_result) /// /// 用于「重新生成」(UX-02):删掉最后一条 AI 回复(可能跨多轮 tool_calls + tool_results /// 紧随其后),保留触发它的 user 消息,以便重跑 agentic loop 再生成。 /// /// 语义:从末尾向前弹出,直到弹出至少一条 assistant 消息;若弹出 assistant 后紧邻的更早 /// 消息仍是 assistant/tool(同一轮多块),继续一并弹出,确保不留半截三元组污染下轮。 /// user 消息作为停止边界(不弹出 user),保证重生成时历史末尾是 user 消息。 pub fn pop_last_assistant_round(&mut self) -> bool { if self.messages.is_empty() { return false; } let mut popped_any = false; // 从尾向前:先弹掉末尾非 user 的消息(assistant / tool),直到遇到 user 或空 while let Some(last) = self.messages.last() { if matches!(last.message.role, MessageRole::User) { break; } let removed = self.messages.pop().expect("just checked non-empty"); self.history_tokens = self.history_tokens.saturating_sub(removed.token_count); if matches!(removed.message.role, MessageRole::Assistant) { popped_any = true; } } popped_any } /// 编辑某条 user 消息后,将其后所有消息标记为 truncated(UX-09 编辑重生成)。 /// /// 软删语义:保留在内存真相源 + DB(可追溯),但 sanitize_messages 过滤后不进 LLM 上下文, /// 前端按 is_active 过滤从视图移除。返回被标 truncated 的条数(0 表示该 user 消息已是末尾,无后续)。 /// /// `target_content` 为该 user 消息的预期内容(用于反向唯一定位:末条 user 消息可能内容相同, /// 故从尾部向前找第一条 role=User 且 content 匹配且仍 active 的消息)。 /// 找不到返回 Err(()),调用方据此报错。 pub fn truncate_after_user_message(&mut self, target_content: &str) -> Result { // 反向找末条 active user 消息且 content 匹配 let pos = self.messages.iter().rposition(|t| { matches!(t.message.role, MessageRole::User) && t.message.content == target_content && t.message.is_active() }); let Some(i) = pos else { return Err(()) }; // i 之后的全部标 truncated(已 truncated 的跳过,只统计本次新标的) let mut count = 0usize; for t in self.messages[i + 1..].iter_mut() { if t.message.is_active() { t.message.status = Some("truncated".to_string()); count += 1; } } Ok(count) } /// 替换末条 active user 消息的 content(UX-09 编辑重生成)。 /// /// 编辑语义:只能编辑最后一条 user 消息(中间编辑语义复杂,拒绝)。返回 Err(()) 表示无 active user 消息。 /// 成功后调用方应紧接着 truncate_after_user_message(new_content) 软删其后续消息。 pub fn replace_last_active_user_content(&mut self, new_content: &str) -> Result<(), ()> { let pos = self.messages.iter().rposition(|t| { matches!(t.message.role, MessageRole::User) && t.message.is_active() }); let Some(i) = pos else { return Err(()) }; let old_tokens = self.messages[i].token_count; self.messages[i].message.content = new_content.to_string(); let new_tokens = self.estimator.estimate_message(&self.messages[i].message); self.messages[i].token_count = new_tokens; self.history_tokens = self.history_tokens.saturating_sub(old_tokens).saturating_add(new_tokens); Ok(()) } /// 只读迭代(兼容 ensure_conversation_title 的 .iter().filter() 等) pub fn iter(&self) -> impl Iterator { self.messages.iter().map(|t| &t.message) } // ── F-15 上下文管理增强辅助方法(阶段1 基础,零行为变化)── // // 这些方法供阶段2 IPC(ai_chat_compress_context)/ 阶段3 agentic loop // 自动压缩调用。本阶段只暴露方法 + 单测,不接 IPC/前端/agentic.rs。 /// 配置(只读视图,供 agentic.rs 计算压缩触发阈值 `config().budget_limit()`) pub fn config(&self) -> &ContextConfig { &self.config } /// 可变消息切片(供阶段2 标记 status="compressed"/"archived_segment" + 调整 token) /// /// 调用方约定:仅改 `message.status` / `message.content`,不增删条目(增删走 /// [`push`] / [`insert_at`]),否则 `history_tokens` 会与实际脱钩。 pub fn messages_mut(&mut self) -> &mut [TrackedMessage] { &mut self.messages } /// 在给定位置插入一条消息(其余向后移),并把它计入 token 预算(active 才计)。 /// /// 供阶段2 在压缩点插入摘要 system 消息。`index` 越界则 panic(对齐 Vec::insert 语义, /// 调用方负责算合法 index,如 `compress_end` 已由 `compress_old_messages` 校验)。 pub fn insert_at(&mut self, index: usize, message: ChatMessage) { let tokens = self.estimator.estimate_message(&message); let group = classify_group(&message); if message.is_active() { self.history_tokens += tokens; } self.messages.insert(index, TrackedMessage { message, token_count: tokens, group, }); } /// 按淘汰单元分组消息范围(三元组原子性),供压缩定位/分段标记复用同一分组逻辑。 /// /// 返回每个单元的右开区间 end + token 总和,保证: /// - 工具调用三元组(Head + Tail* + 紧随的 Standalone Assistant)在同一单元 /// - 保护区 `[protect_start, len)` 内的消息不纳入任何单元 /// /// 公开供阶段2 会话分段(`archived_segment` 按组原子标记)与压缩定位共用。 pub fn build_eviction_units(&self, protect_start: usize) -> Vec { let mut units = Vec::new(); let mut i = 0usize; while i < protect_start { let mut token_sum = 0u32; if self.messages[i].group == MessageGroup::ToolCallHead { token_sum += self.messages[i].token_count; i += 1; while i < protect_start && self.messages[i].group == MessageGroup::ToolResultTail { token_sum += self.messages[i].token_count; i += 1; } if i < protect_start && self.messages[i].group == MessageGroup::Standalone && matches!(self.messages[i].message.role, MessageRole::Assistant) { token_sum += self.messages[i].token_count; i += 1; } } else { token_sum += self.messages[i].token_count; i += 1; } units.push(EvictionUnit { end: i, token_sum }); } units } /// 保护区外是否存在可压缩消息(供 agentic loop 顶部触发判断)。 /// /// "可压缩"= status 为 None/active 的消息(已 compressed/archived_segment/truncated /// 不参与二次压缩,幂等)。`protect_start` 为保护区起点(如 `len - PROTECT_COUNT`)。 pub fn has_compressible_messages(&self, protect_start: usize) -> bool { let end = protect_start.min(self.messages.len()); self.messages[..end] .iter() .any(|t| t.message.is_active()) } /// 把保护区 `[0, compress_end)` 范围内的 active 消息标记为 `status="compressed"`, /// 同步从 `history_tokens` 扣除其 token,返回被压缩消息的克隆(供阶段2 喂 LLM 摘要)。 /// /// **幂等**:已 compressed(或任何 !active)的消息跳过,不会被二次压缩;`history_tokens` /// 也只扣首次标记的 token。返回的 Vec 仅含**本次新标记**的消息(已 compressed 的不返)。 /// /// **单向不可逆**:压缩后 DB 原始消息保留,但 LLM 上下文里被 is_active 白名单隔离 /// (sanitize step0 过滤)。`compress_end` 越界自动 clamp 到 `messages.len()`。 /// /// 返回空 Vec 表示本批次无可压缩消息(全部已 compressed 或范围空),调用方据此跳过 LLM 调用。 pub fn compress_old_messages(&mut self, compress_end: usize) -> Vec { let end = compress_end.min(self.messages.len()); let mut newly_compressed = Vec::new(); for t in self.messages[..end].iter_mut() { if t.message.is_active() { newly_compressed.push(t.message.clone()); t.message.status = Some("compressed".to_string()); self.history_tokens = self.history_tokens.saturating_sub(t.token_count); } } newly_compressed } /// 压缩重入标志(读)。true 表示一次 LLM 压缩正在进行中,触发方应跳过本轮压缩。 pub fn is_compressing(&self) -> bool { self.is_compressing } /// 压缩重入标志(写)。`true`=开始压缩(进入 agentic loop 顶部前置置位), /// `false`=压缩结束(无论成功或降级)。调用方必须成对调用,防止永久卡死。 pub fn set_compressing(&mut self, v: bool) { self.is_compressing = v; } // ── 内部方法 ── // (原 build_eviction_units / classify_group 等纯函数已在上方公开或文件级定义; // ContextManager 的私有辅助如需新增放在这里。) } #[cfg(test)] mod tests { use super::*; fn cfg(max_tokens: u32) -> ContextConfig { ContextConfig { max_tokens, output_reserve: 0, safety_ratio: 1.0, } } #[test] fn sanitize_drops_placeholder_id_triplet() { // 中毒场景:assistant 头带 tool_missing_ 占位 id + 无主 tool_result // → build_for_request 必须剔除两者,否则下次回传触发 provider 500 死循环 let mut mgr = ContextManager::new(cfg(100_000)); mgr.push(ChatMessage::user("调用工具")); mgr.push(ChatMessage::assistant_with_tools( "调", vec![ToolCall::new("tool_missing_0", "read_file", "{}")], )); mgr.push(ChatMessage::tool_result("tool_missing_0", "结果")); let (msgs, _trimmed) = mgr.build_for_request(0); // 中毒的 head + 无主 tool_result 全部剔除,只剩首条 user assert_eq!(msgs.len(), 1, "占位 id 三元组应被 sanitize 剔除"); assert_eq!(msgs[0].content, "调用工具"); // 持久化全量保留(自愈只改发送视图) assert_eq!(mgr.all_messages_clone().len(), 3, "sanitize 不应污染内存全量"); } #[test] fn sanitize_drops_unresolved_tool_call_head() { // 真实 id 但全无 tool_result(流式中断/异常)→ 全未闭合档,整头丢弃 let mut mgr = ContextManager::new(cfg(100_000)); mgr.push(ChatMessage::assistant_with_tools( "调", vec![ToolCall::new("call_a", "fn_a", "{}"), ToolCall::new("call_b", "fn_b", "{}")], )); // 两个 call 都无 result → 整头丢弃 let (msgs, _trimmed) = mgr.build_for_request(0); assert!(msgs.iter().all(|m| !matches!(m.role, MessageRole::Assistant) || m.tool_calls.as_ref().is_none_or(|c| c.is_empty())), "全未闭合的 assistant 头应被剔除"); assert!(msgs.is_empty(), "全未闭合应整条删"); assert_eq!(mgr.all_messages_clone().len(), 1, "全量保留"); } #[test] fn sanitize_partial_resolve_rewrites_head() { // 多 tool_call,部分有匹配部分无 → 保留头但只留已闭合的,丢弃无主 result // provider 协议铁律:head 的每个 tool_call.id 必须都有 tool_result。 // 「不能整条删」——保住已发生的合法工具交互(call_b)历史。 let mut mgr = ContextManager::new(cfg(100_000)); mgr.push(ChatMessage::user("触发工具")); mgr.push(ChatMessage::assistant_with_tools( "调", vec![ToolCall::new("call_a", "fn_a", "{}"), ToolCall::new("call_b", "fn_b", "{}")], )); // 只有 call_b 有 result → call_a 未闭合 mgr.push(ChatMessage::tool_result("call_b", "b 结果")); let (msgs, _trimmed) = mgr.build_for_request(0); // 头被保留(不是整条删),但 tool_calls 重写为只剩已闭合的 call_b let assistant_heads: Vec<_> = msgs .iter() .filter(|m| matches!(m.role, MessageRole::Assistant) && m.tool_calls.is_some()) .collect(); assert_eq!(assistant_heads.len(), 1, "部分闭合头应保留(重写不删)"); let head_calls = assistant_heads[0].tool_calls.as_ref().unwrap(); assert_eq!(head_calls.len(), 1, "重写后只保留 1 个已闭合 tool_call"); assert_eq!(head_calls[0].id, "call_b", "保留的应是已闭合的 call_b"); assert!( head_calls.iter().all(|c| c.id != "call_a"), "未闭合的 call_a 应从重写后的 tool_calls 中移除" ); // tool_result 只剩 call_b 的,user 前置保留 let tool_msgs: Vec<_> = msgs .iter() .filter(|m| matches!(m.role, MessageRole::Tool)) .collect(); assert_eq!(tool_msgs.len(), 1, "只保留 call_b 的 tool_result"); assert_eq!(tool_msgs[0].tool_call_id.as_deref(), Some("call_b")); assert!( msgs.iter() .any(|m| matches!(m.role, MessageRole::User) && m.content == "触发工具"), "无关 user 消息不应被误删" ); assert_eq!(mgr.all_messages_clone().len(), 3, "sanitize 不应污染内存全量"); } #[test] fn sanitize_keeps_well_formed_triplet() { // 正常三元组:user → head(tool_call 有匹配 tool_result) → tool_result,不被剔除。 // 首条必须是 user(step4 Anthropic 协议修复),故前置一条 user 消息。 let mut mgr = ContextManager::new(cfg(100_000)); mgr.push(ChatMessage::user("问题")); mgr.push(ChatMessage::assistant_with_tools( "调", vec![ToolCall::new("call_a", "fn_a", "{}")], )); mgr.push(ChatMessage::tool_result("call_a", "结果")); let (msgs, _trimmed) = mgr.build_for_request(0); assert_eq!(msgs.len(), 3, "正常三元组不应被 sanitize 剔除"); } #[test] fn sanitize_mixed_multi_message_sequence() { // 场景 4:连续多条消息混合—— // [user] [全闭合三元组] [中毒占位 id 三元组] [部分闭合头] [user] // 预期:全闭合原样保留、中毒占位整删、部分闭合头重写为只剩已闭合、user 不动 let mut mgr = ContextManager::new(cfg(100_000)); mgr.push(ChatMessage::user("开场")); // 全闭合三元组 mgr.push(ChatMessage::assistant_with_tools( "调1", vec![ToolCall::new("ok_1", "fn", "{}")], )); mgr.push(ChatMessage::tool_result("ok_1", "ok 结果")); // 中毒占位 id 三元组(全未闭合)→ 整删 mgr.push(ChatMessage::assistant_with_tools( "调2", vec![ToolCall::new("tool_missing_0", "fn", "{}")], )); mgr.push(ChatMessage::tool_result("tool_missing_0", "占位结果")); // 部分闭合头:keep_1 闭合,drop_1 未闭合 → 重写为只剩 keep_1 mgr.push(ChatMessage::assistant_with_tools( "调3", vec![ToolCall::new("keep_1", "fn", "{}"), ToolCall::new("drop_1", "fn", "{}")], )); mgr.push(ChatMessage::tool_result("keep_1", "keep 结果")); mgr.push(ChatMessage::user("收尾")); let (msgs, _trimmed) = mgr.build_for_request(0); // 任何保留头里的 id 都不得含占位前缀、不得含未闭合的 drop_1 let head_ids: Vec> = msgs .iter() .filter(|m| matches!(m.role, MessageRole::Assistant) && m.tool_calls.is_some()) .map(|m| { m.tool_calls .as_ref() .unwrap() .iter() .map(|c| c.id.clone()) .collect() }) .collect(); assert!( head_ids.iter().all(|ids| { ids.iter().all(|id| !id.starts_with("tool_missing_") && id != "drop_1") }), "保留的 assistant 头中不应再有中毒 id 或未闭合 id, 实际 {:?}", head_ids ); let tool_ids: Vec = msgs .iter() .filter(|m| matches!(m.role, MessageRole::Tool)) .filter_map(|m| m.tool_call_id.clone()) .collect(); // ok_1(全闭合保留)、keep_1(部分闭合重写后保留)应在;占位/drop_1 不在 assert!(tool_ids.contains(&"ok_1".to_string()), "全闭合 ok_1 应保留"); assert!(tool_ids.contains(&"keep_1".to_string()), "部分闭合 keep_1 应保留"); assert!( tool_ids.iter().all(|id| !id.starts_with("tool_missing_") && id != "drop_1"), "无主 tool_result 应被剔除, 实际 {:?}", tool_ids ); // 两条 user 消息原样保留 let user_contents: Vec<&str> = msgs .iter() .filter(|m| matches!(m.role, MessageRole::User)) .map(|m| m.content.as_str()) .collect(); assert_eq!(user_contents, vec!["开场", "收尾"], "无关 user 消息不应被误删"); assert_eq!(mgr.all_messages_clone().len(), 8, "sanitize 不应污染内存全量"); } #[test] fn estimate_message_counts_parts_tokens() { // F-260614-05 多模态回归:含图消息的大段 base64 必须计入 token 预算, // 否则 history_tokens 严重低估 → build_for_request 不裁剪 → provider 超限。 let est = TokenEstimator::default(); // 纯文本基线 let text_msg = ChatMessage::user("短文本"); let text_tokens = est.estimate_message(&text_msg); // 同样 content + 含大段 base64 的 parts → token 应显著高于纯文本 let big_base64 = "iVBORw0KGgoAAAANS".repeat(100); // ~1.7k 字符 let multimodal = ChatMessage::user_parts( "短文本", vec![crate::provider::ContentPart::image_base64("image/png", big_base64.clone())], ); let mm_tokens = est.estimate_message(&multimodal); assert!( mm_tokens > text_tokens, "含图消息 token({}) 应高于纯文本({})", mm_tokens, text_tokens ); // base64 字符按 0.35 粗估,约 1.7k * 0.35 ≈ 595 tokens 量级 assert!( mm_tokens > 500, "大 base64 应贡献可观 token,实际 {}", mm_tokens ); // url 模式(无字节)也按 URL 长度估算,不爆 let url_msg = ChatMessage::user_parts( "t", vec![crate::provider::ContentPart::image_url("https://example.com/x.png")], ); let url_tokens = est.estimate_message(&url_msg); assert!(url_tokens > text_tokens, "url 片也应有少量 token 贡献"); } #[test] fn short_history_no_trim() { let mut mgr = ContextManager::new(cfg(100_000)); mgr.push(ChatMessage::user("你好")); mgr.push(ChatMessage::assistant("你好啊")); let (msgs, trimmed) = mgr.build_for_request(10); assert!(!trimmed); assert_eq!(msgs.len(), 2); } #[test] fn over_budget_trims_old() { // 小预算强制裁剪:20 条超预算,触发裁剪且保留保护区 let mut mgr = ContextManager::new(cfg(200)); // user/assistant 交替(真实对话序列;连续 user 会被 ensure_sequence_legal 合并,无法测条数裁剪) for i in 0..20 { if i % 2 == 0 { mgr.push(ChatMessage::user(&format!("这是第 {} 条较长的消息用于撑爆预算", i))); } else { mgr.push(ChatMessage::assistant(&format!("第 {} 条较长的回复用于撑爆预算", i))); } } let (msgs, trimmed) = mgr.build_for_request(0); assert!(trimmed, "超预算应触发裁剪"); assert!(msgs.len() < 20, "应裁掉部分旧消息, 实际 {}", msgs.len()); // 保护区:最新一条必保留(末条 i=19 是 assistant) assert_eq!( msgs.last().unwrap().content, "第 19 条较长的回复用于撑爆预算", "保护区最新消息被误裁" ); // 裁剪是视图:内存全量不变 assert_eq!(mgr.all_messages_clone().len(), 20, "裁剪污染了内存全量"); } #[test] fn tool_triplet_kept_atomic() { // 三元组不可分离:Head 与 Tail 同进同出,永不从中间切断 // 布局:6 旧(淘汰区) + 三元组(裁剪边界) + 6 新(保护区) = 15 条 let mut mgr = ContextManager::new(cfg(95)); for i in 0..6 { mgr.push(ChatMessage::user(&format!("旧消息 {}", i))); } mgr.push(ChatMessage::assistant_with_tools( "调工具", vec![ToolCall::new("tc1", "read_file", "{}")], )); mgr.push(ChatMessage::tool_result("tc1", "文件内容")); mgr.push(ChatMessage::assistant("完成")); for i in 0..6 { mgr.push(ChatMessage::user(&format!("新消息 {}", i))); } // 分支一:预算宽松,三元组整体保留 → Head 在则 Tail 在 let (msgs_keep, trimmed1) = mgr.build_for_request(0); assert!(trimmed1, "分支一应触发裁剪"); assert_eq!( has_head(&msgs_keep), has_tail(&msgs_keep), "分支一三元组被切断: head={} tail={}", has_head(&msgs_keep), has_tail(&msgs_keep) ); // 分支二:预算紧张,三元组整体丢弃 → Head 不在则 Tail 也不在 let (msgs_drop, trimmed2) = mgr.build_for_request(40); assert!(trimmed2, "分支二应触发裁剪"); assert_eq!( has_head(&msgs_drop), has_tail(&msgs_drop), "分支二三元组被切断: head={} tail={}", has_head(&msgs_drop), has_tail(&msgs_drop) ); // 裁剪是视图:两次 build 都不应改变内存全量 assert_eq!( mgr.all_messages_clone().len(), 15, "裁剪污染了内存全量" ); } fn has_head(msgs: &[ChatMessage]) -> bool { msgs.iter() .any(|m| matches!(m.role, MessageRole::Assistant) && m.tool_calls.is_some()) } fn has_tail(msgs: &[ChatMessage]) -> bool { msgs.iter().any(|m| matches!(m.role, MessageRole::Tool)) } #[test] fn replace_tool_result_updates_tokens() { let mut mgr = ContextManager::new(cfg(100_000)); mgr.push(ChatMessage::tool_result("tc1", "短")); let before = mgr.history_tokens(); assert!(mgr.replace_tool_result_content("tc1", "这是一个明显更长的替换内容用于验证 token 重估")); let after = mgr.history_tokens(); assert!(after > before); } #[test] fn restore_rebuilds_token_cache() { let mut mgr = ContextManager::new(cfg(100_000)); let src = vec![ ChatMessage::user("测试消息一"), ChatMessage::assistant("回复一"), ChatMessage::user("测试消息二"), ]; mgr.restore_from_messages(src); assert!(mgr.history_tokens() > 0); assert_eq!(mgr.len(), 3); } #[test] fn push_token_only_active() { // F-15 §3.3:!active 消息(truncated / archived_segment / compressed)仍 push // 到 self.messages(全量保留,持久化/前端视图自管),但不计入 history_tokens, // 避免 build_for_request 误判超预算触发不必要裁剪。 // 1) 直接 push 路径 let mut mgr = ContextManager::new(cfg(100_000)); let active_msg = ChatMessage::user("这条是 active 的"); let active_tokens = TokenEstimator::default().estimate_message(&active_msg); let mut inactive_msg = ChatMessage::assistant("这条被截断了不该计 token"); inactive_msg.status = Some("truncated".to_string()); let inactive_tokens = TokenEstimator::default().estimate_message(&inactive_msg); mgr.push(active_msg); mgr.push(inactive_msg); // 全量保留(两条都在内存) assert_eq!(mgr.len(), 2, "active + !active 都应 push 到 self.messages"); assert_eq!(mgr.all_messages_clone().len(), 2, "持久化全量不受 push 修正影响"); // token 预算只含 active assert_eq!( mgr.history_tokens(), active_tokens, "history_tokens 应只含 active,多算了 {}(inactive 应被忽略)", mgr.history_tokens().saturating_sub(active_tokens) ); assert!( inactive_tokens > 0, "前提:inactive 消息本身确有 token,否则无法证明它被排除" ); // 2) restore_from_messages 路径(调 push,token 同步仅 active) let mut mgr2 = ContextManager::new(cfg(100_000)); let mut a = ChatMessage::user("active 一"); a.status = Some("active".to_string()); let mut b = ChatMessage::user("archived 一"); b.status = Some("archived_segment".to_string()); let mut c = ChatMessage::user("compressed 一"); c.status = Some("compressed".to_string()); mgr2.restore_from_messages(vec![a, b, c]); assert_eq!(mgr2.len(), 3, "restore 后全量保留三条"); // 只 active 一条计 token(b/c 是白名单外状态,is_active 返回 false) let only_active_tokens = TokenEstimator::default() .estimate_message(&ChatMessage::user("active 一")); assert_eq!( mgr2.history_tokens(), only_active_tokens, "restore 后 history_tokens 应只含 active 一条,archived/compressed 不计" ); } #[test] fn empty_history_returns_empty() { let mgr = ContextManager::new(cfg(100_000)); let (msgs, trimmed) = mgr.build_for_request(10); assert!(!trimmed); assert!(msgs.is_empty(), "空历史应返回空列表"); } #[test] fn protect_zone_returns_full_when_untrimmable() { // 消息全在保护区(< PROTECT_COUNT 条)且超预算 → 无可淘汰单元,走 trim_end==0 兜底返回全量 let mut mgr = ContextManager::new(cfg(10)); mgr.push(ChatMessage::user("撑爆小预算的长消息内容")); mgr.push(ChatMessage::user("第二条撑爆预算的长消息")); let (msgs, trimmed) = mgr.build_for_request(0); assert!(!trimmed, "无可淘汰单元应返回 false(兜底)"); assert_eq!(msgs.len(), 2, "兜底返回全部保护区消息"); } #[test] fn system_over_budget_trims_to_protect_zone() { // system prompt 吃光预算 → history 仍尝试裁剪到保护区,不 panic let mut mgr = ContextManager::new(cfg(200)); for i in 0..10 { mgr.push(ChatMessage::user(&format!("消息 {} 撑量", i))); } let (msgs, _trimmed) = mgr.build_for_request(195); assert!( msgs.len() <= PROTECT_COUNT, "system 超预算时裁剪后至多保留保护区 {} 条,实际 {}", PROTECT_COUNT, msgs.len() ); } // ── F-15 阶段1 辅助方法单测 ── #[test] fn compress_old_messages_marks_compressed_and_returns_refs() { // F-15 §4.2/§4.3:compress_old_messages 把 [0, end) 内 active 消息标 compressed, // 同步扣 history_tokens,返回它们的克隆供 LLM 摘要。持久化全量保留。 let mut mgr = ContextManager::new(cfg(100_000)); mgr.push(ChatMessage::user("旧消息1")); mgr.push(ChatMessage::assistant("旧回复1")); mgr.push(ChatMessage::user("新消息2")); let tokens_before = mgr.history_tokens(); assert!(tokens_before > 0); let compressed = mgr.compress_old_messages(2); assert_eq!(compressed.len(), 2, "应压缩前 2 条 active"); assert_eq!(compressed[0].content, "旧消息1"); assert_eq!(compressed[1].content, "旧回复1"); // status 已改 compressed assert_eq!(mgr.messages_mut()[0].message.status.as_deref(), Some("compressed")); assert_eq!(mgr.messages_mut()[1].message.status.as_deref(), Some("compressed")); // 保护区外(本例 index 2)仍 active assert!(mgr.messages_mut()[2].message.is_active(), "保护区外消息不应被动"); // 持久化全量不变 assert_eq!(mgr.all_messages_clone().len(), 3, "compress 不应删消息(单向,全量保留)"); // token 已扣(剩第 3 条的) let only_third_tokens = TokenEstimator::default().estimate_message(&ChatMessage::user("新消息2")); assert_eq!(mgr.history_tokens(), only_third_tokens, "history_tokens 应扣除前两条"); } #[test] fn compress_old_messages_is_idempotent() { // 幂等:已 compressed 不二次压缩,二次调用返回空 Vec 且 history_tokens 不再变。 let mut mgr = ContextManager::new(cfg(100_000)); mgr.push(ChatMessage::user("a")); mgr.push(ChatMessage::user("b")); let first = mgr.compress_old_messages(2); assert_eq!(first.len(), 2); let tokens_after_first = mgr.history_tokens(); let second = mgr.compress_old_messages(2); assert!(second.is_empty(), "二次压缩应返回空(已 compressed 不重压)"); assert_eq!( mgr.history_tokens(), tokens_after_first, "二次压缩 history_tokens 不应再变(幂等)" ); } #[test] fn compress_old_messages_clamps_oversized_end() { // compress_end 越界自动 clamp 到 len,不 panic。 let mut mgr = ContextManager::new(cfg(100_000)); mgr.push(ChatMessage::user("唯一")); let compressed = mgr.compress_old_messages(999); assert_eq!(compressed.len(), 1, "越界 end 应 clamp 到 len(1)"); assert_eq!(mgr.history_tokens(), 0, "全量压缩后 history_tokens 归零"); } #[test] fn compress_old_messages_skips_already_inactive() { // 范围内含 truncated(已 !active)的消息:跳过,不返,不重复扣 token。 let mut mgr = ContextManager::new(cfg(100_000)); let mut truncated = ChatMessage::user("被截断"); truncated.status = Some("truncated".to_string()); mgr.push(truncated); mgr.push(ChatMessage::user("active 一条")); let tokens_before = mgr.history_tokens(); // truncated 已不计 token(见 push_token_only_active),所以 tokens_before 只含 active 一条 let compressed = mgr.compress_old_messages(2); assert_eq!(compressed.len(), 1, "只压缩 active 那条,truncated 跳过"); assert_eq!(mgr.history_tokens(), 0); assert_eq!( mgr.history_tokens(), tokens_before.saturating_sub(tokens_before), "幂等扣除:truncated 本就没计 token,active 扣光" ); // truncated 状态不被改成 compressed(保留原 truncated,语义不混淆) assert_eq!( mgr.messages_mut()[0].message.status.as_deref(), Some("truncated"), "已 truncated 不应被改写为 compressed" ); } #[test] fn has_compressible_messages_respects_protect_zone() { let mut mgr = ContextManager::new(cfg(100_000)); for i in 0..8 { mgr.push(ChatMessage::user(&format!("消息 {}", i))); } // protect_start=6 → [0,6) 内有 active → true assert!(mgr.has_compressible_messages(6)); // protect_start=0 → 空范围 → false assert!(!mgr.has_compressible_messages(0)); // 全部压缩后 → false mgr.compress_old_messages(6); assert!(!mgr.has_compressible_messages(6), "全 compressed 后不应有可压缩消息"); } #[test] fn is_compressing_flag_round_trip() { // 标志位读写 round-trip;clear() 复位。 let mut mgr = ContextManager::new(cfg(100_000)); assert!(!mgr.is_compressing(), "默认 false"); mgr.set_compressing(true); assert!(mgr.is_compressing(), "set true 后应读到 true"); mgr.set_compressing(false); assert!(!mgr.is_compressing(), "set false 后复位"); // clear 复位 mgr.set_compressing(true); mgr.clear(); assert!(!mgr.is_compressing(), "clear() 应复位 is_compressing"); } #[test] fn insert_at_adds_to_budget_when_active() { let mut mgr = ContextManager::new(cfg(100_000)); mgr.push(ChatMessage::user("a")); let tokens_before = mgr.history_tokens(); // 插入 active system 消息 → 计入 token mgr.insert_at(0, ChatMessage::system("## 摘要")); assert!(mgr.history_tokens() > tokens_before, "active 消息应计入 token"); assert_eq!(mgr.len(), 2); assert_eq!(mgr.messages_mut()[0].message.content, "## 摘要"); // 插入 !active 消息 → 不计入 token let tokens_before_inactive = mgr.history_tokens(); let mut inactive = ChatMessage::user("x"); inactive.status = Some("truncated".to_string()); mgr.insert_at(0, inactive); assert_eq!( mgr.history_tokens(), tokens_before_inactive, "!active 消息插入不应计 token" ); } #[test] fn build_eviction_units_keeps_triplet_atomic_public() { // 公开的 build_eviction_units:三元组(Head + Tail + Standalone Assistant)应落同一单元。 let mut mgr = ContextManager::new(cfg(100_000)); mgr.push(ChatMessage::user("前置")); mgr.push(ChatMessage::assistant_with_tools( "调", vec![ToolCall::new("c1", "fn", "{}")], )); mgr.push(ChatMessage::tool_result("c1", "结果")); mgr.push(ChatMessage::assistant("完成")); mgr.push(ChatMessage::user("后置")); // protect_start=5(全部纳入) let units = mgr.build_eviction_units(5); // 第一个单元是前置 Standalone(end=1);第二个单元应包含三元组三件套 + 后置应分开 // 确认三元组的 Head+Tail+Assistant 在同一单元(end 跳过 3) let unit2 = units.iter().find(|u| u.end >= 4).expect("应有跨三元组的单元"); assert!( unit2.end >= 4, "三元组三件套应在同一淘汰单元, end={}", unit2.end ); } }