diff --git a/desktop/Cargo.lock b/desktop/Cargo.lock index 64d98b9a5..cde9821ad 100644 --- a/desktop/Cargo.lock +++ b/desktop/Cargo.lock @@ -2255,6 +2255,15 @@ dependencies = [ "syn 2.0.118", ] +[[package]] +name = "num_threads" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c7398b9c8b70908f6371f47ed36737907c87c52af34c268fed0bf0ceb92ead9" +dependencies = [ + "libc", +] + [[package]] name = "objc2" version = "0.6.4" @@ -4306,7 +4315,9 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "18dfaaeddcb932337b5e7866ee7d0ce9b76d2fd092997146f187ec09b4558a50" dependencies = [ "deranged", + "libc", "num-conv", + "num_threads", "powerfmt", "serde_core", "time-core", diff --git a/desktop/Cargo.toml b/desktop/Cargo.toml index ca34be237..854c8f6f8 100644 --- a/desktop/Cargo.toml +++ b/desktop/Cargo.toml @@ -39,7 +39,7 @@ cookie = "0.18" regex = "1" # cookie 罐过期时间的 RFC3339 解析。time 本就在依赖树里(cookie 的 Expires # 解析),显式声明 + parsing feature,替代手写历算(严格校验月/日范围) -time = { version = "0.3", features = ["parsing"] } +time = { version = "0.3", features = ["parsing", "local-offset"] } [dev-dependencies] # baizhi 集成测试:假服务端 + async 测试运行时 diff --git a/desktop/build.rs b/desktop/build.rs index 32a815cea..90bb86792 100644 --- a/desktop/build.rs +++ b/desktop/build.rs @@ -53,6 +53,7 @@ fn main() { "session_delete", "session_patch", "models_list", + "usage_stats", "session_open", "session_history", "session_outline", diff --git a/desktop/src/driver/fold.rs b/desktop/src/driver/fold.rs index a260483e7..5769336cb 100644 --- a/desktop/src/driver/fold.rs +++ b/desktop/src/driver/fold.rs @@ -201,8 +201,7 @@ pub(super) struct TurnFold { } impl TurnFold { - pub(super) fn push(&mut self, f: &Value) { - let seq = frame_seq(f); + pub(super) fn push(&mut self, f: &Value) { let seq = frame_seq(f); if self.out.is_empty() { self.from = seq; } @@ -254,6 +253,25 @@ impl TurnFold { self.tail = None; } + /// usage 事件(input/output tokens)挂到本轮最后一条 agent_message 帧上, + /// 供 UI 在每条助手消息旁展示其 token 用量。provider 在同一次模型调用 + /// 的头尾可能各发一次 usage,后到者覆盖前值(取该调用的最终计数)。 + pub(super) fn attach_usage(&mut self, input: u64, output: u64) { + for f in self.out.iter_mut().rev() { + let Some(data) = f.get("data") else { continue }; + if session_update(data) != Some("agent_message_chunk") { + continue; + } + if let Some(update) = f.get_mut("data").and_then(|d| d.get_mut("update")).and_then(|u| u.as_object_mut()) { + update.insert( + "usage".into(), + json!({ "input_tokens": input, "output_tokens": output }), + ); + } + break; + } + } + pub(super) fn is_empty(&self) -> bool { self.out.is_empty() } diff --git a/desktop/src/driver/mod.rs b/desktop/src/driver/mod.rs index dbb32cb39..d20e488f6 100644 --- a/desktop/src/driver/mod.rs +++ b/desktop/src/driver/mod.rs @@ -396,6 +396,12 @@ pub async fn models_list(host: State<'_, DriverHost>) -> Result { host.get()?.models_list().await } +/// 本地会话 token 用量统计(按天/会话/模型聚合,usage 事件记账)。 +#[tauri::command] +pub async fn usage_stats(host: State<'_, DriverHost>) -> Result { + Ok(host.get()?.usage_stats()) +} + /// 打开会话:返回尾部回放窗口 `{frames, cursor, has_more}`。历史走返回值 /// 而非 `frames:{sid}` 事件——返回值天生有序,不必依赖"监听先于命令"。 #[tauri::command] diff --git a/desktop/src/driver/normalize.rs b/desktop/src/driver/normalize.rs index 350577198..747a48d03 100644 --- a/desktop/src/driver/normalize.rs +++ b/desktop/src/driver/normalize.rs @@ -415,6 +415,19 @@ impl Inner { if let Some((used, window)) = context_usage_fields(&data) { self.push_usage(&sid, used, window); } + self.record_usage(&sid, &data); + // 挂到本轮最后一条 agent_message 帧:每条助手消息显示其 token + // 用量。usage 事件是每次模型调用的全量,同帧后到覆盖前值。 + let input = data.get("input_tokens").and_then(|v| v.as_u64()).unwrap_or(0); + let output = data.get("output_tokens").and_then(|v| v.as_u64()).unwrap_or(0); + if input > 0 || output > 0 { + if let Some(sess) = self.sess.sessions.lock_ok().get_mut(&sid) { + sess.fold.attach_usage(input, output); + } + // 实时路径:usage 晚于流式帧,单独补发 session-usage 事件, + // UI 据此把用量补到最后一条助手消息与大纲条目上。 + self.emit_session_usage(&sid, input, output); + } } // 会话摘要:引擎每轮用户消息后异步生成一句 ≤60 字的对话摘要 // (随对话演进改写,后一轮覆盖前一轮),只给顶层会话生成。 @@ -469,6 +482,36 @@ impl Inner { _ => {} } } + + /// usage 事件里的 input/output tokens → 用量统计(按天/会话/模型)。 + /// 引擎每次模型调用发一个 usage 事件,input/output 为该调用全量,直接 + /// 累加进对应桶。模型取会话当前 model_name——运行中不可切模型,归属 + /// 可靠。只记 input/output 非 0 的事件(纯 context 快照不记账)。 + pub(super) fn record_usage(&self, sid: &str, data: &Value) { + let input = data.get("input_tokens").and_then(|v| v.as_u64()).unwrap_or(0); + let output = data.get("output_tokens").and_then(|v| v.as_u64()).unwrap_or(0); + if input == 0 && output == 0 { + return; + } + // 跨组嵌套仅 subagents → sessions 一条(见 ohmy.rs::Inner),先取子代理 + let parent = self.sub.subagents.lock_ok().get(sid).map(|r| r.parent_sid.clone()); + let (model, title) = { + let sessions = self.sess.sessions.lock_ok(); + match sessions.get(sid) { + Some(s) => (s.model_name.clone(), s.title.clone()), + None => (String::new(), String::new()), + } + }; + self.stats.record( + &crate::stats::today(), + sid, + &title, + &model, + parent.as_deref(), + input, + output, + ); + } } /// 上下文占用字段防腐层。13c8adc 起所有新入口统一为扁平 diff --git a/desktop/src/driver/ohmy.rs b/desktop/src/driver/ohmy.rs index 217d18059..f3111150f 100644 --- a/desktop/src/driver/ohmy.rs +++ b/desktop/src/driver/ohmy.rs @@ -104,6 +104,11 @@ impl OhmyDriver { crate::wsl::host_fs_view(&w.distro, &w.guest_home).to_string_lossy().into_owned() }) } + + /// 本地会话 token 用量统计快照(按天/会话/模型聚合)。 + pub fn usage_stats(&self) -> Value { + self.0.stats.snapshot() + } } /// WSL 运行环境上下文(kernel_env=wsl:* 时随引擎启动填入;一次 prepare @@ -145,6 +150,8 @@ pub(super) struct Inner { pub(super) chat_workspaces_dir: PathBuf, /// 壳侧审批记忆持久化路径(兼容尾巴,配对 SessionsState::perm_remember) pub(super) perm_persist_path: PathBuf, + /// 本地会话 token 用量统计(按天/会话/模型;usage 事件记账,落盘 usage-stats.json) + pub(super) stats: crate::stats::UsageStats, /// WSL 运行环境上下文(本机模式 None;见 WslCtx) pub(super) wsl: Option, /// 技能库来源(skills.rs):内置(bundle 资源,可缺)与用户目录, diff --git a/desktop/src/driver/ohmy_tests.rs b/desktop/src/driver/ohmy_tests.rs index fb54e477f..3afbfeb34 100644 --- a/desktop/src/driver/ohmy_tests.rs +++ b/desktop/src/driver/ohmy_tests.rs @@ -971,6 +971,7 @@ fn bare_inner_events(tag: &str) -> (Arc, EmittedEvents) { engine_dir: home.join("ohmyagent"), chat_workspaces_dir: home.join("local-data/chat-workspaces"), perm_persist_path: home.join("perm.json"), + stats: crate::stats::UsageStats::new(&home), wsl: None, skills_builtin_dir: None, skills_user_dir: home.join("skills"), @@ -1503,6 +1504,49 @@ fn streaming_usage_updates_emitting_agent_without_parent_leak() { assert_eq!(usage_of("child"), vec![(45_678, 64_000)]); } +/// usage 事件里的 input/output tokens 会记入本地用量统计(按天/会话/模型), +/// 纯 context 快照(无 input/output)不重复记账;模型归属取会话当前 model_name。 +#[test] +fn usage_events_accumulate_into_daily_stats() { + let inner = bare_inner("usage-stats"); + { + let mut sessions = inner.sess.sessions.lock().unwrap(); + let mut main = bare_session("main"); + main.model_name = "glm-4.5".into(); + main.title = "重构登录".into(); + sessions.insert("main".into(), main); + } + + inner.handle_event(json!({ "type": "usage", "session_id": "main", "seq": 1, + "data": { "input_tokens": 900, "output_tokens": 20, + "context_used": 1_234, "context_window": 200_000 } })); + // 同一调用的 message_start/message_delta 第二个事件没有 input/output → 不记账 + inner.handle_event(json!({ "type": "usage", "session_id": "main", "seq": 2, + "data": { "context_used": 1_234, "context_window": 200_000 } })); + inner.handle_event(json!({ "type": "usage", "session_id": "main", "seq": 3, + "data": { "input_tokens": 40_000, "output_tokens": 5_000, + "context_used": 45_678, "context_window": 64_000 } })); + + let snap = inner.stats.snapshot(); + assert_eq!(snap["totals"]["input_tokens"], 40_900, "{snap}"); + assert_eq!(snap["totals"]["output_tokens"], 5_020, "{snap}"); + assert_eq!(snap["totals"]["calls"], 2, "纯 context 快照不应计调用次数: {snap}"); + + let today = crate::stats::today(); + assert_eq!(snap["days"][0]["date"], today, "{snap}"); + assert_eq!(snap["days"][0]["input_tokens"], 40_900, "{snap}"); + + assert_eq!(snap["models"][0]["model"], "glm-4.5", "{snap}"); + assert_eq!(snap["models"][0]["input_tokens"], 40_900, "{snap}"); + + let sess = &snap["sessions"][0]; + assert_eq!(sess["session_id"], "main", "{snap}"); + assert_eq!(sess["title"], "重构登录", "{snap}"); + assert!(sess["parent"].is_null(), "{snap}"); + assert_eq!(sess["input_tokens"], 40_900, "{snap}"); + assert_eq!(sess["days"][0]["output_tokens"], 5_020, "{snap}"); +} + /// 最新 Agent 将 turn/stopped 与 usage 统一成扁平字段；旧版嵌套形状仍 /// 接受，保证桌面壳与已分发 sidecar 的滚动升级兼容。 #[test] diff --git a/desktop/src/driver/session.rs b/desktop/src/driver/session.rs index 42792352d..681823f78 100644 --- a/desktop/src/driver/session.rs +++ b/desktop/src/driver/session.rs @@ -2314,6 +2314,17 @@ impl Inner { ); } + /// 该会话最近一次模型调用的 token 用量(session-event 事件,不落帧): + /// UI 实时把用量补到当前回合的助手消息与大纲条目上(帧管线里 usage 事件 + /// 晚于流式帧,实时路径走这里;回放路径靠 attach_usage 挂在帧上)。 + pub(super) fn emit_session_usage(&self, sid: &str, input: u64, output: u64) { + let title = self.sess.sessions.lock_ok().get(sid).map(|s| s.title.clone()).unwrap_or_default(); + self.app.emit_json( + "session-event", + json!({ "type": "session-usage", "id": sid, "title": title, "input": input, "output": output }), + ); + } + pub(super) fn emit_session_ask(&self, sid: &str, open: bool) { let title = self.sess.sessions.lock_ok().get(sid).map(|s| s.title.clone()).unwrap_or_default(); self.app.emit_json( diff --git a/desktop/src/driver/transport.rs b/desktop/src/driver/transport.rs index e586b1069..7c0f698c9 100644 --- a/desktop/src/driver/transport.rs +++ b/desktop/src/driver/transport.rs @@ -349,6 +349,7 @@ impl OhmyDriver { engine_dir, chat_workspaces_dir, perm_persist_path, + stats: crate::stats::UsageStats::new(&cfg_dir), wsl: wsl_ctx, skills_builtin_dir: app_builtin_skills, skills_user_dir: crate::skills::user_dir(&cfg_dir), diff --git a/desktop/src/main.rs b/desktop/src/main.rs index 0f72e5b4c..2e3acfbc4 100644 --- a/desktop/src/main.rs +++ b/desktop/src/main.rs @@ -20,6 +20,7 @@ mod driver; mod native_pet; mod repo; mod skills; +mod stats; mod telemetry; mod todos; mod uploads; @@ -1425,6 +1426,7 @@ fn main() { driver::session_delete, driver::session_patch, driver::models_list, + driver::usage_stats, driver::session_open, driver::session_history, driver::session_outline, diff --git a/desktop/src/stats.rs b/desktop/src/stats.rs new file mode 100644 index 000000000..810b690eb --- /dev/null +++ b/desktop/src/stats.rs @@ -0,0 +1,242 @@ +//! 本地会话 token 用量统计(按天 → 会话 → 模型聚合)。 +//! +//! 数据源:引擎每收到一次模型调用的 provider usage 就推一个 `usage` 事件, +//! 壳在 normalize.rs 的 usage 分支里把 `input_tokens`/`output_tokens` 记到这里。 +//! 仅统计本机桌面会话,不涉及云端任务。 +//! +//! 落盘:`/usage-stats.json`,结构: +//! `{ "": { "": { "": Record } } }`。 +//! 每次 record 全量原子写(文件小,usage 事件按模型调用频次,量级可接受)。 + +use std::collections::{BTreeMap, HashMap}; +use std::path::{Path, PathBuf}; +use std::sync::Mutex as StdMutex; + +use serde::{Deserialize, Serialize}; +use serde_json::{json, Value}; + +use crate::config::atomic_write_private; +use crate::util::LockExt; + +/// 保留最近多少天的记录,避免落盘文件无限膨胀。 +const KEEP_DAYS: i64 = 366; + +/// 某天某会话某模型的累计消耗(合并幂等:同名桶只累加 token 与调用次数)。 +#[derive(Default, Clone, Debug, Serialize, Deserialize)] +pub(super) struct Record { + pub(super) title: String, + /// 子代理会话的父会话 id(顶层任务为 None,UI 可据此归并) + #[serde(default, skip_serializing_if = "Option::is_none")] + pub(super) parent: Option, + pub(super) input_tokens: u64, + pub(super) output_tokens: u64, + /// 模型调用次数(即收到的 usage 事件数) + pub(super) calls: u64, +} + +type ModelBuckets = HashMap; +type SessionBuckets = HashMap; +type Days = HashMap; + +pub(super) struct UsageStats { + path: PathBuf, + days: StdMutex, +} + +impl UsageStats { + pub(super) fn new(config_dir: &Path) -> Self { + let path = config_dir.join("usage-stats.json"); + let days = std::fs::read_to_string(&path) + .ok() + .and_then(|s| serde_json::from_str::(&s).ok()) + .unwrap_or_default(); + let store = Self { path, days: StdMutex::new(days) }; + store.prune(); + store + } + + /// 记一条模型调用的 token 消耗(usage 事件里的 input/output 全量累加)。 + pub(super) fn record( + &self, + date: &str, + sid: &str, + title: &str, + model: &str, + parent: Option<&str>, + input_tokens: u64, + output_tokens: u64, + ) { + if input_tokens == 0 && output_tokens == 0 { + return; + } + let model = if model.is_empty() { "" } else { model }; + let mut days = self.days.lock_ok(); + let rec = days + .entry(date.to_string()) + .or_default() + .entry(sid.to_string()) + .or_default() + .entry(model.to_string()) + .or_default(); + if !title.is_empty() { + rec.title = title.to_string(); + } + if parent.is_some() { + rec.parent = parent.map(|p| p.to_string()); + } + rec.input_tokens += input_tokens; + rec.output_tokens += output_tokens; + rec.calls += 1; + drop(days); + self.prune(); + self.persist(); + } + + /// 丢掉超过 KEEP_DAYS 的旧天。 + fn prune(&self) { + let cutoff = today_inner_checked_sub(KEEP_DAYS); + let mut days = self.days.lock_ok(); + days.retain(|date, _| date.as_str() >= cutoff.as_str()); + } + + /// 全量原子落盘。 + fn persist(&self) { + let days = self.days.lock_ok(); + let _ = serde_json::to_vec_pretty(&*days) + .map(|data| atomic_write_private(&self.path, &data)); + } + + /// 聚合快照,供 `usage_stats` IPC 命令返回给 UI。 + /// - totals:全部记录汇总 + /// - days:按天汇总(倒序) + /// - models:按模型汇总(用量倒序) + /// - sessions:按会话汇总(用量倒序),每条含 parent 与按天/按模型子拆分 + pub(super) fn snapshot(&self) -> Value { + let days = self.days.lock_ok(); + let mut totals = (0u64, 0u64, 0u64); + let mut by_day: BTreeMap = BTreeMap::new(); + let mut by_model: BTreeMap = BTreeMap::new(); + // sid → (title, parent, total, by_day, by_model) + let mut by_session: BTreeMap = BTreeMap::new(); + + for (date, sessions) in days.iter() { + for (sid, models) in sessions.iter() { + for (model, rec) in models.iter() { + totals.0 += rec.input_tokens; + totals.1 += rec.output_tokens; + totals.2 += rec.calls; + + let d = by_day.entry(date.clone()).or_default(); + d.0 += rec.input_tokens; + d.1 += rec.output_tokens; + d.2 += rec.calls; + + let m = by_model.entry(model.clone()).or_default(); + m.0 += rec.input_tokens; + m.1 += rec.output_tokens; + m.2 += rec.calls; + + let s = by_session.entry(sid.clone()).or_insert_with(|| SessionAgg { + title: rec.title.clone(), + parent: rec.parent.clone(), + ..Default::default() + }); + if !rec.title.is_empty() { + s.title = rec.title.clone(); + } + if rec.parent.is_some() { + s.parent = rec.parent.clone(); + } + s.total.0 += rec.input_tokens; + s.total.1 += rec.output_tokens; + s.total.2 += rec.calls; + let d2 = s.by_day.entry(date.clone()).or_default(); + d2.0 += rec.input_tokens; + d2.1 += rec.output_tokens; + d2.2 += rec.calls; + let m2 = s.by_model.entry(model.clone()).or_default(); + m2.0 += rec.input_tokens; + m2.1 += rec.output_tokens; + m2.2 += rec.calls; + } + } + } + drop(days); + + let sessions_json = { + let mut rows: Vec<(&String, &SessionAgg)> = by_session.iter().collect(); + rows.sort_by(|a, b| { + let ta = a.1.total.0 + a.1.total.1; + let tb = b.1.total.0 + b.1.total.1; + tb.cmp(&ta).then(a.0.cmp(b.0)) + }); + rows.into_iter().map(|(sid, s)| json!({ + "session_id": sid, + "title": s.title, + "parent": s.parent, + "input_tokens": s.total.0, + "output_tokens": s.total.1, + "calls": s.total.2, + "days": s.by_day.iter().rev().map(|(date, d)| json!({ + "date": date, + "input_tokens": d.0, + "output_tokens": d.1, + "calls": d.2, + })).collect::>(), + "models": sorted_bucket_json(s.by_model.clone()), + })).collect::>() + }; + + json!({ + "totals": bucket_json(totals), + "days": by_day.iter().rev().map(|(date, d)| json!({ + "date": date, + "input_tokens": d.0, + "output_tokens": d.1, + "calls": d.2, + })).collect::>(), + "models": sorted_bucket_json(by_model), + "sessions": sessions_json, + }) + } +} + +#[derive(Default)] +struct SessionAgg { + title: String, + parent: Option, + total: (u64, u64, u64), + by_day: BTreeMap, + by_model: BTreeMap, +} + +fn bucket_json(b: (u64, u64, u64)) -> Value { + json!({ "input_tokens": b.0, "output_tokens": b.1, "calls": b.2 }) +} + +fn sorted_bucket_json(map: BTreeMap) -> Vec { + let mut rows: Vec<(&String, &(u64, u64, u64))> = map.iter().collect(); + rows.sort_by(|a, b| (b.1).0.cmp(&(a.1).0).then(b.0.cmp(a.0))); + rows.into_iter() + .map(|(k, v)| json!({ + "model": k, + "input_tokens": v.0, + "output_tokens": v.1, + "calls": v.2, + })) + .collect() +} + +/// 本地时区的今天,`YYYY-MM-DD`。取不到本地时区时退回 UTC。 +pub(super) fn today() -> String { + let now = time::OffsetDateTime::now_local().unwrap_or_else(|_| time::OffsetDateTime::now_utc()); + let d = now.date(); + format!("{:04}-{:02}-{:02}", d.year(), d.month() as u8, d.day()) +} + +/// today() 往前推 n 天(仅用于剪枝,不解析历史日期)。 +fn today_inner_checked_sub(days: i64) -> String { + let now = time::OffsetDateTime::now_local().unwrap_or_else(|_| time::OffsetDateTime::now_utc()); + let d = now.date().checked_sub(time::Duration::days(days)).unwrap_or(now.date()); + format!("{:04}-{:02}-{:02}", d.year(), d.month() as u8, d.day()) +} diff --git a/desktop/ui-next/src/app/App.tsx b/desktop/ui-next/src/app/App.tsx index ec87d0d65..7e18a0fde 100644 --- a/desktop/ui-next/src/app/App.tsx +++ b/desktop/ui-next/src/app/App.tsx @@ -7,7 +7,7 @@ // 侧栏 attention 高亮; // - D8 增量自愈:session-event/意图指向未知 id → 重拉全表再选中; // - H9 意图消费:open-* 事件送达即 takeUiIntent 消费壳侧副本,防刷新重放。 -import { IconAlertCircle, IconCircleCheck, IconCloud, IconFolderCode, IconHelpCircle, IconMessages, IconPlayerStop, IconSend, IconSettings, IconWorld, IconX } from "@tabler/icons-react"; +import { IconAlertCircle, IconChartBar, IconCircleCheck, IconCloud, IconFolderCode, IconHelpCircle, IconMessages, IconPlayerStop, IconSend, IconSettings, IconWorld, IconX } from "@tabler/icons-react"; import { useEffect, useRef, useState } from "react"; import { ChatView } from "@/features/chat/ChatView"; @@ -17,6 +17,7 @@ import { EngineBanner } from "@/features/engine/EngineBanner"; import { NewTaskModal } from "@/features/newtask/NewTaskModal"; import { SettingsView } from "@/features/settings/SettingsView"; import { Sidebar } from "@/features/sidebar/Sidebar"; +import { UsageStatsView } from "@/features/stats/UsageStatsView"; import { useTodos } from "@/features/todo/useTodos"; import { ResizeEdges } from "@/features/titlebar/ResizeEdges"; import { MacWindowControls, TitleBar } from "@/features/titlebar/TitleBar"; @@ -55,6 +56,7 @@ const SPACE_ICONS: Record = { local: IconFolderCode, cloud: IconCloud, chat: IconMessages, + stats: IconChartBar, }; const NOTICE_TONE: Record = { @@ -129,7 +131,7 @@ function SpaceRail({ onToggleSettings: () => void; }) { const { t } = useI18n(); - const labels: Record = { local: t("rail.local"), cloud: t("rail.cloud"), chat: t("rail.chat") }; + const labels: Record = { local: t("rail.local"), cloud: t("rail.cloud"), chat: t("rail.chat"), stats: t("rail.stats") }; return (