From e2841ce791e76f68c3c90cde7f4a6e192557ef75 Mon Sep 17 00:00:00 2001 From: oudecheng <13802883547@139.com> Date: Tue, 18 Aug 2026 22:59:01 +0800 Subject: [PATCH] =?UTF-8?q?perf(topics):=20=E8=AF=9D=E9=A2=98=E5=88=97?= =?UTF-8?q?=E8=A1=A8=E5=88=B7=E6=96=B0=E6=8F=90=E9=80=9F=E2=80=94=E2=80=94?= =?UTF-8?q?token=20=E7=BB=9F=E8=AE=A1=E5=A2=9E=E9=87=8F=E7=BB=B4=E6=8A=A4?= =?UTF-8?q?=E5=88=97=E6=9B=BF=E4=BB=A3=E5=85=A8=E8=A1=A8=E8=81=9A=E5=90=88?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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 个测试：增量维护、子代理隔离、删除/清空重算、替换重算、回填迁移幂等 --- src/bus/mod.rs | 5 +- src/command/handlers/list_topics.rs | 20 +- src/storage/migrations.rs | 129 ++++++++++++ src/storage/mod.rs | 304 +++++++++++++++++++++------- src/storage/tests.rs | 262 ++++++++++++++++++++++++ 5 files changed, 639 insertions(+), 81 deletions(-) diff --git a/src/bus/mod.rs b/src/bus/mod.rs index 4a5abf9..2d52083 100644 --- a/src/bus/mod.rs +++ b/src/bus/mod.rs @@ -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; diff --git a/src/command/handlers/list_topics.rs b/src/command/handlers/list_topics.rs index 848c3b8..4f66c55 100644 --- a/src/command/handlers/list_topics.rs +++ b/src/command/handlers/list_topics.rs @@ -139,12 +139,20 @@ async fn handle_list_topics( session_id: String, ctx: CommandContext, ) -> Result { - let topics = handler - .store - .list_topics(&session_id) - .map_err(|e| CommandError::new("LIST_TOPICS_ERROR", e.to_string()))?; - - let summaries = build_topic_summaries(handler.store.as_ref(), topics)?; + // 同步 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, CommandError> { + let topics = store + .list_topics(&session_id_bg) + .map_err(|e| CommandError::new("LIST_TOPICS_ERROR", e.to_string()))?; + 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()))?; diff --git a/src/storage/migrations.rs b/src/storage/migrations.rs index 2d935db..f8b4842 100644 --- a/src/storage/migrations.rs +++ b/src/storage/migrations.rs @@ -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(()) } diff --git a/src/storage/mod.rs b/src/storage/mod.rs index 0209bce..9a8a35a 100644 --- a/src/storage/mod.rs +++ b/src/storage/mod.rs @@ -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 = 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::>() .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>(5)?.map(|v| v as u32), + context_window_tokens: row.get::<_, Option>(6)?.map(|v| v as u32), + }, + )) })?; let mut stats: HashMap = 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::>() + .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>(1)?, - row.get::<_, Option>(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); - } - Ok(stats) + 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(()) + } + + /// 重算指定 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 = { + 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::>()? + }; + self.recompute_topic_usage_stats(&ids) } /// 查询单个 session 的 token 消耗统计（cost 累计 + context 瞬时）。 @@ -2469,17 +2625,17 @@ fn default_session_db_path() -> Result { 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(()) diff --git a/src/storage/tests.rs b/src/storage/tests.rs index a504847..fefb876 100644 --- a/src/storage/tests.rs +++ b/src/storage/tests.rs @@ -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, +) -> 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 = 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); +}