修复: token 汇总正确性(会话级双计/is_estimated跨端/流式usage传递) + is_estimated SELECT 缺失修复

This commit is contained in:
lxy
2026-08-09 23:59:55 +08:00
parent 11f4978ec1
commit 6bafdcd5a5
16 changed files with 189 additions and 127 deletions
+84 -46
View File
@@ -972,6 +972,11 @@ pub(crate) async fn run_agentic_loop(
// token 累加器:loop 生命周期内各轮叠加,退出时传 save_conversation(累加模式落库)
let mut tokens = TokenAccumulator::default();
// 会话级 token 落库口径:save_conversation 内部做 old+add,各 save 传「自上次 save 的增量」
// (usage_delta_since)。传累计快照会被同 loop 多次 save 重复累加(双计);增量口径同轮幂等,
// 审批暂停→恢复跨 loop 实例仍由 old+add 兜住(新实例 snapshot 归零)。
let mut saved_token_snapshot = df_ai::provider::TokenUsage::default();
// 收敛标志:仅当 LLM 末轮无 tool_calls 自行 break(正常收敛)时置 true;
// 区分"正常收敛退出"与"达 MAX 被截断退出"——后者末轮 tool_calls 仍非空(tool_result 不再回传 LLM),属异常
let mut converged = false;
@@ -1357,7 +1362,9 @@ pub(crate) async fn run_agentic_loop(
// 用户请求停止 或 已被新 loop 接管(stale)→ 收尾退出(已生成文本已在上一轮入库)。
// F1:stale 判定(epoch 不匹配)让旧 loop 在新 loop 启动后立即在此退出,不再继续跑。
if stop_flag.load(Ordering::SeqCst) || loop_epoch_arc.load(Ordering::SeqCst) != my_epoch {
let usage = df_ai::provider::TokenUsage {
// save_usage 增量(自上次 save 后新增,首轮即停为 0);emit_usage 累计快照
let save_usage = usage_delta_since(&tokens, &mut saved_token_snapshot);
let emit_usage = df_ai::provider::TokenUsage {
prompt_tokens: tokens.prompt(),
completion_tokens: tokens.completion(),
total_tokens: tokens.total(),
@@ -1370,11 +1377,11 @@ pub(crate) async fn run_agentic_loop(
// (stale 时 finish_round_exit 内部仅 disarm 跳过 save/reset/emit,不干扰新 loop)
finish_round_exit(
&session_arc, &db, &conv_id,
Some(&usage), None,
Some(&save_usage), None,
true,
&provider_config, &llm_concurrency,
&mut guard,
&usage,
&emit_usage,
// 入口 stop:本轮尚未 stream,last_round_estimated 仍为初始 false(无估算)。
last_round_estimated,
None, None, true,
@@ -1595,6 +1602,8 @@ pub(crate) async fn run_agentic_loop(
// 预估输入 token(兜底:部分 provider 如 GLM 流式 usage 不报 prompt_tokens,后段用它补)
// 注:stream_one_provider 内每次重试重建 request(因 provider.stream 消费 body),
// 此处不再预构建 request(旧 request 变量已废弃),仅保留 messages 供 estimated_prompt。
// 语义说明:每轮对全量历史重估并累加(GLM 类无 usage provider 多轮 prompt 总和偏高),
// 但「多轮累计」本就接近真实总量且标估算展示,保留现状(增量估算收益低,不动)。
let estimated_prompt: u32 = {
let est = TokenEstimator::default();
messages.iter().map(|m| est.estimate_message(m)).sum()
@@ -1802,7 +1811,8 @@ pub(crate) async fn run_agentic_loop(
// G4.3:本轮 token 用量是否估算值(provider 未报 prompt_tokens → estimated_prompt 兜底),
// 供 loop 内/loop 后各退出路径透传 AiCompleted(is_estimated)仅作展示标注。
last_round_estimated = round_usage.prompt_tokens == 0;
// 语义修正:整 loop 只要任一轮估算即标估算(累计总量含估算成分),非仅末轮。
last_round_estimated |= round_usage.prompt_tokens == 0;
// F1 并发 epoch:流返回后若已被新 loop 接管(force_send 等),旧 loop 不再 push 消息 /
// 保文 / 执行工具,立即退出(guard disarm 跳过复位,防 clobber 新 loop 的 Generating)。
@@ -1821,7 +1831,7 @@ pub(crate) async fn run_agentic_loop(
let usage = df_ai::provider::TokenUsage {
prompt_tokens: if round_usage.prompt_tokens == 0 { estimated_prompt } else { round_usage.prompt_tokens },
completion_tokens: round_usage.completion_tokens,
total_tokens: if round_usage.prompt_tokens == 0 { estimated_prompt + round_usage.completion_tokens } else { round_usage.total_tokens },
total_tokens: if round_usage.prompt_tokens == 0 { estimated_prompt.saturating_add(round_usage.completion_tokens) } else { round_usage.total_tokens },
// 分项 token(2026-08-02):cache/reasoning 透传自 round_usage,落库 + 累加器都需
prompt_cache_hit_tokens: round_usage.prompt_cache_hit_tokens,
prompt_cache_miss_tokens: round_usage.prompt_cache_miss_tokens,
@@ -1858,6 +1868,8 @@ pub(crate) async fn run_agentic_loop(
msg.prompt_cache_hit_tokens = Some(usage.prompt_cache_hit_tokens);
msg.prompt_cache_miss_tokens = Some(usage.prompt_cache_miss_tokens);
msg.reasoning_tokens = Some(usage.reasoning_tokens);
// 消息级估算标记:本轮 prompt 是否 estimated 兜底,reload 逐条回显对齐 live 态
msg.is_estimated = Some(round_usage.prompt_tokens == 0);
conv.messages.push(msg);
// 追加系统提示消息:响应因网络中断不完整(对齐决策 a1 系统提示机制)
let mut notice = ChatMessage::system("⚠ 响应因网络中断不完整,以上为已接收的部分内容。可重新发送以获取完整回复。");
@@ -1868,25 +1880,26 @@ pub(crate) async fn run_agentic_loop(
// 统一走 finish_round_exit 收尾(save + spawn_title + reset + emit)。
// 注意:partial 文本+系统提示已先 push(上方 block),此 save 落库含本轮 partial,幂等覆盖。
// emit_usage 用 tokens 快照(tokens.add 已累加本轮):total_tokens/prompt/completion 对齐原 emit 三元组
// save_usage 用增量(tokens 已累加本轮,减上次快照);emit_usage 用 tokens 快照累计
// MidStream 分叉:emit_incomplete=Some(true)(前端标不完整),publish_incomplete=None(总线消费方),
// do_publish=true(publish 走总线)。spawn_title=true(后台标题,失败 extract 兜底)。
let save_usage = usage_delta_since(&tokens, &mut saved_token_snapshot);
let emit_usage = df_ai::provider::TokenUsage {
prompt_tokens: tokens.prompt(),
completion_tokens: tokens.completion(),
total_tokens: usage.total_tokens,
total_tokens: tokens.total(),
prompt_cache_hit_tokens: tokens.cache_hit(),
prompt_cache_miss_tokens: tokens.cache_miss(),
reasoning_tokens: tokens.reasoning(),
};
finish_round_exit(
&session_arc, &db, &conv_id,
Some(&usage), Some(&resolved_model),
Some(&save_usage), Some(&resolved_model),
true,
&provider_config, &llm_concurrency,
&mut guard,
&emit_usage,
round_usage.prompt_tokens == 0,
last_round_estimated,
Some(true), None, true,
&pinned_goals_snapshot,
&app_handle,
@@ -1933,6 +1946,7 @@ pub(crate) async fn run_agentic_loop(
round_prompt, round_usage.completion_tokens,
round_usage.prompt_cache_hit_tokens, round_usage.prompt_cache_miss_tokens,
round_usage.reasoning_tokens,
round_usage.prompt_tokens == 0,
);
if GOAL_PIN_ENABLED {
update_pinned_goals(&mut session, &conv_id, &tool_calls_acc);
@@ -1944,21 +1958,17 @@ pub(crate) async fn run_agentic_loop(
// 或 session lock 竞争)永远到不了出口,用户重启后上轮回复丢失。此处出 push 锁作用域后
// 立即 save,幂等(每轮重复覆盖落库),即使后续工具卡住本轮消息已持久化。
{
let usage = df_ai::provider::TokenUsage {
prompt_tokens: tokens.prompt(),
completion_tokens: tokens.completion(),
total_tokens: tokens.total(),
prompt_cache_hit_tokens: tokens.cache_hit(),
prompt_cache_miss_tokens: tokens.cache_miss(),
reasoning_tokens: tokens.reasoning(),
};
// 落库传「自上次 save 的增量」而非累计快照(累计会被同 loop 多次 save 重复累加双计)
let usage = usage_delta_since(&tokens, &mut saved_token_snapshot);
save_conversation(&session_arc, &db, &conv_id, Some(&usage), Some(&resolved_model), true).await;
}
// 停止信号 或 已被新 loop 接管(stale):已生成文本入库后退出,不再执行后续工具调用。
// F1:stale 时旧 loop 在此退出,不执行工具(防重复工具执行);finish_round_exit 内部跳过 emit。
if stop_flag.load(Ordering::SeqCst) || loop_epoch_arc.load(Ordering::SeqCst) != my_epoch {
let usage = df_ai::provider::TokenUsage {
// save_usage 增量(自上次 save 后新增);emit_usage 累计快照(前端展示本 loop 总量)
let save_usage = usage_delta_since(&tokens, &mut saved_token_snapshot);
let emit_usage = df_ai::provider::TokenUsage {
prompt_tokens: tokens.prompt(),
completion_tokens: tokens.completion(),
total_tokens: tokens.total(),
@@ -1969,12 +1979,12 @@ pub(crate) async fn run_agentic_loop(
// 统一走 finish_round_exit:save(Some usage, Some model) + spawn_title + emit(None,None,publish=true)
finish_round_exit(
&session_arc, &db, &conv_id,
Some(&usage), Some(&resolved_model),
Some(&save_usage), Some(&resolved_model),
true,
&provider_config, &llm_concurrency,
&mut guard,
&usage,
round_usage.prompt_tokens == 0,
&emit_usage,
last_round_estimated,
None, None, true,
&pinned_goals_snapshot,
&app_handle,
@@ -2057,14 +2067,8 @@ pub(crate) async fn run_agentic_loop(
guard.disarm();
return;
}
let usage = df_ai::provider::TokenUsage {
prompt_tokens: tokens.prompt(),
completion_tokens: tokens.completion(),
total_tokens: tokens.total(),
prompt_cache_hit_tokens: tokens.cache_hit(),
prompt_cache_miss_tokens: tokens.cache_miss(),
reasoning_tokens: tokens.reasoning(),
};
// 落库传「自上次 save 的增量」;审批暂停→恢复跨 loop 实例由 save_conversation old+add 兜住
let usage = usage_delta_since(&tokens, &mut saved_token_snapshot);
save_conversation(&session_arc, &db, &conv_id, Some(&usage), Some(&resolved_model), true).await;
// 审批等待 return 前 disarm guard——保持 generating=true 留 try_continue 续生成,
// 同时 Drop 因 done=true 跳过复位 spawn(避免误复位审批态 generating 致 ai_approve→try_continue 不续)
@@ -2088,7 +2092,9 @@ pub(crate) async fn run_agentic_loop(
max_iter = max_iterations,
"[ai] agentic 循环达最大轮次仍未收敛,自动完成(incomplete=true)",
);
let usage = df_ai::provider::TokenUsage {
// save_usage 增量;emit_usage 累计快照(前端展示本 loop 总量)
let save_usage = usage_delta_since(&tokens, &mut saved_token_snapshot);
let emit_usage = df_ai::provider::TokenUsage {
prompt_tokens: tokens.prompt(),
completion_tokens: tokens.completion(),
total_tokens: tokens.total(),
@@ -2101,11 +2107,11 @@ pub(crate) async fn run_agentic_loop(
// (与其他 5 路径不一致是历史现状,本次仅收敛重复代码不改 publish 策略,语义零变更)。
finish_round_exit(
&session_arc, &db, &conv_id,
Some(&usage), Some(&resolved_model),
Some(&save_usage), Some(&resolved_model),
false,
&provider_config, &llm_concurrency,
&mut guard,
&usage,
&emit_usage,
last_round_estimated,
Some(true), None, false,
&pinned_goals_snapshot,
@@ -2122,19 +2128,10 @@ pub(crate) async fn run_agentic_loop(
guard.disarm();
return;
}
let usage = df_ai::provider::TokenUsage {
prompt_tokens: tokens.prompt(),
completion_tokens: tokens.completion(),
total_tokens: tokens.total(),
prompt_cache_hit_tokens: tokens.cache_hit(),
prompt_cache_miss_tokens: tokens.cache_miss(),
reasoning_tokens: tokens.reasoning(),
};
// 落库 + 标题 + 知识提炼打包后台化:不阻塞 generating 复位与 Completed 事件
// save 先行(extract/title 都读已落库消息);extract 内部 fire-and-forget,与 title 可能并发
// (均受 per_conv 信号量约束,读写不同字段互不干扰)
// 并发取舍:与新对话新 loop 的 save 存在低概率并发 upsert,最多丢少量 token 累加(非功能错误,可接受)
let usage_total = usage.total_tokens;
// save_usage 增量(自上次 save 后新增,正常收敛时上轮 save 已落,此处通常为 0);
// emit_usage 用 tokens 累计快照。normal_usage.total 直接取 tokens.total() 保持三字段一致。
let save_usage = usage_delta_since(&tokens, &mut saved_token_snapshot);
let usage_total = tokens.total();
{
let session_arc = session_arc.clone();
let db = db.clone();
@@ -2145,7 +2142,7 @@ pub(crate) async fn run_agentic_loop(
let llm_concurrency = llm_concurrency.clone();
let resolved_model = resolved_model.clone();
tauri::async_runtime::spawn(async move {
save_conversation(&session_arc, &db, &conv_id, Some(&usage), Some(&resolved_model), true).await;
save_conversation(&session_arc, &db, &conv_id, Some(&save_usage), Some(&resolved_model), true).await;
// 知识提炼:需读已落库的对话消息,故在 save 之后
if let Err(e) = maybe_spawn_extraction(&session_arc, &db, &conv_id, &provider_config, &knowledge_config, llm_concurrency.clone()).await {
tracing::warn!("知识提炼触发失败(非阻断): {}", e);
@@ -2707,6 +2704,43 @@ async fn emit_ai_completed_once(
}
}
// ── usage_delta_since: 会话级 token 落库「增量」口径 helper ──
//
// save_conversation 对 usage 做 old+add(accumulate_tokens),故每次 save 只能传
// 「自上次 save 以来的增量」——传累计快照会被同 loop 多次 save(每轮 + 各退出路径)
// 重复累加致 ai_conversations token 双计。本函数取 tokens 当前值减上次快照得增量,
// 并推进快照,同轮多次 save 幂等(第二次增量=0)。
fn usage_delta_since(
tokens: &TokenAccumulator,
last_saved: &mut df_ai::provider::TokenUsage,
) -> df_ai::provider::TokenUsage {
let cur = df_ai::provider::TokenUsage {
prompt_tokens: tokens.prompt(),
completion_tokens: tokens.completion(),
total_tokens: tokens.total(),
prompt_cache_hit_tokens: tokens.cache_hit(),
prompt_cache_miss_tokens: tokens.cache_miss(),
reasoning_tokens: tokens.reasoning(),
};
let delta = df_ai::provider::TokenUsage {
prompt_tokens: cur.prompt_tokens.saturating_sub(last_saved.prompt_tokens),
completion_tokens: cur.completion_tokens.saturating_sub(last_saved.completion_tokens),
total_tokens: 0, // 下方按 prompt+completion 增量重算,保持三字段一致
prompt_cache_hit_tokens: cur
.prompt_cache_hit_tokens
.saturating_sub(last_saved.prompt_cache_hit_tokens),
prompt_cache_miss_tokens: cur
.prompt_cache_miss_tokens
.saturating_sub(last_saved.prompt_cache_miss_tokens),
reasoning_tokens: cur.reasoning_tokens.saturating_sub(last_saved.reasoning_tokens),
};
*last_saved = cur;
df_ai::provider::TokenUsage {
total_tokens: delta.prompt_tokens.saturating_add(delta.completion_tokens),
..delta
}
}
// ── finish_round_exit: run_agentic_loop 收尾统一入口(抽自 5 处退出路径重复代码) ──
//
// 收敛各退出点的「save_conversation + spawn_ensure_title + guard.reset + emit AiCompleted」序列。
@@ -2786,6 +2820,7 @@ async fn finish_round_exit(
// prompt_tokens/completion_tokens: 本轮 LLM 调用 token 用量(消息级持久化,解 reload/压缩/切会话后
// 历史 assistant 消息 token 不显)。两构造分支都设。
// 分项 token(2026-08-02):cache_hit/cache_miss/reasoning 透传自 round_usage,前端分计费展示。
// estimated:本轮 prompt 是否估算兜底(round_usage.prompt_tokens==0),落库供 reload 逐条回显。
fn push_assistant_message(
session: &mut AiSession,
conv_id: &str,
@@ -2799,6 +2834,7 @@ fn push_assistant_message(
cache_hit: u32,
cache_miss: u32,
reasoning: u32,
estimated: bool,
) {
// 根治「空工具轮 assistant 消息落库」:LLM 仅返回 tool_calls 无文本时 full_text 可能为
// 空/纯空白(assistant("")/assistant("\n") 均合法落库),前端渲染空气泡。入口统一 trim:
@@ -2831,6 +2867,7 @@ fn push_assistant_message(
msg.prompt_cache_hit_tokens = Some(cache_hit);
msg.prompt_cache_miss_tokens = Some(cache_miss);
msg.reasoning_tokens = Some(reasoning);
msg.is_estimated = Some(estimated);
session.conv(conv_id).messages.push(msg);
} else if !full_text.is_empty() {
let mut msg = ChatMessage::assistant(full_text);
@@ -2841,6 +2878,7 @@ fn push_assistant_message(
msg.prompt_cache_hit_tokens = Some(cache_hit);
msg.prompt_cache_miss_tokens = Some(cache_miss);
msg.reasoning_tokens = Some(reasoning);
msg.is_estimated = Some(estimated);
session.conv(conv_id).messages.push(msg);
}
}
@@ -106,6 +106,7 @@ pub fn record_to_message(rec: &AiMessageRecord) -> ChatMessage {
prompt_cache_hit_tokens: rec.prompt_cache_hit_tokens,
prompt_cache_miss_tokens: rec.prompt_cache_miss_tokens,
reasoning_tokens: rec.reasoning_tokens,
is_estimated: rec.is_estimated,
}
}
@@ -155,6 +156,7 @@ pub fn message_to_record(
prompt_cache_hit_tokens: msg.prompt_cache_hit_tokens,
prompt_cache_miss_tokens: msg.prompt_cache_miss_tokens,
reasoning_tokens: msg.reasoning_tokens,
is_estimated: msg.is_estimated,
}
}
@@ -836,6 +838,7 @@ mod tests {
prompt_cache_hit_tokens: None,
prompt_cache_miss_tokens: None,
reasoning_tokens: None,
is_estimated: None,
}
}
@@ -706,6 +706,7 @@ mod tests {
prompt_cache_hit_tokens: None,
prompt_cache_miss_tokens: None,
reasoning_tokens: None,
is_estimated: None,
timestamp: None,
}
}
+18 -1
View File
@@ -328,6 +328,11 @@ pub(crate) async fn stream_llm(
),
};
}
// 已收到 finish_reason 仅剩等 usage 帧:慢 provider 拖过 idle 窗口不降 Partial,
// 按正常完成收尾(usage 缺失用默认 0,文本/tool_calls 已完整)。
if finished_received {
break;
}
warn!(
provider = %provider.name(),
conv_id = %conv_id,
@@ -379,6 +384,11 @@ pub(crate) async fn stream_llm(
),
};
}
// 已收到 finish_reason 仅剩等 usage 帧:慢/坏 provider 拖过 15s 不降 Partial,
// 按正常完成收尾(与 idle_deadline 分支同口径,usage 缺失用默认 0)。
if finished_received {
break;
}
warn!(
provider = %provider.name(),
conv_id = %conv_id,
@@ -461,7 +471,14 @@ pub(crate) async fn stream_llm(
}
if chunk.finished {
finished_received = true;
break;
// usage 携带点:OpenAI 兼容流中 usage 挂在 [DONE] 帧(或 usage-only 帧),
// 而 finish_reason 帧(finished=true, usage=None)在其之前到达。
// 若此刻 break 会错过 [DONE] 帧的 usage → 真实 completion_tokens 丢失,
// 前端 token 显示 0。拿到真实 usage 才停;usage 仍 None 则继续读到
// usage-only/[DONE] 帧或通道关闭(Ok(None) 兜底退出,不断连误判)。
if chunk.usage.is_some() {
break;
}
}
}
Ok(Some(Err(err_str))) => {