438 lines
20 KiB
Rust
438 lines
20 KiB
Rust
//! 对话持久化 + Token 累加器
|
|
|
|
use std::sync::Arc;
|
|
|
|
use tokio::sync::Mutex;
|
|
|
|
use df_storage::crud::AiConversationRepo;
|
|
use df_storage::db::Database;
|
|
|
|
use crate::commands::now_millis;
|
|
|
|
use super::commands::message_to_record;
|
|
use super::AiSession;
|
|
|
|
/// Token 用量累加器(agent loop 生命周期内各轮叠加)
|
|
///
|
|
/// 纯结构 + 方法:抽自 run_agentic_loop 的 `total_prompt`/`total_completion` 双计数器,
|
|
/// 保证多轮累加、None 起始、跨 loop 实例叠加语义一致且可单测。
|
|
#[derive(Debug, Clone, Default)]
|
|
pub(crate) struct TokenAccumulator {
|
|
prompt: u32,
|
|
completion: u32,
|
|
}
|
|
|
|
impl TokenAccumulator {
|
|
/// 叠加一轮用量(round_usage 为本轮流式末 chunk 的累计用量)
|
|
///
|
|
/// saturating_add:恶意/异常 provider 返回巨大值时,避免 u32 += 溢出回绕打乱后续 budget 判定。
|
|
pub(crate) fn add(&mut self, prompt: u32, completion: u32) {
|
|
self.prompt = self.prompt.saturating_add(prompt);
|
|
self.completion = self.completion.saturating_add(completion);
|
|
}
|
|
|
|
pub(crate) fn prompt(&self) -> u32 {
|
|
self.prompt
|
|
}
|
|
|
|
pub(crate) fn completion(&self) -> u32 {
|
|
self.completion
|
|
}
|
|
|
|
pub(crate) fn total(&self) -> u32 {
|
|
self.prompt.saturating_add(self.completion)
|
|
}
|
|
}
|
|
|
|
/// 把单轮增量叠加到 DB 的 Option<i64> 字段(读旧值+增量,跨 loop 实例防覆盖)
|
|
///
|
|
/// 纯函数:抽自 save_conversation 的 token 累加逻辑,None 起始当作 0。
|
|
/// saturating_add:长期对话累积接近 i64::MAX 时不再翻负,封顶在 i64::MAX(统计语义安全,
|
|
/// 溢出回绕成负值会污染前端用量展示与计费/限额判定)。
|
|
pub(crate) fn accumulate_tokens(old: Option<i64>, add: u32) -> Option<i64> {
|
|
Some(old.unwrap_or(0).saturating_add(add as i64))
|
|
}
|
|
|
|
/// 持久化截断阈值:超过此长度的消息 content 落库前截断头尾各保 HEAD/TAIL 字符。
|
|
///
|
|
/// 防 read_file 1MB 洞 / list_directory 大体量结果落库后每轮重发累积致 token 暴增
|
|
/// (Sprint 19 实测单对话 in=115万 / 消息体 1.6MB)。仅作用于持久化视图,不污染内存真相源。
|
|
///
|
|
/// B-260619-02:50KB → 8KB 收紧。治 GLM 1214 messages 非法根因——GLM anthropic 端点对
|
|
/// 单条 tool_result content 限制更严(实测 ~10-20KB 即拒 1214),read_file 默认前 500 行源码
|
|
/// 约 11-21KB,旧 50KB 阈值未触发→DB 存全量→reload/下轮重发→GLM 1214 整请求失败。
|
|
/// 8KB 对齐 GLM 单条上限留余量;截断后 LLM 仍能据截断片段推理(优于 1214 直接拒收)。
|
|
pub(crate) const TRUNCATE_THRESHOLD: usize = 8_192;
|
|
const TRUNCATE_HEAD: usize = 3_072;
|
|
const TRUNCATE_TAIL: usize = 3_072;
|
|
|
|
/// 落库前对超长 content 做截断(保留头尾各 ~3KB + 中段标注省略字符数 + 总字节数)。
|
|
/// 8KB 阈值以下原样返回(零开销);按字符而非字节切避免 UTF-8 切坏中文。
|
|
/// 截断提示含原始字节数,让 LLM 明确感知非全文(防它误以为已看全)。
|
|
pub(crate) fn truncate_for_persist(content: &str) -> String {
|
|
let chars: Vec<char> = content.chars().collect();
|
|
if chars.len() <= TRUNCATE_THRESHOLD {
|
|
return content.to_string();
|
|
}
|
|
let head: String = chars.iter().take(TRUNCATE_HEAD).collect();
|
|
let tail: String = chars[chars.len() - TRUNCATE_TAIL..].iter().collect();
|
|
let omitted = chars.len() - TRUNCATE_HEAD - TRUNCATE_TAIL;
|
|
format!(
|
|
"{}\n\n[...省略 {} 字符(已截断,原 {} 字节,仅显示前 8KB,完整内容仅在内存态可读)...]\n\n{}",
|
|
head, omitted, content.len(), tail
|
|
)
|
|
}
|
|
|
|
/// 落库前对 `ChatMessage.parts` 做截断(F-260614-05 Phase 2a)。
|
|
///
|
|
/// - `Text` 片:复用 `truncate_for_persist` 头尾截断语义(单 Text 片 > 8KB 才截)。
|
|
/// - `Image` 片:base64 通常已 8KB+,落库前替换为占位 `Text` 片
|
|
/// `<image: base64 已省略, 共 N 字节>`,避免大体量图把对话 JSON 撑爆
|
|
/// (一张 1MB PNG 的 base64 ≈ 1.3MB 字符串)。原 Image 片仅在内存真相源(ContextManager)
|
|
/// 保留,重发时仍带图——对齐现有「持久化视图不污染内存真相源」约定。
|
|
/// - `Image` url 模式(无 base64):url 本身短,原样保留。
|
|
///
|
|
/// 返回 None 表示 parts 无需保留(全空或全部被替换且原 parts 仅含图);调用方据此清空 parts。
|
|
pub(crate) fn truncate_parts_for_persist(parts: &[df_ai::provider::ContentPart]) -> Option<Vec<df_ai::provider::ContentPart>> {
|
|
use df_ai::provider::ContentPart;
|
|
let mut out: Vec<ContentPart> = Vec::with_capacity(parts.len());
|
|
let mut changed = false;
|
|
for p in parts {
|
|
match p {
|
|
ContentPart::Text { text } => {
|
|
let truncated = truncate_for_persist(text);
|
|
if truncated.len() != text.len() {
|
|
changed = true;
|
|
}
|
|
out.push(ContentPart::Text { text: truncated });
|
|
}
|
|
ContentPart::Image { url, base64, media_type, alt } => {
|
|
// base64 模式:体量大,替换占位 Text 片
|
|
if let Some(b) = base64 {
|
|
let base64_len = b.len(); // 纯 base64 字符串长度（不含 data: 前缀）
|
|
let placeholder = format!("<image: base64 已省略, 共约 {} 字符>", base64_len);
|
|
out.push(ContentPart::Text { text: placeholder });
|
|
changed = true;
|
|
} else {
|
|
// url 模式:url 短,原样保留(含 media_type/alt)
|
|
out.push(ContentPart::Image {
|
|
url: url.clone(),
|
|
base64: None,
|
|
media_type: media_type.clone(),
|
|
alt: alt.clone(),
|
|
});
|
|
}
|
|
}
|
|
}
|
|
}
|
|
if out.is_empty() {
|
|
None
|
|
} else if changed {
|
|
Some(out)
|
|
} else {
|
|
// 无变化:返回克隆的原 parts(保持引用语义一致)
|
|
Some(parts.to_vec())
|
|
}
|
|
}
|
|
|
|
/// 保存对话到数据库（按 conv_id 写库，不受 active_conversation_id 切换影响）
|
|
///
|
|
/// 写 ai_messages(每条消息一行,F-260619-03 拆分存储)+ updated_at + 累加 token 用量 +
|
|
/// 首次落库的 model;标题由 ensure_conversation_title 单独生成。
|
|
/// token 走累加模式:upsert 读旧值叠加,保证审批暂停→恢复跨 loop 实例不覆盖丢失。
|
|
/// model 仅首次落库写入 + 旧记录缺值时补填(不覆盖历史已存值,兼容本次改造前的老对话)。
|
|
///
|
|
/// F-260619-03 批次 B:写路径切 ai_messages(全量重写,replace_conversation)。
|
|
/// 内存 ContextManager 是运行时真相源,save 是同步点——每轮全量回写 ai_messages,
|
|
/// compress/replace/edit/clear_context 内存改 status/content 后下轮 save 自动覆盖。
|
|
/// 旧 messages JSON 列**保留不写**(作备份,防 reload fallback 读旧脏数据),不赋新值也不置空。
|
|
pub(crate) async fn save_conversation(
|
|
session_arc: &Arc<Mutex<AiSession>>,
|
|
db: &Arc<Database>,
|
|
conv_id: &str,
|
|
usage: Option<&df_ai::provider::TokenUsage>,
|
|
model: Option<&str>,
|
|
touch_updated_at: bool,
|
|
) {
|
|
// 取 messages + 懒创建首次落库所需的 provider_id/created_at
|
|
// 工具结果(content)超 8KB(B-260619-02:50KB→8KB,治 GLM 1214)时截断头尾各 ~3KB + 中段
|
|
// 标注(含原字节数),防大体量结果(read_file 1MB 洞 / list_directory 13782 项)落库后每轮重发
|
|
// 累积致 token 暴增。仅影响持久化视图,不污染内存真相源(ContextManager)——build_for_request
|
|
// 仍读全量 messages。
|
|
//
|
|
// F-260616-09 B 批4:per_conv.messages 唯一真相源(conv_id 来源:本函数入参,与 save 落库的 conv 一致)。
|
|
// loop 内 save 由 run_agentic_loop 入参 conv_id 透传;IPC 路径(commands.rs)save 也传 conv_id。
|
|
// conv() 惰性建:save 路径 conv 必然已建(send/regenerate/edit/switch 均先 conv());若极端
|
|
// 未建(如启动恢复无 live conv),conv() 建空 PerConvState,save 空 messages(幂等不污染)。
|
|
let (persist_msgs, provider_id, created_at) = {
|
|
let mut session = session_arc.lock().await;
|
|
let mut msgs = session.conv(conv_id).messages.all_messages_clone();
|
|
for m in &mut msgs {
|
|
m.content = truncate_for_persist(&m.content);
|
|
// F-260614-05 Phase 2a: parts(Image base64) 同样截断(替换占位 Text 片),
|
|
// 防大体量图把对话 JSON 撑爆。仅作用于持久化副本,不污染内存真相源。
|
|
if let Some(parts) = m.parts.as_ref() {
|
|
m.parts = truncate_parts_for_persist(parts);
|
|
}
|
|
}
|
|
(
|
|
msgs,
|
|
session.active_provider_id.clone(),
|
|
session.active_conv_created_at.clone(),
|
|
)
|
|
};
|
|
|
|
// F-260619-03 批次 B:映射 Vec<ChatMessage> → Vec<AiMessageRecord>(带 seq 索引 + conv_id)
|
|
// 全量重写 ai_messages(单事务 DELETE + INSERT OR IGNORE,原子无中间空窗)。
|
|
// created_at 用对话级 created_at(老对话 None 时 now 兜底),保证消息创建时间与对话一致。
|
|
let now = now_millis();
|
|
let msg_created_at = created_at.clone().unwrap_or_else(|| now.clone());
|
|
let records: Vec<df_storage::models::AiMessageRecord> = persist_msgs
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(seq, m)| message_to_record(m, conv_id, seq as i64, &msg_created_at))
|
|
.collect();
|
|
|
|
let conv_repo = AiConversationRepo::new(db);
|
|
let msg_repo = df_storage::crud::AiMessageRepo::new(db);
|
|
match conv_repo.get_by_id(conv_id).await {
|
|
Ok(Some(mut rec)) => {
|
|
// 已落库:更新对话元数据(token/model/updated_at);messages JSON 列**不赋新值**(保留旧值作备份)。
|
|
// updated_at 按 touch_updated_at 条件改(用户活跃=true 反映最后活跃,
|
|
// 系统摘要/压缩=false 防会话时间分组"昨天→今天"跳变);
|
|
// token 累加(读旧值+新值,跨 loop 实例防覆盖)。
|
|
if touch_updated_at {
|
|
rec.updated_at = now_millis();
|
|
}
|
|
if let Some(u) = usage {
|
|
rec.prompt_tokens = accumulate_tokens(rec.prompt_tokens, u.prompt_tokens);
|
|
rec.completion_tokens = accumulate_tokens(rec.completion_tokens, u.completion_tokens);
|
|
}
|
|
// model: 旧记录缺值时补填(不覆盖已有);models: 去重追加用过的所有 model(JSON 数组)
|
|
if let Some(m) = model {
|
|
if rec.model.is_none() { rec.model = Some(m.to_string()); }
|
|
let mut list: Vec<String> = rec.models
|
|
.as_deref()
|
|
.and_then(|s| serde_json::from_str(s).ok())
|
|
.unwrap_or_default();
|
|
if !list.iter().any(|x| x == m) { list.push(m.to_string()); }
|
|
rec.models = Some(serde_json::to_string(&list).unwrap_or_else(|_| "[]".to_string()));
|
|
}
|
|
// update_full 仍写 messages 列(保留旧值,本批不改 messages 字段),写元数据 + updated_at
|
|
if let Err(e) = conv_repo.update_full(&rec).await {
|
|
tracing::warn!("更新对话元数据失败 {conv_id}: {e}");
|
|
}
|
|
// 消息拆分存储:全量重写 ai_messages
|
|
if let Err(e) = msg_repo.replace_conversation(conv_id, records).await {
|
|
tracing::warn!("全量重写 ai_messages 失败 {conv_id}: {e}");
|
|
}
|
|
}
|
|
Ok(None) => {
|
|
// 懒创建首次落库(此为空对话不落库的落库点:走到这里 messages 必非空)
|
|
// messages JSON 列首次落库也写(兼容未跑迁移的老库 fallback 读路径),
|
|
// 同时写 ai_messages(新读路径真相源)。
|
|
let messages_json = serde_json::to_string(&persist_msgs).unwrap_or_else(|_| "[]".to_string());
|
|
let conv_created = created_at.unwrap_or_else(|| now.clone());
|
|
let rec = df_storage::models::AiConversationRecord {
|
|
id: conv_id.to_string(),
|
|
title: None,
|
|
messages: messages_json,
|
|
provider_id,
|
|
model: model.map(|m| m.to_string()),
|
|
models: model.map(|m| serde_json::to_string(&[m]).unwrap_or_else(|_| "[]".to_string())),
|
|
archived: false,
|
|
pinned: false,
|
|
prompt_tokens: usage.map(|u| u.prompt_tokens as i64),
|
|
completion_tokens: usage.map(|u| u.completion_tokens as i64),
|
|
created_at: conv_created.clone(),
|
|
updated_at: now,
|
|
};
|
|
if let Err(e) = conv_repo.insert(rec).await {
|
|
tracing::warn!("落库对话失败 {conv_id}: {e}");
|
|
}
|
|
// 消息拆分存储:首次落库同样全量写 ai_messages
|
|
// (records 已用 conv_created 作 created_at,与对话记录一致)
|
|
if let Err(e) = msg_repo.replace_conversation(conv_id, records).await {
|
|
tracing::warn!("首次写 ai_messages 失败 {conv_id}: {e}");
|
|
}
|
|
}
|
|
Err(e) => tracing::warn!("读取对话 {conv_id} 失败: {e}"),
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
// ---------- TokenAccumulator + accumulate_tokens ----------
|
|
|
|
#[test]
|
|
fn accumulator_starts_zero() {
|
|
let acc = TokenAccumulator::default();
|
|
assert_eq!(acc.prompt(), 0);
|
|
assert_eq!(acc.completion(), 0);
|
|
assert_eq!(acc.total(), 0);
|
|
}
|
|
|
|
#[test]
|
|
fn accumulator_single_add() {
|
|
let mut acc = TokenAccumulator::default();
|
|
acc.add(100, 50);
|
|
assert_eq!(acc.prompt(), 100);
|
|
assert_eq!(acc.completion(), 50);
|
|
assert_eq!(acc.total(), 150);
|
|
}
|
|
|
|
#[test]
|
|
fn accumulator_multi_round_accumulation() {
|
|
// 多轮累加(模拟 agent loop 多次迭代)
|
|
let mut acc = TokenAccumulator::default();
|
|
acc.add(100, 20); // 轮1
|
|
acc.add(200, 40); // 轮2
|
|
acc.add(50, 10); // 轮3
|
|
assert_eq!(acc.prompt(), 350);
|
|
assert_eq!(acc.completion(), 70);
|
|
assert_eq!(acc.total(), 420);
|
|
}
|
|
|
|
#[test]
|
|
fn accumulator_add_zero_is_noop() {
|
|
let mut acc = TokenAccumulator::default();
|
|
acc.add(10, 5);
|
|
acc.add(0, 0);
|
|
assert_eq!(acc.total(), 15);
|
|
}
|
|
|
|
#[test]
|
|
fn accumulate_tokens_from_none() {
|
|
// 新记录(None 起始)落库
|
|
assert_eq!(accumulate_tokens(None, 100), Some(100));
|
|
assert_eq!(accumulate_tokens(None, 0), Some(0));
|
|
}
|
|
|
|
#[test]
|
|
fn accumulate_tokens_adds_to_existing() {
|
|
// 跨 loop 实例叠加:旧值 + 新增不覆盖
|
|
assert_eq!(accumulate_tokens(Some(500), 100), Some(600));
|
|
assert_eq!(accumulate_tokens(Some(0), 42), Some(42));
|
|
}
|
|
|
|
#[test]
|
|
fn accumulate_tokens_multi_round_db_simulation() {
|
|
// 模拟 save_conversation 多次落库累加(审批暂停→恢复跨 loop)
|
|
let mut field: Option<i64> = None;
|
|
field = accumulate_tokens(field, 100); // 首次
|
|
field = accumulate_tokens(field, 200); // 二次
|
|
field = accumulate_tokens(field, 50); // 三次
|
|
assert_eq!(field, Some(350));
|
|
}
|
|
|
|
#[test]
|
|
fn accumulator_and_db_accumulate_are_consistent() {
|
|
// loop 内 TokenAccumulator 与落库 accumulate_tokens 总量语义一致
|
|
let mut acc = TokenAccumulator::default();
|
|
let mut db_prompt: Option<i64> = None;
|
|
let mut db_completion: Option<i64> = None;
|
|
for (p, c) in [(100u32, 20u32), (200, 40), (50, 10)] {
|
|
acc.add(p, c);
|
|
db_prompt = accumulate_tokens(db_prompt, p);
|
|
db_completion = accumulate_tokens(db_completion, c);
|
|
}
|
|
assert_eq!(acc.prompt() as i64, db_prompt.unwrap());
|
|
assert_eq!(acc.completion() as i64, db_completion.unwrap());
|
|
}
|
|
|
|
// ---------- truncate_for_persist ----------
|
|
|
|
#[test]
|
|
fn truncate_short_content_unchanged() {
|
|
// 阈值以下原样返回
|
|
assert_eq!(truncate_for_persist("hello"), "hello");
|
|
assert_eq!(truncate_for_persist(""), "");
|
|
let near_limit: String = "a".repeat(TRUNCATE_THRESHOLD);
|
|
assert_eq!(truncate_for_persist(&near_limit).len(), TRUNCATE_THRESHOLD);
|
|
}
|
|
|
|
#[test]
|
|
fn truncate_long_content_keeps_head_and_tail() {
|
|
// 超阈值:保留头尾各 TRUNCATE_HEAD/TAIL 字符 + 中段标注
|
|
let long: String = "x".repeat(TRUNCATE_THRESHOLD + 1000);
|
|
let result = truncate_for_persist(&long);
|
|
// 头尾各 20k 字符应在结果中
|
|
assert!(result.starts_with(&"x".repeat(TRUNCATE_HEAD)));
|
|
assert!(result.ends_with(&"x".repeat(TRUNCATE_TAIL)));
|
|
// 中段标注存在 + 标注省略字符数(中段 = 总长 - 头 - 尾 = 9192 - 3072 - 3072 = 3048)
|
|
assert!(result.contains("已截断"));
|
|
assert!(result.contains("省略 3048 字符"));
|
|
// B-260619-02:截断提示含原字节数(9192 个 'x' = 9192 字节),让 LLM 感知非全文
|
|
assert!(result.contains("原 9192 字节"));
|
|
// 结果总长应远小于原长(20k 头 + 20k 尾 + 标注)
|
|
assert!(result.chars().count() < TRUNCATE_THRESHOLD + 1000);
|
|
}
|
|
|
|
#[test]
|
|
fn truncate_preserves_utf8_chinese() {
|
|
// 按字符切不切坏 UTF-8 中文
|
|
let chinese: String = "中".repeat(TRUNCATE_THRESHOLD + 500);
|
|
let result = truncate_for_persist(&chinese);
|
|
assert!(result.starts_with('中'));
|
|
assert!(result.ends_with('中'));
|
|
assert!(result.contains("已截断"));
|
|
}
|
|
|
|
// ---------- F-260614-05 Phase 2a truncate_parts_for_persist ----------
|
|
|
|
#[test]
|
|
fn truncate_parts_image_base64_replaced_with_placeholder() {
|
|
use df_ai::provider::ContentPart;
|
|
let parts = vec![
|
|
ContentPart::text("前缀文本"),
|
|
ContentPart::image_base64("image/png", &"A".repeat(60_000)),
|
|
ContentPart::text("后缀"),
|
|
];
|
|
let out = truncate_parts_for_persist(&parts).expect("应有结果");
|
|
// Image base64 被替换为占位 Text 片
|
|
assert!(out.iter().all(|p| !matches!(p, ContentPart::Image { base64: Some(_), .. })));
|
|
let placeholder = out.iter().find_map(|p| match p {
|
|
ContentPart::Text { text } if text.contains("base64 已省略") => Some(text.clone()),
|
|
_ => None,
|
|
}).expect("应含占位 Text 片");
|
|
assert!(placeholder.contains("60000"));
|
|
// 非 base64 的 Text 片保留原值
|
|
assert!(out.iter().any(|p| matches!(p, ContentPart::Text { text } if text == "前缀文本")));
|
|
assert!(out.iter().any(|p| matches!(p, ContentPart::Text { text } if text == "后缀")));
|
|
}
|
|
|
|
#[test]
|
|
fn truncate_parts_image_url_preserved() {
|
|
use df_ai::provider::ContentPart;
|
|
// url 模式:url 短,原样保留(含 media_type)
|
|
let parts = vec![ContentPart::Image {
|
|
url: Some("https://x/a.png".into()),
|
|
base64: None,
|
|
media_type: None,
|
|
alt: None,
|
|
}];
|
|
let out = truncate_parts_for_persist(&parts).expect("应有结果");
|
|
assert!(matches!(out[0], ContentPart::Image { ref url, .. } if url.as_deref() == Some("https://x/a.png")));
|
|
}
|
|
|
|
#[test]
|
|
fn truncate_parts_long_text_part_truncated() {
|
|
use df_ai::provider::ContentPart;
|
|
// 单 Text 片超阈值 → 头尾截断
|
|
let long: String = "x".repeat(TRUNCATE_THRESHOLD + 1000);
|
|
let parts = vec![ContentPart::Text { text: long }];
|
|
let out = truncate_parts_for_persist(&parts).expect("应有结果");
|
|
match &out[0] {
|
|
ContentPart::Text { text } => assert!(text.contains("已截断"), "超长 Text 片应被截断"),
|
|
_ => panic!("应仍是 Text 片"),
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn truncate_parts_empty_returns_none() {
|
|
assert_eq!(truncate_parts_for_persist(&[]), None);
|
|
}
|
|
}
|