perf(topics): 话题列表刷新提速——token 统计增量维护列替代全表聚合
- topics 表新增 stat_* 增量列，append 时同事务累加（排除 sub:% 子代理消息） - batch_topic_token_stats 从 messages 全表 SUM+MAX(seq) 聚合改为读 topics 列，刷新耗时与消息总量解耦 - 变更路径重算自愈：delete_by_ids/replace_topic/compact_topic 按 topic 重算，replace_active/compact_active 按 session 重算，clear_messages 归零 - backfill_topic_usage_stats 一次性回填老数据（user_version=4 守卫，仅首启执行） - list_topics handler 移入 spawn_blocking，不再阻塞 tokio worker - 修复既有 bug：压缩/替换路径的 INSERT SQL 漏写 cached_tokens 列 - 新增 5 个测试：增量维护、子代理隔离、删除/清空重算、替换重算、回填迁移幂等
This commit is contained in:
oudecheng
parent
6037c79b9c
commit
e2841ce791
@ -2,8 +2,9 @@ pub mod message;
|
||||
|
||||
pub use crate::domain::messages::ContentBlock;
|
||||
pub use message::{
|
||||
ChatMessage, InboundMessage, MediaItem, OutboundMessage, SYSTEM_CONTEXT_AGENT_PROMPT,
|
||||
SYSTEM_CONTEXT_HISTORY_COMPACTION, SYSTEM_CONTEXT_SCHEDULED_PROMPT,
|
||||
ChatMessage, InboundMessage, MediaItem, MessageUsage, OutboundMessage,
|
||||
SYSTEM_CONTEXT_AGENT_PROMPT, SYSTEM_CONTEXT_HISTORY_COMPACTION,
|
||||
SYSTEM_CONTEXT_SCHEDULED_PROMPT,
|
||||
};
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
@ -139,12 +139,20 @@ async fn handle_list_topics(
|
||||
session_id: String,
|
||||
ctx: CommandContext,
|
||||
) -> Result<CommandResponse, CommandError> {
|
||||
let topics = handler
|
||||
.store
|
||||
.list_topics(&session_id)
|
||||
// 同步 rusqlite 查询移入 blocking 线程池，避免大库查询阻塞 tokio worker
|
||||
// 饿死同连接上的其他 WS 命令（话题列表刷新是高频操作）。
|
||||
let store = handler.store.clone();
|
||||
let session_id_bg = session_id.clone();
|
||||
let summaries = tokio::task::spawn_blocking(
|
||||
move || -> Result<Vec<TopicSummary>, CommandError> {
|
||||
let topics = store
|
||||
.list_topics(&session_id_bg)
|
||||
.map_err(|e| CommandError::new("LIST_TOPICS_ERROR", e.to_string()))?;
|
||||
|
||||
let summaries = build_topic_summaries(handler.store.as_ref(), topics)?;
|
||||
build_topic_summaries(store.as_ref(), topics)
|
||||
},
|
||||
)
|
||||
.await
|
||||
.map_err(|e| CommandError::new("LIST_TOPICS_ERROR", e.to_string()))??;
|
||||
|
||||
let topics_json = serde_json::to_string(&summaries)
|
||||
.map_err(|e| CommandError::new("SERIALIZE_ERROR", e.to_string()))?;
|
||||
|
||||
@ -110,6 +110,135 @@ pub(super) fn ensure_topics_schema(conn: &Connection) -> Result<(), StorageError
|
||||
if !has_column(conn, "topics", "model")? {
|
||||
add_column_if_missing(conn, "ALTER TABLE topics ADD COLUMN model TEXT")?;
|
||||
}
|
||||
|
||||
// Token 用量增量维护列。topics 列表加载（每次前端刷新）原先对 messages
|
||||
// 全量 SUM + MAX(seq) 聚合，大库下显著变慢。改为与 message_count 相同的
|
||||
// 增量维护模式：append 时累加，列表加载直接读列。
|
||||
// stat_has_usage=0 表示该话题从无 usage 数据（前端不显示 token 标签，
|
||||
// 与旧逻辑"stats_map 无条目 → token_stats=None"语义一致）。
|
||||
if !has_column(conn, "topics", "stat_prompt_tokens")? {
|
||||
add_column_if_missing(
|
||||
conn,
|
||||
"ALTER TABLE topics ADD COLUMN stat_prompt_tokens INTEGER NOT NULL DEFAULT 0",
|
||||
)?;
|
||||
}
|
||||
if !has_column(conn, "topics", "stat_completion_tokens")? {
|
||||
add_column_if_missing(
|
||||
conn,
|
||||
"ALTER TABLE topics ADD COLUMN stat_completion_tokens INTEGER NOT NULL DEFAULT 0",
|
||||
)?;
|
||||
}
|
||||
if !has_column(conn, "topics", "stat_total_tokens")? {
|
||||
add_column_if_missing(
|
||||
conn,
|
||||
"ALTER TABLE topics ADD COLUMN stat_total_tokens INTEGER NOT NULL DEFAULT 0",
|
||||
)?;
|
||||
}
|
||||
if !has_column(conn, "topics", "stat_cached_tokens")? {
|
||||
add_column_if_missing(
|
||||
conn,
|
||||
"ALTER TABLE topics ADD COLUMN stat_cached_tokens INTEGER NOT NULL DEFAULT 0",
|
||||
)?;
|
||||
}
|
||||
if !has_column(conn, "topics", "stat_last_prompt_tokens")? {
|
||||
add_column_if_missing(
|
||||
conn,
|
||||
"ALTER TABLE topics ADD COLUMN stat_last_prompt_tokens INTEGER",
|
||||
)?;
|
||||
}
|
||||
if !has_column(conn, "topics", "stat_context_window_tokens")? {
|
||||
add_column_if_missing(
|
||||
conn,
|
||||
"ALTER TABLE topics ADD COLUMN stat_context_window_tokens INTEGER",
|
||||
)?;
|
||||
}
|
||||
if !has_column(conn, "topics", "stat_has_usage")? {
|
||||
add_column_if_missing(
|
||||
conn,
|
||||
"ALTER TABLE topics ADD COLUMN stat_has_usage INTEGER NOT NULL DEFAULT 0",
|
||||
)?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 从 messages 表一次性回填 topics 的 token 统计增量列（user_version 4，一次性迁移）。
|
||||
///
|
||||
/// 聚合口径与旧 batch_topic_token_stats 完全一致：
|
||||
/// - 仅 role='assistant' 且 session_id NOT LIKE 'sub:%'（子代理消息的 topic_id
|
||||
/// 指向父 topic，必须显式排除，否则父话题统计被污染）；
|
||||
/// - last_prompt_tokens / context_window_tokens 取每个 topic 最新一条
|
||||
/// prompt_tokens IS NOT NULL 的 assistant 消息（按 seq 倒序）。
|
||||
///
|
||||
/// 相关子查询命中 idx_messages_topic_seq(topic_id, seq)，总复杂度 O(全部消息)，
|
||||
/// 仅首启执行一次。事务包裹，失败回滚后可安全重跑。
|
||||
pub(super) fn backfill_topic_usage_stats(conn: &mut Connection) -> Result<(), StorageError> {
|
||||
const TOPIC_USAGE_BACKFILL_VERSION: i64 = 4;
|
||||
|
||||
let current_version: i64 = conn.query_row("PRAGMA user_version", [], |row| row.get(0))?;
|
||||
if current_version >= TOPIC_USAGE_BACKFILL_VERSION {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let tx = conn.transaction()?;
|
||||
|
||||
// SUM 累计列：四个相关子查询共享同一过滤条件
|
||||
tx.execute(
|
||||
"UPDATE topics SET
|
||||
stat_prompt_tokens = COALESCE((
|
||||
SELECT SUM(m.prompt_tokens) FROM messages m
|
||||
WHERE m.topic_id = topics.id AND m.role = 'assistant'
|
||||
AND m.session_id NOT LIKE 'sub:%'), 0),
|
||||
stat_completion_tokens = COALESCE((
|
||||
SELECT SUM(m.completion_tokens) FROM messages m
|
||||
WHERE m.topic_id = topics.id AND m.role = 'assistant'
|
||||
AND m.session_id NOT LIKE 'sub:%'), 0),
|
||||
stat_total_tokens = COALESCE((
|
||||
SELECT SUM(m.total_tokens) FROM messages m
|
||||
WHERE m.topic_id = topics.id AND m.role = 'assistant'
|
||||
AND m.session_id NOT LIKE 'sub:%'), 0),
|
||||
stat_cached_tokens = COALESCE((
|
||||
SELECT SUM(m.cached_tokens) FROM messages m
|
||||
WHERE m.topic_id = topics.id AND m.role = 'assistant'
|
||||
AND m.session_id NOT LIKE 'sub:%'), 0)",
|
||||
[],
|
||||
)?;
|
||||
|
||||
// 最新 usage 瞬时列：每个 topic 最新一条有 prompt_tokens 的 assistant 消息
|
||||
tx.execute(
|
||||
"UPDATE topics SET
|
||||
stat_last_prompt_tokens = (
|
||||
SELECT m.prompt_tokens FROM messages m
|
||||
WHERE m.topic_id = topics.id AND m.role = 'assistant'
|
||||
AND m.prompt_tokens IS NOT NULL
|
||||
AND m.session_id NOT LIKE 'sub:%'
|
||||
ORDER BY m.seq DESC LIMIT 1),
|
||||
stat_context_window_tokens = (
|
||||
SELECT m.context_window_tokens FROM messages m
|
||||
WHERE m.topic_id = topics.id AND m.role = 'assistant'
|
||||
AND m.prompt_tokens IS NOT NULL
|
||||
AND m.session_id NOT LIKE 'sub:%'
|
||||
ORDER BY m.seq DESC LIMIT 1)",
|
||||
[],
|
||||
)?;
|
||||
|
||||
// 有 assistant 消息的话题打标（与旧聚合"GROUP BY 有条目"语义一致，
|
||||
// 前端以此决定是否显示 token 标签）
|
||||
tx.execute(
|
||||
"UPDATE topics SET stat_has_usage = 1
|
||||
WHERE EXISTS (
|
||||
SELECT 1 FROM messages m
|
||||
WHERE m.topic_id = topics.id AND m.role = 'assistant'
|
||||
AND m.session_id NOT LIKE 'sub:%')",
|
||||
[],
|
||||
)?;
|
||||
|
||||
tx.execute(
|
||||
&format!("PRAGMA user_version = {TOPIC_USAGE_BACKFILL_VERSION}"),
|
||||
[],
|
||||
)?;
|
||||
tx.commit()?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
||||
@ -7,7 +7,7 @@ use crate::utils::current_timestamp;
|
||||
|
||||
use r2d2::Pool;
|
||||
use r2d2_sqlite::SqliteConnectionManager;
|
||||
use rusqlite::{Connection, OptionalExtension, TransactionBehavior, params};
|
||||
use rusqlite::{Connection, OptionalExtension, Transaction, TransactionBehavior, params};
|
||||
|
||||
use crate::bus::ChatMessage;
|
||||
|
||||
@ -240,6 +240,7 @@ impl SessionStore {
|
||||
ensure_pending_subagents_schema(&conn)?;
|
||||
repair_session_id_prefix_pollution(&mut conn)?;
|
||||
cleanup_legacy_empty_cli_sessions(&mut conn)?;
|
||||
backfill_topic_usage_stats(&mut conn)?;
|
||||
|
||||
drop(conn);
|
||||
|
||||
@ -696,6 +697,46 @@ impl SessionStore {
|
||||
WHERE id = ?1",
|
||||
params![tid],
|
||||
)?;
|
||||
// token 统计增量列同样按活查询重算（含 has_usage 打标），
|
||||
// 口径与 backfill_topic_usage_stats 一致（排除 sub:% 子代理消息）。
|
||||
conn.execute(
|
||||
"UPDATE topics SET
|
||||
stat_prompt_tokens = COALESCE((
|
||||
SELECT SUM(m.prompt_tokens) FROM messages m
|
||||
WHERE m.topic_id = ?1 AND m.role = 'assistant'
|
||||
AND m.session_id NOT LIKE 'sub:%'), 0),
|
||||
stat_completion_tokens = COALESCE((
|
||||
SELECT SUM(m.completion_tokens) FROM messages m
|
||||
WHERE m.topic_id = ?1 AND m.role = 'assistant'
|
||||
AND m.session_id NOT LIKE 'sub:%'), 0),
|
||||
stat_total_tokens = COALESCE((
|
||||
SELECT SUM(m.total_tokens) FROM messages m
|
||||
WHERE m.topic_id = ?1 AND m.role = 'assistant'
|
||||
AND m.session_id NOT LIKE 'sub:%'), 0),
|
||||
stat_cached_tokens = COALESCE((
|
||||
SELECT SUM(m.cached_tokens) FROM messages m
|
||||
WHERE m.topic_id = ?1 AND m.role = 'assistant'
|
||||
AND m.session_id NOT LIKE 'sub:%'), 0),
|
||||
stat_last_prompt_tokens = (
|
||||
SELECT m.prompt_tokens FROM messages m
|
||||
WHERE m.topic_id = ?1 AND m.role = 'assistant'
|
||||
AND m.prompt_tokens IS NOT NULL
|
||||
AND m.session_id NOT LIKE 'sub:%'
|
||||
ORDER BY m.seq DESC LIMIT 1),
|
||||
stat_context_window_tokens = (
|
||||
SELECT m.context_window_tokens FROM messages m
|
||||
WHERE m.topic_id = ?1 AND m.role = 'assistant'
|
||||
AND m.prompt_tokens IS NOT NULL
|
||||
AND m.session_id NOT LIKE 'sub:%'
|
||||
ORDER BY m.seq DESC LIMIT 1),
|
||||
stat_has_usage = CASE WHEN EXISTS (
|
||||
SELECT 1 FROM messages m
|
||||
WHERE m.topic_id = ?1 AND m.role = 'assistant'
|
||||
AND m.session_id NOT LIKE 'sub:%')
|
||||
THEN 1 ELSE 0 END
|
||||
WHERE id = ?1",
|
||||
params![tid],
|
||||
)?;
|
||||
}
|
||||
}
|
||||
Ok(total)
|
||||
@ -720,6 +761,19 @@ impl SessionStore {
|
||||
",
|
||||
params![session_id, now],
|
||||
)?;
|
||||
// token 统计增量列随消息清空归零（消息已全部删除，无需重算）
|
||||
conn.execute(
|
||||
"UPDATE topics SET
|
||||
stat_prompt_tokens = 0,
|
||||
stat_completion_tokens = 0,
|
||||
stat_total_tokens = 0,
|
||||
stat_cached_tokens = 0,
|
||||
stat_last_prompt_tokens = NULL,
|
||||
stat_context_window_tokens = NULL,
|
||||
stat_has_usage = 0
|
||||
WHERE session_id = ?1",
|
||||
params![session_id],
|
||||
)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@ -802,6 +856,10 @@ impl SessionStore {
|
||||
"UPDATE topics SET message_count = message_count + 1, last_active_at = ?2 WHERE id = ?1",
|
||||
params![tid, now],
|
||||
)?;
|
||||
// token 统计增量维护（与 batch_topic_token_stats 读列口径一致）：
|
||||
// 仅主代理 assistant 消息计入（子代理 session_id 以 'sub:' 开头，
|
||||
// 其消息 topic_id 指向父 topic，不计入父话题统计）。
|
||||
Self::accumulate_topic_usage_in_tx(&tx, tid, session_id, message)?;
|
||||
}
|
||||
|
||||
tx.commit()?;
|
||||
@ -898,12 +956,67 @@ impl SessionStore {
|
||||
"UPDATE topics SET message_count = message_count + ?2, last_active_at = ?3 WHERE id = ?1",
|
||||
params![tid, msg_count, now],
|
||||
)?;
|
||||
// token 统计增量维护：逐条累加（批次内 assistant usage 消息通常仅 1-2 条）。
|
||||
// 批次按 seq 递增追加，最后一条有 usage 的 assistant 消息即"最新"。
|
||||
for message in messages {
|
||||
Self::accumulate_topic_usage_in_tx(&tx, tid, session_id, message)?;
|
||||
}
|
||||
}
|
||||
|
||||
tx.commit()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 在给定事务中将单条消息的 usage 累加到 topic 统计列。
|
||||
///
|
||||
/// 过滤口径与回填迁移/旧聚合查询一致：仅主代理（session_id 不以 'sub:' 开头）
|
||||
/// 的 assistant 消息且携带 usage 时计入。last_prompt_tokens /
|
||||
/// context_window_tokens 仅当该消息 prompt_tokens 非空时覆盖
|
||||
/// （对齐旧 SQL 的 `prompt_tokens IS NOT NULL` 过滤）。
|
||||
fn accumulate_topic_usage_in_tx(
|
||||
tx: &Transaction,
|
||||
topic_id: &str,
|
||||
session_id: &str,
|
||||
message: &ChatMessage,
|
||||
) -> Result<(), StorageError> {
|
||||
if session_id.starts_with("sub:") || message.role != "assistant" {
|
||||
return Ok(());
|
||||
}
|
||||
let Some(usage) = message.usage.as_ref() else {
|
||||
return Ok(());
|
||||
};
|
||||
let prompt_tokens = usage.prompt_tokens as i64;
|
||||
let completion_tokens = usage.completion_tokens as i64;
|
||||
let total_tokens = usage.total_tokens as i64;
|
||||
let cached_tokens = usage.cached_tokens as i64;
|
||||
// 与旧聚合一致：last_* 仅取 prompt_tokens IS NOT NULL 的消息。
|
||||
// usage 存在即代表本条 assistant 消息有 prompt_tokens（MessageUsage
|
||||
// 由 provider 响应构造），故直接覆盖。
|
||||
let last_prompt = prompt_tokens;
|
||||
let last_ctx_window: Option<i64> = usage.context_window_tokens.map(|v| v as i64);
|
||||
tx.execute(
|
||||
"UPDATE topics SET
|
||||
stat_prompt_tokens = stat_prompt_tokens + ?2,
|
||||
stat_completion_tokens = stat_completion_tokens + ?3,
|
||||
stat_total_tokens = stat_total_tokens + ?4,
|
||||
stat_cached_tokens = stat_cached_tokens + ?5,
|
||||
stat_last_prompt_tokens = ?6,
|
||||
stat_context_window_tokens = ?7,
|
||||
stat_has_usage = 1
|
||||
WHERE id = ?1",
|
||||
params![
|
||||
topic_id,
|
||||
prompt_tokens,
|
||||
completion_tokens,
|
||||
total_tokens,
|
||||
cached_tokens,
|
||||
last_prompt,
|
||||
last_ctx_window,
|
||||
],
|
||||
)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn compact_active_history(
|
||||
&self,
|
||||
session_id: &str,
|
||||
@ -992,6 +1105,8 @@ impl SessionStore {
|
||||
)?;
|
||||
|
||||
tx.commit()?;
|
||||
// 会话级压缩重组了消息行（标记/删除/插入摘要），所有 topic 统计可能漂移
|
||||
self.recompute_session_topic_usage_stats(session_id)?;
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
@ -1047,6 +1162,8 @@ impl SessionStore {
|
||||
)?;
|
||||
|
||||
tx.commit()?;
|
||||
// 整会话消息被替换，所有 topic 统计可能漂移——重算自愈
|
||||
self.recompute_session_topic_usage_stats(session_id)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@ -1117,6 +1234,8 @@ impl SessionStore {
|
||||
)?;
|
||||
|
||||
tx.commit()?;
|
||||
// 该 topic 消息被整体替换——重算统计列自愈
|
||||
self.recompute_topic_usage_stats(&[topic_id.to_string()])?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@ -1218,6 +1337,8 @@ impl SessionStore {
|
||||
)?;
|
||||
|
||||
tx.commit()?;
|
||||
// 摘要消息可能携带 usage，且旧摘要被删除——重算该 topic 统计列自愈
|
||||
self.recompute_topic_usage_stats(&[topic_id.to_string()])?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@ -2044,16 +2165,19 @@ impl SessionStore {
|
||||
.map_err(StorageError::from)
|
||||
}
|
||||
|
||||
/// 批量查询多个 topic 的 token 消耗统计（cost 累计 + context 瞬时）。
|
||||
/// 批量读取 topic 的 token 用量统计（cost 累计 + context 瞬时）。
|
||||
///
|
||||
/// 按 `topic_id` 聚合而非 `session_id`：一个 session 可包含多个 topic，
|
||||
/// 若按 session_id 聚合会导致同 session 下的所有 topic 显示相同的总和。
|
||||
/// 直接读取 topics 表的增量维护列（stat_*），不再对 messages 全量聚合——
|
||||
/// 话题列表每次刷新都调用本函数，大库下聚合是显著瓶颈。
|
||||
/// 增量列由 append 路径累加、变更路径（删除/替换/清空）重算自愈、
|
||||
/// 老数据由 backfill_topic_usage_stats 一次性回填。
|
||||
///
|
||||
/// 子代理隔离：子代理消息持久化时 session_id='sub:...'，topic_id=父 topic_id
|
||||
/// （见 task::runtime PersistingEmittedMessageHandler 构造），因此不能仅靠
|
||||
/// topic_id 隔离。此处用 `session_id NOT LIKE 'sub:%'` 显式排除子代理消息，
|
||||
/// 与项目约定一致（session 列表同样过滤 'sub:%'）。子代理 token 不计入父 topic，
|
||||
/// 保持"子代理分别计算"语义。
|
||||
/// 子代理隔离在写入/重算侧保证：子代理消息持久化时 session_id='sub:...'，
|
||||
/// topic_id=父 topic_id，累加与重算均以 `session_id NOT LIKE 'sub:%'`
|
||||
/// 排除，子代理 token 不计入父 topic。
|
||||
///
|
||||
/// `stat_has_usage = 0` 的话题不返回条目，保持旧语义
|
||||
/// （stats_map 无条目 → token_stats=None，前端不显示 token 标签）。
|
||||
pub fn batch_topic_token_stats(
|
||||
&self,
|
||||
topic_ids: &[&str],
|
||||
@ -2067,81 +2191,113 @@ impl SessionStore {
|
||||
.map(|i| format!("?{}", i + 1))
|
||||
.collect::<Vec<_>>()
|
||||
.join(", ");
|
||||
// topic_id IN (...) 自动排除 NULL topic_id 的旧消息；
|
||||
// session_id NOT LIKE 'sub:%' 排除子代理消息（其 topic_id=父 topic_id）。
|
||||
// SUM 列清单与行映射见 USAGE_SUM_COLUMNS / read_usage_sum_row（共享于子代理查询）。
|
||||
let sum_sql = format!(
|
||||
"SELECT topic_id, {USAGE_SUM_COLUMNS} \
|
||||
FROM messages \
|
||||
WHERE topic_id IN ({placeholders}) AND role = 'assistant' \
|
||||
AND session_id NOT LIKE 'sub:%' \
|
||||
GROUP BY topic_id"
|
||||
let sql = format!(
|
||||
"SELECT id, stat_prompt_tokens, stat_completion_tokens, stat_total_tokens, \
|
||||
stat_cached_tokens, stat_last_prompt_tokens, stat_context_window_tokens \
|
||||
FROM topics \
|
||||
WHERE id IN ({placeholders}) AND stat_has_usage = 1"
|
||||
);
|
||||
|
||||
let mut stmt = conn.prepare(&sum_sql)?;
|
||||
let mut stmt = conn.prepare(&sql)?;
|
||||
let params: Vec<&dyn rusqlite::ToSql> = topic_ids
|
||||
.iter()
|
||||
.map(|s| s as &dyn rusqlite::ToSql)
|
||||
.collect();
|
||||
let sum_rows = stmt.query_map(params.as_slice(), |row| {
|
||||
Ok((row.get::<_, String>(0)?, read_usage_sum_row(row, 1)?))
|
||||
let rows = stmt.query_map(params.as_slice(), |row| {
|
||||
Ok((
|
||||
row.get::<_, String>(0)?,
|
||||
SessionTokenStats {
|
||||
prompt_tokens: row.get::<_, i64>(1)? as u64,
|
||||
completion_tokens: row.get::<_, i64>(2)? as u64,
|
||||
total_tokens: row.get::<_, i64>(3)? as u64,
|
||||
cached_tokens: row.get::<_, i64>(4)? as u64,
|
||||
last_prompt_tokens: row.get::<_, Option<i64>>(5)?.map(|v| v as u32),
|
||||
context_window_tokens: row.get::<_, Option<i64>>(6)?.map(|v| v as u32),
|
||||
},
|
||||
))
|
||||
})?;
|
||||
|
||||
let mut stats: HashMap<String, SessionTokenStats> = HashMap::new();
|
||||
for row in sum_rows {
|
||||
for row in rows {
|
||||
let (tid, s) = row?;
|
||||
stats.insert(tid, s);
|
||||
}
|
||||
Ok(stats)
|
||||
}
|
||||
|
||||
// 查找每个 topic 中最新的**有 usage 数据的** assistant 消息，
|
||||
// 读取其 prompt_tokens 和 context_window_tokens。
|
||||
// 过滤 prompt_tokens IS NOT NULL 确保跳过 error/cancel 消息（usage 为 NULL）；
|
||||
// session_id NOT LIKE 'sub:%' 排除子代理消息，避免取到子代理的 context_window。
|
||||
//
|
||||
// 注意：seq 是 session 级递增（见 append_message_with_topic），主 session 与
|
||||
// 子代理 session 各自独立计数，可能存在相同 seq。外层 WHERE 必须再次过滤
|
||||
// session_id NOT LIKE 'sub:%'，否则 JOIN 会同时匹配主消息和子代理消息，
|
||||
// 导致重复行并使 stats.entry(tid) 被覆盖，结果不确定。
|
||||
let last_sql = format!(
|
||||
"SELECT m.topic_id, m.prompt_tokens, m.context_window_tokens \
|
||||
FROM messages m \
|
||||
INNER JOIN ( \
|
||||
SELECT topic_id, MAX(seq) AS max_seq \
|
||||
FROM messages \
|
||||
WHERE topic_id IN ({placeholders}) AND role = 'assistant' \
|
||||
AND prompt_tokens IS NOT NULL \
|
||||
AND session_id NOT LIKE 'sub:%' \
|
||||
GROUP BY topic_id \
|
||||
) latest ON m.topic_id = latest.topic_id AND m.seq = latest.max_seq \
|
||||
WHERE m.session_id NOT LIKE 'sub:%'"
|
||||
);
|
||||
let mut stmt2 = conn.prepare(&last_sql)?;
|
||||
let params2: Vec<&dyn rusqlite::ToSql> = topic_ids
|
||||
/// 从 messages 表重算指定 topic 的 token 统计增量列（自愈漂移）。
|
||||
///
|
||||
/// 用于变更消息行的路径：sanitize 删除、话题历史替换、会话清空等。
|
||||
/// 聚合口径与 backfill_topic_usage_stats 完全一致（排除 sub:% 子代理消息）。
|
||||
pub fn recompute_topic_usage_stats(&self, topic_ids: &[String]) -> Result<(), StorageError> {
|
||||
if topic_ids.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
let conn = self.pool.get()?;
|
||||
let placeholders = (0..topic_ids.len())
|
||||
.map(|i| format!("?{}", i + 1))
|
||||
.collect::<Vec<_>>()
|
||||
.join(", ");
|
||||
let filter = format!("id IN ({placeholders})");
|
||||
let params: Vec<&dyn rusqlite::ToSql> = topic_ids
|
||||
.iter()
|
||||
.map(|s| s as &dyn rusqlite::ToSql)
|
||||
.collect();
|
||||
let last_rows = stmt2.query_map(params2.as_slice(), |row| {
|
||||
Ok((
|
||||
row.get::<_, String>(0)?,
|
||||
row.get::<_, Option<i64>>(1)?,
|
||||
row.get::<_, Option<i64>>(2)?,
|
||||
))
|
||||
})?;
|
||||
for row in last_rows {
|
||||
let (tid, last_prompt, last_ctx_window) = row?;
|
||||
let entry = stats.entry(tid).or_insert(SessionTokenStats {
|
||||
prompt_tokens: 0,
|
||||
completion_tokens: 0,
|
||||
total_tokens: 0,
|
||||
cached_tokens: 0,
|
||||
last_prompt_tokens: None,
|
||||
context_window_tokens: None,
|
||||
});
|
||||
entry.last_prompt_tokens = last_prompt.map(|v| v as u32);
|
||||
entry.context_window_tokens = last_ctx_window.map(|v| v as u32);
|
||||
|
||||
conn.execute(
|
||||
&format!(
|
||||
"UPDATE topics SET
|
||||
stat_prompt_tokens = COALESCE((
|
||||
SELECT SUM(m.prompt_tokens) FROM messages m
|
||||
WHERE m.topic_id = topics.id AND m.role = 'assistant'
|
||||
AND m.session_id NOT LIKE 'sub:%'), 0),
|
||||
stat_completion_tokens = COALESCE((
|
||||
SELECT SUM(m.completion_tokens) FROM messages m
|
||||
WHERE m.topic_id = topics.id AND m.role = 'assistant'
|
||||
AND m.session_id NOT LIKE 'sub:%'), 0),
|
||||
stat_total_tokens = COALESCE((
|
||||
SELECT SUM(m.total_tokens) FROM messages m
|
||||
WHERE m.topic_id = topics.id AND m.role = 'assistant'
|
||||
AND m.session_id NOT LIKE 'sub:%'), 0),
|
||||
stat_cached_tokens = COALESCE((
|
||||
SELECT SUM(m.cached_tokens) FROM messages m
|
||||
WHERE m.topic_id = topics.id AND m.role = 'assistant'
|
||||
AND m.session_id NOT LIKE 'sub:%'), 0),
|
||||
stat_last_prompt_tokens = (
|
||||
SELECT m.prompt_tokens FROM messages m
|
||||
WHERE m.topic_id = topics.id AND m.role = 'assistant'
|
||||
AND m.prompt_tokens IS NOT NULL
|
||||
AND m.session_id NOT LIKE 'sub:%'
|
||||
ORDER BY m.seq DESC LIMIT 1),
|
||||
stat_context_window_tokens = (
|
||||
SELECT m.context_window_tokens FROM messages m
|
||||
WHERE m.topic_id = topics.id AND m.role = 'assistant'
|
||||
AND m.prompt_tokens IS NOT NULL
|
||||
AND m.session_id NOT LIKE 'sub:%'
|
||||
ORDER BY m.seq DESC LIMIT 1),
|
||||
stat_has_usage = CASE WHEN EXISTS (
|
||||
SELECT 1 FROM messages m
|
||||
WHERE m.topic_id = topics.id AND m.role = 'assistant'
|
||||
AND m.session_id NOT LIKE 'sub:%')
|
||||
THEN 1 ELSE 0 END
|
||||
WHERE {filter}"
|
||||
),
|
||||
params.as_slice(),
|
||||
)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
Ok(stats)
|
||||
/// 重算指定 session 下所有 topic 的 token 统计增量列（自愈漂移）。
|
||||
///
|
||||
/// 用于整会话级变更路径（replace_active_history / compact_active_history）：
|
||||
/// 这些路径会删除或重组整个 session 的消息行，所有 topic 的统计都可能漂移。
|
||||
pub fn recompute_session_topic_usage_stats(&self, session_id: &str) -> Result<(), StorageError> {
|
||||
let ids: Vec<String> = {
|
||||
let conn = self.pool.get()?;
|
||||
let mut stmt = conn.prepare("SELECT id FROM topics WHERE session_id = ?1")?;
|
||||
stmt.query_map(params![session_id], |row| row.get::<_, String>(0))?
|
||||
.collect::<Result<_, _>>()?
|
||||
};
|
||||
self.recompute_topic_usage_stats(&ids)
|
||||
}
|
||||
|
||||
/// 查询单个 session 的 token 消耗统计（cost 累计 + context 瞬时）。
|
||||
@ -2469,17 +2625,17 @@ fn default_session_db_path() -> Result<PathBuf, std::io::Error> {
|
||||
const INSERT_MESSAGE_SQL: &str = "
|
||||
INSERT INTO messages (
|
||||
id, session_id, seq, role, content,
|
||||
system_context, reasoning_content, media_refs_json, tool_call_id, tool_name, tool_calls_json, tool_duration_ms, prompt_tokens, completion_tokens, total_tokens, context_window_tokens, created_at
|
||||
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17)
|
||||
system_context, reasoning_content, media_refs_json, tool_call_id, tool_name, tool_calls_json, tool_duration_ms, prompt_tokens, completion_tokens, total_tokens, context_window_tokens, cached_tokens, created_at
|
||||
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18)
|
||||
";
|
||||
|
||||
/// 批量插入消息的预编译 SQL（18 列，含 topic_id，无 cached_tokens）。
|
||||
/// 批量插入消息的预编译 SQL（19 列，含 topic_id 与 cached_tokens）。
|
||||
/// 由 `insert_message_with_topic_seq` 使用，循环写入前 `prepare_cached` 一次复用。
|
||||
const INSERT_MESSAGE_TOPIC_SQL: &str = "
|
||||
INSERT INTO messages (
|
||||
id, session_id, topic_id, seq, role, content,
|
||||
system_context, reasoning_content, media_refs_json, tool_call_id, tool_name, tool_calls_json, tool_duration_ms, prompt_tokens, completion_tokens, total_tokens, context_window_tokens, created_at
|
||||
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18)
|
||||
system_context, reasoning_content, media_refs_json, tool_call_id, tool_name, tool_calls_json, tool_duration_ms, prompt_tokens, completion_tokens, total_tokens, context_window_tokens, cached_tokens, created_at
|
||||
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18, ?19)
|
||||
";
|
||||
|
||||
fn insert_message_with_seq(
|
||||
@ -2514,6 +2670,7 @@ fn insert_message_with_seq(
|
||||
.usage
|
||||
.as_ref()
|
||||
.and_then(|u| u.context_window_tokens.map(|v| v as i64)),
|
||||
message.usage.as_ref().map(|u| u.cached_tokens as i64),
|
||||
message.timestamp,
|
||||
])?;
|
||||
Ok(())
|
||||
@ -2558,6 +2715,7 @@ fn insert_message_with_topic_seq(
|
||||
.usage
|
||||
.as_ref()
|
||||
.and_then(|u| u.context_window_tokens.map(|v| v as i64)),
|
||||
message.usage.as_ref().map(|u| u.cached_tokens as i64),
|
||||
message.timestamp,
|
||||
])?;
|
||||
Ok(())
|
||||
|
||||
@ -986,3 +986,265 @@ fn test_delete_messages_by_ids_removes_only_target_rows() {
|
||||
let session_after = store.get_session(&session.id).unwrap().unwrap();
|
||||
assert_eq!(session_after.message_count, 3);
|
||||
}
|
||||
|
||||
fn assistant_with_usage(
|
||||
content: &str,
|
||||
prompt: u32,
|
||||
completion: u32,
|
||||
cached: u32,
|
||||
ctx_window: Option<u32>,
|
||||
) -> ChatMessage {
|
||||
let mut msg = ChatMessage::assistant(content);
|
||||
msg.usage = Some(crate::bus::MessageUsage {
|
||||
prompt_tokens: prompt,
|
||||
completion_tokens: completion,
|
||||
total_tokens: prompt + completion,
|
||||
cached_tokens: cached,
|
||||
context_window_tokens: ctx_window,
|
||||
});
|
||||
msg
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_topic_token_stats_incremental_maintenance() {
|
||||
let store = SessionStore::in_memory().unwrap();
|
||||
let session = store.create_cli_session(Some("stats")).unwrap();
|
||||
let topic = store.create_topic(&session.id, "topic-stats", None).unwrap();
|
||||
|
||||
// 无 usage 数据时不返回条目（前端不显示 token 标签）
|
||||
store
|
||||
.append_message_with_topic(&session.id, Some(&topic.id), &ChatMessage::user("q1"))
|
||||
.unwrap();
|
||||
let stats = store.batch_topic_token_stats(&[topic.id.as_str()]).unwrap();
|
||||
assert!(stats.is_empty());
|
||||
|
||||
// 两条带 usage 的 assistant 消息：SUM 累加，last_* 取最新一条
|
||||
store
|
||||
.append_message_with_topic(
|
||||
&session.id,
|
||||
Some(&topic.id),
|
||||
&assistant_with_usage("a1", 100, 50, 20, Some(128_000)),
|
||||
)
|
||||
.unwrap();
|
||||
store
|
||||
.append_message_with_topic(
|
||||
&session.id,
|
||||
Some(&topic.id),
|
||||
&assistant_with_usage("a2", 200, 80, 100, Some(200_000)),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let stats = store.batch_topic_token_stats(&[topic.id.as_str()]).unwrap();
|
||||
let s = stats.get(&topic.id).expect("topic stats entry");
|
||||
assert_eq!(s.prompt_tokens, 300);
|
||||
assert_eq!(s.completion_tokens, 130);
|
||||
assert_eq!(s.total_tokens, 430);
|
||||
assert_eq!(s.cached_tokens, 120);
|
||||
assert_eq!(s.last_prompt_tokens, Some(200));
|
||||
assert_eq!(s.context_window_tokens, Some(200_000));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_topic_token_stats_excludes_subagent_messages() {
|
||||
let store = SessionStore::in_memory().unwrap();
|
||||
let session = store.create_cli_session(Some("stats-sub")).unwrap();
|
||||
let topic = store.create_topic(&session.id, "topic-sub", None).unwrap();
|
||||
|
||||
// 主代理消息计入
|
||||
store
|
||||
.append_message_with_topic(
|
||||
&session.id,
|
||||
Some(&topic.id),
|
||||
&assistant_with_usage("main", 100, 50, 0, Some(128_000)),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
// 子代理消息 topic_id 指向父 topic 但 session_id 以 'sub:' 开头——不得计入
|
||||
let sub_session = store
|
||||
.ensure_session("sub:child-1", "cli", "sub-chat", "sub")
|
||||
.unwrap();
|
||||
store
|
||||
.append_message_with_topic(
|
||||
&sub_session.id,
|
||||
Some(&topic.id),
|
||||
&assistant_with_usage("sub", 999, 999, 999, Some(999_999)),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let stats = store.batch_topic_token_stats(&[topic.id.as_str()]).unwrap();
|
||||
let s = stats.get(&topic.id).expect("topic stats entry");
|
||||
assert_eq!(s.prompt_tokens, 100);
|
||||
assert_eq!(s.total_tokens, 150);
|
||||
assert_eq!(s.last_prompt_tokens, Some(100));
|
||||
assert_eq!(s.context_window_tokens, Some(128_000));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_topic_token_stats_recompute_after_delete_and_clear() {
|
||||
let store = SessionStore::in_memory().unwrap();
|
||||
let session = store.create_cli_session(Some("stats-recompute")).unwrap();
|
||||
let topic = store.create_topic(&session.id, "topic-rc", None).unwrap();
|
||||
|
||||
store
|
||||
.append_message_with_topic(
|
||||
&session.id,
|
||||
Some(&topic.id),
|
||||
&assistant_with_usage("a1", 100, 50, 0, Some(128_000)),
|
||||
)
|
||||
.unwrap();
|
||||
store
|
||||
.append_message_with_topic(
|
||||
&session.id,
|
||||
Some(&topic.id),
|
||||
&assistant_with_usage("a2", 200, 80, 0, Some(128_000)),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
// 删除第二条 assistant 消息（sanitize 回写场景）：统计重算为仅第一条
|
||||
let msgs = store.load_messages_for_topic(&topic.id, None).unwrap();
|
||||
let assistant_ids: Vec<String> = msgs
|
||||
.iter()
|
||||
.filter(|m| m.role == "assistant")
|
||||
.map(|m| m.id.clone())
|
||||
.collect();
|
||||
store
|
||||
.delete_messages_by_ids(&session.id, &assistant_ids[1..2])
|
||||
.unwrap();
|
||||
|
||||
let stats = store.batch_topic_token_stats(&[topic.id.as_str()]).unwrap();
|
||||
let s = stats.get(&topic.id).expect("topic stats entry");
|
||||
assert_eq!(s.prompt_tokens, 100);
|
||||
assert_eq!(s.total_tokens, 150);
|
||||
assert_eq!(s.last_prompt_tokens, Some(100));
|
||||
|
||||
// 删除全部 assistant 消息后无条目（has_usage=0）
|
||||
store
|
||||
.delete_messages_by_ids(&session.id, &assistant_ids[0..1])
|
||||
.unwrap();
|
||||
let stats = store.batch_topic_token_stats(&[topic.id.as_str()]).unwrap();
|
||||
assert!(stats.is_empty());
|
||||
|
||||
// 再次写入后恢复；clear_messages 归零
|
||||
store
|
||||
.append_message_with_topic(
|
||||
&session.id,
|
||||
Some(&topic.id),
|
||||
&assistant_with_usage("a3", 300, 100, 0, None),
|
||||
)
|
||||
.unwrap();
|
||||
assert!(!store
|
||||
.batch_topic_token_stats(&[topic.id.as_str()])
|
||||
.unwrap()
|
||||
.is_empty());
|
||||
store.clear_messages(&session.id).unwrap();
|
||||
let stats = store.batch_topic_token_stats(&[topic.id.as_str()]).unwrap();
|
||||
assert!(stats.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_topic_token_stats_recompute_after_replace_topic_history() {
|
||||
let store = SessionStore::in_memory().unwrap();
|
||||
let session = store.create_cli_session(Some("stats-replace")).unwrap();
|
||||
let topic = store.create_topic(&session.id, "topic-rp", None).unwrap();
|
||||
|
||||
store
|
||||
.append_message_with_topic(
|
||||
&session.id,
|
||||
Some(&topic.id),
|
||||
&assistant_with_usage("old", 500, 200, 0, Some(128_000)),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
// 整体替换该 topic 历史（压缩器场景）：统计应反映新消息
|
||||
store
|
||||
.replace_topic_history(
|
||||
&session.id,
|
||||
&topic.id,
|
||||
&[
|
||||
ChatMessage::user("new-q"),
|
||||
assistant_with_usage("new-a", 40, 10, 5, Some(64_000)),
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let stats = store.batch_topic_token_stats(&[topic.id.as_str()]).unwrap();
|
||||
let s = stats.get(&topic.id).expect("topic stats entry");
|
||||
assert_eq!(s.prompt_tokens, 40);
|
||||
assert_eq!(s.completion_tokens, 10);
|
||||
assert_eq!(s.total_tokens, 50);
|
||||
assert_eq!(s.cached_tokens, 5);
|
||||
assert_eq!(s.last_prompt_tokens, Some(40));
|
||||
assert_eq!(s.context_window_tokens, Some(64_000));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_backfill_topic_usage_stats_restores_incremental_columns() {
|
||||
let store = SessionStore::in_memory().unwrap();
|
||||
let session = store.create_cli_session(Some("backfill")).unwrap();
|
||||
let topic = store.create_topic(&session.id, "topic-bf", None).unwrap();
|
||||
|
||||
// 主代理 usage 消息两条（增量列此时已有值）
|
||||
store
|
||||
.append_message_with_topic(
|
||||
&session.id,
|
||||
Some(&topic.id),
|
||||
&assistant_with_usage("a1", 100, 50, 20, Some(128_000)),
|
||||
)
|
||||
.unwrap();
|
||||
store
|
||||
.append_message_with_topic(
|
||||
&session.id,
|
||||
Some(&topic.id),
|
||||
&assistant_with_usage("a2", 200, 80, 0, Some(200_000)),
|
||||
)
|
||||
.unwrap();
|
||||
// 子代理消息（topic_id 指向父 topic）：回填必须排除
|
||||
let sub_session = store
|
||||
.ensure_session("sub:bf-child", "cli", "sub-chat", "sub")
|
||||
.unwrap();
|
||||
store
|
||||
.append_message_with_topic(
|
||||
&sub_session.id,
|
||||
Some(&topic.id),
|
||||
&assistant_with_usage("sub", 999, 999, 999, Some(999_999)),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
// 模拟老库：增量列全部清零，版本回退到 3（尚未回填）
|
||||
let mut conn = store.pool.get().unwrap();
|
||||
conn.execute(
|
||||
"UPDATE topics SET stat_prompt_tokens = 0, stat_completion_tokens = 0, \
|
||||
stat_total_tokens = 0, stat_cached_tokens = 0, \
|
||||
stat_last_prompt_tokens = NULL, stat_context_window_tokens = NULL, \
|
||||
stat_has_usage = 0",
|
||||
[],
|
||||
)
|
||||
.unwrap();
|
||||
conn.execute("PRAGMA user_version = 3", []).unwrap();
|
||||
|
||||
super::migrations::backfill_topic_usage_stats(&mut conn).unwrap();
|
||||
|
||||
// 版本守卫推进到 4
|
||||
let version: i64 = conn
|
||||
.query_row("PRAGMA user_version", [], |row| row.get(0))
|
||||
.unwrap();
|
||||
assert_eq!(version, 4);
|
||||
drop(conn);
|
||||
|
||||
// 回填后与增量维护的值一致（子代理被排除）
|
||||
let stats = store.batch_topic_token_stats(&[topic.id.as_str()]).unwrap();
|
||||
let s = stats.get(&topic.id).expect("topic stats entry");
|
||||
assert_eq!(s.prompt_tokens, 300);
|
||||
assert_eq!(s.completion_tokens, 130);
|
||||
assert_eq!(s.total_tokens, 430);
|
||||
assert_eq!(s.cached_tokens, 20);
|
||||
assert_eq!(s.last_prompt_tokens, Some(200));
|
||||
assert_eq!(s.context_window_tokens, Some(200_000));
|
||||
|
||||
// 幂等：版本守卫使第二次调用直接跳过，结果不变
|
||||
let mut conn = store.pool.get().unwrap();
|
||||
super::migrations::backfill_topic_usage_stats(&mut conn).unwrap();
|
||||
drop(conn);
|
||||
let stats = store.batch_topic_token_stats(&[topic.id.as_str()]).unwrap();
|
||||
assert_eq!(stats.get(&topic.id).unwrap().total_tokens, 430);
|
||||
}
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user