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:
parent
07247ed140
commit
2e4b1931a6
@ -350,6 +350,7 @@ fn filter_images_by_age_and_count(
|
|||||||
tool_state: message.tool_state.clone(),
|
tool_state: message.tool_state.clone(),
|
||||||
tool_duration_ms: message.tool_duration_ms,
|
tool_duration_ms: message.tool_duration_ms,
|
||||||
tool_calls: message.tool_calls.clone(),
|
tool_calls: message.tool_calls.clone(),
|
||||||
|
usage: message.usage.clone(),
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -1214,6 +1215,11 @@ impl AgentLoop {
|
|||||||
if had_streaming {
|
if had_streaming {
|
||||||
assistant_message.id = streaming_message_id;
|
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());
|
messages.push(assistant_message.clone());
|
||||||
emitted_messages.push(assistant_message);
|
emitted_messages.push(assistant_message);
|
||||||
self.emit_live_tool_call_message(
|
self.emit_live_tool_call_message(
|
||||||
@ -1379,6 +1385,11 @@ impl AgentLoop {
|
|||||||
if had_streaming {
|
if had_streaming {
|
||||||
assistant_message.id = streaming_message_id.to_string();
|
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());
|
emitted_messages.push(assistant_message.clone());
|
||||||
self.emit_live_tool_call_message(assistant_message.clone())
|
self.emit_live_tool_call_message(assistant_message.clone())
|
||||||
.await;
|
.await;
|
||||||
@ -1481,12 +1492,17 @@ impl AgentLoop {
|
|||||||
|
|
||||||
match final_result {
|
match final_result {
|
||||||
Ok(response) => {
|
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)
|
ChatMessage::assistant_with_reasoning(response.content, reasoning_content)
|
||||||
} else {
|
} else {
|
||||||
ChatMessage::assistant(response.content)
|
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());
|
emitted_messages.push(assistant_message.clone());
|
||||||
self.emit_live_tool_call_message(assistant_message.clone())
|
self.emit_live_tool_call_message(assistant_message.clone())
|
||||||
.await;
|
.await;
|
||||||
|
|||||||
@ -66,6 +66,38 @@ pub struct ChatMessage {
|
|||||||
pub tool_duration_ms: Option<u64>,
|
pub tool_duration_ms: Option<u64>,
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
pub tool_calls: Option<Vec<ToolCall>>,
|
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 {
|
impl ChatMessage {
|
||||||
@ -83,6 +115,7 @@ impl ChatMessage {
|
|||||||
tool_duration_ms: None,
|
tool_duration_ms: None,
|
||||||
tool_state: None,
|
tool_state: None,
|
||||||
tool_calls: None,
|
tool_calls: None,
|
||||||
|
usage: None,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -100,6 +133,7 @@ impl ChatMessage {
|
|||||||
tool_duration_ms: None,
|
tool_duration_ms: None,
|
||||||
tool_state: None,
|
tool_state: None,
|
||||||
tool_calls: None,
|
tool_calls: None,
|
||||||
|
usage: None,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -117,6 +151,7 @@ impl ChatMessage {
|
|||||||
tool_duration_ms: None,
|
tool_duration_ms: None,
|
||||||
tool_state: None,
|
tool_state: None,
|
||||||
tool_calls: None,
|
tool_calls: None,
|
||||||
|
usage: None,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -146,6 +181,7 @@ impl ChatMessage {
|
|||||||
tool_duration_ms: None,
|
tool_duration_ms: None,
|
||||||
tool_state: None,
|
tool_state: None,
|
||||||
tool_calls: Some(tool_calls),
|
tool_calls: Some(tool_calls),
|
||||||
|
usage: None,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -180,6 +216,7 @@ impl ChatMessage {
|
|||||||
tool_duration_ms: None,
|
tool_duration_ms: None,
|
||||||
tool_state: None,
|
tool_state: None,
|
||||||
tool_calls: None,
|
tool_calls: None,
|
||||||
|
usage: None,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -215,6 +252,7 @@ impl ChatMessage {
|
|||||||
tool_duration_ms: None,
|
tool_duration_ms: None,
|
||||||
tool_state: Some(tool_state),
|
tool_state: Some(tool_state),
|
||||||
tool_calls: None,
|
tool_calls: None,
|
||||||
|
usage: None,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@ -1,7 +1,7 @@
|
|||||||
use crate::command::Command;
|
use crate::command::Command;
|
||||||
use crate::command::context::CommandContext;
|
use crate::command::context::CommandContext;
|
||||||
use crate::command::handler::{CommandHandler, CommandMetadata};
|
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::command::response::{CommandError, CommandResponse, MessageKind};
|
||||||
use crate::gateway::session::SessionManager;
|
use crate::gateway::session::SessionManager;
|
||||||
use crate::storage::SessionStore;
|
use crate::storage::SessionStore;
|
||||||
@ -87,18 +87,7 @@ async fn handle_delete_topic(
|
|||||||
.list_topics(session_id)
|
.list_topics(session_id)
|
||||||
.map_err(|e| CommandError::new("LIST_TOPICS_ERROR", e.to_string()))?;
|
.map_err(|e| CommandError::new("LIST_TOPICS_ERROR", e.to_string()))?;
|
||||||
|
|
||||||
let topic_summaries: Vec<TopicSummary> = topics
|
let topic_summaries = build_topic_summaries(handler.store.as_ref(), 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 topics_json = serde_json::to_string(&topic_summaries)
|
let topics_json = serde_json::to_string(&topic_summaries)
|
||||||
.map_err(|e| CommandError::new("SERIALIZE_ERROR", e.to_string()))?;
|
.map_err(|e| CommandError::new("SERIALIZE_ERROR", e.to_string()))?;
|
||||||
|
|||||||
@ -2,11 +2,27 @@ use crate::command::Command;
|
|||||||
use crate::command::context::CommandContext;
|
use crate::command::context::CommandContext;
|
||||||
use crate::command::handler::{CommandHandler, CommandMetadata};
|
use crate::command::handler::{CommandHandler, CommandMetadata};
|
||||||
use crate::command::response::{CommandError, CommandResponse, MessageKind};
|
use crate::command::response::{CommandError, CommandResponse, MessageKind};
|
||||||
use crate::storage::SessionStore;
|
use crate::storage::{SessionStore, SessionTokenStats, TopicRecord};
|
||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
|
use std::collections::HashMap;
|
||||||
use std::sync::Arc;
|
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<u32>,
|
||||||
|
pub context_window_tokens: u32,
|
||||||
|
}
|
||||||
|
|
||||||
/// Topic 摘要信息
|
/// Topic 摘要信息
|
||||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
pub struct TopicSummary {
|
pub struct TopicSummary {
|
||||||
@ -17,6 +33,64 @@ pub struct TopicSummary {
|
|||||||
pub message_count: i64,
|
pub message_count: i64,
|
||||||
pub created_at: i64,
|
pub created_at: i64,
|
||||||
pub last_active_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 命令处理器
|
/// 列出 Session 的 Topics 命令处理器
|
||||||
@ -66,18 +140,9 @@ async fn handle_list_topics(
|
|||||||
.list_topics(&session_id)
|
.list_topics(&session_id)
|
||||||
.map_err(|e| CommandError::new("LIST_TOPICS_ERROR", e.to_string()))?;
|
.map_err(|e| CommandError::new("LIST_TOPICS_ERROR", e.to_string()))?;
|
||||||
|
|
||||||
let summaries: Vec<TopicSummary> = topics
|
// context_window_tokens 当前未从配置链路注入(前端可从已有 config 接口获取),
|
||||||
.into_iter()
|
// 此处传 0 表示"后端不提供上限",前端按需隐藏百分比。
|
||||||
.map(|t| TopicSummary {
|
let summaries = build_topic_summaries(handler.store.as_ref(), topics)?;
|
||||||
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 topics_json = serde_json::to_string(&summaries)
|
let topics_json = serde_json::to_string(&summaries)
|
||||||
.map_err(|e| CommandError::new("SERIALIZE_ERROR", e.to_string()))?;
|
.map_err(|e| CommandError::new("SERIALIZE_ERROR", e.to_string()))?;
|
||||||
|
|||||||
@ -1,7 +1,7 @@
|
|||||||
use crate::command::Command;
|
use crate::command::Command;
|
||||||
use crate::command::context::CommandContext;
|
use crate::command::context::CommandContext;
|
||||||
use crate::command::handler::{CommandHandler, CommandMetadata};
|
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::command::response::{CommandError, CommandResponse, MessageKind};
|
||||||
use crate::storage::SessionStore;
|
use crate::storage::SessionStore;
|
||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
@ -83,14 +83,16 @@ async fn handle_rename_topic(
|
|||||||
.store
|
.store
|
||||||
.list_topics(session_id)
|
.list_topics(session_id)
|
||||||
.map_err(|e| CommandError::new("LIST_TOPICS_ERROR", e.to_string()))?;
|
.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)
|
return Ok(CommandResponse::success(ctx.request_id)
|
||||||
.with_message(
|
.with_message(
|
||||||
MessageKind::Notification,
|
MessageKind::Notification,
|
||||||
&format!("✓ 话题标题未变化: {}", trimmed_title),
|
&format!("✓ 话题标题未变化: {}", trimmed_title),
|
||||||
)
|
)
|
||||||
.with_metadata("topics", &topic_summaries)
|
.with_metadata("topics", &topic_summaries_json)
|
||||||
.with_metadata("topic_id", &topic_id)
|
.with_metadata("topic_id", &topic_id)
|
||||||
.with_metadata("title", trimmed_title)
|
.with_metadata("title", trimmed_title)
|
||||||
.with_metadata("session_id", session_id));
|
.with_metadata("session_id", session_id));
|
||||||
@ -108,34 +110,20 @@ async fn handle_rename_topic(
|
|||||||
.list_topics(session_id)
|
.list_topics(session_id)
|
||||||
.map_err(|e| CommandError::new("LIST_TOPICS_ERROR", e.to_string()))?;
|
.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);
|
let message = format!("✓ 已重命名话题: {} → {}", old_title, trimmed_title);
|
||||||
|
|
||||||
Ok(CommandResponse::success(ctx.request_id)
|
Ok(CommandResponse::success(ctx.request_id)
|
||||||
.with_message(MessageKind::Notification, &message)
|
.with_message(MessageKind::Notification, &message)
|
||||||
.with_metadata("topics", &topic_summaries)
|
.with_metadata("topics", &topic_summaries_json)
|
||||||
.with_metadata("topic_id", &topic_id)
|
.with_metadata("topic_id", &topic_id)
|
||||||
.with_metadata("title", trimmed_title)
|
.with_metadata("title", trimmed_title)
|
||||||
.with_metadata("session_id", session_id))
|
.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)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
|
|||||||
@ -1,7 +1,7 @@
|
|||||||
use crate::command::Command;
|
use crate::command::Command;
|
||||||
use crate::command::context::CommandContext;
|
use crate::command::context::CommandContext;
|
||||||
use crate::command::handler::{CommandHandler, CommandMetadata};
|
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::command::response::{CommandError, CommandResponse, MessageKind};
|
||||||
use crate::gateway::session::SessionManager;
|
use crate::gateway::session::SessionManager;
|
||||||
use crate::storage::SessionStore;
|
use crate::storage::SessionStore;
|
||||||
@ -109,18 +109,7 @@ async fn handle_create_session(
|
|||||||
.list_topics(session_id)
|
.list_topics(session_id)
|
||||||
.map_err(|e| CommandError::new("LIST_TOPICS_ERROR", e.to_string()))?;
|
.map_err(|e| CommandError::new("LIST_TOPICS_ERROR", e.to_string()))?;
|
||||||
|
|
||||||
let topic_summaries: Vec<TopicSummary> = topics
|
let topic_summaries = build_topic_summaries(handler.store.as_ref(), 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 topics_json = serde_json::to_string(&topic_summaries)
|
let topics_json = serde_json::to_string(&topic_summaries)
|
||||||
.map_err(|e| CommandError::new("SERIALIZE_ERROR", e.to_string()))?;
|
.map_err(|e| CommandError::new("SERIALIZE_ERROR", e.to_string()))?;
|
||||||
|
|||||||
@ -1,20 +1,45 @@
|
|||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
use std::sync::Arc;
|
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::bus::{MessageBus, OutboundMessage};
|
||||||
use crate::channels::base::{Channel, ChannelError};
|
use crate::channels::base::{Channel, ChannelError};
|
||||||
|
|
||||||
/// Consumes outbound messages from MessageBus and dispatches them to channels.
|
/// 每个 channel 的发送队列容量。
|
||||||
pub struct OutboundDispatcher {
|
///
|
||||||
bus: Arc<MessageBus>,
|
/// 略小于 MessageBus 的容量(100),确保 bus 的 `try_send` 丢消息
|
||||||
channels: Arc<RwLock<HashMap<String, Arc<dyn Channel + Send + Sync>>>>,
|
/// 防线仍有效——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.
|
/// Prefix for virtual scheduler chat IDs that should not be sent to external channels.
|
||||||
const SCHEDULER_VIRTUAL_CHAT_ID_PREFIX: &str = "scheduler/";
|
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 {
|
impl OutboundDispatcher {
|
||||||
pub fn new(bus: Arc<MessageBus>) -> Self {
|
pub fn new(bus: Arc<MessageBus>) -> Self {
|
||||||
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>) {
|
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
|
self.channels
|
||||||
.write()
|
.write()
|
||||||
.await
|
.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) {
|
pub async fn run(&self) {
|
||||||
@ -41,6 +117,7 @@ impl OutboundDispatcher {
|
|||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
#[cfg(debug_assertions)]
|
#[cfg(debug_assertions)]
|
||||||
tracing::debug!(
|
tracing::debug!(
|
||||||
channel = %msg.channel,
|
channel = %msg.channel,
|
||||||
@ -52,6 +129,7 @@ impl OutboundDispatcher {
|
|||||||
// Skip messages with virtual scheduler chat IDs (e.g., "scheduler/job_id")
|
// Skip messages with virtual scheduler chat IDs (e.g., "scheduler/job_id")
|
||||||
// These are internal messages from SilentAgentTask that should not be sent externally
|
// These are internal messages from SilentAgentTask that should not be sent externally
|
||||||
if msg.chat_id.starts_with(SCHEDULER_VIRTUAL_CHAT_ID_PREFIX) {
|
if msg.chat_id.starts_with(SCHEDULER_VIRTUAL_CHAT_ID_PREFIX) {
|
||||||
|
#[cfg(debug_assertions)]
|
||||||
tracing::debug!(
|
tracing::debug!(
|
||||||
channel = %msg.channel,
|
channel = %msg.channel,
|
||||||
chat_id = %msg.chat_id,
|
chat_id = %msg.chat_id,
|
||||||
@ -61,12 +139,26 @@ impl OutboundDispatcher {
|
|||||||
}
|
}
|
||||||
|
|
||||||
let channel_name = msg.channel.clone();
|
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 {
|
match sink {
|
||||||
Some(ch) => {
|
Some(sink) => {
|
||||||
if let Err(error) = self.send_with_retry(&*ch, msg).await {
|
// try_send 保证 dispatcher 永不阻塞:队列满时立即丢弃该消息,
|
||||||
tracing::error!(channel = %channel_name, error = %error, "Failed to send message after retries");
|
// 不影响其他 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 => {
|
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(
|
async fn send_with_retry(
|
||||||
&self,
|
|
||||||
channel: &dyn Channel,
|
channel: &dyn Channel,
|
||||||
msg: OutboundMessage,
|
msg: OutboundMessage,
|
||||||
) -> Result<(), ChannelError> {
|
) -> Result<(), ChannelError> {
|
||||||
const DELAYS: [u64; 3] = [1, 2, 4];
|
for (attempt_index, delay) in RETRY_DELAYS_SECS.iter().enumerate() {
|
||||||
|
|
||||||
for (attempt_index, delay) in DELAYS.iter().enumerate() {
|
|
||||||
match channel.send(msg.clone()).await {
|
match channel.send(msg.clone()).await {
|
||||||
Ok(()) => return Ok(()),
|
Ok(()) => return Ok(()),
|
||||||
Err(error) if attempt_index < DELAYS.len() - 1 => {
|
Err(error) if attempt_index < RETRY_DELAYS_SECS.len() - 1 => {
|
||||||
tracing::warn!(
|
tracing::warn!(
|
||||||
attempt = attempt_index + 1,
|
attempt = attempt_index + 1,
|
||||||
delay = delay,
|
delay = delay,
|
||||||
@ -102,3 +202,269 @@ impl OutboundDispatcher {
|
|||||||
unreachable!()
|
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;
|
||||||
|
});
|
||||||
|
|
||||||
|
// 先发一条 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<MessageBus> 克隆,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<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();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@ -42,6 +42,22 @@ pub struct TopicSummary {
|
|||||||
pub last_active_at: i64,
|
pub last_active_at: i64,
|
||||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||||
pub description: Option<String>,
|
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)]
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
|||||||
@ -32,6 +32,8 @@ struct StreamingAccumulator {
|
|||||||
reasoning_content: Option<String>,
|
reasoning_content: Option<String>,
|
||||||
tool_calls: BTreeMap<usize, StreamingToolCall>,
|
tool_calls: BTreeMap<usize, StreamingToolCall>,
|
||||||
response_id: String,
|
response_id: String,
|
||||||
|
/// 流式末帧返回的 usage(需要 stream_options.include_usage=true)
|
||||||
|
usage: Option<OpenAIUsage>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl StreamingAccumulator {
|
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
|
/// 构建最终的 ChatCompletionResponse
|
||||||
fn build_response(self, model: String) -> ChatCompletionResponse {
|
fn build_response(self, model: String) -> ChatCompletionResponse {
|
||||||
let tool_calls: Vec<ToolCall> = self
|
let tool_calls: Vec<ToolCall> = self
|
||||||
@ -116,11 +126,15 @@ impl StreamingAccumulator {
|
|||||||
content: self.content,
|
content: self.content,
|
||||||
reasoning_content: self.reasoning_content,
|
reasoning_content: self.reasoning_content,
|
||||||
tool_calls,
|
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,
|
prompt_tokens: 0,
|
||||||
completion_tokens: 0,
|
completion_tokens: 0,
|
||||||
total_tokens: 0,
|
total_tokens: 0,
|
||||||
},
|
}),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@ -385,6 +399,8 @@ impl OpenAIProvider {
|
|||||||
let mut body = self.build_request_body(request);
|
let mut body = self.build_request_body(request);
|
||||||
// 启用流式输出
|
// 启用流式输出
|
||||||
body["stream"] = json!(true);
|
body["stream"] = json!(true);
|
||||||
|
// 请求在流式末帧返回 usage(DeepSeek/OpenAI 兼容协议)
|
||||||
|
body["stream_options"] = json!({ "include_usage": true });
|
||||||
|
|
||||||
let mut req_builder = self
|
let mut req_builder = self
|
||||||
.client
|
.client
|
||||||
@ -469,6 +485,17 @@ impl OpenAIProvider {
|
|||||||
accumulator.set_response_id(id.to_string());
|
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::<OpenAIUsage>(usage_val.clone())
|
||||||
|
{
|
||||||
|
accumulator.set_usage(u);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// 提取 choices
|
// 提取 choices
|
||||||
if let Some(choices) = json.get("choices").and_then(|c| c.as_array()) {
|
if let Some(choices) = json.get("choices").and_then(|c| c.as_array()) {
|
||||||
for choice in choices {
|
for choice in choices {
|
||||||
@ -582,6 +609,17 @@ impl OpenAIProvider {
|
|||||||
accumulator.set_response_id(id.to_string());
|
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()) {
|
if let Some(choices) = json.get("choices").and_then(|c| c.as_array()) {
|
||||||
for choice in choices {
|
for choice in choices {
|
||||||
// 尝试从 delta 提取
|
// 尝试从 delta 提取
|
||||||
@ -691,6 +729,12 @@ impl OpenAIProvider {
|
|||||||
.collect()
|
.collect()
|
||||||
})
|
})
|
||||||
.unwrap_or_default();
|
.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,
|
arguments: OAIFunctionArguments,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Deserialize, Default)]
|
#[derive(Deserialize, Default, Clone, Debug)]
|
||||||
struct OpenAIUsage {
|
struct OpenAIUsage {
|
||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
prompt_tokens: u32,
|
prompt_tokens: u32,
|
||||||
|
|||||||
@ -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 索引(如果不存在)
|
// 创建 topic_id 索引(如果不存在)
|
||||||
conn.execute(
|
conn.execute(
|
||||||
"CREATE INDEX IF NOT EXISTS idx_messages_topic_seq ON messages(topic_id, seq) WHERE topic_id IS NOT NULL",
|
"CREATE INDEX IF NOT EXISTS idx_messages_topic_seq ON messages(topic_id, seq) WHERE topic_id IS NOT NULL",
|
||||||
|
|||||||
@ -1,6 +1,8 @@
|
|||||||
#[cfg(not(test))]
|
#[cfg(not(test))]
|
||||||
use std::path::{Path, PathBuf};
|
use std::path::{Path, PathBuf};
|
||||||
|
|
||||||
|
use std::collections::HashMap;
|
||||||
|
|
||||||
use r2d2::Pool;
|
use r2d2::Pool;
|
||||||
use r2d2_sqlite::SqliteConnectionManager;
|
use r2d2_sqlite::SqliteConnectionManager;
|
||||||
use rusqlite::{Connection, OptionalExtension, TransactionBehavior, params};
|
use rusqlite::{Connection, OptionalExtension, TransactionBehavior, params};
|
||||||
@ -25,8 +27,8 @@ pub use ports::{
|
|||||||
};
|
};
|
||||||
pub use records::{
|
pub use records::{
|
||||||
ALLOWED_MEMORY_NAMESPACES, GLOBAL_SCOPE_KEY, MemoryRecord, MemoryUpsert, SchedulerJobRecord,
|
ALLOWED_MEMORY_NAMESPACES, GLOBAL_SCOPE_KEY, MemoryRecord, MemoryUpsert, SchedulerJobRecord,
|
||||||
SchedulerJobState, SchedulerJobStatus, SchedulerJobUpsert, SessionRecord, SkillEventRecord,
|
SchedulerJobState, SchedulerJobStatus, SchedulerJobUpsert, SessionRecord, SessionTokenStats,
|
||||||
TodoRecord, TopicRecord, allowed_namespace_names, get_namespace_description,
|
SkillEventRecord, TodoRecord, TopicRecord, allowed_namespace_names, get_namespace_description,
|
||||||
is_valid_namespace,
|
is_valid_namespace,
|
||||||
};
|
};
|
||||||
|
|
||||||
@ -102,6 +104,11 @@ impl SessionStore {
|
|||||||
tool_call_id TEXT,
|
tool_call_id TEXT,
|
||||||
tool_name TEXT,
|
tool_name TEXT,
|
||||||
tool_calls_json 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,
|
created_at INTEGER NOT NULL,
|
||||||
FOREIGN KEY(session_id) REFERENCES sessions(id) ON DELETE CASCADE,
|
FOREIGN KEY(session_id) REFERENCES sessions(id) ON DELETE CASCADE,
|
||||||
FOREIGN KEY(topic_id) REFERENCES topics(id) ON DELETE SET NULL,
|
FOREIGN KEY(topic_id) REFERENCES topics(id) ON DELETE SET NULL,
|
||||||
@ -599,8 +606,8 @@ impl SessionStore {
|
|||||||
"
|
"
|
||||||
INSERT INTO messages (
|
INSERT INTO messages (
|
||||||
id, session_id, topic_id, seq, role, content,
|
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
|
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)
|
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18)
|
||||||
",
|
",
|
||||||
params![
|
params![
|
||||||
message.id,
|
message.id,
|
||||||
@ -616,6 +623,10 @@ impl SessionStore {
|
|||||||
message.tool_name,
|
message.tool_name,
|
||||||
tool_calls_json,
|
tool_calls_json,
|
||||||
message.tool_duration_ms.map(|v| v as i64),
|
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,
|
message.timestamp,
|
||||||
],
|
],
|
||||||
)?;
|
)?;
|
||||||
@ -677,8 +688,8 @@ impl SessionStore {
|
|||||||
INSERT INTO messages (
|
INSERT INTO messages (
|
||||||
id, session_id, topic_id, seq, role, content,
|
id, session_id, topic_id, seq, role, content,
|
||||||
system_context, reasoning_content, media_refs_json,
|
system_context, reasoning_content, media_refs_json,
|
||||||
tool_call_id, tool_name, tool_calls_json, tool_duration_ms, created_at
|
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)
|
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18)
|
||||||
",
|
",
|
||||||
params![
|
params![
|
||||||
message.id,
|
message.id,
|
||||||
@ -694,6 +705,10 @@ impl SessionStore {
|
|||||||
message.tool_name,
|
message.tool_name,
|
||||||
tool_calls_json,
|
tool_calls_json,
|
||||||
message.tool_duration_ms.map(|v| v as i64),
|
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,
|
message.timestamp,
|
||||||
],
|
],
|
||||||
)?;
|
)?;
|
||||||
@ -1551,7 +1566,7 @@ impl SessionStore {
|
|||||||
if let Some(sid) = session_id {
|
if let Some(sid) = session_id {
|
||||||
let mut stmt = conn.prepare(
|
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
|
FROM messages
|
||||||
WHERE topic_id = ?1 AND session_id = ?2
|
WHERE topic_id = ?1 AND session_id = ?2
|
||||||
ORDER BY seq ASC
|
ORDER BY seq ASC
|
||||||
@ -1566,7 +1581,7 @@ impl SessionStore {
|
|||||||
} else {
|
} else {
|
||||||
let mut stmt = conn.prepare(
|
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
|
FROM messages
|
||||||
WHERE topic_id = ?1
|
WHERE topic_id = ?1
|
||||||
ORDER BY seq ASC
|
ORDER BY seq ASC
|
||||||
@ -1615,6 +1630,96 @@ impl SessionStore {
|
|||||||
.map_err(StorageError::from)
|
.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(
|
pub fn replace_todos(
|
||||||
&self,
|
&self,
|
||||||
scope_key: &str,
|
scope_key: &str,
|
||||||
@ -1750,8 +1855,8 @@ fn insert_message_with_seq(
|
|||||||
"
|
"
|
||||||
INSERT INTO messages (
|
INSERT INTO messages (
|
||||||
id, session_id, seq, role, content,
|
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
|
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)
|
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17)
|
||||||
",
|
",
|
||||||
params![
|
params![
|
||||||
message.id,
|
message.id,
|
||||||
@ -1766,6 +1871,10 @@ fn insert_message_with_seq(
|
|||||||
message.tool_name,
|
message.tool_name,
|
||||||
tool_calls_json,
|
tool_calls_json,
|
||||||
message.tool_duration_ms.map(|v| v as i64),
|
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,
|
message.timestamp,
|
||||||
],
|
],
|
||||||
)?;
|
)?;
|
||||||
@ -1794,8 +1903,8 @@ fn insert_message_with_topic_seq(
|
|||||||
"
|
"
|
||||||
INSERT INTO messages (
|
INSERT INTO messages (
|
||||||
id, session_id, topic_id, seq, role, content,
|
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
|
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)
|
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18)
|
||||||
",
|
",
|
||||||
params![
|
params![
|
||||||
message.id,
|
message.id,
|
||||||
@ -1811,6 +1920,10 @@ fn insert_message_with_topic_seq(
|
|||||||
message.tool_name,
|
message.tool_name,
|
||||||
tool_calls_json,
|
tool_calls_json,
|
||||||
message.tool_duration_ms.map(|v| v as i64),
|
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,
|
message.timestamp,
|
||||||
],
|
],
|
||||||
)?;
|
)?;
|
||||||
@ -1831,6 +1944,8 @@ fn clone_message_for_compaction(message: &ChatMessage, timestamp: i64) -> ChatMe
|
|||||||
tool_state: message.tool_state.clone(),
|
tool_state: message.tool_state.clone(),
|
||||||
tool_duration_ms: message.tool_duration_ms,
|
tool_duration_ms: message.tool_duration_ms,
|
||||||
tool_calls: message.tool_calls.clone(),
|
tool_calls: message.tool_calls.clone(),
|
||||||
|
// 压缩克隆不保留 usage:压缩产生的是合成消息,不代表真实 LLM 调用
|
||||||
|
usage: None,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -1842,7 +1957,7 @@ fn load_messages_between(
|
|||||||
) -> Result<Vec<ChatMessage>, StorageError> {
|
) -> Result<Vec<ChatMessage>, StorageError> {
|
||||||
let mut stmt = conn.prepare(
|
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
|
FROM messages
|
||||||
WHERE session_id = ?1 AND seq > ?2 AND seq <= ?3
|
WHERE session_id = ?1 AND seq > ?2 AND seq <= ?3
|
||||||
ORDER BY seq ASC
|
ORDER BY seq ASC
|
||||||
@ -1888,6 +2003,7 @@ fn load_messages_between(
|
|||||||
tool_state: None,
|
tool_state: None,
|
||||||
tool_duration_ms: row.get::<_, Option<i64>>(10)?.map(|v| v as u64),
|
tool_duration_ms: row.get::<_, Option<i64>>(10)?.map(|v| v as u64),
|
||||||
tool_calls,
|
tool_calls,
|
||||||
|
usage: map_usage_row(row, 11, 12, 13, 14)?,
|
||||||
})
|
})
|
||||||
},
|
},
|
||||||
)?;
|
)?;
|
||||||
@ -1906,7 +2022,7 @@ fn load_messages_after(
|
|||||||
) -> Result<Vec<ChatMessage>, StorageError> {
|
) -> Result<Vec<ChatMessage>, StorageError> {
|
||||||
let mut stmt = conn.prepare(
|
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
|
FROM messages
|
||||||
WHERE session_id = ?1 AND seq > ?2
|
WHERE session_id = ?1 AND seq > ?2
|
||||||
ORDER BY seq ASC
|
ORDER BY seq ASC
|
||||||
@ -1949,6 +2065,7 @@ fn load_messages_after(
|
|||||||
tool_state: None,
|
tool_state: None,
|
||||||
tool_duration_ms: row.get::<_, Option<i64>>(10)?.map(|v| v as u64),
|
tool_duration_ms: row.get::<_, Option<i64>>(10)?.map(|v| v as u64),
|
||||||
tool_calls,
|
tool_calls,
|
||||||
|
usage: map_usage_row(row, 11, 12, 13, 14)?,
|
||||||
})
|
})
|
||||||
})?;
|
})?;
|
||||||
|
|
||||||
|
|||||||
@ -111,6 +111,16 @@ pub struct TopicRecord {
|
|||||||
pub message_count: i64,
|
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)]
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
pub struct MemoryRecord {
|
pub struct MemoryRecord {
|
||||||
pub id: String,
|
pub id: String,
|
||||||
|
|||||||
@ -8,12 +8,41 @@
|
|||||||
use rusqlite::{Connection, OptionalExtension, params};
|
use rusqlite::{Connection, OptionalExtension, params};
|
||||||
|
|
||||||
use crate::bus::ChatMessage;
|
use crate::bus::ChatMessage;
|
||||||
|
use crate::bus::message::MessageUsage;
|
||||||
|
|
||||||
use super::{
|
use super::{
|
||||||
MemoryRecord, SchedulerJobRecord, SchedulerJobState, SchedulerJobStatus, SessionRecord,
|
MemoryRecord, SchedulerJobRecord, SchedulerJobState, SchedulerJobStatus, SessionRecord,
|
||||||
SkillEventRecord, StorageError,
|
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(
|
pub(super) fn get_session_with_conn(
|
||||||
conn: &Connection,
|
conn: &Connection,
|
||||||
session_id: &str,
|
session_id: &str,
|
||||||
@ -147,6 +176,7 @@ pub(super) fn map_chat_message_row(row: &rusqlite::Row<'_>) -> rusqlite::Result<
|
|||||||
tool_state: None,
|
tool_state: None,
|
||||||
tool_duration_ms: row.get::<_, Option<i64>>(10)?.map(|v| v as u64),
|
tool_duration_ms: row.get::<_, Option<i64>>(10)?.map(|v| v as u64),
|
||||||
tool_calls,
|
tool_calls,
|
||||||
|
usage: map_usage_row(row, 11, 12, 13, 14)?,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@ -12,8 +12,9 @@ import {
|
|||||||
ChevronLeft,
|
ChevronLeft,
|
||||||
ChevronRight,
|
ChevronRight,
|
||||||
Edit2,
|
Edit2,
|
||||||
|
Coins,
|
||||||
} from 'lucide-react';
|
} from 'lucide-react';
|
||||||
import type { Topic } from '../../types/protocol';
|
import type { Topic, TopicTokenStats } from '../../types/protocol';
|
||||||
|
|
||||||
interface TopicListProps {
|
interface TopicListProps {
|
||||||
sessionId: string | null;
|
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({
|
export function TopicList({
|
||||||
sessionId,
|
sessionId,
|
||||||
topics,
|
topics,
|
||||||
@ -253,7 +276,7 @@ export function TopicList({
|
|||||||
>
|
>
|
||||||
{topic.description || topic.title}
|
{topic.description || topic.title}
|
||||||
</div>
|
</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">
|
<span className="text-xs text-[var(--text-muted)] flex items-center gap-1">
|
||||||
<Hash className="h-3 w-3" />
|
<Hash className="h-3 w-3" />
|
||||||
{topic.message_count} 条消息
|
{topic.message_count} 条消息
|
||||||
@ -262,6 +285,29 @@ export function TopicList({
|
|||||||
<Clock className="h-3 w-3" />
|
<Clock className="h-3 w-3" />
|
||||||
{formatTime(topic.updated_at)}
|
{formatTime(topic.updated_at)}
|
||||||
</span>
|
</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>
|
||||||
</div>
|
</div>
|
||||||
{topic.id === currentTopicId && (
|
{topic.id === currentTopicId && (
|
||||||
|
|||||||
@ -40,6 +40,7 @@ function mapTopicSummaries(summaries: TopicSummary[]): Topic[] {
|
|||||||
message_count: Number(t.message_count),
|
message_count: Number(t.message_count),
|
||||||
created_at: t.created_at,
|
created_at: t.created_at,
|
||||||
updated_at: t.last_active_at,
|
updated_at: t.last_active_at,
|
||||||
|
tokenStats: t.token_stats,
|
||||||
}));
|
}));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@ -149,6 +149,14 @@ export interface SessionSaved {
|
|||||||
filepath: string;
|
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 {
|
export interface TopicSummary {
|
||||||
topic_id: string;
|
topic_id: string;
|
||||||
session_id: string;
|
session_id: string;
|
||||||
@ -157,6 +165,7 @@ export interface TopicSummary {
|
|||||||
message_count: number;
|
message_count: number;
|
||||||
created_at: number;
|
created_at: number;
|
||||||
last_active_at: number;
|
last_active_at: number;
|
||||||
|
token_stats?: TopicTokenStats;
|
||||||
}
|
}
|
||||||
|
|
||||||
export interface TopicList {
|
export interface TopicList {
|
||||||
@ -508,6 +517,7 @@ export interface Topic {
|
|||||||
message_count: number;
|
message_count: number;
|
||||||
created_at: number;
|
created_at: number;
|
||||||
updated_at: number;
|
updated_at: number;
|
||||||
|
tokenStats?: TopicTokenStats;
|
||||||
}
|
}
|
||||||
|
|
||||||
export interface Session {
|
export interface Session {
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user