绝尘
63fe0a6375
- 新建 df-ai/context_helpers.rs(195行): 纯函数/类型/常量(TokenEstimator/ContextConfig/MessageGroup/TrackedMessage/classify_group/EvictionUnit/PROTECT_COUNT/TOOL_MISSING_PREFIX) - context.rs 1332→1178行: 保留ContextManager struct+impl, pub use重导出(外部agentic.rs/mod.rs路径不变) - lib.rs: pub mod context_helpers Rust impl块约束: ContextManager方法未动一个字符, 纯函数搬迁零行为变更 主代兜底: cargo check --workspace 0 + test df-ai 119 + grep抽离项/pub use印证 strategy: 自底向上核心库优先, 单批1-2文件原子操作
1179 lines
52 KiB
Rust
1179 lines
52 KiB
Rust
//! 上下文管理器 — 管理对话上下文和 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<ChatMessage>`）。
|
||
/// 裁剪仅影响发送视图（`build_for_request`），不影响持久化（`all_messages_clone`）。
|
||
pub struct ContextManager {
|
||
messages: Vec<TrackedMessage>,
|
||
/// 当前历史总 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<ChatMessage>, 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<ChatMessage> = 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<ChatMessage>) -> Vec<ChatMessage> {
|
||
use std::collections::HashSet;
|
||
|
||
// step 0（UX-09）：过滤 truncated 软删消息，不进 LLM 上下文。
|
||
// 编辑某条 user 消息后其后续消息标 truncated（保留 DB 可追溯），发送视图必须剔除，
|
||
// 否则被编辑前的旧回复仍进入 LLM 历史，污染重生成语义。落库全量保留不受影响。
|
||
let messages: Vec<ChatMessage> = 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<String> = HashSet::new();
|
||
let mut full_drop_heads = 0u32;
|
||
let mut partial_heads = 0u32;
|
||
// 记录部分闭合头中「应保留的 id」，用于 step3 精确重写
|
||
let mut partial_keep: std::collections::HashMap<usize, Vec<ToolCall>> =
|
||
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<ChatMessage> = if orphaned_ids.is_empty() && partial_keep.is_empty() {
|
||
messages
|
||
} else {
|
||
// step 3：保序过滤 + 部分闭合头重写 tool_calls
|
||
let sanitized: Vec<ChatMessage> = 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<ChatMessage>) -> Vec<ChatMessage> {
|
||
let mut skipped = 0u32;
|
||
let mut merged = 0u32;
|
||
let mut fixed: Vec<ChatMessage> = 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<ChatMessage> {
|
||
self.messages.iter().map(|t| t.message.clone()).collect()
|
||
}
|
||
|
||
/// 从 Vec 恢复（兼容从 DB 加载）
|
||
pub fn restore_from_messages(&mut self, messages: Vec<ChatMessage>) {
|
||
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<usize, ()> {
|
||
// 反向找末条 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<Item = &ChatMessage> {
|
||
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<EvictionUnit> {
|
||
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<ChatMessage> {
|
||
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<Vec<String>> = 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<String> = 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
|
||
);
|
||
}
|
||
}
|