From da5ee053116828bda4eac1bbaa2cf78bbbf0a804 Mon Sep 17 00:00:00 2001 From: xiaoxixi Date: Wed, 5 Aug 2026 16:45:56 +0800 Subject: [PATCH] feat: add session token and context statistics --- AGENTS.md | 1 + Cargo.toml | 2 +- README.md | 3 +- docs/ARCHITECTURE.md | 4 +- .../about-picobot/references/architecture.md | 3 +- .../about-picobot/references/db-schema.md | 2 + src/agent/agent_loop.rs | 16 +- src/agent/context_compressor.rs | 41 ++- src/channels/cli_chat.rs | 19 ++ src/client/mod.rs | 4 +- src/gateway/router.rs | 5 + src/protocol.rs | 47 ++++ src/session/commands.rs | 2 + src/session/events.rs | 2 + src/session/mod.rs | 2 + src/session/persistence.rs | 27 +- src/session/session.rs | 219 +++++++++------ src/session/stats.rs | 142 ++++++++++ src/storage/mod.rs | 262 +++++++++++++++++- src/storage/usage.rs | 28 ++ webui/package-lock.json | 4 +- webui/package.json | 2 +- webui/src/lib/chat.svelte.js | 9 + webui/src/lib/components/ActivitySpine.svelte | 176 +++++++++--- webui/src/pages/ChatPage.svelte | 26 +- 25 files changed, 883 insertions(+), 165 deletions(-) create mode 100644 src/session/stats.rs create mode 100644 src/storage/usage.rs diff --git a/AGENTS.md b/AGENTS.md index 9525427..2489473 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -97,6 +97,7 @@ Scheduler → SessionManager scheduled execution → AgentLoop → Scheduler del - **WebUI styling** uses the local Fluent 2 semantic tokens in `webui/src/styles.css`; page and component styles should consume the aliases instead of introducing independent hard-coded palettes, and must preserve selectable light/dark and brand-color themes, keyboard focus, responsive layout, and reduced-motion behavior. Browser-only appearance settings belong in `lib/theme.js`/`localStorage`, and `public/theme-init.js` must restore them before Svelte mounts - **Gateway config reload** validates and prepares a complete next runtime generation before retiring the old one; close runtime admission before draining inbound lanes, Sessions, Scheduler jobs, and background sub-agents; MCP connects only during activation; host/port, workspace, and effective storage path remain restart-only, and runtime reload must never mutate process environment variables - **WebUI slash completion** must consume the existing `get_slash_commands` WebSocket response; do not duplicate the backend command list in frontend source +- **Session token statistics** persist Provider-reported usage atomically with each completed Turn; WebUI `session_stats` and `/info [--json]` must consume the same SessionStats projection, and context occupancy must use the final request's prompt usage rather than accumulated Turn totals - **WebUI chat rendering** sanitizes Markdown before inserting HTML; durable `turn_committed` deltas calibrate normal terminal Turns without a full history reload, and history must preserve structured tool-call metadata so calls and results remain independently collapsible - **WebUI/TUI file transfer** streams bytes over authenticated HTTP and sends only short-lived upload IDs/attachment metadata over WebSocket; messages persist local media paths without guaranteeing later availability, and client responses must never expose those paths - **WebUI/TUI same-turn media delivery** stages same-session `send_message(files=...)` media on the active Turn and commits it on the final assistant message, after durable tool-call history; safe raster formats should render as an inline preview with download fallback diff --git a/Cargo.toml b/Cargo.toml index 18201bf..7a82c80 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "picobot" -version = "1.3.1" +version = "1.4.0" edition = "2024" [dependencies] diff --git a/README.md b/README.md index ee28a09..f984f14 100644 --- a/README.md +++ b/README.md @@ -12,6 +12,7 @@ PicoBot 是一个用 Rust 编写的个人 AI 助手运行时。它在本地启 - 从脚本或命令行发送一条任务，等待完整的模型/工具循环后只输出最终结果。 - 在 TUI 或浏览器中实时查看正文、思考过程和工具执行状态，并在完成后收敛到持久化历史。 - 在浏览器中查看日志、任务和记忆，修改运行配置与助手档案。 +- 在 WebUI 顶栏查看当前会话的累计输入/输出 Token、上下文窗口和占用比例。 - 复杂任务可创建 session 级 Todo 计划，把不同子项并行委托给多个子 Agent；聊天页侧栏实时显示进度。 - 将同一套 Agent 能力接入飞书/Lark，并可选用单张卡片实时更新回复。 - 让 Agent 使用本地文件、Shell、搜索、HTTP、浏览器、MCP 工具完成任务。 @@ -306,7 +307,7 @@ Session ID 使用三段式： | `/rename ` | 重命名当前 dialog | | `/delete` | 删除当前 dialog 并创建新 dialog | | `/compact` | 手动压缩上下文 | -| `/info` | 查看当前 dialog 信息 | +| `/info [--json]` | 查看当前 dialog、累计 Token 与上下文窗口信息；可选 JSON 输出 | | `/dump` | 导出当前 dialog 为 Markdown | | `/mcp` | 查看 MCP 服务器和工具状态 | | `/health` | 检查 PicoBot 运行依赖 | diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index b67eb39..9f73cd9 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -210,7 +210,7 @@ SessionManager 负责组装会话上下文：系统提示、Skills、召回的 K - 5 秒 busy timeout。 - schema version 迁移。 -持久化范围包括 sessions、messages、memories、task plans/items、scheduled jobs、job runs 和 background tasks。消息保存可展示 `reasoning_content`、Provider 私有回放状态、`turn_id`、iteration 和 completion status；私有 Provider 状态不进入 WebSocket/Channel，且只允许回放给同一 Provider。成功持久化一个 Turn 后，交互 Channel 收到只包含公开字段的 `CommittedTurnDelta`，其中 `history_revision` 是本批次最高 durable sequence；客户端按 revision 幂等合并，正常完成不重新加载整段历史，断线重连和失败/取消仍使用 `SessionHistory` 校准。Provider 诊断和失败审计只记录模型、消息数、工具数等请求摘要，不保存正文、reasoning 或签名 payload。修改 schema 时应： +持久化范围包括 sessions、messages、session turn usage、memories、task plans/items、scheduled jobs、job runs 和 background tasks。消息保存可展示 `reasoning_content`、Provider 私有回放状态、`turn_id`、iteration 和 completion status；私有 Provider 状态不进入 WebSocket/Channel，且只允许回放给同一 Provider。成功 Turn 的 Provider usage 与消息批次在同一事务中写入 `session_turn_usage`，以 `turn_id` 幂等累计会话输入、输出、缓存输入和请求数；升级前的历史没有可归属 usage，统计起点必须显式呈现。成功持久化一个 Turn 后，交互 Channel 收到只包含公开字段的 `CommittedTurnDelta`，其中 `history_revision` 是本批次最高 durable sequence；客户端按 revision 幂等合并，正常完成不重新加载整段历史，断线重连和失败/取消仍使用 `SessionHistory` 校准。Provider 诊断和失败审计只记录模型、消息数、工具数等请求摘要，不保存正文、reasoning 或签名 payload。修改 schema 时应： 1. 更新集中式 schema/迁移逻辑。 2. 保留已有数据库的升级路径。 @@ -247,7 +247,7 @@ Turn delivery 使用 `spawn_graceful`。全局取消发生时先停止读取快 ### WebUI 与管理 API -Gateway 在 `/` 提供随二进制编译的 HTML/CSS/JavaScript，不依赖外部 CDN 或前端运行服务。前端源码位于 `webui/`，由 Svelte 5 + Vite 构建，Bits UI 提供无样式的可访问组件原语。视觉层通过 `webui/src/styles.css` 中的本地 Fluent 2 语义令牌实现浅色/深色表面、六套品牌色、状态色、层级和控件状态；页面组件必须复用语义别名，不能把独立硬编码调色板或外部 Fluent 运行库引入发布产物。明暗模式和品牌色只保存在浏览器 `localStorage`，`theme-init.js` 必须在 Svelte 挂载前恢复 `data-theme` 与 `data-accent`，防止首屏颜色闪烁；这些外观选项不属于 Gateway 配置，也不跨设备同步。`build.rs` 监听前端源码、锁文件和构建配置，增量地将生产资源生成到 Cargo `OUT_DIR`，Rust 再从该目录编译嵌入；前端产物不进入仓库，最终用户使用发布二进制时不需要 Node.js。浏览器聊天继续使用 `/ws` 和 `cli_chat` 渠道，因此复用现有 dialog scope、每会话串行 worker、历史持久化和出站 lane；会话历史帧保留工具调用 ID、名称、参数和工具结果角色，WebUI 在正常完成时合并 `turn_committed` 增量并将工具信息渲染为默认折叠的工具卡片；聊天页的 Todo 侧栏默认隐藏，按 session 保存快照和未读状态，并通过结构化 `session_plan`/`plan_updated` 帧刷新，计划变化不会写入聊天历史；斜杠命令补全通过 `get_slash_commands` 获取 Gateway 的实时命令与别名清单，不在前端重复定义；WebUI 不直接调用 Provider 或 SessionManager。 +Gateway 在 `/` 提供随二进制编译的 HTML/CSS/JavaScript，不依赖外部 CDN 或前端运行服务。前端源码位于 `webui/`，由 Svelte 5 + Vite 构建，Bits UI 提供无样式的可访问组件原语。视觉层通过 `webui/src/styles.css` 中的本地 Fluent 2 语义令牌实现浅色/深色表面、六套品牌色、状态色、层级和控件状态；页面组件必须复用语义别名，不能把独立硬编码调色板或外部 Fluent 运行库引入发布产物。明暗模式和品牌色只保存在浏览器 `localStorage`，`theme-init.js` 必须在 Svelte 挂载前恢复 `data-theme` 与 `data-accent`，防止首屏颜色闪烁；这些外观选项不属于 Gateway 配置，也不跨设备同步。`build.rs` 监听前端源码、锁文件和构建配置，增量地将生产资源生成到 Cargo `OUT_DIR`，Rust 再从该目录编译嵌入；前端产物不进入仓库，最终用户使用发布二进制时不需要 Node.js。浏览器聊天继续使用 `/ws` 和 `cli_chat` 渠道，因此复用现有 dialog scope、每会话串行 worker、历史持久化和出站 lane；会话历史帧保留工具调用 ID、名称、参数和工具结果角色，WebUI 在正常完成时合并 `turn_committed` 增量并将工具信息渲染为默认折叠的工具卡片；聊天页的 Todo 侧栏默认隐藏，按 session 保存快照和未读状态，并通过结构化 `session_plan`/`plan_updated` 帧刷新，计划变化不会写入聊天历史；活动状态栏通过结构化 `session_stats` 展示当前 session 的已提交 Turn 用量与上下文占用，累计量来自 Provider usage，窗口占用明确区分 API 基准上的混合估算与纯字符估算；斜杠命令补全通过 `get_slash_commands` 获取 Gateway 的实时命令与别名清单，不在前端重复定义；WebUI 不直接调用 Provider 或 SessionManager。 WebUI/TUI 文件字节通过受鉴权的 HTTP 接口流式传输，WebSocket 只携带短期 `upload_id` 和结构化附件描述。`UploadRegistry` 在内存中按 `cli_chat` chat scope 校验并消费待发送上传；消息继续以 `media_refs` 保存 Gateway 本地路径，不建立永久附件资产。下载接口必须通过 client、session、message 和附件序号反查路径，不能接受客户端路径。历史附件路径失效属于正常状态，不得影响历史文本读取。待发送但未进入消息的上传由 `TaskSupervisor` 所有的限时清理任务回收。Agent 上下文会为所有媒体注入内部路径清单，客户端响应不得暴露该路径。 diff --git a/resources/skills/about-picobot/references/architecture.md b/resources/skills/about-picobot/references/architecture.md index 582c639..96c2c25 100644 --- a/resources/skills/about-picobot/references/architecture.md +++ b/resources/skills/about-picobot/references/architecture.md @@ -54,6 +54,7 @@ Scheduler → SessionManager.handle_cron_message → AgentLoop → send_message - 复杂任务可使用 `todo` 创建 session 级计划；多个子项可通过 `delegate.plan_item_id` 并行委托，子 Agent 不能修改计划 - WebUI 聊天复用 `/ws` 与 `cli_chat`；同源管理 API 只读取受限的日志、任务、记忆，并对白名单配置文件做原子写入 - WebUI 使用 Svelte 5 + Vite，Bits UI 提供无样式可访问组件；`cargo build` 增量生成前端到 Cargo `OUT_DIR`，再嵌入单二进制，仓库不保存生成产物 +- WebUI 通过 `get_session_stats`/`session_stats` 显示当前会话累计输入输出 Token 和上下文窗口占用；`/info [--json]` 读取同一份 SessionStats ## 关键约束 @@ -293,7 +294,7 @@ Gateway 关停顺序： | `/rename <title>` | 重命名当前对话 | | `/delete` | 删除当前对话 | | `/compact` | 手动触发上下文压缩 | -| `/info` | 显示当前对话信息 | +| `/info [--json]` | 显示当前对话、累计 Token 与上下文窗口信息；可选 JSON 输出 | | `/dump` | 保存当前对话为 markdown | | `/?`, `/help` | 显示帮助 | | `/mcp` | 显示 MCP 状态 | diff --git a/resources/skills/about-picobot/references/db-schema.md b/resources/skills/about-picobot/references/db-schema.md index 7887098..4229e29 100644 --- a/resources/skills/about-picobot/references/db-schema.md +++ b/resources/skills/about-picobot/references/db-schema.md @@ -24,6 +24,8 @@ | `last_consolidated_at` | INTEGER | 上次记忆归并时间 | | `last_compressed_message_at` | INTEGER | 上次上下文压缩边界时间戳 | +`session_turn_usage` 以 `turn_id` 幂等保存已提交 Turn 的 Provider usage，包括累计输入、输出、缓存输入、请求数和最后一次请求的 prompt tokens。它与 Turn 消息批次在同一事务中提交，供 WebUI 状态栏和 `/info` 使用；升级前历史无法可靠回填，因此统计起点以首条 usage 记录为准。 + `(channel, chat_id, dialog_id)` 唯一。普通列表排除 `deleted_at`；是否包含归档记录由查询参数决定。 ## messages 表 diff --git a/src/agent/agent_loop.rs b/src/agent/agent_loop.rs index ceaaf8a..57fd9aa 100644 --- a/src/agent/agent_loop.rs +++ b/src/agent/agent_loop.rs @@ -314,6 +314,10 @@ pub struct AgentProcessResult { pub emitted_messages: Vec<ChatMessage>, pub total_tokens: Option<u32>, pub usage: Option<crate::providers::Usage>, + /// Provider usage for the final successful request in this Turn. This is + /// the correct basis for context-window occupancy; `usage` is accumulated + /// across every tool iteration. + pub last_request_usage: Option<crate::providers::Usage>, } fn merge_usage(total: &mut crate::providers::Usage, next: &crate::providers::Usage) { @@ -691,6 +695,7 @@ impl AgentLoop { let mut emitted_messages = Vec::new(); let mut accumulated_tokens: u32 = 0; let mut accumulated_usage = crate::providers::Usage::default(); + let mut last_request_usage = None; for iteration in 0..self.max_iterations { #[cfg(debug_assertions)] @@ -741,6 +746,7 @@ impl AgentLoop { accumulated_tokens = accumulated_tokens.saturating_add(response.usage.total_tokens); merge_usage(&mut accumulated_usage, &response.usage); + last_request_usage = Some(response.usage.clone()); #[cfg(debug_assertions)] tracing::debug!( @@ -766,6 +772,7 @@ impl AgentLoop { emitted_messages, total_tokens: Some(accumulated_tokens), usage: Some(accumulated_usage), + last_request_usage, }); } @@ -897,6 +904,7 @@ impl AgentLoop { Ok(response) => { accumulated_tokens = accumulated_tokens.saturating_add(response.usage.total_tokens); merge_usage(&mut accumulated_usage, &response.usage); + last_request_usage = Some(response.usage.clone()); let mut assistant_message = ChatMessage::assistant(response.content); assistant_message.reasoning_content = response.reasoning_content; assistant_message.provider_state = response.provider_state; @@ -916,6 +924,7 @@ impl AgentLoop { emitted_messages, total_tokens: Some(accumulated_tokens), usage: Some(accumulated_usage), + last_request_usage, }) } Err(e) => { @@ -944,10 +953,8 @@ impl AgentLoop { Self::annotate_message(&mut final_message, turn.as_ref(), summary_iteration, true); emitted_messages.push(final_message.clone()); let turn_usage = (accumulated_usage.total_tokens > 0).then_some(&accumulated_usage); - crate::observability::metrics::global_metrics().record_turn( - turn_usage, - turn_start.elapsed().as_millis() as u64, - ); + crate::observability::metrics::global_metrics() + .record_turn(turn_usage, turn_start.elapsed().as_millis() as u64); Ok(AgentProcessResult { final_response: final_message, emitted_messages, @@ -957,6 +964,7 @@ impl AgentLoop { None }, usage: (accumulated_usage.total_tokens > 0).then_some(accumulated_usage), + last_request_usage, }) } } diff --git a/src/agent/context_compressor.rs b/src/agent/context_compressor.rs index 14ba089..37cc9ca 100644 --- a/src/agent/context_compressor.rs +++ b/src/agent/context_compressor.rs @@ -70,8 +70,8 @@ pub struct ContextCompressor { session_id: Option<String>, /// Message count sent in the last LLM call (used to split known/new history). last_sent_message_count: Option<usize>, - /// Real total_tokens from the last API response. - last_api_total_tokens: Option<u32>, + /// Real prompt_tokens from the final API request in the last completed Turn. + last_api_prompt_tokens: Option<u32>, } /// Result of context compression. @@ -85,7 +85,7 @@ pub struct TokenInfo { pub context_window: usize, pub threshold: usize, pub estimated_tokens: usize, - pub last_api_tokens: Option<u32>, + pub last_prompt_tokens: Option<u32>, pub cache_active: bool, } @@ -104,7 +104,7 @@ impl ContextCompressor { memory, session_id: None, last_sent_message_count: None, - last_api_total_tokens: None, + last_api_prompt_tokens: None, } } @@ -123,7 +123,7 @@ impl ContextCompressor { memory, session_id: None, last_sent_message_count: None, - last_api_total_tokens: None, + last_api_prompt_tokens: None, } } @@ -137,24 +137,28 @@ impl ContextCompressor { self.context_window = window; } + pub fn context_window(&self) -> usize { + self.context_window + } + /// Record the API's reported token usage from the last completed turn. /// `msg_count`: number of messages sent to LLM in that call. - /// `tokens`: `total_tokens` from the API response. + /// `tokens`: `prompt_tokens` from the final API request in the Turn. pub fn set_last_api_info(&mut self, msg_count: usize, tokens: Option<u32>) { self.last_sent_message_count = Some(msg_count); - self.last_api_total_tokens = tokens; + self.last_api_prompt_tokens = tokens; } /// Invalidate the cached API token info — called after compression modifies messages. fn invalidate_token_cache(&mut self) { self.last_sent_message_count = None; - self.last_api_total_tokens = None; + self.last_api_prompt_tokens = None; } /// Hybrid token estimation: API-reported tokens for known history + /// char/4 estimate for new messages since last API call. fn token_estimate_with_history(&self, messages: &[ChatMessage]) -> usize { - match (self.last_api_total_tokens, self.last_sent_message_count) { + match (self.last_api_prompt_tokens, self.last_sent_message_count) { (Some(known), Some(known_count)) if messages.len() > known_count => { let delta = &messages[known_count..]; known as usize + estimate_tokens(delta) @@ -175,8 +179,8 @@ impl ContextCompressor { context_window: self.context_window, threshold: self.threshold(), estimated_tokens: self.token_estimate_with_history(messages), - last_api_tokens: self.last_api_total_tokens, - cache_active: self.last_api_total_tokens.is_some(), + last_prompt_tokens: self.last_api_prompt_tokens, + cache_active: self.last_api_prompt_tokens.is_some(), } } @@ -797,6 +801,21 @@ mod tests { assert_eq!(compressor.threshold(), 89_600); } + #[test] + fn api_prompt_usage_is_the_base_for_new_history_estimates() { + let mut compressor = + ContextCompressor::new(mock_provider(), 128_000, test_memory_manager()); + compressor.set_last_api_info(1, Some(100)); + let messages = vec![ + ChatMessage::user("known"), + ChatMessage::assistant("new response"), + ]; + + let info = compressor.token_info(&messages); + assert_eq!(info.last_prompt_tokens, Some(100)); + assert_eq!(info.estimated_tokens, 100 + estimate_tokens(&messages[1..])); + } + #[tokio::test] async fn test_compress_if_needed_fast_trims_tool_results() { // context_window=200 → threshold=100. diff --git a/src/channels/cli_chat.rs b/src/channels/cli_chat.rs index 2005094..ec63aff 100644 --- a/src/channels/cli_chat.rs +++ b/src/channels/cli_chat.rs @@ -475,6 +475,25 @@ impl CliChatChannel { None => return Err(ChannelError::Other("Control channel closed".to_string())), } } + WsInbound::GetSessionStats { session_id } => { + let unified_id = Self::parse_client_session(&client, &session_id)?; + let (reply_tx, mut reply_rx) = mpsc::channel(1); + bus.publish_control(ControlMessage { + op: SessionCommand::GetSessionStats { + session_id: unified_id, + }, + reply_tx, + }) + .await?; + match reply_rx.recv().await { + Some(Ok(SessionEvent::SessionStats { stats })) => { + let _ = client.sender.send(WsOutbound::SessionStats { stats }).await; + } + Some(Ok(_)) => {} + Some(Err(error)) => return Err(error), + None => return Err(ChannelError::Other("Control channel closed".to_string())), + } + } WsInbound::RenameSession { session_id, title } => { let target = session_id .or(current_session_guard.clone()) diff --git a/src/client/mod.rs b/src/client/mod.rs index 8966059..14846b4 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -380,7 +380,9 @@ async fn handle_ws_message(app: &mut App, outbound: WsOutbound) { } => app.set_history(&session_id, messages), // The first Todo UI is WebUI-only. CLI keeps receiving ordinary task // notifications and may inspect plans through /todo. - WsOutbound::SessionPlan { .. } | WsOutbound::PlanUpdated { .. } => {} + WsOutbound::SessionPlan { .. } + | WsOutbound::SessionStats { .. } + | WsOutbound::PlanUpdated { .. } => {} WsOutbound::SessionRenamed { session_id, title } => { if let Some(session) = app .sessions diff --git a/src/gateway/router.rs b/src/gateway/router.rs index e9e4295..f69a9f8 100644 --- a/src/gateway/router.rs +++ b/src/gateway/router.rs @@ -325,6 +325,11 @@ async fn handle_control_message(session_manager: &SessionManager, message: Contr .await .map(|plan| SessionEvent::TaskPlan { session_id, plan }) .map_err(|error| ChannelError::Other(error.to_string())), + GetSessionStats { session_id } => session_manager + .get_session_stats(&session_id) + .await + .map(|stats| SessionEvent::SessionStats { stats }) + .map_err(|error| ChannelError::Other(error.to_string())), RenameDialog { session_id, title } => session_manager .rename_dialog(&session_id, &title) .await diff --git a/src/protocol.rs b/src/protocol.rs index a408eac..1741b3d 100644 --- a/src/protocol.rs +++ b/src/protocol.rs @@ -172,6 +172,8 @@ pub enum WsInbound { }, #[serde(rename = "get_session_plan")] GetSessionPlan { session_id: String }, + #[serde(rename = "get_session_stats")] + GetSessionStats { session_id: String }, #[serde(rename = "rename_session")] RenameSession { #[serde(default, skip_serializing_if = "Option::is_none")] @@ -249,6 +251,8 @@ pub enum WsOutbound { session_id: String, plan: Option<crate::work::TaskPlan>, }, + #[serde(rename = "session_stats")] + SessionStats { stats: crate::session::SessionStats }, #[serde(rename = "plan_updated")] PlanUpdated { session_id: String, @@ -362,4 +366,47 @@ mod tests { crate::bus::CompletionStatus::Completed ); } + + #[test] + fn session_stats_request_and_response_use_structured_frames() { + let inbound = + parse_inbound(r#"{"type":"get_session_stats","session_id":"cli_chat:client:dialog"}"#) + .unwrap(); + assert!(matches!(inbound, WsInbound::GetSessionStats { .. })); + + let stats = crate::session::SessionStats { + session_id: "cli_chat:client:dialog".into(), + title: "stats".into(), + provider: "provider".into(), + model: "model".into(), + user_message_count: 1, + history_message_count: 2, + lifetime_usage: crate::session::LifetimeUsage { + input_tokens: 100, + output_tokens: 20, + total_tokens: 120, + cached_input_tokens: Some(40), + request_count: 1, + turn_count: 1, + tracked_since: Some(1), + }, + context: crate::session::ContextUsage { + configured_window_tokens: 128_000, + effective_window_tokens: 128_000, + used_tokens: 100, + remaining_tokens: 127_900, + compression_threshold_tokens: 89_600, + source: crate::session::ContextUsageSource::Hybrid, + last_observed_prompt_tokens: Some(90), + observed_at: Some(1), + }, + created_at: 1, + last_active_at: 2, + updated_at: 3, + }; + let value = serde_json::to_value(WsOutbound::SessionStats { stats }).unwrap(); + assert_eq!(value["type"], "session_stats"); + assert_eq!(value["stats"]["context"]["source"], "hybrid"); + assert_eq!(value["stats"]["lifetime_usage"]["input_tokens"], 100); + } } diff --git a/src/session/commands.rs b/src/session/commands.rs index 24b35db..6ec9cc0 100644 --- a/src/session/commands.rs +++ b/src/session/commands.rs @@ -28,6 +28,8 @@ pub enum SessionCommand { }, /// Load the active task plan for a dialog. GetTaskPlan { session_id: UnifiedSessionId }, + /// Load token totals and context-window state for a dialog. + GetSessionStats { session_id: UnifiedSessionId }, /// Get the current dialog for a chat GetCurrentDialog { channel: String, chat_id: String }, /// Rename a dialog diff --git a/src/session/events.rs b/src/session/events.rs index ec3a024..dfc791e 100644 --- a/src/session/events.rs +++ b/src/session/events.rs @@ -41,6 +41,8 @@ pub enum SessionEvent { session_id: UnifiedSessionId, plan: Option<crate::work::TaskPlan>, }, + /// Provider usage totals and current context-window state. + SessionStats { stats: crate::session::SessionStats }, /// Dialog renamed DialogRenamed { session_id: UnifiedSessionId, diff --git a/src/session/mod.rs b/src/session/mod.rs index 08b9c1b..cbf2a70 100644 --- a/src/session/mod.rs +++ b/src/session/mod.rs @@ -8,6 +8,7 @@ mod turn_input; #[allow(clippy::module_inception)] pub mod session; pub mod session_id; +pub mod stats; pub mod turn; pub use commands::SessionCommand; @@ -15,6 +16,7 @@ pub use error::SessionError; pub use events::{DialogInfo, SessionEvent}; pub use session::{SLASH_COMMANDS, Session, SessionManager, SessionManagerServices, SlashCommand}; pub use session_id::UnifiedSessionId; +pub use stats::{ContextUsage, ContextUsageSource, LifetimeUsage, SessionStats}; pub use turn::{ BlockId, ToolStatus, TurnBlock, TurnController, TurnId, TurnPhase, TurnSnapshot, TurnState, TurnStatus, diff --git a/src/session/persistence.rs b/src/session/persistence.rs index e3f51eb..7596283 100644 --- a/src/session/persistence.rs +++ b/src/session/persistence.rs @@ -10,6 +10,7 @@ use crate::{providers::Usage, session::TurnController}; async fn persist_added_messages( snapshots: Vec<Option<MessagePersistSnapshot>>, + usage: Option<&crate::storage::TurnUsageRecord>, ) -> Result<(), StorageError> { let mut storage = None; let mut session_id = None; @@ -35,9 +36,15 @@ async fn persist_added_messages( else { return Ok(()); }; - storage - .persist_message_batch_with_retry(&session_id, &messages, &final_meta) - .await + if let Some(usage) = usage { + storage + .persist_turn_batch_with_retry(&session_id, &messages, &final_meta, usage) + .await + } else { + storage + .persist_message_batch_with_retry(&session_id, &messages, &final_meta) + .await + } } pub(super) async fn append_persisted_messages( @@ -69,6 +76,7 @@ pub(super) async fn append_active_turn_message( session, vec![message], VersionPolicy::PreserveForOwnedTurn(turn_id), + None, ) .await .map(|_| ()) @@ -78,13 +86,22 @@ pub(super) async fn append_persisted_messages_with_meta( session: &Arc<Mutex<Session>>, messages: Vec<ChatMessage>, ) -> Result<Vec<crate::storage::message::MessageMeta>, StorageError> { - append_persisted_messages_inner(session, messages, VersionPolicy::Advance).await + append_persisted_messages_inner(session, messages, VersionPolicy::Advance, None).await +} + +pub(super) async fn append_persisted_turn_messages( + session: &Arc<Mutex<Session>>, + messages: Vec<ChatMessage>, + usage: crate::storage::TurnUsageRecord, +) -> Result<Vec<crate::storage::message::MessageMeta>, StorageError> { + append_persisted_messages_inner(session, messages, VersionPolicy::Advance, Some(usage)).await } async fn append_persisted_messages_inner( session: &Arc<Mutex<Session>>, messages: Vec<ChatMessage>, version_policy: VersionPolicy, + usage: Option<crate::storage::TurnUsageRecord>, ) -> Result<Vec<crate::storage::message::MessageMeta>, StorageError> { if messages.is_empty() { return Ok(Vec::new()); @@ -113,7 +130,7 @@ async fn append_persisted_messages_inner( .map(|(_, _, message, _)| message.clone()) .collect(); - if let Err(error) = persist_added_messages(snapshots).await { + if let Err(error) = persist_added_messages(snapshots, usage.as_ref()).await { session .lock() .await diff --git a/src/session/session.rs b/src/session/session.rs index af454cf..93cc109 100644 --- a/src/session/session.rs +++ b/src/session/session.rs @@ -4,7 +4,8 @@ use std::sync::Arc; use tokio::sync::{Mutex, mpsc, oneshot}; use super::persistence::{ - append_persisted_messages, append_persisted_messages_with_meta, finalize_turn_after_persistence, + append_persisted_messages, append_persisted_messages_with_meta, append_persisted_turn_messages, + finalize_turn_after_persistence, }; use super::turn::{TurnBlock, TurnController, TurnSnapshot}; use super::turn_input::prepare_turn_input; @@ -498,6 +499,7 @@ mod cancelled_partial_tests { ], total_tokens: None, usage: None, + last_request_usage: None, }; attach_pending_turn_deliveries( @@ -1113,7 +1115,7 @@ impl Session { }; let mut compressor = ContextCompressor::with_config( self.provider.clone(), - self.provider_config.token_limit, + self.compressor.context_window(), compressor_config, self.memory_manager.clone(), ); @@ -1155,7 +1157,7 @@ impl Session { self.provider_config.workspace_dir.clone(), self.provider_config.input_types.clone(), ) - .with_context_window(self.provider_config.token_limit)) + .with_context_window(self.compressor.context_window())) } /// 构建系统提示词（包含 AgentLoop 的基础提示词 + skills + memory） @@ -1856,87 +1858,15 @@ impl SessionManager { } "info" => { if let Some(sid) = current_session_id { - let session = self.get_or_create_session(sid).await?; - let session_guard = session.lock().await; - let history = session_guard.get_history(); - let message_count = history.len(); - let session_id_str = session_guard.session_id(); - let title = &session_guard.title; - let model_name = &session_guard.provider_config.name; - let created_at = - chrono::DateTime::from_timestamp_millis(session_guard.created_at) - .map(|dt| { - dt.with_timezone(&chrono::Local) - .format("%Y-%m-%d %H:%M:%S") - .to_string() - }) - .unwrap_or_default(); - let last_active_at = - chrono::DateTime::from_timestamp_millis(session_guard.last_active_at) - .map(|dt| { - dt.with_timezone(&chrono::Local) - .format("%Y-%m-%d %H:%M:%S") - .to_string() - }) - .unwrap_or_default(); - let token_info = session_guard.compressor.token_info(history); - let cache_info = if token_info.cache_active { - format!( - "API精确: {} tokens", - token_info.last_api_tokens.unwrap_or(0) - ) + let stats = self.get_session_stats(sid).await?; + let output = if args.is_some_and(|value| value.trim() == "--json") { + serde_json::to_string_pretty(&stats).map_err(|error| { + AgentError::Other(format!("failed to serialize session info: {error}")) + })? } else { - "无API精确缓存".to_string() + stats.render_text() }; - let threshold_pct = if token_info.context_window > 0 { - (token_info.threshold as f64 / token_info.context_window as f64 * 100.0) - as usize - } else { - 0 - }; - let usage_pct = if token_info.context_window > 0 { - (token_info.estimated_tokens as f64 / token_info.context_window as f64 - * 100.0) - .min(100.0) as usize - } else { - 0 - }; - let usage_bar = if token_info.context_window > 0 { - format!( - "{}/{} tokens ({}%)", - token_info.estimated_tokens, token_info.context_window, usage_pct - ) - } else { - "未设置".to_string() - }; - let compression_status = if token_info.estimated_tokens > token_info.threshold { - "[即将压缩]" - } else { - "[正常]" - }; - let ctx_info = format!( - "[窗口] {} [阈值] {}/{} ({}) [状态] {} {}", - usage_bar, - token_info.threshold, - token_info.context_window, - threshold_pct, - compression_status, - cache_info, - ); - Ok(( - None, - format!( - "对话标题: {}\nSession ID: {}\n模型: {}\n用户消息: {} / 总消息: {}\n创建时间: {}\n最后活跃: {}\n\n上下文: {}", - title, - session_id_str, - model_name, - session_guard.message_count, - message_count, - created_at, - last_active_at, - ctx_info, - ), - )) + Ok((None, output)) } else { Ok((None, "No active session.".to_string())) } @@ -2430,6 +2360,81 @@ impl SessionManager { .map_err(|error| AgentError::Other(format!("failed to load task plan: {error}"))) } + pub async fn get_session_stats( + &self, + session_id: &UnifiedSessionId, + ) -> Result<crate::session::SessionStats, AgentError> { + let session = self.get_or_create_session(session_id).await?; + let ( + title, + provider, + model, + user_message_count, + history_message_count, + created_at, + last_active_at, + configured_window, + token_info, + ) = { + let guard = session.lock().await; + let history = guard.get_history(); + ( + guard.title.clone(), + guard.provider_config.name.clone(), + guard.provider_config.model_id.clone(), + u64::try_from(guard.message_count).unwrap_or_default(), + u64::try_from(history.len()).unwrap_or(u64::MAX), + guard.created_at, + guard.last_active_at, + guard.provider_config.token_limit, + guard.compressor.token_info(history), + ) + }; + let totals = self + .storage + .get_session_usage_totals(&session_id.to_string()) + .await + .map_err(|error| AgentError::Other(format!("failed to load session usage: {error}")))?; + let effective_window = u64::try_from(token_info.context_window).unwrap_or(u64::MAX); + let used_tokens = u64::try_from(token_info.estimated_tokens).unwrap_or(u64::MAX); + + Ok(crate::session::SessionStats { + session_id: session_id.to_string(), + title, + provider, + model, + user_message_count, + history_message_count, + lifetime_usage: crate::session::LifetimeUsage { + input_tokens: totals.prompt_tokens, + output_tokens: totals.completion_tokens, + total_tokens: totals.total_tokens, + cached_input_tokens: totals.cached_input_tokens, + request_count: totals.request_count, + turn_count: totals.turn_count, + tracked_since: totals.tracked_since, + }, + context: crate::session::ContextUsage { + configured_window_tokens: u64::try_from(configured_window).unwrap_or(u64::MAX), + effective_window_tokens: effective_window, + used_tokens, + remaining_tokens: effective_window.saturating_sub(used_tokens), + compression_threshold_tokens: u64::try_from(token_info.threshold) + .unwrap_or(u64::MAX), + source: if token_info.cache_active { + crate::session::ContextUsageSource::Hybrid + } else { + crate::session::ContextUsageSource::Estimated + }, + last_observed_prompt_tokens: totals.last_prompt_tokens, + observed_at: totals.last_observed_at, + }, + created_at, + last_active_at, + updated_at: chrono::Utc::now().timestamp_millis(), + }) + } + pub async fn list_dialogs( &self, channel: &str, @@ -3246,9 +3251,12 @@ fn spawn_agent_worker( let pending = take_current_turn_deliveries(); attach_pending_turn_deliveries(&mut result, pending); let response_content = result.final_response.content; - let total_tokens = result.total_tokens; let usage = result.usage; - { + let last_prompt_tokens = result + .last_request_usage + .as_ref() + .map(|value| value.prompt_tokens); + let (provider_name, model_name) = { let guard = session2.lock().await; if guard.worker_generation != worker_gen || guard.state_version != base_version @@ -3258,21 +3266,48 @@ fn spawn_agent_worker( )); return; } - } + ( + guard.provider_config.name.clone(), + guard.provider_config.model_id.clone(), + ) + }; + let usage_record = usage.as_ref().map(|turn_usage| { + crate::storage::TurnUsageRecord { + session_id: response_session_id.clone(), + turn_id: agent_turn.turn_id.clone(), + provider: provider_name, + model: model_name, + usage: turn_usage.clone(), + last_prompt_tokens: last_prompt_tokens.unwrap_or_default(), + created_at: chrono::Utc::now().timestamp_millis(), + } + }); + let emitted_messages = result.emitted_messages; let response = match finalize_turn_after_persistence( turn_lifecycle, - usage, - append_persisted_messages_with_meta( - &session2, - result.emitted_messages, - ), + usage.clone(), + async { + if let Some(usage_record) = usage_record { + append_persisted_turn_messages( + &session2, + emitted_messages, + usage_record, + ) + .await + } else { + append_persisted_messages_with_meta(&session2, emitted_messages) + .await + } + }, ) .await { Ok(committed_messages) => { let mut guard = session2.lock().await; - let sent_count = guard.messages.len(); - guard.compressor.set_last_api_info(sent_count, total_tokens); + let prompt_message_count = guard.messages.len().saturating_sub(1); + guard + .compressor + .set_last_api_info(prompt_message_count, last_prompt_tokens); Some((response_content, committed_messages)) } Err(e) => { diff --git a/src/session/stats.rs b/src/session/stats.rs new file mode 100644 index 0000000..f8bc949 --- /dev/null +++ b/src/session/stats.rs @@ -0,0 +1,142 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct SessionStats { + pub session_id: String, + pub title: String, + pub provider: String, + pub model: String, + pub user_message_count: u64, + pub history_message_count: u64, + pub lifetime_usage: LifetimeUsage, + pub context: ContextUsage, + pub created_at: i64, + pub last_active_at: i64, + pub updated_at: i64, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct LifetimeUsage { + pub input_tokens: u64, + pub output_tokens: u64, + pub total_tokens: u64, + pub cached_input_tokens: Option<u64>, + pub request_count: u64, + pub turn_count: u64, + pub tracked_since: Option<i64>, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct ContextUsage { + pub configured_window_tokens: u64, + pub effective_window_tokens: u64, + pub used_tokens: u64, + pub remaining_tokens: u64, + pub compression_threshold_tokens: u64, + pub source: ContextUsageSource, + pub last_observed_prompt_tokens: Option<u64>, + pub observed_at: Option<i64>, +} + +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum ContextUsageSource { + Hybrid, + Estimated, +} + +impl ContextUsageSource { + pub fn label(self) -> &'static str { + match self { + Self::Hybrid => "混合估算", + Self::Estimated => "字符估算", + } + } +} + +impl SessionStats { + pub fn render_text(&self) -> String { + let percent = if self.context.effective_window_tokens == 0 { + 0.0 + } else { + self.context.used_tokens as f64 / self.context.effective_window_tokens as f64 * 100.0 + }; + let created_at = format_timestamp(self.created_at); + let last_active_at = format_timestamp(self.last_active_at); + let tracked_since = self + .lifetime_usage + .tracked_since + .map(format_timestamp) + .unwrap_or_else(|| "尚无已完成模型请求".to_string()); + let cached = self + .lifetime_usage + .cached_input_tokens + .map(format_tokens) + .unwrap_or_else(|| "—".to_string()); + let observed = self + .context + .last_observed_prompt_tokens + .map(format_tokens) + .unwrap_or_else(|| "—".to_string()); + + format!( + "会话\n 标题 {}\n ID {}\n 模型 {} / {}\n 消息 {} 条用户消息，{} 条历史消息\n 创建 {}\n 最后活跃 {}\n\nToken 用量 · 已提交 Turns\n 输入 {}\n 输出 {}\n 合计 {}\n 缓存输入 {}\n 请求 {}\n Turns {}\n 统计起点 {}\n\n上下文窗口 · {}\n 占用 {} / {}（{:.1}%）\n 剩余 {}\n 压缩阈值 {}（70%）\n 最近实测 {}", + self.title, + self.session_id, + self.provider, + self.model, + self.user_message_count, + self.history_message_count, + created_at, + last_active_at, + format_tokens(self.lifetime_usage.input_tokens), + format_tokens(self.lifetime_usage.output_tokens), + format_tokens(self.lifetime_usage.total_tokens), + cached, + format_tokens(self.lifetime_usage.request_count), + format_tokens(self.lifetime_usage.turn_count), + tracked_since, + self.context.source.label(), + format_tokens(self.context.used_tokens), + format_tokens(self.context.effective_window_tokens), + percent, + format_tokens(self.context.remaining_tokens), + format_tokens(self.context.compression_threshold_tokens), + observed, + ) + } +} + +fn format_timestamp(value: i64) -> String { + chrono::DateTime::from_timestamp_millis(value) + .map(|timestamp| { + timestamp + .with_timezone(&chrono::Local) + .format("%Y-%m-%d %H:%M:%S") + .to_string() + }) + .unwrap_or_else(|| "—".to_string()) +} + +fn format_tokens(value: u64) -> String { + let digits = value.to_string(); + let mut formatted = String::with_capacity(digits.len() + digits.len() / 3); + for (index, ch) in digits.chars().enumerate() { + if index > 0 && (digits.len() - index).is_multiple_of(3) { + formatted.push(','); + } + formatted.push(ch); + } + formatted +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn formats_large_token_counts() { + assert_eq!(format_tokens(1_234_567), "1,234,567"); + assert_eq!(format_tokens(12), "12"); + } +} diff --git a/src/storage/mod.rs b/src/storage/mod.rs index 6520855..37d4f79 100644 --- a/src/storage/mod.rs +++ b/src/storage/mod.rs @@ -4,10 +4,12 @@ pub mod memory; pub mod message; pub mod scheduler; pub mod session; +pub mod usage; pub use background_task::BackgroundTask; pub use error::StorageError; pub use scheduler::{DeliveryPolicy, JobKind, JobRun, ScheduledJob}; +pub use usage::{SessionUsageTotals, TurnUsageRecord}; use sqlx::sqlite::{ SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions, SqliteRow, SqliteSynchronous, @@ -16,7 +18,7 @@ use sqlx::{Pool, Row, Sqlite}; use std::path::Path; use tokio::time::{Duration, sleep}; -const SCHEMA_VERSION: i64 = 4; +const SCHEMA_VERSION: i64 = 5; const INSERT_MESSAGE_SQL: &str = r#" INSERT INTO messages ( id, session_id, seq, role, content, reasoning_content, provider_state, @@ -362,6 +364,33 @@ impl Storage { .execute(&self.pool) .await?; + sqlx::query( + r#" + CREATE TABLE IF NOT EXISTS session_turn_usage ( + turn_id TEXT PRIMARY KEY, + session_id TEXT NOT NULL, + provider TEXT NOT NULL, + model TEXT NOT NULL, + prompt_tokens INTEGER NOT NULL, + completion_tokens INTEGER NOT NULL, + total_tokens INTEGER NOT NULL, + cached_input_tokens INTEGER, + request_count INTEGER NOT NULL, + last_prompt_tokens INTEGER NOT NULL, + created_at INTEGER NOT NULL, + FOREIGN KEY (session_id) REFERENCES sessions(id) ON DELETE CASCADE + ) + "#, + ) + .execute(&self.pool) + .await?; + + sqlx::query( + "CREATE INDEX IF NOT EXISTS idx_session_turn_usage_session_created ON session_turn_usage(session_id, created_at)", + ) + .execute(&self.pool) + .await?; + Self::init_scheduler_schema(&self.pool).await?; self.migrate_schema().await?; @@ -467,6 +496,31 @@ impl Storage { ) .execute(&mut *tx) .await?; + sqlx::query( + r#" + CREATE TABLE IF NOT EXISTS session_turn_usage ( + turn_id TEXT PRIMARY KEY, + session_id TEXT NOT NULL, + provider TEXT NOT NULL, + model TEXT NOT NULL, + prompt_tokens INTEGER NOT NULL, + completion_tokens INTEGER NOT NULL, + total_tokens INTEGER NOT NULL, + cached_input_tokens INTEGER, + request_count INTEGER NOT NULL, + last_prompt_tokens INTEGER NOT NULL, + created_at INTEGER NOT NULL, + FOREIGN KEY (session_id) REFERENCES sessions(id) ON DELETE CASCADE + ) + "#, + ) + .execute(&mut *tx) + .await?; + sqlx::query( + "CREATE INDEX IF NOT EXISTS idx_session_turn_usage_session_created ON session_turn_usage(session_id, created_at)", + ) + .execute(&mut *tx) + .await?; sqlx::query(sqlx::AssertSqlSafe(format!( "PRAGMA user_version = {SCHEMA_VERSION}" ))) @@ -803,11 +857,12 @@ impl Storage { /// Atomically persist all messages produced by one logical turn together /// with the resulting session metadata. A turn is either fully visible /// after restart or not visible at all. - pub async fn persist_message_batch( + async fn persist_message_batch_inner( &self, session_id: &str, msgs: &[crate::storage::message::MessageMeta], meta: &crate::storage::session::SessionMeta, + usage: Option<&crate::storage::TurnUsageRecord>, ) -> Result<(), StorageError> { let mut tx = self.pool.begin().await?; @@ -848,10 +903,63 @@ impl Storage { .execute(&mut *tx) .await?; + if let Some(usage) = usage { + debug_assert_eq!(session_id, usage.session_id); + let request_count = msgs + .iter() + .filter_map(|message| message.iteration) + .max() + .map_or(1_i64, |iteration| iteration.saturating_add(1)); + sqlx::query( + r#" + INSERT INTO session_turn_usage ( + turn_id, session_id, provider, model, prompt_tokens, + completion_tokens, total_tokens, cached_input_tokens, + request_count, last_prompt_tokens, created_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(turn_id) DO NOTHING + "#, + ) + .bind(&usage.turn_id) + .bind(&usage.session_id) + .bind(&usage.provider) + .bind(&usage.model) + .bind(i64::from(usage.usage.prompt_tokens)) + .bind(i64::from(usage.usage.completion_tokens)) + .bind(i64::from(usage.usage.total_tokens)) + .bind(usage.usage.cached_tokens.map(i64::from)) + .bind(request_count) + .bind(i64::from(usage.last_prompt_tokens)) + .bind(usage.created_at) + .execute(&mut *tx) + .await?; + } + tx.commit().await?; Ok(()) } + pub async fn persist_message_batch( + &self, + session_id: &str, + msgs: &[crate::storage::message::MessageMeta], + meta: &crate::storage::session::SessionMeta, + ) -> Result<(), StorageError> { + self.persist_message_batch_inner(session_id, msgs, meta, None) + .await + } + + pub async fn persist_turn_batch( + &self, + session_id: &str, + msgs: &[crate::storage::message::MessageMeta], + meta: &crate::storage::session::SessionMeta, + usage: &crate::storage::TurnUsageRecord, + ) -> Result<(), StorageError> { + self.persist_message_batch_inner(session_id, msgs, meta, Some(usage)) + .await + } + /// Persist a turn with bounded retry. Retrying the whole transaction keeps /// message rows and metadata consistent on transient SQLite failures. pub async fn persist_message_batch_with_retry( @@ -874,6 +982,80 @@ impl Storage { unreachable!() } + pub async fn persist_turn_batch_with_retry( + &self, + session_id: &str, + msgs: &[crate::storage::message::MessageMeta], + meta: &crate::storage::session::SessionMeta, + usage: &crate::storage::TurnUsageRecord, + ) -> Result<(), StorageError> { + let delays = [100, 200, 300]; + for (attempt, delay) in delays.iter().enumerate() { + match self.persist_turn_batch(session_id, msgs, meta, usage).await { + Ok(()) => return Ok(()), + Err(error) if attempt < delays.len() - 1 && error.is_transient() => { + tracing::warn!(attempt = attempt + 1, error = %error, "Turn persistence failed; retrying"); + sleep(Duration::from_millis(*delay)).await; + } + Err(error) => return Err(error), + } + } + unreachable!() + } + + pub async fn get_session_usage_totals( + &self, + session_id: &str, + ) -> Result<crate::storage::SessionUsageTotals, StorageError> { + let row = sqlx::query( + r#" + SELECT + COALESCE(SUM(prompt_tokens), 0) AS prompt_tokens, + COALESCE(SUM(completion_tokens), 0) AS completion_tokens, + COALESCE(SUM(total_tokens), 0) AS total_tokens, + SUM(cached_input_tokens) AS cached_input_tokens, + COALESCE(SUM(request_count), 0) AS request_count, + COUNT(*) AS turn_count, + MIN(created_at) AS tracked_since + FROM session_turn_usage + WHERE session_id = ? + "#, + ) + .bind(session_id) + .fetch_one(self.pool()) + .await?; + + let last = sqlx::query( + r#" + SELECT last_prompt_tokens, created_at + FROM session_turn_usage + WHERE session_id = ? + ORDER BY created_at DESC, rowid DESC + LIMIT 1 + "#, + ) + .bind(session_id) + .fetch_optional(self.pool()) + .await?; + + Ok(crate::storage::SessionUsageTotals { + prompt_tokens: u64::try_from(row.get::<i64, _>("prompt_tokens")).unwrap_or_default(), + completion_tokens: u64::try_from(row.get::<i64, _>("completion_tokens")) + .unwrap_or_default(), + total_tokens: u64::try_from(row.get::<i64, _>("total_tokens")).unwrap_or_default(), + cached_input_tokens: row + .get::<Option<i64>, _>("cached_input_tokens") + .and_then(|value| u64::try_from(value).ok()), + request_count: u64::try_from(row.get::<i64, _>("request_count")).unwrap_or_default(), + turn_count: u64::try_from(row.get::<i64, _>("turn_count")).unwrap_or_default(), + tracked_since: row.get("tracked_since"), + last_prompt_tokens: last + .as_ref() + .and_then(|value| u64::try_from(value.get::<i64, _>("last_prompt_tokens")).ok()), + last_observed_at: last.map(|value| value.get("created_at")), + }) + } + pub async fn load_messages( &self, session_id: &str, @@ -1382,6 +1564,80 @@ mod tests { assert!(orphan.is_err()); } + #[tokio::test] + async fn committed_turn_usage_is_aggregated_and_idempotent() { + let (storage, _dir) = create_test_storage().await; + let meta = crate::storage::session::SessionMeta { + id: "cli_chat:chat:dialog".to_string(), + channel: "cli_chat".to_string(), + chat_id: "chat".to_string(), + dialog_id: "dialog".to_string(), + title: "usage".to_string(), + created_at: 1, + last_active_at: 2, + message_count: 1, + routing_info: None, + archived_at: None, + deleted_at: None, + last_consolidated_at: None, + last_compressed_message_at: None, + }; + let first = crate::storage::TurnUsageRecord { + session_id: meta.id.clone(), + turn_id: "turn-1".to_string(), + provider: "test".to_string(), + model: "model".to_string(), + usage: crate::providers::Usage { + prompt_tokens: 100, + completion_tokens: 20, + total_tokens: 120, + cached_tokens: Some(40), + cache_read_input_tokens: None, + cache_creation_input_tokens: None, + }, + last_prompt_tokens: 75, + created_at: 10, + }; + storage + .persist_turn_batch(&meta.id, &[], &meta, &first) + .await + .unwrap(); + storage + .persist_turn_batch(&meta.id, &[], &meta, &first) + .await + .unwrap(); + + let second = crate::storage::TurnUsageRecord { + turn_id: "turn-2".to_string(), + usage: crate::providers::Usage { + prompt_tokens: 50, + completion_tokens: 10, + total_tokens: 60, + cached_tokens: None, + cache_read_input_tokens: None, + cache_creation_input_tokens: None, + }, + last_prompt_tokens: 45, + created_at: 20, + ..first + }; + storage + .persist_turn_batch(&meta.id, &[], &meta, &second) + .await + .unwrap(); + + let totals = storage.get_session_usage_totals(&meta.id).await.unwrap(); + assert_eq!(totals.prompt_tokens, 150); + assert_eq!(totals.completion_tokens, 30); + assert_eq!(totals.total_tokens, 180); + assert_eq!(totals.cached_input_tokens, Some(40)); + assert_eq!(totals.request_count, 2); + assert_eq!(totals.turn_count, 2); + assert_eq!(totals.tracked_since, Some(10)); + assert_eq!(totals.last_prompt_tokens, Some(45)); + assert_eq!(totals.last_observed_at, Some(20)); + } + #[tokio::test] async fn reopening_database_does_not_rebuild_existing_fts_index() { let dir = tempfile::tempdir().unwrap(); @@ -1629,7 +1885,7 @@ mod tests { .await .unwrap(); assert_eq!(schema_version, SCHEMA_VERSION); - for table in ["task_plans", "task_items"] { + for table in ["task_plans", "task_items", "session_turn_usage"] { let exists: i64 = sqlx::query_scalar( "SELECT COUNT(*) FROM sqlite_master WHERE type = 'table' AND name = ?", ) diff --git a/src/storage/usage.rs b/src/storage/usage.rs new file mode 100644 index 0000000..53196aa --- /dev/null +++ b/src/storage/usage.rs @@ -0,0 +1,28 @@ +use crate::providers::Usage; + +/// Provider-reported usage committed with one durable assistant Turn. +#[derive(Debug, Clone)] +pub struct TurnUsageRecord { + pub session_id: String, + pub turn_id: String, + pub provider: String, + pub model: String, + pub usage: Usage, + /// Prompt usage from the final provider request in the Turn. Unlike + /// `usage.prompt_tokens`, this is not accumulated across tool iterations. + pub last_prompt_tokens: u32, + pub created_at: i64, +} + +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct SessionUsageTotals { + pub prompt_tokens: u64, + pub completion_tokens: u64, + pub total_tokens: u64, + pub cached_input_tokens: Option<u64>, + pub request_count: u64, + pub turn_count: u64, + pub tracked_since: Option<i64>, + pub last_prompt_tokens: Option<u64>, + pub last_observed_at: Option<i64>, +} diff --git a/webui/package-lock.json b/webui/package-lock.json index ff9bb5a..1815367 100644 --- a/webui/package-lock.json +++ b/webui/package-lock.json @@ -1,12 +1,12 @@ { "name": "picobot-webui", - "version": "1.3.1", + "version": "1.4.0", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "picobot-webui", - "version": "1.3.1", + "version": "1.4.0", "dependencies": { "bits-ui": "^2.0.0", "dompurify": "^3.4.12", diff --git a/webui/package.json b/webui/package.json index 9ad5b9f..39802de 100644 --- a/webui/package.json +++ b/webui/package.json @@ -1,7 +1,7 @@ { "name": "picobot-webui", "private": true, - "version": "1.3.1", + "version": "1.4.0", "type": "module", "engines": { "node": ">=20" diff --git a/webui/src/lib/chat.svelte.js b/webui/src/lib/chat.svelte.js index b5adb83..d9d5300 100644 --- a/webui/src/lib/chat.svelte.js +++ b/webui/src/lib/chat.svelte.js @@ -3,6 +3,8 @@ import { clientId } from "./api.js"; class ChatClient { connected = $state(false); turn = $state(null); // 最新 turn 快照（任意 session），供活动脊 + currentSessionId = $state(null); + statsBySession = $state({}); #socket = null; #handlers = new Set(); #reconnectTimer = null; @@ -33,6 +35,9 @@ class ChatClient { let frame; try { frame = JSON.parse(event.data); } catch { return; } if (frame.type === "turn_updated" && frame.snapshot) this.turn = frame.snapshot; + if (frame.type === "session_stats" && frame.stats?.session_id) { + this.statsBySession[frame.stats.session_id] = frame.stats; + } this.#dispatch(frame); }; } @@ -53,6 +58,10 @@ class ChatClient { return false; } + get currentStats() { + return this.currentSessionId ? this.statsBySession[this.currentSessionId] ?? null : null; + } + subscribe(handler) { this.#handlers.add(handler); return () => this.#handlers.delete(handler); diff --git a/webui/src/lib/components/ActivitySpine.svelte b/webui/src/lib/components/ActivitySpine.svelte index a883751..4f40dd6 100644 --- a/webui/src/lib/components/ActivitySpine.svelte +++ b/webui/src/lib/components/ActivitySpine.svelte @@ -2,51 +2,151 @@ import { chat } from "../chat.svelte.js"; let { version = "" } = $props(); - let lastTokens = null; // { at, completion } — plain bookkeeping, NOT reactive - let rate = $state(null); - $effect(() => { - const turn = chat.turn; - if (!turn || turn.status !== "running") { rate = null; lastTokens = null; return; } - const completion = turn.usage?.completion_tokens; - const now = Date.now(); - if (completion != null && lastTokens && now > lastTokens.at) { - const delta = completion - lastTokens.completion; - const secs = (now - lastTokens.at) / 1000; - if (delta >= 0 && secs > 0) rate = Math.round(delta / secs); - } - if (completion != null) lastTokens = { at: now, completion }; - }); + const running = $derived(chat.turn?.status === "running" && chat.turn?.session_id === chat.currentSessionId); + const turnLabel = $derived(running ? `Turn ${String(chat.turn?.id ?? "").slice(0, 6).toUpperCase()}` : ""); + const stats = $derived(chat.currentStats); + const context = $derived(stats?.context ?? null); + const lifetime = $derived(stats?.lifetime_usage ?? null); + const percent = $derived(context?.effective_window_tokens + ? context.used_tokens / context.effective_window_tokens * 100 + : 0); + const boundedPercent = $derived(Math.max(0, Math.min(percent, 100))); + const pressure = $derived(percent >= 90 ? "danger" : percent >= 70 ? "warning" : "normal"); - const running = $derived(chat.turn?.status === "running"); - const turnLabel = $derived(chat.turn ? `Turn ${String(chat.turn.id ?? "").slice(0, 6).toUpperCase()}` : ""); - const ctx = $derived(chat.turn?.usage?.prompt_tokens != null - ? `${(chat.turn.usage.prompt_tokens / 1000).toFixed(1)}k` : null); + function compactTokens(value) { + if (!Number.isFinite(value)) return "—"; + if (value >= 1_000_000) return `${(value / 1_000_000).toFixed(value >= 10_000_000 ? 1 : 2)}M`; + if (value >= 1_000) return `${(value / 1_000).toFixed(value >= 100_000 ? 0 : 1)}k`; + return String(value); + } + + function exactTokens(value) { + return Number.isFinite(value) ? new Intl.NumberFormat("zh-CN").format(value) : "—"; + } + + function sourceLabel(value) { + return value === "hybrid" ? "混合估算" : "字符估算"; + } + + function trackedSince(value) { + if (!Number.isFinite(value)) return "尚无已提交 Turn"; + return `自 ${new Date(value).toLocaleString("zh-CN", { hour12: false })}`; + } </script> -<div class="spine mono" title={version}> - {#if running} - <span class="spine-turn active"><i class="pulse-dot active"></i>{turnLabel} · 生成中</span> - {#if rate != null}<span class="spine-rate">▲ {rate} tok/s</span>{/if} - {#if ctx}<span>ctx {ctx}</span>{/if} - {:else if chat.turn} - <span class="spine-turn idle"><i class="pulse-dot idle"></i>空闲</span> - <span class="spine-context">最近 {turnLabel}</span> - {:else} - <span class="spine-turn idle"><i class="pulse-dot idle"></i>就绪</span> +<details class="spine-shell mono" data-pressure={pressure} title={version}> + <summary class="spine"> + {#if running} + <span class="spine-turn active"><i class="pulse-dot active"></i>{turnLabel} · 生成中</span> + {:else} + <span class="spine-turn idle"><i class="pulse-dot idle"></i>{chat.connected ? "空闲" : "重连中"}</span> + {/if} + + {#if context} + <span class="context-readout"> + <span class="context-label">上下文</span> + <span + class="context-rail" + role="progressbar" + aria-label="当前上下文窗口占用" + aria-valuemin="0" + aria-valuemax="100" + aria-valuenow={Math.round(boundedPercent)} + style={`--context-fill: ${boundedPercent}%`} + ><i></i><b></b></span> + <span class="context-value">{compactTokens(context.used_tokens)} / {compactTokens(context.effective_window_tokens)}</span> + <strong class="context-percent">{percent.toFixed(1)}%</strong> + </span> + <span class="usage-summary"><span>↑ {compactTokens(lifetime?.input_tokens)}</span><span>↓ {compactTokens(lifetime?.output_tokens)}</span></span> + {:else} + <span class="context-pending">等待会话统计</span> + {/if} + + <span class="spine-right" class:spine-ok={chat.connected} class:spine-down={!chat.connected}> + {chat.connected ? "已连接" : "重连中"} + </span> + </summary> + + {#if stats && context && lifetime} + <section class="stats-panel" aria-label="会话 Token 与上下文详情"> + <header> + <div><span>当前上下文</span><strong>{percent.toFixed(1)}%</strong></div> + <small>{sourceLabel(context.source)}</small> + </header> + <div class="detail-rail" role="presentation" style={`--context-fill: ${boundedPercent}%`}><i></i><b></b></div> + <dl class="context-grid"> + <div><dt>占用</dt><dd>{exactTokens(context.used_tokens)} / {exactTokens(context.effective_window_tokens)}</dd></div> + <div><dt>剩余</dt><dd>{exactTokens(context.remaining_tokens)}</dd></div> + <div><dt>压缩阈值</dt><dd>{exactTokens(context.compression_threshold_tokens)} · 70%</dd></div> + <div><dt>最近实测</dt><dd>{exactTokens(context.last_observed_prompt_tokens)}</dd></div> + </dl> + + <div class="usage-heading"><span>会话累计</span><small>Provider 报告 · 已提交 Turns</small></div> + <dl class="usage-grid"> + <div><dt>输入</dt><dd>{exactTokens(lifetime.input_tokens)}</dd></div> + <div><dt>输出</dt><dd>{exactTokens(lifetime.output_tokens)}</dd></div> + <div><dt>缓存输入</dt><dd>{exactTokens(lifetime.cached_input_tokens)}</dd></div> + <div><dt>请求 / Turns</dt><dd>{exactTokens(lifetime.request_count)} / {exactTokens(lifetime.turn_count)}</dd></div> + </dl> + <footer>{stats.provider} / {stats.model}<span>{trackedSince(lifetime.tracked_since)}</span></footer> + </section> {/if} - <span class="spine-right"> - <span class:spine-ok={chat.connected} class:spine-down={!chat.connected}>{chat.connected ? "已连接" : "重连中"}</span> - </span> -</div> +</details> <style> - .spine { display: flex; align-items: center; gap: 10px; min-height: 32px; margin-left: auto; padding: 0 11px; border: 1px solid var(--line); border-radius: 6px; font-size: 11.5px; color: var(--text-soft); background: var(--panel-2); font-variant-numeric: tabular-nums; } - .spine-turn { display: inline-flex; align-items: center; gap: 7px; font-weight: 700; } - .spine-turn.active, .spine-rate { color: var(--accent); } + .spine-shell { position: relative; margin-left: auto; min-width: 0; } + .spine-shell summary { list-style: none; } + .spine-shell summary::-webkit-details-marker { display: none; } + .spine { display: flex; align-items: center; gap: 10px; min-height: 32px; padding: 0 11px; border: 1px solid var(--line); border-radius: 6px; font-size: 11.5px; color: var(--text-soft); background: var(--panel-2); font-variant-numeric: tabular-nums; cursor: pointer; user-select: none; } + .spine:hover { border-color: var(--line-strong); background: var(--color-neutral-background-4); } + .spine:focus-visible { outline: 2px solid var(--accent); outline-offset: 2px; } + .spine-turn { display: inline-flex; align-items: center; gap: 7px; font-weight: 700; white-space: nowrap; } + .spine-turn.active { color: var(--accent); } .spine-turn.idle, .spine-ok { color: var(--signal); } - .spine-right { display: inline-flex; gap: 10px; padding-left: 10px; border-left: 1px solid var(--line); color: var(--muted); } + .context-readout { display: inline-flex; align-items: center; gap: 7px; white-space: nowrap; } + .context-label { color: var(--muted); } + .context-rail, .detail-rail { position: relative; overflow: hidden; background: var(--color-neutral-background-4); } + .context-rail { width: 74px; height: 6px; border-radius: 99px; } + .context-rail i, .detail-rail i { position: absolute; inset: 0 auto 0 0; width: var(--context-fill); background: var(--accent); } + .context-rail b, .detail-rail b { position: absolute; inset: 0 auto 0 70%; width: 1px; background: var(--warning); } + .context-value { color: var(--text-soft); } + .context-percent { color: var(--accent); font-weight: 700; } + [data-pressure="warning"] .context-percent, [data-pressure="warning"] .context-rail i, [data-pressure="warning"] .detail-rail i { color: var(--warning); background: var(--warning); } + [data-pressure="danger"] .context-percent, [data-pressure="danger"] .context-rail i, [data-pressure="danger"] .detail-rail i { color: var(--danger); background: var(--danger); } + .usage-summary { display: inline-flex; gap: 8px; padding-left: 10px; border-left: 1px solid var(--line); color: var(--muted); white-space: nowrap; } + .spine-right { padding-left: 10px; border-left: 1px solid var(--line); white-space: nowrap; } .spine-down { color: var(--danger); } - @media (max-width: 1040px) { .spine-context, .spine-rate, .spine > span:not(.spine-turn):not(.spine-right) { display: none; } } - @media (max-width: 680px) { .spine { display: none; } } + .context-pending { color: var(--muted); } + + .stats-panel { position: absolute; z-index: 30; top: calc(100% + 8px); right: 0; width: min(390px, calc(100vw - 28px)); padding: 17px; border: 1px solid var(--line-strong); border-radius: 10px; color: var(--text); background: var(--overlay); box-shadow: var(--shadow-16); font-family: var(--font-ui); } + .stats-panel header, .usage-heading, .stats-panel footer { display: flex; align-items: baseline; justify-content: space-between; gap: 12px; } + .stats-panel header div { display: flex; align-items: baseline; gap: 10px; } + .stats-panel header span, .usage-heading span { font-size: 12px; font-weight: 700; } + .stats-panel header strong { font-family: var(--font-mono); font-size: 21px; letter-spacing: -0.04em; } + .stats-panel small, .stats-panel footer { color: var(--muted); font-size: 10.5px; } + .detail-rail { height: 8px; margin: 12px 0 15px; border-radius: 99px; } + .context-grid, .usage-grid { display: grid; grid-template-columns: 1fr 1fr; gap: 1px; margin: 0; overflow: hidden; border: 1px solid var(--line); border-radius: 7px; background: var(--line); } + .context-grid div, .usage-grid div { min-width: 0; padding: 9px 10px; background: var(--panel); } + dt { color: var(--muted); font-size: 10.5px; } + dd { margin: 3px 0 0; overflow: hidden; font-family: var(--font-mono); font-size: 11.5px; font-variant-numeric: tabular-nums; text-overflow: ellipsis; white-space: nowrap; } + .usage-heading { margin: 16px 0 8px; } + .stats-panel footer { margin-top: 12px; padding-top: 11px; border-top: 1px solid var(--line); } + .stats-panel footer span { text-align: right; } + + @media (max-width: 1180px) { + .context-value, .usage-summary { display: none; } + .context-rail { width: 58px; } + } + @media (max-width: 760px) { + .spine-shell { margin-left: 0; } + .spine { padding: 0 9px; } + .spine-turn, .context-label, .context-rail, .spine-right, .context-pending { display: none; } + .context-readout { gap: 4px; } + .context-percent::before { content: "上下文 "; color: var(--muted); font-weight: 400; } + .stats-panel { position: fixed; top: auto; right: 14px; bottom: 14px; left: 14px; width: auto; } + } + @media (prefers-reduced-motion: reduce) { + .context-rail i, .detail-rail i { transition: none; } + } </style> diff --git a/webui/src/pages/ChatPage.svelte b/webui/src/pages/ChatPage.svelte index e3a402d..1a13ba7 100644 --- a/webui/src/pages/ChatPage.svelte +++ b/webui/src/pages/ChatPage.svelte @@ -55,14 +55,22 @@ function handleFrame(frame) { switch (frame.type) { - case "session_established": currentId = frame.session_id; activeTurn = null; historyRevision = 0; break; + case "session_established": + currentId = frame.session_id; + chat.currentSessionId = currentId; + activeTurn = null; + historyRevision = 0; + send({ type: "get_session_stats", session_id: currentId }); + break; case "session_list": sessions = frame.sessions || []; if (frame.current_session_id) currentId = frame.current_session_id; + chat.currentSessionId = currentId; if (currentId && messages.length === 0) loadSession(currentId); break; case "session_created": currentId = frame.session_id; + chat.currentSessionId = currentId; messages = []; activeTurn = null; historyRevision = 0; @@ -70,10 +78,12 @@ break; case "session_loaded": currentId = frame.session_id; + chat.currentSessionId = currentId; activeTurn = null; historyRevision = 0; send({ type: "get_session_history", session_id: currentId, limit: 1000 }); send({ type: "get_session_plan", session_id: currentId }); + send({ type: "get_session_stats", session_id: currentId }); break; case "session_history": if (frame.session_id === currentId) { @@ -118,13 +128,16 @@ if (currentId) send({ type: "get_session_history", session_id: currentId, limit: 1000 }); } send({ type: "list_sessions", include_archived: false }); + if (currentId) send({ type: "get_session_stats", session_id: currentId }); break; case "turn_updated": { const next = frame.snapshot; if (!next || next.session_id !== currentId) break; if (activeTurn?.id === next.id && activeTurn.revision >= next.revision) break; + const firstSnapshotForTurn = activeTurn?.id !== next.id; activeTurn = next; thinking = next.status === "running"; + if (firstSnapshotForTurn) send({ type: "get_session_stats", session_id: currentId }); if (next.status !== "running" && messages.some((message) => message.id === next.message_id)) { activeTurn = null; } @@ -148,12 +161,18 @@ activeTurn = null; } scrollToBottom(); + send({ type: "get_session_stats", session_id: currentId }); break; } + case "session_stats": break; case "system_notification": if (!frame.session_id || frame.session_id === currentId) appendMessage("assistant", frame.content); break; - case "command_executed": thinking = false; appendMessage("assistant", frame.message); break; + case "command_executed": + thinking = false; + appendMessage("assistant", frame.message); + if (currentId) send({ type: "get_session_stats", session_id: currentId }); + break; case "error": thinking = false; notify(frame.message || frame.code, true); break; } } @@ -171,6 +190,7 @@ function loadSession(id) { if (!id) return; currentId = id; + chat.currentSessionId = id; historyRevision = 0; clearPendingUploads(); messages = []; @@ -399,7 +419,7 @@ </button> {/if} <Tooltip.Root> - <Tooltip.Trigger class="icon-button" aria-label="刷新会话" onclick={() => { send({ type: "list_sessions", include_archived: false }); if (currentId) send({ type: "get_session_plan", session_id: currentId }); }}><Icon name="refresh" size={18} /></Tooltip.Trigger> + <Tooltip.Trigger class="icon-button" aria-label="刷新会话" onclick={() => { send({ type: "list_sessions", include_archived: false }); if (currentId) { send({ type: "get_session_plan", session_id: currentId }); send({ type: "get_session_stats", session_id: currentId }); } }}><Icon name="refresh" size={18} /></Tooltip.Trigger> <Tooltip.Portal><Tooltip.Content class="tooltip" sideOffset={7}>刷新会话<Tooltip.Arrow class="tooltip-arrow" /></Tooltip.Content></Tooltip.Portal> </Tooltip.Root> </div>