优化: 消息落库增量写(Phase1:全量 DELETE+INSERT → 增量 INSERT + dirty 标记)

This commit is contained in:
lxy
2026-08-01 12:42:49 +08:00
parent 484080ac12
commit c34ad062e7
4 changed files with 390 additions and 11 deletions
+220
View File
@@ -717,3 +717,223 @@ fn topic_marker_take_is_idempotent_and_clears() {
// 第二次 take 必 None
assert!(mgr.take_topic_marker().is_none(), "marker 一次性消费, 二次 take 必空");
}
// ============================================================
// msg-split-phase1:增量写路径基线(persisted_msg_count)+ 全量重写标志(needs_full_rewrite)
// ============================================================
//
// 这些测试验证 save_conversation 增量优化的核心不变量:
// - 纯 push(append)不置 needs_full_rewrite,persisted_msg_count 由 save 推进(此处只测标志语义)
// - 任何修改既有消息 / 缩短 / 中段插入 / 清空的入口都置 needs_full_rewrite=true
// (漏置位 = 静默 DB 不同步,是本优化最危险的回归点)
// - restore_from_messages(DB 加载)后基线对齐到 len + 清 needs_full_rewrite
//
// persisted_msg_count 的推进/读取由 save_conversation(src-tauri) 调 advance_persisted_count 完成,
// 此处仅测 ContextManager 侧的字段初值 + advance 语义(max 防回退 + 清标志)。
#[test]
fn phase1_new_manager_zero_baseline_no_rewrite_needed() {
// 新建 ContextManager:基线 0,无需全量重写(append-only 起点)
let mgr = ContextManager::new(cfg(100_000));
assert_eq!(mgr.persisted_msg_count(), 0);
assert!(!mgr.needs_full_rewrite(), "新 manager 不应需全量重写");
}
#[test]
fn phase1_push_only_keeps_clean_no_dirty_flag() {
// 纯 append(push 多条)不置 needs_full_rewrite —— 这是常态(agentic loop 每轮追加),
// 必须走增量路径只插新消息,不能误触发全量重写。
let mut mgr = ContextManager::new(cfg(100_000));
mgr.push(ChatMessage::user("问题1"));
assert!(!mgr.needs_full_rewrite(), "push 不应置 dirty");
mgr.push(ChatMessage::assistant("回答1"));
mgr.push(ChatMessage::user("问题2"));
assert!(!mgr.needs_full_rewrite(), "连续 push 不应置 dirty");
assert_eq!(mgr.len(), 3);
}
#[test]
fn phase1_advance_persisted_count_sets_baseline_and_clears_dirty() {
// advance_persisted_count(new_count):推进基线 + 清 needs_full_rewrite(DB 已与内存一致)
let mut mgr = ContextManager::new(cfg(100_000));
mgr.push(ChatMessage::user("x"));
mgr.push(ChatMessage::user("y"));
// 模拟 save 成功后推进
mgr.advance_persisted_count(2);
assert_eq!(mgr.persisted_msg_count(), 2);
assert!(!mgr.needs_full_rewrite());
}
#[test]
fn phase1_advance_persisted_count_max_no_regression() {
// 防回退:advance 传更小值不回退基线(防误传 / 并发致基线丢失 → 下次增量漏写)
let mut mgr = ContextManager::new(cfg(100_000));
mgr.advance_persisted_count(5);
mgr.advance_persisted_count(3); // 误传更小
assert_eq!(mgr.persisted_msg_count(), 5, "基线不可回退");
}
#[test]
fn phase1_compress_old_messages_sets_dirty() {
// compress 改既有消息 status(active→compressed) → 必须全量重写收敛
let mut mgr = ContextManager::new(cfg(100_000));
mgr.push(ChatMessage::user("旧1"));
mgr.push(ChatMessage::user("旧2"));
mgr.push(ChatMessage::user("保护区"));
assert!(!mgr.needs_full_rewrite());
let compressed = mgr.compress_old_messages(2);
assert_eq!(compressed.len(), 2, "应压缩 2 条");
assert!(mgr.needs_full_rewrite(), "compress 后必须全量重写");
}
#[test]
fn phase1_compress_no_op_does_not_set_dirty() {
// 边界:compress_end 范围内无可压缩(全已 compressed)→ 状态未变,不应置 dirty
// (避免无意义全量重写,compress 重试场景常见)
let mut mgr = ContextManager::new(cfg(100_000));
mgr.push(ChatMessage::user("a"));
let _ = mgr.compress_old_messages(1);
assert!(mgr.needs_full_rewrite(), "首次压缩应置 dirty");
mgr.advance_persisted_count(1); // 模拟 save 后清 dirty
// 二次压缩同范围:全已 compressed,无新标 → 不应再置 dirty
let compressed = mgr.compress_old_messages(1);
assert!(compressed.is_empty());
// 注:此处 dirty 在二次 compress 时未变(未置位),但 advance 已清,故仍 false
assert!(!mgr.needs_full_rewrite(), "无新压缩不应置 dirty");
}
#[test]
fn phase1_replace_tool_result_content_sets_dirty() {
// 审批 replace_tool_result 改既有 tool 消息 content → 必须全量重写
let mut mgr = ContextManager::new(cfg(100_000));
mgr.push(ChatMessage::assistant_with_tools(
"调用工具",
vec![ToolCall {
id: "call_1".into(),
call_type: "function".into(),
function: crate::provider::ToolCallFunction {
name: "read_file".into(),
arguments: "{}".into(),
},
}],
));
mgr.push(ChatMessage::tool_result("call_1", "占位结果"));
assert!(!mgr.needs_full_rewrite());
let replaced = mgr.replace_tool_result_content("call_1", "真实结果");
assert!(replaced, "应找到并替换");
assert!(mgr.needs_full_rewrite(), "replace content 后必须全量重写");
}
#[test]
fn phase1_pop_last_assistant_round_sets_dirty() {
// 重新生成 pop 末尾 assistant 轮 → len 缩短 → 必须全量重写(删 DB 尾行)
let mut mgr = ContextManager::new(cfg(100_000));
mgr.push(ChatMessage::user(""));
mgr.push(ChatMessage::assistant(""));
assert!(!mgr.needs_full_rewrite());
let popped = mgr.pop_last_assistant_round();
assert!(popped);
assert!(mgr.needs_full_rewrite(), "pop 后必须全量重写");
assert_eq!(mgr.len(), 1, "应只剩 user");
}
#[test]
fn phase1_pop_no_op_does_not_set_dirty() {
// 边界:无可弹 assistant(末尾是 user)→ 不变,不应置 dirty
let mut mgr = ContextManager::new(cfg(100_000));
mgr.push(ChatMessage::user("只有 user"));
let popped = mgr.pop_last_assistant_round();
assert!(!popped);
assert!(!mgr.needs_full_rewrite(), "pop no-op 不应置 dirty");
}
#[test]
fn phase1_truncate_after_user_message_sets_dirty() {
// 编辑重生成 truncate 后续消息 status → 必须全量重写
let mut mgr = ContextManager::new(cfg(100_000));
mgr.push(ChatMessage::user("原始问题"));
mgr.push(ChatMessage::assistant("答1"));
mgr.push(ChatMessage::assistant("答2"));
assert!(!mgr.needs_full_rewrite());
let count = mgr.truncate_after_user_message("原始问题").expect("应找到");
assert_eq!(count, 2);
assert!(mgr.needs_full_rewrite(), "truncate 后必须全量重写");
}
#[test]
fn phase1_replace_last_active_user_content_sets_dirty() {
// 编辑 user 消息 content → 必须全量重写
let mut mgr = ContextManager::new(cfg(100_000));
mgr.push(ChatMessage::user("旧问题"));
assert!(!mgr.needs_full_rewrite());
mgr.replace_last_active_user_content("新问题").expect("应替换");
assert!(mgr.needs_full_rewrite(), "replace user content 后必须全量重写");
}
#[test]
fn phase1_insert_at_sets_dirty() {
// 中段插入(压缩摘要/话题标记)改变 seq → 必须全量重写
let mut mgr = ContextManager::new(cfg(100_000));
mgr.push(ChatMessage::user("a"));
mgr.push(ChatMessage::user("b"));
assert!(!mgr.needs_full_rewrite());
mgr.insert_at(0, ChatMessage::system("[摘要]"));
assert!(mgr.needs_full_rewrite(), "insert_at 后必须全量重写( seq 错位)");
assert_eq!(mgr.len(), 3);
}
#[test]
fn phase1_clear_resets_baseline_and_sets_dirty() {
// clear:基线重置 0 + 置 dirty(下次 save 全量重写空列表清 DB)
let mut mgr = ContextManager::new(cfg(100_000));
mgr.push(ChatMessage::user("a"));
mgr.advance_persisted_count(1);
assert_eq!(mgr.persisted_msg_count(), 1);
mgr.clear();
assert_eq!(mgr.persisted_msg_count(), 0, "clear 应重置基线到 0");
assert!(mgr.needs_full_rewrite(), "clear 应置 dirty");
assert_eq!(mgr.len(), 0);
}
#[test]
fn phase1_restore_from_messages_aligns_baseline_clears_dirty() {
// DB 加载(restore_from_messages)后:基线对齐到 len + 清 dirty
// (DB 已是该真相源,append 基线应对齐当前长度,不需全量重写)
let mut mgr = ContextManager::new(cfg(100_000));
mgr.push(ChatMessage::user("旧内存")); // 制造一个 dirty + 非零状态
mgr.mark_needs_full_rewrite();
// 模拟从 DB 加载 3 条消息
let loaded = vec![
ChatMessage::user("db1"),
ChatMessage::assistant("db2"),
ChatMessage::user("db3"),
];
mgr.restore_from_messages(loaded);
assert_eq!(mgr.persisted_msg_count(), 3, "基线应对齐到 DB 加载的消息数");
assert!(!mgr.needs_full_rewrite(), "DB 加载后不应需全量重写");
assert_eq!(mgr.len(), 3);
}
#[test]
fn phase1_mark_needs_full_rewrite_idempotent() {
// 幂等:重复 mark 不报警不改变语义(供 messages_mut 调用方多次调用兜底)
let mut mgr = ContextManager::new(cfg(100_000));
mgr.mark_needs_full_rewrite();
mgr.mark_needs_full_rewrite();
mgr.mark_needs_full_rewrite();
assert!(mgr.needs_full_rewrite());
}
#[test]
fn phase1_append_after_dirty_stays_dirty_until_save() {
// dirty 后再 push(append):dirty 保持(下次 save 仍走全量重写收敛,不会因 push 重置)
// 验证 push 不清 dirty(否则 compress 后追加新消息会漏收敛旧行)
let mut mgr = ContextManager::new(cfg(100_000));
mgr.push(ChatMessage::user("a"));
mgr.mark_needs_full_rewrite(); // 模拟 compress
mgr.push(ChatMessage::user("b")); // 追加
assert!(mgr.needs_full_rewrite(), "dirty 后 push 不应清 dirty");
// 模拟 save(全量重写)成功后推进 + 清
mgr.advance_persisted_count(2);
assert!(!mgr.needs_full_rewrite());
}
+91
View File
@@ -54,6 +54,30 @@ pub struct ContextManager {
/// 都非 None 且不同(双高置信),置位本字段,格式 "old|new"。agentic loop 顶部读并
/// 消费(insert 系统标记后清空)。纯内存态,不落库。保守:任一 topic 为 None 不置位(宁可漏报)。
pending_topic_marker: Option<String>,
/// msg-split-phase1:已落库(ai_messages 表)的消息条数。
///
/// 增量写路径用:save_conversation 据此判定 `[persisted_msg_count..len)` 是"自上次 save
/// 后新增的消息"(纯 append),走 `insert_batch`(INSERT OR IGNORE)只插新行,跳过昂贵的
/// `replace_conversation`(DELETE 整对话 + 全量重插)。
///
/// 仅 append-only 路径才有效——一旦 [`needs_full_rewrite`] 被置位(DB 与内存可能脱钩:
/// compress 改 status / replace_tool_result_content 改 content / pop_last_assistant_round
/// 缩短 / insert_at 中段插入改 seq / clear 清空),save 必须走全量重写收敛,本字段随之重置。
///
/// restore_from_messages(DB 加载)后设为 len(DB 已是该真相源,append 基线对齐到当前长度)。
/// 纯内存态,不落库(下次进程重启从 DB reload 时由 restore_from_messages 重置)。
persisted_msg_count: usize,
/// msg-split-phase1:全量重写需求标志。
///
/// `true` = 自上次 save 后发生过修改既有消息 / 缩短 / 中段插入 / 清空,内存与 DB 可能脱钩,
/// 下次 save 必须走 `replace_conversation` 全量重写收敛(DELETE + INSERT 全部行)。
/// save 全量重写成功后清零;append-only save(insert_batch)成功后保持 false。
///
/// 任何修改既有消息(非纯 push append)的入口置位:clear / pop_last_assistant_round /
/// truncate_after_user_message / replace_tool_result_content / replace_last_active_user_content /
/// insert_at / compress_old_messages / messages_mut(可变借用兜底,调用方可能改 status/content)。
/// 误置位零代价(只是下次 save 多走一次全量重写,正确性不变);漏置位才会丢更新——故保守置位。
needs_full_rewrite: bool,
}
impl ContextManager {
@@ -65,6 +89,8 @@ impl ContextManager {
estimator: TokenEstimator::default(),
is_compressing: false,
pending_topic_marker: None,
persisted_msg_count: 0,
needs_full_rewrite: false,
}
}
@@ -122,6 +148,11 @@ impl ContextManager {
self.history_tokens = 0;
self.is_compressing = false;
self.pending_topic_marker = None;
// msg-split-phase1:全清后 DB 必须也清(调用方 chat.rs:1149 走 delete_range),
// 全量重写基线重置到 0。needs_full_rewrite=true 保险(若 save 先于 delete_range 触发,
// 全量重写空列表也会清 DB)。
self.persisted_msg_count = 0;
self.needs_full_rewrite = true;
}
/// 消息数量
@@ -246,10 +277,16 @@ impl ContextManager {
/// 从 Vec 恢复(兼容从 DB 加载)
pub fn restore_from_messages(&mut self, messages: Vec<ChatMessage>) {
// clear() 会置 needs_full_rewrite=true + persisted_msg_count=0;但本入口是"DB 刚加载",
// 加载后的 messages 与 DB 完全一致(它们就是从 DB 来的),append 基线应对齐到当前 len,
// 且不需要全量重写(DB 已是真相源)。故 clear 后立即覆盖这两个字段。
let len = messages.len();
self.clear();
for msg in messages {
self.push(msg);
}
self.persisted_msg_count = len;
self.needs_full_rewrite = false;
}
/// 就地替换某条 tool_result 的内容(兼容审批 replace_tool_result
@@ -271,6 +308,8 @@ impl ContextManager {
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);
// msg-split-phase1:改既有消息 content,DB 需全量重写收敛(append-only 路径不覆盖旧行)。
self.needs_full_rewrite = true;
true
}
@@ -298,6 +337,11 @@ impl ContextManager {
popped_any = true;
}
}
// msg-split-phase1:从末尾弹出 → len 缩短,append-only 基线(persisted_msg_count)
// 会大于新 len。下次 save 必须全量重写(删 DB 中已不存在的尾行)。
if popped_any {
self.needs_full_rewrite = true;
}
popped_any
}
@@ -325,6 +369,10 @@ impl ContextManager {
count += 1;
}
}
// msg-split-phase1:改既有消息 status(truncated),DB 需全量重写收敛。
if count > 0 {
self.needs_full_rewrite = true;
}
Ok(count)
}
@@ -342,6 +390,8 @@ impl ContextManager {
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);
// msg-split-phase1:改既有 user 消息 content,DB 需全量重写收敛。
self.needs_full_rewrite = true;
Ok(())
}
@@ -421,10 +471,42 @@ impl ContextManager {
///
/// 调用方约定:仅改 `message.status` / `message.content`,不增删条目(增删走
/// [`push`] / [`insert_at`]),否则 `history_tokens` 会与实际脱钩。
///
/// msg-split-phase1:**若调用方借此句柄改了既有消息(status/content),必须紧接着调
/// [`mark_needs_full_rewrite`]**,否则下次 save 走 append-only 增量路径会漏更新旧行,
/// DB 与内存脱钩。仅读(过滤后 clone)不需调。本方法无法自行置位(借用冲突:返回 &mut
/// 切片时不能再持有 &mut self 标志)。
pub fn messages_mut(&mut self) -> &mut [TrackedMessage] {
&mut self.messages
}
/// msg-split-phase1:显式标记"下次 save 须全量重写"。
///
/// 供 [`messages_mut`] 的调用方在改完既有消息后调用(本 struct 无法在返回 &mut 切片时
/// 自行置位)。也可供任何绕过本 impl 直接改 messages 的路径兜底。幂等(重复置 true 无害)。
pub fn mark_needs_full_rewrite(&mut self) {
self.needs_full_rewrite = true;
}
/// msg-split-phase1:读取已落库消息条数(append-only 增量写路径的基线)。供 save_conversation
/// 判定 `[persisted_msg_count..len)` 是否为"自上次 save 后新增"。
pub fn persisted_msg_count(&self) -> usize {
self.persisted_msg_count
}
/// msg-split-phase1:推进已落库消息条数(insert_batch/replace_conversation 成功后调用)。
/// 同时清 needs_full_rewrite(DB 已与内存一致)。供 save_conversation 在写库成功后调用。
pub fn advance_persisted_count(&mut self, new_count: usize) {
// 防回退:若调用方误传更小值(理论上不应发生),取 max 保不丢基线(下次 save 仍能收敛)。
self.persisted_msg_count = self.persisted_msg_count.max(new_count);
self.needs_full_rewrite = false;
}
/// msg-split-phase1:是否需要全量重写(改过既有消息 / 缩短 / 中段插入 / 清空)。
pub fn needs_full_rewrite(&self) -> bool {
self.needs_full_rewrite
}
/// 在给定位置插入一条消息(其余向后移),并把它计入 token 预算(active 才计)。
///
/// 供压缩点插入摘要 system 消息。`index` 越界则 panic(对齐 Vec::insert 语义,
@@ -441,6 +523,9 @@ impl ContextManager {
group,
topic: None,
});
// msg-split-phase1:中段插入会改变 index 及之后所有消息的 seq(append-only 基线
// persisted_msg_count 按"末尾追加"语义计算 seq,中段插入后 seq 错位),DB 必须全量重写。
self.needs_full_rewrite = true;
}
/// 按淘汰单元分组消息范围(三元组原子性),供压缩定位/分段标记复用同一分组逻辑。
@@ -517,6 +602,12 @@ impl ContextManager {
self.history_tokens = self.history_tokens.saturating_sub(t.token_count);
}
}
// msg-split-phase1:改了既有消息 status(active → compressed),DB 需全量重写收敛
// (append-only 增量路径不覆盖旧行)。仅在确实标了新 compressed 时置位(无可压缩时
// newly_compressed 为空,状态未变,无需全量重写)。
if !newly_compressed.is_empty() {
self.needs_full_rewrite = true;
}
newly_compressed
}