feat(tokens): 添加 topic 维度 token 消耗与上下文窗口占用统计

后端:
- 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)
This commit is contained in:
oudecheng 2026-08-04 17:01:54 +08:00
parent 07247ed140
commit 2e4b1931a6
16 changed files with 842 additions and 97 deletions

View File

@ -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;

View File

@ -66,6 +66,38 @@ pub struct ChatMessage {
pub tool_duration_ms: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tool_calls: Option<Vec<ToolCall>>,
/// LLM 调用 usage仅 assistant 消息有值,来自 provider 响应)
#[serde(default, skip_serializing_if = "Option::is_none")]
pub usage: Option<MessageUsage>,
}
/// 单次 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<u32>,
}
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,
}
}

View File

@ -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<TopicSummary> = 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()))?;

View File

@ -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_tokenscontext 占用瞬时值)
/// - `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<u32>,
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<TopicTokenStats>,
}
/// 构建 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<TopicRecord>,
) -> Result<Vec<TopicSummary>, CommandError> {
if topics.is_empty() {
return Ok(Vec::new());
}
// 收集所有 topic 的 session_id去重
let mut session_ids: Vec<String> = 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<String, SessionTokenStats> = 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<TopicSummary> = 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()))?;

View File

@ -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<TopicSummary> = 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::*;

View File

@ -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<TopicSummary> = 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()))?;

View File

@ -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<MessageBus>,
channels: Arc<RwLock<HashMap<String, Arc<dyn Channel + Send + Sync>>>>,
}
/// 每个 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<OutboundMessage>,
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<MessageBus>,
channels: Arc<RwLock<HashMap<String, ChannelSink>>>,
}
impl OutboundDispatcher {
pub fn new(bus: Arc<MessageBus>) -> 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<dyn Channel + Send + Sync>) {
let (tx, rx) = mpsc::channel::<OutboundMessage>(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<dyn Channel + Send + Sync>,
mut rx: mpsc::Receiver<OutboundMessage>,
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<AtomicU32>,
delay_ms: u64,
fail_first_n: u32,
call_count: Arc<AtomicU32>,
}
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<MessageBus>) -> 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;
});
// 先发一条 slow500ms 延迟),紧接着发一条 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<MessageBus> 克隆drop(bus) 不会关闭 bus。
// 直接 abort dispatcher 及其 sender task 即可清理。
dispatcher_handle.abort();
}
#[tokio::test]
async fn test_retry_does_not_block_other_channel() {
// 验证channel A 重试 sleep1+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<MessageBus> 克隆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<MessageBus> 克隆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<MessageBus> 克隆drop(bus) 不会关闭 bus。
// 直接 abort dispatcher 及其 sender task 即可清理。
dispatcher_handle.abort();
}
}

View File

@ -42,6 +42,22 @@ pub struct TopicSummary {
pub last_active_at: i64,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub description: Option<String>,
/// Token 用量统计(与 command::handlers::TopicSummary 对应)。
/// 老消息或未触发 LLM 调用的 topic 为 None。
#[serde(default, skip_serializing_if = "Option::is_none")]
pub token_stats: Option<TopicTokenStats>,
}
/// 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<u32>,
#[serde(default)]
pub context_window_tokens: u32,
}
#[derive(Debug, Clone, Serialize, Deserialize)]

View File

@ -32,6 +32,8 @@ struct StreamingAccumulator {
reasoning_content: Option<String>,
tool_calls: BTreeMap<usize, StreamingToolCall>,
response_id: String,
/// 流式末帧返回的 usage需要 stream_options.include_usage=true
usage: Option<OpenAIUsage>,
}
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<ToolCall> = 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);
// 请求在流式末帧返回 usageDeepSeek/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());
}
// 提取流式末帧的 usagestream_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::<OpenAIUsage>(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::<OpenAIUsage>(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,

View File

@ -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",

View File

@ -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<HashMap<String, SessionTokenStats>, 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::<Vec<_>>()
.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<String, SessionTokenStats> = 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<i64>>(1)?,
row.get::<_, Option<i64>>(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<Vec<ChatMessage>, 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<i64>>(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<Vec<ChatMessage>, 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<i64>>(10)?.map(|v| v as u64),
tool_calls,
usage: map_usage_row(row, 11, 12, 13, 14)?,
})
})?;

View File

@ -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<u32>,
pub context_window_tokens: Option<u32>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MemoryRecord {
pub id: String,

View File

@ -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<Option<MessageUsage>> {
let prompt: Option<i64> = row.get(prompt_idx)?;
let completion: Option<i64> = row.get(completion_idx)?;
let total: Option<i64> = row.get(total_idx)?;
let context_window: Option<i64> = 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<i64>>(10)?.map(|v| v as u64),
tool_calls,
usage: map_usage_row(row, 11, 12, 13, 14)?,
})
}

View File

@ -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}
</div>
<div className="flex items-center gap-3 mt-1.5">
<div className="flex flex-wrap items-center gap-x-3 gap-y-1 mt-1.5">
<span className="text-xs text-[var(--text-muted)] flex items-center gap-1">
<Hash className="h-3 w-3" />
{topic.message_count}
@ -262,6 +285,29 @@ export function TopicList({
<Clock className="h-3 w-3" />
{formatTime(topic.updated_at)}
</span>
{topic.tokenStats && topic.tokenStats.total_tokens > 0 && (
<>
<span
className="text-xs text-[var(--text-muted)] flex items-center gap-1"
title={`输入 ${formatTokenCount(topic.tokenStats.prompt_tokens)} / 输出 ${formatTokenCount(topic.tokenStats.completion_tokens)}`}
>
<Coins className="h-3 w-3" />
{formatTokenCount(topic.tokenStats.total_tokens)}
</span>
{(() => {
const pct = contextOccupancyPct(topic.tokenStats!);
if (pct == null) return null;
return (
<span
className={`text-xs flex items-center gap-1 ${occupancyColor(pct)}`}
title={`上下文窗口占用 ${pct}%${formatTokenCount(topic.tokenStats!.last_prompt_tokens!)} / ${formatTokenCount(topic.tokenStats!.context_window_tokens)}`}
>
ctx {pct}%
</span>
);
})()}
</>
)}
</div>
</div>
{topic.id === currentTopicId && (

View File

@ -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,
}));
}

View File

@ -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 {