From 76779b6e14115ab5032812ec34ad36513667e262 Mon Sep 17 00:00:00 2001 From: fengpanhe Date: Thu, 22 Oct 2026 16:43:14 +0800 Subject: [PATCH 1/3] =?UTF-8?q?feat:=20=E4=B8=BA=20TRAE=20CLI=20=E4=BC=9A?= =?UTF-8?q?=E8=AF=9D=E6=B7=BB=E5=8A=A0=E5=8E=9F=E7=94=9F=20token=20?= =?UTF-8?q?=E4=BD=BF=E7=94=A8=E9=87=8F=E7=9A=84=E5=AF=BC=E5=85=A5=E5=92=8C?= =?UTF-8?q?=E7=BB=9F=E8=AE=A1?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 新增 trae_usage 解析模块，用于从 TRAE CLI 原生 rollout 文件的 token_usage_record 行中读取精确的使用量数据，并将其补充到 ACP 会话的历史记录中。同时修改相关服务和解析逻辑，将 schema 版本升级到 7，新增 TRAE CLI 会话目录的解析路径，支持内置和自定义的 TRAE CLI 代理的使用量统计，解决了之前仅使用 ACP 上下文占用数据不准确的问题。 --- src-tauri/src/acp/registry.rs | 5 +- src-tauri/src/acp/trae.rs | 96 ++++- src-tauri/src/commands/token_usage.rs | 8 +- src-tauri/src/db/service/import_service.rs | 3 +- src-tauri/src/parsers/acp_native.rs | 187 +++++++++- src-tauri/src/parsers/mod.rs | 11 +- src-tauri/src/parsers/trae_usage.rs | 394 +++++++++++++++++++++ 7 files changed, 684 insertions(+), 20 deletions(-) create mode 100644 src-tauri/src/parsers/trae_usage.rs diff --git a/src-tauri/src/acp/registry.rs b/src-tauri/src/acp/registry.rs index 4286aa838de..fb76bda59f2 100644 --- a/src-tauri/src/acp/registry.rs +++ b/src-tauri/src/acp/registry.rs @@ -5514,8 +5514,9 @@ pub fn get_agent_meta(agent_type: AgentType) -> AcpAgentMeta { // install and pins none; `~/.local/bin` is probed by the system // resolver even when a Finder-launched app's PATH lacks it. // - // History is codeg's own ACP transcript - // (`AgentType::uses_acp_transcript`), not TRAE CLI's rollouts — + // Message history is codeg's own ACP transcript + // (`AgentType::uses_acp_transcript`); only exact + // `token_usage_record` rows are read from TRAE CLI's rollouts — // see `acp::trae` for why. MCP servers // come from `/traecli.toml`, which the CLI loads itself // and codeg's MCP settings manage, so codeg does not forward the diff --git a/src-tauri/src/acp/trae.rs b/src-tauri/src/acp/trae.rs index 33e65c932a3..b91cb4f457d 100644 --- a/src-tauri/src/acp/trae.rs +++ b/src-tauri/src/acp/trae.rs @@ -3,13 +3,15 @@ //! TRAE CLI 2.0 is ByteDance's terminal coding agent. It is a fork of OpenAI's //! Codex CLI (the binary its installer ships is `traex`, linked into //! `~/.local/bin` as `traecli`) with its own ACP server, `traecli acp serve`. -//! codeg launches the user's own install (see the registry entry) and keeps no -//! parser for its store: TRAE CLI writes Codex-derived rollouts in a paginated -//! format of its own (`history_mutation` / `world_state` records under +//! codeg launches the user's own install (see the registry entry). TRAE CLI +//! writes Codex-derived rollouts in a paginated format of its own +//! (`history_mutation` / `world_state` records under //! `/cli/sessions/YYYY/MM/DD/rollout-*.jsonl`, plus SQLite state -//! beside them) that no public contract describes, so history comes from -//! codeg's own ACP transcript (`parsers::acp_native`), like ZCode and custom -//! agents. +//! beside them) that no public contract describes, so message history comes +//! from codeg's own ACP transcript (`parsers::acp_native`), like ZCode and +//! custom agents. The stable `token_usage_record` rows in those rollouts are +//! read separately by `parsers::trae_usage` to enrich that history with exact +//! usage. //! //! Everything here was verified against `traecli 0.209.1` (public channel) //! with an isolated `HOME`: @@ -36,13 +38,18 @@ use std::collections::BTreeMap; use std::ffi::OsString; -use std::path::PathBuf; +use std::path::{Path, PathBuf}; +use crate::acp::registry::AgentDistribution; +use crate::models::AgentType; use crate::parsers::AgentEnv; /// Overrides the `~/.trae` directory (config, skills, sessions). pub(crate) const TRAE_HOME_ENV: &str = "TRAE_HOME"; +/// Overrides TRAE CLI's runtime directory below [`TRAE_HOME_ENV`]. +pub(crate) const TRAECLI_HOME_ENV: &str = "TRAECLI_HOME"; + /// TRAE CLI's personal access token: the sign-in for a machine without a /// browser (a server or Docker install). TRAE CLI reads it only when its PAT /// auth method runs (`traecli login --with-trae-pat`, or ACP `authenticate` @@ -72,6 +79,19 @@ pub(crate) fn resolve_trae_home_dir_for_env(env: &AgentEnv) -> PathBuf { resolve_trae_home_dir_from(env.var_os(TRAE_HOME_ENV), dirs::home_dir()) } +/// The runtime directory that owns TRAE's native session rollouts. +pub(crate) fn resolve_trae_cli_home_dir_for_env(env: &AgentEnv) -> PathBuf { + env.var_os(TRAECLI_HOME_ENV) + .filter(|value| !value.is_empty()) + .map(PathBuf::from) + .unwrap_or_else(|| resolve_trae_home_dir_for_env(env).join("cli")) +} + +/// `/sessions`, where `token_usage_record` rows are written. +pub(crate) fn resolve_trae_sessions_dir_for_env(env: &AgentEnv) -> PathBuf { + resolve_trae_cli_home_dir_for_env(env).join("sessions") +} + fn resolve_trae_home_dir_from(trae_home: Option, home: Option) -> PathBuf { if let Some(dir) = trae_home.filter(|value| !value.is_empty()) { return PathBuf::from(dir); @@ -79,6 +99,39 @@ fn resolve_trae_home_dir_from(trae_home: Option, home: Option home.unwrap_or_default().join(".trae") } +/// Whether this agent is backed by TRAE CLI's native rollout store. +/// +/// Besides the built-in, this recognizes installations created before TRAE +/// became built in, when users registered `traecli` as a custom ACP agent. +pub(crate) fn uses_trae_rollouts(agent_type: AgentType) -> bool { + match agent_type { + AgentType::Trae => true, + AgentType::Custom(id) => { + matches!( + id.to_ascii_lowercase().as_str(), + "trae" | "traecli" | "trae-cli" + ) || crate::acp::custom_registry::get(id) + .is_some_and(|meta| distribution_runs_trae(&meta.distribution)) + } + _ => false, + } +} + +fn distribution_runs_trae(distribution: &AgentDistribution) -> bool { + let command = match distribution { + AgentDistribution::Npx { cmd, .. } + | AgentDistribution::Binary { cmd, .. } + | AgentDistribution::Uvx { cmd, .. } + | AgentDistribution::System { cmd, .. } => cmd, + }; + let trimmed = command.trim(); + let basename = Path::new(trimmed) + .file_stem() + .and_then(|name| name.to_str()) + .unwrap_or(trimmed); + basename.eq_ignore_ascii_case("traecli") || basename.eq_ignore_ascii_case("traex") +} + /// `/traecli.toml`, the user config TRAE CLI loads at startup (and /// where its `[mcp_servers]` live). pub(crate) fn trae_config_toml_path() -> PathBuf { @@ -128,6 +181,35 @@ mod tests { ); } + #[test] + fn traecli_home_relocates_the_native_sessions_store() { + let env = AgentEnv::with_overrides(&[ + (TRAE_HOME_ENV, "/srv/trae"), + (TRAECLI_HOME_ENV, "/data/trae-runtime"), + ]); + assert_eq!( + resolve_trae_sessions_dir_for_env(&env), + PathBuf::from("/data/trae-runtime/sessions") + ); + + let env = AgentEnv::with_overrides(&[(TRAE_HOME_ENV, "/srv/trae")]); + assert_eq!( + resolve_trae_sessions_dir_for_env(&env), + PathBuf::from("/srv/trae/cli/sessions") + ); + } + + #[test] + fn builtin_and_legacy_custom_trae_agents_use_native_rollouts() { + assert!(uses_trae_rollouts(AgentType::Trae)); + assert!(uses_trae_rollouts( + AgentType::custom("traecli").unwrap() + )); + assert!(!uses_trae_rollouts( + AgentType::custom("unrelated-agent").unwrap() + )); + } + #[test] fn pat_sign_in_follows_what_the_child_will_see() { let with = |pairs: &[(&str, &str)]| -> BTreeMap { diff --git a/src-tauri/src/commands/token_usage.rs b/src-tauri/src/commands/token_usage.rs index 2094f6f1aed..348e5fa0d5f 100644 --- a/src-tauri/src/commands/token_usage.rs +++ b/src-tauri/src/commands/token_usage.rs @@ -135,6 +135,12 @@ const TOP_CONVERSATIONS: usize = 8; /// them again. Forks are on by default in the terminal since CLI 2.1.232 /// and in every Agent SDK session since 2.1.302, so unchanged transcripts /// read differently now. See `parsers::claude::is_fork_directive_record`. +/// * `7` — TRAE CLI usage counted from its native `token_usage_record` rows. +/// TRAE's ACP `usage_update` is only context occupancy, so both the built-in +/// and pre-integration custom `traecli` sessions previously wrote no facts. +/// Rollout `thread_id` is the ACP session id, which makes the native counters +/// an exact enrichment of codeg's ACP-projected history. See +/// `parsers::trae_usage`. /// /// Only the accounting stored in `token_usage_turn` counts: the four token /// counters, the duration and the timestamp. A conversation's context WINDOW is @@ -149,7 +155,7 @@ const TOP_CONVERSATIONS: usize = 8; /// counters predate it, since a custom/BYO model has always exposed them and /// the parser reads every session under `~/.qoder/projects`, not just the ones /// codeg launched. -const FACT_SCHEMA_VERSION: &str = "6"; +const FACT_SCHEMA_VERSION: &str = "7"; const FACT_SCHEMA_VERSION_KEY: &str = "token_usage_fact_schema_version"; diff --git a/src-tauri/src/db/service/import_service.rs b/src-tauri/src/db/service/import_service.rs index b66a46b4318..a010173f765 100644 --- a/src-tauri/src/db/service/import_service.rs +++ b/src-tauri/src/db/service/import_service.rs @@ -29,7 +29,8 @@ const ALL_PARSER_AGENTS: [AgentType; 19] = [ AgentType::Kiro, // Codeg-owned ACP history only; no scan of ZCode's native SQLite store. AgentType::ZCode, - // Codeg-owned ACP history only; TRAE CLI's own rollouts are not read. + // Codeg-owned ACP message history, enriched with usage from TRAE's native + // token_usage_record rows. AgentType::Trae, // Codeg-owned ACP history only; Command Code's own session trees are not // read. diff --git a/src-tauri/src/parsers/acp_native.rs b/src-tauri/src/parsers/acp_native.rs index 42d10b67ea9..f2457e9b93b 100644 --- a/src-tauri/src/parsers/acp_native.rs +++ b/src-tauri/src/parsers/acp_native.rs @@ -57,24 +57,50 @@ use crate::acp_transcript::{self, EntryKind, Transcript, TranscriptEntry}; use crate::models::agent::AgentType; use crate::models::conversation::{ConversationDetail, ConversationSummary, SessionStats}; use crate::models::message::{ContentBlock, ImageData, MessageTurn, TurnRole, TurnUsage}; -use crate::parsers::{user_turn_block, user_turn_block_from_wire, AgentParser, ParseError}; +use crate::parsers::{ + user_turn_block, user_turn_block_from_wire, AgentEnv, AgentParser, ParseError, +}; pub struct AcpNativeParser { agent_type: AgentType, root: PathBuf, + trae_sessions: Option, } impl AcpNativeParser { pub fn new(agent_type: AgentType) -> Self { + Self::for_env(agent_type, &AgentEnv::inherited()) + } + + pub fn for_env(agent_type: AgentType, env: &AgentEnv) -> Self { Self { agent_type, root: crate::paths::codeg_acp_transcripts_root(), + trae_sessions: crate::acp::trae::uses_trae_rollouts(agent_type) + .then(|| crate::acp::trae::resolve_trae_sessions_dir_for_env(env)), } } /// Root-injectable constructor for tests. pub fn new_in(agent_type: AgentType, root: PathBuf) -> Self { - Self { agent_type, root } + Self { + agent_type, + root, + trae_sessions: None, + } + } + + #[cfg(test)] + pub fn new_in_with_trae_sessions( + agent_type: AgentType, + root: PathBuf, + trae_sessions: PathBuf, + ) -> Self { + Self { + agent_type, + root, + trae_sessions: Some(trae_sessions), + } } /// Directory name for this agent's transcripts: its ACP registry id. @@ -113,6 +139,12 @@ impl AgentParser for AcpNativeParser { fn get_conversation(&self, conversation_id: &str) -> Result { let dir = self.agent_dir(); + let mut native_session_ids = vec![conversation_id.to_string()]; + native_session_ids.extend(acp_transcript::continuation_ancestors_in( + &self.root, + dir, + conversation_id, + )); // Reads the whole continuation chain: when the agent had forgotten the // session and codeg started a fresh one, the earlier turns still live // under the previous session id. @@ -120,10 +152,23 @@ impl AgentParser for AcpNativeParser { if transcript.header.is_none() && transcript.is_empty() { return Err(ParseError::ConversationNotFound(conversation_id.to_string())); } - let turns = project_turns(&transcript.entries); + let mut turns = project_turns(&transcript.entries); + let native_usage = self.trae_sessions.as_deref().and_then(|sessions| { + super::trae_usage::enrich_turns(sessions, &native_session_ids, &mut turns) + }); let (used, size) = latest_context_window(&transcript.entries); + let stats = session_stats(&turns); + let stats = if stats + .as_ref() + .and_then(|stats| stats.total_usage.as_ref()) + .is_none() + { + session_stats_with_native_usage(stats, native_usage) + } else { + stats + }; let session_stats = crate::parsers::merge_context_window_stats( - session_stats(&turns), + stats, used, size, ); @@ -1052,6 +1097,36 @@ fn session_stats(turns: &[MessageTurn]) -> Option { }) } +/// Preserve native usage from a TRAE turn that produced no renderable +/// assistant message. Normally `trae_usage::enrich_turns` attaches every +/// native turn to an ACP turn and the regular aggregation above wins. +fn session_stats_with_native_usage( + stats: Option, + native_usage: Option, +) -> Option { + let usage = native_usage?; + let total_tokens = usage + .input_tokens + .saturating_add(usage.output_tokens) + .saturating_add(usage.cache_creation_input_tokens) + .saturating_add(usage.cache_read_input_tokens); + match stats { + Some(mut stats) => { + stats.total_usage = Some(usage); + stats.total_tokens = Some(total_tokens); + Some(stats) + } + None => Some(SessionStats { + total_usage: Some(usage), + total_tokens: Some(total_tokens), + total_duration_ms: 0, + context_window_used_tokens: None, + context_window_max_tokens: None, + context_window_usage_percent: None, + }), + } +} + #[cfg(test)] mod tests { use super::*; @@ -2372,6 +2447,110 @@ mod tests { let _ = std::fs::remove_dir_all(&root); } + #[tokio::test] + async fn custom_traecli_history_is_enriched_from_its_native_usage_record() { + let transcript_dir = tempfile::tempdir().unwrap(); + let native_dir = tempfile::tempdir().unwrap(); + let session_id = "01a13673-cfd5-7c23-bc00-ce5be75fd42b"; + let agent = AgentType::custom("traecli").unwrap(); + let agent_dir = crate::acp::registry::registry_id_for(agent); + + let header = TranscriptHeader::new( + "custom:traecli", + session_id, + "/workspace", + 1_791_914_922_197, + ); + acp_transcript::record_header_in(transcript_dir.path(), agent_dir, &header) + .await + .unwrap(); + for (kind, payload) in [ + ( + EntryKind::Prompt, + serde_json::json!([{"type":"text","text":"hello"}]), + ), + ( + EntryKind::Update, + serde_json::json!({ + "sessionUpdate":"agent_message_chunk", + "content":{"type":"text","text":"done"} + }), + ), + ( + EntryKind::Update, + serde_json::json!({ + "sessionUpdate":"usage_update", + "used":20_584, + "size":760_000 + }), + ), + ( + EntryKind::TurnEnd, + serde_json::json!({ + "stopReason":"end_turn", + "durationMs":3_497, + "model":"gpt-5.6-sol" + }), + ), + ] { + acp_transcript::record_entry_in( + transcript_dir.path(), + agent_dir, + session_id, + kind, + payload, + ) + .await + .unwrap(); + } + + let day = native_dir.path().join("2026/10/14"); + std::fs::create_dir_all(&day).unwrap(); + let rollout = day.join(format!( + "rollout-2026-10-14T02-08-42-{session_id}.jsonl" + )); + std::fs::write( + rollout, + serde_json::json!({ + "timestamp":"2026-10-13T18:08:45.732Z", + "type":"token_usage_record", + "payload":{ + "thread_id":session_id, + "session_id":session_id, + "turn_id":"turn-1", + "response_id":"resp-1", + "usage":{ + "input_tokens":20_454, + "cache_creation_input_tokens":128, + "cached_input_tokens":19_968, + "output_tokens":130, + "reasoning_output_tokens":24, + "total_tokens":20_584 + } + } + }) + .to_string() + + "\n", + ) + .unwrap(); + + let parser = AcpNativeParser::new_in_with_trae_sessions( + agent, + transcript_dir.path().to_path_buf(), + native_dir.path().to_path_buf(), + ); + let detail = parser.get_conversation(session_id).unwrap(); + let usage = detail.turns[1].usage.as_ref().unwrap(); + assert_eq!(usage.input_tokens, 358); + assert_eq!(usage.cache_creation_input_tokens, 128); + assert_eq!(usage.cache_read_input_tokens, 19_968); + assert_eq!(usage.output_tokens, 130); + let stats = detail.session_stats.unwrap(); + assert_eq!(stats.total_tokens, Some(20_584)); + assert_eq!(stats.context_window_used_tokens, Some(20_584)); + assert_eq!(stats.context_window_max_tokens, Some(760_000)); + } + #[test] fn a_session_opened_but_never_used_is_not_listed() { // Opening a conversation writes the ACP header before the user has diff --git a/src-tauri/src/parsers/mod.rs b/src-tauri/src/parsers/mod.rs index 6557b5e4d90..493cbeaa91f 100644 --- a/src-tauri/src/parsers/mod.rs +++ b/src-tauri/src/parsers/mod.rs @@ -20,6 +20,7 @@ mod opencode_context_window; pub mod pi; pub mod qoder; mod summary_cache; +mod trae_usage; use std::collections::{HashMap, HashSet}; use std::path::PathBuf; @@ -419,12 +420,12 @@ pub fn build_agent_parser(agent_type: AgentType, env: &AgentEnv) -> Box Box::new(antigravity::AntigravityParser::for_env(env)), AgentType::Kiro => Box::new(kiro::KiroParser::for_env(env)), // Custom ACP agents have no native store to reverse-engineer; their - // history is codeg's own ACP transcript. TRAE CLI's rollouts use an - // undocumented paginated format of its own (see `acp::trae`), and so - // do Command Code's session trees (see `acp::command_code`), so both - // take the same route. + // history is codeg's own ACP transcript. TRAE CLI takes the same route + // for message history, but AcpNativeParser enriches those turns from + // the precise token_usage_record rows in TRAE's native rollouts. + // Command Code's undocumented session tree remains ACP-only. AgentType::Custom(_) | AgentType::ZCode | AgentType::Trae | AgentType::CommandCode => { - Box::new(acp_native::AcpNativeParser::new(agent_type)) + Box::new(acp_native::AcpNativeParser::for_env(agent_type, env)) } }; Box::new(RouteSanitized(inner)) diff --git a/src-tauri/src/parsers/trae_usage.rs b/src-tauri/src/parsers/trae_usage.rs new file mode 100644 index 00000000000..2943fcbf52c --- /dev/null +++ b/src-tauri/src/parsers/trae_usage.rs @@ -0,0 +1,394 @@ +//! Token usage enrichment for TRAE CLI sessions. +//! +//! TRAE's ACP `usage_update` is context occupancy, not billable usage. The +//! precise per-request counters live in the CLI's native rollout as +//! `token_usage_record` rows. ACP session ids and rollout thread ids are the +//! same, so this module can enrich codeg's ACP-projected turns without parsing +//! TRAE's private message/history format. + +use std::collections::{HashMap, HashSet}; +use std::fs::File; +use std::io::{BufRead, BufReader}; +use std::path::{Path, PathBuf}; + +use chrono::{DateTime, Utc}; +use walkdir::WalkDir; + +use crate::models::{MessageTurn, TurnRole, TurnUsage}; + +#[derive(Debug)] +struct UsageRecord { + turn_id: Option, + response_id: String, + occurred_at: Option>, + usage: TurnUsage, +} + +/// Read every usage record belonging to `session_ids`, attach each native turn +/// to the closest ACP assistant turn, and return the native session total. +/// +/// The returned total remains useful when a failed/empty turn produced usage +/// but no assistant turn for the ACP projection to carry. +pub(crate) fn enrich_turns( + sessions_dir: &Path, + session_ids: &[String], + turns: &mut [MessageTurn], +) -> Option { + let records = read_usage_records(sessions_dir, session_ids); + if records.is_empty() { + return None; + } + + let total = records.iter().fold(TurnUsage::default(), |sum, record| { + add_usage(&sum, &record.usage) + }); + let assistant_indices: Vec = turns + .iter() + .enumerate() + .filter_map(|(index, turn)| matches!(turn.role, TurnRole::Assistant).then_some(index)) + .collect(); + if assistant_indices.is_empty() { + return Some(total); + } + + // One TRAE turn may make several model requests around tool calls. Keep + // those requests together before choosing an ACP turn, then write their sum + // once so the conversation footer and dashboard cannot double-count them. + let mut groups: Vec<(String, Vec<&UsageRecord>)> = Vec::new(); + let mut group_index: HashMap = HashMap::new(); + for record in &records { + let key = record + .turn_id + .clone() + .unwrap_or_else(|| format!("response:{}", record.response_id)); + let index = match group_index.get(&key).copied() { + Some(index) => index, + None => { + let index = groups.len(); + group_index.insert(key.clone(), index); + groups.push((key, Vec::new())); + index + } + }; + groups[index].1.push(record); + } + + let mut assigned: HashMap = HashMap::new(); + for (sequence, (_, group)) in groups.into_iter().enumerate() { + let occurred_at = group.iter().find_map(|record| record.occurred_at); + let target = occurred_at + .map(|at| closest_assistant_turn(turns, &assistant_indices, at)) + .unwrap_or_else(|| assistant_indices[sequence.min(assistant_indices.len() - 1)]); + let group_usage = group.iter().fold(TurnUsage::default(), |sum, record| { + add_usage(&sum, &record.usage) + }); + assigned + .entry(target) + .and_modify(|usage| *usage = add_usage(usage, &group_usage)) + .or_insert(group_usage); + } + + // Native rollout counters are authoritative for the turns they cover. A + // future TRAE ACP release may also report those counters at turn end; an + // overwrite here prevents the two copies from being added together. + for (index, usage) in assigned { + turns[index].usage = Some(usage); + } + Some(total) +} + +fn read_usage_records(sessions_dir: &Path, session_ids: &[String]) -> Vec { + if session_ids.is_empty() || !sessions_dir.exists() { + return Vec::new(); + } + let wanted: HashSet<&str> = session_ids.iter().map(String::as_str).collect(); + let mut seen: HashSet<(String, String)> = HashSet::new(); + let mut records = Vec::new(); + + for path in rollout_paths(sessions_dir, &wanted) { + let Ok(file) = File::open(&path) else { + continue; + }; + let mut reader = BufReader::new(file); + let mut bytes = Vec::new(); + let mut line_no = 0usize; + loop { + bytes.clear(); + let read = match reader.read_until(b'\n', &mut bytes) { + Ok(read) => read, + Err(_) => break, + }; + if read == 0 { + break; + } + line_no += 1; + let line = String::from_utf8_lossy(&bytes); + let Ok(row) = serde_json::from_str::(&line) else { + continue; + }; + if row.get("type").and_then(serde_json::Value::as_str) != Some("token_usage_record") { + continue; + } + let Some(payload) = row.get("payload") else { + continue; + }; + let Some(thread_id) = string_id(payload.get("thread_id")) + .or_else(|| string_id(payload.get("session_id"))) + else { + continue; + }; + if !wanted.contains(thread_id.as_str()) { + continue; + } + let response_id = string_id(payload.get("response_id")) + .unwrap_or_else(|| format!("{}:{line_no}", path.to_string_lossy())); + if !seen.insert((thread_id, response_id.clone())) { + continue; + } + let Some(usage) = payload.get("usage").and_then(usage_from_value) else { + continue; + }; + records.push(UsageRecord { + turn_id: string_id(payload.get("turn_id")) + .or_else(|| string_id(payload.get("root_turn_id"))), + response_id, + occurred_at: row + .get("timestamp") + .and_then(serde_json::Value::as_str) + .and_then(parse_timestamp), + usage, + }); + } + } + + records.sort_by_key(|record| record.occurred_at); + records +} + +fn rollout_paths(sessions_dir: &Path, session_ids: &HashSet<&str>) -> Vec { + WalkDir::new(sessions_dir) + .into_iter() + .filter_map(Result::ok) + .filter(|entry| entry.file_type().is_file()) + .map(walkdir::DirEntry::into_path) + .filter(|path| { + let Some(name) = path.file_name().and_then(|name| name.to_str()) else { + return false; + }; + if !name.starts_with("rollout-") || !name.ends_with(".jsonl") { + return false; + } + session_ids.iter().any(|id| { + name.ends_with(&format!("-{id}.jsonl")) || name.contains(&format!("-{id}_")) + }) + }) + .collect() +} + +fn string_id(value: Option<&serde_json::Value>) -> Option { + let value = value?; + if let Some(value) = value + .as_str() + .map(str::trim) + .filter(|value| !value.is_empty()) + { + return Some(value.to_string()); + } + (value.is_number() || value.is_boolean()).then(|| value.to_string()) +} + +fn parse_timestamp(raw: &str) -> Option> { + DateTime::parse_from_rfc3339(raw) + .ok() + .map(|value| value.with_timezone(&Utc)) +} + +/// TRAE follows Codex semantics: `input_tokens` is the whole prompt and both +/// cache counters are subsets of it. Codeg stores disjoint counters, so split +/// the cache portions out. `reasoning_output_tokens` is already included in +/// TRAE's `output_tokens` (and consequently in `total_tokens`), so it must not +/// be added again. +fn usage_from_value(value: &serde_json::Value) -> Option { + let field = |name: &str| { + value + .get(name) + .and_then(serde_json::Value::as_u64) + .unwrap_or(0) + }; + let input = field("input_tokens"); + let cache_read = field("cached_input_tokens"); + let cache_write = field("cache_creation_input_tokens").min(input.saturating_sub(cache_read)); + let usage = TurnUsage { + input_tokens: input.saturating_sub(cache_read).saturating_sub(cache_write), + output_tokens: field("output_tokens"), + cache_creation_input_tokens: cache_write, + cache_read_input_tokens: cache_read, + }; + (usage.input_tokens > 0 + || usage.output_tokens > 0 + || usage.cache_creation_input_tokens > 0 + || usage.cache_read_input_tokens > 0) + .then_some(usage) +} + +fn add_usage(left: &TurnUsage, right: &TurnUsage) -> TurnUsage { + TurnUsage { + input_tokens: left.input_tokens.saturating_add(right.input_tokens), + output_tokens: left.output_tokens.saturating_add(right.output_tokens), + cache_creation_input_tokens: left + .cache_creation_input_tokens + .saturating_add(right.cache_creation_input_tokens), + cache_read_input_tokens: left + .cache_read_input_tokens + .saturating_add(right.cache_read_input_tokens), + } +} + +fn closest_assistant_turn( + turns: &[MessageTurn], + assistant_indices: &[usize], + occurred_at: DateTime, +) -> usize { + assistant_indices + .iter() + .copied() + .min_by_key(|index| { + let turn = &turns[*index]; + let end = turn.completed_at.unwrap_or(turn.timestamp); + if occurred_at >= turn.timestamp && occurred_at <= end { + 0 + } else if occurred_at < turn.timestamp { + (turn.timestamp - occurred_at) + .num_milliseconds() + .unsigned_abs() + } else { + (occurred_at - end).num_milliseconds().unsigned_abs() + } + }) + .unwrap_or(assistant_indices[0]) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::models::{ContentBlock, MessageTurn}; + use tempfile::TempDir; + + fn at(raw: &str) -> DateTime { + DateTime::parse_from_rfc3339(raw) + .unwrap() + .with_timezone(&Utc) + } + + fn assistant(start: &str, end: &str) -> MessageTurn { + MessageTurn { + id: start.to_string(), + role: TurnRole::Assistant, + blocks: vec![ContentBlock::Text { + text: "answer".to_string(), + }], + timestamp: at(start), + usage: None, + duration_ms: None, + model: Some("gpt-5.6-sol".to_string()), + completed_at: Some(at(end)), + agent_message_id: None, + } + } + + fn write_rollout(root: &Path, session_id: &str, lines: &[serde_json::Value]) { + let dir = root.join("2026/10/13"); + std::fs::create_dir_all(&dir).unwrap(); + let path = dir.join(format!("rollout-2026-10-13T10-00-00-{session_id}.jsonl")); + let content = lines + .iter() + .map(serde_json::Value::to_string) + .collect::>() + .join("\n"); + std::fs::write(path, format!("{content}\n")).unwrap(); + } + + #[test] + fn trae_usage_splits_cached_input_without_double_counting_reasoning() { + let usage = usage_from_value(&serde_json::json!({ + "input_tokens": 20_454, + "cache_creation_input_tokens": 128, + "cached_input_tokens": 19_968, + "output_tokens": 130, + "reasoning_output_tokens": 24, + "total_tokens": 20_584 + })) + .unwrap(); + assert_eq!(usage.input_tokens, 358); + assert_eq!(usage.cache_creation_input_tokens, 128); + assert_eq!(usage.cache_read_input_tokens, 19_968); + assert_eq!(usage.output_tokens, 130); + assert_eq!( + usage.input_tokens + + usage.cache_creation_input_tokens + + usage.cache_read_input_tokens + + usage.output_tokens, + 20_584 + ); + } + + #[test] + fn records_are_deduped_and_grouped_onto_their_acp_turns() { + let dir = TempDir::new().unwrap(); + let session_id = "01a13673-cfd5-7c23-bc00-ce5be75fd42b"; + let record = |timestamp: &str, + turn_id: &str, + response_id: &str, + input: u64, + cached: u64, + output: u64| { + serde_json::json!({ + "timestamp": timestamp, + "type": "token_usage_record", + "payload": { + "thread_id": session_id, + "session_id": session_id, + "turn_id": turn_id, + "response_id": response_id, + "usage": { + "input_tokens": input, + "cache_creation_input_tokens": 0, + "cached_input_tokens": cached, + "output_tokens": output, + "reasoning_output_tokens": 1, + "total_tokens": input + output + } + } + }) + }; + write_rollout( + dir.path(), + session_id, + &[ + record("2026-10-13T10:00:02Z", "turn-1", "resp-1", 100, 80, 5), + // Same request repeated by a replay/pagination boundary. + record("2026-10-13T10:00:02Z", "turn-1", "resp-1", 100, 80, 5), + record("2026-10-13T10:00:03Z", "turn-1", "resp-2", 120, 90, 7), + record("2026-10-13T10:01:02Z", "turn-2", "resp-3", 200, 150, 9), + ], + ); + + let mut turns = vec![ + assistant("2026-10-13T10:00:00Z", "2026-10-13T10:00:10Z"), + assistant("2026-10-13T10:01:00Z", "2026-10-13T10:01:10Z"), + ]; + let total = enrich_turns(dir.path(), &[session_id.to_string()], &mut turns).unwrap(); + + let first = turns[0].usage.as_ref().unwrap(); + assert_eq!(first.input_tokens, 50); + assert_eq!(first.cache_read_input_tokens, 170); + assert_eq!(first.output_tokens, 12); + let second = turns[1].usage.as_ref().unwrap(); + assert_eq!(second.input_tokens, 50); + assert_eq!(second.cache_read_input_tokens, 150); + assert_eq!(second.output_tokens, 9); + assert_eq!(total.input_tokens, 100); + assert_eq!(total.cache_read_input_tokens, 320); + assert_eq!(total.output_tokens, 21); + } +} From 43bd7b06f2a5f9f7ec597c178303458a34b7364c Mon Sep 17 00:00:00 2001 From: xintaofei Date: Mon, 26 Oct 2026 23:02:41 +0800 Subject: [PATCH 2/3] fix(trae): count sub-agent spend in TRAE CLI token usage A TRAE CLI sub-agent runs on a thread of its own and writes its own rollout, named by that thread's id, so reading only the session's rollout dropped everything its sub-agents spent. Follow the spawn events recorded on the session's own threads to the child rollouts, at any depth, and group every row under its root turn so a sub-agent's spend lands on the turn that launched it, even when it reports after that turn ended. A forked session's rollout opens with a copy of the parent's rows and spawn events, still stamped with the parent's thread; those stay out of the fork's numbers, and a sub-agent row naming another root session is never counted. Parse only the lines that are usage rows or spawn events, walk the sessions tree once per read without descending into the per-rollout .artifacts directories, and resolve the continuation chain only for TRAE. --- src-tauri/src/acp/trae.rs | 7 +- src-tauri/src/commands/token_usage.rs | 11 +- src-tauri/src/parsers/acp_native.rs | 13 +- src-tauri/src/parsers/trae_usage.rs | 555 ++++++++++++++++++++++---- 4 files changed, 502 insertions(+), 84 deletions(-) diff --git a/src-tauri/src/acp/trae.rs b/src-tauri/src/acp/trae.rs index b91cb4f457d..35d7d59c117 100644 --- a/src-tauri/src/acp/trae.rs +++ b/src-tauri/src/acp/trae.rs @@ -9,7 +9,7 @@ //! `/cli/sessions/YYYY/MM/DD/rollout-*.jsonl`, plus SQLite state //! beside them) that no public contract describes, so message history comes //! from codeg's own ACP transcript (`parsers::acp_native`), like ZCode and -//! custom agents. The stable `token_usage_record` rows in those rollouts are +//! custom agents. The flat `token_usage_record` rows in those rollouts are //! read separately by `parsers::trae_usage` to enrich that history with exact //! usage. //! @@ -20,6 +20,11 @@ //! tree: `traecli.toml`, `skills/`, and the `cli/` runtime dir (sessions, //! SQLite state, logs). The value is taken verbatim — `TRAE_HOME=~/x` is NOT //! expanded, and a path that does not exist is a startup error. +//! `TRAECLI_HOME` relocates the `cli/` runtime dir alone. +//! * A rollout's thread id is the ACP session id. Every model request writes +//! one `token_usage_record` row; a sub-agent's rows go to its own rollout, +//! named by its own thread id. The prompt response's `usage` covers the +//! session's own requests only, not its sub-agents'. //! * MCP servers live in `traecli.toml` as Codex's `[mcp_servers.]` //! tables, and `traecli acp serve` loads them itself. Servers sent over ACP //! `session/new` are loaded too (both reach the model as `mcp____*`), diff --git a/src-tauri/src/commands/token_usage.rs b/src-tauri/src/commands/token_usage.rs index 348e5fa0d5f..d9c304146b6 100644 --- a/src-tauri/src/commands/token_usage.rs +++ b/src-tauri/src/commands/token_usage.rs @@ -136,10 +136,13 @@ const TOP_CONVERSATIONS: usize = 8; /// and in every Agent SDK session since 2.1.302, so unchanged transcripts /// read differently now. See `parsers::claude::is_fork_directive_record`. /// * `7` — TRAE CLI usage counted from its native `token_usage_record` rows. -/// TRAE's ACP `usage_update` is only context occupancy, so both the built-in -/// and pre-integration custom `traecli` sessions previously wrote no facts. -/// Rollout `thread_id` is the ACP session id, which makes the native counters -/// an exact enrichment of codeg's ACP-projected history. See +/// TRAE's ACP `usage_update` is only context occupancy, and codeg records +/// none of the counters its prompt response carries (which leave out +/// sub-agents anyway), so both the built-in and pre-integration custom +/// `traecli` sessions previously wrote no facts. Rollout `thread_id` is the +/// ACP session id, which makes the native counters an exact enrichment of +/// codeg's ACP-projected history; each sub-agent's own rollout is read too, +/// and its spend lands on the turn that launched it. See /// `parsers::trae_usage`. /// /// Only the accounting stored in `token_usage_turn` counts: the four token diff --git a/src-tauri/src/parsers/acp_native.rs b/src-tauri/src/parsers/acp_native.rs index f2457e9b93b..a5c5180e5c6 100644 --- a/src-tauri/src/parsers/acp_native.rs +++ b/src-tauri/src/parsers/acp_native.rs @@ -139,12 +139,6 @@ impl AgentParser for AcpNativeParser { fn get_conversation(&self, conversation_id: &str) -> Result { let dir = self.agent_dir(); - let mut native_session_ids = vec![conversation_id.to_string()]; - native_session_ids.extend(acp_transcript::continuation_ancestors_in( - &self.root, - dir, - conversation_id, - )); // Reads the whole continuation chain: when the agent had forgotten the // session and codeg started a fresh one, the earlier turns still live // under the previous session id. @@ -154,6 +148,13 @@ impl AgentParser for AcpNativeParser { } let mut turns = project_turns(&transcript.entries); let native_usage = self.trae_sessions.as_deref().and_then(|sessions| { + // Every session of the chain spent tokens under its own id. + let mut native_session_ids = vec![conversation_id.to_string()]; + native_session_ids.extend(acp_transcript::continuation_ancestors_in( + &self.root, + dir, + conversation_id, + )); super::trae_usage::enrich_turns(sessions, &native_session_ids, &mut turns) }); let (used, size) = latest_context_window(&transcript.entries); diff --git a/src-tauri/src/parsers/trae_usage.rs b/src-tauri/src/parsers/trae_usage.rs index 2943fcbf52c..f4d7c2500be 100644 --- a/src-tauri/src/parsers/trae_usage.rs +++ b/src-tauri/src/parsers/trae_usage.rs @@ -1,12 +1,36 @@ //! Token usage enrichment for TRAE CLI sessions. //! -//! TRAE's ACP `usage_update` is context occupancy, not billable usage. The -//! precise per-request counters live in the CLI's native rollout as +//! TRAE's ACP `usage_update` is context occupancy, not billable usage. Its +//! prompt response does carry the turn's counters (`usage`, and +//! `_meta._trae/tokenUsage`), but codeg does not record them, and they cover +//! only the session's own model requests: what the sub-agents it launched +//! spent is not in them (verified against `traecli 0.209.1`). +//! +//! The complete per-request counters live in the CLI's native rollouts as //! `token_usage_record` rows. ACP session ids and rollout thread ids are the //! same, so this module can enrich codeg's ACP-projected turns without parsing //! TRAE's private message/history format. +//! +//! ## Sub-agents +//! +//! A sub-agent runs on a thread of its own and writes its own rollout, named +//! by that thread's id. Its rows carry the child thread as `thread_id`, the +//! root session as `session_id` and the root turn that launched it as +//! `root_turn_id` — at every depth. The parent's rollout names each child in a +//! spawn event (`SubAgentActivity`, `CollabAgentToolCall`), so the reader +//! follows those events to the child rollouts, and from there to their own +//! children. +//! +//! ## Copied history +//! +//! A forked session's rollout opens with a copy of its parent's history: the +//! parent's `token_usage_record` rows and spawn events, still stamped with the +//! parent's thread id. That spend belongs to the parent's conversation, so a +//! row counts only when its thread is one this conversation runs on and its +//! root session is one of this conversation's, and a spawn is followed only +//! when it happened on one of this conversation's threads. -use std::collections::{HashMap, HashSet}; +use std::collections::{HashMap, HashSet, VecDeque}; use std::fs::File; use std::io::{BufRead, BufReader}; use std::path::{Path, PathBuf}; @@ -18,12 +42,25 @@ use crate::models::{MessageTurn, TurnRole, TurnUsage}; #[derive(Debug)] struct UsageRecord { - turn_id: Option, + /// The ROOT turn that spent this: the session's own turn for its own + /// requests, and the turn that launched a sub-agent for the sub-agent's. + root_turn_id: Option, response_id: String, occurred_at: Option>, usage: TurnUsage, } +/// Substrings a rollout line must contain to be worth parsing: a usage row or +/// a spawn event. Everything else in a rollout — history mutations, world +/// state, tool output — is skipped unparsed, so reading the counters costs one +/// scan of the file rather than a JSON parse of every line of it. +const SCANNED_LINE_MARKERS: [&str; 4] = [ + "token_usage_record", + "SubAgentActivity", + "CollabAgentToolCall", + "sub_agent_activity", +]; + /// Read every usage record belonging to `session_ids`, attach each native turn /// to the closest ACP assistant turn, and return the native session total. /// @@ -51,14 +88,16 @@ pub(crate) fn enrich_turns( return Some(total); } - // One TRAE turn may make several model requests around tool calls. Keep - // those requests together before choosing an ACP turn, then write their sum - // once so the conversation footer and dashboard cannot double-count them. + // One TRAE turn may make several model requests around tool calls, and the + // sub-agents it launches make their own. Keep all of them together under + // the root turn before choosing an ACP turn, then write their sum once so + // the conversation footer and dashboard cannot double-count them — and a + // sub-agent that finishes after its turn ended still bills that turn. let mut groups: Vec<(String, Vec<&UsageRecord>)> = Vec::new(); let mut group_index: HashMap = HashMap::new(); for record in &records { let key = record - .turn_id + .root_turn_id .clone() .unwrap_or_else(|| format!("response:{}", record.response_id)); let index = match group_index.get(&key).copied() { @@ -88,9 +127,9 @@ pub(crate) fn enrich_turns( .or_insert(group_usage); } - // Native rollout counters are authoritative for the turns they cover. A - // future TRAE ACP release may also report those counters at turn end; an - // overwrite here prevents the two copies from being added together. + // Native rollout counters are authoritative for the turns they cover. Were + // codeg to record the prompt response's own counters at turn end, an + // overwrite here keeps the two copies from being added together. for (index, usage) in assigned { turns[index].usage = Some(usage); } @@ -101,63 +140,36 @@ fn read_usage_records(sessions_dir: &Path, session_ids: &[String]) -> Vec = session_ids.iter().map(String::as_str).collect(); + // The conversation's own root sessions: its session and the ones it + // continues. Every row it may count carries one of them as `session_id`. + let roots: HashSet<&str> = session_ids.iter().map(String::as_str).collect(); + let rollouts = rollout_files(sessions_dir); + // The threads it runs on: the roots, then every sub-agent thread a spawn + // event on one of them names. Grows while the rollouts are read. + let mut threads: HashSet = session_ids.iter().cloned().collect(); + let mut pending: VecDeque = session_ids.iter().cloned().collect(); + let mut scanned: HashSet<&Path> = HashSet::new(); let mut seen: HashSet<(String, String)> = HashSet::new(); let mut records = Vec::new(); - for path in rollout_paths(sessions_dir, &wanted) { - let Ok(file) = File::open(&path) else { - continue; - }; - let mut reader = BufReader::new(file); - let mut bytes = Vec::new(); - let mut line_no = 0usize; - loop { - bytes.clear(); - let read = match reader.read_until(b'\n', &mut bytes) { - Ok(read) => read, - Err(_) => break, - }; - if read == 0 { - break; - } - line_no += 1; - let line = String::from_utf8_lossy(&bytes); - let Ok(row) = serde_json::from_str::(&line) else { - continue; - }; - if row.get("type").and_then(serde_json::Value::as_str) != Some("token_usage_record") { - continue; - } - let Some(payload) = row.get("payload") else { - continue; - }; - let Some(thread_id) = string_id(payload.get("thread_id")) - .or_else(|| string_id(payload.get("session_id"))) - else { - continue; - }; - if !wanted.contains(thread_id.as_str()) { + while let Some(thread) = pending.pop_front() { + let own_file = format!("-{thread}.jsonl"); + let reverted_file = format!("-{thread}_"); + for (path, name) in &rollouts { + if !(name.ends_with(&own_file) || name.contains(&reverted_file)) { continue; } - let response_id = string_id(payload.get("response_id")) - .unwrap_or_else(|| format!("{}:{line_no}", path.to_string_lossy())); - if !seen.insert((thread_id, response_id.clone())) { + if !scanned.insert(path.as_path()) { continue; } - let Some(usage) = payload.get("usage").and_then(usage_from_value) else { - continue; - }; - records.push(UsageRecord { - turn_id: string_id(payload.get("turn_id")) - .or_else(|| string_id(payload.get("root_turn_id"))), - response_id, - occurred_at: row - .get("timestamp") - .and_then(serde_json::Value::as_str) - .and_then(parse_timestamp), - usage, - }); + scan_rollout( + path, + &roots, + &mut threads, + &mut pending, + &mut seen, + &mut records, + ); } } @@ -165,22 +177,151 @@ fn read_usage_records(sessions_dir: &Path, session_ids: &[String]) -> Vec) -> Vec { +/// Collect one rollout's usage rows into `records`, and queue every sub-agent +/// thread its spawn events name that is not known yet. +fn scan_rollout( + path: &Path, + roots: &HashSet<&str>, + threads: &mut HashSet, + pending: &mut VecDeque, + seen: &mut HashSet<(String, String)>, + records: &mut Vec, +) { + let Ok(file) = File::open(path) else { + return; + }; + let mut reader = BufReader::new(file); + let mut bytes = Vec::new(); + let mut line_no = 0usize; + loop { + bytes.clear(); + let read = match reader.read_until(b'\n', &mut bytes) { + Ok(read) => read, + Err(_) => break, + }; + if read == 0 { + break; + } + line_no += 1; + let line = String::from_utf8_lossy(&bytes); + if !SCANNED_LINE_MARKERS + .iter() + .any(|marker| line.contains(marker)) + { + continue; + } + let Ok(row) = serde_json::from_str::(&line) else { + continue; + }; + let Some(payload) = row.get("payload") else { + continue; + }; + match row.get("type").and_then(serde_json::Value::as_str) { + Some("token_usage_record") => {} + Some("event_msg") => { + // A spawn counts only where it happened: copied history keeps + // the thread it was recorded on, which is not one of ours. + let on_our_thread = string_id(payload.get("thread_id")) + .is_none_or(|thread| threads.contains(&thread)); + if on_our_thread { + for child in spawned_threads(payload) { + if threads.insert(child.clone()) { + pending.push_back(child); + } + } + } + continue; + } + _ => continue, + } + let Some(thread_id) = string_id(payload.get("thread_id")) + .or_else(|| string_id(payload.get("session_id"))) + else { + continue; + }; + if !threads.contains(&thread_id) { + continue; + } + // A sub-agent's row names the root session it ran under; one naming + // another session's root is that session's spend, whatever led here. + if !roots.contains(thread_id.as_str()) + && string_id(payload.get("session_id")) + .is_some_and(|root| !roots.contains(root.as_str())) + { + continue; + } + let Some(usage) = payload.get("usage").and_then(usage_from_value) else { + continue; + }; + let response_id = string_id(payload.get("response_id")) + .unwrap_or_else(|| format!("{}:{line_no}", path.to_string_lossy())); + if !seen.insert((thread_id, response_id.clone())) { + continue; + } + records.push(UsageRecord { + root_turn_id: string_id(payload.get("root_turn_id")) + .or_else(|| string_id(payload.get("turn_id"))), + response_id, + occurred_at: row + .get("timestamp") + .and_then(serde_json::Value::as_str) + .and_then(parse_timestamp), + usage, + }); + } +} + +/// The sub-agent threads an `event_msg` names, in either shape TRAE inherits +/// from Codex: a `SubAgentActivity` item (nested under `item_completed`, or the +/// older flat `sub_agent_activity` event), and the spawn's own +/// `CollabAgentToolCall` item with its `receiver_thread_ids`. +fn spawned_threads(payload: &serde_json::Value) -> Vec { + let event = payload.get("type").and_then(serde_json::Value::as_str); + let source = match event { + Some("sub_agent_activity") => payload, + Some("item_completed") => match payload.get("item") { + Some(item) => item, + None => return Vec::new(), + }, + _ => return Vec::new(), + }; + let kind = (event == Some("sub_agent_activity")) + .then_some("SubAgentActivity") + .or_else(|| source.get("type").and_then(serde_json::Value::as_str)); + match kind { + Some("SubAgentActivity") => string_id(source.get("agent_thread_id")) + .into_iter() + .collect(), + Some("CollabAgentToolCall") => source + .get("receiver_thread_ids") + .and_then(serde_json::Value::as_array) + .into_iter() + .flatten() + .filter_map(|id| string_id(Some(id))) + .collect(), + _ => Vec::new(), + } +} + +/// Every rollout file under `sessions_dir`, with its file name. Walked once per +/// read; the per-rollout `.artifacts` directories (background task output) are +/// not descended into. +fn rollout_files(sessions_dir: &Path) -> Vec<(PathBuf, String)> { WalkDir::new(sessions_dir) .into_iter() + .filter_entry(|entry| { + !(entry.file_type().is_dir() + && entry + .file_name() + .to_str() + .is_some_and(|name| name.ends_with(".artifacts"))) + }) .filter_map(Result::ok) .filter(|entry| entry.file_type().is_file()) - .map(walkdir::DirEntry::into_path) - .filter(|path| { - let Some(name) = path.file_name().and_then(|name| name.to_str()) else { - return false; - }; - if !name.starts_with("rollout-") || !name.ends_with(".jsonl") { - return false; - } - session_ids.iter().any(|id| { - name.ends_with(&format!("-{id}.jsonl")) || name.contains(&format!("-{id}_")) - }) + .filter_map(|entry| { + let name = entry.file_name().to_str()?.to_string(); + (name.starts_with("rollout-") && name.ends_with(".jsonl")) + .then(|| (entry.into_path(), name)) }) .collect() } @@ -391,4 +532,272 @@ mod tests { assert_eq!(total.cache_read_input_tokens, 320); assert_eq!(total.output_tokens, 21); } + + /// Write `lines` as the rollout `file_name` under one day directory. + fn write_named_rollout(root: &Path, file_name: &str, lines: &[serde_json::Value]) { + let dir = root.join("2026/10/26"); + std::fs::create_dir_all(&dir).unwrap(); + let content = lines + .iter() + .map(serde_json::Value::to_string) + .collect::>() + .join("\n"); + std::fs::write(dir.join(file_name), format!("{content}\n")).unwrap(); + } + + /// A `token_usage_record` row in the shape `traecli 0.209.1` writes: + /// `thread_id` is the thread that made the request, `session_id` its root + /// session, `root_turn_id` the root turn it was made for. + #[allow(clippy::too_many_arguments)] + fn usage_row( + at: &str, + thread: &str, + session: &str, + root_turn: &str, + response: &str, + input: u64, + cached: u64, + output: u64, + ) -> serde_json::Value { + serde_json::json!({ + "timestamp": at, + "type": "token_usage_record", + "payload": { + "thread_id": thread, + "turn_id": format!("{thread}:{root_turn}"), + "session_id": session, + "root_turn_id": root_turn, + "response_id": response, + "usage": { + "input_tokens": input, + "cache_creation_input_tokens": 0, + "cached_input_tokens": cached, + "output_tokens": output, + "reasoning_output_tokens": 0, + "total_tokens": input + output + } + } + }) + } + + /// The `SubAgentActivity` event a thread records when it spawns `child`. + fn spawn_activity(at: &str, thread: &str, child: &str) -> serde_json::Value { + serde_json::json!({ + "timestamp": at, + "type": "event_msg", + "payload": { + "type": "item_completed", + "thread_id": thread, + "item": { + "type": "SubAgentActivity", + "id": "subagent-activity:call_1", + "kind": "started", + "agent_thread_id": child + } + } + }) + } + + /// The spawn's own `CollabAgentToolCall` item, naming its receivers. + fn spawn_call(at: &str, thread: &str, children: &[&str]) -> serde_json::Value { + serde_json::json!({ + "timestamp": at, + "type": "event_msg", + "payload": { + "type": "item_completed", + "thread_id": thread, + "item": { + "type": "CollabAgentToolCall", + "id": "call_2", + "tool": "spawn_agent", + "sender_thread_id": thread, + "receiver_thread_ids": children + } + } + }) + } + + fn summed(usage: &TurnUsage) -> u64 { + usage.input_tokens + + usage.output_tokens + + usage.cache_creation_input_tokens + + usage.cache_read_input_tokens + } + + #[test] + fn sub_agent_spend_lands_on_the_turn_that_launched_it() { + let dir = TempDir::new().unwrap(); + let root = "01a17899-f92f-7a10-8a25-10e504be3950"; + let child = "01a17899-f9a9-75f0-893b-81de0432dcef"; + let grandchild = "01a17899-f9ff-7cf0-b812-55fecc2c3a97"; + let reverted = "01a17899-ffff-7000-8000-000000000001"; + write_named_rollout( + dir.path(), + &format!("rollout-2026-10-26T10-00-00-{root}.jsonl"), + &[ + usage_row("2026-10-26T10:00:01Z", root, root, "turn-1", "resp_1", 2000, 0, 20), + spawn_activity("2026-10-26T10:00:01Z", root, child), + usage_row("2026-10-26T10:00:05Z", root, root, "turn-1", "resp_5", 1234, 1000, 56), + ], + ); + // The child spawns a grandchild of its own, which the root's stream + // never mentions; and it reports once more after the root turn ended. + write_named_rollout( + dir.path(), + &format!("rollout-2026-10-26T10-00-01-{child}.jsonl"), + &[ + usage_row("2026-10-26T10:00:02Z", child, root, "turn-1", "resp_2", 2000, 0, 20), + spawn_call("2026-10-26T10:00:02Z", child, &[grandchild]), + usage_row("2026-10-26T10:00:04Z", child, root, "turn-1", "resp_4", 1234, 1000, 56), + usage_row("2026-10-26T10:00:40Z", child, root, "turn-1", "resp_7", 300, 0, 10), + ], + ); + write_named_rollout( + dir.path(), + &format!("rollout-2026-10-26T10-00-02-{grandchild}.jsonl"), + &[usage_row("2026-10-26T10:00:03Z", grandchild, root, "turn-1", "resp_3", 777, 0, 33)], + ); + // A later turn the thread recorded after moving onto a new rollout. + write_named_rollout( + dir.path(), + &format!("rollout-2026-10-26T10-01-00-{root}_{reverted}.jsonl"), + &[usage_row("2026-10-26T10:01:02Z", root, root, "turn-2", "resp_6", 500, 400, 5)], + ); + + let mut turns = vec![ + assistant("2026-10-26T10:00:00Z", "2026-10-26T10:00:10Z"), + assistant("2026-10-26T10:01:00Z", "2026-10-26T10:01:10Z"), + ]; + let total = enrich_turns(dir.path(), &[root.to_string()], &mut turns).unwrap(); + + // Root 2020 + 1290, child 2020 + 1290 + 310, grandchild 810. + let first = turns[0].usage.as_ref().unwrap(); + assert_eq!(summed(first), 7740); + assert_eq!(first.cache_read_input_tokens, 2000); + assert_eq!(first.output_tokens, 20 + 56 + 20 + 56 + 10 + 33); + let second = turns[1].usage.as_ref().unwrap(); + assert_eq!(summed(second), 505); + assert_eq!(summed(&total), 7740 + 505); + } + + #[test] + fn a_fork_counts_only_its_own_spend() { + let dir = TempDir::new().unwrap(); + let parent = "01a178a3-5646-7e53-af99-5c0063df738b"; + let parent_child = "01a178a3-56de-7f72-8abc-6d309f5fbf25"; + let fork = "01a178a3-5f1d-7aa2-b4cb-daf7491a84f6"; + let fork_child = "01a178a3-5f77-71a0-9554-d2998025e6bb"; + let parent_rows = [ + usage_row("2026-10-26T10:00:01Z", parent, parent, "p-turn", "resp_12", 2000, 0, 20), + spawn_activity("2026-10-26T10:00:01Z", parent, parent_child), + spawn_call("2026-10-26T10:00:01Z", parent, &[parent_child]), + usage_row("2026-10-26T10:00:03Z", parent, parent, "p-turn", "resp_14", 1234, 1000, 56), + ]; + write_named_rollout( + dir.path(), + &format!("rollout-2026-10-26T10-00-00-{parent}.jsonl"), + &parent_rows, + ); + write_named_rollout( + dir.path(), + &format!("rollout-2026-10-26T10-00-01-{parent_child}.jsonl"), + &[usage_row( + "2026-10-26T10:00:02Z", + parent_child, + parent, + "p-turn", + "resp_13", + 777, + 0, + 33, + )], + ); + // The fork's rollout opens with a copy of the parent's history — its + // rows and its spawns, still stamped with the parent's thread. + let mut fork_rows = parent_rows.to_vec(); + fork_rows.extend([ + usage_row("2026-10-26T10:05:01Z", fork, fork, "f-turn", "resp_15", 2000, 0, 20), + spawn_activity("2026-10-26T10:05:01Z", fork, fork_child), + usage_row("2026-10-26T10:05:03Z", fork, fork, "f-turn", "resp_17", 1234, 1000, 56), + ]); + write_named_rollout( + dir.path(), + &format!("rollout-2026-10-26T10-05-00-{fork}.jsonl"), + &fork_rows, + ); + write_named_rollout( + dir.path(), + &format!("rollout-2026-10-26T10-05-01-{fork_child}.jsonl"), + &[usage_row( + "2026-10-26T10:05:02Z", + fork_child, + fork, + "f-turn", + "resp_16", + 777, + 0, + 33, + )], + ); + + let mut fork_turns = vec![assistant("2026-10-26T10:05:00Z", "2026-10-26T10:05:10Z")]; + let fork_total = enrich_turns(dir.path(), &[fork.to_string()], &mut fork_turns).unwrap(); + assert_eq!(summed(&fork_total), 2020 + 1290 + 810); + assert_eq!(summed(fork_turns[0].usage.as_ref().unwrap()), 2020 + 1290 + 810); + + let mut parent_turns = vec![assistant("2026-10-26T10:00:00Z", "2026-10-26T10:00:10Z")]; + let parent_total = + enrich_turns(dir.path(), &[parent.to_string()], &mut parent_turns).unwrap(); + assert_eq!(summed(&parent_total), 2020 + 1290 + 810); + } + + #[test] + fn a_spawned_thread_of_another_session_is_never_counted() { + let dir = TempDir::new().unwrap(); + let ours = "01a17899-0000-7000-8000-00000000000a"; + let stranger = "01a17899-0000-7000-8000-00000000000b"; + let theirs = "01a17899-0000-7000-8000-00000000000c"; + write_named_rollout( + dir.path(), + &format!("rollout-2026-10-26T10-00-00-{ours}.jsonl"), + &[ + usage_row("2026-10-26T10:00:01Z", ours, ours, "turn-1", "resp_1", 100, 0, 10), + spawn_activity("2026-10-26T10:00:01Z", ours, stranger), + ], + ); + // A thread running under another root session. + write_named_rollout( + dir.path(), + &format!("rollout-2026-10-26T10-00-01-{stranger}.jsonl"), + &[usage_row("2026-10-26T10:00:02Z", stranger, theirs, "turn-x", "resp_2", 900, 0, 90)], + ); + + let mut turns = vec![assistant("2026-10-26T10:00:00Z", "2026-10-26T10:00:10Z")]; + let total = enrich_turns(dir.path(), &[ours.to_string()], &mut turns).unwrap(); + assert_eq!(summed(&total), 110); + } + + #[test] + fn only_usage_rows_and_spawn_events_are_parsed() { + // Everything else is skipped before it is parsed; a line that merely + // mentions a marker inside some other record is parsed and dropped. + let dir = TempDir::new().unwrap(); + let session = "01a17899-0000-7000-8000-0000000000aa"; + let noise = serde_json::json!({ + "timestamp": "2026-10-26T10:00:01Z", + "type": "history_mutation", + "payload": {"items": [{"type": "message", "content": "token_usage_record"}]} + }); + write_named_rollout( + dir.path(), + &format!("rollout-2026-10-26T10-00-00-{session}.jsonl"), + &[ + noise, + usage_row("2026-10-26T10:00:02Z", session, session, "turn-1", "resp_1", 40, 0, 2), + ], + ); + let mut turns = vec![assistant("2026-10-26T10:00:00Z", "2026-10-26T10:00:10Z")]; + let total = enrich_turns(dir.path(), &[session.to_string()], &mut turns).unwrap(); + assert_eq!(summed(&total), 42); + } } From 40c4caedc70838ad14a2fa4fe851d45d64a33bfc Mon Sep 17 00:00:00 2001 From: xintaofei Date: Mon, 26 Oct 2026 23:17:16 +0800 Subject: [PATCH 3/3] fix(acp): keep duration-only session stats when there is no native usage The TRAE native-usage fallback returned no stats at all when it had no native usage to add, so every ACP-native conversation whose turns report a duration but no tokens (custom agents, ZCode, Command Code, and TRAE without a rollout) lost its session duration; session details fell back to wall-clock time and exports dropped the row. Leave the stats as they are in that case. --- src-tauri/src/parsers/acp_native.rs | 70 ++++++++++++++++++++++++++++- 1 file changed, 69 insertions(+), 1 deletion(-) diff --git a/src-tauri/src/parsers/acp_native.rs b/src-tauri/src/parsers/acp_native.rs index a5c5180e5c6..cf8cc6b342b 100644 --- a/src-tauri/src/parsers/acp_native.rs +++ b/src-tauri/src/parsers/acp_native.rs @@ -1101,11 +1101,16 @@ fn session_stats(turns: &[MessageTurn]) -> Option { /// Preserve native usage from a TRAE turn that produced no renderable /// assistant message. Normally `trae_usage::enrich_turns` attaches every /// native turn to an ACP turn and the regular aggregation above wins. +/// +/// Without native usage the stats stand as they are: most ACP agents report +/// no tokens, and their footer still has the turns' durations to show. fn session_stats_with_native_usage( stats: Option, native_usage: Option, ) -> Option { - let usage = native_usage?; + let Some(usage) = native_usage else { + return stats; + }; let total_tokens = usage .input_tokens .saturating_add(usage.output_tokens) @@ -2552,6 +2557,69 @@ mod tests { assert_eq!(stats.context_window_max_tokens, Some(760_000)); } + /// A turn end that reports a duration but no tokens — what most ACP agents + /// send — keeps its duration in the session footer, whether or not the + /// agent has a native store to read usage from. + #[tokio::test] + async fn a_duration_only_turn_keeps_its_session_stats() { + async fn record_duration_only_turn(root: &Path, agent: AgentType, session_id: &str) { + let agent_dir = crate::acp::registry::registry_id_for(agent); + let header = + TranscriptHeader::new(&agent.as_wire(), session_id, "/workspace", 1_791_914_922_197); + acp_transcript::record_header_in(root, agent_dir, &header) + .await + .unwrap(); + for (kind, payload) in [ + ( + EntryKind::Prompt, + serde_json::json!([{"type":"text","text":"hello"}]), + ), + ( + EntryKind::Update, + serde_json::json!({ + "sessionUpdate":"agent_message_chunk", + "content":{"type":"text","text":"done"} + }), + ), + ( + EntryKind::TurnEnd, + serde_json::json!({"stopReason":"end_turn","durationMs":3_000}), + ), + ] { + acp_transcript::record_entry_in(root, agent_dir, session_id, kind, payload) + .await + .unwrap(); + } + } + + let session_id = "01a13673-cfd5-7c23-bc00-ce5be75fd42c"; + let plain = AgentType::custom("plain-agent").unwrap(); + let traecli = AgentType::custom("traecli").unwrap(); + let transcripts = tempfile::tempdir().unwrap(); + // TRAE's native store, holding no rollout for this session. + let native = tempfile::tempdir().unwrap(); + record_duration_only_turn(transcripts.path(), plain, session_id).await; + record_duration_only_turn(transcripts.path(), traecli, session_id).await; + + for parser in [ + AcpNativeParser::new_in(plain, transcripts.path().to_path_buf()), + AcpNativeParser::new_in_with_trae_sessions( + traecli, + transcripts.path().to_path_buf(), + native.path().to_path_buf(), + ), + ] { + let stats = parser + .get_conversation(session_id) + .unwrap() + .session_stats + .expect("a duration-only turn still has session stats"); + assert_eq!(stats.total_duration_ms, 3_000); + assert!(stats.total_usage.is_none()); + assert_eq!(stats.total_tokens, None); + } + } + #[test] fn a_session_opened_but_never_used_is_not_listed() { // Opening a conversation writes the ACP header before the user has