From 2e4b1931a685308ca8c5b02d78d6aff3f00885b3 Mon Sep 17 00:00:00 2001 From: oudecheng <13802883547@139.com> Date: Tue, 4 Aug 2026 17:01:54 +0800 Subject: [PATCH] =?UTF-8?q?feat(tokens):=20=E6=B7=BB=E5=8A=A0=20topic=20?= =?UTF-8?q?=E7=BB=B4=E5=BA=A6=20token=20=E6=B6=88=E8=80=97=E4=B8=8E?= =?UTF-8?q?=E4=B8=8A=E4=B8=8B=E6=96=87=E7=AA=97=E5=8F=A3=E5=8D=A0=E7=94=A8?= =?UTF-8?q?=E7=BB=9F=E8=AE=A1?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 后端： - openai provider 流式响应启用 stream_options.include_usage，从 SSE 末帧捕获 usage - ChatMessage 新增 MessageUsage（prompt/completion/total/context_window_tokens） - agent_loop 三处 LLM 调用点持久化 usage，含 context_window_tokens - storage 新增 context_window_tokens 列及 migration，batch_session_token_stats 聚合查询 - last_prompt_tokens 查询过滤 prompt_tokens IS NOT NULL，跳过 error/cancel 消息 - build_topic_summaries 简化签名，context_window_tokens 从消息记录读取 - protocol 扩展 TopicTokenStats，按 session_id 聚合天然隔离子代理 前端： - protocol.ts 新增 TopicTokenStats 类型 - useTopics 映射 token_stats 到 Topic - TopicList 展示总 token 消耗与上下文占用百分比（绿/黄/红三色） 测试：outbound_dispatcher 4 个测试修复（drop(bus) 不关闭 bus，改用 abort） --- src/agent/agent_loop.rs | 18 +- src/bus/message.rs | 38 +++ src/command/handlers/delete_topic.rs | 15 +- src/command/handlers/list_topics.rs | 91 +++++- src/command/handlers/rename_topic.rs | 30 +- src/command/handlers/session.rs | 15 +- src/gateway/outbound_dispatcher.rs | 400 ++++++++++++++++++++++- src/protocol/mod.rs | 16 + src/providers/openai.rs | 50 ++- src/storage/migrations.rs | 20 ++ src/storage/mod.rs | 145 +++++++- src/storage/records.rs | 10 + src/storage/row_mapping.rs | 30 ++ web/src/components/Sidebar/TopicList.tsx | 50 ++- web/src/hooks/chat/useTopics.ts | 1 + web/src/types/protocol.ts | 10 + 16 files changed, 842 insertions(+), 97 deletions(-) diff --git a/src/agent/agent_loop.rs b/src/agent/agent_loop.rs index a03314f..e137614 100644 --- a/src/agent/agent_loop.rs +++ b/src/agent/agent_loop.rs @@ -350,6 +350,7 @@ fn filter_images_by_age_and_count( tool_state: message.tool_state.clone(), tool_duration_ms: message.tool_duration_ms, tool_calls: message.tool_calls.clone(), + usage: message.usage.clone(), }); } @@ -1214,6 +1215,11 @@ impl AgentLoop { if had_streaming { assistant_message.id = streaming_message_id; } + // 记录本次 LLM 调用的 token 用量（中间带 tool_calls 的调用也算消耗） + assistant_message.usage = Some( + crate::bus::message::MessageUsage::from_provider_usage(response.usage.clone()) + .with_context_window(self.runtime_config.context_window_tokens), + ); messages.push(assistant_message.clone()); emitted_messages.push(assistant_message); self.emit_live_tool_call_message( @@ -1379,6 +1385,11 @@ impl AgentLoop { if had_streaming { assistant_message.id = streaming_message_id.to_string(); } + // 记录本次 LLM 调用的 token 用量（cost 累计 + context 瞬时） + assistant_message.usage = Some( + crate::bus::message::MessageUsage::from_provider_usage(response.usage) + .with_context_window(self.runtime_config.context_window_tokens), + ); emitted_messages.push(assistant_message.clone()); self.emit_live_tool_call_message(assistant_message.clone()) .await; @@ -1481,12 +1492,17 @@ impl AgentLoop { match final_result { Ok(response) => { - let assistant_message = if let Some(reasoning_content) = response.reasoning_content + let mut assistant_message = if let Some(reasoning_content) = response.reasoning_content { ChatMessage::assistant_with_reasoning(response.content, reasoning_content) } else { ChatMessage::assistant(response.content) }; + // 记录 summary 调用的 token 用量 + assistant_message.usage = Some( + crate::bus::message::MessageUsage::from_provider_usage(response.usage) + .with_context_window(self.runtime_config.context_window_tokens), + ); emitted_messages.push(assistant_message.clone()); self.emit_live_tool_call_message(assistant_message.clone()) .await; diff --git a/src/bus/message.rs b/src/bus/message.rs index c8b2472..0d85a7b 100644 --- a/src/bus/message.rs +++ b/src/bus/message.rs @@ -66,6 +66,38 @@ pub struct ChatMessage { pub tool_duration_ms: Option, #[serde(skip_serializing_if = "Option::is_none")] pub tool_calls: Option>, + /// LLM 调用 usage（仅 assistant 消息有值，来自 provider 响应） + #[serde(default, skip_serializing_if = "Option::is_none")] + pub usage: Option, +} + +/// 单次 LLM 调用的 token 用量 +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +pub struct MessageUsage { + pub prompt_tokens: u32, + pub completion_tokens: u32, + pub total_tokens: u32, + /// 本次调用所用模型的上下文窗口大小（来自 AgentRuntimeConfig）。 + /// 与 prompt_tokens 一起持久化，用于计算上下文占用率。 + #[serde(default, skip_serializing_if = "Option::is_none")] + pub context_window_tokens: Option, +} + +impl MessageUsage { + pub fn from_provider_usage(u: crate::providers::Usage) -> Self { + Self { + prompt_tokens: u.prompt_tokens, + completion_tokens: u.completion_tokens, + total_tokens: u.total_tokens, + context_window_tokens: None, + } + } + + /// 链式设置 context_window_tokens（来自 AgentRuntimeConfig） + pub fn with_context_window(mut self, ctx: usize) -> Self { + self.context_window_tokens = Some(ctx as u32); + self + } } impl ChatMessage { @@ -83,6 +115,7 @@ impl ChatMessage { tool_duration_ms: None, tool_state: None, tool_calls: None, + usage: None, } } @@ -100,6 +133,7 @@ impl ChatMessage { tool_duration_ms: None, tool_state: None, tool_calls: None, + usage: None, } } @@ -117,6 +151,7 @@ impl ChatMessage { tool_duration_ms: None, tool_state: None, tool_calls: None, + usage: None, } } @@ -146,6 +181,7 @@ impl ChatMessage { tool_duration_ms: None, tool_state: None, tool_calls: Some(tool_calls), + usage: None, } } @@ -180,6 +216,7 @@ impl ChatMessage { tool_duration_ms: None, tool_state: None, tool_calls: None, + usage: None, } } @@ -215,6 +252,7 @@ impl ChatMessage { tool_duration_ms: None, tool_state: Some(tool_state), tool_calls: None, + usage: None, } } diff --git a/src/command/handlers/delete_topic.rs b/src/command/handlers/delete_topic.rs index 1d1e452..6000fae 100644 --- a/src/command/handlers/delete_topic.rs +++ b/src/command/handlers/delete_topic.rs @@ -1,7 +1,7 @@ use crate::command::Command; use crate::command::context::CommandContext; use crate::command::handler::{CommandHandler, CommandMetadata}; -use crate::command::handlers::list_topics::TopicSummary; +use crate::command::handlers::list_topics::build_topic_summaries; use crate::command::response::{CommandError, CommandResponse, MessageKind}; use crate::gateway::session::SessionManager; use crate::storage::SessionStore; @@ -87,18 +87,7 @@ async fn handle_delete_topic( .list_topics(session_id) .map_err(|e| CommandError::new("LIST_TOPICS_ERROR", e.to_string()))?; - let topic_summaries: Vec = topics - .into_iter() - .map(|t| TopicSummary { - topic_id: t.id, - session_id: t.session_id, - title: t.title, - description: t.description.filter(|d| !d.is_empty()), - message_count: t.message_count, - created_at: t.created_at, - last_active_at: t.last_active_at, - }) - .collect(); + let topic_summaries = build_topic_summaries(handler.store.as_ref(), topics)?; let topics_json = serde_json::to_string(&topic_summaries) .map_err(|e| CommandError::new("SERIALIZE_ERROR", e.to_string()))?; diff --git a/src/command/handlers/list_topics.rs b/src/command/handlers/list_topics.rs index e99bc5a..c57dc68 100644 --- a/src/command/handlers/list_topics.rs +++ b/src/command/handlers/list_topics.rs @@ -2,11 +2,27 @@ use crate::command::Command; use crate::command::context::CommandContext; use crate::command::handler::{CommandHandler, CommandMetadata}; use crate::command::response::{CommandError, CommandResponse, MessageKind}; -use crate::storage::SessionStore; +use crate::storage::{SessionStore, SessionTokenStats, TopicRecord}; use async_trait::async_trait; use serde::{Deserialize, Serialize}; +use std::collections::HashMap; use std::sync::Arc; +/// Topic 维度的 token 统计（cost 累计 + context 瞬时）。 +/// +/// - `prompt_tokens` / `completion_tokens` / `total_tokens`：累计求和（cost 维度） +/// - `last_prompt_tokens`：最后一条 assistant 消息的 prompt_tokens（context 占用瞬时值） +/// - `context_window_tokens`：当前 session 使用的模型上下文窗口上限（来自配置）； +/// 为 0 表示未配置，前端不显示百分比。 +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct TopicTokenStats { + pub prompt_tokens: u64, + pub completion_tokens: u64, + pub total_tokens: u64, + pub last_prompt_tokens: Option, + pub context_window_tokens: u32, +} + /// Topic 摘要信息 #[derive(Debug, Clone, Serialize, Deserialize)] pub struct TopicSummary { @@ -17,6 +33,64 @@ pub struct TopicSummary { pub message_count: i64, pub created_at: i64, pub last_active_at: i64, + /// Token 用量统计。老 topic 或无 LLM 调用时为 None，前端不显示 token 标签。 + #[serde(skip_serializing_if = "Option::is_none")] + pub token_stats: Option, +} + +/// 构建 TopicSummary 列表的公共函数。 +/// +/// 一次批量查询所有 topic 对应 session 的 token 统计，避免 N 次 RTT。 +/// 按 `session_id` 聚合天然分离主 agent 与子 agent（子代理 session_id 形如 +/// `sub:...`，不在此查询的主 session_id 列表中）。 +/// +/// `context_window_tokens` 来自最新 assistant 消息记录（LLM 调用时持久化）， +/// 无需从配置链路注入，保持存储层与配置解耦。 +pub fn build_topic_summaries( + store: &SessionStore, + topics: Vec, +) -> Result, CommandError> { + if topics.is_empty() { + return Ok(Vec::new()); + } + + // 收集所有 topic 的 session_id（去重） + let mut session_ids: Vec = Vec::new(); + for t in &topics { + if !session_ids.contains(&t.session_id) { + session_ids.push(t.session_id.clone()); + } + } + let session_id_refs: Vec<&str> = session_ids.iter().map(|s| s.as_str()).collect(); + + // 一次批量查询 token 统计 + let stats_map: HashMap = store + .batch_session_token_stats(&session_id_refs) + .map_err(|e| CommandError::new("TOKEN_STATS_ERROR", e.to_string()))?; + + let summaries = topics + .into_iter() + .map(|t| { + let token_stats = stats_map.get(&t.session_id).map(|s| TopicTokenStats { + prompt_tokens: s.prompt_tokens, + completion_tokens: s.completion_tokens, + total_tokens: s.total_tokens, + last_prompt_tokens: s.last_prompt_tokens, + context_window_tokens: s.context_window_tokens.unwrap_or(0), + }); + TopicSummary { + topic_id: t.id, + session_id: t.session_id, + title: t.title, + description: t.description.filter(|d| !d.is_empty()), + message_count: t.message_count, + created_at: t.created_at, + last_active_at: t.last_active_at, + token_stats, + } + }) + .collect(); + Ok(summaries) } /// 列出 Session 的 Topics 命令处理器 @@ -66,18 +140,9 @@ async fn handle_list_topics( .list_topics(&session_id) .map_err(|e| CommandError::new("LIST_TOPICS_ERROR", e.to_string()))?; - let summaries: Vec = topics - .into_iter() - .map(|t| TopicSummary { - topic_id: t.id, - session_id: t.session_id, - title: t.title, - description: t.description.filter(|d| !d.is_empty()), - message_count: t.message_count, - created_at: t.created_at, - last_active_at: t.last_active_at, - }) - .collect(); + // context_window_tokens 当前未从配置链路注入（前端可从已有 config 接口获取）， + // 此处传 0 表示"后端不提供上限"，前端按需隐藏百分比。 + let summaries = build_topic_summaries(handler.store.as_ref(), topics)?; let topics_json = serde_json::to_string(&summaries) .map_err(|e| CommandError::new("SERIALIZE_ERROR", e.to_string()))?; diff --git a/src/command/handlers/rename_topic.rs b/src/command/handlers/rename_topic.rs index 004bbd8..0db12b7 100644 --- a/src/command/handlers/rename_topic.rs +++ b/src/command/handlers/rename_topic.rs @@ -1,7 +1,7 @@ use crate::command::Command; use crate::command::context::CommandContext; use crate::command::handler::{CommandHandler, CommandMetadata}; -use crate::command::handlers::list_topics::TopicSummary; +use crate::command::handlers::list_topics::build_topic_summaries; use crate::command::response::{CommandError, CommandResponse, MessageKind}; use crate::storage::SessionStore; use async_trait::async_trait; @@ -83,14 +83,16 @@ async fn handle_rename_topic( .store .list_topics(session_id) .map_err(|e| CommandError::new("LIST_TOPICS_ERROR", e.to_string()))?; - let topic_summaries = serialize_summaries(&topics); + let topic_summaries = build_topic_summaries(handler.store.as_ref(), topics)?; + let topic_summaries_json = serde_json::to_string(&topic_summaries) + .map_err(|e| CommandError::new("SERIALIZE_ERROR", e.to_string()))?; return Ok(CommandResponse::success(ctx.request_id) .with_message( MessageKind::Notification, &format!("✓ 话题标题未变化: {}", trimmed_title), ) - .with_metadata("topics", &topic_summaries) + .with_metadata("topics", &topic_summaries_json) .with_metadata("topic_id", &topic_id) .with_metadata("title", trimmed_title) .with_metadata("session_id", session_id)); @@ -108,34 +110,20 @@ async fn handle_rename_topic( .list_topics(session_id) .map_err(|e| CommandError::new("LIST_TOPICS_ERROR", e.to_string()))?; - let topic_summaries = serialize_summaries(&topics); + let topic_summaries = build_topic_summaries(handler.store.as_ref(), topics)?; + let topic_summaries_json = serde_json::to_string(&topic_summaries) + .map_err(|e| CommandError::new("SERIALIZE_ERROR", e.to_string()))?; let message = format!("✓ 已重命名话题: {} → {}", old_title, trimmed_title); Ok(CommandResponse::success(ctx.request_id) .with_message(MessageKind::Notification, &message) - .with_metadata("topics", &topic_summaries) + .with_metadata("topics", &topic_summaries_json) .with_metadata("topic_id", &topic_id) .with_metadata("title", trimmed_title) .with_metadata("session_id", session_id)) } -fn serialize_summaries(topics: &[crate::storage::TopicRecord]) -> String { - let summaries: Vec = topics - .iter() - .map(|t| TopicSummary { - topic_id: t.id.clone(), - session_id: t.session_id.clone(), - title: t.title.clone(), - description: t.description.clone().filter(|d| !d.is_empty()), - message_count: t.message_count, - created_at: t.created_at, - last_active_at: t.last_active_at, - }) - .collect(); - serde_json::to_string(&summaries).unwrap_or_else(|_| "[]".to_string()) -} - #[cfg(test)] mod tests { use super::*; diff --git a/src/command/handlers/session.rs b/src/command/handlers/session.rs index 7da9549..8368c93 100644 --- a/src/command/handlers/session.rs +++ b/src/command/handlers/session.rs @@ -1,7 +1,7 @@ use crate::command::Command; use crate::command::context::CommandContext; use crate::command::handler::{CommandHandler, CommandMetadata}; -use crate::command::handlers::list_topics::TopicSummary; +use crate::command::handlers::list_topics::build_topic_summaries; use crate::command::response::{CommandError, CommandResponse, MessageKind}; use crate::gateway::session::SessionManager; use crate::storage::SessionStore; @@ -109,18 +109,7 @@ async fn handle_create_session( .list_topics(session_id) .map_err(|e| CommandError::new("LIST_TOPICS_ERROR", e.to_string()))?; - let topic_summaries: Vec = topics - .into_iter() - .map(|t| TopicSummary { - topic_id: t.id, - session_id: t.session_id, - title: t.title, - description: t.description.filter(|d| !d.is_empty()), - message_count: t.message_count, - created_at: t.created_at, - last_active_at: t.last_active_at, - }) - .collect(); + let topic_summaries = build_topic_summaries(handler.store.as_ref(), topics)?; let topics_json = serde_json::to_string(&topic_summaries) .map_err(|e| CommandError::new("SERIALIZE_ERROR", e.to_string()))?; diff --git a/src/gateway/outbound_dispatcher.rs b/src/gateway/outbound_dispatcher.rs index 3f6c604..0aeac41 100644 --- a/src/gateway/outbound_dispatcher.rs +++ b/src/gateway/outbound_dispatcher.rs @@ -1,20 +1,45 @@ use std::collections::HashMap; use std::sync::Arc; -use tokio::sync::RwLock; +use tokio::sync::{RwLock, mpsc}; +use tokio_util::sync::CancellationToken; use crate::bus::{MessageBus, OutboundMessage}; use crate::channels::base::{Channel, ChannelError}; -/// Consumes outbound messages from MessageBus and dispatches them to channels. -pub struct OutboundDispatcher { - bus: Arc, - channels: Arc>>>, -} +/// 每个 channel 的发送队列容量。 +/// +/// 略小于 MessageBus 的容量（100），确保 bus 的 `try_send` 丢消息 +/// 防线仍有效——channel 队列满时 dispatcher 立即丢弃该消息， +/// 不会阻塞路由循环影响其他 channel。 +const PER_CHANNEL_QUEUE_CAPACITY: usize = 64; + +/// 单个 channel 的发送重试间隔（秒）。 +const RETRY_DELAYS_SECS: [u64; 3] = [1, 2, 4]; /// Prefix for virtual scheduler chat IDs that should not be sent to external channels. const SCHEDULER_VIRTUAL_CHAT_ID_PREFIX: &str = "scheduler/"; +/// 单个 channel 的发送上下文：独立 mpsc 队列 + sender task。 +/// +/// dispatcher 将消息 `try_send` 到 `tx`，`sender_task` 串行消费并调用 +/// `Channel::send`（含重试）。channel 之间完全隔离——某个 channel 的 +/// 慢发送或重试 sleep 不会阻塞其他 channel 的消息投递。 +#[derive(Clone)] +struct ChannelSink { + tx: mpsc::Sender, + cancel: CancellationToken, +} + +/// Consumes outbound messages from MessageBus and dispatches them to channels. +/// +/// 架构：dispatcher 主循环只负责路由（O(1) try_send），不参与发送。 +/// 每个 channel 拥有独立的 sender task 和有界队列，实现 channel 级隔离。 +pub struct OutboundDispatcher { + bus: Arc, + channels: Arc>>, +} + impl OutboundDispatcher { pub fn new(bus: Arc) -> Self { Self { @@ -23,11 +48,62 @@ impl OutboundDispatcher { } } + /// 注册 channel 并启动其独立 sender task。 + /// + /// sender task 生命周期与 dispatcher 一致：dispatcher `run()` 退出时 + /// 通过 cancel token 终止所有 sender task。 pub async fn register_channel(&self, name: &str, channel: Arc) { + let (tx, rx) = mpsc::channel::(PER_CHANNEL_QUEUE_CAPACITY); + let cancel = CancellationToken::new(); + + let channel_name = name.to_string(); + let cancel_for_task = cancel.clone(); + + tokio::spawn(async move { + Self::run_sender_task(&channel_name, channel, rx, cancel_for_task).await; + }); + self.channels .write() .await - .insert(name.to_string(), channel); + .insert(name.to_string(), ChannelSink { tx, cancel }); + } + + /// sender task：串行消费 channel 队列，调用 `Channel::send` 并重试。 + /// + /// 重试 sleep 只阻塞当前 channel 的 task，不影响其他 channel。 + async fn run_sender_task( + channel_name: &str, + channel: Arc, + mut rx: mpsc::Receiver, + cancel: CancellationToken, + ) { + tracing::info!(channel = %channel_name, "Channel sender task started"); + + loop { + tokio::select! { + // 接收下一条待发送消息 + msg = rx.recv() => { + let Some(msg) = msg else { + tracing::info!(channel = %channel_name, "Channel queue closed, sender task stopping"); + break; + }; + + if let Err(error) = Self::send_with_retry(&*channel, msg).await { + tracing::error!( + channel = %channel_name, + error = %error, + "Failed to send message after retries" + ); + } + } + // dispatcher 退出时取消所有 sender task + _ = cancel.cancelled() => { + tracing::info!(channel = %channel_name, "Sender task cancelled, stopping"); + break; + } + } + } } pub async fn run(&self) { @@ -41,6 +117,7 @@ impl OutboundDispatcher { break; } }; + #[cfg(debug_assertions)] tracing::debug!( channel = %msg.channel, @@ -52,6 +129,7 @@ impl OutboundDispatcher { // Skip messages with virtual scheduler chat IDs (e.g., "scheduler/job_id") // These are internal messages from SilentAgentTask that should not be sent externally if msg.chat_id.starts_with(SCHEDULER_VIRTUAL_CHAT_ID_PREFIX) { + #[cfg(debug_assertions)] tracing::debug!( channel = %msg.channel, chat_id = %msg.chat_id, @@ -61,12 +139,26 @@ impl OutboundDispatcher { } let channel_name = msg.channel.clone(); - let channel = self.channels.read().await.get(&channel_name).cloned(); + let sink = self.channels.read().await.get(&channel_name).cloned(); - match channel { - Some(ch) => { - if let Err(error) = self.send_with_retry(&*ch, msg).await { - tracing::error!(channel = %channel_name, error = %error, "Failed to send message after retries"); + match sink { + Some(sink) => { + // try_send 保证 dispatcher 永不阻塞：队列满时立即丢弃该消息， + // 不影响其他 channel 的投递。与 bus.publish_outbound 策略一致。 + match sink.tx.try_send(msg) { + Ok(()) => {} + Err(mpsc::error::TrySendError::Full(_)) => { + tracing::warn!( + channel = %channel_name, + "Channel queue full, dropping message" + ); + } + Err(mpsc::error::TrySendError::Closed(_)) => { + tracing::warn!( + channel = %channel_name, + "Channel queue closed, dropping message" + ); + } } } None => { @@ -74,19 +166,27 @@ impl OutboundDispatcher { } } } + + // 通知所有 sender task 退出 + let sinks = self.channels.write().await; + for (name, sink) in sinks.iter() { + sink.cancel.cancel(); + tracing::debug!(channel = %name, "Cancelled sender task"); + } } + /// 发送消息，失败时按 `[1, 2, 4]` 秒间隔重试。 + /// + /// 仅在单个 channel 的 sender task 内执行——重试 sleep 只阻塞 + /// 该 channel 的发送，不影响其他 channel。 async fn send_with_retry( - &self, channel: &dyn Channel, msg: OutboundMessage, ) -> Result<(), ChannelError> { - const DELAYS: [u64; 3] = [1, 2, 4]; - - for (attempt_index, delay) in DELAYS.iter().enumerate() { + for (attempt_index, delay) in RETRY_DELAYS_SECS.iter().enumerate() { match channel.send(msg.clone()).await { Ok(()) => return Ok(()), - Err(error) if attempt_index < DELAYS.len() - 1 => { + Err(error) if attempt_index < RETRY_DELAYS_SECS.len() - 1 => { tracing::warn!( attempt = attempt_index + 1, delay = delay, @@ -102,3 +202,269 @@ impl OutboundDispatcher { unreachable!() } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::bus::OutboundMessage; + use async_trait::async_trait; + use std::sync::atomic::{AtomicU32, Ordering}; + use std::time::Duration; + + /// 测试用 channel：记录所有收到的消息内容，可配置人为延迟和失败。 + struct TestChannel { + name: String, + received: Arc, + delay_ms: u64, + fail_first_n: u32, + call_count: Arc, + } + + impl TestChannel { + fn new(name: &str) -> Self { + Self { + name: name.to_string(), + received: Arc::new(AtomicU32::new(0)), + delay_ms: 0, + fail_first_n: 0, + call_count: Arc::new(AtomicU32::new(0)), + } + } + + fn with_delay(mut self, ms: u64) -> Self { + self.delay_ms = ms; + self + } + + fn with_fail_first_n(mut self, n: u32) -> Self { + self.fail_first_n = n; + self + } + } + + #[async_trait] + impl Channel for TestChannel { + fn name(&self) -> &str { + &self.name + } + + fn is_running(&self) -> bool { + true + } + + async fn start(&self, _bus: Arc) -> Result<(), ChannelError> { + Ok(()) + } + + async fn stop(&self) -> Result<(), ChannelError> { + Ok(()) + } + + async fn send(&self, _msg: OutboundMessage) -> Result<(), ChannelError> { + let count = self.call_count.fetch_add(1, Ordering::SeqCst); + if (count as u32) < self.fail_first_n { + return Err(ChannelError::SendError("simulated failure".to_string())); + } + + if self.delay_ms > 0 { + tokio::time::sleep(Duration::from_millis(self.delay_ms)).await; + } + + self.received.fetch_add(1, Ordering::SeqCst); + Ok(()) + } + } + + fn make_message(channel: &str, chat_id: &str, content: &str) -> OutboundMessage { + OutboundMessage::assistant( + channel, + chat_id, + None, + content, + None, + std::collections::HashMap::new(), + ) + } + + #[tokio::test] + async fn test_fast_channel_not_blocked_by_slow_channel() { + // 验证核心目标：channel A 慢发送不应阻塞 channel B 的消息投递 + let bus = MessageBus::new(16); + let dispatcher = OutboundDispatcher::new(bus.clone()); + + let slow = Arc::new(TestChannel::new("slow").with_delay(500)); + let fast = Arc::new(TestChannel::new("fast")); + + let slow_received = slow.received.clone(); + let fast_received = fast.received.clone(); + + dispatcher.register_channel("slow", slow).await; + dispatcher.register_channel("fast", fast).await; + + let dispatcher_handle = tokio::spawn(async move { + dispatcher.run().await; + }); + + // 先发一条 slow（500ms 延迟），紧接着发一条 fast + bus.publish_outbound(make_message("slow", "chat-1", "slow-msg")).await.unwrap(); + bus.publish_outbound(make_message("fast", "chat-2", "fast-msg")).await.unwrap(); + + // 等待 fast 消息被投递（远早于 slow 完成） + tokio::time::timeout(Duration::from_millis(200), async { + while fast_received.load(Ordering::SeqCst) == 0 { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("fast channel should receive message within 200ms, but was blocked by slow channel"); + + // 等待 slow 消息完成 + tokio::time::timeout(Duration::from_secs(2), async { + while slow_received.load(Ordering::SeqCst) == 0 { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("slow channel should eventually receive message"); + + assert_eq!(fast_received.load(Ordering::SeqCst), 1); + assert_eq!(slow_received.load(Ordering::SeqCst), 1); + + // 关闭 bus 让 dispatcher 退出 + // dispatcher 持有 Arc 克隆，drop(bus) 不会关闭 bus。 + // 直接 abort dispatcher 及其 sender task 即可清理。 + dispatcher_handle.abort(); + } + + #[tokio::test] + async fn test_retry_does_not_block_other_channel() { + // 验证：channel A 重试 sleep（1+2=3秒）期间，channel B 正常投递 + let bus = MessageBus::new(16); + let dispatcher = OutboundDispatcher::new(bus.clone()); + + let flaky = Arc::new( + TestChannel::new("flaky") + .with_fail_first_n(2) // 前 2 次失败，触发 1+2 秒重试 + .with_delay(0), + ); + let stable = Arc::new(TestChannel::new("stable")); + + let flaky_received = flaky.received.clone(); + let stable_received = stable.received.clone(); + + dispatcher.register_channel("flaky", flaky).await; + dispatcher.register_channel("stable", stable).await; + + let dispatcher_handle = tokio::spawn(async move { + dispatcher.run().await; + }); + + // 先发 flaky（会重试 3 秒），紧接着发 stable + bus.publish_outbound(make_message("flaky", "chat-1", "flaky-msg")).await.unwrap(); + bus.publish_outbound(make_message("stable", "chat-2", "stable-msg")).await.unwrap(); + + // stable 应在 200ms 内收到，远早于 flaky 的 3 秒重试完成 + tokio::time::timeout(Duration::from_millis(200), async { + while stable_received.load(Ordering::SeqCst) == 0 { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("stable channel should not be blocked by flaky channel's retry sleep"); + + // 等待 flaky 重试成功（第 3 次尝试） + tokio::time::timeout(Duration::from_secs(5), async { + while flaky_received.load(Ordering::SeqCst) == 0 { + tokio::time::sleep(Duration::from_millis(50)).await; + } + }) + .await + .expect("flaky channel should eventually succeed after retries"); + + assert_eq!(stable_received.load(Ordering::SeqCst), 1); + assert_eq!(flaky_received.load(Ordering::SeqCst), 1); + + // dispatcher 持有 Arc 克隆，drop(bus) 不会关闭 bus。 + // 直接 abort dispatcher 及其 sender task 即可清理。 + dispatcher_handle.abort(); + } + + #[tokio::test] + async fn test_scheduler_virtual_chat_id_skipped() { + // 验证：scheduler/ 前缀的 chat_id 不被投递到任何 channel + let bus = MessageBus::new(16); + let dispatcher = OutboundDispatcher::new(bus.clone()); + + let channel = Arc::new(TestChannel::new("test")); + let received = channel.received.clone(); + + dispatcher.register_channel("test", channel).await; + + let dispatcher_handle = tokio::spawn(async move { + dispatcher.run().await; + }); + + // scheduler 虚拟消息应被跳过 + bus.publish_outbound(make_message("test", "scheduler/job-1", "internal")) + .await + .unwrap(); + // 正常消息应被投递 + bus.publish_outbound(make_message("test", "chat-1", "normal")) + .await + .unwrap(); + + tokio::time::timeout(Duration::from_secs(2), async { + while received.load(Ordering::SeqCst) == 0 { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("normal message should be delivered"); + + // 只收到 1 条（scheduler 虚拟消息被跳过） + assert_eq!(received.load(Ordering::SeqCst), 1); + + // dispatcher 持有 Arc 克隆，drop(bus) 不会关闭 bus。 + // 直接 abort dispatcher 及其 sender task 即可清理。 + dispatcher_handle.abort(); + } + + #[tokio::test] + async fn test_unknown_channel_warns_and_continues() { + // 验证：未知 channel 的消息被跳过，不影响后续消息投递 + let bus = MessageBus::new(16); + let dispatcher = OutboundDispatcher::new(bus.clone()); + + let channel = Arc::new(TestChannel::new("known")); + let received = channel.received.clone(); + + dispatcher.register_channel("known", channel).await; + + let dispatcher_handle = tokio::spawn(async move { + dispatcher.run().await; + }); + + // 发往未知 channel 的消息 + bus.publish_outbound(make_message("unknown", "chat-1", "lost")) + .await + .unwrap(); + // 发往已知 channel 的消息 + bus.publish_outbound(make_message("known", "chat-2", "delivered")) + .await + .unwrap(); + + tokio::time::timeout(Duration::from_secs(2), async { + while received.load(Ordering::SeqCst) == 0 { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("known channel should receive message despite preceding unknown channel message"); + + assert_eq!(received.load(Ordering::SeqCst), 1); + + // dispatcher 持有 Arc 克隆，drop(bus) 不会关闭 bus。 + // 直接 abort dispatcher 及其 sender task 即可清理。 + dispatcher_handle.abort(); + } +} diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs index 75812fb..ed14bc8 100644 --- a/src/protocol/mod.rs +++ b/src/protocol/mod.rs @@ -42,6 +42,22 @@ pub struct TopicSummary { pub last_active_at: i64, #[serde(default, skip_serializing_if = "Option::is_none")] pub description: Option, + /// Token 用量统计（与 command::handlers::TopicSummary 对应）。 + /// 老消息或未触发 LLM 调用的 topic 为 None。 + #[serde(default, skip_serializing_if = "Option::is_none")] + pub token_stats: Option, +} + +/// Topic 维度的 token 统计（与 command::handlers::TopicTokenStats 对应）。 +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct TopicTokenStats { + pub prompt_tokens: u64, + pub completion_tokens: u64, + pub total_tokens: u64, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub last_prompt_tokens: Option, + #[serde(default)] + pub context_window_tokens: u32, } #[derive(Debug, Clone, Serialize, Deserialize)] diff --git a/src/providers/openai.rs b/src/providers/openai.rs index 40d84ee..bc0cf0e 100644 --- a/src/providers/openai.rs +++ b/src/providers/openai.rs @@ -32,6 +32,8 @@ struct StreamingAccumulator { reasoning_content: Option, tool_calls: BTreeMap, response_id: String, + /// 流式末帧返回的 usage（需要 stream_options.include_usage=true） + usage: Option, } impl StreamingAccumulator { @@ -89,6 +91,14 @@ impl StreamingAccumulator { } } + /// 设置 usage（来自流式末帧的 usage 字段） + /// 跳过 total_tokens=0 的占位帧，避免覆盖真实值。 + fn set_usage(&mut self, usage: OpenAIUsage) { + if usage.total_tokens > 0 { + self.usage = Some(usage); + } + } + /// 构建最终的 ChatCompletionResponse fn build_response(self, model: String) -> ChatCompletionResponse { let tool_calls: Vec = self @@ -116,11 +126,15 @@ impl StreamingAccumulator { content: self.content, reasoning_content: self.reasoning_content, tool_calls, - usage: Usage { + usage: self.usage.clone().map(|u| Usage { + prompt_tokens: u.prompt_tokens, + completion_tokens: u.completion_tokens, + total_tokens: u.total_tokens, + }).unwrap_or(Usage { prompt_tokens: 0, completion_tokens: 0, total_tokens: 0, - }, + }), } } } @@ -385,6 +399,8 @@ impl OpenAIProvider { let mut body = self.build_request_body(request); // 启用流式输出 body["stream"] = json!(true); + // 请求在流式末帧返回 usage（DeepSeek/OpenAI 兼容协议） + body["stream_options"] = json!({ "include_usage": true }); let mut req_builder = self .client @@ -469,6 +485,17 @@ impl OpenAIProvider { accumulator.set_response_id(id.to_string()); } + // 提取流式末帧的 usage（stream_options.include_usage=true 时返回） + if let Some(usage_val) = json.get("usage") { + if !usage_val.is_null() { + if let Ok(u) = + serde_json::from_value::(usage_val.clone()) + { + accumulator.set_usage(u); + } + } + } + // 提取 choices if let Some(choices) = json.get("choices").and_then(|c| c.as_array()) { for choice in choices { @@ -582,6 +609,17 @@ impl OpenAIProvider { accumulator.set_response_id(id.to_string()); } + // 提取流式末帧的 usage（与主循环一致） + if let Some(usage_val) = json.get("usage") { + if !usage_val.is_null() { + if let Ok(u) = + serde_json::from_value::(usage_val.clone()) + { + accumulator.set_usage(u); + } + } + } + if let Some(choices) = json.get("choices").and_then(|c| c.as_array()) { for choice in choices { // 尝试从 delta 提取 @@ -691,6 +729,12 @@ impl OpenAIProvider { .collect() }) .unwrap_or_default(); + // 回退场景下也从非流式响应提取 usage + response.usage = Usage { + prompt_tokens: openai_resp.usage.prompt_tokens, + completion_tokens: openai_resp.usage.completion_tokens, + total_tokens: openai_resp.usage.total_tokens, + }; } } } @@ -1033,7 +1077,7 @@ struct OAIFunction { arguments: OAIFunctionArguments, } -#[derive(Deserialize, Default)] +#[derive(Deserialize, Default, Clone, Debug)] struct OpenAIUsage { #[serde(default)] prompt_tokens: u32, diff --git a/src/storage/migrations.rs b/src/storage/migrations.rs index 109f8c6..8d5f2f6 100644 --- a/src/storage/migrations.rs +++ b/src/storage/migrations.rs @@ -51,6 +51,26 @@ pub(super) fn ensure_messages_schema(conn: &Connection) -> Result<(), StorageErr )?; } + // Token usage 字段（仅 assistant 消息有值，来自 LLM 响应） + if !has_column(conn, "messages", "prompt_tokens")? { + add_column_if_missing(conn, "ALTER TABLE messages ADD COLUMN prompt_tokens INTEGER")?; + } + if !has_column(conn, "messages", "completion_tokens")? { + add_column_if_missing( + conn, + "ALTER TABLE messages ADD COLUMN completion_tokens INTEGER", + )?; + } + if !has_column(conn, "messages", "total_tokens")? { + add_column_if_missing(conn, "ALTER TABLE messages ADD COLUMN total_tokens INTEGER")?; + } + if !has_column(conn, "messages", "context_window_tokens")? { + add_column_if_missing( + conn, + "ALTER TABLE messages ADD COLUMN context_window_tokens INTEGER", + )?; + } + // 创建 topic_id 索引（如果不存在） conn.execute( "CREATE INDEX IF NOT EXISTS idx_messages_topic_seq ON messages(topic_id, seq) WHERE topic_id IS NOT NULL", diff --git a/src/storage/mod.rs b/src/storage/mod.rs index ad9ae7b..4226a10 100644 --- a/src/storage/mod.rs +++ b/src/storage/mod.rs @@ -1,6 +1,8 @@ #[cfg(not(test))] use std::path::{Path, PathBuf}; +use std::collections::HashMap; + use r2d2::Pool; use r2d2_sqlite::SqliteConnectionManager; use rusqlite::{Connection, OptionalExtension, TransactionBehavior, params}; @@ -25,8 +27,8 @@ pub use ports::{ }; pub use records::{ ALLOWED_MEMORY_NAMESPACES, GLOBAL_SCOPE_KEY, MemoryRecord, MemoryUpsert, SchedulerJobRecord, - SchedulerJobState, SchedulerJobStatus, SchedulerJobUpsert, SessionRecord, SkillEventRecord, - TodoRecord, TopicRecord, allowed_namespace_names, get_namespace_description, + SchedulerJobState, SchedulerJobStatus, SchedulerJobUpsert, SessionRecord, SessionTokenStats, + SkillEventRecord, TodoRecord, TopicRecord, allowed_namespace_names, get_namespace_description, is_valid_namespace, }; @@ -102,6 +104,11 @@ impl SessionStore { tool_call_id TEXT, tool_name TEXT, tool_calls_json TEXT, + tool_duration_ms INTEGER, + prompt_tokens INTEGER, + completion_tokens INTEGER, + total_tokens INTEGER, + context_window_tokens INTEGER, created_at INTEGER NOT NULL, FOREIGN KEY(session_id) REFERENCES sessions(id) ON DELETE CASCADE, FOREIGN KEY(topic_id) REFERENCES topics(id) ON DELETE SET NULL, @@ -599,8 +606,8 @@ impl SessionStore { " 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, created_at - ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14) + 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) ", params![ message.id, @@ -616,6 +623,10 @@ impl SessionStore { message.tool_name, tool_calls_json, message.tool_duration_ms.map(|v| v as i64), + message.usage.as_ref().map(|u| u.prompt_tokens as i64), + message.usage.as_ref().map(|u| u.completion_tokens as i64), + message.usage.as_ref().map(|u| u.total_tokens as i64), + message.usage.as_ref().and_then(|u| u.context_window_tokens.map(|v| v as i64)), message.timestamp, ], )?; @@ -677,8 +688,8 @@ impl SessionStore { 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, created_at - ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14) + 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) ", params![ message.id, @@ -694,6 +705,10 @@ impl SessionStore { message.tool_name, tool_calls_json, message.tool_duration_ms.map(|v| v as i64), + message.usage.as_ref().map(|u| u.prompt_tokens as i64), + message.usage.as_ref().map(|u| u.completion_tokens as i64), + message.usage.as_ref().map(|u| u.total_tokens as i64), + message.usage.as_ref().and_then(|u| u.context_window_tokens.map(|v| v as i64)), message.timestamp, ], )?; @@ -1551,7 +1566,7 @@ impl SessionStore { if let Some(sid) = session_id { let mut stmt = conn.prepare( " - SELECT id, role, content, system_context, reasoning_content, media_refs_json, created_at, tool_call_id, tool_name, tool_calls_json, tool_duration_ms + SELECT id, role, content, system_context, reasoning_content, media_refs_json, created_at, tool_call_id, tool_name, tool_calls_json, tool_duration_ms, prompt_tokens, completion_tokens, total_tokens, context_window_tokens FROM messages WHERE topic_id = ?1 AND session_id = ?2 ORDER BY seq ASC @@ -1566,7 +1581,7 @@ impl SessionStore { } else { let mut stmt = conn.prepare( " - SELECT id, role, content, system_context, reasoning_content, media_refs_json, created_at, tool_call_id, tool_name, tool_calls_json, tool_duration_ms + SELECT id, role, content, system_context, reasoning_content, media_refs_json, created_at, tool_call_id, tool_name, tool_calls_json, tool_duration_ms, prompt_tokens, completion_tokens, total_tokens, context_window_tokens FROM messages WHERE topic_id = ?1 ORDER BY seq ASC @@ -1615,6 +1630,96 @@ impl SessionStore { .map_err(StorageError::from) } + /// 批量查询多个 session 的 token 消耗统计（cost 累计 + context 瞬时）。 + pub fn batch_session_token_stats( + &self, + session_ids: &[&str], + ) -> Result, StorageError> { + if session_ids.is_empty() { + return Ok(HashMap::new()); + } + let conn = self.pool.get()?; + + let placeholders = (0..session_ids.len()) + .map(|i| format!("?{}", i + 1)) + .collect::>() + .join(", "); + let sum_sql = format!( + "SELECT session_id, \ + COALESCE(SUM(prompt_tokens), 0) AS sum_prompt, \ + COALESCE(SUM(completion_tokens), 0) AS sum_completion, \ + COALESCE(SUM(total_tokens), 0) AS sum_total \ + FROM messages \ + WHERE session_id IN ({placeholders}) AND role = 'assistant' \ + GROUP BY session_id" + ); + + let mut stmt = conn.prepare(&sum_sql)?; + let params: Vec<&dyn rusqlite::ToSql> = session_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)?, + 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, + last_prompt_tokens: None, + context_window_tokens: None, + }, + )) + })?; + + let mut stats: HashMap = HashMap::new(); + for row in sum_rows { + let (sid, s) = row?; + stats.insert(sid, s); + } + + // 查找每个 session 中最新的**有 usage 数据的** assistant 消息， + // 读取其 prompt_tokens 和 context_window_tokens。 + // 过滤 prompt_tokens IS NOT NULL 确保跳过 error/cancel 消息（usage 为 NULL）。 + let last_sql = format!( + "SELECT m.session_id, m.prompt_tokens, m.context_window_tokens \ + FROM messages m \ + INNER JOIN ( \ + SELECT session_id, MAX(seq) AS max_seq \ + FROM messages \ + WHERE session_id IN ({placeholders}) AND role = 'assistant' \ + AND prompt_tokens IS NOT NULL \ + GROUP BY session_id \ + ) latest ON m.session_id = latest.session_id AND m.seq = latest.max_seq" + ); + let mut stmt2 = conn.prepare(&last_sql)?; + let params2: Vec<&dyn rusqlite::ToSql> = session_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 (sid, last_prompt, last_ctx_window) = row?; + let entry = stats.entry(sid).or_insert(SessionTokenStats { + prompt_tokens: 0, + completion_tokens: 0, + total_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) + } + pub fn replace_todos( &self, scope_key: &str, @@ -1750,8 +1855,8 @@ fn insert_message_with_seq( " 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, created_at - ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13) + 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) ", params![ message.id, @@ -1766,6 +1871,10 @@ fn insert_message_with_seq( message.tool_name, tool_calls_json, message.tool_duration_ms.map(|v| v as i64), + message.usage.as_ref().map(|u| u.prompt_tokens as i64), + message.usage.as_ref().map(|u| u.completion_tokens as i64), + message.usage.as_ref().map(|u| u.total_tokens as i64), + message.usage.as_ref().and_then(|u| u.context_window_tokens.map(|v| v as i64)), message.timestamp, ], )?; @@ -1794,8 +1903,8 @@ fn insert_message_with_topic_seq( " 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, created_at - ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14) + 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) ", params![ message.id, @@ -1811,6 +1920,10 @@ fn insert_message_with_topic_seq( message.tool_name, tool_calls_json, message.tool_duration_ms.map(|v| v as i64), + message.usage.as_ref().map(|u| u.prompt_tokens as i64), + message.usage.as_ref().map(|u| u.completion_tokens as i64), + message.usage.as_ref().map(|u| u.total_tokens as i64), + message.usage.as_ref().and_then(|u| u.context_window_tokens.map(|v| v as i64)), message.timestamp, ], )?; @@ -1831,6 +1944,8 @@ fn clone_message_for_compaction(message: &ChatMessage, timestamp: i64) -> ChatMe tool_state: message.tool_state.clone(), tool_duration_ms: message.tool_duration_ms, tool_calls: message.tool_calls.clone(), + // 压缩克隆不保留 usage：压缩产生的是合成消息，不代表真实 LLM 调用 + usage: None, } } @@ -1842,7 +1957,7 @@ fn load_messages_between( ) -> Result, StorageError> { let mut stmt = conn.prepare( " - SELECT id, role, content, system_context, reasoning_content, media_refs_json, created_at, tool_call_id, tool_name, tool_calls_json, tool_duration_ms + SELECT id, role, content, system_context, reasoning_content, media_refs_json, created_at, tool_call_id, tool_name, tool_calls_json, tool_duration_ms, prompt_tokens, completion_tokens, total_tokens, context_window_tokens FROM messages WHERE session_id = ?1 AND seq > ?2 AND seq <= ?3 ORDER BY seq ASC @@ -1888,6 +2003,7 @@ fn load_messages_between( tool_state: None, tool_duration_ms: row.get::<_, Option>(10)?.map(|v| v as u64), tool_calls, + usage: map_usage_row(row, 11, 12, 13, 14)?, }) }, )?; @@ -1906,7 +2022,7 @@ fn load_messages_after( ) -> Result, StorageError> { let mut stmt = conn.prepare( " - SELECT id, role, content, system_context, reasoning_content, media_refs_json, created_at, tool_call_id, tool_name, tool_calls_json, tool_duration_ms + SELECT id, role, content, system_context, reasoning_content, media_refs_json, created_at, tool_call_id, tool_name, tool_calls_json, tool_duration_ms, prompt_tokens, completion_tokens, total_tokens, context_window_tokens FROM messages WHERE session_id = ?1 AND seq > ?2 ORDER BY seq ASC @@ -1949,6 +2065,7 @@ fn load_messages_after( tool_state: None, tool_duration_ms: row.get::<_, Option>(10)?.map(|v| v as u64), tool_calls, + usage: map_usage_row(row, 11, 12, 13, 14)?, }) })?; diff --git a/src/storage/records.rs b/src/storage/records.rs index 90127a1..5087353 100644 --- a/src/storage/records.rs +++ b/src/storage/records.rs @@ -111,6 +111,16 @@ pub struct TopicRecord { pub message_count: i64, } +/// 单个 session 的 token 用量统计（聚合结果）。 +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct SessionTokenStats { + pub prompt_tokens: u64, + pub completion_tokens: u64, + pub total_tokens: u64, + pub last_prompt_tokens: Option, + pub context_window_tokens: Option, +} + #[derive(Debug, Clone, Serialize, Deserialize)] pub struct MemoryRecord { pub id: String, diff --git a/src/storage/row_mapping.rs b/src/storage/row_mapping.rs index 4747920..3db8766 100644 --- a/src/storage/row_mapping.rs +++ b/src/storage/row_mapping.rs @@ -8,12 +8,41 @@ use rusqlite::{Connection, OptionalExtension, params}; use crate::bus::ChatMessage; +use crate::bus::message::MessageUsage; use super::{ MemoryRecord, SchedulerJobRecord, SchedulerJobState, SchedulerJobStatus, SessionRecord, SkillEventRecord, StorageError, }; +/// 从指定列索引读取 token usage 四元组（含 context_window_tokens）。 +pub(super) fn map_usage_row( + row: &rusqlite::Row<'_>, + prompt_idx: usize, + completion_idx: usize, + total_idx: usize, + context_window_idx: usize, +) -> rusqlite::Result> { + let prompt: Option = row.get(prompt_idx)?; + let completion: Option = row.get(completion_idx)?; + let total: Option = row.get(total_idx)?; + let context_window: Option = row.get(context_window_idx)?; + if prompt.is_none() + && completion.is_none() + && total.is_none() + && context_window.is_none() + { + Ok(None) + } else { + Ok(Some(MessageUsage { + prompt_tokens: prompt.unwrap_or(0) as u32, + completion_tokens: completion.unwrap_or(0) as u32, + total_tokens: total.unwrap_or(0) as u32, + context_window_tokens: context_window.map(|v| v as u32), + })) + } +} + pub(super) fn get_session_with_conn( conn: &Connection, session_id: &str, @@ -147,6 +176,7 @@ pub(super) fn map_chat_message_row(row: &rusqlite::Row<'_>) -> rusqlite::Result< tool_state: None, tool_duration_ms: row.get::<_, Option>(10)?.map(|v| v as u64), tool_calls, + usage: map_usage_row(row, 11, 12, 13, 14)?, }) } diff --git a/web/src/components/Sidebar/TopicList.tsx b/web/src/components/Sidebar/TopicList.tsx index 9f695af..e954df2 100644 --- a/web/src/components/Sidebar/TopicList.tsx +++ b/web/src/components/Sidebar/TopicList.tsx @@ -12,8 +12,9 @@ import { ChevronLeft, ChevronRight, Edit2, + Coins, } from 'lucide-react'; -import type { Topic } from '../../types/protocol'; +import type { Topic, TopicTokenStats } from '../../types/protocol'; interface TopicListProps { sessionId: string | null; @@ -43,6 +44,28 @@ function formatTime(timestamp: number): string { } } +/** 紧凑格式化 token 数量：1234 -> "1.2K"，1234567 -> "1.2M" */ +function formatTokenCount(n: number): string { + if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(1)}M`; + if (n >= 1_000) return `${(n / 1_000).toFixed(1)}K`; + return String(n); +} + +/** 计算上下文窗口占用百分比：last_prompt_tokens / context_window_tokens */ +function contextOccupancyPct(stats: TopicTokenStats): number | null { + if (!stats.context_window_tokens || stats.context_window_tokens === 0) return null; + const last = stats.last_prompt_tokens; + if (last == null) return null; + return Math.min(100, Math.round((last / stats.context_window_tokens) * 100)); +} + +/** 根据占用率返回颜色 class */ +function occupancyColor(pct: number): string { + if (pct >= 80) return 'text-red-400'; + if (pct >= 50) return 'text-amber-400'; + return 'text-emerald-400'; +} + export function TopicList({ sessionId, topics, @@ -253,7 +276,7 @@ export function TopicList({ > {topic.description || topic.title} -
+
{topic.message_count} 条消息 @@ -262,6 +285,29 @@ export function TopicList({ {formatTime(topic.updated_at)} + {topic.tokenStats && topic.tokenStats.total_tokens > 0 && ( + <> + + + {formatTokenCount(topic.tokenStats.total_tokens)} + + {(() => { + const pct = contextOccupancyPct(topic.tokenStats!); + if (pct == null) return null; + return ( + + ctx {pct}% + + ); + })()} + + )}
{topic.id === currentTopicId && ( diff --git a/web/src/hooks/chat/useTopics.ts b/web/src/hooks/chat/useTopics.ts index 152516c..098769a 100644 --- a/web/src/hooks/chat/useTopics.ts +++ b/web/src/hooks/chat/useTopics.ts @@ -40,6 +40,7 @@ function mapTopicSummaries(summaries: TopicSummary[]): Topic[] { message_count: Number(t.message_count), created_at: t.created_at, updated_at: t.last_active_at, + tokenStats: t.token_stats, })); } diff --git a/web/src/types/protocol.ts b/web/src/types/protocol.ts index f69fd96..611b3a9 100644 --- a/web/src/types/protocol.ts +++ b/web/src/types/protocol.ts @@ -149,6 +149,14 @@ export interface SessionSaved { filepath: string; } +export interface TopicTokenStats { + prompt_tokens: number; + completion_tokens: number; + total_tokens: number; + last_prompt_tokens?: number; + context_window_tokens: number; +} + export interface TopicSummary { topic_id: string; session_id: string; @@ -157,6 +165,7 @@ export interface TopicSummary { message_count: number; created_at: number; last_active_at: number; + token_stats?: TopicTokenStats; } export interface TopicList { @@ -508,6 +517,7 @@ export interface Topic { message_count: number; created_at: number; updated_at: number; + tokenStats?: TopicTokenStats; } export interface Session {