ooodc/PicoBot
ooodc/PicoBot
1
0
Fork 0
Code Pull Requests Activity

feat(web): token 统计从侧边栏迁移到右侧面板，支持子代理每轮刷新

Browse Source 
- 新增 TopicTokenStatsPanel 面板，显示五字段 + 上下文占用进度条

- 移除 TopicList 侧边栏的 token 统计显示

- 后端新增 get_session_token_stats，按 session_id 精确查询子代理 token

- TaskMessagesLoaded 协议扩展 token_stats 字段，ws.rs 转发至前端

- useMessages/useSubAgentView 每轮 assistant_response/execution_completed 触发刷新

- App.tsx 根据当前视图分派 list_topics（主代理）或 load_task_messages（子代理）

- 抽取 utils/tokenStats.ts 共享格式化与占用率计算函数
This commit is contained in:
oudecheng 2026-08-05 12:09:40 +08:00
parent beb7b581a6
commit bb2774c6b5
13 changed files with 262 additions and 60 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
src/command/handlers/load_task_messages.rs
View File

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

5
src/gateway/ws.rs
View File

@ -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::<crate::protocol::TopicTokenStats>(json).ok());
let _ = sender
.send(WsOutbound::TaskMessagesLoaded {
@ -616,6 +620,7 @@ async fn handle_inbound(
subagent_type,
status,
summary,
token_stats,
})
.await;
}

4
src/protocol/mod.rs
View File

@ -288,6 +288,10 @@ pub enum WsOutbound {
status: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
summary: Option<String>,
/// 子代理 session 的 token 用量统计（cost 累计 + context 瞬时）。
/// 无 assistant 消息时为 None。前端用于在子代理视图下显示 token 统计。
#[serde(default, skip_serializing_if = "Option::is_none")]
token_stats: Option<TopicTokenStats>,
},
#[serde(rename = "scheduler_job_list")]
SchedulerJobList { jobs: Vec<SchedulerJobSummary> },

70
src/storage/mod.rs
View File

@ -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<Option<SessionTokenStats>, 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<i64>>(0)?,
row.get::<_, Option<i64>>(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,

27
web/src/App.tsx
View File

@ -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(() => {
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() {
</button>
) : (
<div className="w-80 h-full flex flex-col">
{/* Token 用量面板（主代理/子代理共用，按视图切换数据源） */}
<TopicTokenStatsPanel
tokenStats={subAgentView ? subAgentView.tokenStats : currentTopic?.tokenStats}
/>
{/* Tab 栏 */}
<div className="shrink-0 flex border-b border-[var(--border-color)]">
<button

92
web/src/components/Panel/TopicTokenStatsPanel.tsx Normal file
View File

@ -0,0 +1,92 @@
import { Coins } from 'lucide-react';
import type { TopicTokenStats } from '../../types/protocol';
import { formatTokenCount, contextOccupancyPct, occupancyColor } from '../../utils/tokenStats';
interface TopicTokenStatsPanelProps {
tokenStats?: TopicTokenStats | null;
}
export function TopicTokenStatsPanel({ tokenStats }: TopicTokenStatsPanelProps) {
// 无数据态
if (!tokenStats || tokenStats.total_tokens === 0) {
return (
<div className="shrink-0 border-b border-[var(--border-color)] p-3">
<div className="flex items-center gap-1.5 text-xs font-medium text-[var(--text-muted)] mb-2">
<Coins className="h-3.5 w-3.5" />
Token 用量
</div>
<p className="text-xs text-[var(--text-muted)]">暂无 token 数据</p>
</div>
);
}
const pct = contextOccupancyPct(tokenStats);
return (
<div className="shrink-0 border-b border-[var(--border-color)] p-3">
{/* 标题 */}
<div className="flex items-center gap-1.5 text-xs font-medium text-[var(--text-secondary)] mb-2">
<Coins className="h-3.5 w-3.5 text-[var(--accent-cyan)]" />
Token 用量
</div>
{/* 五字段网格 */}
<div className="grid grid-cols-2 gap-x-3 gap-y-1 text-xs">
<div className="flex justify-between">
<span className="text-[var(--text-muted)]">输入</span>
<span className="text-[var(--text-secondary)] font-mono">
{formatTokenCount(tokenStats.prompt_tokens)}
</span>
</div>
<div className="flex justify-between">
<span className="text-[var(--text-muted)]">输出</span>
<span className="text-[var(--text-secondary)] font-mono">
{formatTokenCount(tokenStats.completion_tokens)}
</span>
</div>
<div className="flex justify-between">
<span className="text-[var(--text-muted)]">总计</span>
<span className="text-[var(--text-primary)] font-mono font-semibold">
{formatTokenCount(tokenStats.total_tokens)}
</span>
</div>
<div className="flex justify-between">
<span className="text-[var(--text-muted)]">上次输入</span>
<span className="text-[var(--text-secondary)] font-mono">
{tokenStats.last_prompt_tokens != null
? formatTokenCount(tokenStats.last_prompt_tokens)
: '—'}
</span>
</div>
<div className="flex justify-between col-span-2">
<span className="text-[var(--text-muted)]">上下文窗口</span>
<span className="text-[var(--text-secondary)] font-mono">
{formatTokenCount(tokenStats.context_window_tokens)}
</span>
</div>
</div>
{/* 上下文占用百分比 + 进度条 */}
{pct != null && (
<div className="mt-2">
<div className="flex items-center justify-between text-xs mb-1">
<span className="text-[var(--text-muted)]">上下文占用</span>
<span className={`font-mono font-medium ${occupancyColor(pct)}`}>ctx {pct}%</span>
</div>
<div className="h-1.5 rounded-full bg-[var(--overlay-subtle)] overflow-hidden">
<div
className={`h-full rounded-full transition-all ${
pct >= 80
? 'bg-red-400'
: pct >= 50
? 'bg-amber-400'
: 'bg-emerald-400'
}`}
style={{ width: `${pct}%` }}
/>
</div>
</div>
)}
</div>
);
}

48
web/src/components/Sidebar/TopicList.tsx
View File

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

3
web/src/hooks/chat/types.ts
View File

@ -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;
}
/** 定时任务执行对话查看视图 */

11
web/src/hooks/chat/useMessages.ts
View File

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

13
web/src/hooks/chat/useSubAgentView.ts
View File

@ -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<SubAgentView[]>([]);
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;
});

1
web/src/hooks/useChat.ts
View File

@ -136,6 +136,7 @@ export function useChat(): UseChatReturn {
const subAgent = useSubAgentView({
sendCommand: conn.sendCommand,
requestSubAgentTodoList: sideData.requestSubAgentTodoList,
bumpTopicRefreshTrigger: topics.bumpTopicRefreshTrigger,
});
const scheduler = useSchedulerView();

1
web/src/types/protocol.ts
View File

@ -272,6 +272,7 @@ export interface TaskMessagesLoaded {
subagent_type: string;
status: string;
summary?: string;
token_stats?: TopicTokenStats;
}
export interface ExecutionCancelled {

23
web/src/utils/tokenStats.ts Normal file
View File

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