diff --git a/src/command/handlers/load_task_messages.rs b/src/command/handlers/load_task_messages.rs index 798af18..0d49c80 100644 --- a/src/command/handlers/load_task_messages.rs +++ b/src/command/handlers/load_task_messages.rs @@ -1,6 +1,7 @@ use crate::command::Command; use crate::command::context::CommandContext; use crate::command::handler::{CommandHandler, CommandMetadata}; +use crate::command::handlers::list_topics::TopicTokenStats; use crate::command::response::{CommandError, CommandResponse}; use crate::storage::SessionStore; use crate::tools::task::repository::TaskRepository; @@ -100,6 +101,19 @@ async fn handle_load_task_messages( let status = format!("{:?}", task.state).to_lowercase(); + // 查询子代理 session 的 token 统计（按 session_id 精确匹配，不过滤 sub:%） + let token_stats = handler + .store + .get_session_token_stats(&task.session_id) + .map_err(|e| CommandError::new("TOKEN_STATS_ERROR", e.to_string()))? + .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), + }); + let mut response = CommandResponse::success(ctx.request_id) .with_metadata("task_session_id", &task.session_id) .with_metadata("task_id", &task.id) @@ -111,6 +125,12 @@ async fn handle_load_task_messages( response = response.with_metadata("task_summary", summary); } + if let Some(ref stats) = token_stats { + let stats_json = serde_json::to_string(stats) + .map_err(|e| CommandError::new("SERIALIZE_ERROR", e.to_string()))?; + response = response.with_metadata("task_token_stats", &stats_json); + } + Ok(response) } diff --git a/src/gateway/ws.rs b/src/gateway/ws.rs index 960040d..d755ea0 100644 --- a/src/gateway/ws.rs +++ b/src/gateway/ws.rs @@ -608,6 +608,10 @@ async fn handle_inbound( .cloned() .unwrap_or_default(); let summary = response.metadata.get("task_summary").cloned(); + let token_stats = response + .metadata + .get("task_token_stats") + .and_then(|json| serde_json::from_str::(json).ok()); let _ = sender .send(WsOutbound::TaskMessagesLoaded { @@ -616,6 +620,7 @@ async fn handle_inbound( subagent_type, status, summary, + token_stats, }) .await; } diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs index ed14bc8..86966da 100644 --- a/src/protocol/mod.rs +++ b/src/protocol/mod.rs @@ -288,6 +288,10 @@ pub enum WsOutbound { status: String, #[serde(default, skip_serializing_if = "Option::is_none")] summary: Option, + /// 子代理 session 的 token 用量统计（cost 累计 + context 瞬时）。 + /// 无 assistant 消息时为 None。前端用于在子代理视图下显示 token 统计。 + #[serde(default, skip_serializing_if = "Option::is_none")] + token_stats: Option, }, #[serde(rename = "scheduler_job_list")] SchedulerJobList { jobs: Vec }, diff --git a/src/storage/mod.rs b/src/storage/mod.rs index b1243e7..68800da 100644 --- a/src/storage/mod.rs +++ b/src/storage/mod.rs @@ -1740,6 +1740,76 @@ impl SessionStore { Ok(stats) } + /// 查询单个 session 的 token 消耗统计（cost 累计 + context 瞬时）。 + /// + /// 按 `session_id` 精确匹配查询，**不过滤** `sub:%`——专门用于子代理 session + /// （session_id = `sub:...`）的 token 统计。子代理没有 topic，一个 session 即 + /// 一个完整执行单元，因此按 session 维度聚合而非 topic 维度。 + /// + /// 与 `batch_topic_token_stats` 共享同一套 SQL 模式（SUM + last），差异仅在 + /// WHERE 条件：单 session 精确匹配，无 `NOT LIKE 'sub:%'` 过滤，无 topic 维度。 + pub fn get_session_token_stats( + &self, + session_id: &str, + ) -> Result, StorageError> { + let conn = self.pool.get()?; + + // 1. SUM 查询：累计 prompt/completion/total + let sum_sql = "SELECT \ + COALESCE(SUM(prompt_tokens), 0), \ + COALESCE(SUM(completion_tokens), 0), \ + COALESCE(SUM(total_tokens), 0) \ + FROM messages \ + WHERE session_id = ?1 AND role = 'assistant'"; + let mut stmt = conn.prepare(sum_sql)?; + let sum_row = stmt.query_row(params![session_id], |row| { + Ok(( + row.get::<_, i64>(0)? as u64, + row.get::<_, i64>(1)? as u64, + row.get::<_, i64>(2)? as u64, + )) + })?; + let (prompt_tokens, completion_tokens, total_tokens) = sum_row; + + // 无 assistant 消息时直接返回 None + if total_tokens == 0 && prompt_tokens == 0 && completion_tokens == 0 { + // 需要二次确认是否真的没有 assistant 消息（usage 全 0 也可能是合法的） + let count_sql = "SELECT COUNT(*) FROM messages WHERE session_id = ?1 AND role = 'assistant'"; + let count: i64 = conn.query_row(count_sql, params![session_id], |row| row.get(0))?; + if count == 0 { + return Ok(None); + } + } + + // 2. last 查询：最新有 usage 的 assistant 消息的 prompt_tokens + context_window_tokens + let last_sql = "SELECT prompt_tokens, context_window_tokens \ + FROM messages \ + WHERE session_id = ?1 AND role = 'assistant' AND prompt_tokens IS NOT NULL \ + ORDER BY seq DESC LIMIT 1"; + let mut stmt2 = conn.prepare(last_sql)?; + let last_row = stmt2 + .query_row(params![session_id], |row| { + Ok(( + row.get::<_, Option>(0)?, + row.get::<_, Option>(1)?, + )) + }) + .optional()?; + + let (last_prompt_tokens, context_window_tokens) = match last_row { + Some((lp, lcw)) => (lp.map(|v| v as u32), lcw.map(|v| v as u32)), + None => (None, None), + }; + + Ok(Some(SessionTokenStats { + prompt_tokens, + completion_tokens, + total_tokens, + last_prompt_tokens, + context_window_tokens, + })) + } + pub fn replace_todos( &self, scope_key: &str, diff --git a/web/src/App.tsx b/web/src/App.tsx index 86955f4..8a2c9ad 100644 --- a/web/src/App.tsx +++ b/web/src/App.tsx @@ -18,6 +18,7 @@ import { SchedulerJobList } from './components/Sidebar/SchedulerJobList'; import { MemoryPanel } from './components/Panel/MemoryPanel'; import { SkillList } from './components/Panel/SkillList'; import { TodoPanel } from './components/Panel/TodoPanel'; +import { TopicTokenStatsPanel } from './components/Panel/TopicTokenStatsPanel'; import { getGatewaySettings, buildWsUrl, @@ -232,20 +233,34 @@ function App() { } }, [sessionId, status, handleCommand, sendMessage, requestTopicList]); - // 话题描述异步生成后自动刷新话题列表 + // token 统计刷新：根据当前视图分派 list_topics（主代理）或 load_task_messages（子代理） useEffect(() => { if (topicRefreshTrigger === 0) return; if (status !== 'connected') return; - const topicCmd = requestTopicList(); - if (!topicCmd) return; const timer = setTimeout(() => { - handleCommand(topicCmd); - sendMessage({ type: 'command', payload: JSON.stringify(topicCmd) }); + if (subAgentView) { + // 子代理视图：发 load_task_messages 刷新子代理 token_stats + const cmd = { type: 'load_task_messages' as const, task_id: subAgentView.taskId }; + handleCommand(cmd); + sendMessage({ type: 'command', payload: JSON.stringify(cmd) }); + } else { + // 主代理视图：发 list_topics 刷新 topic 列表 + token_stats + const topicCmd = requestTopicList(); + if (!topicCmd) return; + handleCommand(topicCmd); + sendMessage({ type: 'command', payload: JSON.stringify(topicCmd) }); + } }, 500); return () => clearTimeout(timer); - }, [topicRefreshTrigger, status, handleCommand, sendMessage, requestTopicList]); + }, [topicRefreshTrigger, status, subAgentView, handleCommand, sendMessage, requestTopicList]); + + // 当前选中 topic（用于右侧 Sidebar token 统计面板） + const currentTopic = useMemo( + () => topics.find((t) => t.id === selectedTopic), + [topics, selectedTopic], + ); // Topics 加载后，自动选择第一个（仅当用户尚未手动选择 topic 时） useEffect(() => { @@ -920,6 +935,10 @@ function App() { ) : (
+ {/* Token 用量面板（主代理/子代理共用，按视图切换数据源） */} + {/* Tab 栏 */}
+ ); + } + + const pct = contextOccupancyPct(tokenStats); + + return ( +
+ {/* 标题 */} +
+ + Token 用量 +
+ + {/* 五字段网格 */} +
+
+ 输入 + + {formatTokenCount(tokenStats.prompt_tokens)} + +
+
+ 输出 + + {formatTokenCount(tokenStats.completion_tokens)} + +
+
+ 总计 + + {formatTokenCount(tokenStats.total_tokens)} + +
+
+ 上次输入 + + {tokenStats.last_prompt_tokens != null + ? formatTokenCount(tokenStats.last_prompt_tokens) + : '—'} + +
+
+ 上下文窗口 + + {formatTokenCount(tokenStats.context_window_tokens)} + +
+
+ + {/* 上下文占用百分比 + 进度条 */} + {pct != null && ( +
+
+ 上下文占用 + ctx {pct}% +
+
+
= 80 + ? 'bg-red-400' + : pct >= 50 + ? 'bg-amber-400' + : 'bg-emerald-400' + }`} + style={{ width: `${pct}%` }} + /> +
+
+ )} +
+ ); +} diff --git a/web/src/components/Sidebar/TopicList.tsx b/web/src/components/Sidebar/TopicList.tsx index 752b033..6fe39ec 100644 --- a/web/src/components/Sidebar/TopicList.tsx +++ b/web/src/components/Sidebar/TopicList.tsx @@ -12,9 +12,8 @@ import { ChevronLeft, ChevronRight, Edit2, - Coins, } from 'lucide-react'; -import type { Topic, TopicTokenStats } from '../../types/protocol'; +import type { Topic } from '../../types/protocol'; interface TopicListProps { sessionId: string | null; @@ -44,28 +43,6 @@ 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, @@ -286,29 +263,6 @@ export function TopicList({ {formatTime(topic.updated_at)} - {topic.tokenStats && topic.tokenStats.total_tokens > 0 && ( - <> - - - {formatTokenCount(topic.tokenStats.total_tokens)} - - {(() => { - const pct = contextOccupancyPct(topic.tokenStats!); - if (pct == null) return null; - return ( - - ctx {pct}% - - ); - })()} - - )}
{topic.id === currentTopicId && ( diff --git a/web/src/hooks/chat/types.ts b/web/src/hooks/chat/types.ts index 8f4ed33..42a9309 100644 --- a/web/src/hooks/chat/types.ts +++ b/web/src/hooks/chat/types.ts @@ -1,4 +1,4 @@ -import type { ChatMessage } from '../../types/protocol'; +import type { ChatMessage, TopicTokenStats } from '../../types/protocol'; /** 子智能体视图（栈中的一层） */ export interface SubAgentView { @@ -8,6 +8,7 @@ export interface SubAgentView { status: string; summary?: string; messages: ChatMessage[]; + tokenStats?: TopicTokenStats; } /** 定时任务执行对话查看视图 */ diff --git a/web/src/hooks/chat/useMessages.ts b/web/src/hooks/chat/useMessages.ts index 66879e6..6df67f3 100644 --- a/web/src/hooks/chat/useMessages.ts +++ b/web/src/hooks/chat/useMessages.ts @@ -167,11 +167,14 @@ export function useMessages(options: UseMessagesOptions): UseMessagesReturn { case 'execution_completed': { const msg = message as ExecutionCompleted; - if (getSubagentTaskId(message)) return true; + if (getSubagentTaskId(message)) { + // 子代理执行完成：bump 统一 trigger，App.tsx 根据 subAgentView 分派 load_task_messages + bumpTopicRefreshTrigger(); + return true; + } if (msg.topic_id && msg.topic_id !== selectedTopicRef.current) return true; setIsLoading(false); // 主代理本次执行结束：刷新 topic 列表以更新 token 统计 - // （LLM usage 在消息持久化后已写入 DB，list_topics 会读到最新值） bumpTopicRefreshTrigger(); return true; } @@ -198,9 +201,9 @@ export function useMessages(options: UseMessagesOptions): UseMessagesReturn { } return [...prev, newMsg]; }); - // 当前话题无描述时，可能刚触发了异步生成，标记需要刷新 + // 每轮 assistant 响应到达即刷新 token 统计（复用 500ms 防抖） const currentTopic = topicsRef.current.find((t) => t.id === selectedTopicRef.current); - if (currentTopic && !currentTopic.description) { + if (currentTopic) { bumpTopicRefreshTrigger(); } if (msg.user_message_id) applyUserMessageId(msg.user_message_id); diff --git a/web/src/hooks/chat/useSubAgentView.ts b/web/src/hooks/chat/useSubAgentView.ts index 7560c60..49e44ed 100644 --- a/web/src/hooks/chat/useSubAgentView.ts +++ b/web/src/hooks/chat/useSubAgentView.ts @@ -26,6 +26,8 @@ interface UseSubAgentViewOptions { sendCommand: (cmd: Command) => void; /** 构建子代理待办刷新命令 */ requestSubAgentTodoList: (subTaskId: string) => Command; + /** 刷新 token 统计（子代理 assistant_response 到达时触发） */ + bumpTopicRefreshTrigger: () => void; } export interface UseSubAgentViewReturn { @@ -42,7 +44,7 @@ export interface UseSubAgentViewReturn { } export function useSubAgentView(options: UseSubAgentViewOptions): UseSubAgentViewReturn { - const { sendCommand, requestSubAgentTodoList } = options; + const { sendCommand, requestSubAgentTodoList, bumpTopicRefreshTrigger } = options; const [subAgentStack, setSubAgentStack] = useState([]); const subAgentView = useMemo( () => (subAgentStack.length > 0 ? subAgentStack[subAgentStack.length - 1] : null), @@ -64,6 +66,10 @@ export function useSubAgentView(options: UseSubAgentViewOptions): UseSubAgentVie // 追加消息到栈顶视图（含流式累加） const appendToSubAgentViewMessage = useCallback((message: WsOutbound) => { + // 子代理 assistant_response 到达时刷新 token 统计 + if (message.type === 'assistant_response') { + bumpTopicRefreshTrigger(); + } // stream_delta: accumulate into existing message by ID, or create new if (message.type === 'stream_delta') { const msg = message as StreamDelta; @@ -104,6 +110,8 @@ export function useSubAgentView(options: UseSubAgentViewOptions): UseSubAgentVie newStack[newStack.length - 1] = { ...top, status: 'completed' }; return newStack; }); + // 兜底刷新：防止最后一轮 assistant_response 丢失导致 token 统计停留旧值 + bumpTopicRefreshTrigger(); return; } // error: 更新栈顶 status 为 error，并追加错误消息 @@ -160,7 +168,7 @@ export function useSubAgentView(options: UseSubAgentViewOptions): UseSubAgentVie return newStack; }); } - }, []); + }, [bumpTopicRefreshTrigger]); // 追加消息到栈中非栈顶的匹配层（按 taskId 匹配） const appendToSubAgentLayerMessage = useCallback((taskId: string, message: WsOutbound) => { @@ -313,6 +321,7 @@ export function useSubAgentView(options: UseSubAgentViewOptions): UseSubAgentVie subagentType: msg.subagent_type, status: msg.status, summary: msg.summary, + tokenStats: msg.token_stats, }; return newStack; }); diff --git a/web/src/hooks/useChat.ts b/web/src/hooks/useChat.ts index d96cb35..5f9f2e0 100644 --- a/web/src/hooks/useChat.ts +++ b/web/src/hooks/useChat.ts @@ -136,6 +136,7 @@ export function useChat(): UseChatReturn { const subAgent = useSubAgentView({ sendCommand: conn.sendCommand, requestSubAgentTodoList: sideData.requestSubAgentTodoList, + bumpTopicRefreshTrigger: topics.bumpTopicRefreshTrigger, }); const scheduler = useSchedulerView(); diff --git a/web/src/types/protocol.ts b/web/src/types/protocol.ts index 611b3a9..5780495 100644 --- a/web/src/types/protocol.ts +++ b/web/src/types/protocol.ts @@ -272,6 +272,7 @@ export interface TaskMessagesLoaded { subagent_type: string; status: string; summary?: string; + token_stats?: TopicTokenStats; } export interface ExecutionCancelled { diff --git a/web/src/utils/tokenStats.ts b/web/src/utils/tokenStats.ts new file mode 100644 index 0000000..636f7aa --- /dev/null +++ b/web/src/utils/tokenStats.ts @@ -0,0 +1,23 @@ +import type { TopicTokenStats } from '../types/protocol'; + +/** 紧凑格式化 token 数量：1234 -> "1.2K"，1234567 -> "1.2M" */ +export 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 */ +export 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 */ +export function occupancyColor(pct: number): string { + if (pct >= 80) return 'text-red-400'; + if (pct >= 50) return 'text-amber-400'; + return 'text-emerald-400'; +}