feat: 为 TRAE CLI 会话添加原生 token 使用量的导入和统计

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 `<TRAE_HOME>/traecli.toml`, which the CLI loads itself
             // and codeg's MCP settings manage, so codeg does not forward the

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
 //! `<TRAE_HOME>/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 flat `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`:
@@ -18,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.<name>]`
 //!   tables, and `traecli acp serve` loads them itself. Servers sent over ACP
 //!   `session/new` are loaded too (both reach the model as `mcp__<server>__*`),
@@ -36,13 +43,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,13 +84,59 @@ 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"))
+}
+
+/// `<TRAECLI_HOME>/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<OsString>, home: Option<PathBuf>) -> PathBuf {
     if let Some(dir) = trae_home.filter(|value| !value.is_empty()) {
         return PathBuf::from(dir);
     }
     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")
+}
+
 /// `<TRAE_HOME>/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 +186,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<String, String> {

src-tauri/src/commands/token_usage.rs

@@ -135,6 +135,15 @@ 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, 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
 /// counters, the duration and the timestamp. A conversation's context WINDOW is
@@ -149,7 +158,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";
 

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.