diff --git a/.github/assets/client-cursor.jpg b/.github/assets/client-cursor.jpg deleted file mode 100644 index 447a9edc8..000000000 Binary files a/.github/assets/client-cursor.jpg and /dev/null differ diff --git a/.github/assets/client-trae.png b/.github/assets/client-trae.png deleted file mode 100644 index e46c38271..000000000 Binary files a/.github/assets/client-trae.png and /dev/null differ diff --git a/AGENTS.md b/AGENTS.md index dcf5fe52a..f8e16bebd 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -162,5 +162,5 @@ For ordinary validation, prefer source-build checks: ```bash bun install bun run build:core -bun run cli -- --no-spinner --light +bun run cli -- models --no-spinner ``` diff --git a/Cargo.lock b/Cargo.lock index 92776038c..3b0c2da8a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -24,17 +24,6 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" -[[package]] -name = "aes" -version = "0.8.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b169f7a6d4742236a0a00c541b845991d0ac43e546831af1249753ab4c3aa3a0" -dependencies = [ - "cfg-if", - "cipher", - "cpufeatures", -] - [[package]] name = "ahash" version = "0.8.12" @@ -327,15 +316,6 @@ dependencies = [ "generic-array", ] -[[package]] -name = "block-padding" -version = "0.3.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a8894febbff9f758034a5b8e12d87918f56dfc64a8e1fe757d65e29041538d93" -dependencies = [ - "generic-array", -] - [[package]] name = "bstr" version = "1.12.1" @@ -398,15 +378,6 @@ dependencies = [ "rustversion", ] -[[package]] -name = "cbc" -version = "0.1.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "26b52a9543ae338f279b96b0b9fed9c8093744685043739079ce85cd58f289a6" -dependencies = [ - "cipher", -] - [[package]] name = "cc" version = "1.2.55" @@ -472,16 +443,6 @@ dependencies = [ "half", ] -[[package]] -name = "cipher" -version = "0.4.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "773f3b9af64447d2ce9850330c473515014aa235e6a783b02db81ff39e4a3dad" -dependencies = [ - "crypto-common", - "inout", -] - [[package]] name = "clap" version = "4.5.57" @@ -1602,16 +1563,6 @@ dependencies = [ "rustversion", ] -[[package]] -name = "inout" -version = "0.1.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "879f10e63c20629ecabbb64a8010319738c66a5cd0c29b02d63d272b03751d01" -dependencies = [ - "block-padding", - "generic-array", -] - [[package]] name = "instability" version = "0.3.11" @@ -3377,12 +3328,9 @@ name = "tokscale-cli" version = "4.5.0" dependencies = [ "ab_glyph", - "aes", "anyhow", "arboard", "assert_cmd", - "base64", - "cbc", "chrono", "clap", "colored", @@ -3408,6 +3356,7 @@ dependencies = [ "sha2", "signal-hook", "tempfile", + "thiserror 2.0.18", "tokio", "tokscale-core", "toml", diff --git a/Cargo.toml b/Cargo.toml index d00f56d52..c61b0a37e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -85,11 +85,6 @@ codspeed-criterion-compat = "4.7.0" tempfile = "3" serial_test = "3" -# Trae iCubeAuthInfo decryption: AES-128-CBC + base64 -aes = "0.8" -cbc = { version = "0.1", features = ["std"] } -base64 = "0.22" - # Clipboard arboard = "3" diff --git a/README.md b/README.md index db8464444..8b7585e6d 100644 --- a/README.md +++ b/README.md @@ -70,7 +70,7 @@ Run the local wrapper: bun run cli # Script-friendly report -bun run cli -- --no-spinner --light +bun run cli -- models --no-spinner # Inspect detected clients and scan locations bun run cli -- clients @@ -90,23 +90,23 @@ npm install -g @juya-ai/tokscale # TUI tokscale tokscale tui -tokscale models -tokscale monthly -tokscale hourly +tokscale tui --tab models -# Scriptable reports -tokscale --no-spinner --light +# Deterministic table and JSON reports +tokscale models --no-spinner +tokscale monthly --no-spinner +tokscale hourly --no-spinner tokscale models --no-spinner --json tokscale graph --no-spinner --output graph.json # Filters -tokscale --client opencode,claude --week +tokscale tui --client opencode,claude --week tokscale models --since 2026-01-01 --until 2026-01-31 tokscale models --group-by client,provider,model --json # Pricing catalog lookup -tokscale pricing claude-sonnet-4-5 --no-spinner -tokscale pricing list-overrides --json +tokscale pricing lookup claude-sonnet-4-5 --no-spinner +tokscale pricing overrides --json ``` When running from source, replace `tokscale` with `bun run cli --`. @@ -119,10 +119,10 @@ The canonical client identity list lives in Current catalog entries include: -OpenCode, Claude Code, Codex CLI, Cursor, Gemini CLI, Amp, Droid, OpenClaw, +OpenCode, Claude Code, Codex CLI, Gemini CLI, Amp, Droid, OpenClaw, Pi, OMP, Kimi, Qwen CLI, Roo Code, KiloCode, Mux, Kilo CLI, Hermes Agent, Copilot, Goose, Codebuff, CodeBuddy, Antigravity, Zed Agent, -ZCode, Kiro, Junie, Trae, Warp, Cline, Command Code, and Grok Build. +ZCode, Kiro, Junie, Warp, Cline, Command Code, and Grok Build. Some catalog entries have explicit boundaries: @@ -130,11 +130,7 @@ Some catalog entries have explicit boundaries: Tokscale applies the fixed total-only bucket allocation from ADR 0017. - `commandcode` is transcript-estimated usage, not authoritative vendor token accounting. -- `cursor` reads a local API cache. Logged-in local reports and the TUI can - refresh a stale cache automatically when no `--home` override is used, Cursor - is in scope, and the cache is older than five minutes; `tokscale cursor sync` - forces a refresh. -- `antigravity` and `trae` use local caches refreshed by explicit sync commands. +- `antigravity` uses a local cache refreshed by an explicit sync command. ## Data and pricing semantics diff --git a/README.zh-cn.md b/README.zh-cn.md index f6bab7439..45a47580d 100644 --- a/README.zh-cn.md +++ b/README.zh-cn.md @@ -57,7 +57,7 @@ bun run build:core bun run cli # 适合脚本的报表 -bun run cli -- --no-spinner --light +bun run cli -- models --no-spinner # 查看检测到的客户端和扫描位置 bun run cli -- clients @@ -71,23 +71,23 @@ bun run cli -- clients # TUI tokscale tokscale tui -tokscale models -tokscale monthly -tokscale hourly +tokscale tui --tab models -# 脚本化报表 -tokscale --no-spinner --light +# 报表 +tokscale models --no-spinner tokscale models --no-spinner --json +tokscale monthly --no-spinner +tokscale hourly --no-spinner tokscale graph --no-spinner --output graph.json # 过滤 -tokscale --client opencode,claude --week +tokscale tui --client opencode,claude --week tokscale models --since 2026-01-01 --until 2026-01-31 tokscale models --group-by client,provider,model --json # 查询价格目录 -tokscale pricing claude-sonnet-4-5 --no-spinner -tokscale pricing list-overrides --json +tokscale pricing lookup claude-sonnet-4-5 --no-spinner +tokscale pricing overrides --json ``` 从源码运行时，把 `tokscale` 替换成 `bun run cli --`。 @@ -99,14 +99,13 @@ tokscale pricing list-overrides --json 当前 catalog 包括： -OpenCode、Claude Code、Codex CLI、Cursor、Gemini CLI、Amp、Droid、OpenClaw、Pi、OMP、Kimi、Qwen CLI、Roo Code、KiloCode、Mux、Kilo CLI、Hermes Agent、Copilot、Goose、Codebuff、Antigravity、Zed Agent、ZCode、Kiro、Junie、Trae、Warp、Cline、Command Code 和 Grok Build。 +OpenCode、Claude Code、Codex CLI、Gemini CLI、Amp、Droid、OpenClaw、Pi、OMP、Kimi、Qwen CLI、Roo Code、KiloCode、Mux、Kilo CLI、Hermes Agent、Copilot、Goose、Codebuff、Antigravity、Zed Agent、ZCode、Kiro、Junie、Warp、Cline、Command Code 和 Grok Build。 部分 catalog 条目有明确边界： - `grok` 和本地 `warp.sqlite` 只提供没有 bucket 拆分的 token 总数，因此 Tokscale 使用 ADR 0017 定义的固定 bucket 分配。 - `commandcode` 是基于 transcript 的估算用量，不是供应商权威 token 记账。 -- `cursor` 读取本地 API 缓存。已登录时，如果没有使用 `--home`，且 Cursor 在客户端范围内，并且缓存超过五分钟，普通本地报表和 TUI 可以自动刷新过期缓存；`tokscale cursor sync` 用于强制刷新。 -- `antigravity` 和 `trae` 使用显式 sync 命令刷新的本地缓存。 +- `antigravity` 使用显式 sync 命令刷新的本地缓存。 ## 数据和定价语义 diff --git a/crates/tokscale-cli/Cargo.toml b/crates/tokscale-cli/Cargo.toml index 21d49b46d..601669b11 100644 --- a/crates/tokscale-cli/Cargo.toml +++ b/crates/tokscale-cli/Cargo.toml @@ -19,6 +19,7 @@ ratatui = { workspace = true } crossterm = { workspace = true } tokio = { workspace = true } anyhow = { workspace = true } +thiserror = { workspace = true } comfy-table = { workspace = true } colored = { workspace = true } indicatif = { workspace = true } @@ -43,11 +44,6 @@ sha2 = "0.10" csv = "1.3" unicode-width = "0.2" -# Trae iCubeAuthInfo decryption (Electron globalStorage) -aes = { workspace = true } -cbc = { workspace = true } -base64 = { workspace = true } - [target.'cfg(unix)'.dependencies] signal-hook = { workspace = true } libc = { workspace = true } diff --git a/crates/tokscale-cli/src/cli.rs b/crates/tokscale-cli/src/cli.rs new file mode 100644 index 000000000..13cfc875b --- /dev/null +++ b/crates/tokscale-cli/src/cli.rs @@ -0,0 +1,1070 @@ +use std::ffi::OsString; +use std::io::IsTerminal; +use std::path::PathBuf; +use std::time::Duration; + +use anyhow::Result; +use chrono::NaiveDate; +use clap::{error::ErrorKind, Arg, Args, Command, CommandFactory, Parser, Subcommand, ValueEnum}; +use tokscale_core::{ClientId, GroupBy}; + +use crate::commands::shared::{ + build_client_filter, build_date_filter, normalize_year_filter, parse_client_id_arg, +}; +use crate::failure::CliFailure; +use crate::tui::{self, Tab}; + +#[derive(Parser, Debug)] +#[command(name = "tokscale")] +#[command(author, version, about = "AI token usage analytics")] +pub(crate) struct Cli { + #[command(subcommand)] + pub(crate) command: Option, +} + +impl Cli { + /// Parse the process arguments without accepting compatibility aliases. + /// Known v4 spellings still get one actionable migration hint after Clap + /// rejects them, so a breaking change does not turn into a guessing game. + pub(crate) fn parse_from_env() -> Self { + let args = std::env::args_os().collect::>(); + match Self::try_parse_from(args.clone()) { + Ok(cli) => cli, + Err(error) => { + let exit_code = error.exit_code(); + let show_hint = !matches!( + error.kind(), + ErrorKind::DisplayHelp | ErrorKind::DisplayVersion + ); + if let Err(print_error) = error.print() { + eprintln!("error: failed to print CLI error: {print_error}"); + } + if show_hint { + let arguments = args + .into_iter() + .skip(1) + .map(|argument| argument.to_string_lossy().into_owned()) + .collect::>(); + if let Some(hint) = legacy_invocation_hint(&arguments) { + eprintln!("\nhint: {hint}"); + } + } + std::process::exit(exit_code); + } + } + } +} + +pub(crate) fn legacy_invocation_hint(arguments: &[String]) -> Option { + let first = arguments.first()?.as_str(); + + if first == "pricing" { + match arguments.get(1).map(String::as_str) { + Some("list-overrides") => { + let mut replacement = vec!["pricing".to_string(), "overrides".to_string()]; + replacement.extend(arguments.iter().skip(2).cloned()); + return valid_replacement_hint(replacement); + } + Some("lookup") if contains_long_option(arguments, "provider") => { + return Some("replace `--provider` with `--source`".to_string()); + } + Some(value) if !value.starts_with('-') && value != "lookup" && value != "overrides" => { + let mut replacement = vec!["pricing".to_string(), "lookup".to_string()]; + replacement.extend(arguments.iter().skip(1).cloned()); + return valid_replacement_hint(replacement); + } + _ => {} + } + } + + if first == "headless" && !arguments.iter().any(|argument| argument == "--") { + return Some( + "separate Tokscale options from the child command with `--`, for example `tokscale headless codex --format jsonl -- codex exec ...`" + .to_string(), + ); + } + + if contains_long_option(arguments, "write-cache") + || contains_long_option(arguments, "no-write-cache") + { + return Some( + "use `tokscale cache warm` to build the TUI aggregate cache explicitly".to_string(), + ); + } + + let command = Cli::command(); + let command_index = top_level_command_index(&command, arguments); + if let Some(command_index @ 1..) = command_index { + let mut replacement = vec![arguments[command_index].clone()]; + replacement.extend( + arguments + .iter() + .enumerate() + .filter(|(index, argument)| { + *index != command_index && argument.as_str() != "--light" + }) + .map(|(_, argument)| argument.clone()), + ); + if let Some(hint) = valid_replacement_hint(replacement.clone()) { + return Some(hint); + } + return explicit_command_migration_hint(&command, replacement); + } + + if command_index == Some(0) && arguments.iter().any(|argument| argument == "--light") { + let replacement = arguments + .iter() + .filter(|argument| argument.as_str() != "--light") + .cloned() + .collect::>(); + if let Some(hint) = valid_replacement_hint(replacement) { + return Some(hint); + } + } + + if command_index == Some(0) { + return explicit_command_migration_hint(&command, arguments.to_vec()); + } + + let target = migration_target(&command, arguments)?; + let migrated = arguments + .iter() + .filter(|argument| argument.as_str() != "--light") + .cloned() + .collect::>(); + valid_replacement_hint( + std::iter::once(target.to_string()) + .chain(migrated) + .collect(), + ) +} + +fn explicit_command_migration_hint(command: &Command, arguments: Vec) -> Option { + let current = arguments.first()?.as_str(); + + // Graph already emits JSON, so its removed --json flag is redundant. A + // migration hint must preserve the graph product rather than redirecting + // the user to a syntactically valid but unrelated report command. + if current == "graph" && arguments.iter().any(|argument| argument == "--json") { + let replacement = arguments + .iter() + .filter(|argument| argument.as_str() != "--json") + .cloned() + .collect::>(); + return valid_replacement_hint(replacement); + } + + // Cross-command migration is an explicit product decision, not something + // Clap ownership can prove. The only intentional explicit-command move is + // from the interactive TUI to the canonical Models report. + if current != "tui" || migration_target(command, &arguments)? != "models" { + return None; + } + let replacement = std::iter::once("models".to_string()) + .chain(arguments.into_iter().skip(1)) + .collect(); + valid_replacement_hint(replacement) +} + +#[derive(Clone, Copy)] +enum OptionName<'a> { + Long(&'a str), + Short(char), +} + +impl OptionName<'_> { + fn matches(self, argument: &Arg) -> bool { + match self { + Self::Long(name) => argument.get_long() == Some(name), + Self::Short(name) => argument.get_short() == Some(name), + } + } + + fn is_long(self, expected: &str) -> bool { + matches!(self, Self::Long(name) if name == expected) + } +} + +#[derive(Clone, Copy)] +struct OptionToken<'a> { + name: OptionName<'a>, + has_inline_value: bool, +} + +fn option_token(argument: &str) -> Option> { + if argument == "--" { + return None; + } + if let Some(body) = argument.strip_prefix("--") { + let (name, has_inline_value) = match body.split_once('=') { + Some((name, _)) => (name, true), + None => (body, false), + }; + return (!name.is_empty()).then_some(OptionToken { + name: OptionName::Long(name), + has_inline_value, + }); + } + let body = argument.strip_prefix('-')?; + let mut characters = body.chars(); + let name = characters.next()?; + Some(OptionToken { + name: OptionName::Short(name), + has_inline_value: characters.next().is_some(), + }) +} + +fn contains_long_option(arguments: &[String], expected: &str) -> bool { + arguments + .iter() + .take_while(|argument| argument.as_str() != "--") + .filter_map(|argument| option_token(argument)) + .any(|option| option.name.is_long(expected)) +} + +fn command_defines_option(command: &Command, option: OptionName<'_>) -> bool { + command + .get_arguments() + .any(|argument| option.matches(argument)) +} + +fn command_tree_option_takes_value(command: &Command, option: OptionName<'_>) -> bool { + command + .get_arguments() + .any(|argument| option.matches(argument) && argument.get_action().takes_values()) + || command + .get_subcommands() + .any(|subcommand| command_tree_option_takes_value(subcommand, option)) +} + +fn top_level_command_index(command: &Command, arguments: &[String]) -> Option { + let mut option_takes_next_value = false; + for (index, argument) in arguments.iter().enumerate() { + if argument == "--" { + break; + } + if option_takes_next_value { + option_takes_next_value = false; + continue; + } + if let Some(option) = option_token(argument) { + option_takes_next_value = + !option.has_inline_value && command_tree_option_takes_value(command, option.name); + continue; + } + if command + .get_subcommands() + .any(|subcommand| subcommand.get_name() == argument) + { + return Some(index); + } + } + None +} + +fn migration_target<'a>(command: &'a Command, arguments: &[String]) -> Option<&'a str> { + // TUI and Models are the canonical v5 destinations for the former root + // options. Reading their actual Clap arguments keeps ownership in one + // place; options shared by both retain the old root command's TUI default. + let tui = command + .find_subcommand("tui") + .expect("Clap command tree must contain tui"); + let models = command + .find_subcommand("models") + .expect("Clap command tree must contain models"); + let mut models_only = false; + let mut tui_only = false; + let mut shared = false; + let mut option_takes_next_value = false; + + for argument in arguments { + if argument == "--" { + break; + } + if option_takes_next_value { + option_takes_next_value = false; + continue; + } + let Some(option) = option_token(argument) else { + continue; + }; + option_takes_next_value = + !option.has_inline_value && command_tree_option_takes_value(command, option.name); + + let belongs_to_tui = command_defines_option(tui, option.name); + let belongs_to_models = + command_defines_option(models, option.name) || option.name.is_long("light"); + match (belongs_to_tui, belongs_to_models) { + (true, true) => shared = true, + (true, false) => tui_only = true, + (false, true) => models_only = true, + (false, false) => {} + } + } + + match (models_only, tui_only, shared) { + (true, false, _) => Some(models.get_name()), + (false, true, _) | (false, false, true) => Some(tui.get_name()), + _ => None, + } +} + +fn valid_replacement_hint(replacement: Vec) -> Option { + let mut argv = vec!["tokscale".to_string()]; + argv.extend(replacement.iter().cloned()); + Cli::try_parse_from(argv) + .is_ok() + .then(|| format!("use `tokscale {}`", replacement.join(" "))) +} + +#[derive(Subcommand, Debug)] +pub(crate) enum Commands { + #[command(about = "Launch the interactive terminal interface")] + Tui(TuiArgs), + #[command(about = "Show model usage report")] + Models(ModelsArgs), + #[command(about = "Show monthly usage report")] + Monthly(ReportArgs), + #[command(about = "Show hourly usage report")] + Hourly(ReportArgs), + #[command(about = "Show session time metrics")] + TimeMetrics(ReportArgs), + #[command(about = "Show local scan locations and session counts")] + Clients(ClientsArgs), + #[command(about = "Export contribution graph data as JSON")] + Graph(GraphArgs), + #[command(about = "Query model pricing")] + Pricing { + #[command(subcommand)] + subcommand: PricingSubcommand, + }, + #[command(about = "Show subscription usage and quota for AI providers")] + Usage { + #[arg(long, help = "Output as JSON")] + json: bool, + }, + #[command(about = "Generate year-in-review wrapped image")] + Wrapped(WrappedArgs), + #[command(about = "Capture subprocess output for token usage tracking")] + Headless(HeadlessArgs), + #[command(about = "Maintain local Tokscale caches")] + Cache { + #[command(subcommand)] + subcommand: CacheSubcommand, + }, + #[command(about = "Antigravity integration commands")] + Antigravity { + #[command(subcommand)] + subcommand: AntigravitySubcommand, + }, + #[command(about = "Warp/Oz aggregate usage integration commands")] + Warp { + #[command(subcommand)] + subcommand: WarpSubcommand, + }, +} + +#[derive(Args, Debug, Default)] +pub(crate) struct TuiArgs { + #[arg(long, value_enum, help = "Open a specific tab")] + pub(crate) tab: Option, + #[arg(short, long, value_parser = parse_theme_arg)] + pub(crate) theme: Option, + #[arg( + short, + long, + value_name = "SECONDS", + value_parser = parse_positive_u64, + conflicts_with = "no_refresh" + )] + pub(crate) refresh: Option, + #[arg(long, conflicts_with = "refresh", help = "Disable automatic refresh")] + pub(crate) no_refresh: bool, + #[arg(long)] + pub(crate) debug: bool, + #[command(flatten)] + pub(crate) source: SourceScopeArgs, + #[command(flatten)] + pub(crate) date: DateRangeFlags, +} + +#[derive(Args, Debug)] +pub(crate) struct ModelsArgs { + #[command(flatten)] + pub(crate) report: ReportArgs, + #[arg( + long, + value_name = "STRATEGY", + default_value = "client,model", + help = "Grouping strategy: model, client,model, client,provider,model, workspace,model, session,model, client,session,model" + )] + pub(crate) group_by: GroupBy, +} + +#[derive(Args, Debug)] +pub(crate) struct ReportArgs { + #[arg(long, help = "Output as JSON")] + pub(crate) json: bool, + #[command(flatten)] + pub(crate) source: SourceScopeArgs, + #[command(flatten)] + pub(crate) date: DateRangeFlags, + #[arg(long, help = "Write processing time to stderr")] + pub(crate) benchmark: bool, + #[arg(long, help = "Disable progress animation")] + pub(crate) no_spinner: bool, +} + +#[derive(Args, Debug)] +pub(crate) struct ClientsArgs { + #[arg(long, help = "Output as JSON")] + pub(crate) json: bool, + #[command(flatten)] + pub(crate) source: SourceScopeArgs, +} + +#[derive(Args, Debug)] +pub(crate) struct GraphArgs { + #[arg(long, value_name = "PATH", help = "Write JSON to a file")] + pub(crate) output: Option, + #[command(flatten)] + pub(crate) source: SourceScopeArgs, + #[command(flatten)] + pub(crate) date: DateRangeFlags, + #[arg(long, help = "Write processing time to stderr")] + pub(crate) benchmark: bool, + #[arg(long, help = "Disable progress animation")] + pub(crate) no_spinner: bool, +} + +#[derive(Args, Debug)] +pub(crate) struct WrappedArgs { + #[arg(long, value_name = "PATH", help = "Output file path")] + pub(crate) output: Option, + #[arg(long, value_parser = parse_year_arg, help = "Year to generate")] + pub(crate) year: Option, + #[command(flatten)] + pub(crate) source: SourceScopeArgs, + #[arg(long, help = "Display total tokens in abbreviated format")] + pub(crate) short: bool, + #[arg( + long, + value_enum, + help = "Choose the ranking panel instead of automatic selection" + )] + pub(crate) ranking: Option, + #[arg(long, help = "Disable pinning of Sisyphus agents in rankings")] + pub(crate) disable_pinned: bool, + #[arg(long, help = "Disable progress animation")] + pub(crate) no_spinner: bool, +} + +#[derive(ValueEnum, Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum WrappedRankingArg { + Agents, + Clients, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum WrappedRanking { + Auto, + Agents, + Clients, +} + +impl From for WrappedRanking { + fn from(value: WrappedRankingArg) -> Self { + match value { + WrappedRankingArg::Agents => Self::Agents, + WrappedRankingArg::Clients => Self::Clients, + } + } +} + +#[derive(Args, Debug)] +pub(crate) struct HeadlessArgs { + #[arg(value_enum, help = "Usage adapter for the captured process")] + pub(crate) source: HeadlessSource, + #[arg(long, value_enum, help = "Captured output format")] + pub(crate) format: Option, + #[arg(long, value_name = "PATH", help = "Write captured output to this file")] + pub(crate) output: Option, + #[arg(long, help = "Do not add source-specific structured-output flags")] + pub(crate) no_auto_flags: bool, + #[arg( + last = true, + required = true, + num_args = 1.., + value_name = "COMMAND", + help = "Child command and arguments after `--`" + )] + pub(crate) command: Vec, +} + +#[derive(Args, Clone, Debug, Default)] +pub(crate) struct SourceScopeArgs { + #[arg( + long, + value_name = "PATH", + value_parser = parse_home_arg, + help = "Read local session data from this existing home directory" + )] + pub(crate) home: Option, + #[command(flatten)] + pub(crate) clients: ClientFlags, +} + +#[derive(Args, Clone, Debug, Default)] +pub(crate) struct ClientFlags { + /// Canonical client filter. Repeatable or comma-separated. + #[arg( + long = "client", + short = 'c', + value_parser = parse_client_id_arg, + value_delimiter = ',', + action = clap::ArgAction::Append, + help = "Filter by client. Repeatable or comma-separated" + )] + pub(crate) clients: Vec, +} + +#[derive(Args, Clone, Debug, Default)] +pub(crate) struct DateRangeFlags { + #[arg( + long, + conflicts_with_all = ["week", "month", "year", "since", "until"], + help = "Show only today's usage" + )] + pub(crate) today: bool, + #[arg( + long, + conflicts_with_all = ["today", "month", "year", "since", "until"], + help = "Show the last seven days" + )] + pub(crate) week: bool, + #[arg( + long, + conflicts_with_all = ["today", "week", "year", "since", "until"], + help = "Show the current month" + )] + pub(crate) month: bool, + #[arg( + long, + value_parser = parse_date_arg, + conflicts_with_all = ["today", "week", "month", "year"], + help = "Inclusive start date (YYYY-MM-DD)" + )] + pub(crate) since: Option, + #[arg( + long, + value_parser = parse_date_arg, + conflicts_with_all = ["today", "week", "month", "year"], + help = "Inclusive end date (YYYY-MM-DD)" + )] + pub(crate) until: Option, + #[arg( + long, + value_parser = parse_year_arg, + conflicts_with_all = ["today", "week", "month", "since", "until"], + help = "Filter by year (YYYY)" + )] + pub(crate) year: Option, +} + +#[derive(Subcommand, Debug)] +pub(crate) enum PricingSubcommand { + #[command(about = "Look up pricing for a model")] + Lookup { + #[arg(help = "Model ID to look up")] + model_id: String, + #[arg(long, help = "Output as JSON")] + json: bool, + #[arg(long, value_enum, help = "Use one pricing data source")] + source: Option, + #[arg(long, help = "Disable progress animation")] + no_spinner: bool, + }, + #[command(about = "List custom pricing overrides")] + Overrides { + #[arg(long, help = "Output as JSON")] + json: bool, + }, +} + +#[derive(Subcommand, Debug)] +pub(crate) enum CacheSubcommand { + #[command(about = "Build the TUI aggregate cache for a source scope")] + Warm { + #[command(flatten)] + source: SourceScopeArgs, + }, + #[command(about = "Remove orphaned and superseded source-message cache shards")] + Prune, +} + +#[derive(Subcommand, Debug)] +pub(crate) enum AntigravitySubcommand { + #[command(about = "Sync usage from running Antigravity language servers")] + Sync, + #[command(about = "Show Antigravity sync status")] + Status { + #[arg(long, help = "Output as JSON")] + json: bool, + }, + #[command(about = "Delete cached Antigravity usage artifacts")] + PurgeCache, +} + +#[derive(Subcommand, Debug)] +pub(crate) enum WarpSubcommand { + #[command(about = "Save Warp GraphQL authentication")] + Login { + #[arg(long, help = "Warp bearer token or cookie header value")] + token: Option, + #[arg(long, help = "Treat token as a Cookie header")] + cookie: bool, + }, + #[command(about = "Remove cached Warp credentials")] + Logout { + #[arg(long, help = "Also delete cached Warp aggregate usage")] + purge_cache: bool, + }, + #[command(about = "Show Warp aggregate sync status")] + Status { + #[arg(long, help = "Output as JSON")] + json: bool, + }, + #[command(about = "Sync Warp aggregate usage into local cache")] + Sync { + #[arg(long, help = "Output as JSON")] + json: bool, + }, +} + +#[derive(ValueEnum, Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum TuiTab { + Overview, + Models, + Monthly, + Weekly, + Daily, + Hourly, + Stats, + Agents, + Issues, + Usage, +} + +impl From for Tab { + fn from(value: TuiTab) -> Self { + match value { + TuiTab::Overview => Tab::Overview, + TuiTab::Models => Tab::Models, + TuiTab::Monthly => Tab::Monthly, + TuiTab::Weekly => Tab::Weekly, + TuiTab::Daily => Tab::Daily, + TuiTab::Hourly => Tab::Hourly, + TuiTab::Stats => Tab::Stats, + TuiTab::Agents => Tab::Agents, + TuiTab::Issues => Tab::Issues, + TuiTab::Usage => Tab::Usage, + } + } +} + +#[derive(ValueEnum, Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum PricingSource { + Custom, + Litellm, + Openrouter, + #[value(name = "models.dev")] + ModelsDev, +} + +impl PricingSource { + pub(crate) const fn as_str(self) -> &'static str { + match self { + Self::Custom => "custom", + Self::Litellm => "litellm", + Self::Openrouter => "openrouter", + Self::ModelsDev => "models.dev", + } + } +} + +#[derive(ValueEnum, Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum HeadlessSource { + Codex, +} + +impl HeadlessSource { + pub(crate) const fn as_str(self) -> &'static str { + match self { + Self::Codex => "codex", + } + } +} + +#[derive(ValueEnum, Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum HeadlessFormat { + Json, + Jsonl, +} + +impl HeadlessFormat { + pub(crate) const fn as_str(self) -> &'static str { + match self { + Self::Json => "json", + Self::Jsonl => "jsonl", + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct TerminalState { + pub(crate) stdin: bool, + pub(crate) stdout: bool, +} + +impl TerminalState { + pub(crate) fn detect() -> Self { + Self { + stdin: std::io::stdin().is_terminal(), + stdout: std::io::stdout().is_terminal(), + } + } + + fn interactive(self) -> bool { + self.stdin && self.stdout + } +} + +#[derive(Debug)] +pub(crate) struct ResolvedSourceScope { + pub(crate) home: Option, + pub(crate) clients: Option>, +} + +#[derive(Debug)] +pub(crate) struct ResolvedDateRange { + pub(crate) today: bool, + pub(crate) week: bool, + pub(crate) month: bool, + pub(crate) since: Option, + pub(crate) until: Option, + pub(crate) year: Option, +} + +#[derive(Debug)] +pub(crate) struct LocalReportPlan { + pub(crate) json: bool, + pub(crate) source: ResolvedSourceScope, + pub(crate) date: ResolvedDateRange, + pub(crate) benchmark: bool, + pub(crate) no_spinner: bool, +} + +#[derive(Debug)] +pub(crate) struct ModelsPlan { + pub(crate) report: LocalReportPlan, + pub(crate) group_by: GroupBy, +} + +#[derive(Debug)] +pub(crate) struct TuiPlan { + pub(crate) theme: Option, + pub(crate) refresh: Option, + pub(crate) no_refresh: bool, + pub(crate) debug: bool, + pub(crate) source: ResolvedSourceScope, + pub(crate) date: ResolvedDateRange, + pub(crate) initial_tab: Option, +} + +#[derive(Debug)] +pub(crate) struct ClientsPlan { + pub(crate) json: bool, + pub(crate) source: ResolvedSourceScope, +} + +#[derive(Debug)] +pub(crate) struct GraphPlan { + pub(crate) output: Option, + pub(crate) source: ResolvedSourceScope, + pub(crate) date: ResolvedDateRange, + pub(crate) benchmark: bool, + pub(crate) no_spinner: bool, +} + +#[derive(Debug)] +pub(crate) struct WrappedPlan { + pub(crate) output: Option, + pub(crate) year: Option, + pub(crate) source: ResolvedSourceScope, + pub(crate) short: bool, + pub(crate) ranking: WrappedRanking, + pub(crate) disable_pinned: bool, + pub(crate) no_spinner: bool, +} + +#[derive(Debug)] +pub(crate) struct HeadlessPlan { + pub(crate) source: HeadlessSource, + pub(crate) command: Vec, + pub(crate) format: Option, + pub(crate) output: Option, + pub(crate) no_auto_flags: bool, + pub(crate) timeout: Duration, +} + +#[derive(Debug)] +pub(crate) enum ExecutionPlan { + Tui(TuiPlan), + Models(ModelsPlan), + Monthly(LocalReportPlan), + Hourly(LocalReportPlan), + TimeMetrics(LocalReportPlan), + Clients(ClientsPlan), + Graph(GraphPlan), + Pricing(PricingSubcommand), + Usage { json: bool }, + Wrapped(WrappedPlan), + Headless(HeadlessPlan), + CachePrune, + CacheWarm(ResolvedSourceScope), + Antigravity(AntigravitySubcommand), + Warp(WarpSubcommand), +} + +impl ExecutionPlan { + pub(crate) fn resolve(cli: Cli, terminal: TerminalState) -> Result { + match cli.command.unwrap_or(Commands::Tui(TuiArgs::default())) { + Commands::Tui(args) => resolve_tui(args, terminal).map(Self::Tui), + Commands::Models(args) => Ok(Self::Models(ModelsPlan { + report: resolve_report(args.report)?, + group_by: args.group_by, + })), + Commands::Monthly(args) => resolve_report(args).map(Self::Monthly), + Commands::Hourly(args) => resolve_report(args).map(Self::Hourly), + Commands::TimeMetrics(args) => resolve_report(args).map(Self::TimeMetrics), + Commands::Clients(args) => Ok(Self::Clients(ClientsPlan { + json: args.json, + source: resolve_source(args.source)?, + })), + Commands::Graph(args) => Ok(Self::Graph(GraphPlan { + output: args.output, + source: resolve_source(args.source)?, + date: resolve_date(args.date)?, + benchmark: args.benchmark, + no_spinner: args.no_spinner, + })), + Commands::Pricing { subcommand } => Ok(Self::Pricing(subcommand)), + Commands::Usage { json } => Ok(Self::Usage { json }), + Commands::Wrapped(args) => resolve_wrapped(args).map(Self::Wrapped), + Commands::Headless(args) => resolve_headless(args).map(Self::Headless), + Commands::Cache { subcommand } => match subcommand { + CacheSubcommand::Prune => Ok(Self::CachePrune), + CacheSubcommand::Warm { source } => resolve_source(source).map(Self::CacheWarm), + }, + Commands::Antigravity { subcommand } => Ok(Self::Antigravity(subcommand)), + Commands::Warp { subcommand } => Ok(Self::Warp(subcommand)), + } + } +} + +fn resolve_wrapped(args: WrappedArgs) -> Result { + let source = resolve_source(args.source)?; + let ranking = args + .ranking + .map(WrappedRanking::from) + .unwrap_or(WrappedRanking::Auto); + + if ranking == WrappedRanking::Agents + && source.clients.as_ref().is_some_and(|clients| { + !clients + .iter() + .any(|client| client == ClientId::OpenCode.as_str()) + }) + { + return Err(CliFailure::invalid_message( + "--ranking agents requires `opencode` in the --client scope".to_string(), + )); + } + + if ranking == WrappedRanking::Clients && args.disable_pinned { + return Err(CliFailure::invalid_message( + "--disable-pinned does not apply to --ranking clients".to_string(), + )); + } + + Ok(WrappedPlan { + output: args.output, + year: args.year, + source, + short: args.short, + ranking, + disable_pinned: args.disable_pinned, + no_spinner: args.no_spinner, + }) +} + +fn resolve_tui(args: TuiArgs, terminal: TerminalState) -> Result { + if !terminal.interactive() { + return Err(CliFailure::invalid_message( + "TUI requires an interactive terminal\nhint: use `tokscale models --json` for structured output" + .to_string(), + )); + } + + let source = resolve_source(args.source)?; + let initial_tab = args.tab.map(Tab::from); + if initial_tab == Some(Tab::Usage) { + let settings = tui::settings::Settings::load_for_home_override( + source.home.as_deref().map(std::path::Path::new), + )?; + if !settings.usage_tab_enabled { + return Err(CliFailure::invalid_message( + "TUI tab `usage` is disabled in settings.json".to_string(), + )); + } + } + + Ok(TuiPlan { + theme: args.theme, + refresh: args.refresh, + no_refresh: args.no_refresh, + debug: args.debug, + source, + date: resolve_date(args.date)?, + initial_tab, + }) +} + +fn resolve_report(args: ReportArgs) -> Result { + Ok(LocalReportPlan { + json: args.json, + source: resolve_source(args.source)?, + date: resolve_date(args.date)?, + benchmark: args.benchmark, + no_spinner: args.no_spinner, + }) +} + +fn resolve_headless(args: HeadlessArgs) -> Result { + if args + .command + .first() + .is_none_or(|program| program.trim().is_empty()) + { + return Err(CliFailure::invalid_message( + "headless child command must start with a non-empty executable".to_string(), + )); + } + + let settings = tui::settings::Settings::load()?; + let timeout = settings.get_native_timeout()?; + + Ok(HeadlessPlan { + source: args.source, + command: args.command, + format: args.format, + output: args.output, + no_auto_flags: args.no_auto_flags, + timeout, + }) +} + +fn resolve_source(args: SourceScopeArgs) -> Result { + let home = args.home.map(|path| path.to_string_lossy().into_owned()); + let clients = build_client_filter(args.clients, &home)?; + Ok(ResolvedSourceScope { home, clients }) +} + +fn resolve_date(date: DateRangeFlags) -> Result { + if let (Some(since), Some(until)) = (&date.since, &date.until) { + let since_date = NaiveDate::parse_from_str(since, "%Y-%m-%d") + .expect("Clap date parser must validate --since"); + let until_date = NaiveDate::parse_from_str(until, "%Y-%m-%d") + .expect("Clap date parser must validate --until"); + if since_date > until_date { + return Err(CliFailure::invalid_message(format!( + "--since ({since}) must not be later than --until ({until})" + ))); + } + } + + let (since, until) = + build_date_filter(date.today, date.week, date.month, date.since, date.until); + let year = normalize_year_filter(date.today, date.week, date.month, date.year); + Ok(ResolvedDateRange { + today: date.today, + week: date.week, + month: date.month, + since, + until, + year, + }) +} + +fn parse_home_arg(raw: &str) -> Result { + let raw = raw.trim(); + if raw.is_empty() { + return Err("--home must not be empty".to_string()); + } + let path = PathBuf::from(raw); + if !path.is_dir() { + return Err(format!( + "--home must be an existing directory: {}", + path.display() + )); + } + path.canonicalize().map_err(|error| { + format!( + "failed to canonicalize --home `{}`: {error}", + path.display() + ) + }) +} + +fn parse_date_arg(raw: &str) -> Result { + NaiveDate::parse_from_str(raw, "%Y-%m-%d") + .map(|date| date.format("%Y-%m-%d").to_string()) + .map_err(|_| format!("invalid date `{raw}`; expected YYYY-MM-DD")) +} + +fn parse_year_arg(raw: &str) -> Result { + if raw.len() != 4 || !raw.bytes().all(|byte| byte.is_ascii_digit()) { + return Err(format!("invalid year `{raw}`; expected YYYY")); + } + let year = raw + .parse::() + .map_err(|_| format!("invalid year `{raw}`; expected YYYY"))?; + NaiveDate::from_ymd_opt(year, 1, 1) + .ok_or_else(|| format!("invalid year `{raw}`; expected YYYY"))?; + Ok(raw.to_string()) +} + +fn parse_positive_u64(raw: &str) -> Result { + let value = raw + .parse::() + .map_err(|_| format!("invalid refresh interval `{raw}`"))?; + if value == 0 { + return Err("--refresh must be greater than zero".to_string()); + } + Ok(value) +} + +fn parse_theme_arg(raw: &str) -> Result { + raw.parse::() + .map(|_| raw.to_string()) + .map_err(|_| { + format!( + "invalid theme `{raw}`; expected one of: {}", + crate::tui::ThemeName::all() + .iter() + .map(crate::tui::ThemeName::as_str) + .collect::>() + .join(", ") + ) + }) +} diff --git a/crates/tokscale-cli/src/commands/cache.rs b/crates/tokscale-cli/src/commands/cache.rs index 370fb80f4..7b1df7af4 100644 --- a/crates/tokscale-cli/src/commands/cache.rs +++ b/crates/tokscale-cli/src/commands/cache.rs @@ -1,107 +1,7 @@ -use crate::commands::shared::{parse_client_id_set, parse_default_client_filters}; -use crate::tui; +use crate::commands::shared::parse_client_id_set; use anyhow::Result; use tokscale_core::ClientId; -/// Resolve the filter set used by a no-`--client`-flag TUI launch. -/// -/// Mirrors the resolution that `build_client_filter` + `tui::run` perform -/// when the user passes no CLI client flag: -/// -/// 1. If `defaultClients` from `~/.config/tokscale/settings.json` is -/// set, use it after validating every id. -/// 2. Otherwise use every catalog client. -/// -/// This **must** stay in lockstep with the resolution that -/// `tui::run(.., clients = None, ..)` would compute. If it drifts, the -/// local warm cache uses one filter set while the next no-flag TUI launch -/// wants another, the cache key mismatches, and the warming becomes a -/// wasted background scan. -pub(crate) fn resolve_default_tui_filter_set() -> Result> { - let configured = tui::settings::load_default_clients()?; - resolve_default_tui_filter_set_with(&configured) -} - -/// Pure variant of `resolve_default_tui_filter_set` for unit-testable -/// resolution. `configured` is the (raw, pre-validation) list of ids -/// from settings.json. -pub(crate) fn resolve_default_tui_filter_set_with( - configured: &[String], -) -> Result> { - let parsed = parse_default_client_filters(configured)?; - if parsed.is_empty() { - Ok(ClientId::iter().collect()) - } else { - Ok(parsed.into_iter().collect()) - } -} - -pub(crate) fn resolve_should_write_cache( - cli_write: bool, - cli_no_write: bool, - settings: &tui::settings::Settings, -) -> bool { - if cli_no_write { - return false; - } - if cli_write { - return true; - } - settings.light.write_cache -} - -pub(crate) fn resolve_light_cache_filter_set( - clients: &Option>, -) -> std::collections::HashSet { - if let Some(clients) = clients { - parse_client_id_set(clients) - } else { - ClientId::iter().collect() - } -} - -pub(crate) fn write_light_cache( - clients: &Option>, - since: &Option, - until: &Option, - year: &Option, - group_by: &tokscale_core::GroupBy, -) -> Result<()> { - use crate::tui::{save_cached_data, CacheReportScope, DataLoader}; - - let enabled_set = resolve_light_cache_filter_set(clients); - let mut scan_clients: Vec = enabled_set.iter().copied().collect(); - scan_clients.sort_by_key(|client| *client as usize); - - let loader = DataLoader::with_filters(None, since.clone(), until.clone(), year.clone()); - let report_scope = CacheReportScope::new(since.clone(), until.clone(), year.clone()); - let result = loader.load_with_diagnostics(&scan_clients, group_by)?; - save_cached_data( - &result.data, - &enabled_set, - group_by, - &report_scope, - result.source_inventory_signature, - )?; - Ok(()) -} - -pub(crate) fn validate_light_cache_write(home_dir: &Option) -> Result<()> { - // The TUI cache key includes date filters, but not `--home`. Validate this - // before scanning or rendering so a rejected write intent cannot emit a - // successful-looking report first. - if !can_write_light_cache(home_dir) { - anyhow::bail!( - "--write-cache cannot be combined with --home because the TUI cache key does not include that filter" - ); - } - Ok(()) -} - -pub(crate) fn can_write_light_cache(home_dir: &Option) -> bool { - home_dir.is_none() -} - pub(crate) fn run_source_cache_prune() -> Result<()> { let stats = tokscale_core::prune_source_message_cache()?; println!( @@ -111,32 +11,33 @@ pub(crate) fn run_source_cache_prune() -> Result<()> { Ok(()) } -pub(crate) fn run_warm_tui_cache() -> Result<()> { +pub(crate) fn run_warm_tui_cache( + home_dir: Option, + clients: Option>, +) -> Result<()> { use crate::tui::{save_cached_data, CacheReportScope, DataLoader, TUI_DEFAULT_GROUP_BY}; - use tokscale_core::ClientId; - // Warm the cache using the same default filter set the TUI uses on a - // no-flag launch. Going through `resolve_default_tui_filter_set()` keeps - // these two paths in lockstep, including the user's `defaultClients` - // setting. - // - // The `group_by` MUST be `TUI_DEFAULT_GROUP_BY`, NOT - // `GroupBy::default()`. Using `GroupBy::default()` here is the bug - // that motivated this constant — the TUI's cache reader keys on - // `TUI_DEFAULT_GROUP_BY` (= `GroupBy::Model`) while - // `GroupBy::default()` is `GroupBy::ClientModel`, so the warm cache - // was written under a key the TUI never queried. - let enabled_set = resolve_default_tui_filter_set()?; + let enabled_set: std::collections::HashSet = clients + .as_ref() + .map(|clients| parse_client_id_set(clients)) + .unwrap_or_else(|| ClientId::iter().collect()); let mut scan_clients: Vec = enabled_set.iter().copied().collect(); scan_clients.sort_by_key(|client| *client as usize); - let loader = DataLoader::with_filters(None, None, None, None); + let report_scope = CacheReportScope::for_request(home_dir.clone(), None, None, None)?; + let loader = DataLoader::with_filters( + home_dir.clone().map(std::path::PathBuf::from), + None, + None, + None, + ); let result = loader.load_with_diagnostics(&scan_clients, &TUI_DEFAULT_GROUP_BY)?; save_cached_data( &result.data, &enabled_set, &TUI_DEFAULT_GROUP_BY, - &CacheReportScope::default(), + &report_scope, result.source_inventory_signature, )?; + println!("TUI cache warmed."); Ok(()) } diff --git a/crates/tokscale-cli/src/commands/clients.rs b/crates/tokscale-cli/src/commands/clients.rs index cd2be1b3c..229324f3f 100644 --- a/crates/tokscale-cli/src/commands/clients.rs +++ b/crates/tokscale-cli/src/commands/clients.rs @@ -1,11 +1,15 @@ use crate::claude_diagnostics; use crate::commands::render::format_number; -use crate::commands::shared::use_env_roots; +use crate::commands::shared::{use_env_roots, ReportEnvelope}; use crate::tui; use anyhow::Result; use std::path::{Path, PathBuf}; -pub(crate) fn run_clients_command(json: bool, home_dir: Option) -> Result<()> { +pub(crate) fn run_clients_command( + json: bool, + home_dir: Option, + clients: Option>, +) -> Result<()> { use tokscale_core::scanner::{ built_in_extra_scan_paths_for, copilot_exporter_path_with_env_strategy, discover_opencode_dbs, extra_scan_paths_for, opencode_data_dir_with_env_strategy, @@ -16,6 +20,19 @@ pub(crate) fn run_clients_command(json: bool, home_dir: Option) -> Resul LocalParseOptions, }; + let start = std::time::Instant::now(); + let selected_clients: std::collections::HashSet = clients + .as_ref() + .map(|clients| { + clients + .iter() + .map(|client| { + ClientId::from_str(client) + .expect("resolved client scope must contain canonical client ids") + }) + .collect() + }) + .unwrap_or_else(|| ClientId::iter().collect()); let explicit_home_dir = home_dir; let use_env_roots = use_env_roots(&explicit_home_dir); let scanner_settings = tui::settings::load_scanner_settings_for_home(&explicit_home_dir)?; @@ -30,7 +47,7 @@ pub(crate) fn run_clients_command(json: bool, home_dir: Option) -> Resul use_env_roots, clients: Some( ClientId::iter() - .filter(|client| client.parse_local()) + .filter(|client| selected_clients.contains(client)) .map(|client| client.as_str().to_string()) .collect(), ), @@ -39,7 +56,7 @@ pub(crate) fn run_clients_command(json: bool, home_dir: Option) -> Resul year: None, scanner_settings: scanner_settings.clone(), }) - .map_err(|e| anyhow::anyhow!(e))?; + .map_err(anyhow::Error::new)?; let mut health = client_counts.health.clone(); let headless_roots = @@ -99,7 +116,7 @@ pub(crate) fn run_clients_command(json: bool, home_dir: Option) -> Resul source: String, } - let all_clients: std::collections::HashSet = ClientId::iter().collect(); + let all_clients = selected_clients.clone(); let extra_dirs_val = if use_env_roots { match std::env::var("TOKSCALE_EXTRA_DIRS") { Ok(value) => value, @@ -125,16 +142,21 @@ pub(crate) fn run_clients_command(json: bool, home_dir: Option) -> Resul let settings_extra_dirs = extra_scan_paths_for(&scanner_settings, &all_clients)?; let copilot_exporter_path = copilot_exporter_path_with_env_strategy(use_env_roots); let opencode_data_root = opencode_data_dir_with_env_strategy(&home_dir_str, use_env_roots); - let opencode_auto_dbs = match discover_opencode_dbs(&opencode_data_root) { - Ok(paths) => paths, - Err(_) => { - health.record_unavailable_source(ClientId::OpenCode.as_str()); - Vec::new() + let opencode_auto_dbs = if selected_clients.contains(&ClientId::OpenCode) { + match discover_opencode_dbs(&opencode_data_root) { + Ok(paths) => paths, + Err(_) => { + health.record_unavailable_source(ClientId::OpenCode.as_str()); + Vec::new() + } } + } else { + Vec::new() }; let clients: Vec = ClientId::iter() + .filter(|client| selected_clients.contains(client)) .map(|client| { let warp_default_roots = if client == ClientId::Warp { warp_sqlite_roots_with_env_strategy(&home_dir_str, use_env_roots) @@ -288,22 +310,21 @@ pub(crate) fn run_clients_command(json: bool, home_dir: Option) -> Resul if json { #[derive(serde::Serialize)] #[serde(rename_all = "camelCase")] - struct Output<'a> { + struct ClientsData { headless_roots: Vec, clients: Vec, note: String, - health: &'a tokscale_core::source_health::HealthReport, } - let output = Output { + let data = ClientsData { headless_roots: headless_roots .iter() .map(|p| p.to_string_lossy().to_string()) .collect(), clients, note: "Headless capture is supported for Codex CLI only.".to_string(), - health: &health, }; + let output = ReportEnvelope::new(data, health, start.elapsed().as_millis() as u64); println!("{}", serde_json::to_string_pretty(&output)?); } else { diff --git a/crates/tokscale-cli/src/commands/graph.rs b/crates/tokscale-cli/src/commands/graph.rs index e0d8c1ef0..06f6454e6 100644 --- a/crates/tokscale-cli/src/commands/graph.rs +++ b/crates/tokscale-cli/src/commands/graph.rs @@ -1,11 +1,8 @@ use crate::commands::render::format_currency; -use crate::commands::shared::{ - auto_sync_cursor_for_local_report, client_filter_explicitly_requests_cursor, - emit_cursor_setup_warnings, emit_cursor_sync_warning, has_cursor_usage_cache_for_report, - setup_warnings_for_report, use_env_roots, -}; +use crate::commands::shared::{use_env_roots, ReportEnvelope}; use crate::tui; use anyhow::Result; +use std::path::PathBuf; #[derive(serde::Serialize)] #[serde(rename_all = "camelCase")] @@ -103,7 +100,6 @@ pub(crate) struct GraphExportData { contributions: Vec, #[serde(skip_serializing_if = "Option::is_none")] time_metrics: Option, - health: tokscale_core::source_health::HealthReport, } pub(crate) fn to_graph_export_data(graph: &tokscale_core::GraphResult) -> GraphExportData { @@ -188,13 +184,12 @@ pub(crate) fn to_graph_export_data(graph: &tokscale_core::GraphResult) -> GraphE max_concurrent_sessions: tm.max_concurrent_sessions, session_count: tm.session_count, }), - health: graph.health.clone(), } } #[allow(clippy::too_many_arguments)] pub(crate) fn run_graph_command( - output: Option, + output: Option, home_dir: Option, clients: Option>, since: Option, @@ -208,10 +203,6 @@ pub(crate) fn run_graph_command( use tokscale_core::{generate_local_graph_report, GroupBy, ReportOptions}; let show_progress = output.is_some() && !no_spinner; - let had_cursor_cache = has_cursor_usage_cache_for_report(&home_dir); - let explicit_cursor_filter = client_filter_explicitly_requests_cursor(&clients); - let cursor_sync_result = auto_sync_cursor_for_local_report(&home_dir, &clients); - let cursor_setup_warnings = setup_warnings_for_report(&home_dir, &clients); if show_progress { eprintln!(" Scanning session data..."); @@ -238,33 +229,32 @@ pub(crate) fn run_graph_command( }) .await }) - .map_err(|e| anyhow::anyhow!(e))?; - emit_cursor_sync_warning( - cursor_sync_result.as_ref(), - had_cursor_cache, - explicit_cursor_filter, - ); + .map_err(anyhow::Error::new)?; super::shared::emit_health_summary(&graph_result.health); - emit_cursor_setup_warnings(&cursor_setup_warnings); let processing_time_ms = start.elapsed().as_millis() as u32; let output_data = to_graph_export_data(&graph_result); - let json_output = serde_json::to_string_pretty(&output_data)?; + let output_document = ReportEnvelope::new( + output_data, + graph_result.health.clone(), + processing_time_ms as u64, + ); + let json_output = serde_json::to_string_pretty(&output_document)?; if let Some(output_path) = output { std::fs::write(&output_path, json_output)?; eprintln!( "{}", - format!("✓ Graph data written to {}", output_path).green() + format!("✓ Graph data written to {}", output_path.display()).green() ); eprintln!( "{}", format!( " {} days, {} clients, {} models", - output_data.contributions.len(), - output_data.summary.clients.len(), - output_data.summary.models.len() + output_document.data.contributions.len(), + output_document.data.summary.clients.len(), + output_document.data.summary.models.len() ) .bright_black() ); @@ -272,37 +262,22 @@ pub(crate) fn run_graph_command( "{}", format!( " Total: {}", - format_currency(output_data.summary.total_cost) + format_currency(output_document.data.summary.total_cost) ) .bright_black() ); - - if benchmark { - eprintln!( - "{}", - format!(" Processing time: {}ms (Rust native)", processing_time_ms).bright_black() - ); - if let Some(sync) = cursor_sync_result { - if sync.synced { - eprintln!( - "{}", - format!( - " Cursor: {} usage events synced (full lifetime data)", - sync.rows - ) - .bright_black() - ); - } else if let Some(err) = sync.error { - if had_cursor_cache { - eprintln!("{}", format!(" Cursor: sync failed - {}", err).yellow()); - } - } - } - } + println!("{}", output_path.display()); } else { println!("{}", json_output); } + if benchmark { + eprintln!( + "{}", + format!(" Processing time: {}ms (Rust native)", processing_time_ms).bright_black() + ); + } + Ok(()) } @@ -345,7 +320,12 @@ mod tests { }, }; - let json = serde_json::to_value(to_graph_export_data(&graph)).unwrap(); + let json = serde_json::to_value(ReportEnvelope::new( + to_graph_export_data(&graph), + graph.health, + 0_u64, + )) + .unwrap(); assert_eq!(json["health"]["complete"], false); assert_eq!(json["health"]["cleanSources"], 4); diff --git a/crates/tokscale-cli/src/commands/headless.rs b/crates/tokscale-cli/src/commands/headless.rs index b3fc0b328..12e392e86 100644 --- a/crates/tokscale-cli/src/commands/headless.rs +++ b/crates/tokscale-cli/src/commands/headless.rs @@ -1,4 +1,3 @@ -use crate::tui; use anyhow::Result; use std::path::{Path, PathBuf}; use std::time::Duration; @@ -96,31 +95,31 @@ pub(crate) fn run_capture_command( pub(crate) fn run_headless_command( source: &str, - args: Vec, - format: Option, + command: Vec, + format: Option<&str>, output: Option, no_auto_flags: bool, + timeout: Duration, ) -> Result<()> { use chrono::Utc; use uuid::Uuid; let source_lower = source.to_lowercase(); - if source_lower != "codex" { - eprintln!("\n Error: Unknown headless source '{}'.", source); - eprintln!(" Currently only 'codex' is supported.\n"); - std::process::exit(1); - } + anyhow::ensure!( + source_lower == "codex", + "unsupported headless source `{source}`" + ); + let (program, child_args) = command + .split_first() + .ok_or_else(|| anyhow::anyhow!("headless child command must not be empty"))?; let resolved_format = match format { - Some(f) if f == "json" || f == "jsonl" => f, - Some(f) => { - eprintln!("\n Error: Invalid format '{}'. Use json or jsonl.\n", f); - std::process::exit(1); - } + Some(f) if f == "json" || f == "jsonl" => f.to_string(), + Some(f) => anyhow::bail!("invalid headless format `{f}`"), None => "jsonl".to_string(), }; - let mut final_args = args.clone(); + let mut final_args = child_args.to_vec(); if !no_auto_flags && source_lower == "codex" && !final_args.contains(&"--json".to_string()) { final_args.push("--json".to_string()); } @@ -158,23 +157,20 @@ pub(crate) fn run_headless_command( dir.join(filename) }; - let settings = tui::settings::Settings::load()?; - let timeout = settings.get_native_timeout()?; - use colored::Colorize; - println!("\n {}", "Headless capture".cyan()); - println!(" {}", format!("source: {}", source_lower).bright_black()); - println!( + eprintln!("\n {}", "Headless capture".cyan()); + eprintln!(" {}", format!("source: {}", source_lower).bright_black()); + eprintln!( " {}", format!("output: {}", output_path.display()).bright_black() ); - println!( + eprintln!( " {}", format!("timeout: {}s", timeout.as_secs()).bright_black() ); - println!(); + eprintln!(); - let outcome = run_capture_command(&source_lower, &final_args, &output_path, timeout)?; + let outcome = run_capture_command(program, &final_args, &output_path, timeout)?; if outcome.timed_out { eprintln!( @@ -182,15 +178,11 @@ pub(crate) fn run_headless_command( format!("\n Subprocess timed out after {}s", timeout.as_secs()).red() ); eprintln!("{}", " Partial output saved. Increase timeout with TOKSCALE_NATIVE_TIMEOUT_MS or settings.json".bright_black()); - println!(); std::process::exit(124); } - println!( - "{}", - format!("✓ Saved headless output to {}", output_path.display()).green() - ); - println!(); + eprintln!("{}", "✓ Headless output saved".green()); + println!("{}", output_path.display()); if outcome.exit_code != 0 { std::process::exit(outcome.exit_code); diff --git a/crates/tokscale-cli/src/commands/hourly.rs b/crates/tokscale-cli/src/commands/hourly.rs index 042338701..80b65c0b9 100644 --- a/crates/tokscale-cli/src/commands/hourly.rs +++ b/crates/tokscale-cli/src/commands/hourly.rs @@ -2,11 +2,7 @@ use crate::commands::render::{ dim_borders, format_currency, format_tokens_with_commas, formatted_unique_model_names, LightSpinner, TABLE_PRESET, }; -use crate::commands::shared::{ - auto_sync_cursor_for_local_report, client_filter_explicitly_requests_cursor, - emit_cursor_setup_warnings, emit_cursor_sync_warning, get_date_range_label, - has_cursor_usage_cache_for_report, setup_warnings_for_report, use_env_roots, -}; +use crate::commands::shared::{get_date_range_label, use_env_roots, ReportEnvelope}; use crate::tui::{self, get_client_display_name}; use anyhow::Result; use std::io::IsTerminal; @@ -42,15 +38,11 @@ pub(crate) fn run_hourly_report( let date_range = get_date_range_label(today, week, month_flag, &since, &until, &year); - let had_cursor_cache = has_cursor_usage_cache_for_report(&home_dir); - let explicit_cursor_filter = client_filter_explicitly_requests_cursor(&clients); let spinner = if no_spinner { None } else { Some(LightSpinner::start("Scanning session data...")) }; - let cursor_sync_result = auto_sync_cursor_for_local_report(&home_dir, &clients); - let cursor_setup_warnings = setup_warnings_for_report(&home_dir, &clients); let use_env_roots = use_env_roots(&home_dir); let scanner_settings = tui::settings::load_scanner_settings_for_home(&home_dir)?; let start = Instant::now(); @@ -69,16 +61,11 @@ pub(crate) fn run_hourly_report( }) .await }) - .map_err(|e| anyhow::anyhow!(e))?; + .map_err(anyhow::Error::new)?; if let Some(spinner) = spinner { spinner.stop(); } - emit_cursor_sync_warning( - cursor_sync_result.as_ref(), - had_cursor_cache, - explicit_cursor_filter, - ); super::shared::emit_health_summary(&report.health); let processing_time_ms = start.elapsed().as_millis(); @@ -101,16 +88,14 @@ pub(crate) fn run_hourly_report( #[derive(serde::Serialize)] #[serde(rename_all = "camelCase")] - struct HourlyReportJson { + struct HourlyReportData { entries: Vec, total_cost: f64, - processing_time_ms: u32, - #[serde(skip_serializing_if = "Vec::is_empty")] - warnings: Vec, - health: tokscale_core::source_health::HealthReport, } - let output = HourlyReportJson { + let health = report.health.clone(); + let report_processing_time_ms = report.processing_time_ms; + let data = HourlyReportData { entries: report .entries .into_iter() @@ -128,16 +113,12 @@ pub(crate) fn run_hourly_report( }) .collect(), total_cost: report.total_cost, - processing_time_ms: report.processing_time_ms, - warnings: cursor_setup_warnings, - health: report.health.clone(), }; + let output = ReportEnvelope::new(data, health, report_processing_time_ms as u64); println!("{}", serde_json::to_string_pretty(&output)?); } else { use comfy_table::{Cell, CellAlignment, Color, ContentArrangement, Table}; - - emit_cursor_setup_warnings(&cursor_setup_warnings); let term_width = crossterm::terminal::size() .map(|(w, _)| w as usize) .unwrap_or(120); @@ -291,13 +272,14 @@ pub(crate) fn run_hourly_report( "Total:".bold(), format_currency(report.total_cost).green().bold() ); + } - if benchmark { - println!( - "{}", - format!(" Processing time: {}ms (Rust native)", processing_time_ms).bright_black() - ); - } + if benchmark { + use colored::Colorize; + eprintln!( + "{}", + format!(" Processing time: {}ms (Rust native)", processing_time_ms).bright_black() + ); } Ok(()) diff --git a/crates/tokscale-cli/src/commands/integrations.rs b/crates/tokscale-cli/src/commands/integrations.rs index 9f2b2058e..0a9f243c8 100644 --- a/crates/tokscale-cli/src/commands/integrations.rs +++ b/crates/tokscale-cli/src/commands/integrations.rs @@ -1,36 +1,7 @@ -use crate::{ - antigravity, commands, cursor, trae, warp, AntigravitySubcommand, CodexSubcommand, - CursorSubcommand, TraeSubcommand, WarpSubcommand, -}; +use crate::cli::{AntigravitySubcommand, WarpSubcommand}; +use crate::{antigravity, warp}; use anyhow::Result; -pub(crate) fn run_codex_command(subcommand: CodexSubcommand) -> Result<()> { - match subcommand { - CodexSubcommand::Import { name } => commands::usage::codex::run_codex_import(name), - CodexSubcommand::Accounts { json } => commands::usage::codex::run_codex_accounts(json), - CodexSubcommand::Switch { name } => commands::usage::codex::run_codex_switch(&name), - CodexSubcommand::Remove { name } => commands::usage::codex::run_codex_remove(&name), - CodexSubcommand::Status { name, json } => { - commands::usage::codex::run_codex_status(name, json) - } - } -} - -pub(crate) fn run_cursor_command(subcommand: CursorSubcommand) -> Result<()> { - match subcommand { - CursorSubcommand::Login { name } => cursor::run_cursor_login(name), - CursorSubcommand::Logout { - name, - all, - purge_cache, - } => cursor::run_cursor_logout(name, all, purge_cache), - CursorSubcommand::Status { name } => cursor::run_cursor_status(name), - CursorSubcommand::Accounts { json } => cursor::run_cursor_accounts(json), - CursorSubcommand::Sync { json } => cursor::run_cursor_sync(json), - CursorSubcommand::Switch { name } => cursor::run_cursor_switch(&name), - } -} - pub(crate) fn run_antigravity_command(subcommand: AntigravitySubcommand) -> Result<()> { match subcommand { AntigravitySubcommand::Sync => antigravity::run_antigravity_sync(), @@ -39,150 +10,6 @@ pub(crate) fn run_antigravity_command(subcommand: AntigravitySubcommand) -> Resu } } -/// Parse `--variant` into a typed value. -/// -/// Returns: -/// - `Ok(Some(v))` when a recognized value was provided -/// - `Ok(None)` when the flag was omitted entirely -/// - `Err` when an unrecognized value was provided -/// -/// The earlier version returned `Option<_>` and merged the "unrecognized" and -/// "omitted" cases, which let callers silently fall through to "all variants" -/// when the user typed something like `--variant slo` — they got every variant -/// touched instead of an error. -pub(crate) fn parse_variant_arg(arg: Option<&str>) -> Result> { - match arg { - Some("solo") => Ok(Some(trae::auth::TraeVariant::Solo)), - Some("ide") => Ok(Some(trae::auth::TraeVariant::Ide)), - Some(other) => anyhow::bail!("unknown variant: {other}, valid values: solo, ide"), - None => Ok(None), - } -} - -pub(crate) fn run_trae_command(subcommand: TraeSubcommand) -> Result<()> { - use colored::Colorize; - - match subcommand { - TraeSubcommand::Login { manual, variant } => { - if manual { - use std::io::{self, Write}; - // Default to international Solo when `--variant` is omitted. - let selected = - parse_variant_arg(variant.as_deref())?.unwrap_or(trae::auth::TraeVariant::Solo); - println!(); - println!(" {}", "Trae Manual Token Login".cyan()); - println!( - " {}", - "Paste your JWT access token from the browser DevTools:".bright_black() - ); - println!( - " {}", - "1. Open https://www.trae.ai/account-setting#usage".bright_black() - ); - println!( - " {}", - "2. F12 → Network → filter 'query_user_usage' → copy Authorization value" - .bright_black() - ); - print!(" Token: "); - io::stdout().flush()?; - let mut token = String::new(); - io::stdin().read_line(&mut token)?; - let token = token.trim().to_string(); - if token.is_empty() { - anyhow::bail!("token must not be empty"); - } - trae::auth::save_manual_token(selected, token, None)?; - println!( - "\n {}", - format!("Token saved for {}", selected.client_str()).green() - ); - } else { - let rt = tokio::runtime::Runtime::new()?; - let variants: Vec = - match parse_variant_arg(variant.as_deref())? { - Some(v) => vec![v], - None => trae::auth::all_variants().to_vec(), - }; - - let mut any_success = false; - for v in variants { - match rt.block_on(trae::auth::resolve_token(v)) { - Ok(_) => { - println!(" {} logged in (auto-detected)", v.client_str().green()); - any_success = true; - } - Err(e) => { - println!(" {} auto-login failed: {}", v.client_str().yellow(), e); - } - } - } - if !any_success { - println!( - " {}", - "No Trae credentials found. Use --manual to paste a token by hand." - .yellow() - ); - } - } - Ok(()) - } - TraeSubcommand::Logout { variant } => { - let variants: Vec = - match parse_variant_arg(variant.as_deref())? { - Some(v) => vec![v], - None => trae::auth::all_variants().to_vec(), - }; - for v in variants { - trae::auth::logout(v)?; - println!(" {} logged out", v.client_str().green()); - } - Ok(()) - } - TraeSubcommand::Status { json } => { - let mut status = serde_json::Map::new(); - for v in trae::auth::all_variants() { - let has = trae::auth::has_credentials(v); - if json { - status.insert(v.client_str().to_string(), serde_json::Value::Bool(has)); - } else { - println!( - " {}: {}", - v.client_str(), - if has { - "authenticated".green() - } else { - "not authenticated".yellow() - } - ); - } - } - if json { - println!("{}", serde_json::to_string_pretty(&status)?); - } - Ok(()) - } - TraeSubcommand::Sync { since, include_aux } => { - let days = since.unwrap_or(30); - // Negative `days` would compute `now - (negative * 86400)` → a - // future `start_time`, and zero collapses the query window to an - // empty range. Reject both at the CLI boundary instead of - // forwarding garbage to the sync layer. - if days <= 0 { - anyhow::bail!("--since must be a positive number of days (got {days})"); - } - // Trae IDE and Trae Solo share account-level usage data, so we - // always sync once using whichever credential source is available. - let variants: Vec = trae::auth::all_variants() - .into_iter() - .filter(|v| trae::auth::has_credentials(*v)) - .collect(); - let rt = tokio::runtime::Runtime::new()?; - rt.block_on(trae::sync::run_trae_sync(&variants, days, include_aux)) - } - } -} - pub(crate) fn run_warp_command(subcommand: WarpSubcommand) -> Result<()> { match subcommand { WarpSubcommand::Login { token, cookie } => warp::run_warp_login(token, cookie), diff --git a/crates/tokscale-cli/src/commands/models.rs b/crates/tokscale-cli/src/commands/models.rs index 1a283f15f..55b645223 100644 --- a/crates/tokscale-cli/src/commands/models.rs +++ b/crates/tokscale-cli/src/commands/models.rs @@ -1,16 +1,11 @@ use crate::claude_diagnostics; -use crate::commands::cache::{ - resolve_should_write_cache, validate_light_cache_write, write_light_cache, -}; use crate::commands::render::{ aggregate_model_report_performance, dim_borders, format_currency, format_model_name, format_ms_per_1k, format_tokens_with_commas, LightSpinner, TABLE_PRESET, }; use crate::commands::shared::{ - auto_sync_cursor_for_local_report, client_filter_explicitly_requests_cursor, - emit_client_diagnostics, emit_cursor_setup_warnings, emit_cursor_sync_warning, - get_date_range_label, has_cursor_usage_cache_for_report, model_usage_includes_client, - resolve_effective_home_dir, setup_warnings_for_report, use_env_roots, + emit_client_diagnostics, get_date_range_label, model_usage_includes_client, + resolve_effective_home_dir, use_env_roots, ReportEnvelope, }; use crate::tui::{ self, get_client_display_name, get_provider_display_name, truncate_model_display_name, @@ -52,8 +47,6 @@ pub(crate) fn run_models_report( week: bool, month_flag: bool, group_by: tokscale_core::GroupBy, - cli_write_cache: bool, - cli_no_write_cache: bool, ) -> Result<()> { use std::time::Instant; use tokio::runtime::Runtime; @@ -62,30 +55,14 @@ pub(crate) fn run_models_report( if !json { tui::config::TokscaleConfig::initialize()?; } - let should_write_cache = if json { - false - } else { - let settings = tui::settings::Settings::load()?; - let should_write = - resolve_should_write_cache(cli_write_cache, cli_no_write_cache, &settings); - if should_write { - validate_light_cache_write(&home_dir)?; - } - should_write - }; - let date_range = get_date_range_label(today, week, month_flag, &since, &until, &year); let effective_home_dir = resolve_effective_home_dir(&home_dir); - let had_cursor_cache = has_cursor_usage_cache_for_report(&home_dir); - let explicit_cursor_filter = client_filter_explicitly_requests_cursor(&clients); let spinner = if no_spinner { None } else { Some(LightSpinner::start("Scanning session data...")) }; - let cursor_sync_result = auto_sync_cursor_for_local_report(&home_dir, &clients); - let cursor_setup_warnings = setup_warnings_for_report(&home_dir, &clients); let use_env_roots = use_env_roots(&home_dir); let scanner_settings = tui::settings::load_scanner_settings_for_home(&home_dir)?; let start = Instant::now(); @@ -104,16 +81,11 @@ pub(crate) fn run_models_report( }) .await }) - .map_err(|e| anyhow::anyhow!(e))?; + .map_err(anyhow::Error::new)?; if let Some(spinner) = spinner { spinner.stop(); } - emit_cursor_sync_warning( - cursor_sync_result.as_ref(), - had_cursor_cache, - explicit_cursor_filter, - ); super::shared::emit_health_summary(&report.health); let processing_time_ms = start.elapsed().as_millis(); let claude_message_count = report @@ -140,6 +112,7 @@ pub(crate) fn run_models_report( report.total_cache_read, report.total_cache_write, ]); + emit_client_diagnostics(&diagnostics); if json { #[derive(serde::Serialize)] @@ -166,7 +139,7 @@ pub(crate) fn run_models_report( #[derive(serde::Serialize)] #[serde(rename_all = "camelCase")] - struct ModelReportJson { + struct ModelReportData { group_by: String, entries: Vec, total_input: i64, @@ -176,16 +149,11 @@ pub(crate) fn run_models_report( total_tokens: i64, total_messages: i32, total_cost: f64, - processing_time_ms: u32, - #[serde(skip_serializing_if = "Vec::is_empty")] - warnings: Vec, - #[serde(skip_serializing_if = "Vec::is_empty")] - diagnostics: Vec, - health: tokscale_core::source_health::HealthReport, } let health = report.health.clone(); - let output = ModelReportJson { + let report_processing_time_ms = report.processing_time_ms; + let data = ModelReportData { group_by: group_by.to_string(), entries: report .entries @@ -234,17 +202,11 @@ pub(crate) fn run_models_report( total_tokens: report_token_total, total_messages: report.total_messages, total_cost: report.total_cost, - processing_time_ms: report.processing_time_ms, - warnings: cursor_setup_warnings, - diagnostics, - health, }; + let output = ReportEnvelope::new(data, health, report_processing_time_ms as u64); println!("{}", serde_json::to_string_pretty(&output)?); } else { use comfy_table::{Attribute, Cell, CellAlignment, Color, ContentArrangement, Table}; - emit_client_diagnostics(&diagnostics); - - emit_cursor_setup_warnings(&cursor_setup_warnings); let total_performance = aggregate_model_report_performance(&report.entries); let term_width = crossterm::terminal::size() .map(|(w, _)| w as usize) @@ -770,19 +732,15 @@ pub(crate) fn run_models_report( format_currency(report.total_cost) ); - if benchmark { - use colored::Colorize; - println!( - "{}", - format!(" Processing time: {}ms (Rust native)", processing_time_ms).bright_black() - ); - } - io::stdout().flush()?; + } - if should_write_cache { - write_light_cache(&clients, &since, &until, &year, &group_by)?; - } + if benchmark { + use colored::Colorize; + eprintln!( + "{}", + format!(" Processing time: {}ms (Rust native)", processing_time_ms).bright_black() + ); } Ok(()) diff --git a/crates/tokscale-cli/src/commands/monthly.rs b/crates/tokscale-cli/src/commands/monthly.rs index f545f61f3..fa20791ac 100644 --- a/crates/tokscale-cli/src/commands/monthly.rs +++ b/crates/tokscale-cli/src/commands/monthly.rs @@ -2,11 +2,7 @@ use crate::commands::render::{ dim_borders, format_currency, format_tokens_with_commas, formatted_unique_model_names, LightSpinner, TABLE_PRESET, }; -use crate::commands::shared::{ - auto_sync_cursor_for_local_report, client_filter_explicitly_requests_cursor, - emit_cursor_setup_warnings, emit_cursor_sync_warning, get_date_range_label, - has_cursor_usage_cache_for_report, setup_warnings_for_report, use_env_roots, -}; +use crate::commands::shared::{get_date_range_label, use_env_roots, ReportEnvelope}; use crate::tui; use anyhow::Result; use std::io::IsTerminal; @@ -51,15 +47,11 @@ pub(crate) fn run_monthly_report( let date_range = get_date_range_label(today, week, month_flag, &since, &until, &year); - let had_cursor_cache = has_cursor_usage_cache_for_report(&home_dir); - let explicit_cursor_filter = client_filter_explicitly_requests_cursor(&clients); let spinner = if no_spinner { None } else { Some(LightSpinner::start("Scanning session data...")) }; - let cursor_sync_result = auto_sync_cursor_for_local_report(&home_dir, &clients); - let cursor_setup_warnings = setup_warnings_for_report(&home_dir, &clients); let use_env_roots = use_env_roots(&home_dir); let scanner_settings = tui::settings::load_scanner_settings_for_home(&home_dir)?; let start = Instant::now(); @@ -78,16 +70,11 @@ pub(crate) fn run_monthly_report( }) .await }) - .map_err(|e| anyhow::anyhow!(e))?; + .map_err(anyhow::Error::new)?; if let Some(spinner) = spinner { spinner.stop(); } - emit_cursor_sync_warning( - cursor_sync_result.as_ref(), - had_cursor_cache, - explicit_cursor_filter, - ); super::shared::emit_health_summary(&report.health); let processing_time_ms = start.elapsed().as_millis(); @@ -108,16 +95,14 @@ pub(crate) fn run_monthly_report( #[derive(serde::Serialize)] #[serde(rename_all = "camelCase")] - struct MonthlyReportJson { + struct MonthlyReportData { entries: Vec, total_cost: f64, - processing_time_ms: u32, - #[serde(skip_serializing_if = "Vec::is_empty")] - warnings: Vec, - health: tokscale_core::source_health::HealthReport, } - let output = MonthlyReportJson { + let health = report.health.clone(); + let report_processing_time_ms = report.processing_time_ms; + let data = MonthlyReportData { entries: report .entries .into_iter() @@ -136,16 +121,12 @@ pub(crate) fn run_monthly_report( }) .collect(), total_cost: report.total_cost, - processing_time_ms: report.processing_time_ms, - warnings: cursor_setup_warnings, - health: report.health.clone(), }; + let output = ReportEnvelope::new(data, health, report_processing_time_ms as u64); println!("{}", serde_json::to_string_pretty(&output)?); } else { use comfy_table::{Attribute, Cell, CellAlignment, Color, ContentArrangement, Table}; - - emit_cursor_setup_warnings(&cursor_setup_warnings); let term_width = crossterm::terminal::size() .map(|(w, _)| w as usize) .unwrap_or(120); @@ -302,14 +283,14 @@ pub(crate) fn run_monthly_report( "\x1b[90m\n Total Cost: \x1b[32m{}\x1b[90m\x1b[0m", format_currency(report.total_cost) ); + } - if benchmark { - use colored::Colorize; - println!( - "{}", - format!(" Processing time: {}ms (Rust native)", processing_time_ms).bright_black() - ); - } + if benchmark { + use colored::Colorize; + eprintln!( + "{}", + format!(" Processing time: {}ms (Rust native)", processing_time_ms).bright_black() + ); } Ok(()) diff --git a/crates/tokscale-cli/src/commands/pricing.rs b/crates/tokscale-cli/src/commands/pricing.rs index 7dd8b1cc8..4fd98e888 100644 --- a/crates/tokscale-cli/src/commands/pricing.rs +++ b/crates/tokscale-cli/src/commands/pricing.rs @@ -3,7 +3,7 @@ use anyhow::Result; pub(crate) fn run_pricing_lookup( model_id: &str, json: bool, - provider: Option<&str>, + source: Option<&str>, no_spinner: bool, ) -> Result<()> { use colored::Colorize; @@ -12,16 +12,14 @@ pub(crate) fn run_pricing_lookup( use tokio::runtime::Runtime; use tokscale_core::pricing::PricingService; - if model_id.eq_ignore_ascii_case("list-overrides") { - return run_pricing_list_overrides(json); - } - - let provider_normalized = provider.map(|p| p.to_lowercase()); + let source_normalized = source.map(|value| value.to_lowercase()); let spinner = if no_spinner { None } else { - let provider_label = provider.map(|p| format!(" from {}", p)).unwrap_or_default(); + let provider_label = source + .map(|value| format!(" from {}", value)) + .unwrap_or_default(); let pb = ProgressBar::new_spinner(); pb.set_style(ProgressStyle::default_spinner()); pb.set_message(format!("Fetching pricing data{}...", provider_label)); @@ -32,29 +30,13 @@ pub(crate) fn run_pricing_lookup( let rt = Runtime::new()?; let result = match rt.block_on(async { let svc = PricingService::get_or_init().await?; - Ok::<_, String>(svc.lookup_with_source(model_id, provider_normalized.as_deref())) + Ok::<_, String>(svc.lookup_with_source(model_id, source_normalized.as_deref())) }) { Ok(result) => result, Err(err) => { if let Some(pb) = spinner { pb.finish_and_clear(); } - if json { - #[derive(serde::Serialize)] - #[serde(rename_all = "camelCase")] - struct ErrorOutput { - error: String, - model_id: String, - } - println!( - "{}", - serde_json::to_string_pretty(&ErrorOutput { - error: err, - model_id: model_id.to_string(), - })? - ); - std::process::exit(1); - } return Err(anyhow::anyhow!(err)); } }; @@ -103,20 +85,7 @@ pub(crate) fn run_pricing_lookup( println!("{}", serde_json::to_string_pretty(&output)?); } None => { - #[derive(serde::Serialize)] - #[serde(rename_all = "camelCase")] - struct ErrorOutput { - error: String, - model_id: String, - } - - let output = ErrorOutput { - error: "Model not found".to_string(), - model_id: model_id.to_string(), - }; - - println!("{}", serde_json::to_string_pretty(&output)?); - std::process::exit(1); + return Err(anyhow::anyhow!("Model not found: {model_id}")); } } } else { @@ -152,8 +121,10 @@ pub(crate) fn run_pricing_lookup( println!(); } None => { - println!("\n {}\n", format!("Model not found: {}", model_id).red()); - std::process::exit(1); + return Err(anyhow::anyhow!( + "{}", + format!("Model not found: {model_id}").red() + )); } } } diff --git a/crates/tokscale-cli/src/commands/shared.rs b/crates/tokscale-cli/src/commands/shared.rs index 636af782e..c4eebd5a9 100644 --- a/crates/tokscale-cli/src/commands/shared.rs +++ b/crates/tokscale-cli/src/commands/shared.rs @@ -1,4 +1,6 @@ -use crate::{claude_diagnostics, cursor, tui, ClientFlags}; +use crate::cli::ClientFlags; +use crate::failure::InvalidConfiguration; +use crate::{claude_diagnostics, tui}; use anyhow::Result; use std::path::PathBuf; use tokscale_core::ClientId; @@ -86,11 +88,12 @@ pub(crate) fn parse_default_client_filters(defaults: &[String]) -> Result Option { @@ -110,311 +113,6 @@ pub(crate) fn parse_client_id_set(clients: &[String]) -> std::collections::HashS .collect() } -pub(crate) fn client_filter_includes_cursor(clients: &Option>) -> bool { - clients - .as_ref() - .is_none_or(|sources| sources.iter().any(|source| source == "cursor")) -} - -pub(crate) fn client_filter_explicitly_requests_cursor(clients: &Option>) -> bool { - clients - .as_ref() - .is_some_and(|sources| sources.iter().any(|source| source == "cursor")) -} - -#[derive(Debug)] -pub(crate) struct CursorSetupState { - has_credentials: bool, - has_cache: bool, - cache_glob: String, - home_override: bool, -} - -pub(crate) fn cursor_setup_state(home_dir: &Option) -> Option { - let (home_path, home_override) = match home_dir { - Some(home) => (PathBuf::from(home), true), - None => (dirs::home_dir()?, false), - }; - let has_credentials = if home_override { - cursor::has_active_credentials_in_home(&home_path) - } else { - cursor::is_cursor_logged_in() - }; - let has_cache = cursor::has_cursor_usage_cache_in_home(&home_path); - let cache_glob = if home_override { - home_path - .join(".config/tokscale/cursor-cache/usage*.csv") - .to_string_lossy() - .to_string() - } else { - "~/.config/tokscale/cursor-cache/usage*.csv".to_string() - }; - - Some(CursorSetupState { - has_credentials, - has_cache, - cache_glob, - home_override, - }) -} - -pub(crate) fn has_cursor_usage_cache_for_report(home_dir: &Option) -> bool { - cursor_setup_state(home_dir).is_some_and(|state| state.has_cache) -} - -pub(crate) fn cursor_setup_warnings_for_report( - home_dir: &Option, - clients: &Option>, -) -> Vec { - if !client_filter_explicitly_requests_cursor(clients) { - return Vec::new(); - } - - let Some(state) = cursor_setup_state(home_dir) else { - return vec![ - "Cursor usage requires Tokscale's Cursor API cache, but the home directory could not be resolved. Run `tokscale cursor login` and `tokscale cursor sync --json`. Tokscale does not parse local `~/.cursor` session data.".to_string(), - ]; - }; - if state.has_cache { - return Vec::new(); - } - - let action = if state.home_override { - "run `tokscale cursor login` and `tokscale cursor sync --json`, or populate that cache before running a report with --home" - } else if state.has_credentials { - "run `tokscale cursor sync --json`" - } else { - "run `tokscale cursor login` and `tokscale cursor sync --json`" - }; - - vec![format!( - "Cursor usage requires Tokscale's Cursor API cache at `{}`; {}. Tokscale does not parse local `~/.cursor` session data.", - state.cache_glob, action - )] -} - -pub(crate) fn emit_cursor_setup_warnings(warnings: &[String]) { - if warnings.is_empty() { - return; - } - - use colored::Colorize; - for warning in warnings { - eprintln!("{}", format!(" Warning: {}", warning).yellow()); - } -} - -pub(crate) fn setup_warnings_for_report( - home_dir: &Option, - clients: &Option>, -) -> Vec { - cursor_setup_warnings_for_report(home_dir, clients) -} - -pub(crate) fn should_auto_sync_cursor_for_local_report( - home_dir: &Option, - clients: &Option>, -) -> bool { - home_dir.is_none() && client_filter_includes_cursor(clients) -} - -pub(crate) fn auto_sync_cursor_for_local_report( - home_dir: &Option, - clients: &Option>, -) -> Option { - if !should_auto_sync_cursor_for_local_report(home_dir, clients) - || !cursor::is_cursor_logged_in() - { - return None; - } - - // Skip the implicit refresh when each expected Cursor account cache is - // recent enough — running `tokscale models` 30× in a script must not - // produce 30 Cursor API calls. The manual `tokscale cursor sync` command - // bypasses this gate. - if cursor::cursor_usage_cache_is_fresh(cursor::CURSOR_AUTO_SYNC_FRESHNESS) { - return None; - } - - Some(run_best_effort_cursor_sync_with_runtime_factory( - tokio::runtime::Runtime::new, - )) -} - -pub(crate) fn run_best_effort_cursor_sync_with_runtime_factory( - build_runtime: F, -) -> cursor::SyncCursorResult -where - F: FnOnce() -> std::io::Result, -{ - match build_runtime() { - Ok(rt) => rt.block_on(async { cursor::sync_cursor_cache().await }), - Err(error) => cursor::SyncCursorResult { - synced: false, - rows: 0, - error: Some(format!( - "Failed to initialize Cursor sync runtime: {}", - error - )), - }, - } -} - -pub(crate) fn auto_sync_cursor_before_tui( - home_dir: &Option, - clients: &Option>, -) -> Result<()> { - let had_cursor_cache = has_cursor_usage_cache_for_report(home_dir); - let explicit_cursor_filter = client_filter_explicitly_requests_cursor(clients); - let cursor_sync_result = auto_sync_cursor_for_local_report(home_dir, clients); - emit_cursor_sync_warning( - cursor_sync_result.as_ref(), - had_cursor_cache, - explicit_cursor_filter, - ); - let cursor_setup_warnings = setup_warnings_for_report(home_dir, clients); - emit_cursor_setup_warnings(&cursor_setup_warnings); - Ok(()) -} - -pub(crate) fn emit_cursor_sync_warning( - sync: Option<&cursor::SyncCursorResult>, - had_cursor_cache: bool, - explicit_cursor_filter: bool, -) { - let Some(sync) = sync else { - return; - }; - let Some(error) = sync.error.as_ref() else { - return; - }; - if sync.synced || had_cursor_cache || explicit_cursor_filter { - use colored::Colorize; - let prefix = if sync.synced { - "Cursor sync warning" - } else if had_cursor_cache { - "Cursor sync failed; using cached data" - } else { - "Cursor sync failed" - }; - eprintln!("{}", format!(" {}: {}", prefix, error).yellow()); - } -} - -pub(crate) fn reject_unsupported_home_override( - home_dir: &Option, - command: &str, -) -> Result<()> { - if home_dir.is_some() { - return Err(anyhow::anyhow!( - "--home is currently supported only for local report commands. It is not supported for `{}`.", - command - )); - } - - Ok(()) -} - -#[derive(Clone, Copy)] -pub(crate) struct UsageParentFlag { - pub(crate) id: &'static str, - pub(crate) display: &'static str, -} - -pub(crate) const USAGE_PARENT_FLAGS: [UsageParentFlag; 17] = [ - UsageParentFlag { - id: "json", - display: "--json", - }, - UsageParentFlag { - id: "light", - display: "--light", - }, - UsageParentFlag { - id: "write_cache", - display: "--write-cache", - }, - UsageParentFlag { - id: "no_write_cache", - display: "--no-write-cache", - }, - UsageParentFlag { - id: "clients", - display: "--client", - }, - UsageParentFlag { - id: "today", - display: "--today", - }, - UsageParentFlag { - id: "week", - display: "--week", - }, - UsageParentFlag { - id: "month", - display: "--month", - }, - UsageParentFlag { - id: "since", - display: "--since", - }, - UsageParentFlag { - id: "until", - display: "--until", - }, - UsageParentFlag { - id: "year", - display: "--year", - }, - UsageParentFlag { - id: "benchmark", - display: "--benchmark", - }, - UsageParentFlag { - id: "group_by", - display: "--group-by", - }, - UsageParentFlag { - id: "no_spinner", - display: "--no-spinner", - }, - UsageParentFlag { - id: "theme", - display: "--theme", - }, - UsageParentFlag { - id: "refresh", - display: "--refresh", - }, - UsageParentFlag { - id: "debug", - display: "--debug", - }, -]; - -pub(crate) fn reject_usage_parent_flags(matches: &clap::ArgMatches) -> Result<()> { - use clap::parser::ValueSource; - - let flags = USAGE_PARENT_FLAGS - .into_iter() - .filter_map(|flag| { - matches - .value_source(flag.id) - .is_some_and(|source| source == ValueSource::CommandLine) - .then_some(flag.display) - }) - .collect::>(); - - if flags.is_empty() { - return Ok(()); - } - - Err(anyhow::anyhow!( - "`usage` does not support parent flag(s): {}. Use `tokscale usage` or `tokscale usage --json`.", - flags.join(", ") - )) -} - pub(crate) fn use_env_roots(home_dir: &Option) -> bool { home_dir.is_none() } @@ -449,16 +147,6 @@ pub(crate) fn emit_client_diagnostics(diagnostics: &[claude_diagnostics::ClientD } } -pub(crate) fn ensure_home_supported_for_tui(home_dir: &Option) -> Result<()> { - if home_dir.is_some() { - return Err(anyhow::anyhow!( - "--home is currently supported for local report commands only. Use `--json`, `--light`, `models`, `monthly`, or `graph` instead of TUI mode." - )); - } - - Ok(()) -} - pub(crate) fn build_date_filter( today: bool, week: bool, @@ -596,3 +284,34 @@ pub(crate) fn emit_health_summary(health: &tokscale_core::source_health::HealthR .yellow() ); } + +/// Stable JSON envelope shared by every local report command. +#[derive(Debug, serde::Serialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct ReportEnvelope { + pub(crate) data: T, + pub(crate) health: tokscale_core::source_health::HealthReport, + pub(crate) metadata: ReportMetadata, +} + +#[derive(Debug, serde::Serialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct ReportMetadata { + pub(crate) processing_time_ms: u64, +} + +impl ReportEnvelope { + pub(crate) fn new( + data: T, + health: tokscale_core::source_health::HealthReport, + processing_time_ms: impl Into, + ) -> Self { + Self { + data, + health, + metadata: ReportMetadata { + processing_time_ms: processing_time_ms.into(), + }, + } + } +} diff --git a/crates/tokscale-cli/src/commands/time_metrics.rs b/crates/tokscale-cli/src/commands/time_metrics.rs index 0e93de1bf..f2abefbdd 100644 --- a/crates/tokscale-cli/src/commands/time_metrics.rs +++ b/crates/tokscale-cli/src/commands/time_metrics.rs @@ -1,9 +1,5 @@ use crate::commands::render::LightSpinner; -use crate::commands::shared::{ - auto_sync_cursor_for_local_report, client_filter_explicitly_requests_cursor, - emit_cursor_setup_warnings, emit_cursor_sync_warning, has_cursor_usage_cache_for_report, - setup_warnings_for_report, use_env_roots, -}; +use crate::commands::shared::{use_env_roots, ReportEnvelope}; use crate::tui; use anyhow::Result; @@ -15,20 +11,17 @@ pub(crate) fn run_time_metrics_report( since: Option, until: Option, year: Option, + benchmark: bool, no_spinner: bool, ) -> Result<()> { use tokio::runtime::Runtime; use tokscale_core::{get_time_metrics_report, GroupBy, ReportOptions}; - let had_cursor_cache = has_cursor_usage_cache_for_report(&home_dir); - let explicit_cursor_filter = client_filter_explicitly_requests_cursor(&clients); let spinner = if no_spinner { None } else { Some(LightSpinner::start("Computing time metrics...")) }; - let cursor_sync_result = auto_sync_cursor_for_local_report(&home_dir, &clients); - let cursor_setup_warnings = setup_warnings_for_report(&home_dir, &clients); let use_env_roots = use_env_roots(&home_dir); let scanner_settings = tui::settings::load_scanner_settings_for_home(&home_dir)?; let rt = Runtime::new()?; @@ -46,16 +39,11 @@ pub(crate) fn run_time_metrics_report( }) .await }) - .map_err(|e| anyhow::anyhow!(e))?; + .map_err(anyhow::Error::new)?; if let Some(spinner) = spinner { spinner.stop(); } - emit_cursor_sync_warning( - cursor_sync_result.as_ref(), - had_cursor_cache, - explicit_cursor_filter, - ); super::shared::emit_health_summary(&report.health); let m = &report.metrics; @@ -63,23 +51,20 @@ pub(crate) fn run_time_metrics_report( if json { #[derive(serde::Serialize)] #[serde(rename_all = "camelCase")] - struct TimeMetricsReportJson<'a> { + struct TimeMetricsData<'a> { metrics: &'a tokscale_core::TimeMetrics, - processing_time_ms: u32, - #[serde(skip_serializing_if = "Vec::is_empty")] - warnings: Vec, - health: &'a tokscale_core::source_health::HealthReport, } - let output = TimeMetricsReportJson { + let data = TimeMetricsData { metrics: &report.metrics, - processing_time_ms: report.processing_time_ms, - warnings: cursor_setup_warnings, - health: &report.health, }; + let output = ReportEnvelope::new( + data, + report.health.clone(), + report.processing_time_ms as u64, + ); println!("{}", serde_json::to_string_pretty(&output)?); } else { - emit_cursor_setup_warnings(&cursor_setup_warnings); println!("Session Time Metrics"); println!("===================="); println!( @@ -96,7 +81,10 @@ pub(crate) fn run_time_metrics_report( ); println!("Max concurrent sessions: {}", m.max_concurrent_sessions); println!("Total sessions: {}", m.session_count); - println!("Processing time: {}ms", report.processing_time_ms); + } + + if benchmark { + eprintln!("Processing time: {}ms", report.processing_time_ms); } Ok(()) diff --git a/crates/tokscale-cli/src/commands/usage/codex.rs b/crates/tokscale-cli/src/commands/usage/codex.rs index e86e65b9a..502ecb7d2 100644 --- a/crates/tokscale-cli/src/commands/usage/codex.rs +++ b/crates/tokscale-cli/src/commands/usage/codex.rs @@ -1,30 +1,20 @@ use anyhow::{Context, Result}; use chrono::{TimeZone, Utc}; -use serde::{Deserialize, Serialize}; -use sha2::{Digest, Sha256}; -use std::collections::HashMap; +use serde::Deserialize; use std::path::{Path, PathBuf}; use super::helpers::capitalize; -use super::{UsageAccount, UsageMetric, UsageOutput}; +use super::{UsageMetric, UsageOutput}; -const CLIENT_ID: &str = "app_EMoamEEZ73f0CkXaXp7hrann"; - -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Clone, Deserialize)] struct Auth { tokens: Option, } -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Clone, Deserialize)] struct Tokens { - #[serde(skip_serializing_if = "Option::is_none")] access_token: Option, - #[serde(skip_serializing_if = "Option::is_none")] - refresh_token: Option, - #[serde(skip_serializing_if = "Option::is_none")] account_id: Option, - #[serde(skip_serializing_if = "Option::is_none")] - id_token: Option, } #[derive(Debug, Deserialize)] @@ -49,63 +39,11 @@ struct Window { reset_at: Option, } -#[derive(Debug, Deserialize)] -struct Refresh { - access_token: Option, - refresh_token: Option, - #[allow(dead_code)] - expires_in: Option, -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct CodexAccount { - tokens: Tokens, - #[serde(rename = "createdAt")] - created_at: String, - #[serde(skip_serializing_if = "Option::is_none")] - label: Option, -} - -#[derive(Debug, Serialize, Deserialize)] -pub struct CodexCredentialsStore { - version: i32, - #[serde(rename = "activeAccountId")] - active_account_id: String, - accounts: HashMap, -} - -#[derive(Debug, Clone, Serialize)] -pub struct CodexAccountInfo { - pub id: String, - #[serde(skip_serializing_if = "Option::is_none")] - pub label: Option, - #[serde(rename = "accountId", skip_serializing_if = "Option::is_none")] - pub account_id: Option, - #[serde(rename = "createdAt")] - pub created_at: String, - #[serde(rename = "isActive")] - pub is_active: bool, -} - -#[derive(Debug, Clone)] -enum CredentialSource { - File(PathBuf), - Keychain, - Store(String), -} - -fn codex_store_path() -> PathBuf { - crate::paths::get_config_dir().join("codex-credentials.json") -} - -fn current_auth_paths() -> Vec { - let home = dirs::home_dir().unwrap_or_else(|| PathBuf::from(".")); +fn current_auth_paths_for_home(home: &Path, codex_home: Option<&str>) -> Vec { let mut paths = Vec::new(); - if let Ok(codex_home) = std::env::var("CODEX_HOME") { - if !codex_home.trim().is_empty() { - paths.push(PathBuf::from(codex_home).join("auth.json")); - } + if let Some(codex_home) = codex_home.map(str::trim).filter(|value| !value.is_empty()) { + paths.push(PathBuf::from(codex_home).join("auth.json")); } paths.push(home.join(".config").join("codex").join("auth.json")); @@ -113,41 +51,30 @@ fn current_auth_paths() -> Vec { paths } -/// Where `switch` writes the codex CLI auth. Derived from -/// [`current_auth_paths`]: an explicit `CODEX_HOME` always wins (even if no -/// auth.json exists there yet); otherwise the first existing path, falling -/// back to the modern config location. -fn auth_write_path() -> Result { - let paths = current_auth_paths(); - let has_codex_home = std::env::var("CODEX_HOME") - .map(|home| !home.trim().is_empty()) - .unwrap_or(false); - - if !has_codex_home { - if let Some(existing) = paths.iter().find(|path| path.exists()) { - return Ok(existing.clone()); - } - } +fn current_auth_paths() -> Vec { + let home = dirs::home_dir().unwrap_or_else(|| PathBuf::from(".")); + let codex_home = std::env::var("CODEX_HOME").ok(); + current_auth_paths_for_home(&home, codex_home.as_deref()) +} - paths - .into_iter() - .next() - .context("Could not determine Codex auth path") +fn parse_auth_file(path: &Path) -> Result> { + let content = std::fs::read_to_string(path) + .with_context(|| format!("Failed to read Codex auth from {}", path.display()))?; + let auth = serde_json::from_str::(&content) + .with_context(|| format!("Failed to parse Codex auth from {}", path.display()))?; + Ok(auth + .tokens + .as_ref() + .and_then(|tokens| tokens.access_token.as_deref()) + .is_some_and(|token| !token.trim().is_empty()) + .then_some(auth)) } -fn read_current_credentials() -> Result<(Auth, CredentialSource)> { - for p in current_auth_paths() { - if p.exists() { - let content = std::fs::read_to_string(&p)?; - if let Ok(auth) = serde_json::from_str::(&content) { - if auth - .tokens - .as_ref() - .and_then(|t| t.access_token.as_ref()) - .is_some() - { - return Ok((auth, CredentialSource::File(p))); - } +fn read_current_credentials() -> Result { + for path in current_auth_paths() { + if path.exists() { + if let Some(auth) = parse_auth_file(&path)? { + return Ok(auth); } } } @@ -157,473 +84,27 @@ fn read_current_credentials() -> Result<(Auth, CredentialSource)> { if auth .tokens .as_ref() - .and_then(|t| t.access_token.as_ref()) - .is_some() + .and_then(|tokens| tokens.access_token.as_deref()) + .is_some_and(|token| !token.trim().is_empty()) { - return Ok((auth, CredentialSource::Keychain)); - } - } - } - - anyhow::bail!("No Codex credentials found. Run 'codex' to log in.") -} - -fn auth_document(tokens: &Tokens) -> serde_json::Value { - serde_json::json!({ - "tokens": tokens, - "last_refresh": chrono::Utc::now().to_rfc3339(), - }) -} - -fn save_auth_tokens(path: &Path, tokens: &Tokens) -> Result<()> { - let content = serde_json::to_string_pretty(&auth_document(tokens))?; - super::helpers::atomic_write_secret(path, content.as_bytes()) - .with_context(|| format!("Failed to write Codex auth to {}", path.display())) -} - -fn persist_tokens(source: &CredentialSource, tokens: &Tokens) { - match source { - CredentialSource::File(path) => { - if let Err(e) = save_auth_tokens(path, tokens) { - eprintln!("warning: failed to save Codex credentials: {e}"); - } - } - CredentialSource::Store(account_id) => { - if let Err(e) = update_account_tokens(account_id, tokens.clone()) { - eprintln!("warning: failed to save Codex account credentials: {e}"); - } - } - CredentialSource::Keychain => {} - } -} - -fn hash_token(token: &str) -> String { - let digest = Sha256::digest(token.as_bytes()); - digest - .iter() - .take(8) - .map(|b| format!("{b:02x}")) - .collect::() -} - -fn derive_account_id(tokens: &Tokens) -> String { - if let Some(account_id) = tokens.account_id.as_deref() { - let trimmed = account_id.trim(); - if !trimmed.is_empty() { - return trimmed.to_string(); - } - } - - if let Some(id_token) = tokens.id_token.as_deref() { - let trimmed = id_token.trim(); - if !trimmed.is_empty() { - return format!("id-{}", hash_token(trimmed)); - } - } - - tokens - .access_token - .as_deref() - .map(|token| format!("token-{}", hash_token(token))) - .unwrap_or_else(|| "account".to_string()) -} - -fn normalized_token_field(value: Option<&str>) -> Option<&str> { - value.map(str::trim).filter(|value| !value.is_empty()) -} - -/// Compares one identity field; `None` means the field is not present on both -/// sides and the next field should decide. -fn field_identity(a: Option<&str>, b: Option<&str>) -> Option { - match (normalized_token_field(a), normalized_token_field(b)) { - (Some(a), Some(b)) => Some(a == b), - _ => None, - } -} - -fn same_token_identity(a: &Tokens, b: &Tokens) -> bool { - field_identity(a.account_id.as_deref(), b.account_id.as_deref()) - .or_else(|| field_identity(a.id_token.as_deref(), b.id_token.as_deref())) - .or_else(|| field_identity(a.access_token.as_deref(), b.access_token.as_deref())) - .unwrap_or(false) -} - -fn next_available_account_id(store: &CodexCredentialsStore, base_id: &str) -> String { - if !store.accounts.contains_key(base_id) { - return base_id.to_string(); - } - - for suffix in 2usize.. { - let candidate = format!("{base_id}-{suffix}"); - if !store.accounts.contains_key(&candidate) { - return candidate; - } - } - - unreachable!("unbounded suffix search must eventually find an unused Codex account id") -} - -fn validate_label_available( - store: &CodexCredentialsStore, - account_id: &str, - label: Option<&str>, -) -> Result<()> { - let Some(label) = label.map(str::trim).filter(|label| !label.is_empty()) else { - return Ok(()); - }; - let needle = label.to_lowercase(); - - for (id, account) in &store.accounts { - if id == account_id { - continue; - } - if account - .label - .as_deref() - .map(str::trim) - .map(str::to_lowercase) - .as_deref() - == Some(needle.as_str()) - { - anyhow::bail!("Codex account label already exists: {label}"); - } - } - - Ok(()) -} - -pub fn load_credentials_store() -> Option { - load_credentials_store_from_path(&codex_store_path()) -} - -fn load_credentials_store_from_path(path: &Path) -> Option { - load_credentials_store_for_update(path).ok().flatten() -} - -/// Loads the store while distinguishing "no usable store" (`Ok(None)`) from a -/// store written by a newer tokscale (`Err`). Write paths must propagate the -/// error instead of silently clobbering a future-version store; read paths can -/// treat both as "nothing usable". -fn load_credentials_store_for_update(path: &Path) -> Result> { - let Ok(content) = std::fs::read_to_string(path) else { - return Ok(None); - }; - bail_on_unknown_store_version(path, &content)?; - let Ok(mut store) = serde_json::from_str::(&content) else { - return Ok(None); - }; - - if store.accounts.is_empty() { - return Ok(None); - } - - if !store.accounts.contains_key(&store.active_account_id) { - if let Some(first_id) = first_account_id(&store) { - store.active_account_id = first_id; - let _ = save_credentials_store_at_path(path, &store); - } - } - - Ok(Some(store)) -} - -/// A future-version store may not even deserialize into the current struct, so -/// the version is checked on the raw JSON before the typed parse. -fn bail_on_unknown_store_version(path: &Path, content: &str) -> Result<()> { - let Ok(value) = serde_json::from_str::(content) else { - return Ok(()); - }; - let Some(version) = value.get("version").and_then(serde_json::Value::as_i64) else { - return Ok(()); - }; - if version != 1 { - anyhow::bail!( - "Unsupported Codex account store version {version} at {} (this tokscale supports version 1); refusing to modify it", - path.display() - ); - } - Ok(()) -} - -fn save_credentials_store(store: &CodexCredentialsStore) -> Result<()> { - save_credentials_store_at_path(&codex_store_path(), store) -} - -fn save_credentials_store_at_path(path: &Path, store: &CodexCredentialsStore) -> Result<()> { - let json = serde_json::to_string_pretty(store)?; - super::helpers::atomic_write_secret(path, json.as_bytes()) - .with_context(|| format!("Failed to write Codex account store to {}", path.display())) -} - -fn resolve_account_id(store: &CodexCredentialsStore, name_or_id: &str) -> Option { - let needle = name_or_id.trim(); - if needle.is_empty() { - return None; - } - - if store.accounts.contains_key(needle) { - return Some(needle.to_string()); - } - - let needle_lower = needle.to_lowercase(); - for (id, account) in &store.accounts { - if account - .label - .as_deref() - .map(str::trim) - .map(str::to_lowercase) - .as_deref() - == Some(needle_lower.as_str()) - { - return Some(id.clone()); - } - } - - None -} - -fn account_info( - store: &CodexCredentialsStore, - account_id: &str, - account: &CodexAccount, -) -> CodexAccountInfo { - CodexAccountInfo { - id: account_id.to_string(), - label: account.label.clone(), - account_id: account.tokens.account_id.clone(), - created_at: account.created_at.clone(), - is_active: account_id == store.active_account_id, - } -} - -/// Case-insensitive sort key shared by every place that orders accounts: -/// the label when present, falling back to the account id. -fn account_sort_key(label: Option<&str>, id: &str) -> String { - label.unwrap_or(id).to_lowercase() -} - -fn first_account_id(store: &CodexCredentialsStore) -> Option { - store - .accounts - .iter() - .min_by_key(|(id, account)| { - ( - account_sort_key(account.label.as_deref(), id), - (*id).clone(), - ) - }) - .map(|(id, _)| id.clone()) -} - -fn remove_account_from_store( - store: &mut CodexCredentialsStore, - name_or_id: &str, -) -> Result { - let resolved = resolve_account_id(store, name_or_id) - .ok_or_else(|| anyhow::anyhow!("Codex account not found: {name_or_id}"))?; - let removed_was_active = store.active_account_id == resolved; - let account = store - .accounts - .remove(&resolved) - .ok_or_else(|| anyhow::anyhow!("Codex account not found: {resolved}"))?; - let removed = CodexAccountInfo { - id: resolved, - label: account.label, - account_id: account.tokens.account_id.clone(), - created_at: account.created_at, - is_active: removed_was_active, - }; - - if removed_was_active { - if let Some(next_id) = first_account_id(store) { - store.active_account_id = next_id; - } else { - store.active_account_id.clear(); - } - } - - Ok(removed) -} - -pub fn list_accounts() -> Vec { - let store = match load_credentials_store() { - Some(store) => store, - None => return Vec::new(), - }; - - let mut accounts: Vec<_> = store - .accounts - .iter() - .map(|(id, account)| account_info(&store, id, account)) - .collect(); - - accounts.sort_by_key(|account| { - ( - !account.is_active, - account_sort_key(account.label.as_deref(), &account.id), - ) - }); - - accounts -} - -fn save_account_from_auth(auth: Auth, label: Option<&str>) -> Result { - save_account_from_auth_at_path(&codex_store_path(), auth, label, true) -} - -fn save_account_from_auth_at_path( - store_path: &Path, - auth: Auth, - label: Option<&str>, - make_active: bool, -) -> Result { - let tokens = auth - .tokens - .ok_or_else(|| anyhow::anyhow!("No Codex tokens."))?; - if tokens - .access_token - .as_deref() - .unwrap_or("") - .trim() - .is_empty() - { - anyhow::bail!("No Codex access token."); - } - - let base_account_id = derive_account_id(&tokens); - let mut store = - load_credentials_store_for_update(store_path)?.unwrap_or_else(|| CodexCredentialsStore { - version: 1, - active_account_id: base_account_id.clone(), - accounts: HashMap::new(), - }); - - // Scan every stored account (not just the base-id key) so an account that - // was stored under a collision-suffixed id (e.g. `acct_x-2`) is updated in - // place instead of re-importing as `acct_x-3`, `acct_x-4`, ... - let existing_identity_id = store - .accounts - .iter() - .find(|(_, existing)| same_token_identity(&existing.tokens, &tokens)) - .map(|(id, _)| id.clone()); - - if let Some(existing_id) = existing_identity_id { - validate_label_available(&store, &existing_id, label)?; - if let Some(existing) = store.accounts.get_mut(&existing_id) { - existing.tokens = tokens; - if let Some(label) = label.map(str::trim).filter(|s| !s.is_empty()) { - existing.label = Some(label.to_string()); + return Ok(auth); } } - if make_active { - store.active_account_id = existing_id.clone(); - } - save_credentials_store_at_path(store_path, &store)?; - - let account = store - .accounts - .get(&existing_id) - .ok_or_else(|| anyhow::anyhow!("Failed to save Codex account"))?; - return Ok(account_info(&store, &existing_id, account)); } - let account_id = if store.accounts.contains_key(&base_account_id) { - next_available_account_id(&store, &base_account_id) - } else { - base_account_id - }; - - validate_label_available(&store, &account_id, label)?; - - let account = CodexAccount { - tokens, - created_at: chrono::Utc::now().to_rfc3339(), - label: label - .map(str::trim) - .filter(|s| !s.is_empty()) - .map(str::to_string), - }; - - store.accounts.insert(account_id.clone(), account); - if make_active || store.active_account_id.trim().is_empty() { - store.active_account_id = account_id.clone(); - } - save_credentials_store_at_path(store_path, &store)?; - - let account = store - .accounts - .get(&account_id) - .ok_or_else(|| anyhow::anyhow!("Failed to save Codex account"))?; - Ok(account_info(&store, &account_id, account)) -} - -fn update_account_tokens(account_id: &str, tokens: Tokens) -> Result<()> { - let mut store = - load_credentials_store().ok_or_else(|| anyhow::anyhow!("No saved Codex accounts"))?; - let account = store - .accounts - .get_mut(account_id) - .ok_or_else(|| anyhow::anyhow!("Codex account not found: {account_id}"))?; - account.tokens = tokens; - save_credentials_store(&store) -} - -fn load_account(name_or_id: Option<&str>) -> Result<(String, CodexAccount, CodexAccountInfo)> { - let store = - load_credentials_store().ok_or_else(|| anyhow::anyhow!("No saved Codex accounts"))?; - let resolved = match name_or_id { - Some(name) => resolve_account_id(&store, name) - .ok_or_else(|| anyhow::anyhow!("Codex account not found: {name}"))?, - None => store.active_account_id.clone(), - }; - let account = store - .accounts - .get(&resolved) - .cloned() - .ok_or_else(|| anyhow::anyhow!("Codex account not found: {resolved}"))?; - let info = account_info(&store, &resolved, &account); - Ok((resolved, account, info)) -} - -fn auth_from_account(account: &CodexAccount) -> Auth { - Auth { - tokens: Some(account.tokens.clone()), - } + anyhow::bail!("No Codex credentials found. Run `codex login` to authenticate.") } pub fn has_credentials() -> bool { - if load_credentials_store() - .map(|store| !store.accounts.is_empty()) - .unwrap_or(false) - { - return true; - } - read_current_credentials().is_ok() } -async fn refresh_token(client: &reqwest::Client, rt: &str) -> Result { - let resp = client - .post("https://auth.openai.com/oauth/token") - .form(&[ - ("grant_type", "refresh_token"), - ("client_id", CLIENT_ID), - ("refresh_token", rt), - ]) - .send() - .await?; - if !resp.status().is_success() { - anyhow::bail!("Codex token refresh failed (HTTP {})", resp.status()); - } - Ok(resp.json().await?) -} - async fn fetch_usage( client: &reqwest::Client, token: &str, account_id: Option<&str>, ) -> Result { - let mut req = client + let mut request = client .get("https://chatgpt.com/backend-api/wham/usage") .header("Authorization", format!("Bearer {token}")) .header("Accept", "application/json") @@ -631,817 +112,147 @@ async fn fetch_usage( "User-Agent", "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7)", ); - if let Some(id) = account_id { - req = req.header("ChatGPT-Account-Id", id); + if let Some(account_id) = account_id { + request = request.header("ChatGPT-Account-Id", account_id); } - let resp = req.send().await?; - let status = resp.status(); + + let response = request.send().await?; + let status = response.status(); if status == reqwest::StatusCode::UNAUTHORIZED || status == reqwest::StatusCode::FORBIDDEN { - anyhow::bail!("NEEDS_AUTH"); + anyhow::bail!( + "Codex credentials were rejected. Run `codex login` to refresh the provider-owned authentication." + ); } if !status.is_success() { anyhow::bail!("Codex usage request failed (HTTP {status})"); } - let body = resp.text().await?; + + let body = response.text().await?; if body.trim().starts_with('<') { - anyhow::bail!("NEEDS_AUTH"); + anyhow::bail!( + "Codex usage returned an authentication page. Run `codex login` to refresh the provider-owned authentication." + ); } Ok(serde_json::from_str(&body)?) } fn metric_from_window(label: &str, window: &Window) -> UsageMetric { - let pct = window.used_percent.unwrap_or(0).clamp(0, 100) as f64; + let used_percent = window.used_percent.unwrap_or(0).clamp(0, 100) as f64; UsageMetric { label: label.into(), - used_percent: pct, - remaining_percent: 100.0 - pct, + used_percent, + remaining_percent: 100.0 - used_percent, remaining_label: None, resets_at: window .reset_at - .and_then(|ts| Utc.timestamp_opt(ts, 0).single()) - .map(|dt| dt.to_rfc3339()), + .and_then(|timestamp| Utc.timestamp_opt(timestamp, 0).single()) + .map(|date| date.to_rfc3339()), } } -async fn fetch_with_auth_async( - auth: Auth, - source: CredentialSource, - provider_name: String, - account: Option, -) -> Result { +async fn fetch_async(auth: Auth) -> Result { let tokens = auth .tokens .ok_or_else(|| anyhow::anyhow!("No Codex tokens."))?; let access_token = tokens .access_token - .clone() + .as_deref() + .map(str::trim) + .filter(|token| !token.is_empty()) .ok_or_else(|| anyhow::anyhow!("No Codex access token."))?; - let client = reqwest::Client::new(); - let resp = match fetch_usage(&client, &access_token, tokens.account_id.as_deref()).await { - Ok(r) => r, - Err(e) if e.to_string().contains("NEEDS_AUTH") => { - let rt_str = tokens - .refresh_token - .as_ref() - .ok_or_else(|| anyhow::anyhow!("No refresh token."))?; - let refreshed = refresh_token(&client, rt_str).await?; - let new = refreshed - .access_token - .clone() - .ok_or_else(|| anyhow::anyhow!("Refresh returned no token."))?; - - let mut updated_tokens = tokens.clone(); - updated_tokens.access_token = Some(new.clone()); - if let Some(new_rt) = refreshed.refresh_token { - updated_tokens.refresh_token = Some(new_rt); - } - persist_tokens(&source, &updated_tokens); + let response = fetch_usage( + &reqwest::Client::new(), + access_token, + tokens.account_id.as_deref(), + ) + .await?; - fetch_usage(&client, &new, updated_tokens.account_id.as_deref()).await? - } - Err(e) => return Err(e), - }; - - let plan = resp.plan_type.as_deref().map(capitalize); let mut metrics = Vec::new(); - if let Some(ref rl) = resp.rate_limit { - if let Some(ref w) = rl.primary_window { - metrics.push(metric_from_window("Session", w)); + if let Some(rate_limit) = &response.rate_limit { + if let Some(window) = &rate_limit.primary_window { + metrics.push(metric_from_window("Session", window)); } - if let Some(ref w) = rl.secondary_window { - metrics.push(metric_from_window("Weekly", w)); + if let Some(window) = &rate_limit.secondary_window { + metrics.push(metric_from_window("Weekly", window)); } } Ok(UsageOutput { - provider: provider_name, - account, - plan, - email: resp.email, + provider: "Codex".into(), + account: None, + plan: response.plan_type.as_deref().map(capitalize), + email: response.email, metrics, }) } -fn fetch_with_auth( - auth: Auth, - source: CredentialSource, - provider_name: String, - account: Option, -) -> Result { - let rt = tokio::runtime::Builder::new_current_thread() +pub fn fetch() -> Result { + let auth = read_current_credentials()?; + let runtime = tokio::runtime::Builder::new_current_thread() .enable_all() .build()?; - rt.block_on(fetch_with_auth_async(auth, source, provider_name, account)) -} - -pub fn fetch() -> Result { - let (auth, source) = read_current_credentials()?; - fetch_with_auth(auth, source, "Codex".into(), None) -} - -fn usage_account_from_saved( - store: &CodexCredentialsStore, - account_id: &str, - account: &CodexAccount, -) -> UsageAccount { - UsageAccount { - id: account_id.to_string(), - label: account.label.clone(), - is_active: account_id == store.active_account_id, - } -} - -pub fn fetch_all() -> Result> { - let Some(store) = load_credentials_store() else { - return fetch().map(|output| vec![output]); - }; - - if store.accounts.is_empty() { - return fetch().map(|output| vec![output]); - } - - let mut account_ids: Vec<_> = store.accounts.keys().cloned().collect(); - account_ids.sort_by_key(|id| { - ( - id != &store.active_account_id, - account_sort_key( - store - .accounts - .get(id) - .and_then(|account| account.label.as_deref()), - id, - ), - ) - }); - - let mut outputs = Vec::new(); - let mut first_error = None; - for account_id in account_ids { - let Some(account) = store.accounts.get(&account_id) else { - continue; - }; - let usage_account = usage_account_from_saved(&store, &account_id, account); - match fetch_with_auth( - auth_from_account(account), - CredentialSource::Store(account_id.clone()), - "Codex".into(), - Some(usage_account), - ) { - Ok(output) => outputs.push(output), - Err(e) if first_error.is_none() => first_error = Some(e), - Err(_) => {} - } - } - - if outputs.is_empty() { - if let Some(error) = first_error { - Err(error) - } else { - Ok(outputs) - } - } else { - Ok(outputs) - } -} - -fn fetch_saved_account(name_or_id: Option<&str>) -> Result<(CodexAccountInfo, UsageOutput)> { - let (account_id, account, info) = load_account(name_or_id)?; - let usage_account = UsageAccount { - id: info.id.clone(), - label: info.label.clone(), - is_active: info.is_active, - }; - let usage = fetch_with_auth( - auth_from_account(&account), - CredentialSource::Store(account_id), - "Codex".into(), - Some(usage_account), - )?; - Ok((info, usage)) -} - -pub fn import_current_account(label: Option<&str>) -> Result { - let (auth, _) = read_current_credentials()?; - save_account_from_auth(auth, label) -} - -pub fn switch_active_account(name_or_id: &str) -> Result { - let mut store = - load_credentials_store().ok_or_else(|| anyhow::anyhow!("No saved Codex accounts"))?; - let resolved = resolve_account_id(&store, name_or_id) - .ok_or_else(|| anyhow::anyhow!("Codex account not found: {name_or_id}"))?; - let account = store - .accounts - .get(&resolved) - .cloned() - .ok_or_else(|| anyhow::anyhow!("Codex account not found: {resolved}"))?; - - let path = auth_write_path()?; - save_auth_tokens(&path, &account.tokens)?; - - store.active_account_id = resolved.clone(); - save_credentials_store(&store)?; - - Ok(account_info(&store, &resolved, &account)) -} - -/// Removes an account from tokscale's store only. The codex CLI's own -/// `auth.json` is intentionally left untouched: rewriting it would silently -/// re-log the codex CLI into a different account (or log it out entirely). -pub fn remove_account(name_or_id: &str) -> Result { - let mut store = - load_credentials_store().ok_or_else(|| anyhow::anyhow!("No saved Codex accounts"))?; - let removed = remove_account_from_store(&mut store, name_or_id)?; - save_credentials_store(&store)?; - Ok(removed) -} - -pub fn run_codex_import(name: Option) -> Result<()> { - use colored::Colorize; - - let info = import_current_account(name.as_deref())?; - let display = info.label.as_deref().unwrap_or(&info.id); - - println!("\n {}\n", "Codex - Import".cyan()); - println!( - " {}", - format!("Imported Codex account {}", display.bold()).green() - ); - println!("{}", format!(" Account ID: {}", info.id).bright_black()); - println!(); - - Ok(()) -} - -pub fn run_codex_accounts(json: bool) -> Result<()> { - use colored::Colorize; - - let accounts = list_accounts(); - if json { - #[derive(Serialize)] - struct Output { - accounts: Vec, - } - println!("{}", serde_json::to_string_pretty(&Output { accounts })?); - return Ok(()); - } - - if accounts.is_empty() { - println!("\n {}\n", "No saved Codex accounts.".yellow()); - return Ok(()); - } - - println!("{}", "\n Codex - Accounts\n".cyan()); - for account in &accounts { - let name = if let Some(label) = &account.label { - format!("{} ({})", label, account.id) - } else { - account.id.clone() - }; - let marker = if account.is_active { "*" } else { "-" }; - let marker_colored = if account.is_active { - marker.green().to_string() - } else { - marker.bright_black().to_string() - }; - println!(" {} {}", marker_colored, name); - if let Some(account_id) = &account.account_id { - println!( - "{}", - format!(" Account ID: {}", account_id).bright_black() - ); - } - } - println!(); - - Ok(()) -} - -pub fn run_codex_switch(name: &str) -> Result<()> { - use colored::Colorize; - - let info = switch_active_account(name)?; - let display = info.label.as_deref().unwrap_or(&info.id); - - println!( - "\n {}\n", - format!("Active Codex account set to {}", display.bold()).green() - ); - - Ok(()) -} - -pub fn run_codex_remove(name: &str) -> Result<()> { - use colored::Colorize; - - let info = remove_account(name)?; - let display = info.label.as_deref().unwrap_or(&info.id); - - println!( - "\n {}", - format!("Stopped tracking Codex account {}", display.bold()).green() - ); - println!( - "{}\n", - " The codex CLI login was not changed.".bright_black() - ); - - Ok(()) -} - -pub fn run_codex_status(name: Option, json: bool) -> Result<()> { - use colored::Colorize; - - let result = if name.is_some() || load_credentials_store().is_some() { - fetch_saved_account(name.as_deref()).map(|(account, usage)| (Some(account), usage)) - } else { - fetch().map(|usage| (None, usage)) - }; - - if json { - #[derive(Serialize)] - struct Output { - #[serde(skip_serializing_if = "Option::is_none")] - account: Option, - #[serde(skip_serializing_if = "Option::is_none")] - usage: Option, - #[serde(skip_serializing_if = "Option::is_none")] - error: Option, - } - let output = match result { - Ok((account, usage)) => Output { - account, - usage: Some(usage), - error: None, - }, - Err(e) => Output { - account: None, - usage: None, - error: Some(e.to_string()), - }, - }; - println!("{}", serde_json::to_string_pretty(&output)?); - return Ok(()); - } - - println!("\n {}\n", "Codex - Status".cyan()); - match result { - Ok((account, usage)) => { - if let Some(account) = account { - let display = account.label.as_deref().unwrap_or(&account.id); - println!("{}", format!(" Account: {}", display).white()); - if let Some(account_id) = account.account_id { - println!("{}", format!(" Account ID: {}", account_id).bright_black()); - } - } - if let Some(email) = usage.email { - println!("{}", format!(" Email: {}", email).white()); - } - if let Some(plan) = usage.plan { - println!("{}", format!(" Plan: {}", plan).white()); - } - if usage.metrics.is_empty() { - println!("{}", " No quota metrics returned.".yellow()); - } else { - for metric in usage.metrics { - let remaining = metric - .remaining_label - .unwrap_or_else(|| format!("{:.0}% left", metric.remaining_percent)); - println!( - " {} {}", - format!("{:<10}", metric.label).bright_black(), - remaining - ); - } - } - } - Err(e) => { - println!(" {}", format!("Status failed: {e}").red()); - } - } - println!(); - - Ok(()) + runtime.block_on(fetch_async(auth)) } #[cfg(test)] mod tests { use super::*; - use tempfile::TempDir; - - fn test_store_path(tmp: &TempDir) -> PathBuf { - tmp.path().join("codex-credentials.json") - } - - fn tokens(access: &str, account_id: Option<&str>) -> Tokens { - Tokens { - access_token: Some(access.to_string()), - refresh_token: Some("refresh".to_string()), - account_id: account_id.map(str::to_string), - id_token: None, - } - } - - fn tokens_with_id_token(access: &str, account_id: Option<&str>, id_token: &str) -> Tokens { - Tokens { - access_token: Some(access.to_string()), - refresh_token: Some("refresh".to_string()), - account_id: account_id.map(str::to_string), - id_token: Some(id_token.to_string()), - } - } - - #[test] - fn derive_account_id_prefers_account_id() { - let tokens = tokens("access-token", Some("acct_work")); - assert_eq!(derive_account_id(&tokens), "acct_work"); - } - - #[test] - fn derive_account_id_falls_back_to_stable_token_hash() { - let id = derive_account_id(&tokens("access-token", None)); - assert!(id.starts_with("token-")); - assert_eq!(id, derive_account_id(&tokens("access-token", None))); - } - - #[test] - fn same_token_identity_prefers_account_id_over_rotating_id_token() { - let a = tokens_with_id_token("access-a", Some("acct_shared"), "id-token-a"); - let b = tokens_with_id_token("access-b", Some("acct_shared"), "id-token-b"); - - assert!(same_token_identity(&a, &b)); - } #[test] - fn load_credentials_store_repairs_missing_active_account() -> Result<()> { - let tmp = TempDir::new()?; - let mut accounts = HashMap::new(); - accounts.insert( - "acct_a".to_string(), - CodexAccount { - tokens: tokens("access-a", Some("acct_a")), - created_at: "2026-01-01T00:00:00Z".to_string(), - label: Some("zulu".to_string()), - }, - ); - accounts.insert( - "acct_b".to_string(), - CodexAccount { - tokens: tokens("access-b", Some("acct_b")), - created_at: "2026-01-02T00:00:00Z".to_string(), - label: Some("alpha".to_string()), - }, - ); - let store = CodexCredentialsStore { - version: 1, - active_account_id: "missing".to_string(), - accounts, - }; - let store_path = test_store_path(&tmp); - save_credentials_store_at_path(&store_path, &store)?; - - let loaded = load_credentials_store_from_path(&store_path).unwrap(); - assert_eq!(loaded.active_account_id, "acct_b"); - Ok(()) - } - - #[test] - fn resolve_account_id_matches_label_case_insensitively() { - let mut accounts = HashMap::new(); - accounts.insert( - "acct_a".to_string(), - CodexAccount { - tokens: tokens("access-a", Some("acct_a")), - created_at: "2026-01-01T00:00:00Z".to_string(), - label: Some("Work".to_string()), - }, - ); - let store = CodexCredentialsStore { - version: 1, - active_account_id: "acct_a".to_string(), - accounts, - }; + fn credential_paths_only_reference_provider_owned_auth() { + let home = Path::new("/home/tester"); + let paths = current_auth_paths_for_home(home, Some("/tmp/codex-home")); + assert_eq!(paths[0], PathBuf::from("/tmp/codex-home/auth.json")); assert_eq!( - resolve_account_id(&store, "work").as_deref(), - Some("acct_a") + paths[1], + PathBuf::from("/home/tester/.config/codex/auth.json") ); + assert_eq!(paths[2], PathBuf::from("/home/tester/.codex/auth.json")); + assert!(paths + .iter() + .all(|path| !path.ends_with("tokscale/codex-credentials.json"))); } #[test] - fn save_account_from_auth_at_path_imports_tokens_without_touching_real_home() -> Result<()> { - let tmp = TempDir::new()?; - let store_path = test_store_path(&tmp); - let info = save_account_from_auth_at_path( - &store_path, - Auth { - tokens: Some(tokens("access-a", Some("acct_a"))), - }, - Some("work"), - true, - )?; - - assert_eq!(info.id, "acct_a"); - assert_eq!(info.label.as_deref(), Some("work")); - assert!(info.is_active); - - let loaded = load_credentials_store_from_path(&store_path).unwrap(); - assert_eq!(loaded.active_account_id, "acct_a"); - assert!(loaded.accounts.contains_key("acct_a")); - Ok(()) - } - - #[test] - fn save_account_from_auth_at_path_preserves_label_when_updating_same_account() -> Result<()> { - let tmp = TempDir::new()?; - let store_path = test_store_path(&tmp); - save_account_from_auth_at_path( - &store_path, - Auth { - tokens: Some(tokens("access-a", Some("acct_a"))), - }, - Some("work"), - true, - )?; - - let info = save_account_from_auth_at_path( - &store_path, - Auth { - tokens: Some(tokens("access-b", Some("acct_a"))), - }, - None, - true, - )?; - - assert_eq!(info.id, "acct_a"); - assert_eq!(info.label.as_deref(), Some("work")); - - let loaded = load_credentials_store_from_path(&store_path).unwrap(); - assert_eq!(loaded.accounts.len(), 1); - let account = loaded.accounts.get("acct_a").unwrap(); - assert_eq!(account.label.as_deref(), Some("work")); - assert_eq!(account.tokens.access_token.as_deref(), Some("access-b")); - Ok(()) - } - - #[test] - fn save_account_from_auth_at_path_keeps_existing_account_on_identity_collision() -> Result<()> { - let tmp = TempDir::new()?; - let store_path = test_store_path(&tmp); - let mut accounts = HashMap::new(); - accounts.insert( - "acct_shared".to_string(), - CodexAccount { - tokens: tokens_with_id_token("access-a", Some("acct_other"), "id-token-a"), - created_at: "2026-01-01T00:00:00Z".to_string(), - label: Some("work".to_string()), - }, - ); - save_credentials_store_at_path( - &store_path, - &CodexCredentialsStore { - version: 1, - active_account_id: "acct_shared".to_string(), - accounts, - }, - )?; - - let info = save_account_from_auth_at_path( - &store_path, - Auth { - tokens: Some(tokens_with_id_token( - "access-b", - Some("acct_shared"), - "id-token-b", - )), - }, - None, - true, - )?; - - assert_eq!(info.id, "acct_shared-2"); - - let loaded = load_credentials_store_from_path(&store_path).unwrap(); - assert_eq!(loaded.accounts.len(), 2); - assert_eq!(loaded.active_account_id, "acct_shared-2"); + fn blank_codex_home_is_not_a_credential_root() { + let paths = current_auth_paths_for_home(Path::new("/home/tester"), Some(" ")); + assert_eq!(paths.len(), 2); assert_eq!( - loaded - .accounts - .get("acct_shared") - .and_then(|account| account.label.as_deref()), - Some("work") - ); - assert!(loaded.accounts.contains_key("acct_shared-2")); - Ok(()) - } - - #[test] - fn save_account_from_auth_at_path_can_add_without_changing_active_account() -> Result<()> { - let tmp = TempDir::new()?; - let store_path = test_store_path(&tmp); - save_account_from_auth_at_path( - &store_path, - Auth { - tokens: Some(tokens("access-a", Some("acct_a"))), - }, - Some("work"), - true, - )?; - - let info = save_account_from_auth_at_path( - &store_path, - Auth { - tokens: Some(tokens("access-b", Some("acct_b"))), - }, - Some("personal"), - false, - )?; - - assert_eq!(info.id, "acct_b"); - assert!(!info.is_active); - - let loaded = load_credentials_store_from_path(&store_path).unwrap(); - assert_eq!(loaded.active_account_id, "acct_a"); - assert!(loaded.accounts.contains_key("acct_a")); - assert!(loaded.accounts.contains_key("acct_b")); - Ok(()) - } - - #[test] - fn remove_account_from_store_keeps_active_when_removing_inactive() -> Result<()> { - let mut accounts = HashMap::new(); - accounts.insert( - "acct_a".to_string(), - CodexAccount { - tokens: tokens("access-a", Some("acct_a")), - created_at: "2026-01-01T00:00:00Z".to_string(), - label: Some("Work".to_string()), - }, - ); - accounts.insert( - "acct_b".to_string(), - CodexAccount { - tokens: tokens("access-b", Some("acct_b")), - created_at: "2026-01-02T00:00:00Z".to_string(), - label: Some("Personal".to_string()), - }, - ); - let mut store = CodexCredentialsStore { - version: 1, - active_account_id: "acct_a".to_string(), - accounts, - }; - - let removed = remove_account_from_store(&mut store, "personal")?; - - assert_eq!(removed.id, "acct_b"); - assert!(!removed.is_active); - assert_eq!(store.active_account_id, "acct_a"); - assert!(!store.accounts.contains_key("acct_b")); - Ok(()) - } - - #[test] - fn remove_account_from_store_selects_next_active_when_removing_active() -> Result<()> { - let mut accounts = HashMap::new(); - accounts.insert( - "acct_a".to_string(), - CodexAccount { - tokens: tokens("access-a", Some("acct_a")), - created_at: "2026-01-01T00:00:00Z".to_string(), - label: Some("Work".to_string()), - }, - ); - accounts.insert( - "acct_b".to_string(), - CodexAccount { - tokens: tokens("access-b", Some("acct_b")), - created_at: "2026-01-02T00:00:00Z".to_string(), - label: Some("Personal".to_string()), - }, + paths[0], + PathBuf::from("/home/tester/.config/codex/auth.json") ); - let mut store = CodexCredentialsStore { - version: 1, - active_account_id: "acct_a".to_string(), - accounts, - }; - - let removed = remove_account_from_store(&mut store, "work")?; - - assert_eq!(removed.id, "acct_a"); - assert!(removed.is_active); - assert_eq!(store.active_account_id, "acct_b"); - Ok(()) - } - - #[test] - fn remove_account_from_store_clears_active_when_last_account_removed() -> Result<()> { - let mut accounts = HashMap::new(); - accounts.insert( - "acct_a".to_string(), - CodexAccount { - tokens: tokens("access-a", Some("acct_a")), - created_at: "2026-01-01T00:00:00Z".to_string(), - label: Some("Work".to_string()), - }, - ); - let mut store = CodexCredentialsStore { - version: 1, - active_account_id: "acct_a".to_string(), - accounts, - }; - - let removed = remove_account_from_store(&mut store, "acct_a")?; - - assert_eq!(removed.id, "acct_a"); - assert!(removed.is_active); - assert!(store.accounts.is_empty()); - assert!(store.active_account_id.is_empty()); - Ok(()) } #[test] - fn save_account_from_auth_reuses_suffixed_account_with_same_identity() -> Result<()> { - let tmp = TempDir::new()?; - let store_path = tmp.path().join("codex-credentials.json"); - let mut accounts = HashMap::new(); - accounts.insert( - "acct_shared".to_string(), - CodexAccount { - tokens: tokens_with_id_token("access-a", Some("acct_other"), "id-token-a"), - created_at: "2026-01-01T00:00:00Z".to_string(), - label: Some("work".to_string()), - }, - ); - accounts.insert( - "acct_shared-2".to_string(), - CodexAccount { - tokens: tokens_with_id_token("access-b", Some("acct_shared"), "id-token-b"), - created_at: "2026-01-02T00:00:00Z".to_string(), - label: Some("personal".to_string()), - }, - ); - save_credentials_store_at_path( - &store_path, - &CodexCredentialsStore { - version: 1, - active_account_id: "acct_shared".to_string(), - accounts, - }, - )?; - - let info = save_account_from_auth_at_path( - &store_path, - Auth { - tokens: Some(tokens_with_id_token( - "access-c", - Some("acct_shared"), - "id-token-c", - )), - }, - None, - true, - )?; - - assert_eq!(info.id, "acct_shared-2"); + fn parse_auth_file_reads_only_the_fields_needed_for_usage() { + let temp = tempfile::TempDir::new().unwrap(); + let path = temp.path().join("auth.json"); + std::fs::write( + &path, + r#"{ + "tokens": { + "access_token": "access", + "refresh_token": "must-not-be-deserialized", + "id_token": "must-not-be-deserialized", + "account_id": "account" + } + }"#, + ) + .unwrap(); - let loaded = load_credentials_store_from_path(&store_path).unwrap(); - assert_eq!(loaded.accounts.len(), 2); - assert!(!loaded.accounts.contains_key("acct_shared-3")); - assert_eq!( - loaded - .accounts - .get("acct_shared-2") - .and_then(|account| account.tokens.access_token.as_deref()), - Some("access-c") - ); - Ok(()) + let auth = parse_auth_file(&path).unwrap().unwrap(); + let tokens = auth.tokens.unwrap(); + assert_eq!(tokens.access_token.as_deref(), Some("access")); + assert_eq!(tokens.account_id.as_deref(), Some("account")); } #[test] - fn save_account_from_auth_refuses_to_overwrite_future_store_version() -> Result<()> { - let tmp = TempDir::new()?; - let store_path = tmp.path().join("codex-credentials.json"); - let future_store = - r#"{"version":2,"vaults":[{"id":"acct_a","sealed":"0xdeadbeef"}],"accounts":{}}"#; - std::fs::write(&store_path, future_store)?; - - let result = save_account_from_auth_at_path( - &store_path, - Auth { - tokens: Some(tokens("access-a", Some("acct_a"))), - }, - None, - true, - ); - - let error = result.expect_err("future-version store must not be overwritten"); - assert!( - error.to_string().contains("version 2"), - "unexpected error: {error}" - ); - assert_eq!(std::fs::read_to_string(&store_path)?, future_store); - Ok(()) + fn parse_auth_file_rejects_blank_access_token() { + let temp = tempfile::TempDir::new().unwrap(); + let path = temp.path().join("auth.json"); + std::fs::write(&path, r#"{"tokens":{"access_token":" "}}"#).unwrap(); + assert!(parse_auth_file(&path).unwrap().is_none()); } } diff --git a/crates/tokscale-cli/src/commands/usage/mod.rs b/crates/tokscale-cli/src/commands/usage/mod.rs index 3afc68401..a92c8c94e 100644 --- a/crates/tokscale-cli/src/commands/usage/mod.rs +++ b/crates/tokscale-cli/src/commands/usage/mod.rs @@ -232,14 +232,12 @@ pub fn load_cache() -> Option> { #[derive(Clone, Copy)] enum Fetch { Single(fn() -> Result), - Multi(fn() -> Result>), } impl Fetch { fn call(self) -> Result> { match self { Fetch::Single(fetch) => fetch().map(|output| vec![output]), - Fetch::Multi(fetch) => fetch(), } } } @@ -267,7 +265,7 @@ fn all_providers() -> Vec { label: "Codex", is_available: codex::has_credentials, unavailable_message: "enabled in usageProviders but no Codex OAuth credentials were found", - fetch: Fetch::Multi(codex::fetch_all), + fetch: Fetch::Single(codex::fetch), }, UsageProvider { id: UsageProviderId::Zai, diff --git a/crates/tokscale-cli/src/commands/wrapped.rs b/crates/tokscale-cli/src/commands/wrapped.rs index 167c2dfa7..6fbb76712 100644 --- a/crates/tokscale-cli/src/commands/wrapped.rs +++ b/crates/tokscale-cli/src/commands/wrapped.rs @@ -1,4 +1,3 @@ -use crate::cursor; use ab_glyph::{point, Font, FontArc, GlyphId, PxScale, ScaleFont}; use anyhow::{Context, Result}; use chrono::{Datelike, Duration, Local, NaiveDate}; @@ -15,6 +14,8 @@ use tokscale_core::{ ReportOptions, ViewSet, }; +use crate::cli::WrappedRanking; + const SCALE: i32 = 2; const IMAGE_WIDTH: i32 = 1200 * SCALE; const IMAGE_HEIGHT: i32 = 1200 * SCALE; @@ -59,9 +60,10 @@ const COLOR_SISYPHUS: Rgba = Rgba([0x00, 0xCE, 0xD1, 0xFF]); pub struct WrappedOptions { pub output: Option, pub year: Option, + pub home_dir: Option, pub clients: Option>, pub short: bool, - pub include_agents: bool, + pub ranking: WrappedRanking, pub pin_sisyphus: bool, } @@ -117,10 +119,16 @@ struct FontSet { #[derive(Debug, Clone)] struct RenderOptions { short: bool, - include_agents: bool, + ranking: RenderRanking, pin_sisyphus: bool, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum RenderRanking { + Agents, + Clients, +} + pub fn run(options: WrappedOptions) -> Result { let rt = Runtime::new()?; rt.block_on(async move { generate_wrapped(options).await }) @@ -130,7 +138,6 @@ async fn generate_wrapped(options: WrappedOptions) -> Result { let data = load_wrapped_data(&options).await?; crate::commands::shared::emit_health_summary(&data.health); - let agents_requested = options.include_agents; let has_agent_data = data .top_agents .as_ref() @@ -140,17 +147,27 @@ async fn generate_wrapped(options: WrappedOptions) -> Result { .clients .as_ref() .is_none_or(|clients| clients.iter().any(|s| s == "opencode")); - let effective_include_agents = agents_requested && has_agent_data; + let render_ranking = select_render_ranking(options.ranking, has_agent_data); - if agents_requested && opencode_enabled && !has_agent_data { - println!( + if options.ranking == WrappedRanking::Auto && opencode_enabled && !has_agent_data { + eprintln!( + "{}", + format!("\n ⚠ No OpenCode agent data found for {}.", data.year).yellow() + ); + eprintln!( + "{}", + " Automatic ranking selected the clients view. Use --ranking clients to select it explicitly.\n" + .bright_black() + ); + } else if options.ranking == WrappedRanking::Agents && !has_agent_data { + eprintln!( "{}", format!("\n ⚠ No OpenCode agent data found for {}.", data.year).yellow() ); - println!("{}", " Falling back to clients view.".bright_black()); - println!( + eprintln!( "{}", - " Use --clients to always show clients view.\n".bright_black() + " Rendering the requested agents view with an explicit empty-state panel.\n" + .bright_black() ); } @@ -158,7 +175,7 @@ async fn generate_wrapped(options: WrappedOptions) -> Result { &data, &RenderOptions { short: options.short, - include_agents: effective_include_agents, + ranking: render_ranking, pin_sisyphus: options.pin_sisyphus, }, ) @@ -182,74 +199,31 @@ async fn generate_wrapped(options: WrappedOptions) -> Result { Ok(absolute.to_string_lossy().to_string()) } +fn select_render_ranking(ranking: WrappedRanking, has_agent_data: bool) -> RenderRanking { + match ranking { + WrappedRanking::Auto if has_agent_data => RenderRanking::Agents, + WrappedRanking::Auto => RenderRanking::Clients, + WrappedRanking::Agents => RenderRanking::Agents, + WrappedRanking::Clients => RenderRanking::Clients, + } +} + async fn load_wrapped_data(options: &WrappedOptions) -> Result { let year = options .year .clone() .unwrap_or_else(|| Local::now().year().to_string()); let clients = options.clients.clone().unwrap_or_else(default_clients); - let local_clients: Vec = clients - .iter() - .filter(|src| src.as_str() != ClientId::Cursor.as_str()) - .cloned() - .collect(); - let include_cursor = clients.iter().any(|src| src == ClientId::Cursor.as_str()); - let explicit_cursor = options - .clients - .as_ref() - .is_some_and(|sources| sources.iter().any(|src| src == ClientId::Cursor.as_str())); + let include_agent_view = options.ranking != WrappedRanking::Clients + && clients + .iter() + .any(|client| client == ClientId::OpenCode.as_str()); let since = format!("{}-01-01", year); let until = format!("{}-12-31", year); - let has_cursor_cache = cursor::has_cursor_usage_cache(); - let cursor_logged_in = cursor::is_cursor_logged_in(); - let mut cursor_sync_result: Option = None; - - if include_cursor && cursor_logged_in { - cursor_sync_result = Some(cursor::sync_cursor_cache().await); - } - - if let Some(sync) = cursor_sync_result.as_ref() { - if let Some(error) = sync.error.as_ref() { - if sync.synced || has_cursor_cache { - let prefix = if sync.synced { - "Cursor sync warning" - } else { - "Cursor sync failed; using cached data" - }; - println!("{}", format!(" {}: {}", prefix, error).yellow()); - } - } - } - - let include_cursor_in_graph = if include_cursor { - let synced = cursor_sync_result - .as_ref() - .map(|sync| sync.synced) - .unwrap_or(false); - synced || has_cursor_cache - } else { - false - }; - if let Some(warning) = - cursor_setup_warning_for_wrapped(explicit_cursor, include_cursor_in_graph, cursor_logged_in) - { - eprintln!("{}", format!(" Warning: {warning}").yellow()); - } - - let graph_clients = if include_cursor && !include_cursor_in_graph { - clients - .iter() - .filter(|src| src.as_str() != ClientId::Cursor.as_str()) - .cloned() - .collect::>() - } else { - clients.clone() - }; - let mut views = ViewSet::GRAPH | ViewSet::TIME_METRICS; - if options.include_agents && !local_clients.is_empty() { + if include_agent_view { views |= ViewSet::AGENTS; } @@ -258,19 +232,21 @@ async fn load_wrapped_data(options: &WrappedOptions) -> Result { .map_err(anyhow::Error::msg)?; let aggregated = load_aggregated_views_with_pricing( &ReportOptions { - home_dir: None, - use_env_roots: true, - clients: Some(graph_clients), + home_dir: options.home_dir.clone(), + use_env_roots: crate::commands::shared::use_env_roots(&options.home_dir), + clients: Some(clients), since: Some(since), until: Some(until), year: Some(year.clone()), group_by: GroupBy::default(), - scanner_settings: crate::tui::settings::load_scanner_settings()?, + scanner_settings: crate::tui::settings::load_scanner_settings_for_home( + &options.home_dir, + )?, }, views, Some(pricing.as_ref()), ) - .map_err(anyhow::Error::msg)?; + .map_err(anyhow::Error::new)?; let health = wrapped_health_report(&aggregated); let graph = aggregated.graph.expect("graph view requested"); @@ -294,7 +270,7 @@ async fn load_wrapped_data(options: &WrappedOptions) -> Result { top_clients.sort_by(|a, b| b.cost.partial_cmp(&a.cost).unwrap_or(Ordering::Equal)); top_clients.truncate(3); - let top_agents = if options.include_agents { + let top_agents = if include_agent_view { aggregated .agent_usage .as_ref() @@ -583,7 +559,7 @@ async fn generate_wrapped_image(data: &WrappedData, options: &RenderOptions) -> } y_pos += 40 * SCALE; - if options.include_agents { + if options.ranking == RenderRanking::Agents { draw_text_mut_baseline( &mut canvas, &fonts.regular, @@ -596,6 +572,17 @@ async fn generate_wrapped_image(data: &WrappedData, options: &RenderOptions) -> y_pos += 48 * SCALE; let agents = data.top_agents.clone().unwrap_or_default(); + if agents.is_empty() { + draw_text_mut_baseline( + &mut canvas, + &fonts.regular, + (28 * SCALE) as f32, + COLOR_TEXT_SECONDARY, + PADDING, + y_pos, + "No OpenCode agent data", + ); + } let mut rank_index = 1; for agent in agents { @@ -1717,7 +1704,6 @@ fn capitalize_word(word: &str) -> String { fn default_clients() -> Vec { ClientId::iter() - .filter(|client| client.parse_local()) .map(|client| client.as_str().to_string()) .collect() } @@ -1730,6 +1716,30 @@ mod tests { use tempfile::TempDir; use tokscale_core::{DataHealth, RejectionSummary, SourceFailure, SourceHealth, SourceStatus}; + #[test] + fn automatic_ranking_uses_agents_only_when_agent_data_exists() { + assert_eq!( + select_render_ranking(WrappedRanking::Auto, true), + RenderRanking::Agents + ); + assert_eq!( + select_render_ranking(WrappedRanking::Auto, false), + RenderRanking::Clients + ); + } + + #[test] + fn explicit_agents_ranking_never_changes_to_clients() { + assert_eq!( + select_render_ranking(WrappedRanking::Agents, false), + RenderRanking::Agents + ); + assert_eq!( + select_render_ranking(WrappedRanking::Clients, true), + RenderRanking::Clients + ); + } + fn restore_env_var(key: &str, value: Option) { unsafe { match value { @@ -2327,17 +2337,15 @@ mod tests { } #[test] - fn default_clients_use_local_parse_policy_from_catalog() { + fn default_clients_use_the_complete_catalog() { let clients = default_clients(); let expected = ClientId::iter() - .filter(|client| client.parse_local()) .map(|client| client.as_str().to_string()) .collect::>(); assert_eq!(clients, expected); assert!(clients.iter().any(|client| client == "grok")); assert!(clients.iter().any(|client| client == "kiro")); - assert!(clients.iter().any(|client| client == "trae")); assert!(clients.iter().any(|client| client == "warp")); } @@ -2742,48 +2750,3 @@ mod tests { assert_eq!(longest, 4); } } - -fn cursor_setup_warning_for_wrapped( - explicit_cursor: bool, - include_cursor_in_graph: bool, - cursor_logged_in: bool, -) -> Option { - if !explicit_cursor || include_cursor_in_graph { - return None; - } - - let action = if cursor_logged_in { - "run `tokscale cursor sync --json`" - } else { - "run `tokscale cursor login` and `tokscale cursor sync --json`" - }; - - Some(format!( - "Cursor usage requires Tokscale's Cursor API cache at `~/.config/tokscale/cursor-cache/usage*.csv`; {action}. Tokscale does not parse local `~/.cursor` session data." - )) -} - -#[cfg(test)] -mod cursor_setup_warning_tests { - use super::cursor_setup_warning_for_wrapped; - - #[test] - fn wrapped_cursor_warning_suggests_login_when_not_authenticated() { - let warning = cursor_setup_warning_for_wrapped(true, false, false).unwrap(); - assert!(warning.contains("tokscale cursor login")); - assert!(warning.contains("tokscale cursor sync --json")); - } - - #[test] - fn wrapped_cursor_warning_suggests_sync_only_when_authenticated() { - let warning = cursor_setup_warning_for_wrapped(true, false, true).unwrap(); - assert!(!warning.contains("tokscale cursor login")); - assert!(warning.contains("tokscale cursor sync --json")); - } - - #[test] - fn wrapped_cursor_warning_is_suppressed_without_explicit_missing_cursor() { - assert!(cursor_setup_warning_for_wrapped(false, false, false).is_none()); - assert!(cursor_setup_warning_for_wrapped(true, true, false).is_none()); - } -} diff --git a/crates/tokscale-cli/src/cursor.rs b/crates/tokscale-cli/src/cursor.rs deleted file mode 100644 index 47755720b..000000000 --- a/crates/tokscale-cli/src/cursor.rs +++ /dev/null @@ -1,2016 +0,0 @@ -use anyhow::{Context, Result}; -use serde::{Deserialize, Serialize}; -use sha2::{Digest, Sha256}; -use std::collections::HashMap; -use std::fs; -use std::io::Write; -use std::path::{Path, PathBuf}; -use std::time::{Duration, SystemTime}; - -/// Timeout for every Cursor HTTP request. Picked to bound the worst case for -/// auto-sync (which runs synchronously before local reports and the TUI) while -/// still tolerating routine API latency. If the network is hung, the report -/// proceeds against cached data after this timeout instead of stalling forever. -const CURSOR_HTTP_TIMEOUT: Duration = Duration::from_secs(8); - -/// Skip implicit pre-report sync when every expected Cursor account cache file -/// was modified within this window. Prevents `tokscale models` (and its -/// siblings) from issuing a Cursor API call on every invocation. The manual -/// `tokscale cursor sync` command bypasses this — explicit user intent is -/// always honored. -pub const CURSOR_AUTO_SYNC_FRESHNESS: Duration = Duration::from_secs(5 * 60); - -fn build_cursor_http_client() -> Result { - reqwest::Client::builder() - .timeout(CURSOR_HTTP_TIMEOUT) - .build() - .context("Failed to build Cursor HTTP client") -} - -fn home_dir() -> Result { - dirs::home_dir().context("Could not determine home directory") -} - -fn cursor_credentials_path(home_dir: &Path) -> PathBuf { - home_dir.join(".config/tokscale/cursor-credentials.json") -} - -fn old_cursor_credentials_path(home_dir: &Path) -> PathBuf { - home_dir.join(".tokscale/cursor-credentials.json") -} - -fn cursor_cache_dir(home_dir: &Path) -> PathBuf { - home_dir.join(".config/tokscale/cursor-cache") -} - -fn old_cursor_cache_dir(home_dir: &Path) -> PathBuf { - home_dir.join(".tokscale/cursor-cache") -} - -const USAGE_CSV_ENDPOINT: &str = - "https://cursor.com/api/dashboard/export-usage-events-csv?strategy=tokens"; -const USAGE_SUMMARY_ENDPOINT: &str = "https://cursor.com/api/usage-summary"; - -/// Marker file touched at the end of every `sync_cursor_cache` run (even when -/// some accounts fail). Its mtime gates secondary-account freshness checks so -/// a permanently-stale secondary (expired token, removed account, network -/// partition) does not force an implicit sync on every invocation. -const CURSOR_SYNC_ATTEMPT_MARKER: &str = "usage.last-sync-attempt"; - -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct CursorCredentials { - #[serde(rename = "sessionToken")] - pub session_token: String, - #[serde(rename = "userId", skip_serializing_if = "Option::is_none")] - pub user_id: Option, - #[serde(rename = "createdAt")] - pub created_at: String, - #[serde(rename = "expiresAt", skip_serializing_if = "Option::is_none")] - pub expires_at: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub label: Option, -} - -#[derive(Debug, Serialize, Deserialize)] -pub struct CursorCredentialsStore { - pub version: i32, - #[serde(rename = "activeAccountId")] - pub active_account_id: String, - pub accounts: HashMap, -} - -#[derive(Debug, Serialize)] -pub struct AccountInfo { - pub id: String, - #[serde(skip_serializing_if = "Option::is_none")] - pub label: Option, - #[serde(rename = "userId", skip_serializing_if = "Option::is_none")] - pub user_id: Option, - #[serde(rename = "createdAt")] - pub created_at: String, - #[serde(rename = "isActive")] - pub is_active: bool, -} - -#[derive(Debug, Serialize)] -pub struct SyncCursorResult { - pub synced: bool, - pub rows: usize, - pub error: Option, -} - -pub fn get_cursor_credentials_path() -> Result { - Ok(cursor_credentials_path(&home_dir()?)) -} - -pub fn get_cursor_cache_dir() -> Result { - Ok(cursor_cache_dir(&home_dir()?)) -} - -fn migrate_cache_dir_from_old_path_in_home(home_dir: &Path) { - let old_dir = old_cursor_cache_dir(home_dir); - let new_dir = cursor_cache_dir(home_dir); - if !new_dir.exists() - && old_dir.exists() - && fs::create_dir_all(&new_dir).is_ok() - && copy_dir_recursive(&old_dir, &new_dir).is_ok() - { - let _ = fs::remove_dir_all(&old_dir); - } -} - -fn copy_dir_recursive(src: &std::path::Path, dst: &std::path::Path) -> Result<()> { - for entry in fs::read_dir(src)? { - let entry = entry?; - let path = entry.path(); - let target = dst.join(entry.file_name()); - if path.is_dir() { - fs::create_dir_all(&target)?; - copy_dir_recursive(&path, &target)?; - } else { - fs::copy(&path, &target)?; - } - } - Ok(()) -} - -fn build_cursor_headers(session_token: &str) -> reqwest::header::HeaderMap { - use reqwest::header::HeaderValue; - - let mut headers = reqwest::header::HeaderMap::new(); - headers.insert("Accept", HeaderValue::from_static("*/*")); - headers.insert( - "Accept-Language", - HeaderValue::from_static("en-US,en;q=0.9"), - ); - if let Ok(cookie) = format!("WorkosCursorSessionToken={}", session_token).parse() { - headers.insert("Cookie", cookie); - } - headers.insert( - "Referer", - HeaderValue::from_static("https://www.cursor.com/settings"), - ); - headers.insert( - "User-Agent", - HeaderValue::from_static("Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36"), - ); - headers -} - -fn count_cursor_csv_rows(csv_text: &str) -> usize { - let mut reader = csv::ReaderBuilder::new() - .has_headers(true) - .flexible(true) - .from_reader(csv_text.as_bytes()); - reader.records().filter_map(|r| r.ok()).count() -} - -fn atomic_write_file(path: &std::path::Path, contents: &str) -> Result<()> { - tokscale_core::fs_atomic::write_atomic(path, contents.as_bytes()) - .with_context(|| format!("Failed to persist Cursor cache file: {}", path.display())) -} - -fn ensure_config_dir_in_home(home_dir: &Path) -> Result<()> { - let config_dir = home_dir.join(".config/tokscale"); - - if !config_dir.exists() { - fs::create_dir_all(&config_dir)?; - #[cfg(unix)] - { - use std::os::unix::fs::PermissionsExt; - fs::set_permissions(&config_dir, fs::Permissions::from_mode(0o700))?; - } - } - Ok(()) -} - -fn extract_user_id_from_session_token(token: &str) -> Option { - let token = token.trim(); - if token.contains("%3A%3A") { - let user_id = token.split("%3A%3A").next()?.trim(); - if user_id.is_empty() { - return None; - } - return Some(user_id.to_string()); - } - if token.contains("::") { - let user_id = token.split("::").next()?.trim(); - if user_id.is_empty() { - return None; - } - return Some(user_id.to_string()); - } - None -} - -fn derive_account_id(session_token: &str) -> String { - if let Some(user_id) = extract_user_id_from_session_token(session_token) { - return user_id; - } - let mut hasher = Sha256::new(); - hasher.update(session_token.as_bytes()); - let hash = hasher.finalize(); - let hex = format!("{:x}", hash); - format!("anon-{}", &hex[..12]) -} - -fn sanitize_account_id_for_filename(account_id: &str) -> String { - let sanitized: String = account_id - .trim() - .to_lowercase() - .chars() - .map(|c| { - if c.is_ascii_alphanumeric() || c == '.' || c == '_' || c == '-' { - c - } else { - '-' - } - }) - .collect(); - let trimmed = sanitized.trim_matches('-'); - let result = if trimmed.len() > 80 { - &trimmed[..80] - } else { - trimmed - }; - if result.is_empty() { - "account".to_string() - } else { - result.to_string() - } -} - -pub fn load_credentials_store() -> Option { - let home_dir = home_dir().ok()?; - load_credentials_store_from_home(&home_dir) -} - -fn load_credentials_store_from_home(home_dir: &Path) -> Option { - let path = cursor_credentials_path(home_dir); - let old_path = old_cursor_credentials_path(home_dir); - let read_path = if path.exists() { - path.clone() - } else if old_path.exists() { - old_path - } else { - return None; - }; - - let content = fs::read_to_string(&read_path).ok()?; - - if let Ok(mut store) = serde_json::from_str::(&content) { - if store.version == 1 && !store.accounts.is_empty() { - let mut changed = false; - if !store.accounts.contains_key(&store.active_account_id) { - if let Some(first_id) = store.accounts.keys().next().cloned() { - store.active_account_id = first_id; - changed = true; - } - } - if changed || read_path != path { - let _ = save_credentials_store_in_home(home_dir, &store); - } - if read_path != path { - let _ = fs::remove_file(old_cursor_credentials_path(home_dir)); - } - return Some(store); - } - } - - if let Ok(single) = serde_json::from_str::(&content) { - let account_id = derive_account_id(&single.session_token); - let mut accounts = HashMap::new(); - accounts.insert(account_id.clone(), single); - let migrated = CursorCredentialsStore { - version: 1, - active_account_id: account_id, - accounts, - }; - - let _ = save_credentials_store_in_home(home_dir, &migrated); - if read_path != path { - let _ = fs::remove_file(old_cursor_credentials_path(home_dir)); - } - return Some(migrated); - } - - None -} - -pub fn save_credentials_store(store: &CursorCredentialsStore) -> Result<()> { - save_credentials_store_in_home(&home_dir()?, store) -} - -fn save_credentials_store_in_home(home_dir: &Path, store: &CursorCredentialsStore) -> Result<()> { - ensure_config_dir_in_home(home_dir)?; - let path = cursor_credentials_path(home_dir); - let json = serde_json::to_string_pretty(store)?; - atomic_write_file(&path, &json)?; - - #[cfg(unix)] - { - use std::os::unix::fs::PermissionsExt; - fs::set_permissions(&path, fs::Permissions::from_mode(0o600))?; - } - - Ok(()) -} - -fn resolve_account_id(store: &CursorCredentialsStore, name_or_id: &str) -> Option { - let needle = name_or_id.trim(); - if needle.is_empty() { - return None; - } - - if store.accounts.contains_key(needle) { - return Some(needle.to_string()); - } - - let needle_lower = needle.to_lowercase(); - for (id, acct) in &store.accounts { - if let Some(label) = &acct.label { - if label.to_lowercase() == needle_lower { - return Some(id.clone()); - } - } - } - - None -} - -pub fn list_accounts() -> Vec { - let store = match load_credentials_store() { - Some(s) => s, - None => return vec![], - }; - - let mut accounts: Vec = store - .accounts - .iter() - .map(|(id, acct)| AccountInfo { - id: id.clone(), - label: acct.label.clone(), - user_id: acct.user_id.clone(), - created_at: acct.created_at.clone(), - is_active: id == &store.active_account_id, - }) - .collect(); - - accounts.sort_by(|a, b| { - if a.is_active != b.is_active { - return if a.is_active { - std::cmp::Ordering::Less - } else { - std::cmp::Ordering::Greater - }; - } - let la = a.label.as_deref().unwrap_or(&a.id).to_lowercase(); - let lb = b.label.as_deref().unwrap_or(&b.id).to_lowercase(); - la.cmp(&lb) - }); - - accounts -} - -pub fn find_account(name_or_id: &str) -> Option { - let store = load_credentials_store()?; - let resolved = resolve_account_id(&store, name_or_id)?; - let acct = store.accounts.get(&resolved)?; - - Some(AccountInfo { - id: resolved.clone(), - label: acct.label.clone(), - user_id: acct.user_id.clone(), - created_at: acct.created_at.clone(), - is_active: resolved == store.active_account_id, - }) -} - -pub fn save_credentials(token: &str, label: Option<&str>) -> Result { - let account_id = derive_account_id(token); - let user_id = extract_user_id_from_session_token(token); - - let mut store = load_credentials_store().unwrap_or_else(|| CursorCredentialsStore { - version: 1, - active_account_id: account_id.clone(), - accounts: HashMap::new(), - }); - - if let Some(lbl) = label { - let needle = lbl.trim().to_lowercase(); - if !needle.is_empty() { - for (id, acct) in &store.accounts { - if id == &account_id { - continue; - } - if let Some(existing_label) = &acct.label { - if existing_label.trim().to_lowercase() == needle { - anyhow::bail!("Cursor account label already exists: {}", lbl); - } - } - } - } - } - - let credentials = CursorCredentials { - session_token: token.to_string(), - user_id, - created_at: chrono::Utc::now().to_rfc3339(), - expires_at: None, - label: label.map(|s| s.to_string()), - }; - - store.accounts.insert(account_id.clone(), credentials); - store.active_account_id = account_id.clone(); - - save_credentials_store(&store)?; - - Ok(account_id) -} - -pub fn remove_account(name_or_id: &str, purge_cache: bool) -> Result<()> { - let mut store = - load_credentials_store().ok_or_else(|| anyhow::anyhow!("No saved Cursor accounts"))?; - - let resolved = resolve_account_id(&store, name_or_id) - .ok_or_else(|| anyhow::anyhow!("Account not found: {}", name_or_id))?; - - let was_active = resolved == store.active_account_id; - - let cache_dir = get_cursor_cache_dir()?; - if cache_dir.exists() { - let per_account = cache_dir.join(format!( - "usage.{}.csv", - sanitize_account_id_for_filename(&resolved) - )); - if per_account.exists() { - if purge_cache { - let _ = fs::remove_file(&per_account); - } else { - let _ = archive_cache_file(&per_account, &format!("usage.{}", resolved)); - } - } - if was_active { - let active_file = cache_dir.join("usage.csv"); - if active_file.exists() { - if purge_cache { - let _ = fs::remove_file(&active_file); - } else { - let _ = archive_cache_file(&active_file, &format!("usage.active.{}", resolved)); - } - } - } - } - - store.accounts.remove(&resolved); - - if store.accounts.is_empty() { - let path = get_cursor_credentials_path()?; - if path.exists() { - fs::remove_file(path)?; - } - return Ok(()); - } - - if was_active { - if let Some(first_id) = store.accounts.keys().next().cloned() { - let new_account_file = cache_dir.join(format!( - "usage.{}.csv", - sanitize_account_id_for_filename(&first_id) - )); - let active_file = cache_dir.join("usage.csv"); - if new_account_file.exists() { - let _ = fs::rename(&new_account_file, &active_file); - } - store.active_account_id = first_id; - } - } - - save_credentials_store(&store)?; - Ok(()) -} - -pub fn remove_all_accounts(purge_cache: bool) -> Result<()> { - let cache_dir = get_cursor_cache_dir()?; - if cache_dir.exists() { - if let Ok(entries) = fs::read_dir(&cache_dir) { - for entry in entries.flatten() { - let name = entry.file_name().to_string_lossy().to_string(); - if name.starts_with("usage") && name.ends_with(".csv") { - if purge_cache { - let _ = fs::remove_file(entry.path()); - } else { - let _ = archive_cache_file(&entry.path(), &format!("usage.all.{}", name)); - } - } - } - } - } - - let path = get_cursor_credentials_path()?; - if path.exists() { - fs::remove_file(path)?; - } - Ok(()) -} - -pub fn set_active_account(name_or_id: &str) -> Result<()> { - let mut store = - load_credentials_store().ok_or_else(|| anyhow::anyhow!("No saved Cursor accounts"))?; - - let resolved = resolve_account_id(&store, name_or_id) - .ok_or_else(|| anyhow::anyhow!("Account not found: {}", name_or_id))?; - - let old_active_id = store.active_account_id.clone(); - - if resolved != old_active_id { - let _ = reconcile_cache_files(&old_active_id, &resolved); - } - - store.active_account_id = resolved; - save_credentials_store(&store)?; - - Ok(()) -} - -fn reconcile_cache_files(old_account_id: &str, new_account_id: &str) -> Result<()> { - let cache_dir = get_cursor_cache_dir()?; - if !cache_dir.exists() { - return Ok(()); - } - - let active_file = cache_dir.join("usage.csv"); - let old_account_file = cache_dir.join(format!( - "usage.{}.csv", - sanitize_account_id_for_filename(old_account_id) - )); - let new_account_file = cache_dir.join(format!( - "usage.{}.csv", - sanitize_account_id_for_filename(new_account_id) - )); - - if active_file.exists() { - if old_account_file.exists() { - let _ = archive_cache_file(&old_account_file, old_account_id); - } - fs::rename(&active_file, &old_account_file)?; - } - - if new_account_file.exists() { - if active_file.exists() { - let _ = archive_cache_file(&active_file, "usage.active"); - } - fs::rename(&new_account_file, &active_file)?; - } - - Ok(()) -} - -pub fn load_active_credentials() -> Option { - let store = load_credentials_store()?; - store.accounts.get(&store.active_account_id).cloned() -} - -pub fn has_active_credentials_in_home(home_dir: &Path) -> bool { - load_credentials_store_from_home(home_dir) - .and_then(|store| store.accounts.get(&store.active_account_id).cloned()) - .is_some() -} - -fn is_cursor_usage_csv_filename(name: &str) -> bool { - if name == "usage.csv" { - return true; - } - if !name.starts_with("usage.") || !name.ends_with(".csv") { - return false; - } - if name.starts_with("usage.backup") { - return false; - } - let stem = name.trim_start_matches("usage.").trim_end_matches(".csv"); - !stem.is_empty() - && stem - .chars() - .all(|c| c.is_ascii_alphanumeric() || c == '.' || c == '_' || c == '-') -} - -pub fn has_cursor_usage_cache_in_home(home_dir: &Path) -> bool { - migrate_cache_dir_from_old_path_in_home(home_dir); - let cache_dir = cursor_cache_dir(home_dir); - if !cache_dir.exists() { - return false; - } - - match fs::read_dir(cache_dir) { - Ok(entries) => entries - .filter_map(|entry| entry.ok()) - .filter_map(|entry| entry.file_name().into_string().ok()) - .any(|name| is_cursor_usage_csv_filename(&name)), - Err(_) => false, - } -} - -pub fn has_cursor_usage_cache() -> bool { - let home_dir = match home_dir() { - Ok(home_dir) => home_dir, - Err(_) => return false, - }; - has_cursor_usage_cache_in_home(&home_dir) -} - -fn expected_cursor_usage_cache_paths_in(home_dir: &Path) -> Vec { - let cache_dir = cursor_cache_dir(home_dir); - - if let Some(store) = load_credentials_store_from_home(home_dir) { - if !store.accounts.is_empty() { - let mut paths = store - .accounts - .keys() - .map(|account_id| { - if account_id == &store.active_account_id { - cache_dir.join("usage.csv") - } else { - cache_dir.join(format!( - "usage.{}.csv", - sanitize_account_id_for_filename(account_id) - )) - } - }) - .collect::>(); - paths.sort_unstable(); - paths.dedup(); - return paths; - } - } - - vec![cache_dir.join("usage.csv")] -} - -fn cursor_usage_cache_file_is_fresh(path: &Path, max_age: Duration) -> bool { - let Ok(mtime) = path.metadata().and_then(|meta| meta.modified()) else { - return false; - }; - match SystemTime::now().duration_since(mtime) { - Ok(age) => age < max_age, - // mtime is in the future (clock skew) — treat as fresh; a clock-skew - // cache is no less authoritative than a freshly-fetched one, and we'd - // rather not thrash the API while the system clock recovers. - Err(_) => true, - } -} - -fn cursor_usage_cache_is_fresh_in(home_dir: &Path, max_age: Duration) -> bool { - let cache_dir = cursor_cache_dir(home_dir); - if !cache_dir.exists() { - return false; - } - - // The active account's cache is non-negotiable: if it is stale or missing, - // implicit sync must run so reports read current data. - let active_path = cache_dir.join("usage.csv"); - if !cursor_usage_cache_file_is_fresh(&active_path, max_age) { - return false; - } - - // For secondaries, a fresh sync-attempt marker is sufficient. This avoids - // forcing a sync on every invocation when a secondary account is - // permanently stale (expired token, removed account, persistent API - // failure). Without the marker, `.all(...)` would return `false` forever. - let marker_fresh = - cursor_usage_cache_file_is_fresh(&cache_dir.join(CURSOR_SYNC_ATTEMPT_MARKER), max_age); - - expected_cursor_usage_cache_paths_in(home_dir) - .iter() - .filter(|p| *p != &active_path) - .all(|p| cursor_usage_cache_file_is_fresh(p, max_age) || marker_fresh) -} - -/// True when the active cursor usage cache (`usage.csv`) was refreshed within -/// `max_age` AND every secondary account cache is either fresh or a recent -/// sync-attempt marker exists. The active cache is unconditionally required — -/// a stale active means reports would show out-of-date data. Secondaries are -/// best-effort: when a secondary is permanently stale (expired token, removed -/// account, persistent API failure) the marker short-circuits the check so we -/// don't force an implicit sync on every invocation. Used by the implicit -/// pre-report sync path to avoid hitting the Cursor API on every invocation. -/// The manual `tokscale cursor sync` CLI bypasses this — explicit user intent -/// is always honored. -pub fn cursor_usage_cache_is_fresh(max_age: Duration) -> bool { - let Ok(home_dir) = home_dir() else { - return false; - }; - cursor_usage_cache_is_fresh_in(&home_dir, max_age) -} - -pub fn is_cursor_logged_in() -> bool { - load_active_credentials().is_some() -} - -pub fn load_credentials_for(name_or_id: &str) -> Option { - let store = load_credentials_store()?; - let resolved = resolve_account_id(&store, name_or_id)?; - store.accounts.get(&resolved).cloned() -} - -#[derive(Debug)] -pub struct ValidateSessionResult { - pub valid: bool, - pub membership_type: Option, - pub error: Option, -} - -pub async fn validate_cursor_session(token: &str) -> ValidateSessionResult { - let client = match build_cursor_http_client() { - Ok(client) => client, - Err(e) => { - return ValidateSessionResult { - valid: false, - membership_type: None, - error: Some(format!("Failed to build HTTP client: {}", e)), - }; - } - }; - let response = match client - .get(USAGE_SUMMARY_ENDPOINT) - .headers(build_cursor_headers(token)) - .send() - .await - { - Ok(resp) => resp, - Err(e) => { - return ValidateSessionResult { - valid: false, - membership_type: None, - error: Some(format!("Failed to connect: {}", e)), - }; - } - }; - - if response.status() == reqwest::StatusCode::UNAUTHORIZED - || response.status() == reqwest::StatusCode::FORBIDDEN - { - return ValidateSessionResult { - valid: false, - membership_type: None, - error: Some("Session token expired or invalid".to_string()), - }; - } - - if !response.status().is_success() { - return ValidateSessionResult { - valid: false, - membership_type: None, - error: Some(format!("API returned status {}", response.status())), - }; - } - - let data: serde_json::Value = match response.json().await { - Ok(d) => d, - Err(e) => { - return ValidateSessionResult { - valid: false, - membership_type: None, - error: Some(format!("Failed to parse response: {}", e)), - }; - } - }; - - let has_billing_start = data - .get("billingCycleStart") - .and_then(|v| v.as_str()) - .is_some(); - let has_billing_end = data - .get("billingCycleEnd") - .and_then(|v| v.as_str()) - .is_some(); - - if has_billing_start && has_billing_end { - let membership_type = data - .get("membershipType") - .and_then(|v| v.as_str()) - .map(|s| s.to_string()); - ValidateSessionResult { - valid: true, - membership_type, - error: None, - } - } else { - ValidateSessionResult { - valid: false, - membership_type: None, - error: Some("Invalid response format".to_string()), - } - } -} - -pub async fn fetch_cursor_usage_csv(session_token: &str) -> Result { - let client = build_cursor_http_client()?; - let response = client - .get(USAGE_CSV_ENDPOINT) - .headers(build_cursor_headers(session_token)) - .send() - .await?; - - if response.status() == reqwest::StatusCode::UNAUTHORIZED - || response.status() == reqwest::StatusCode::FORBIDDEN - { - anyhow::bail!( - "Cursor session expired. Please run 'bunx tokscale@latest cursor login' to re-authenticate." - ); - } - - if !response.status().is_success() { - anyhow::bail!("Cursor API returned status {}", response.status()); - } - - let text = response.text().await?; - - if !text.starts_with("Date,") { - anyhow::bail!("Invalid response from Cursor API - expected CSV format"); - } - - Ok(text) -} - -async fn sync_cursor_cache_with_fetcher(fetch_usage_csv: F) -> SyncCursorResult -where - F: Fn(String) -> Fut, - Fut: std::future::Future>, -{ - let home_dir = match home_dir() { - Ok(home_dir) => home_dir, - Err(e) => { - return SyncCursorResult { - synced: false, - rows: 0, - error: Some(format!("Failed to get home dir: {}", e)), - }; - } - }; - - sync_cursor_cache_with_fetcher_in_home(&home_dir, fetch_usage_csv).await -} - -async fn sync_cursor_cache_with_fetcher_in_home( - home_dir: &Path, - fetch_usage_csv: F, -) -> SyncCursorResult -where - F: Fn(String) -> Fut, - Fut: std::future::Future>, -{ - migrate_cache_dir_from_old_path_in_home(home_dir); - - let store = match load_credentials_store_from_home(home_dir) { - Some(s) => s, - None => { - return SyncCursorResult { - synced: false, - rows: 0, - error: Some("Not authenticated".to_string()), - }; - } - }; - - if store.accounts.is_empty() { - return SyncCursorResult { - synced: false, - rows: 0, - error: Some("Not authenticated".to_string()), - }; - } - - let cache_dir = cursor_cache_dir(home_dir); - if let Err(e) = fs::create_dir_all(&cache_dir) { - return SyncCursorResult { - synced: false, - rows: 0, - error: Some(format!("Failed to create cache dir: {}", e)), - }; - } - #[cfg(unix)] - { - use std::os::unix::fs::PermissionsExt; - let _ = fs::set_permissions(&cache_dir, fs::Permissions::from_mode(0o700)); - } - - let active_dup = cache_dir.join(format!( - "usage.{}.csv", - sanitize_account_id_for_filename(&store.active_account_id) - )); - if active_dup.exists() { - let _ = fs::remove_file(&active_dup); - } - - let mut total_rows = 0; - let mut success_count = 0; - let mut errors: Vec = Vec::new(); - - for (account_id, credentials) in &store.accounts { - let is_active = account_id == &store.active_account_id; - - match fetch_usage_csv(credentials.session_token.clone()).await { - Ok(csv_text) => { - let file_path = if is_active { - cache_dir.join("usage.csv") - } else { - cache_dir.join(format!( - "usage.{}.csv", - sanitize_account_id_for_filename(account_id) - )) - }; - - let row_count = count_cursor_csv_rows(&csv_text); - - if let Err(e) = atomic_write_file(&file_path, &csv_text) { - errors.push(format!("{}: {}", account_id, e)); - } else { - total_rows += row_count; - success_count += 1; - } - } - Err(e) => { - errors.push(format!("{}: {}", account_id, e)); - } - } - } - - // Touch the sync-attempt marker unconditionally after the per-account loop - // (regardless of partial failures). The marker's mtime short-circuits the - // secondary-account freshness check so a permanently-stale secondary - // doesn't force an implicit sync on every invocation. We ignore errors - // here — if the marker can't be written (e.g. disk full) the gate simply - // falls through to the CSV-freshness check as before. - let _ = std::fs::OpenOptions::new() - .create(true) - .truncate(true) - .write(true) - .open(cache_dir.join(CURSOR_SYNC_ATTEMPT_MARKER)); - - if success_count == 0 { - return SyncCursorResult { - synced: false, - rows: 0, - error: Some( - errors - .first() - .cloned() - .unwrap_or_else(|| "Cursor sync failed".to_string()), - ), - }; - } - - SyncCursorResult { - synced: true, - rows: total_rows, - error: if errors.is_empty() { - None - } else { - Some(format!( - "Some accounts failed to sync ({}/{})", - errors.len(), - store.accounts.len() - )) - }, - } -} - -pub async fn sync_cursor_cache() -> SyncCursorResult { - sync_cursor_cache_with_fetcher(|session_token| async move { - fetch_cursor_usage_csv(&session_token).await - }) - .await -} - -fn archive_cache_file(file_path: &std::path::Path, label: &str) -> Result<()> { - let cache_dir = get_cursor_cache_dir()?; - let archive_dir = cache_dir.join("archive"); - if !archive_dir.exists() { - fs::create_dir_all(&archive_dir)?; - #[cfg(unix)] - { - use std::os::unix::fs::PermissionsExt; - fs::set_permissions(&archive_dir, fs::Permissions::from_mode(0o700))?; - } - } - - let safe_label = sanitize_account_id_for_filename(label); - let ts = chrono::Utc::now().format("%Y-%m-%dT%H-%M-%S").to_string(); - let dest = archive_dir.join(format!("{}-{}.csv", safe_label, ts)); - fs::rename(file_path, dest)?; - Ok(()) -} - -pub fn run_cursor_login(name: Option) -> Result<()> { - use colored::Colorize; - use tokio::runtime::Runtime; - - let rt = Runtime::new()?; - - println!("\n {}\n", "Cursor IDE - Login".cyan()); - - if let Some(ref label) = name { - if find_account(label).is_some() { - println!( - " {}", - format!( - "Account '{}' already exists. Use 'bunx tokscale@latest cursor logout --name {}' first.", - label, label - ) - .yellow() - ); - println!(); - return Ok(()); - } - } - - print!(" Enter Cursor WorkosCursorSessionToken value: "); - std::io::stdout().flush()?; - let token = rpassword::read_password().context("Failed to read session token")?; - let token = token.trim().to_string(); - - if token.is_empty() { - println!("\n {}\n", "No token provided.".yellow()); - return Ok(()); - } - - println!(); - println!("{}", " Validating session token...".bright_black()); - - let result = rt.block_on(async { validate_cursor_session(&token).await }); - - if !result.valid { - let msg = result - .error - .unwrap_or_else(|| "Invalid session token".to_string()); - println!( - "\n {}\n", - format!("{}. Please check and try again.", msg).red() - ); - std::process::exit(1); - } - - let account_id = save_credentials(&token, name.as_deref())?; - - let display_name = name.as_deref().unwrap_or(&account_id); - println!( - "\n {}", - format!( - "Successfully logged in to Cursor as {}", - display_name.bold() - ) - .green() - ); - println!("{}", format!(" Account ID: {}", account_id).bright_black()); - println!(); - - Ok(()) -} - -pub fn run_cursor_logout(name: Option, all: bool, purge_cache: bool) -> Result<()> { - use colored::Colorize; - - if all { - let accounts = list_accounts(); - if accounts.is_empty() { - println!("\n {}\n", "No saved Cursor accounts.".yellow()); - return Ok(()); - } - - remove_all_accounts(purge_cache)?; - println!("\n {}\n", "Logged out from all Cursor accounts.".green()); - return Ok(()); - } - - if let Some(ref account_name) = name { - remove_account(account_name, purge_cache)?; - println!( - "\n {}\n", - format!("Logged out from Cursor account '{}'.", account_name).green() - ); - return Ok(()); - } - - let Some(store) = load_credentials_store() else { - println!("\n {}\n", "No saved Cursor accounts.".yellow()); - return Ok(()); - }; - let active_id = store.active_account_id.clone(); - let display = store - .accounts - .get(&active_id) - .and_then(|a| a.label.clone()) - .unwrap_or_else(|| active_id.clone()); - - remove_account(&active_id, purge_cache)?; - println!( - "\n {}\n", - format!("Logged out from Cursor account '{}'.", display).green() - ); - - Ok(()) -} - -pub fn run_cursor_status(name: Option) -> Result<()> { - use colored::Colorize; - use tokio::runtime::Runtime; - - let rt = Runtime::new()?; - - let credentials = if let Some(ref account_name) = name { - load_credentials_for(account_name) - } else { - load_active_credentials() - }; - - let credentials = match credentials { - Some(c) => c, - None => { - if let Some(ref account_name) = name { - println!( - "\n {}\n", - format!("Account not found: {}", account_name).red() - ); - } else { - println!("\n {}", "No saved Cursor accounts.".yellow()); - println!( - "{}", - " Run 'bunx tokscale@latest cursor login' to authenticate.\n".bright_black() - ); - } - return Ok(()); - } - }; - - println!("\n {}\n", "Cursor IDE - Status".cyan()); - - let display_name = credentials.label.as_deref().unwrap_or("(no label)"); - println!("{}", format!(" Account: {}", display_name).white()); - if let Some(ref uid) = credentials.user_id { - println!("{}", format!(" User ID: {}", uid).bright_black()); - } - - println!("{}", " Validating session...".bright_black()); - - let result = rt.block_on(async { validate_cursor_session(&credentials.session_token).await }); - - if result.valid { - println!(" {}", "Session: Valid".green()); - if let Some(membership) = result.membership_type { - println!("{}", format!(" Membership: {}", membership).bright_black()); - } - } else { - let msg = result - .error - .unwrap_or_else(|| "Invalid / Expired".to_string()); - println!(" {}", format!("Session: {}", msg).red()); - } - println!(); - - Ok(()) -} - -pub fn run_cursor_accounts(json: bool) -> Result<()> { - use colored::Colorize; - - let accounts = list_accounts(); - - if json { - #[derive(Serialize)] - struct Output { - accounts: Vec, - } - let output = Output { accounts }; - println!("{}", serde_json::to_string_pretty(&output)?); - return Ok(()); - } - - if accounts.is_empty() { - println!("\n {}\n", "No saved Cursor accounts.".yellow()); - return Ok(()); - } - - println!("{}", "\n Cursor IDE - Accounts\n".cyan()); - for acct in &accounts { - let name = if let Some(ref label) = acct.label { - format!("{} ({})", label, acct.id) - } else { - acct.id.clone() - }; - let marker = if acct.is_active { "*" } else { "-" }; - let marker_colored = if acct.is_active { - marker.green().to_string() - } else { - marker.bright_black().to_string() - }; - println!(" {} {}", marker_colored, name); - } - println!(); - - Ok(()) -} - -pub fn run_cursor_sync(json: bool) -> Result<()> { - use colored::Colorize; - use tokio::runtime::Runtime; - - let rt = Runtime::new()?; - let result = rt.block_on(sync_cursor_cache()); - - if json { - println!("{}", serde_json::to_string_pretty(&result)?); - return Ok(()); - } - - println!("\n {}\n", "Cursor IDE - Sync".cyan()); - if result.synced { - println!( - "{}", - format!(" Synced {} Cursor usage event(s).", result.rows).green() - ); - if let Some(error) = result.error { - println!("{}", format!(" Warning: {}", error).yellow()); - } - } else if let Some(error) = result.error { - println!("{}", format!(" Sync failed: {}", error).red()); - } else { - println!("{}", " Sync failed.".red()); - } - println!(); - - Ok(()) -} - -pub fn run_cursor_switch(name: &str) -> Result<()> { - use colored::Colorize; - - set_active_account(name)?; - println!( - "\n {}\n", - format!("Active Cursor account set to {}", name.bold()).green() - ); - - Ok(()) -} - -#[cfg(test)] -mod tests { - use super::*; - use std::collections::HashMap; - use tempfile::TempDir; - - #[test] - fn test_extract_user_id_from_session_token_with_url_encoding() { - // Test URL-encoded separator (%3A%3A) - assert_eq!( - extract_user_id_from_session_token("user123%3A%3Atoken456"), - Some("user123".to_string()) - ); - assert_eq!( - extract_user_id_from_session_token(" user123%3A%3Atoken456 "), - Some("user123".to_string()) - ); - } - - #[test] - fn test_extract_user_id_from_session_token_with_double_colon() { - // Test plain :: separator - assert_eq!( - extract_user_id_from_session_token("user456::token789"), - Some("user456".to_string()) - ); - assert_eq!( - extract_user_id_from_session_token(" user456::token789 "), - Some("user456".to_string()) - ); - } - - #[test] - fn test_extract_user_id_from_session_token_invalid() { - // No separator - assert_eq!(extract_user_id_from_session_token("invalidtoken"), None); - // Empty user ID - assert_eq!(extract_user_id_from_session_token("%3A%3Atoken"), None); - assert_eq!(extract_user_id_from_session_token("::token"), None); - // Empty string - assert_eq!(extract_user_id_from_session_token(""), None); - // Whitespace only - assert_eq!(extract_user_id_from_session_token(" "), None); - } - - #[test] - fn test_derive_account_id_with_user_id() { - // Should extract user ID when present - let account_id = derive_account_id("user123%3A%3Atoken456"); - assert_eq!(account_id, "user123"); - - let account_id = derive_account_id("user456::token789"); - assert_eq!(account_id, "user456"); - } - - #[test] - fn test_derive_account_id_without_user_id() { - // Should generate anon-{hash} when no user ID - let account_id = derive_account_id("randomtoken"); - assert!(account_id.starts_with("anon-")); - assert_eq!(account_id.len(), 17); // "anon-" + 12 hex chars - - // Same token should produce same hash - let account_id2 = derive_account_id("randomtoken"); - assert_eq!(account_id, account_id2); - - // Different tokens should produce different hashes - let account_id3 = derive_account_id("differenttoken"); - assert_ne!(account_id, account_id3); - } - - #[test] - fn test_sanitize_account_id_for_filename_basic() { - // Alphanumeric, dots, underscores, hyphens should be preserved - assert_eq!(sanitize_account_id_for_filename("user123"), "user123"); - assert_eq!( - sanitize_account_id_for_filename("user.name_123-test"), - "user.name_123-test" - ); - } - - #[test] - fn test_sanitize_account_id_for_filename_unsafe_chars() { - // Unsafe characters should be replaced with hyphens - assert_eq!( - sanitize_account_id_for_filename("user@example.com"), - "user-example.com" - ); - assert_eq!( - sanitize_account_id_for_filename("user/name\\test"), - "user-name-test" - ); - assert_eq!(sanitize_account_id_for_filename("user name"), "user-name"); - } - - #[test] - fn test_sanitize_account_id_for_filename_edge_cases() { - // Uppercase should be lowercased - assert_eq!( - sanitize_account_id_for_filename("UserName123"), - "username123" - ); - - // Leading/trailing hyphens should be trimmed - assert_eq!(sanitize_account_id_for_filename("---user---"), "user"); - - // Empty after sanitization should return "account" - assert_eq!(sanitize_account_id_for_filename("@@@"), "account"); - assert_eq!(sanitize_account_id_for_filename(""), "account"); - - // Whitespace only should return "account" - assert_eq!(sanitize_account_id_for_filename(" "), "account"); - } - - #[test] - fn test_sanitize_account_id_for_filename_length_limit() { - // Should truncate to 80 characters - let long_id = "a".repeat(100); - let sanitized = sanitize_account_id_for_filename(&long_id); - assert_eq!(sanitized.len(), 80); - assert_eq!(sanitized, "a".repeat(80)); - - // Should preserve exactly 80 characters - let exactly_80 = "b".repeat(80); - let sanitized = sanitize_account_id_for_filename(&exactly_80); - assert_eq!(sanitized.len(), 80); - } - - #[test] - fn test_build_cursor_http_client_applies_timeout() { - // Constructing the client must succeed and surface no panics; the - // configured timeout is the property the HIGH finding flagged. - let client = build_cursor_http_client().expect("client builds"); - // reqwest::Client doesn't expose its timeout publicly, but we can at - // least confirm the const wired into the builder is the documented - // 8s value — a future change to the constant must be deliberate. - assert_eq!(CURSOR_HTTP_TIMEOUT, std::time::Duration::from_secs(8)); - // Use the client briefly to ensure it's structurally valid. - let _ = client.get("https://example.invalid").build(); - } - - #[test] - fn test_cursor_usage_cache_is_fresh_returns_false_when_cache_missing() { - let temp = tempfile::tempdir().unwrap(); - // No cache dir created yet. - assert!(!cursor_usage_cache_is_fresh_in( - temp.path(), - Duration::from_secs(300) - )); - } - - #[test] - fn test_cursor_usage_cache_is_fresh_returns_false_when_no_csv_files() { - let temp = tempfile::tempdir().unwrap(); - let cache_dir = cursor_cache_dir(temp.path()); - fs::create_dir_all(&cache_dir).unwrap(); - // Unrelated file present, but no usage*.csv. - fs::write(cache_dir.join("README.txt"), "noise").unwrap(); - assert!(!cursor_usage_cache_is_fresh_in( - temp.path(), - Duration::from_secs(300) - )); - } - - #[test] - fn test_cursor_usage_cache_is_fresh_returns_true_for_recent_file() { - let temp = tempfile::tempdir().unwrap(); - let cache_dir = cursor_cache_dir(temp.path()); - fs::create_dir_all(&cache_dir).unwrap(); - fs::write(cache_dir.join("usage.csv"), "Date,Model\n").unwrap(); - // Just-written file is fresh under any reasonable window. - assert!(cursor_usage_cache_is_fresh_in( - temp.path(), - Duration::from_secs(300) - )); - } - - #[test] - fn test_cursor_usage_cache_is_fresh_returns_false_for_old_file() { - let temp = tempfile::tempdir().unwrap(); - let cache_dir = cursor_cache_dir(temp.path()); - fs::create_dir_all(&cache_dir).unwrap(); - let path = cache_dir.join("usage.csv"); - fs::write(&path, "Date,Model\n").unwrap(); - // Backdate the mtime by an hour. Skip the test if the platform refuses - // to set mtime (rare on POSIX/Windows but possible on exotic FS). - let f = std::fs::OpenOptions::new().write(true).open(&path).unwrap(); - let Ok(()) = f.set_modified(SystemTime::now() - Duration::from_secs(3600)) else { - return; - }; - drop(f); - assert!(!cursor_usage_cache_is_fresh_in( - temp.path(), - Duration::from_secs(300) - )); - } - - #[test] - fn test_cursor_usage_cache_is_fresh_requires_active_usage_csv_when_secondary_is_fresh() { - // A recently-synced secondary account must not mask a stale active - // account cache. The implicit sync gate should refresh the cache that - // local reports read from `usage.csv`. - let temp = tempfile::tempdir().unwrap(); - let cache_dir = cursor_cache_dir(temp.path()); - fs::create_dir_all(&cache_dir).unwrap(); - let stale_path = cache_dir.join("usage.csv"); - fs::write(&stale_path, "Date,Model\n").unwrap(); - let stale = std::fs::OpenOptions::new() - .write(true) - .open(&stale_path) - .unwrap(); - let Ok(()) = stale.set_modified(SystemTime::now() - Duration::from_secs(3600)) else { - return; - }; - drop(stale); - // Secondary account written just now. - fs::write(cache_dir.join("usage.team-a.csv"), "Date,Model\n").unwrap(); - assert!(!cursor_usage_cache_is_fresh_in( - temp.path(), - Duration::from_secs(300) - )); - } - - #[test] - fn test_cursor_usage_cache_is_fresh_returns_false_when_active_cache_missing() { - // A fresh secondary account cache alone is not enough: without the - // active account's `usage.csv`, the next report would use stale/missing - // active data unless the implicit sync runs. - let temp = tempfile::tempdir().unwrap(); - let cache_dir = cursor_cache_dir(temp.path()); - fs::create_dir_all(&cache_dir).unwrap(); - fs::write(cache_dir.join("usage.team-a.csv"), "Date,Model\n").unwrap(); - assert!(!cursor_usage_cache_is_fresh_in( - temp.path(), - Duration::from_secs(300) - )); - } - - #[test] - fn test_cursor_usage_cache_is_fresh_requires_all_expected_account_files() -> Result<()> { - let temp_dir = TempDir::new()?; - let mut accounts = HashMap::new(); - accounts.insert( - "active-account".to_string(), - CursorCredentials { - session_token: "token-active".to_string(), - user_id: Some("active-account".to_string()), - created_at: "2026-01-01T00:00:00Z".to_string(), - expires_at: None, - label: Some("work".to_string()), - }, - ); - accounts.insert( - "team/account".to_string(), - CursorCredentials { - session_token: "token-secondary".to_string(), - user_id: Some("team/account".to_string()), - created_at: "2026-01-01T00:00:00Z".to_string(), - expires_at: None, - label: Some("personal".to_string()), - }, - ); - save_credentials_store_in_home( - temp_dir.path(), - &CursorCredentialsStore { - version: 1, - active_account_id: "active-account".to_string(), - accounts, - }, - )?; - - let cache_dir = cursor_cache_dir(temp_dir.path()); - fs::create_dir_all(&cache_dir)?; - fs::write(cache_dir.join("usage.csv"), "Date,Model\n")?; - - assert!(!cursor_usage_cache_is_fresh_in( - temp_dir.path(), - Duration::from_secs(300) - )); - - fs::write(cache_dir.join("usage.team-account.csv"), "Date,Model\n")?; - assert!(cursor_usage_cache_is_fresh_in( - temp_dir.path(), - Duration::from_secs(300) - )); - - Ok(()) - } - - #[test] - fn test_cursor_expected_cache_paths_dedupes_sanitized_account_collisions() { - let temp_dir = TempDir::new().unwrap(); - let mut accounts = HashMap::new(); - accounts.insert( - "active-account".to_string(), - CursorCredentials { - session_token: "token-active".to_string(), - user_id: Some("active-account".to_string()), - created_at: "2026-01-01T00:00:00Z".to_string(), - expires_at: None, - label: Some("active".to_string()), - }, - ); - accounts.insert( - "team/account-a".to_string(), - CursorCredentials { - session_token: "token-team-a".to_string(), - user_id: Some("team/account-a".to_string()), - created_at: "2026-01-01T00:00:00Z".to_string(), - expires_at: None, - label: Some("team-a".to_string()), - }, - ); - accounts.insert( - "team@account-a".to_string(), - CursorCredentials { - session_token: "token-team-b".to_string(), - user_id: Some("team@account-a".to_string()), - created_at: "2026-01-01T00:00:00Z".to_string(), - expires_at: None, - label: Some("team-b".to_string()), - }, - ); - save_credentials_store_in_home( - temp_dir.path(), - &CursorCredentialsStore { - version: 1, - active_account_id: "active-account".to_string(), - accounts, - }, - ) - .unwrap(); - - let paths = expected_cursor_usage_cache_paths_in(temp_dir.path()); - let cache_dir = cursor_cache_dir(temp_dir.path()); - let expected = vec![ - cache_dir.join("usage.csv"), - cache_dir.join("usage.team-account-a.csv"), - ]; - assert_eq!(paths, expected); - } - - #[test] - fn test_count_cursor_csv_rows_valid() { - // Valid CSV with header - let csv = "Date,Model,Tokens\n2024-01-01,gpt-4,100\n2024-01-02,gpt-4,200\n"; - assert_eq!(count_cursor_csv_rows(csv), 2); - - // Single row - let csv = "Date,Model,Tokens\n2024-01-01,gpt-4,100\n"; - assert_eq!(count_cursor_csv_rows(csv), 1); - } - - #[test] - fn test_count_cursor_csv_rows_empty() { - // Header only - let csv = "Date,Model,Tokens\n"; - assert_eq!(count_cursor_csv_rows(csv), 0); - - // Empty string - let csv = ""; - assert_eq!(count_cursor_csv_rows(csv), 0); - } - - #[test] - fn test_count_cursor_csv_rows_malformed() { - // CSV reader with flexible=true accepts rows with different column counts - // This test verifies the actual behavior: all parseable rows are counted - let csv = "Date,Model,Tokens\n2024-01-01,gpt-4,100\ninvalid,row\n2024-01-02,gpt-4,200\n"; - assert_eq!(count_cursor_csv_rows(csv), 3); - } - - #[test] - fn test_sync_cursor_cache_writes_active_and_secondary_account_files() -> Result<()> { - let temp_dir = TempDir::new()?; - - let mut accounts = HashMap::new(); - accounts.insert( - "active-account".to_string(), - CursorCredentials { - session_token: "token-active".to_string(), - user_id: Some("active-account".to_string()), - created_at: "2026-01-01T00:00:00Z".to_string(), - expires_at: None, - label: Some("work".to_string()), - }, - ); - accounts.insert( - "team/account".to_string(), - CursorCredentials { - session_token: "token-secondary".to_string(), - user_id: Some("team/account".to_string()), - created_at: "2026-01-01T00:00:00Z".to_string(), - expires_at: None, - label: Some("personal".to_string()), - }, - ); - save_credentials_store_in_home( - temp_dir.path(), - &CursorCredentialsStore { - version: 1, - active_account_id: "active-account".to_string(), - accounts, - }, - )?; - - let runtime = tokio::runtime::Runtime::new()?; - let result = runtime.block_on(sync_cursor_cache_with_fetcher_in_home( - temp_dir.path(), - |session_token| { - let csv = match session_token.as_str() { - "token-active" => "Date,Model,Tokens\n2026-01-01,gpt-5,100\n", - "token-secondary" => { - "Date,Model,Tokens\n2026-01-02,gpt-5,200\n2026-01-03,gpt-5,300\n" - } - _ => "Date,Model,Tokens\n", - } - .to_string(); - async move { Ok(csv) } - }, - )); - - assert!(result.synced); - assert_eq!(result.rows, 3); - assert_eq!(result.error, None); - - let cache_dir = cursor_cache_dir(temp_dir.path()); - assert_eq!( - fs::read_to_string(cache_dir.join("usage.csv"))?, - "Date,Model,Tokens\n2026-01-01,gpt-5,100\n" - ); - assert_eq!( - fs::read_to_string(cache_dir.join("usage.team-account.csv"))?, - "Date,Model,Tokens\n2026-01-02,gpt-5,200\n2026-01-03,gpt-5,300\n" - ); - assert!(!cache_dir.join("usage.active-account.csv").exists()); - - Ok(()) - } - - #[test] - fn test_atomic_write_file_basic() -> Result<()> { - let temp_dir = TempDir::new()?; - let file_path = temp_dir.path().join("test.txt"); - let contents = "Hello, world!"; - - atomic_write_file(&file_path, contents)?; - - // Verify file was created and contains correct content - assert!(file_path.exists()); - let read_contents = fs::read_to_string(&file_path)?; - assert_eq!(read_contents, contents); - - Ok(()) - } - - #[test] - fn test_atomic_write_file_creates_parent_dirs() -> Result<()> { - let temp_dir = TempDir::new()?; - let nested_path = temp_dir - .path() - .join("a") - .join("b") - .join("c") - .join("test.txt"); - let contents = "Nested file"; - - atomic_write_file(&nested_path, contents)?; - - // Verify parent directories were created - assert!(nested_path.exists()); - let read_contents = fs::read_to_string(&nested_path)?; - assert_eq!(read_contents, contents); - - Ok(()) - } - - #[test] - fn test_atomic_write_file_overwrites_existing() -> Result<()> { - let temp_dir = TempDir::new()?; - let file_path = temp_dir.path().join("test.txt"); - - // Write initial content - atomic_write_file(&file_path, "Initial")?; - assert_eq!(fs::read_to_string(&file_path)?, "Initial"); - - // Overwrite with new content - atomic_write_file(&file_path, "Updated")?; - assert_eq!(fs::read_to_string(&file_path)?, "Updated"); - - Ok(()) - } - - #[test] - #[cfg(unix)] - fn test_atomic_write_file_permissions() -> Result<()> { - use std::os::unix::fs::PermissionsExt; - - let temp_dir = TempDir::new()?; - let file_path = temp_dir.path().join("test.txt"); - - atomic_write_file(&file_path, "Secret")?; - - // Verify file has 0o600 permissions (owner read/write only) - let metadata = fs::metadata(&file_path)?; - let permissions = metadata.permissions(); - assert_eq!(permissions.mode() & 0o777, 0o600); - - Ok(()) - } - - #[test] - fn test_copy_dir_recursive_basic() -> Result<()> { - let temp_dir = TempDir::new()?; - let src = temp_dir.path().join("src"); - let dst = temp_dir.path().join("dst"); - - // Create source directory structure - fs::create_dir_all(&src)?; - fs::write(src.join("file1.txt"), "Content 1")?; - fs::write(src.join("file2.txt"), "Content 2")?; - - // Create destination directory - fs::create_dir_all(&dst)?; - - // Copy recursively - copy_dir_recursive(&src, &dst)?; - - // Verify files were copied - assert!(dst.join("file1.txt").exists()); - assert!(dst.join("file2.txt").exists()); - assert_eq!(fs::read_to_string(dst.join("file1.txt"))?, "Content 1"); - assert_eq!(fs::read_to_string(dst.join("file2.txt"))?, "Content 2"); - - Ok(()) - } - - #[test] - fn test_copy_dir_recursive_nested() -> Result<()> { - let temp_dir = TempDir::new()?; - let src = temp_dir.path().join("src"); - let dst = temp_dir.path().join("dst"); - - // Create nested source directory structure - fs::create_dir_all(src.join("subdir1").join("subdir2"))?; - fs::write(src.join("root.txt"), "Root")?; - fs::write(src.join("subdir1").join("file1.txt"), "File 1")?; - fs::write( - src.join("subdir1").join("subdir2").join("file2.txt"), - "File 2", - )?; - - // Create destination directory - fs::create_dir_all(&dst)?; - - // Copy recursively - copy_dir_recursive(&src, &dst)?; - - // Verify nested structure was copied - assert!(dst.join("root.txt").exists()); - assert!(dst.join("subdir1").join("file1.txt").exists()); - assert!(dst - .join("subdir1") - .join("subdir2") - .join("file2.txt") - .exists()); - assert_eq!(fs::read_to_string(dst.join("root.txt"))?, "Root"); - assert_eq!( - fs::read_to_string(dst.join("subdir1").join("file1.txt"))?, - "File 1" - ); - assert_eq!( - fs::read_to_string(dst.join("subdir1").join("subdir2").join("file2.txt"))?, - "File 2" - ); - - Ok(()) - } - - #[test] - fn test_copy_dir_recursive_empty_dir() -> Result<()> { - let temp_dir = TempDir::new()?; - let src = temp_dir.path().join("src"); - let dst = temp_dir.path().join("dst"); - - // Create empty source directory - fs::create_dir_all(&src)?; - fs::create_dir_all(&dst)?; - - // Copy recursively (should succeed with no files) - copy_dir_recursive(&src, &dst)?; - - // Verify destination exists but is empty - assert!(dst.exists()); - assert_eq!(fs::read_dir(&dst)?.count(), 0); - - Ok(()) - } - - /// Helper: build a two-account credentials store in `home_dir`. - fn setup_two_account_store(home_dir: &std::path::Path) -> Result<()> { - let mut accounts = HashMap::new(); - accounts.insert( - "active-account".to_string(), - CursorCredentials { - session_token: "token-active".to_string(), - user_id: Some("active-account".to_string()), - created_at: "2026-01-01T00:00:00Z".to_string(), - expires_at: None, - label: Some("work".to_string()), - }, - ); - accounts.insert( - "team/account".to_string(), - CursorCredentials { - session_token: "token-secondary".to_string(), - user_id: Some("team/account".to_string()), - created_at: "2026-01-01T00:00:00Z".to_string(), - expires_at: None, - label: Some("personal".to_string()), - }, - ); - save_credentials_store_in_home( - home_dir, - &CursorCredentialsStore { - version: 1, - active_account_id: "active-account".to_string(), - accounts, - }, - ) - } - - /// Helper: backdate a file's mtime by `secs` seconds. Returns `false` if - /// the platform refuses to set mtime (exotic FS), signalling the caller to - /// skip the test. - fn backdate_file(path: &std::path::Path, secs: u64) -> bool { - let f = match std::fs::OpenOptions::new().write(true).open(path) { - Ok(f) => f, - Err(_) => return false, - }; - f.set_modified(SystemTime::now() - Duration::from_secs(secs)) - .is_ok() - } - - #[test] - fn test_freshness_gate_passes_when_active_fresh_and_marker_fresh_despite_stale_secondary( - ) -> Result<()> { - // Active CSV fresh + stale secondary CSV + fresh marker → gate passes. - // This is the key scenario: a permanently-stale secondary must not - // thrash implicit sync when the marker proves we already tried recently. - let temp_dir = TempDir::new()?; - setup_two_account_store(temp_dir.path())?; - - let cache_dir = cursor_cache_dir(temp_dir.path()); - fs::create_dir_all(&cache_dir)?; - - // Fresh active cache. - fs::write(cache_dir.join("usage.csv"), "Date,Model\n")?; - - // Stale secondary cache. - let secondary = cache_dir.join("usage.team-account.csv"); - fs::write(&secondary, "Date,Model\n")?; - if !backdate_file(&secondary, 3600) { - return Ok(()); // platform can't set mtime — skip - } - - // Fresh sync-attempt marker. - fs::write(cache_dir.join(CURSOR_SYNC_ATTEMPT_MARKER), "")?; - - assert!( - cursor_usage_cache_is_fresh_in(temp_dir.path(), Duration::from_secs(300)), - "fresh marker should short-circuit stale secondary" - ); - Ok(()) - } - - #[test] - fn test_freshness_gate_fails_when_active_fresh_but_no_marker_and_stale_secondary() -> Result<()> - { - // Active CSV fresh + stale secondary CSV + NO marker → gate fails so - // an implicit sync is triggered to try fetching the secondary again. - let temp_dir = TempDir::new()?; - setup_two_account_store(temp_dir.path())?; - - let cache_dir = cursor_cache_dir(temp_dir.path()); - fs::create_dir_all(&cache_dir)?; - - fs::write(cache_dir.join("usage.csv"), "Date,Model\n")?; - - let secondary = cache_dir.join("usage.team-account.csv"); - fs::write(&secondary, "Date,Model\n")?; - if !backdate_file(&secondary, 3600) { - return Ok(()); - } - - // No marker written. - - assert!( - !cursor_usage_cache_is_fresh_in(temp_dir.path(), Duration::from_secs(300)), - "without marker, stale secondary should trigger sync" - ); - Ok(()) - } - - #[test] - fn test_freshness_gate_fails_when_active_stale_even_with_fresh_marker() -> Result<()> { - // Stale active CSV + fresh marker → gate still fails. The marker must - // never mask a stale active cache — the active data is what reports - // read from. - let temp_dir = TempDir::new()?; - setup_two_account_store(temp_dir.path())?; - - let cache_dir = cursor_cache_dir(temp_dir.path()); - fs::create_dir_all(&cache_dir)?; - - // Stale active cache. - let active = cache_dir.join("usage.csv"); - fs::write(&active, "Date,Model\n")?; - if !backdate_file(&active, 3600) { - return Ok(()); - } - - // Fresh secondary and fresh marker. - fs::write(cache_dir.join("usage.team-account.csv"), "Date,Model\n")?; - fs::write(cache_dir.join(CURSOR_SYNC_ATTEMPT_MARKER), "")?; - - assert!( - !cursor_usage_cache_is_fresh_in(temp_dir.path(), Duration::from_secs(300)), - "stale active cache must always trigger sync regardless of marker" - ); - Ok(()) - } - - #[test] - fn test_sync_writes_attempt_marker() -> Result<()> { - // After sync_cursor_cache_with_fetcher_in_home completes (even with a - // partial failure), the marker file must exist in the cache dir. - let temp_dir = TempDir::new()?; - setup_two_account_store(temp_dir.path())?; - - let runtime = tokio::runtime::Runtime::new()?; - let _result = runtime.block_on(sync_cursor_cache_with_fetcher_in_home( - temp_dir.path(), - |session_token| { - // Secondary deliberately fails to simulate a broken account. - let result: Result = match session_token.as_str() { - "token-active" => Ok("Date,Model,Tokens\n2026-01-01,gpt-5,10\n".to_string()), - _ => Err(anyhow::anyhow!("simulated fetch failure")), - }; - async move { result } - }, - )); - - let cache_dir = cursor_cache_dir(temp_dir.path()); - assert!( - cache_dir.join(CURSOR_SYNC_ATTEMPT_MARKER).exists(), - "marker must be written even when a secondary account fetch fails" - ); - Ok(()) - } -} diff --git a/crates/tokscale-cli/src/failure.rs b/crates/tokscale-cli/src/failure.rs new file mode 100644 index 000000000..17ff3a6ef --- /dev/null +++ b/crates/tokscale-cli/src/failure.rs @@ -0,0 +1,169 @@ +use std::fmt; + +use tokscale_core::LocalReportError; + +use crate::tui::settings::{NativeTimeoutError, SettingsLoadError}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum FailureClass { + InvalidInvocation, + Operational, +} + +#[derive(Debug)] +pub(crate) struct CliFailure { + class: FailureClass, + error: anyhow::Error, +} + +impl CliFailure { + pub(crate) fn invalid_message(message: impl Into) -> Self { + Self { + class: FailureClass::InvalidInvocation, + error: anyhow::anyhow!(message.into()), + } + } + + pub(crate) const fn class(&self) -> FailureClass { + self.class + } + + pub(crate) const fn exit_code(&self) -> i32 { + match self.class { + FailureClass::InvalidInvocation => 2, + FailureClass::Operational => 1, + } + } + + fn classify(error: &anyhow::Error) -> FailureClass { + if error.is::() + || error.is::() + || error + .downcast_ref::() + .is_some_and(LocalReportError::is_invalid_invocation) + || error + .downcast_ref::() + .is_some_and(SettingsLoadError::is_invalid_environment) + { + FailureClass::InvalidInvocation + } else { + FailureClass::Operational + } + } +} + +impl From for CliFailure { + fn from(error: anyhow::Error) -> Self { + Self { + class: Self::classify(&error), + error, + } + } +} + +impl From for CliFailure { + fn from(error: SettingsLoadError) -> Self { + Self::from(anyhow::Error::new(error)) + } +} + +impl From for CliFailure { + fn from(error: NativeTimeoutError) -> Self { + Self::from(anyhow::Error::new(error)) + } +} + +impl fmt::Display for CliFailure { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + if self.error.is::() { + for (index, cause) in self.error.chain().enumerate() { + if index > 0 { + formatter.write_str(": ")?; + } + write!(formatter, "{cause}")?; + if cause.is::() { + break; + } + } + return Ok(()); + } + write!(formatter, "{:#}", self.error) + } +} + +impl std::error::Error for CliFailure { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + Some(self.error.as_ref()) + } +} + +#[derive(Debug, thiserror::Error)] +#[error("{message}")] +pub(crate) struct InvalidConfiguration { + message: String, +} + +impl InvalidConfiguration { + pub(crate) fn new(message: impl Into) -> Self { + Self { + message: message.into(), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::path::PathBuf; + + #[test] + fn typed_invalid_configuration_survives_anyhow_context() { + let error = anyhow::Error::new(InvalidConfiguration::new("bad settings value")) + .context("resolve source scope"); + let failure = CliFailure::from(error); + + assert_eq!(failure.class(), FailureClass::InvalidInvocation); + assert_eq!(failure.exit_code(), 2); + assert!(failure.to_string().contains("resolve source scope")); + assert!(failure.to_string().contains("bad settings value")); + } + + #[test] + fn typed_settings_parse_error_survives_anyhow_context() { + let source = serde_json::from_str::("{").unwrap_err(); + let error = SettingsLoadError::Parse { + path: PathBuf::from("settings.json"), + source, + }; + let failure = CliFailure::from(anyhow::Error::new(error).context("load command settings")); + + assert_eq!(failure.class(), FailureClass::InvalidInvocation); + assert_eq!(failure.exit_code(), 2); + } + + #[test] + fn typed_settings_io_error_remains_operational_through_anyhow_context() { + let error = SettingsLoadError::Read { + path: PathBuf::from("settings.json"), + source: std::io::Error::new(std::io::ErrorKind::PermissionDenied, "denied"), + }; + let failure = CliFailure::from(anyhow::Error::new(error).context("load command settings")); + + assert_eq!(failure.class(), FailureClass::Operational); + assert_eq!(failure.exit_code(), 1); + } + + #[test] + fn typed_operational_local_report_error_remains_operational() { + let error = LocalReportError::from("source cache unavailable".to_string()); + let failure = + CliFailure::from(anyhow::Error::new(error).context("generate local model report")); + + assert_eq!(failure.class(), FailureClass::Operational); + assert_eq!(failure.exit_code(), 1); + assert_eq!( + failure.to_string(), + "generate local model report: source cache unavailable" + ); + } +} diff --git a/crates/tokscale-cli/src/main.rs b/crates/tokscale-cli/src/main.rs index 91bce8a79..4e9540cc0 100644 --- a/crates/tokscale-cli/src/main.rs +++ b/crates/tokscale-cli/src/main.rs @@ -1,888 +1,206 @@ mod antigravity; mod claude_diagnostics; +mod cli; mod commands; -mod cursor; +mod failure; mod paths; -mod trae; mod tui; mod warp; +use anyhow::Result; +use cli::{ + Cli, ExecutionPlan, HeadlessFormat, PricingSource, PricingSubcommand, TerminalState, + WrappedPlan, +}; use commands::cache::{run_source_cache_prune, run_warm_tui_cache}; use commands::clients::run_clients_command; use commands::graph::run_graph_command; use commands::headless::run_headless_command; use commands::hourly::run_hourly_report; -use commands::integrations::{ - run_antigravity_command, run_codex_command, run_cursor_command, run_trae_command, - run_warp_command, -}; +use commands::integrations::{run_antigravity_command, run_warp_command}; use commands::models::run_models_report; use commands::monthly::run_monthly_report; -use commands::pricing::run_pricing_lookup; -use commands::shared::{ - auto_sync_cursor_before_tui, build_client_filter, build_date_filter, - ensure_home_supported_for_tui, normalize_year_filter, parse_client_id_arg, - reject_unsupported_home_override, reject_usage_parent_flags, -}; +use commands::pricing::{run_pricing_list_overrides, run_pricing_lookup}; use commands::time_metrics::run_time_metrics_report; - -use anyhow::Result; -use clap::{Args, CommandFactory, FromArgMatches, Parser, Subcommand}; -use tokscale_core::ClientId; -use tui::Tab; - -fn parse_home_arg(raw: &str) -> Result { - if raw.trim().is_empty() { - return Err("--home must not be empty".to_string()); +use failure::{CliFailure, FailureClass}; + +fn main() { + match run() { + Ok(ExecutionOutcome::Completed) => {} + Ok(ExecutionOutcome::Interrupted) => std::process::exit(130), + Err(error) => { + let prefix = match error.class() { + FailureClass::InvalidInvocation => "error", + FailureClass::Operational => "Error", + }; + eprintln!("{prefix}: {error}"); + std::process::exit(error.exit_code()); + } } - Ok(raw.to_string()) } -fn parse_pricing_provider(raw: &str) -> Result { - match raw.to_lowercase().as_str() { - "custom" | "litellm" | "openrouter" | "models.dev" => Ok(raw.to_string()), - _ => Err(format!( - "invalid provider: {raw}. Valid providers: custom, litellm, openrouter, models.dev" - )), - } +fn run() -> std::result::Result { + let cli = Cli::parse_from_env(); + let plan = ExecutionPlan::resolve(cli, TerminalState::detect())?; + execute(plan) } -#[derive(Parser)] -#[command(name = "tokscale")] -#[command(author, version, about = "AI token usage analytics")] -struct Cli { - #[command(subcommand)] - command: Option, - - #[arg(short, long)] - theme: Option, - - #[arg(short, long, default_value = "0")] - refresh: u64, - - #[arg(long)] - debug: bool, - - #[arg(long, help = "Output as JSON")] - json: bool, - - #[arg(long, help = "Use legacy CLI table output")] - light: bool, - - #[arg( - long = "write-cache", - requires = "light", - conflicts_with = "no_write_cache", - help = "After --light renders, atomically overwrite the TUI cache with this report's data so the next `tokscale tui` starts from fresh data. Persists across invocations via settings.json `light.writeCache`." - )] - write_cache: bool, - - #[arg( - long = "no-write-cache", - requires = "light", - conflicts_with = "write_cache", - help = "Skip cache write even if settings.json `light.writeCache` is true. Only valid with --light." - )] - no_write_cache: bool, - - #[command(flatten)] - clients: ClientFlags, - - #[command(flatten)] - date: DateRangeFlags, - - #[arg( - long, - value_name = "PATH", - global = true, - value_parser = parse_home_arg, - help = "Read local session data from this home directory for local report commands" - )] - home: Option, - - #[arg(long, help = "Show processing time")] - benchmark: bool, - - #[arg( - long, - value_name = "STRATEGY", - default_value = "client,model", - help = "Grouping strategy for --light and --json output: model, client,model, client,provider,model, workspace,model, session,model, client,session,model" - )] - group_by: String, - - #[arg(long, help = "Disable spinner (for AI agents and scripts)")] - no_spinner: bool, +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum ExecutionOutcome { + Completed, + Interrupted, } -#[derive(Subcommand)] -enum Commands { - #[command(about = "Show model usage report")] - Models { - #[arg(long)] - json: bool, - #[arg(long)] - light: bool, - #[command(flatten)] - clients: ClientFlags, - #[command(flatten)] - date: DateRangeFlags, - #[arg(long, help = "Show processing time")] - benchmark: bool, - #[arg( - long, - value_name = "STRATEGY", - default_value = "client,model", - help = "Grouping strategy for --light and --json output: model, client,model, client,provider,model, workspace,model, session,model, client,session,model" - )] - group_by: String, - #[arg( - long = "write-cache", - requires = "light", - conflicts_with = "no_write_cache", - help = "After --light renders, atomically overwrite the TUI cache with this report's data so the next `tokscale tui` starts from fresh data. Persists across invocations via settings.json `light.writeCache`." - )] - write_cache: bool, - #[arg( - long = "no-write-cache", - requires = "light", - conflicts_with = "write_cache", - help = "Skip cache write even if settings.json `light.writeCache` is true. Only valid with --light." - )] - no_write_cache: bool, - #[arg(long, help = "Disable spinner")] - no_spinner: bool, - }, - #[command(about = "Show monthly usage report")] - Monthly { - #[arg(long)] - json: bool, - #[arg(long)] - light: bool, - #[command(flatten)] - clients: ClientFlags, - #[command(flatten)] - date: DateRangeFlags, - #[arg(long, help = "Show processing time")] - benchmark: bool, - #[arg(long, help = "Disable spinner")] - no_spinner: bool, - }, - #[command(about = "Show hourly usage report")] - Hourly { - #[arg(long)] - json: bool, - #[arg(long)] - light: bool, - #[command(flatten)] - clients: ClientFlags, - #[command(flatten)] - date: DateRangeFlags, - #[arg(long, help = "Show processing time")] - benchmark: bool, - #[arg(long, help = "Disable spinner")] - no_spinner: bool, - }, - #[command(about = "Show pricing for a model")] - Pricing { - #[arg(help = "Model ID to look up, or `list-overrides`")] - model_id: String, - #[arg(long, help = "Output as JSON")] - json: bool, - #[arg( - long, - value_parser = parse_pricing_provider, - help = "Force specific pricing source (custom, litellm, openrouter, or models.dev)" - )] - provider: Option, - #[arg(long, help = "Disable spinner")] - no_spinner: bool, - }, - #[command(about = "Show local scan locations and session counts")] - Clients { - #[arg(long, help = "Output as JSON")] - json: bool, - }, - #[command(about = "Export contribution graph data as JSON")] - Graph { - #[arg(long, help = "Write to file instead of stdout")] - output: Option, - #[command(flatten)] - clients: ClientFlags, - #[command(flatten)] - date: DateRangeFlags, - #[arg(long, help = "Show processing time")] - benchmark: bool, - #[arg(long, help = "Disable spinner")] - no_spinner: bool, - }, - #[command(about = "Launch interactive TUI with optional filters")] - Tui { - #[command(flatten)] - clients: ClientFlags, - #[command(flatten)] - date: DateRangeFlags, - }, - #[command(about = "Capture subprocess output for token usage tracking")] - Headless { - #[arg(help = "Source CLI (currently only 'codex' supported)")] - source: String, - #[arg(trailing_var_arg = true, allow_hyphen_values = true)] - args: Vec, - #[arg(long, help = "Override output format (json or jsonl)")] - format: Option, - #[arg(long, help = "Write captured output to file")] - output: Option, - #[arg(long, help = "Do not auto-add JSON output flags")] - no_auto_flags: bool, - }, - #[command(about = "Generate year-in-review wrapped image")] - Wrapped { - #[arg(long, help = "Output file path (default: tokscale-{year}-wrapped.png)")] - output: Option, - #[arg(long, help = "Year to generate (default: current year)")] - year: Option, - #[command(flatten)] - client_flags: ClientFlags, - #[arg( - long, - help = "Display total tokens in abbreviated format (e.g., 7.14B)" - )] - short: bool, - #[arg(long, help = "Show Top OpenCode Agents (default)")] - agents: bool, - #[arg( - long = "clients", - help = "Show Top Clients instead of Top OpenCode Agents" - )] - show_clients: bool, - #[arg(long, help = "Disable pinning of Sisyphus agents in rankings")] - disable_pinned: bool, - #[arg(long, help = "Disable loading spinner (for scripting)")] - no_spinner: bool, - }, - #[command(about = "Show subscription usage and quota for AI providers")] - Usage { - #[arg(long, help = "Output as JSON")] - json: bool, - }, - #[command(about = "Maintain local Tokscale caches")] - Cache { - #[command(subcommand)] - subcommand: CacheSubcommand, - }, - #[command(about = "Codex account integration commands")] - Codex { - #[command(subcommand)] - subcommand: CodexSubcommand, - }, - #[command(about = "Cursor API cache integration commands")] - Cursor { - #[command(subcommand)] - subcommand: CursorSubcommand, - }, - #[command(about = "Antigravity integration commands")] - Antigravity { - #[command(subcommand)] - subcommand: AntigravitySubcommand, - }, - #[command(about = "Trae IDE integration commands")] - Trae { - #[command(subcommand)] - subcommand: TraeSubcommand, - }, - #[command(about = "Warp/Oz aggregate usage integration commands")] - Warp { - #[command(subcommand)] - subcommand: WarpSubcommand, - }, - #[command( - about = "Show session time metrics (usage time, longest continuous, max concurrent)" - )] - TimeMetrics { - #[arg(long)] - json: bool, - #[command(flatten)] - clients: ClientFlags, - #[command(flatten)] - date: DateRangeFlags, - #[arg(long, help = "Disable spinner")] - no_spinner: bool, - }, - #[command(about = "Warm TUI cache in background (internal)", hide = true)] - WarmTuiCache, -} - -#[derive(Subcommand)] -enum CacheSubcommand { - #[command(about = "Remove orphaned and superseded source-message cache shards")] - Prune, -} - -#[derive(Subcommand)] -enum CodexSubcommand { - #[command(about = "Import the current Codex OAuth credentials as a saved account")] - Import { - #[arg(long, help = "Label for this Codex account (e.g., work, personal)")] - name: Option, - }, - #[command(about = "List saved Codex accounts")] - Accounts { - #[arg(long, help = "Output as JSON")] - json: bool, - }, - #[command(about = "Switch active Codex account and write Codex auth.json")] - Switch { - #[arg(help = "Account label or id")] - name: String, - }, - #[command(about = "Remove a saved Codex account")] - Remove { - #[arg(help = "Account label or id")] - name: String, - }, - #[command(about = "Check Codex subscription usage for an account")] - Status { - #[arg(long, help = "Account label or id")] - name: Option, - #[arg(long, help = "Output as JSON")] - json: bool, - }, -} - -#[derive(Subcommand)] -enum CursorSubcommand { - #[command(about = "Login to Cursor with a browser session token")] - Login { - #[arg(long, help = "Label for this Cursor account (e.g., work, personal)")] - name: Option, - }, - #[command(about = "Logout from a Cursor account")] - Logout { - #[arg(long, help = "Account label or id")] - name: Option, - #[arg(long, help = "Logout from all Cursor accounts")] - all: bool, - #[arg(long, help = "Also delete cached Cursor usage")] - purge_cache: bool, - }, - #[command(about = "Check Cursor authentication status")] - Status { - #[arg(long, help = "Account label or id")] - name: Option, - }, - #[command(about = "List saved Cursor accounts")] - Accounts { - #[arg(long, help = "Output as JSON")] - json: bool, - }, - #[command(about = "Sync Cursor API usage into cursor-cache/usage*.csv")] - Sync { - #[arg(long, help = "Output as JSON")] - json: bool, - }, - #[command(about = "Switch active Cursor account")] - Switch { - #[arg(help = "Account label or id")] - name: String, - }, -} - -#[derive(Subcommand)] -enum AntigravitySubcommand { - #[command(about = "Sync usage from running Antigravity language servers")] - Sync, - #[command(about = "Show Antigravity sync status")] - Status { - #[arg(long, help = "Output as JSON")] - json: bool, - }, - #[command(about = "Delete cached Antigravity usage artifacts")] - PurgeCache, -} - -#[derive(Subcommand)] -enum TraeSubcommand { - #[command(about = "Authenticate Trae — auto-detect from desktop client or paste JWT")] - Login { - #[arg(long, help = "Paste access token directly (for manual fallback)")] - manual: bool, - #[arg(long, help = "Target Trae variant (solo, ide)")] - variant: Option, - }, - #[command(about = "Remove cached Trae credentials")] - Logout { - #[arg(long, help = "Target Trae variant (solo, ide)")] - variant: Option, - }, - #[command(about = "Show Trae authentication status")] - Status { - #[arg(long, help = "Output as JSON")] - json: bool, - }, - #[command(about = "Sync Trae usage data into local cache")] - Sync { - #[arg(long, help = "Number of days to sync (default: 30)")] - since: Option, - #[arg(long, help = "Include auxiliary usage types (not just main chat)")] - include_aux: bool, - }, -} - -#[derive(Subcommand)] -enum WarpSubcommand { - #[command(about = "Save Warp GraphQL authentication for aggregate usage sync")] - Login { - #[arg(long, help = "Warp bearer token or cookie header value")] - token: Option, - #[arg( - long, - help = "Treat token as a Cookie header instead of a bearer token" - )] - cookie: bool, - }, - #[command(about = "Remove cached Warp credentials")] - Logout { - #[arg(long, help = "Also delete cached Warp aggregate usage")] - purge_cache: bool, - }, - #[command(about = "Show Warp aggregate sync status")] - Status { - #[arg(long, help = "Output as JSON")] - json: bool, - }, - #[command(about = "Sync Warp aggregate usage into local cache")] - Sync { - #[arg(long, help = "Output as JSON")] - json: bool, - }, +impl From for ExecutionOutcome { + fn from(exit: tui::TuiExit) -> Self { + match exit { + tui::TuiExit::Quit => Self::Completed, + tui::TuiExit::Interrupted => Self::Interrupted, + } + } } -fn main() -> Result<()> { - use std::io::IsTerminal; - - let matches = Cli::command().get_matches(); - let cli = Cli::from_arg_matches(&matches).unwrap_or_else(|error| error.exit()); - let can_use_tui = std::io::stdin().is_terminal() && std::io::stdout().is_terminal(); - - match cli.command { - Some(Commands::Models { - json, - light, - clients, - date, - benchmark, - group_by, - write_cache, - no_write_cache, - no_spinner, - }) => { - use tokscale_core::GroupBy; - - let group_by: GroupBy = group_by.parse().unwrap_or_else(|e| { - eprintln!("Error: {}", e); - std::process::exit(1); - }); - let today = date.today; - let week = date.week; - let month = date.month; - let (since, until) = build_date_filter(today, week, month, date.since, date.until); - let year = normalize_year_filter(today, week, month, date.year); - let clients = build_client_filter(clients, &cli.home)?; - if json || light || !can_use_tui { - run_models_report( - json, - cli.home.clone(), - clients, - since, - until, - year, - benchmark, - no_spinner || !can_use_tui, - today, - week, - month, - group_by, - write_cache, - no_write_cache, - ) - } else { - ensure_home_supported_for_tui(&cli.home)?; - auto_sync_cursor_before_tui(&cli.home, &clients)?; - tui::run( - cli.theme.as_deref(), - cli.refresh, - cli.debug, - clients, - since, - until, - year, - Some(Tab::Models), - ) - } - } - Some(Commands::Monthly { - json, - light, - clients, - date, - benchmark, - no_spinner, - }) => { - let today = date.today; - let week = date.week; - let month = date.month; - let (since, until) = build_date_filter(today, week, month, date.since, date.until); - let year = normalize_year_filter(today, week, month, date.year); - let clients = build_client_filter(clients, &cli.home)?; - if json || light || !can_use_tui { - run_monthly_report( - json, - cli.home.clone(), - clients, - since, - until, - year, - benchmark, - no_spinner || !can_use_tui, - today, - week, - month, - ) - } else { - ensure_home_supported_for_tui(&cli.home)?; - auto_sync_cursor_before_tui(&cli.home, &clients)?; - tui::run( - cli.theme.as_deref(), - cli.refresh, - cli.debug, - clients, - since, - until, - year, - Some(Tab::Monthly), - ) - } - } - Some(Commands::Hourly { - json, - light, - clients, - date, - benchmark, - no_spinner, - }) => { - let today = date.today; - let week = date.week; - let month = date.month; - let (since, until) = build_date_filter(today, week, month, date.since, date.until); - let year = normalize_year_filter(today, week, month, date.year); - let clients = build_client_filter(clients, &cli.home)?; - if json || light || !can_use_tui { - run_hourly_report( - json, - cli.home.clone(), - clients, - since, - until, - year, - benchmark, - no_spinner || !can_use_tui, - today, - week, - month, - ) - } else { - ensure_home_supported_for_tui(&cli.home)?; - auto_sync_cursor_before_tui(&cli.home, &clients)?; - tui::run( - cli.theme.as_deref(), - cli.refresh, - cli.debug, - clients, - since, - until, - year, - Some(Tab::Hourly), - ) - } - } - Some(Commands::Pricing { - model_id, - json, - provider, - no_spinner, - }) => { - reject_unsupported_home_override(&cli.home, "pricing")?; - run_pricing_lookup(&model_id, json, provider.as_deref(), no_spinner) - } - Some(Commands::Clients { json }) => run_clients_command(json, cli.home.clone()), - Some(Commands::Graph { - output, - clients, - date, - benchmark, - no_spinner, - }) => { - let today = date.today; - let week = date.week; - let month = date.month; - let (since, until) = build_date_filter(today, week, month, date.since, date.until); - let year = normalize_year_filter(today, week, month, date.year); - let clients = build_client_filter(clients, &cli.home)?; - run_graph_command( - output, - cli.home.clone(), - clients, - since, - until, - year, - benchmark, - no_spinner, +fn execute(plan: ExecutionPlan) -> std::result::Result { + match plan { + ExecutionPlan::Tui(plan) => { + return tui::run( + plan.theme.as_deref(), + plan.refresh, + plan.no_refresh, + plan.debug, + plan.source.home, + plan.source.clients, + plan.date.since, + plan.date.until, + plan.date.year, + plan.initial_tab, ) - } - Some(Commands::Tui { clients, date }) => { - ensure_home_supported_for_tui(&cli.home)?; - let today = date.today; - let week = date.week; - let month = date.month; - let (since, until) = build_date_filter(today, week, month, date.since, date.until); - let year = normalize_year_filter(today, week, month, date.year); - let clients = build_client_filter(clients, &cli.home)?; - auto_sync_cursor_before_tui(&cli.home, &clients)?; - tui::run( - cli.theme.as_deref(), - cli.refresh, - cli.debug, - clients, - since, - until, - year, - None, - ) - } - Some(Commands::Headless { - source, - args, - format, - output, - no_auto_flags, - }) => { - reject_unsupported_home_override(&cli.home, "headless")?; - run_headless_command(&source, args, format, output, no_auto_flags) - } - Some(Commands::Wrapped { - output, - year, - client_flags, - short, - agents, - show_clients, - disable_pinned, - no_spinner: _, - }) => { - reject_unsupported_home_override(&cli.home, "wrapped")?; - let client_filter = build_client_filter(client_flags, &cli.home)?; - run_wrapped_command( - output, - year, - client_filter, - short, - agents, - show_clients, - disable_pinned, + .map(ExecutionOutcome::from) + .map_err(CliFailure::from); + } + ExecutionPlan::Models(plan) => { + let report = plan.report; + let no_spinner = effective_no_spinner(report.json, report.no_spinner); + run_models_report( + report.json, + report.source.home, + report.source.clients, + report.date.since, + report.date.until, + report.date.year, + report.benchmark, + no_spinner, + report.date.today, + report.date.week, + report.date.month, + plan.group_by, ) } - Some(Commands::Cursor { subcommand }) => { - reject_unsupported_home_override(&cli.home, "cursor")?; - run_cursor_command(subcommand) - } - Some(Commands::Antigravity { subcommand }) => { - reject_unsupported_home_override(&cli.home, "antigravity")?; - run_antigravity_command(subcommand) - } - Some(Commands::Usage { json }) => { - reject_unsupported_home_override(&cli.home, "usage")?; - reject_usage_parent_flags(&matches)?; - commands::usage::run(json) - } - Some(Commands::Cache { subcommand }) => { - reject_unsupported_home_override(&cli.home, "cache")?; - match subcommand { - CacheSubcommand::Prune => run_source_cache_prune(), - } - } - Some(Commands::Codex { subcommand }) => { - reject_unsupported_home_override(&cli.home, "codex")?; - run_codex_command(subcommand) - } - Some(Commands::Trae { subcommand }) => { - reject_unsupported_home_override(&cli.home, "trae")?; - run_trae_command(subcommand) - } - Some(Commands::Warp { subcommand }) => { - reject_unsupported_home_override(&cli.home, "warp")?; - run_warp_command(subcommand) - } - Some(Commands::TimeMetrics { - json, - clients, - date, - no_spinner, - }) => { - let today = date.today; - let week = date.week; - let month = date.month; - let (since, until) = build_date_filter(today, week, month, date.since, date.until); - let year = normalize_year_filter(today, week, month, date.year); - let clients = build_client_filter(clients, &cli.home)?; - run_time_metrics_report( + ExecutionPlan::Monthly(plan) => run_monthly_report( + plan.json, + plan.source.home, + plan.source.clients, + plan.date.since, + plan.date.until, + plan.date.year, + plan.benchmark, + effective_no_spinner(plan.json, plan.no_spinner), + plan.date.today, + plan.date.week, + plan.date.month, + ), + ExecutionPlan::Hourly(plan) => run_hourly_report( + plan.json, + plan.source.home, + plan.source.clients, + plan.date.since, + plan.date.until, + plan.date.year, + plan.benchmark, + effective_no_spinner(plan.json, plan.no_spinner), + plan.date.today, + plan.date.week, + plan.date.month, + ), + ExecutionPlan::TimeMetrics(plan) => run_time_metrics_report( + plan.json, + plan.source.home, + plan.source.clients, + plan.date.since, + plan.date.until, + plan.date.year, + plan.benchmark, + effective_no_spinner(plan.json, plan.no_spinner), + ), + ExecutionPlan::Clients(plan) => { + run_clients_command(plan.json, plan.source.home, plan.source.clients) + } + ExecutionPlan::Graph(plan) => run_graph_command( + plan.output, + plan.source.home, + plan.source.clients, + plan.date.since, + plan.date.until, + plan.date.year, + plan.benchmark, + plan.no_spinner, + ), + ExecutionPlan::Pricing(subcommand) => match subcommand { + PricingSubcommand::Lookup { + model_id, json, - cli.home.clone(), - clients, - since, - until, - year, + source, no_spinner, - ) - } - Some(Commands::WarmTuiCache) => run_warm_tui_cache(), - None => { - let today = cli.date.today; - let week = cli.date.week; - let month = cli.date.month; - let clients = build_client_filter(cli.clients, &cli.home)?; - let (since, until) = - build_date_filter(today, week, month, cli.date.since, cli.date.until); - let year = normalize_year_filter(today, week, month, cli.date.year); - let group_by: tokscale_core::GroupBy = cli.group_by.parse().unwrap_or_else(|e| { - eprintln!("Error: {}", e); - std::process::exit(1); - }); - - if cli.json { - run_models_report( - cli.json, - cli.home.clone(), - clients, - since, - until, - year, - cli.benchmark, - cli.no_spinner || cli.json, - today, - week, - month, - group_by, - cli.write_cache, - cli.no_write_cache, - ) - } else if cli.light || !can_use_tui { - run_models_report( - false, - cli.home.clone(), - clients, - since, - until, - year, - cli.benchmark, - cli.no_spinner || !can_use_tui, - today, - week, - month, - group_by, - cli.write_cache, - cli.no_write_cache, - ) - } else { - ensure_home_supported_for_tui(&cli.home)?; - auto_sync_cursor_before_tui(&cli.home, &clients)?; - tui::run( - cli.theme.as_deref(), - cli.refresh, - cli.debug, - clients, - since, - until, - year, - None, - ) - } - } - } -} - -#[derive(Args, Clone, Debug, Default)] -pub struct ClientFlags { - /// Canonical client filter. Repeatable or comma-separated. - /// Example: `--client opencode,claude` or `-c opencode -c claude`. - #[arg( - long = "client", - short = 'c', - value_parser = parse_client_id_arg, - value_delimiter = ',', - action = clap::ArgAction::Append, - help = "Filter by client(s). Repeatable or comma-separated (e.g. -c opencode,claude)." - )] - pub clients: Vec, + } => run_pricing_lookup( + &model_id, + json, + source.map(PricingSource::as_str), + effective_no_spinner(json, no_spinner), + ), + PricingSubcommand::Overrides { json } => run_pricing_list_overrides(json), + }, + ExecutionPlan::Usage { json } => commands::usage::run(json), + ExecutionPlan::Wrapped(plan) => run_wrapped_command(plan), + ExecutionPlan::Headless(args) => run_headless_command( + args.source.as_str(), + args.command, + args.format.map(HeadlessFormat::as_str), + args.output, + args.no_auto_flags, + args.timeout, + ), + ExecutionPlan::CachePrune => run_source_cache_prune(), + ExecutionPlan::CacheWarm(source) => run_warm_tui_cache(source.home, source.clients), + ExecutionPlan::Antigravity(subcommand) => run_antigravity_command(subcommand), + ExecutionPlan::Warp(subcommand) => run_warp_command(subcommand), + }?; + + Ok(ExecutionOutcome::Completed) } -#[derive(Args, Clone, Debug, Default)] -pub struct DateRangeFlags { - #[arg(long, help = "Show only today's usage")] - pub today: bool, - #[arg(long, help = "Show last 7 days")] - pub week: bool, - #[arg(long, help = "Show current month")] - pub month: bool, - #[arg(long, help = "Start date (YYYY-MM-DD)")] - pub since: Option, - #[arg(long, help = "End date (YYYY-MM-DD)")] - pub until: Option, - #[arg(long, help = "Filter by year (YYYY)")] - pub year: Option, +const fn effective_no_spinner(json: bool, explicit_no_spinner: bool) -> bool { + json || explicit_no_spinner } -fn run_wrapped_command( - output: Option, - year: Option, - client_filter: Option>, - short: bool, - agents: bool, - show_clients: bool, - disable_pinned: bool, -) -> Result<()> { +fn run_wrapped_command(plan: WrappedPlan) -> Result<()> { use colored::Colorize; - println!("{}", "\n Tokscale - Generate Wrapped Image\n".cyan()); - - println!("{}", " Generating wrapped image...".bright_black()); - println!(); + if !plan.no_spinner { + eprintln!("{}", "Generating wrapped image...".bright_black()); + } - let include_agents = !show_clients || agents; let wrapped_options = commands::wrapped::WrappedOptions { - output, - year, - clients: client_filter, - short, - include_agents, - pin_sisyphus: !disable_pinned, + output: plan.output, + year: plan.year, + home_dir: plan.source.home, + clients: plan.source.clients, + short: plan.short, + ranking: plan.ranking, + pin_sisyphus: !plan.disable_pinned, }; - match commands::wrapped::run(wrapped_options) { - Ok(output_path) => { - println!( - "{}", - format!("\n ✓ Generated wrapped image: {}\n", output_path).green() - ); - } - Err(err) => { - eprintln!("{}", "\nError generating wrapped image:".red()); - eprintln!(" {}\n", err); - std::process::exit(1); - } - } - + let output_path = commands::wrapped::run(wrapped_options)?; + println!("{output_path}"); Ok(()) } diff --git a/crates/tokscale-cli/src/main_tests.rs b/crates/tokscale-cli/src/main_tests.rs index b8e795eb7..9daa2f167 100644 --- a/crates/tokscale-cli/src/main_tests.rs +++ b/crates/tokscale-cli/src/main_tests.rs @@ -1,45 +1,23 @@ -use super::*; -use crate::commands::cache::*; +use crate::cli::*; use crate::commands::clients::*; -use crate::commands::integrations::*; use crate::commands::render::*; use crate::commands::shared::*; use clap::Parser; use std::path::{Path, PathBuf}; +use tokscale_core::ClientId; #[test] -fn test_parse_variant_arg_accepts_known_values() { +fn tui_exit_maps_to_process_execution_outcome() { assert_eq!( - parse_variant_arg(Some("solo")).unwrap(), - Some(trae::auth::TraeVariant::Solo) + super::ExecutionOutcome::from(crate::tui::TuiExit::Quit), + super::ExecutionOutcome::Completed ); assert_eq!( - parse_variant_arg(Some("ide")).unwrap(), - Some(trae::auth::TraeVariant::Ide) + super::ExecutionOutcome::from(crate::tui::TuiExit::Interrupted), + super::ExecutionOutcome::Interrupted ); } -#[test] -fn test_parse_variant_arg_none_when_omitted() { - assert_eq!(parse_variant_arg(None).unwrap(), None); -} - -#[test] -fn test_parse_variant_arg_rejects_unknown_value() { - // The earlier `Option`-returning version converted this to `None` - // and the caller fell through to "all variants" — a typo like - // `--variant slo` would log out every variant. Now we error out. - let err = parse_variant_arg(Some("slo")).unwrap_err(); - let msg = err.to_string(); - assert!(msg.contains("unknown variant"), "got: {msg}"); - assert!(msg.contains("slo"), "got: {msg}"); -} - -#[test] -fn test_parse_variant_arg_rejects_empty_string() { - assert!(parse_variant_arg(Some("")).is_err()); -} - // Tests below call `build_client_filter_with_defaults` directly with // an explicit `defaults` slice instead of `build_client_filter`, which // reads from `~/.config/tokscale/settings.json`. @@ -62,6 +40,14 @@ fn test_parse_client_id_arg_rejects_unknown_ids() { ); } +#[test] +fn removed_clients_are_not_valid_source_ids() { + for client in ["cursor", "trae"] { + let error = parse_client_id_arg(client).unwrap_err(); + assert!(error.contains(client), "unexpected error: {error}"); + } +} + #[test] fn test_build_client_filter_no_flags_no_defaults_returns_none() { let flags = ClientFlags::default(); @@ -98,67 +84,6 @@ fn test_build_client_filter_canonical_dedups_repeats() { ); } -#[test] -fn test_resolve_default_tui_filter_set_uses_configured_defaults() { - // When `defaultClients` is set, the warm-cache resolver must use - // it verbatim — otherwise the warm cache would store every real - // client while the next no-flag TUI launch wants only the configured - // ones, producing a guaranteed cache miss. - let configured = vec!["opencode".to_string(), "claude".to_string()]; - let set = resolve_default_tui_filter_set_with(&configured).unwrap(); - let mut expected = std::collections::HashSet::new(); - expected.insert(ClientId::OpenCode); - expected.insert(ClientId::Claude); - assert_eq!(set, expected); -} - -#[test] -fn test_resolve_default_tui_filter_set_uses_catalog_when_empty() { - // No defaultClients configured → use the complete accepted catalog. - let set = resolve_default_tui_filter_set_with(&[]).unwrap(); - let expected = ClientId::iter().collect(); - assert_eq!(set, expected); - assert!(set.contains(&ClientId::Cursor)); -} - -#[test] -fn test_light_cache_no_filter_uses_catalog() { - let set = resolve_light_cache_filter_set(&None); - let expected = ClientId::iter().collect(); - assert_eq!(set, expected); - assert!(set.contains(&ClientId::Cursor)); -} - -#[test] -fn test_resolve_default_tui_filter_set_rejects_unknown_ids() { - let configured = vec!["opencode".to_string(), "not-a-real-client".to_string()]; - let err = resolve_default_tui_filter_set_with(&configured).unwrap_err(); - assert!( - err.to_string().contains("not-a-real-client"), - "unexpected error: {err}" - ); -} - -#[test] -fn test_resolve_default_tui_filter_set_rejects_all_unknown_ids() { - let configured = vec!["not-real".to_string(), "also-fake".to_string()]; - let err = resolve_default_tui_filter_set_with(&configured).unwrap_err(); - assert!( - err.to_string().contains("not-real, also-fake"), - "unexpected error: {err}" - ); -} - -#[test] -fn test_resolve_default_tui_filter_set_rejects_removed_synthetic_id() { - let configured = vec!["claude".to_string(), "synthetic".to_string()]; - let err = resolve_default_tui_filter_set_with(&configured).unwrap_err(); - assert!( - err.to_string().contains("synthetic"), - "unexpected error: {err}" - ); -} - #[test] fn test_build_client_filter_with_defaults_when_no_flags() { // No CLI flags + a defaultClients list → defaults apply. @@ -228,68 +153,109 @@ fn test_build_client_filter_defaults_dedup_preserves_order() { fn test_client_flags_parses_canonical_form() { // End-to-end smoke test: ensure clap derives accept the new // `--client a,b` and `-c a -c b` shapes through the CLI parser. - let cli = Cli::try_parse_from(["tokscale", "--client", "opencode,claude"]).expect("parse ok"); + let cli = Cli::try_parse_from(["tokscale", "models", "--client", "opencode,claude"]) + .expect("parse ok"); + let Some(Commands::Models(args)) = cli.command else { + panic!("expected models command"); + }; assert_eq!( - cli.clients.clients, + args.report.source.clients.clients, vec![ClientId::OpenCode, ClientId::Claude] ); - let cli = - Cli::try_parse_from(["tokscale", "-c", "opencode", "-c", "claude"]).expect("parse ok"); + let cli = Cli::try_parse_from(["tokscale", "tui", "-c", "opencode", "-c", "claude"]) + .expect("parse ok"); + let Some(Commands::Tui(args)) = cli.command else { + panic!("expected tui command"); + }; assert_eq!( - cli.clients.clients, + args.source.clients.clients, vec![ClientId::OpenCode, ClientId::Claude] ); } #[test] -fn test_wrapped_parses_clients_view_flag() { +fn wrapped_ranking_is_one_typed_selection() { let cli = Cli::try_parse_from(["tokscale", "wrapped"]).expect("parse ok"); - let Some(Commands::Wrapped { - show_clients, - agents, - .. - }) = cli.command - else { + let Some(Commands::Wrapped(args)) = cli.command else { panic!("expected wrapped command"); }; - assert!(!show_clients); - assert!(!agents); + assert_eq!(args.ranking, None); - let cli = Cli::try_parse_from(["tokscale", "wrapped", "--clients"]).expect("parse ok"); - let Some(Commands::Wrapped { show_clients, .. }) = cli.command else { + let cli = Cli::try_parse_from(["tokscale", "wrapped", "--ranking", "agents"]) + .expect("agents ranking parses"); + let Some(Commands::Wrapped(args)) = cli.command else { panic!("expected wrapped command"); }; - assert!(show_clients); + assert_eq!(args.ranking, Some(WrappedRankingArg::Agents)); + + let cli = Cli::try_parse_from(["tokscale", "wrapped", "--ranking", "clients"]) + .expect("clients ranking parses"); + let Some(Commands::Wrapped(args)) = cli.command else { + panic!("expected wrapped command"); + }; + assert_eq!(args.ranking, Some(WrappedRankingArg::Clients)); + + for removed in ["--agents", "--clients"] { + assert!(Cli::try_parse_from(["tokscale", "wrapped", removed]).is_err()); + } } #[test] -fn test_wrapped_client_filter_coexists_with_clients_view_flag() { - let cli = - Cli::try_parse_from(["tokscale", "wrapped", "--client", "opencode"]).expect("parse ok"); - let Some(Commands::Wrapped { - client_flags, - show_clients, - .. - }) = cli.command - else { - panic!("expected wrapped command"); +fn wrapped_ranking_resolves_without_boolean_precedence() { + let resolve = |args: &[&str]| { + let cli = Cli::try_parse_from(args).expect("wrapped arguments parse"); + ExecutionPlan::resolve( + cli, + TerminalState { + stdin: false, + stdout: false, + }, + ) }; - assert_eq!(client_flags.clients, vec![ClientId::OpenCode]); - assert!(!show_clients); - let cli = Cli::try_parse_from(["tokscale", "wrapped", "--clients", "--client", "opencode"]) - .expect("parse ok"); - let Some(Commands::Wrapped { - client_flags, - show_clients, - .. - }) = cli.command + let ExecutionPlan::Wrapped(plan) = + resolve(&["tokscale", "wrapped"]).expect("default ranking resolves") else { - panic!("expected wrapped command"); + panic!("expected wrapped plan"); + }; + assert_eq!(plan.ranking, WrappedRanking::Auto); + + let ExecutionPlan::Wrapped(plan) = resolve(&[ + "tokscale", + "wrapped", + "--ranking", + "agents", + "--client", + "opencode", + ]) + .expect("agents ranking resolves") else { + panic!("expected wrapped plan"); }; - assert_eq!(client_flags.clients, vec![ClientId::OpenCode]); - assert!(show_clients); + assert_eq!(plan.ranking, WrappedRanking::Agents); + + let error = resolve(&[ + "tokscale", + "wrapped", + "--ranking", + "agents", + "--client", + "claude", + ]) + .expect_err("agents ranking without OpenCode must fail during resolve"); + assert_eq!(error.exit_code(), 2); + assert!(error.to_string().contains("requires `opencode`")); + + let error = resolve(&[ + "tokscale", + "wrapped", + "--ranking", + "clients", + "--disable-pinned", + ]) + .expect_err("client ranking cannot accept an ignored agent option"); + assert_eq!(error.exit_code(), 2); + assert!(error.to_string().contains("does not apply")); } #[test] @@ -300,25 +266,31 @@ fn test_legacy_client_flags_are_removed() { #[test] fn test_client_flag_accepts_uppercase() { - let cli = Cli::try_parse_from(["tokscale", "--client", "OPENCODE"]).expect("uppercase parses"); - assert_eq!(cli.clients.clients, vec![ClientId::OpenCode]); + let cli = Cli::try_parse_from(["tokscale", "models", "--client", "OPENCODE"]) + .expect("uppercase parses"); + let Some(Commands::Models(args)) = cli.command else { + panic!("expected models command"); + }; + assert_eq!(args.report.source.clients.clients, vec![ClientId::OpenCode]); - let cli = - Cli::try_parse_from(["tokscale", "-c", "Codebuff,Antigravity"]).expect("mixed-case parses"); + let cli = Cli::try_parse_from(["tokscale", "models", "-c", "Codebuff,Antigravity"]) + .expect("mixed-case parses"); + let Some(Commands::Models(args)) = cli.command else { + panic!("expected models command"); + }; assert_eq!( - cli.clients.clients, + args.report.source.clients.clients, vec![ClientId::Codebuff, ClientId::Antigravity] ); } #[test] fn test_client_flag_rejects_unknown_and_empty_values() { - assert!(Cli::try_parse_from(["tokscale", "--client", "unknown"]).is_err()); - assert!(Cli::try_parse_from(["tokscale", "--client", ""]).is_err()); + assert!(Cli::try_parse_from(["tokscale", "models", "--client", "unknown"]).is_err()); + assert!(Cli::try_parse_from(["tokscale", "models", "--client", ""]).is_err()); - let error = Cli::try_parse_from(["tokscale", "--client", "crush"]) - .err() - .expect("excluded clients must not remain valid CLI values") + let error = Cli::try_parse_from(["tokscale", "models", "--client", "crush"]) + .expect_err("excluded clients must not remain valid CLI values") .to_string(); assert!(error.contains("invalid client id `crush`"), "{error}"); assert!(!error.contains("does not support local parsing"), "{error}"); @@ -326,36 +298,51 @@ fn test_client_flag_rejects_unknown_and_empty_values() { #[test] fn test_home_arg_rejects_empty_and_blank_values() { - assert!(Cli::try_parse_from(["tokscale", "--home", "", "--light"]).is_err()); - assert!(Cli::try_parse_from(["tokscale", "--home", " ", "--light"]).is_err()); + assert!(Cli::try_parse_from(["tokscale", "models", "--home", ""]).is_err()); + assert!(Cli::try_parse_from(["tokscale", "models", "--home", " "]).is_err()); } #[test] -fn test_pricing_provider_accepts_known_values_case_insensitive() { - let cli = Cli::try_parse_from(["tokscale", "pricing", "gpt-4o", "--provider", "OpenRouter"]) - .expect("provider parses"); - let Some(Commands::Pricing { provider, .. }) = cli.command else { +fn test_pricing_source_accepts_known_values() { + let cli = Cli::try_parse_from([ + "tokscale", + "pricing", + "lookup", + "gpt-4o", + "--source", + "openrouter", + ]) + .expect("source parses"); + let Some(Commands::Pricing { + subcommand: PricingSubcommand::Lookup { source, .. }, + }) = cli.command + else { panic!("expected pricing command"); }; - assert_eq!(provider.as_deref(), Some("OpenRouter")); + assert_eq!(source, Some(PricingSource::Openrouter)); } #[test] -fn test_pricing_provider_rejects_unknown_values() { - assert!( - Cli::try_parse_from(["tokscale", "pricing", "gpt-4o", "--provider", "unknown",]).is_err() - ); +fn test_pricing_source_rejects_unknown_values() { assert!(Cli::try_parse_from([ - "tokscale", - "pricing", - "gpt-4o", - "--json", - "--provider", - "unknown", + "tokscale", "pricing", "lookup", "gpt-4o", "--source", "unknown", ]) .is_err()); } +#[test] +fn pricing_overrides_is_the_registered_user_facing_subcommand() { + let cli = Cli::try_parse_from(["tokscale", "pricing", "overrides", "--json"]) + .expect("documented pricing overrides command must parse"); + + assert!(matches!( + cli.command, + Some(Commands::Pricing { + subcommand: PricingSubcommand::Overrides { json: true } + }) + )); +} + #[test] fn test_build_client_filter_with_defaults_empty_defaults_returns_none() { let flags = ClientFlags::default(); @@ -644,70 +631,269 @@ fn test_light_spinner_scanner_state_cycle_wrap() { } #[test] -fn resolve_cli_write_overrides_settings_false() { - let settings = tui::settings::Settings { - light: tui::settings::LightSettings { write_cache: false }, - ..tui::settings::Settings::default() - }; - assert!(resolve_should_write_cache(true, false, &settings)); -} - -#[test] -fn resolve_cli_no_write_overrides_settings_true() { - let settings = tui::settings::Settings { - light: tui::settings::LightSettings { write_cache: true }, - ..tui::settings::Settings::default() - }; - assert!(!resolve_should_write_cache(false, true, &settings)); +fn root_rejects_business_options() { + for args in [ + vec!["tokscale", "--json"], + vec!["tokscale", "--light"], + vec!["tokscale", "--client", "codex"], + vec!["tokscale", "--week"], + vec!["tokscale", "--json", "models"], + ] { + assert!(Cli::try_parse_from(args).is_err()); + } } #[test] -fn resolve_settings_true_with_no_cli_flag() { - let settings = tui::settings::Settings { - light: tui::settings::LightSettings { write_cache: true }, - ..tui::settings::Settings::default() +fn legacy_v4_invocations_get_one_migration_hint_without_becoming_aliases() { + let strings = |values: &[&str]| { + values + .iter() + .map(|value| (*value).to_string()) + .collect::>() }; - assert!(resolve_should_write_cache(false, false, &settings)); -} -#[test] -fn resolve_settings_false_with_no_cli_flag() { - let settings = tui::settings::Settings { - light: tui::settings::LightSettings { write_cache: false }, - ..tui::settings::Settings::default() - }; - assert!(!resolve_should_write_cache(false, false, &settings)); + assert_eq!( + legacy_invocation_hint(&strings(&["--json", "models"])).as_deref(), + Some("use `tokscale models --json`") + ); + assert_eq!( + legacy_invocation_hint(&strings(&["--light"])).as_deref(), + Some("use `tokscale models`") + ); + assert_eq!( + legacy_invocation_hint(&strings(&["models", "--light"])).as_deref(), + Some("use `tokscale models`") + ); + assert_eq!( + legacy_invocation_hint(&strings(&["--client", "codex"])).as_deref(), + Some("use `tokscale tui --client codex`") + ); + assert_eq!( + legacy_invocation_hint(&strings(&["--client=codex"])).as_deref(), + Some("use `tokscale tui --client=codex`") + ); + assert_eq!( + legacy_invocation_hint(&strings(&["tui", "--json"])).as_deref(), + Some("use `tokscale models --json`") + ); + assert_eq!( + legacy_invocation_hint(&strings(&["graph", "--json"])).as_deref(), + Some("use `tokscale graph`") + ); + assert_eq!( + legacy_invocation_hint(&strings(&["--json", "graph"])).as_deref(), + Some("use `tokscale graph`") + ); + assert_eq!( + legacy_invocation_hint(&strings(&["graph", "--json", "--output", "graph.json"])).as_deref(), + Some("use `tokscale graph --output graph.json`") + ); + assert_eq!( + legacy_invocation_hint(&strings(&["--client=codex", "models"])).as_deref(), + Some("use `tokscale models --client=codex`") + ); + assert_eq!( + legacy_invocation_hint(&strings(&["pricing", "list-overrides"])).as_deref(), + Some("use `tokscale pricing overrides`") + ); + assert_eq!( + legacy_invocation_hint(&strings(&["pricing", "list-overrides", "--json"])).as_deref(), + Some("use `tokscale pricing overrides --json`") + ); + assert_eq!( + legacy_invocation_hint(&strings(&["pricing", "gpt-5", "--json"])).as_deref(), + Some("use `tokscale pricing lookup gpt-5 --json`") + ); + assert_eq!( + legacy_invocation_hint(&strings(&["--json", "pricing", "gpt-5"])), + None, + "migration hints must never suggest another invalid invocation" + ); + assert_eq!( + legacy_invocation_hint(&strings(&["headless", "codex", "--", "tui", "--json"])), + None, + "child-process arguments after -- must not influence migration hints" + ); + for unrelated in [ + &["wrapped", "--json"][..], + &["clients", "--benchmark"], + &["pricing", "--json"], + &["graph", "--group-by", "model"], + ] { + assert_eq!( + legacy_invocation_hint(&strings(unrelated)), + None, + "migration hints must not change an explicit command's product" + ); + } } #[test] -fn resolve_settings_default_returns_false() { - assert!(!resolve_should_write_cache( - false, - false, - &tui::settings::Settings::default() - )); +fn misplaced_and_equals_form_options_remain_parse_errors() { + for args in [ + vec!["tokscale", "tui", "--json"], + vec!["tokscale", "--client=codex"], + vec!["tokscale", "graph", "--json"], + ] { + let error = Cli::try_parse_from(args).expect_err("legacy invocation must be rejected"); + assert_eq!(error.exit_code(), 2); + } } #[test] -fn clap_rejects_write_cache_without_light() { - assert!(Cli::try_parse_from(["tokscale", "--write-cache"]).is_err()); +fn report_execution_plan_does_not_depend_on_terminal_state() { + for terminal in [ + TerminalState { + stdin: true, + stdout: true, + }, + TerminalState { + stdin: false, + stdout: false, + }, + ] { + let cli = Cli::try_parse_from(["tokscale", "models", "--client", "opencode", "--json"]) + .expect("models command parses"); + let plan = ExecutionPlan::resolve(cli, terminal).expect("models plan resolves"); + let ExecutionPlan::Models(plan) = plan else { + panic!("models must never resolve to a TUI plan"); + }; + assert!(plan.report.json); + assert_eq!( + plan.report.source.clients, + Some(vec!["opencode".to_string()]) + ); + } } #[test] -fn clap_rejects_no_write_cache_without_light() { - assert!(Cli::try_parse_from(["tokscale", "--no-write-cache"]).is_err()); -} +fn json_report_plan_preserves_explicit_no_spinner() { + let home = tempfile::TempDir::new().unwrap(); + let home = home.path().to_str().unwrap(); + let resolve = |explicit_no_spinner: bool| { + let mut argv = vec!["tokscale", "models", "--home", home, "--json"]; + if explicit_no_spinner { + argv.push("--no-spinner"); + } + let cli = Cli::try_parse_from(argv).expect("models command parses"); + let plan = ExecutionPlan::resolve( + cli, + TerminalState { + stdin: false, + stdout: false, + }, + ) + .expect("models plan resolves"); + let ExecutionPlan::Models(plan) = plan else { + panic!("expected models plan"); + }; + plan.report + }; -#[test] -fn clap_rejects_both_write_flags_together() { - assert!( - Cli::try_parse_from(["tokscale", "--light", "--write-cache", "--no-write-cache",]).is_err() - ); + let implicit = resolve(false); + let explicit = resolve(true); + assert!(implicit.json); + assert!(explicit.json); + assert!(!implicit.no_spinner); + assert!(explicit.no_spinner); +} + +#[test] +fn effective_spinner_policy_keeps_json_quiet_without_erasing_explicit_intent() { + assert!(!super::effective_no_spinner(false, false)); + assert!(super::effective_no_spinner(false, true)); + assert!(super::effective_no_spinner(true, false)); + assert!(super::effective_no_spinner(true, true)); +} + +#[test] +fn headless_plan_rejects_a_blank_child_executable() { + let cli = Cli::try_parse_from(["tokscale", "headless", "codex", "--", ""]) + .expect("Clap accepts the present but empty COMMAND token"); + let error = ExecutionPlan::resolve( + cli, + TerminalState { + stdin: false, + stdout: false, + }, + ) + .expect_err("a blank executable must not enter the execution plan"); + + assert_eq!(error.exit_code(), 2); + assert!(error.to_string().contains("non-empty executable")); +} + +#[test] +fn tui_execution_plan_requires_both_interactive_streams() { + for terminal in [ + TerminalState { + stdin: false, + stdout: true, + }, + TerminalState { + stdin: true, + stdout: false, + }, + ] { + let cli = Cli::try_parse_from(["tokscale", "tui"]).expect("TUI command parses"); + let error = ExecutionPlan::resolve(cli, terminal).expect_err("non-TTY TUI must fail"); + assert_eq!(error.exit_code(), 2); + assert!(error.to_string().contains("interactive terminal")); + } } #[test] -fn clap_accepts_models_light_write_cache_after_subcommand() { - assert!(Cli::try_parse_from(["tokscale", "models", "--light", "--write-cache"]).is_ok()); +fn tui_execution_plan_rejects_disabled_optional_tab() { + let home = tempfile::TempDir::new().unwrap(); + let cli = Cli::try_parse_from([ + "tokscale", + "tui", + "--home", + home.path().to_str().unwrap(), + "--tab", + "usage", + ]) + .expect("TUI command parses"); + let error = ExecutionPlan::resolve( + cli, + TerminalState { + stdin: true, + stdout: true, + }, + ) + .expect_err("disabled explicit tab must fail before entering the TUI"); + assert_eq!(error.exit_code(), 2); + assert!(error.to_string().contains("disabled in settings.json")); +} + +#[test] +fn resolve_rejects_reversed_custom_date_range() { + let cli = Cli::try_parse_from([ + "tokscale", + "models", + "--since", + "2026-07-15", + "--until", + "2026-07-14", + ]) + .expect("individually valid dates parse"); + let error = ExecutionPlan::resolve( + cli, + TerminalState { + stdin: false, + stdout: false, + }, + ) + .expect_err("reversed range must fail"); + assert_eq!(error.exit_code(), 2); + assert!(error.to_string().contains("must not be later")); +} + +#[test] +fn removed_report_and_cache_flags_are_rejected() { + for flag in ["--light", "--write-cache", "--no-write-cache"] { + assert!(Cli::try_parse_from(["tokscale", "models", flag]).is_err()); + } } #[test] @@ -722,9 +908,24 @@ fn clap_accepts_source_cache_prune_command() { } #[test] -fn clap_accepts_cursor_sync_command() { - assert!(Cli::try_parse_from(["tokscale", "cursor", "sync"]).is_ok()); - assert!(Cli::try_parse_from(["tokscale", "cursor", "sync", "--json"]).is_ok()); +fn clap_accepts_explicit_cache_warm_scope() { + let cli = Cli::try_parse_from(["tokscale", "cache", "warm", "--client", "codex"]) + .expect("cache warm parses"); + assert!(matches!( + cli.command, + Some(Commands::Cache { + subcommand: CacheSubcommand::Warm { .. } + }) + )); +} + +#[test] +fn cli_rejects_removed_account_management_namespaces() { + assert!(Cli::try_parse_from(["tokscale", "cursor", "sync"]).is_err()); + assert!(Cli::try_parse_from(["tokscale", "cursor", "logout", "--all"]).is_err()); + assert!(Cli::try_parse_from(["tokscale", "codex", "accounts"]).is_err()); + assert!(Cli::try_parse_from(["tokscale", "codex", "switch", "work"]).is_err()); + assert!(Cli::try_parse_from(["tokscale", "trae", "status"]).is_err()); } #[test] @@ -742,41 +943,6 @@ fn clap_accepts_usage_without_light_flag() { assert!(Cli::try_parse_from(["tokscale", "usage", "--light"]).is_err()); } -#[test] -fn usage_rejects_parent_flags() { - for flag in USAGE_PARENT_FLAGS { - let args = usage_parent_flag_test_args(flag.display); - let matches = Cli::command().try_get_matches_from(args).unwrap(); - let error = reject_usage_parent_flags(&matches) - .expect_err("usage should reject parent CLI flags") - .to_string(); - assert!( - error.contains(flag.display), - "expected error to mention {} but got {error}", - flag.display - ); - } - - let matches = Cli::command() - .try_get_matches_from(["tokscale", "usage", "--json"]) - .unwrap(); - assert!(reject_usage_parent_flags(&matches).is_ok()); -} - -fn usage_parent_flag_test_args(flag: &'static str) -> Vec<&'static str> { - match flag { - "--write-cache" | "--no-write-cache" => vec!["tokscale", "--light", flag, "usage"], - "--client" => vec!["tokscale", flag, "claude", "usage"], - "--since" => vec!["tokscale", flag, "2026-01-01", "usage"], - "--until" => vec!["tokscale", flag, "2026-01-02", "usage"], - "--year" => vec!["tokscale", flag, "2026", "usage"], - "--group-by" => vec!["tokscale", flag, "model", "usage"], - "--theme" => vec!["tokscale", flag, "red", "usage"], - "--refresh" => vec!["tokscale", flag, "1", "usage"], - _ => vec!["tokscale", flag, "usage"], - } -} - #[test] fn client_id_parses_warp() { assert_eq!(ClientId::from_str("warp"), Some(ClientId::Warp)); @@ -791,8 +957,8 @@ fn client_id_parses_grok() { #[test] fn clap_rejects_antigravity_cli_as_separate_client() { - assert!(Cli::try_parse_from(["tokscale", "--client", "antigravity"]).is_ok()); - assert!(Cli::try_parse_from(["tokscale", "--client", "antigravity-cli"]).is_err()); + assert!(Cli::try_parse_from(["tokscale", "models", "--client", "antigravity"]).is_ok()); + assert!(Cli::try_parse_from(["tokscale", "models", "--client", "antigravity-cli"]).is_err()); } #[test] @@ -860,60 +1026,3 @@ fn headless_roots_trim_env_override() { None => unsafe { std::env::remove_var("TOKSCALE_HEADLESS_DIR") }, } } - -#[test] -fn cursor_auto_sync_enabled_for_default_report() { - assert!(should_auto_sync_cursor_for_local_report(&None, &None)); -} - -#[test] -fn cursor_auto_sync_enabled_when_cursor_filter_is_explicit() { - assert!(should_auto_sync_cursor_for_local_report( - &None, - &Some(vec!["cursor".to_string()]) - )); -} - -#[test] -fn cursor_auto_sync_disabled_when_filter_excludes_cursor() { - assert!(!should_auto_sync_cursor_for_local_report( - &None, - &Some(vec!["codex".to_string()]) - )); -} - -#[test] -fn cursor_auto_sync_disabled_for_home_override() { - assert!(!should_auto_sync_cursor_for_local_report( - &Some("/tmp/other-home".to_string()), - &None - )); - assert!(!should_auto_sync_cursor_for_local_report( - &Some("/tmp/other-home".to_string()), - &Some(vec!["cursor".to_string()]) - )); -} - -#[test] -fn cursor_auto_sync_runtime_init_failure_is_best_effort() { - let result = run_best_effort_cursor_sync_with_runtime_factory(|| { - Err(std::io::Error::other("runtime unavailable")) - }); - - assert!(!result.synced); - assert_eq!(result.rows, 0); - assert!(result - .error - .as_deref() - .is_some_and(|error| error.contains("runtime unavailable"))); -} - -#[test] -fn light_cache_write_allows_default_home() { - assert!(can_write_light_cache(&None)); -} - -#[test] -fn light_cache_write_refuses_when_home_dir_set() { - assert!(!can_write_light_cache(&Some("/tmp/fake-home".to_string()))); -} diff --git a/crates/tokscale-cli/src/trae.rs b/crates/tokscale-cli/src/trae.rs deleted file mode 100644 index dd09d807f..000000000 --- a/crates/tokscale-cli/src/trae.rs +++ /dev/null @@ -1,1558 +0,0 @@ -//! Trae (ByteDance AI IDE) integration: credential decryption, token refresh, -//! usage sync, and CLI commands. -//! -//! Two international variants: -//! -//! | Variant | Product line | App Support directory | -//! |---------|--------------|-----------------------| -//! | Solo | Solo | TRAE SOLO | -//! | Ide | IDE | Trae | -//! -//! Variants are credential sources only. Trae IDE and Trae Solo expose the same -//! account-level usage data through the international usage API, so synced -//! reports are stored once under tokscale's single `trae` client. -//! -//! The China variants (trae.com.cn) are intentionally not integrated: -//! the CN backend does not expose a session-level usage query API. -//! They will be added if/when upstream releases an official endpoint. - -pub mod auth { - //! Trae credential management with automatic token refresh. - //! - //! Two international Trae variants: - //! - //! | Variant | Product line | App Support directory | credential label | - //! |---------|--------------|-----------------------|------------------| - //! | Solo | Solo | TRAE SOLO | trae-solo | - //! | Ide | IDE | Trae | trae | - //! - //! Shared backend (api-sg-central.trae.ai); variants only decide where - //! credentials are discovered. Usage sync stores account-level API results - //! once under the single `trae` report client. - //! - //! The China variants (trae.com.cn) are intentionally not integrated: - //! the CN backend does not expose a session-level usage query API. - //! They will be added if/when upstream releases an official endpoint. - //! - //! Credential lifecycle (highest to lowest priority): - //! 1. Cached `access_token` still valid → use it directly. - //! 2. Cached `refresh_token` still valid → call `ExchangeToken` to mint a new - //! pair and write it back to disk. - //! 3. `refresh_token` also expired → decrypt the Trae desktop client's - //! `storage.json` (`iCubeAuthInfo` entry). - //! 4. Decryption fails / no `storage.json` → fall back to - //! `trae login --manual` (paste a JWT). - //! - //! Cache filenames: `credentials-solo.json` (Solo) and `credentials-ide.json` (Ide). - //! These are fixed names, not derived from the report client id (`client_str()`). - //! Cache directory: `/trae-cache/`, where the - //! config dir is resolved by [`paths::get_config_dir`] and honors - //! `TOKSCALE_CONFIG_DIR` plus XDG defaults (typically - //! `~/.config/tokscale` on Linux/macOS). - - use crate::trae::safestorage; - use anyhow::{Context, Result}; - use base64::Engine; - use chrono::{DateTime, Utc}; - use serde::{Deserialize, Serialize}; - use std::path::PathBuf; - - // ── API endpoints (constants, not secrets) ───────────────────────────── - - /// Solo / IDE international API host. Read from the `host` field inside - /// `iCubeAuthInfo` when available; this constant is the hardcoded fallback. - pub const INTL_HOST: &str = "https://api-sg-central.trae.ai"; - - /// Solo / IDE international ClientID. Extracted from `main.js`'s `QE()` - /// function and verified end-to-end against `ExchangeToken`. - pub const INTL_CLIENT_ID: &str = "en1oxy7wnw8j9n"; - - const EXCHANGE_TOKEN_PATH: &str = "/cloudide/api/v3/trae/oauth/ExchangeToken"; - - // ── Storage ──────────────────────────────────────────────────────────── - - /// Which Trae variant (international, 2 of them). - #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] - pub enum TraeVariant { - /// TRAE SOLO — international Solo product - Solo, - /// Trae — international IDE product - Ide, - } - - impl TraeVariant { - /// tokscale client id string for this variant. - pub fn client_str(&self) -> &'static str { - match self { - Self::Solo => "trae-solo", - Self::Ide => "trae", - } - } - - /// macOS Application Support directory name. - pub fn app_dir_name(&self) -> &'static str { - match self { - Self::Solo => "TRAE SOLO", - Self::Ide => "Trae", - } - } - - pub fn cli_arg(&self) -> &'static str { - match self { - Self::Solo => "solo", - Self::Ide => "ide", - } - } - - fn default_host(&self) -> &'static str { - INTL_HOST - } - - fn default_client_id(&self) -> &'static str { - INTL_CLIENT_ID - } - - fn credentials_filename(&self) -> &'static str { - match self { - Self::Solo => "credentials-solo.json", - Self::Ide => "credentials-ide.json", - } - } - } - - // ── Cache paths ──────────────────────────────────────────────────────── - - /// Root directory for Trae sync cache (credentials + sessions + manifest). - pub fn get_trae_cache_dir() -> PathBuf { - crate::paths::get_config_dir().join("trae-cache") - } - - /// How the token was obtained. - #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] - pub enum TokenSource { - /// Auto-decrypted from the Trae desktop client. - Auto, - /// Pasted manually by the user. - Manual, - } - - /// Cached credentials (persisted to `trae-cache/credentials-*.json`). - #[derive(Debug, Clone, Serialize, Deserialize)] - pub struct CachedCredentials { - pub variant: TraeVariant, - pub token: String, - pub refresh_token: String, - /// `access_token` expiration (ISO 8601, UTC). - pub expired_at: String, - /// `refresh_token` expiration (ISO 8601, UTC). - pub refresh_expired_at: String, - pub host: String, - pub client_id: String, - pub source: TokenSource, - pub user_id: Option, - } - - impl CachedCredentials { - /// Whether `access_token` is expired (with a 5-minute safety margin). - fn is_token_expired(&self) -> bool { - parse_iso_to_timestamp(&self.expired_at) - .map(|ts| Utc::now().timestamp_millis() > ts - 300_000) - .unwrap_or(true) - } - - /// Whether `refresh_token` is expired (with a 1-day safety margin). - fn is_refresh_expired(&self) -> bool { - parse_iso_to_timestamp(&self.refresh_expired_at) - .map(|ts| Utc::now().timestamp_millis() > ts - 86_400_000) - .unwrap_or(true) - } - } - - fn parse_iso_to_timestamp(s: &str) -> Option { - DateTime::parse_from_rfc3339(s) - .ok() - .map(|dt| dt.timestamp_millis()) - } - - fn creds_path(variant: TraeVariant) -> PathBuf { - get_trae_cache_dir().join(variant.credentials_filename()) - } - - fn ensure_cache_dir() -> Result<()> { - let dir = get_trae_cache_dir(); - if !dir.exists() { - std::fs::create_dir_all(&dir)?; - #[cfg(unix)] - { - use std::os::unix::fs::PermissionsExt; - std::fs::set_permissions(&dir, std::fs::Permissions::from_mode(0o700))?; - } - } - Ok(()) - } - - fn save_credentials(creds: &CachedCredentials) -> Result<()> { - ensure_cache_dir()?; - let path = creds_path(creds.variant); - let json = serde_json::to_string_pretty(creds)?; - tokscale_core::fs_atomic::write_atomic(&path, json.as_bytes())?; - Ok(()) - } - - fn load_credentials(variant: TraeVariant) -> Option { - let path = creds_path(variant); - if !path.exists() { - return None; - } - let content = std::fs::read_to_string(path).ok()?; - serde_json::from_str(&content).ok() - } - - fn clear_credentials(variant: TraeVariant) -> Result<()> { - let path = creds_path(variant); - if path.exists() { - std::fs::remove_file(path)?; - } - Ok(()) - } - - // ── Credential resolution (full priority chain) ──────────────────────── - - /// Resolve the active `access_token` for the given variant. Handles - /// expiration → refresh → fall back to decrypting `storage.json`. - pub async fn resolve_token(variant: TraeVariant) -> Result { - // 1. Cache hit & still valid. - if let Some(mut creds) = load_credentials(variant) { - if !creds.is_token_expired() { - return Ok(creds.token); - } - // 2. `access_token` expired but `refresh_token` still valid. - if !creds.is_refresh_expired() { - match exchange_token( - creds.host.clone(), - creds.client_id.clone(), - &creds.refresh_token, - &creds.token, - ) - .await - { - Ok(new) => { - creds.token = new.token; - creds.refresh_token = new.refresh_token; - creds.expired_at = new.expired_at; - creds.refresh_expired_at = new.refresh_expired_at; - save_credentials(&creds)?; - return Ok(creds.token); - } - Err(e) => { - eprintln!( - " Trae {} refresh token failed: {e}; falling back to storage.json decryption", - variant.client_str() - ); - } - } - } - } - - // 3. Decrypt from the Trae desktop client's `storage.json`. - // - // The on-disk credentials can be stale even on the very first read: - // the desktop client may have shipped them with an expired - // `access_token` if it hasn't been opened recently. Run the same - // expiry → refresh dance as the cache-hit branch so the caller - // doesn't immediately hit a 401. - match decrypt_from_storage(variant) { - Ok(mut creds) => { - if creds.is_token_expired() { - if creds.is_refresh_expired() { - eprintln!( - " Trae {} decrypted credentials are fully expired; falling through to manual login", - variant.client_str() - ); - } else { - match exchange_token( - creds.host.clone(), - creds.client_id.clone(), - &creds.refresh_token, - &creds.token, - ) - .await - { - Ok(new) => { - creds.token = new.token; - creds.refresh_token = new.refresh_token; - creds.expired_at = new.expired_at; - creds.refresh_expired_at = new.refresh_expired_at; - save_credentials(&creds)?; - return Ok(creds.token); - } - Err(e) => { - eprintln!( - " Trae {} decrypted token is stale and refresh failed: {e}", - variant.client_str() - ); - } - } - } - } else { - save_credentials(&creds)?; - return Ok(creds.token); - } - } - Err(e) => { - eprintln!(" Trae {} auto-decrypt failed: {e}", variant.client_str()); - } - } - - // 4. Everything failed. - Err(anyhow::anyhow!( - "Could not obtain a Trae {} access token. Run `tokscale trae login --manual --variant {}` to paste a JWT manually.", - variant.client_str(), - variant.cli_arg() - )) - } - - // ── storage.json decryption ──────────────────────────────────────────── - - fn decrypt_from_storage(variant: TraeVariant) -> Result { - let home = dirs::home_dir().context("could not determine home directory")?; - let app_dir = home - .join("Library/Application Support") - .join(variant.app_dir_name()); - let storage = app_dir.join("User/globalStorage/storage.json"); - - if !storage.exists() { - return Err(anyhow::anyhow!( - "storage.json for {} not found (expected at: {})", - variant.app_dir_name(), - storage.display() - )); - } - - let obj: serde_json::Value = serde_json::from_str(&std::fs::read_to_string(&storage)?)?; - let b64 = obj - .get("iCubeAuthInfo://icube.cloudide") - .and_then(|v| v.as_str()) - .or_else(|| { - // Walk every key starting with iCubeAuthInfo and take the longest - // value (in practice that's the cloudide entry). - obj.as_object().and_then(|o| { - o.iter() - .filter(|(k, _)| k.starts_with("iCubeAuthInfo")) - .map(|(_, v)| v) - .filter_map(|v| v.as_str()) - .max_by_key(|s| s.len()) - }) - }) - .context("no iCubeAuthInfo entry found in storage.json")?; - - let json = safestorage::decrypt_base64_blob(b64)?; - let raw: serde_json::Value = - serde_json::from_str(&json).context("decrypted iCubeAuthInfo is not valid JSON")?; - - let token = raw["token"] - .as_str() - .context("missing `token` field")? - .to_string(); - let refresh_token = raw["refreshToken"] - .as_str() - .context("missing `refreshToken` field")? - .to_string(); - let expired_at = raw["expiredAt"] - .as_str() - .map(String::from) - .unwrap_or_default(); - let refresh_expired_at = raw["refreshExpiredAt"] - .as_str() - .map(String::from) - .unwrap_or_default(); - let host = raw["host"] - .as_str() - .map(String::from) - .unwrap_or_else(|| variant.default_host().to_string()); - let user_id = raw["userId"].as_str().map(String::from); - - Ok(CachedCredentials { - variant, - token, - refresh_token, - expired_at, - refresh_expired_at, - host, - client_id: variant.default_client_id().to_string(), - source: TokenSource::Auto, - user_id, - }) - } - - // ── HTTP operations ──────────────────────────────────────────────────── - - /// ExchangeToken response (nested under `Result`). - #[derive(Debug, Deserialize)] - struct ExchangeResult { - #[serde(rename = "Token")] - token: String, - #[serde(rename = "RefreshToken")] - refresh_token: String, - #[serde(rename = "TokenExpireAt")] - token_expire_at: i64, // epoch milliseconds - #[serde(rename = "RefreshExpireAt")] - refresh_expire_at: i64, - } - - #[derive(Debug, Deserialize)] - struct ExchangeResponse { - #[serde(rename = "Result")] - result: ExchangeResult, - } - - struct TokenPair { - token: String, - refresh_token: String, - expired_at: String, - refresh_expired_at: String, - } - - /// Call the `ExchangeToken` endpoint to mint a new `access_token` from a - /// `refresh_token`. - async fn exchange_token( - host: String, - client_id: String, - refresh_token: &str, - current_token: &str, - ) -> Result { - let client = reqwest::Client::new(); - let url = format!("{}{}", host, EXCHANGE_TOKEN_PATH); - let resp = client - .post(&url) - .header("Content-Type", "application/json") - .header("x-cloudide-token", current_token) - .json(&serde_json::json!({ - "ClientID": client_id, - "RefreshToken": refresh_token, - "ClientSecret": "-", - "UserID": "" - })) - .timeout(std::time::Duration::from_secs(15)) - .send() - .await?; - - let status = resp.status(); - if !status.is_success() { - let body = resp.text().await.unwrap_or_default(); - return Err(anyhow::anyhow!( - "ExchangeToken returned {}: {}", - status, - body - )); - } - let data: ExchangeResponse = resp.json().await?; - let r = data.result; - Ok(TokenPair { - token: r.token, - refresh_token: r.refresh_token, - expired_at: epoch_ms_to_iso(r.token_expire_at), - refresh_expired_at: epoch_ms_to_iso(r.refresh_expire_at), - }) - } - - fn epoch_ms_to_iso(ms: i64) -> String { - match chrono::DateTime::from_timestamp_millis(ms) { - Some(dt) => dt.to_rfc3339(), - None => String::new(), - } - } - - // ── Public utility ───────────────────────────────────────────────────── - - /// Obtain the access token + host for a variant. Host is read from the - /// credentials cache, with a hardcoded fallback. - pub async fn get_token_and_host(variant: TraeVariant) -> Result<(String, String)> { - let token = resolve_token(variant).await?; - // `resolve_token` writes the cache file on success, but a concurrent - // `logout`, partial write, or corrupted JSON could still leave us - // unable to load it back. Fall back to the variant's default host - // instead of panicking on `unwrap()`. - let host = load_credentials(variant) - .map(|c| c.host) - .unwrap_or_else(|| variant.default_host().to_string()); - Ok((token, host)) - } - - /// Decode the JWT payload (second `.`-separated segment) as JSON. - /// - /// JWTs use unpadded base64url per RFC 7519. The previous implementation - /// appended `===` and then stripped all `=` with `trim_end_matches`, - /// which silently produced bad input for any payload whose length wasn't - /// already a multiple of 4 — every such token then looked instantly - /// expired because `exp` / `iat` couldn't be read. - fn decode_jwt_payload(token: &str) -> Option { - let parts: Vec<&str> = token.split('.').collect(); - if parts.len() < 2 { - return None; - } - // Some encoders still emit trailing `=`; URL_SAFE_NO_PAD rejects - // padding so strip it here. - let raw = parts[1].trim_end_matches('='); - let bytes = base64::engine::general_purpose::URL_SAFE_NO_PAD - .decode(raw) - .ok()?; - serde_json::from_slice(&bytes).ok() - } - - /// Persist a user-pasted JWT (no live expiration check — the expirations - /// are decoded straight from the JWT payload). - pub fn save_manual_token( - variant: TraeVariant, - token: String, - host: Option, - ) -> Result<()> { - ensure_cache_dir()?; - // Decode `exp` and `iat` from the JWT payload. - let (expired_at, refresh_expired_at) = if let Some(payload) = decode_jwt_payload(&token) { - let exp = payload["exp"].as_i64().map_or_else( - || (Utc::now() + chrono::Duration::days(14)).to_rfc3339(), - epoch_secs_to_iso, - ); - // Estimate `refresh_expired_at` as `iat + 180 days`. - let iat = payload["iat"].as_i64().map_or_else( - || (Utc::now() + chrono::Duration::days(180)).to_rfc3339(), - |ts| { - chrono::DateTime::from_timestamp(ts + 180 * 86400, 0) - .map(|dt| dt.to_rfc3339()) - .unwrap_or_default() - }, - ); - (exp, iat) - } else { - (String::new(), String::new()) - }; - - let host = host.unwrap_or_else(|| variant.default_host().to_string()); - let creds = CachedCredentials { - variant, - token, - refresh_token: String::new(), - expired_at, - refresh_expired_at, - host, - client_id: variant.default_client_id().to_string(), - source: TokenSource::Manual, - user_id: None, - }; - save_credentials(&creds) - } - - fn epoch_secs_to_iso(secs: i64) -> String { - chrono::DateTime::from_timestamp(secs, 0) - .map(|dt| dt.to_rfc3339()) - .unwrap_or_default() - } - - /// Clear the cached credentials for a variant. - pub fn logout(variant: TraeVariant) -> Result<()> { - clear_credentials(variant) - } - - /// Whether the variant has usable credentials (no network calls). - pub fn has_credentials(variant: TraeVariant) -> bool { - load_credentials(variant).is_some() - } - - /// Iterator over all 2 supported variants. - pub fn all_variants() -> [TraeVariant; 2] { - [TraeVariant::Solo, TraeVariant::Ide] - } - - #[cfg(test)] - mod tests { - use super::*; - - #[test] - fn test_variant_client_str() { - assert_eq!(TraeVariant::Solo.client_str(), "trae-solo"); - assert_eq!(TraeVariant::Ide.client_str(), "trae"); - } - - #[test] - fn test_variant_default_host() { - assert_eq!(TraeVariant::Solo.default_host(), INTL_HOST); - assert_eq!(TraeVariant::Ide.default_host(), INTL_HOST); - } - - #[test] - fn test_variant_serialize() { - let solo = serde_json::to_string(&TraeVariant::Solo).unwrap(); - assert_eq!(solo, r#""Solo""#); - let ide = serde_json::to_string(&TraeVariant::Ide).unwrap(); - assert_eq!(ide, r#""Ide""#); - } - - #[test] - fn test_epoch_ms_to_iso() { - let dt = chrono::DateTime::parse_from_rfc3339("2026-05-21T11:29:04.295Z").unwrap(); - let ms = dt.timestamp_millis(); - let iso = epoch_ms_to_iso(ms); - assert!(iso.contains("2026-05-21")); - } - - #[test] - fn test_all_variants_count() { - assert_eq!(all_variants().len(), 2); - } - - fn encode_jwt_payload(payload: &serde_json::Value) -> String { - let header = base64::engine::general_purpose::URL_SAFE_NO_PAD - .encode(b"{\"alg\":\"HS256\",\"typ\":\"JWT\"}"); - let body = base64::engine::general_purpose::URL_SAFE_NO_PAD - .encode(serde_json::to_vec(payload).unwrap()); - let sig = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(b"sig"); - format!("{header}.{body}.{sig}") - } - - #[test] - fn test_decode_jwt_payload_extracts_exp_and_iat() { - // Payload length 27 → not a multiple of 4. The previous - // implementation stripped padding and choked on this. - let payload = serde_json::json!({ "exp": 1900000000_i64, "iat": 1780000000_i64 }); - let token = encode_jwt_payload(&payload); - let decoded = decode_jwt_payload(&token).expect("decode succeeds"); - assert_eq!(decoded["exp"].as_i64(), Some(1900000000)); - assert_eq!(decoded["iat"].as_i64(), Some(1780000000)); - } - - #[test] - fn test_decode_jwt_payload_handles_url_safe_chars() { - // Force a payload that exercises base64url `-` / `_` substitution. - // `>` and `?` map to `+` / `/` in standard base64; in URL-safe - // they map to `-` / `_`. Using a string with a known byte that - // produces `_` in the encoding ensures we don't regress to - // STANDARD-only decoding. - let payload = serde_json::json!({ "data": "??>>??" }); - let token = encode_jwt_payload(&payload); - assert!(token.contains('_') || token.contains('-')); - let decoded = decode_jwt_payload(&token).expect("decode succeeds"); - assert_eq!(decoded["data"].as_str(), Some("??>>??")); - } - - #[test] - fn test_decode_jwt_payload_rejects_malformed_token() { - assert!(decode_jwt_payload("not-a-jwt").is_none()); - assert!(decode_jwt_payload("").is_none()); - assert!(decode_jwt_payload("badbase64!@#.badbase64!@#").is_none()); - } - } -} - -pub mod safestorage { - //! Decrypt the `iCubeAuthInfo://*` blobs stored in the Trae desktop - //! client's `globalStorage/storage.json`. - //! - //! The algorithm is a faithful Rust port of the Trae client's - //! `byteCrypto.js` module (the `V8e()` decrypt path). Field meanings, - //! constants, and the overall flow are kept identical to the source — - //! any changes must be re-validated against the byte offsets. - //! - //! ## Blob layout - //! - //! ```text - //! [magic 6 bytes: "tc\x05\x10\x00\x00"] - //! [salt 32 bytes, random] - //! [ciphertext N×16 bytes, AES-128-CBC + PKCS7 padding] - //! ``` - //! - //! ## Key derivation (`BG` function) - //! - //! ```text - //! hardcoded_pw = JG XOR KG // 64 bytes, hardcoded (obfuscated) in client source - //! kdf_buf = SHA-512(salt) || hardcoded_pw // 128 bytes - //! kdf_out = SHA-512(kdf_buf) // 64 bytes - //! aes_key = kdf_out[0..16] - //! iv = kdf_out[16..32] - //! ``` - //! - //! ## Plaintext integrity check (tail of `V8e`) - //! - //! ```text - //! plaintext = [hash 64 bytes] [data N bytes] - //! Requires SHA-512(data) == hash; otherwise treated as corrupt. - //! ``` - - use aes::Aes128; - use anyhow::{anyhow, Context, Result}; - use base64::{engine::general_purpose::STANDARD as B64, Engine as _}; - use cbc::cipher::{block_padding::Pkcs7, BlockDecryptMut, KeyIvInit}; - use sha2::{Digest, Sha512}; - - type Aes128CbcDec = cbc::Decryptor; - - const MAGIC: [u8; 6] = [b't', b'c', 0x05, 0x10, 0x00, 0x00]; - const SALT_LEN: usize = 32; - const HASH_LEN: usize = 64; - const AES_KEY_LEN: usize = 16; - const IV_LEN: usize = 16; - - /// `JG` array — extracted from `byteCrypto.js`; the first 64 bytes of the - /// AES inverse S-box (obfuscation material). - const JG: [u8; 64] = [ - 82, 9, 106, 213, 48, 54, 165, 56, 191, 64, 163, 158, 129, 243, 215, 251, 124, 227, 57, 130, - 155, 47, 255, 135, 52, 142, 67, 68, 196, 222, 233, 203, 84, 123, 148, 50, 166, 194, 35, 61, - 238, 76, 149, 11, 66, 250, 195, 78, 8, 46, 161, 102, 40, 217, 36, 178, 118, 91, 162, 73, - 109, 139, 209, 37, - ]; - /// `KG` array — extracted from `byteCrypto.js`; the other half of the - /// hardcoded "password". - const KG: [u8; 64] = [ - 31, 221, 168, 51, 136, 7, 199, 49, 177, 18, 16, 89, 39, 128, 236, 95, 96, 81, 127, 169, 25, - 181, 74, 13, 45, 229, 122, 159, 147, 201, 156, 239, 160, 224, 59, 77, 174, 42, 245, 176, - 200, 235, 187, 60, 131, 83, 153, 97, 23, 43, 4, 126, 186, 119, 214, 38, 225, 105, 20, 99, - 85, 33, 12, 125, - ]; - - fn hardcoded_password() -> [u8; 64] { - let mut pw = [0u8; 64]; - for i in 0..64 { - pw[i] = JG[i] ^ KG[i]; - } - pw - } - - fn derive_key_iv(salt: &[u8]) -> Result<([u8; AES_KEY_LEN], [u8; IV_LEN])> { - if salt.len() != SALT_LEN { - return Err(anyhow!( - "wrong salt length: expected {}, got {}", - SALT_LEN, - salt.len() - )); - } - let sha_salt = Sha512::digest(salt); - let mut kdf_buf = [0u8; 128]; - kdf_buf[..HASH_LEN].copy_from_slice(&sha_salt); - kdf_buf[HASH_LEN..].copy_from_slice(&hardcoded_password()); - let kdf_out = Sha512::digest(kdf_buf); - let mut key = [0u8; AES_KEY_LEN]; - let mut iv = [0u8; IV_LEN]; - key.copy_from_slice(&kdf_out[..AES_KEY_LEN]); - iv.copy_from_slice(&kdf_out[AES_KEY_LEN..AES_KEY_LEN + IV_LEN]); - Ok((key, iv)) - } - - /// Decrypt a base64-encoded blob and return the plaintext UTF-8 string - /// (typically JSON). - pub fn decrypt_base64_blob(b64: &str) -> Result { - let raw = B64 - .decode(b64.trim()) - .context("failed to base64-decode iCubeAuthInfo")?; - let pt = decrypt_blob(&raw)?; - String::from_utf8(pt).context("decrypted plaintext is not valid UTF-8") - } - - /// Decrypt a raw-bytes blob and return the plaintext (with the SHA-512 - /// integrity prefix stripped). - pub fn decrypt_blob(blob: &[u8]) -> Result> { - let min_len = MAGIC.len() + SALT_LEN + 16; - if blob.len() < min_len { - return Err(anyhow!( - "blob too short: {} bytes (need at least {} bytes)", - blob.len(), - min_len - )); - } - if blob[..MAGIC.len()] != MAGIC { - return Err(anyhow!( - "magic header mismatch: got {:02x?}, expected {:02x?}", - &blob[..MAGIC.len()], - MAGIC - )); - } - let salt = &blob[MAGIC.len()..MAGIC.len() + SALT_LEN]; - let ciphertext = &blob[MAGIC.len() + SALT_LEN..]; - if !ciphertext.len().is_multiple_of(16) { - return Err(anyhow!( - "ciphertext length {} is not a multiple of the AES block size (16)", - ciphertext.len() - )); - } - let (key, iv) = derive_key_iv(salt)?; - let mut buf = ciphertext.to_vec(); - let plaintext_len = Aes128CbcDec::new(&key.into(), &iv.into()) - .decrypt_padded_mut::(&mut buf) - .map_err(|e| anyhow!("AES-CBC decryption failed: {e}"))? - .len(); - buf.truncate(plaintext_len); - if buf.len() < HASH_LEN { - return Err(anyhow!( - "plaintext too short: {} bytes (need at least {} bytes for the hash prefix)", - buf.len(), - HASH_LEN - )); - } - let (expected_hash, data) = buf.split_at(HASH_LEN); - let actual_hash = Sha512::digest(data); - if expected_hash != actual_hash.as_slice() { - return Err(anyhow!("SHA-512 integrity check failed")); - } - Ok(data.to_vec()) - } - - #[cfg(test)] - mod tests { - use super::*; - use aes::Aes128; - use cbc::cipher::{block_padding::Pkcs7, BlockEncryptMut, KeyIvInit}; - - type Aes128CbcEnc = cbc::Encryptor; - - /// Encrypt with the same algorithm — used only by tests. Production - /// blobs are encrypted by the Trae client itself. - fn encrypt_for_test(plain_text: &[u8]) -> Vec { - // Fixed salt so the test is deterministic. - let salt = [0xAAu8; SALT_LEN]; - let (key, iv) = derive_key_iv(&salt).unwrap(); - // hash || data - let hash = Sha512::digest(plain_text); - let mut buf = Vec::with_capacity(HASH_LEN + plain_text.len() + 16); - buf.extend_from_slice(&hash); - buf.extend_from_slice(plain_text); - // PKCS7 padding happens inside `encrypt_padded_mut`; reserve space. - let unpadded_len = buf.len(); - buf.resize(unpadded_len + 16, 0); - let ct_len = Aes128CbcEnc::new(&key.into(), &iv.into()) - .encrypt_padded_mut::(&mut buf, unpadded_len) - .unwrap() - .len(); - buf.truncate(ct_len); - let mut blob = Vec::with_capacity(MAGIC.len() + SALT_LEN + ct_len); - blob.extend_from_slice(&MAGIC); - blob.extend_from_slice(&salt); - blob.extend_from_slice(&buf); - blob - } - - #[test] - fn test_hardcoded_password_constant() { - // This value must not change — changing it breaks decryption of - // every existing Trae client blob. - let pw = hardcoded_password(); - assert_eq!(pw.len(), 64); - // First byte = JG[0] ^ KG[0] = 82 ^ 31 = 77 = 0x4d - assert_eq!(pw[0], 0x4d); - // Last byte = 37 ^ 125 = 88 = 0x58 - assert_eq!(pw[63], 0x58); - } - - #[test] - fn test_round_trip_simple_json() { - let plain = br#"{"token":"abc","refreshToken":"xyz"}"#; - let blob = encrypt_for_test(plain); - let decrypted = decrypt_blob(&blob).expect("decrypt succeeds"); - assert_eq!(&decrypted, plain); - } - - #[test] - fn test_round_trip_unicode() { - let plain = "Hello, world 🌍 — Unicode test".as_bytes(); - let blob = encrypt_for_test(plain); - let decrypted = decrypt_blob(&blob).expect("decrypt succeeds"); - assert_eq!(&decrypted, plain); - } - - #[test] - fn test_decrypt_base64_blob_round_trip() { - let plain = br#"{"hello":"world"}"#; - let blob = encrypt_for_test(plain); - let b64 = B64.encode(&blob); - let s = decrypt_base64_blob(&b64).expect("decrypt + utf8 ok"); - assert_eq!(s, r#"{"hello":"world"}"#); - } - - #[test] - fn test_wrong_magic_rejected() { - let mut blob = encrypt_for_test(b"test"); - blob[0] = b'x'; - let err = decrypt_blob(&blob).unwrap_err(); - assert!(err.to_string().contains("magic")); - } - - #[test] - fn test_too_short_rejected() { - let blob = vec![0u8; 16]; - let err = decrypt_blob(&blob).unwrap_err(); - assert!(err.to_string().contains("blob too short")); - } - - #[test] - fn test_tampered_ciphertext_caught_by_hash_check() { - let plain = br#"{"token":"abc"}"#; - let mut blob = encrypt_for_test(plain); - // Flip a byte in the ciphertext — PKCS7 unpadding may still - // succeed (depending on the block), but SHA-512 will fail. - let last = blob.len() - 1; - blob[last - 16] ^= 0x01; - let err = decrypt_blob(&blob); - assert!(err.is_err(), "expected decrypt to fail on tampered blob"); - } - - #[test] - fn test_ciphertext_not_block_aligned_rejected() { - // Deliberately construct a blob whose ciphertext is misaligned. - let mut blob = Vec::new(); - blob.extend_from_slice(&MAGIC); - blob.extend_from_slice(&[0u8; SALT_LEN]); - blob.extend_from_slice(&[0u8; 17]); // 17 bytes — not a multiple of 16 - let err = decrypt_blob(&blob).unwrap_err(); - assert!(err.to_string().contains("AES block size")); - } - } -} - -pub mod sync { - //! Trae usage sync: paginated pulls from the official API, persisted to a - //! local cache with a manifest. - //! - //! Mirrors the Antigravity (manifest + lock) and Cursor (HTTP API) patterns. - //! - //! Trae IDE and Trae Solo share account usage data. The variant only - //! controls where credentials are discovered; synced usage is stored once - //! under the single `trae` client cache. - - use super::auth::{self, get_trae_cache_dir, TraeVariant}; - use anyhow::{Context, Result}; - use chrono::Utc; - use serde::{Deserialize, Serialize}; - use std::collections::HashMap; - use std::path::PathBuf; - use std::time::Duration; - - const PAGE_SIZE: i32 = 20; - const API_PAGE_DELAY_MS: u64 = 300; - const OVERLAP_MARGIN_SECS: i64 = 7200; // 2-hour overlap buffer - const MANIFEST_VERSION: i32 = 1; - - fn should_replace_session_entry( - existing: &TraeSessionEntry, - incoming: &TraeSessionEntry, - ) -> bool { - incoming.usage_time > existing.usage_time - || (incoming.usage_time == existing.usage_time - && incoming.artifact_path > existing.artifact_path) - } - - fn upsert_manifest_entry( - manifest_sessions: &mut HashMap, - incoming: TraeSessionEntry, - ) { - if let Some(existing) = manifest_sessions.get_mut(&incoming.session_id) { - if should_replace_session_entry(existing, &incoming) { - *existing = incoming; - } - return; - } - - manifest_sessions.insert(incoming.session_id.clone(), incoming); - } - - fn merge_manifest_sessions( - existing: Vec, - incoming: Vec, - ) -> Vec { - let mut entries: HashMap = existing - .into_iter() - .map(|entry| (entry.session_id.clone(), entry)) - .collect(); - - for incoming_entry in incoming { - upsert_manifest_entry(&mut entries, incoming_entry); - } - - let mut merged: Vec = entries.into_values().collect(); - merged.sort_unstable_by(|a, b| a.session_id.cmp(&b.session_id)); - merged - } - - fn manifest_references_artifact(sessions: &[TraeSessionEntry], artifact_path: &str) -> bool { - sessions - .iter() - .any(|entry| entry.artifact_path == artifact_path) - } - - // ── Manifest ─────────────────────────────────────────────────────────── - - #[derive(Debug, Clone, Serialize, Deserialize)] - pub struct TraeManifest { - pub version: i32, - pub last_synced_at: i64, // epoch seconds - pub sessions: Vec, - } - - #[derive(Debug, Clone, Serialize, Deserialize)] - pub struct TraeSessionEntry { - pub session_id: String, - pub usage_time: i64, // epoch seconds - pub artifact_path: String, - } - - fn manifest_path() -> PathBuf { - get_trae_cache_dir().join("manifest.json") - } - - fn sessions_dir() -> PathBuf { - get_trae_cache_dir().join("sessions") - } - - fn load_manifest() -> Result { - let path = manifest_path(); - if !path.exists() { - return Ok(TraeManifest { - version: MANIFEST_VERSION, - last_synced_at: 0, - sessions: Vec::new(), - }); - } - let content = std::fs::read_to_string(&path)?; - serde_json::from_str(&content).context("failed to parse manifest JSON") - } - - fn save_manifest(manifest: &TraeManifest) -> Result<()> { - let dir = get_trae_cache_dir(); - if !dir.exists() { - std::fs::create_dir_all(&dir)?; - } - let json = serde_json::to_string_pretty(manifest)?; - tokscale_core::fs_atomic::write_atomic(&manifest_path(), json.as_bytes())?; - Ok(()) - } - - fn ensure_sessions_dir() -> Result { - let dir = sessions_dir(); - if !dir.exists() { - std::fs::create_dir_all(&dir)?; - } - Ok(dir) - } - - // ── HTTP client ──────────────────────────────────────────────────────── - - #[derive(Debug, Deserialize)] - struct UsageResponse { - #[serde(rename = "user_usage_group_by_sessions")] - sessions: Option>, - total: Option, - } - - /// Paginated call to the usage API; returns the raw JSON session list. - async fn fetch_usage_pages( - host: &str, - token: &str, - start_time: i64, - end_time: i64, - usage_types: &[i32], - ) -> Result> { - let client = reqwest::Client::new(); - let url = format!("{}/trae/api/v1/pay/query_user_usage_group_by_session", host); - let mut all = Vec::new(); - let mut page = 1; - - loop { - let payload = serde_json::json!({ - "start_time": start_time, - "end_time": end_time, - "page_size": PAGE_SIZE, - "page_num": page, - "usage_type": usage_types, - }); - - let resp = client - .post(&url) - .header("Content-Type", "application/json") - .header("authorization", format!("Cloud-IDE-JWT {}", token)) - .timeout(Duration::from_secs(30)) - .json(&payload) - .send() - .await?; - - let status = resp.status(); - if !status.is_success() { - let body = resp.text().await.unwrap_or_default(); - return Err(anyhow::anyhow!("usage API returned {}: {}", status, body)); - } - - let data: UsageResponse = resp.json().await?; - let sessions = data.sessions.unwrap_or_default(); - let batch = sessions.len(); - all.extend(sessions); - - // When the API omits `total`, keep paginating until an empty - // page. `unwrap_or(0)` would produce `total=0` which makes - // `all.len() >= total` immediately true, silently truncating - // data after the first non-empty page. - if batch == 0 || data.total.is_some_and(|t| all.len() >= t as usize) { - break; - } - page += 1; - tokio::time::sleep(Duration::from_millis(API_PAGE_DELAY_MS)).await; - } - - Ok(all) - } - - // ── Sync lock ────────────────────────────────────────────────────────── - - const SYNC_LOCK_ACQUIRE_ATTEMPTS: usize = 3; - - #[derive(Debug)] - struct SyncLockGuard { - path: PathBuf, - } - - impl SyncLockGuard { - fn acquire(cache_dir: &std::path::Path) -> Result { - let lock_path = cache_dir.join("sync.lock"); - if !cache_dir.exists() { - std::fs::create_dir_all(cache_dir)?; - } - let mut stale_recoveries = 0usize; - loop { - match std::fs::OpenOptions::new() - .create_new(true) - .write(true) - .open(&lock_path) - { - Ok(mut file) => { - use std::io::Write; - let _ = writeln!(file, "{} {}", std::process::id(), Utc::now().timestamp()); - return Ok(Self { path: lock_path }); - } - Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => { - // Only evict the lock when its owner is provably - // dead. Live syncs MUST keep exclusive access for as - // long as they hold the PID — otherwise two - // processes overlap on the manifest and delete each - // other's session artifacts. - if let Some((existing_pid, _)) = read_sync_lock(&lock_path) { - if pid_is_alive(existing_pid) { - return Err(anyhow::anyhow!( - "another trae sync is in progress (pid {existing_pid}); aborting" - )); - } - } - if stale_recoveries >= SYNC_LOCK_ACQUIRE_ATTEMPTS { - return Err(anyhow::anyhow!( - "could not acquire trae sync lock after {SYNC_LOCK_ACQUIRE_ATTEMPTS} stale-lock recoveries; another process keeps recreating the lock file" - )); - } - stale_recoveries += 1; - let _ = std::fs::remove_file(&lock_path); - continue; - } - Err(e) => { - return Err(anyhow::Error::new(e).context("failed to acquire sync lock")); - } - } - } - } - } - - impl Drop for SyncLockGuard { - fn drop(&mut self) { - let _ = std::fs::remove_file(&self.path); - } - } - - fn read_sync_lock(path: &std::path::Path) -> Option<(u32, u64)> { - let contents = std::fs::read_to_string(path).ok()?; - let mut parts = contents.split_whitespace(); - let pid = parts.next()?.parse::().ok()?; - let timestamp = parts.next()?.parse::().ok()?; - Some((pid, timestamp)) - } - - fn pid_is_alive(pid: u32) -> bool { - if pid == 0 { - return false; - } - #[cfg(unix)] - { - // `kill(pid, 0)` is a signal-free liveness probe. EPERM (errno - // 1) still means the process exists, just that we lack - // permission to signal it. - let result = unsafe { libc_kill(pid as i32, 0) }; - result == 0 || std::io::Error::last_os_error().raw_os_error() == Some(1) - } - #[cfg(not(unix))] - { - let _ = pid; - // No portable PID probe available. Treat the lock as stale so a - // crashed previous run doesn't permanently block subsequent - // syncs. This matches the policy used by Antigravity sync and - // accepts a small concurrent-corruption risk on Windows; that - // risk is acceptable because tokscale is a single-user CLI and - // overlapping syncs are rare in practice. - false - } - } - - #[cfg(unix)] - extern "C" { - #[link_name = "kill"] - fn libc_kill(pid: i32, sig: i32) -> i32; - } - - // ── Main sync logic ──────────────────────────────────────────────────── - - /// Run a single incremental sync for one variant. - pub async fn sync_variant( - variant: TraeVariant, - since_days: i64, - usage_types: &[i32], - ) -> Result { - let _lock = SyncLockGuard::acquire(&get_trae_cache_dir())?; - - let (token, host) = auth::get_token_and_host(variant) - .await - .context("failed to obtain Trae access token")?; - - let now = Utc::now().timestamp(); - let manifest = load_manifest()?; - // Incremental start point: take the earlier of (user-requested, manifest - // record) so we never lose data across runs. - let since_user = now - since_days * 86400; - let since_manifest = if manifest.last_synced_at > 0 { - manifest.last_synced_at - OVERLAP_MARGIN_SECS - } else { - 0 - }; - let start_time = if since_manifest > 0 { - since_user.min(since_manifest) - } else { - since_user - }; - let end_time = now; - - let sessions = fetch_usage_pages(&host, &token, start_time, end_time, usage_types).await?; - - if sessions.is_empty() { - return Ok(0); - } - - let dir = ensure_sessions_dir()?; - let mut next_manifest = TraeManifest { - version: MANIFEST_VERSION, - last_synced_at: now, - sessions: manifest.sessions.clone(), - }; - - // Write the whole batch only if at least one fetched session wins the - // manifest merge. Repeated overlapping syncs often fetch older copies - // of already-seen sessions; writing those losing batches creates a - // file that GC immediately removes and can overwrite a still-referenced - // artifact if two syncs happen in the same timestamp bucket. - let batch_ts = Utc::now().format("%Y%m%dT%H%M%S%.3f").to_string(); - let artifact_filename = format!("usage-{batch_ts}.json"); - let manifest_session_path = format!("sessions/{artifact_filename}"); - let artifact_path = dir.join(&artifact_filename); - - let incoming_sessions: Vec = sessions - .iter() - .filter_map(|s| { - let session_id = s["session_id"].as_str()?.to_string(); - if session_id.is_empty() { - return None; - } - let usage_time = s["usage_time"].as_i64().unwrap_or(0); - Some(TraeSessionEntry { - session_id, - usage_time, - artifact_path: manifest_session_path.clone(), - }) - }) - .collect(); - - next_manifest.sessions = merge_manifest_sessions(next_manifest.sessions, incoming_sessions); - - let batch_wins_manifest = - manifest_references_artifact(&next_manifest.sessions, &manifest_session_path); - - if batch_wins_manifest { - let json = serde_json::to_string_pretty(&sessions)?; - tokscale_core::fs_atomic::write_atomic(&artifact_path, json.as_bytes())?; - } - - let valid_paths: std::collections::HashSet = next_manifest - .sessions - .iter() - .map(|e| e.artifact_path.clone()) - .collect(); - - if let Ok(entries) = std::fs::read_dir(&dir) { - for entry in entries.flatten() { - if let Some(name) = entry.file_name().to_str() { - let rel = format!("sessions/{}", name); - if !valid_paths.contains(&rel) && name.ends_with(".json") { - let _ = std::fs::remove_file(entry.path()); - } - } - } - } - - save_manifest(&next_manifest)?; - - Ok(sessions.len()) - } - - /// CLI entry point: sync Trae usage once using a selected credential source. - /// - /// `variants` comes from CLI flags. Because Trae IDE and Trae Solo return - /// the same account-level usage data, only one variant needs to succeed. - /// We iterate through every variant with credentials and stop at the first - /// success; if a variant fails (expired refresh token, unreadable cache, - /// transient HTTP error), we fall over to the next one before giving up. - pub async fn run_trae_sync( - variants: &[TraeVariant], - since_days: i64, - include_aux: bool, - ) -> Result<()> { - let usage_types: Vec = if include_aux { - vec![1, 2, 3, 4, 5, 6, 7, 8] - } else { - vec![5, 6] - }; - - let credentialed: Vec = variants - .iter() - .copied() - .filter(|v| auth::has_credentials(*v)) - .collect(); - - if credentialed.is_empty() { - if variants.is_empty() { - println!(" No Trae credentials found. Run `tokscale trae login` first."); - } else { - for variant in variants { - println!( - " Trae {}: no credentials — run `tokscale trae login --variant {}` first", - variant.client_str(), - variant.cli_arg() - ); - } - } - return Ok(()); - } - - let mut last_err: Option = None; - for variant in &credentialed { - match sync_variant(*variant, since_days, &usage_types).await { - Ok(n) => { - println!( - " Trae: synced {n} sessions (using {} credentials)", - variant.client_str() - ); - return Ok(()); - } - Err(e) => { - eprintln!( - " Trae sync failed using {} credentials: {e}", - variant.client_str() - ); - last_err = Some(e); - } - } - } - - // Every credentialed variant has been tried and failed. Surface the - // last error so the user sees a non-zero exit instead of a silent - // success after the per-variant `eprintln!` lines scroll off. - match last_err { - Some(e) => Err(e.context("all Trae credential sources failed")), - None => Ok(()), - } - } - - #[cfg(test)] - mod tests { - use super::*; - - #[test] - fn test_read_sync_lock_parses_pid_and_timestamp() { - let tmp = tempfile::tempdir().unwrap(); - let path = tmp.path().join("sync.lock"); - std::fs::write(&path, "12345 1776000000\n").unwrap(); - let (pid, ts) = read_sync_lock(&path).expect("readable"); - assert_eq!(pid, 12345); - assert_eq!(ts, 1776000000); - } - - #[test] - fn test_read_sync_lock_returns_none_on_malformed() { - let tmp = tempfile::tempdir().unwrap(); - let path = tmp.path().join("sync.lock"); - std::fs::write(&path, "not a pid\n").unwrap(); - assert!(read_sync_lock(&path).is_none()); - } - - #[test] - fn test_pid_is_alive_zero_is_dead() { - assert!(!pid_is_alive(0)); - } - - #[test] - #[cfg(unix)] - fn test_pid_is_alive_current_process_is_alive() { - let me = std::process::id(); - assert!(pid_is_alive(me)); - } - - #[test] - fn test_acquire_recovers_from_stale_lock_with_dead_pid() { - // PID 0 is reserved by the kernel and never represents a live - // user-space process — perfect stand-in for a crashed sync. - let tmp = tempfile::tempdir().unwrap(); - let cache_dir = tmp.path(); - std::fs::write(cache_dir.join("sync.lock"), "0 1\n").unwrap(); - // Should evict the stale lock and acquire a fresh one. - let guard = SyncLockGuard::acquire(cache_dir).expect("stale lock is recovered"); - assert!(cache_dir.join("sync.lock").exists()); - drop(guard); - // Drop releases the lock. - assert!(!cache_dir.join("sync.lock").exists()); - } - - #[test] - #[cfg(unix)] - fn test_acquire_rejects_when_owner_is_alive() { - // Use our own PID — guaranteed alive. acquire() must refuse. - let tmp = tempfile::tempdir().unwrap(); - let cache_dir = tmp.path(); - let alive_pid = std::process::id(); - std::fs::write(cache_dir.join("sync.lock"), format!("{alive_pid} 1\n")).unwrap(); - let err = SyncLockGuard::acquire(cache_dir).unwrap_err(); - assert!(err.to_string().contains("another trae sync is in progress")); - // Lock file must remain untouched so the live owner can release it. - assert!(cache_dir.join("sync.lock").exists()); - } - - #[test] - fn test_acquire_writes_pid_and_timestamp() { - let tmp = tempfile::tempdir().unwrap(); - let cache_dir = tmp.path(); - let guard = SyncLockGuard::acquire(cache_dir).expect("first acquire"); - let (pid, _) = read_sync_lock(&cache_dir.join("sync.lock")).expect("readable"); - assert_eq!(pid, std::process::id()); - drop(guard); - } - - #[test] - fn test_merge_manifest_upsert_prefers_newer_usage_time() { - let existing = vec![ - TraeSessionEntry { - session_id: "session-stable".to_string(), - usage_time: 1_700_000_000, - artifact_path: "sessions/old.json".to_string(), - }, - TraeSessionEntry { - session_id: "session-older".to_string(), - usage_time: 1_600_000_000, - artifact_path: "sessions/older.json".to_string(), - }, - ]; - - let incoming = vec![ - TraeSessionEntry { - session_id: "session-stable".to_string(), - usage_time: 1_700_000_001, - artifact_path: "sessions/newer.json".to_string(), - }, - TraeSessionEntry { - session_id: "session-older".to_string(), - usage_time: 1_500_000_000, - artifact_path: "sessions/should-not-win.json".to_string(), - }, - TraeSessionEntry { - session_id: "session-new".to_string(), - usage_time: 1_800_000_000, - artifact_path: "sessions/new.json".to_string(), - }, - ]; - - let merged = merge_manifest_sessions(existing, incoming); - merged.iter().for_each(|entry| { - if entry.session_id == "session-stable" { - assert_eq!(entry.usage_time, 1_700_000_001); - assert_eq!(entry.artifact_path, "sessions/newer.json"); - } - if entry.session_id == "session-older" { - assert_eq!(entry.usage_time, 1_600_000_000); - assert_eq!(entry.artifact_path, "sessions/older.json"); - } - }); - assert_eq!(merged.len(), 3); - } - - #[test] - fn test_merge_manifest_session_batch_dedups_same_session() { - let existing = vec![]; - let incoming = vec![ - TraeSessionEntry { - session_id: "session-dupe".to_string(), - usage_time: 1_000, - artifact_path: "sessions/first.json".to_string(), - }, - TraeSessionEntry { - session_id: "session-dupe".to_string(), - usage_time: 1_200, - artifact_path: "sessions/second.json".to_string(), - }, - TraeSessionEntry { - session_id: "session-dupe".to_string(), - usage_time: 1_200, - artifact_path: "sessions/zzz.json".to_string(), - }, - ]; - let merged = merge_manifest_sessions(existing, incoming); - assert_eq!(merged.len(), 1); - assert_eq!(merged[0].session_id, "session-dupe"); - assert_eq!(merged[0].usage_time, 1_200); - assert_eq!(merged[0].artifact_path, "sessions/zzz.json"); - } - - #[test] - fn test_manifest_reference_is_absent_when_batch_loses_merge() { - let current_batch = "sessions/current.json"; - let merged = merge_manifest_sessions( - vec![TraeSessionEntry { - session_id: "session-stable".to_string(), - usage_time: 2_000, - artifact_path: "sessions/previous.json".to_string(), - }], - vec![TraeSessionEntry { - session_id: "session-stable".to_string(), - usage_time: 1_000, - artifact_path: current_batch.to_string(), - }], - ); - - assert!(!manifest_references_artifact(&merged, current_batch)); - assert_eq!(merged[0].artifact_path, "sessions/previous.json"); - } - - #[test] - fn test_manifest_reference_is_present_when_batch_wins_merge() { - let current_batch = "sessions/current.json"; - let merged = merge_manifest_sessions( - vec![TraeSessionEntry { - session_id: "session-stable".to_string(), - usage_time: 1_000, - artifact_path: "sessions/previous.json".to_string(), - }], - vec![TraeSessionEntry { - session_id: "session-stable".to_string(), - usage_time: 2_000, - artifact_path: current_batch.to_string(), - }], - ); - - assert!(manifest_references_artifact(&merged, current_batch)); - assert_eq!(merged[0].artifact_path, current_batch); - } - } -} diff --git a/crates/tokscale-cli/src/tui/app.rs b/crates/tokscale-cli/src/tui/app.rs index acad3966e..4ade1d0b5 100644 --- a/crates/tokscale-cli/src/tui/app.rs +++ b/crates/tokscale-cli/src/tui/app.rs @@ -28,7 +28,8 @@ use super::ui::dialog::{ClientPickerDialog, DialogStack}; pub struct TuiConfig { pub theme: Option, pub refresh: u64, - pub sessions_path: Option, + pub no_refresh: bool, + pub home_dir: Option, pub clients: Option>, pub since: Option, pub until: Option, @@ -36,6 +37,18 @@ pub struct TuiConfig { pub initial_tab: Option, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TuiExit { + Quit, + Interrupted, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum KeyEventOutcome { + Continue, + Exit(TuiExit), +} + #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub enum Tab { Overview, @@ -335,7 +348,6 @@ fn sort_detail_rows(rows: &mut [DetailRow], field: SortField, direction: SortDir } pub struct App { - pub should_quit: bool, pub current_tab: Tab, pub theme: Theme, pub settings: Settings, @@ -419,7 +431,8 @@ pub struct App { impl App { pub fn new_with_cached_data(config: TuiConfig, cached_data: Option) -> Result { - let settings = Settings::load()?; + let settings = + Settings::load_for_home_override(config.home_dir.as_deref().map(std::path::Path::new))?; Self::new_with_cached_data_and_settings(config, cached_data, settings) } @@ -463,13 +476,17 @@ impl App { Duration::from_secs(30) }; - let auto_refresh = config.refresh > 0 || settings.auto_refresh_enabled; + let auto_refresh = if config.no_refresh { + false + } else { + config.refresh > 0 || settings.auto_refresh_enabled + }; let usage_tab_enabled = settings.usage_tab_enabled; let subscription_provider_ids = crate::commands::usage::parse_provider_settings(&settings.usage_providers); let data_loader = DataLoader::with_filters( - config.sessions_path.map(std::path::PathBuf::from), + config.home_dir.map(std::path::PathBuf::from), config.since, config.until, config.year, @@ -480,15 +497,16 @@ impl App { let dialog_stack = DialogStack::new(theme.clone()); let dialog_needs_reload = Rc::new(RefCell::new(false)); let requested_tab = config.initial_tab.unwrap_or(Tab::Overview); - let current_tab = if Self::tab_visible(&settings, requested_tab) { - requested_tab - } else { - Tab::Overview - }; + if !Self::tab_visible(&settings, requested_tab) { + anyhow::bail!( + "TUI tab `{}` is disabled in settings.json", + requested_tab.as_str().to_ascii_lowercase() + ); + } + let current_tab = requested_tab; let (sort_field, sort_direction) = Self::default_sort_for_tab(current_tab); let mut app = Self { - should_quit: false, current_tab, theme, settings, @@ -774,28 +792,26 @@ impl App { } } - pub fn handle_key_event(&mut self, key: KeyEvent) -> bool { + pub(crate) fn handle_key_event(&mut self, key: KeyEvent) -> KeyEventOutcome { if key.code == KeyCode::Char('c') && key.modifiers.contains(KeyModifiers::CONTROL) { - self.should_quit = true; - return true; + return KeyEventOutcome::Exit(TuiExit::Interrupted); } if self.dialog_stack.is_active() { self.dialog_stack.handle_key(key); self.consume_dialog_reload_if_ready(); - return false; + return KeyEventOutcome::Continue; } if let Some(command) = move_command_from_key(key.code) { if self.apply_text_viewport_move(command) { - return false; + return KeyEventOutcome::Continue; } } match key.code { KeyCode::Char('q') => { - self.should_quit = true; - return true; + return KeyEventOutcome::Exit(TuiExit::Quit); } KeyCode::Tab => { let next = self.next_visible_tab(); @@ -922,7 +938,7 @@ impl App { } _ => {} } - false + KeyEventOutcome::Continue } pub fn fetch_subscription_usage(&mut self) { @@ -2148,7 +2164,8 @@ mod tests { TuiConfig { theme: theme.map(str::to_string), refresh: 0, - sessions_path: None, + no_refresh: false, + home_dir: None, clients: None, since: None, until: None, @@ -2277,7 +2294,8 @@ mod tests { let config = TuiConfig { theme: Some("blue".to_string()), refresh: 0, - sessions_path: None, + no_refresh: false, + home_dir: None, clients: None, since: None, until: None, @@ -2326,7 +2344,8 @@ mod tests { let config = TuiConfig { theme: Some("blue".to_string()), refresh: 0, - sessions_path: None, + no_refresh: false, + home_dir: None, clients: None, since: None, until: None, @@ -2375,7 +2394,8 @@ mod tests { let config = TuiConfig { theme: Some("blue".to_string()), refresh: 0, - sessions_path: None, + no_refresh: false, + home_dir: None, clients: None, since: None, until: None, @@ -2415,7 +2435,8 @@ mod tests { let config = TuiConfig { theme: Some("blue".to_string()), refresh: 0, - sessions_path: None, + no_refresh: false, + home_dir: None, clients: None, since: None, until: None, @@ -2444,23 +2465,6 @@ mod tests { assert_eq!(app.sort_direction, SortDirection::Descending); } - #[test] - fn test_should_quit() { - let config = TuiConfig { - theme: Some("blue".to_string()), - refresh: 0, - sessions_path: None, - clients: None, - since: None, - until: None, - year: None, - initial_tab: None, - }; - let app = App::new_with_cached_data(config, None).unwrap(); - - assert!(!app.should_quit); - } - // ── Helper ────────────────────────────────────────────────────── fn test_settings() -> Settings { @@ -2475,7 +2479,8 @@ mod tests { let config = TuiConfig { theme: Some("blue".to_string()), refresh: 0, - sessions_path: None, + no_refresh: false, + home_dir: None, clients: None, since: None, until: None, @@ -2505,7 +2510,8 @@ mod tests { let config = TuiConfig { theme: Some("blue".to_string()), refresh: 0, - sessions_path: None, + no_refresh: false, + home_dir: None, clients: None, since: None, until: None, @@ -2524,7 +2530,7 @@ mod tests { actual, expected, "no-filter TUI must select exactly the accepted client catalog" ); - assert!(actual.contains(&ClientId::Cursor)); + assert!(actual.contains(&ClientId::Claude)); } fn make_app_with_models(n: usize) -> App { @@ -2623,17 +2629,15 @@ mod tests { #[test] fn test_handle_key_quit_q() { let mut app = make_app(); - let quit = app.handle_key_event(key(KeyCode::Char('q'))); - assert!(quit); - assert!(app.should_quit); + let outcome = app.handle_key_event(key(KeyCode::Char('q'))); + assert_eq!(outcome, KeyEventOutcome::Exit(TuiExit::Quit)); } #[test] fn test_handle_key_quit_ctrl_c() { let mut app = make_app(); - let quit = app.handle_key_event(key_with_mod(KeyCode::Char('c'), KeyModifiers::CONTROL)); - assert!(quit); - assert!(app.should_quit); + let outcome = app.handle_key_event(key_with_mod(KeyCode::Char('c'), KeyModifiers::CONTROL)); + assert_eq!(outcome, KeyEventOutcome::Exit(TuiExit::Interrupted)); } #[test] @@ -2641,10 +2645,9 @@ mod tests { let mut app = make_app(); app.open_client_picker(); - let quit = app.handle_key_event(key_with_mod(KeyCode::Char('c'), KeyModifiers::CONTROL)); + let outcome = app.handle_key_event(key_with_mod(KeyCode::Char('c'), KeyModifiers::CONTROL)); - assert!(quit); - assert!(app.should_quit); + assert_eq!(outcome, KeyEventOutcome::Exit(TuiExit::Interrupted)); assert!(!*app.dialog_needs_reload.borrow()); } @@ -2739,24 +2742,54 @@ mod tests { } #[test] - fn test_initial_usage_tab_clamps_to_overview_when_flag_off() { + fn cli_no_refresh_overrides_enabled_setting_for_this_run() { let config = TuiConfig { theme: Some("blue".to_string()), refresh: 0, - sessions_path: None, + no_refresh: true, + home_dir: None, + clients: None, + since: None, + until: None, + year: None, + initial_tab: None, + }; + let settings = Settings { + auto_refresh_enabled: true, + ..Settings::default() + }; + + let app = + App::new_with_cached_data_and_settings(config, Some(UsageData::default()), settings) + .unwrap(); + + assert!(!app.auto_refresh); + assert!(app.settings.auto_refresh_enabled); + } + + #[test] + fn test_initial_usage_tab_fails_when_flag_off() { + let config = TuiConfig { + theme: Some("blue".to_string()), + refresh: 0, + no_refresh: false, + home_dir: None, clients: None, since: None, until: None, year: None, initial_tab: Some(Tab::Usage), }; - let app = App::new_with_cached_data_and_settings( + let result = App::new_with_cached_data_and_settings( config, Some(UsageData::default()), Settings::default(), - ) - .unwrap(); - assert_eq!(app.current_tab, Tab::Overview); + ); + let error = match result { + Ok(_) => panic!("a disabled explicit tab must not silently fall back"), + Err(error) => error, + }; + assert!(error.to_string().contains("disabled in settings.json")); } #[test] @@ -3407,7 +3440,8 @@ mod tests { let config = TuiConfig { theme: Some("blue".to_string()), refresh: 0, - sessions_path: None, + no_refresh: false, + home_dir: None, clients: None, since: None, until: None, @@ -3475,7 +3509,8 @@ mod tests { let config = TuiConfig { theme: Some("blue".to_string()), refresh: 0, - sessions_path: None, + no_refresh: false, + home_dir: None, clients: None, since: None, until: None, @@ -3885,11 +3920,10 @@ mod tests { } #[test] - fn test_handle_key_unrecognized_returns_false() { + fn test_handle_key_unrecognized_continues() { let mut app = make_app(); - let result = app.handle_key_event(key(KeyCode::F(12))); - assert!(!result); - assert!(!app.should_quit); + let outcome = app.handle_key_event(key(KeyCode::F(12))); + assert_eq!(outcome, KeyEventOutcome::Continue); } #[test] diff --git a/crates/tokscale-cli/src/tui/cache.rs b/crates/tokscale-cli/src/tui/cache.rs index 50de3c106..04ae63b61 100644 --- a/crates/tokscale-cli/src/tui/cache.rs +++ b/crates/tokscale-cli/src/tui/cache.rs @@ -23,19 +23,59 @@ use super::data::{ /// Cache staleness threshold: 5 minutes (matches TS implementation) const CACHE_STALE_THRESHOLD_MS: u64 = 5 * 60 * 1000; -const CACHE_SCHEMA_VERSION: u32 = 35; +const CACHE_SCHEMA_VERSION: u32 = 37; #[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct CacheReportScope { + pub resolved_home_dir: String, + pub use_env_roots: bool, pub since: Option, pub until: Option, pub year: Option, } impl CacheReportScope { - pub fn new(since: Option, until: Option, year: Option) -> Self { - Self { since, until, year } + fn new( + resolved_home_dir: String, + use_env_roots: bool, + since: Option, + until: Option, + year: Option, + ) -> Self { + Self { + resolved_home_dir, + use_env_roots, + since, + until, + year, + } + } + + pub fn for_request( + home_dir: Option, + since: Option, + until: Option, + year: Option, + ) -> anyhow::Result { + let (resolved_home_dir, use_env_roots) = match home_dir { + Some(home_dir) => (home_dir, false), + None => ( + dirs::home_dir() + .ok_or_else(|| anyhow::anyhow!("Could not find home directory"))? + .to_string_lossy() + .into_owned(), + true, + ), + }; + + Ok(Self::new( + resolved_home_dir, + use_env_roots, + since, + until, + year, + )) } } @@ -1257,13 +1297,13 @@ mod tests { }, ); - let mut cursor_daily_models = BTreeMap::new(); - cursor_daily_models.insert( - "cursor-model".to_string(), + let mut gemini_daily_models = BTreeMap::new(); + gemini_daily_models.insert( + "gemini-model".to_string(), DailyModelInfo { - provider: "cursor".to_string(), - display_name: "Cursor Model".to_string(), - color_key: "cursor-model".to_string(), + provider: "google".to_string(), + display_name: "Gemini Model".to_string(), + color_key: "gemini-model".to_string(), tokens: token_breakdown(33), cost: 3.3, messages: 9, @@ -1271,11 +1311,11 @@ mod tests { ); let mut source_breakdown = BTreeMap::new(); source_breakdown.insert( - "cursor".to_string(), + "gemini".to_string(), DailySourceInfo { tokens: token_breakdown(22), cost: 2.2, - models: cursor_daily_models, + models: gemini_daily_models, }, ); source_breakdown.insert( @@ -1308,7 +1348,7 @@ mod tests { cost: 5.2, }, ); - let hourly_clients = ["claude".to_string(), "cursor".to_string()] + let hourly_clients = ["claude".to_string(), "gemini".to_string()] .into_iter() .collect(); @@ -1387,8 +1427,10 @@ mod tests { env::remove_var("TOKSCALE_CONFIG_DIR"); } - let clients = make_filters(&[ClientId::Cursor, ClientId::Claude]); + let clients = make_filters(&[ClientId::Gemini, ClientId::Claude]); let scope = CacheReportScope::new( + temp_dir.path().to_string_lossy().into_owned(), + true, Some("2026-07-01".to_string()), Some("2026-07-11".to_string()), Some("2026".to_string()), @@ -1415,7 +1457,7 @@ mod tests { assert_eq!(value["schemaVersion"], CACHE_SCHEMA_VERSION); assert_eq!( value["enabledClients"], - serde_json::json!(["claude", "cursor"]) + serde_json::json!(["claude", "gemini"]) ); assert_eq!(value["data"]["daily"][0]["date"], "2026-07-11"); assert_eq!( @@ -1432,7 +1474,7 @@ mod tests { assert_eq!(value["data"]["health"]["complete"], true); assert_eq!( tuple_array_keys(&value["data"]["daily"][0]["sourceBreakdown"]), - vec!["claude", "cursor"] + vec!["claude", "gemini"] ); assert_eq!( tuple_array_keys(&value["data"]["daily"][0]["sourceBreakdown"][0][1]["models"]), @@ -1458,7 +1500,7 @@ mod tests { ); assert_eq!( ordered.field("reportScope").keys(), - vec!["since", "until", "year"] + vec!["resolvedHomeDir", "useEnvRoots", "since", "until", "year",] ); let ordered_data = ordered.field("data"); @@ -1896,6 +1938,8 @@ mod tests { let clients = make_filters(&[ClientId::Claude]); let filtered_scope = CacheReportScope::new( + temp_dir.path().to_string_lossy().into_owned(), + true, Some("2026-05-01".to_string()), Some("2026-05-07".to_string()), None, @@ -1919,12 +1963,75 @@ mod tests { CacheResult::Fresh(_, _) )); + let other_home_scope = CacheReportScope::new( + "/tmp/other-tokscale-home".to_string(), + false, + Some("2026-05-01".to_string()), + Some("2026-05-07".to_string()), + None, + ); + assert!(matches!( + load_cache(&clients, &GroupBy::Model, &other_home_scope), + CacheResult::Miss + )); + match previous_home { Some(home) => unsafe { env::set_var("HOME", home) }, None => unsafe { env::remove_var("HOME") }, } } + #[test] + #[serial] + fn implicit_home_and_environment_roots_participate_in_cache_scope() { + let root = TempDir::new().unwrap(); + let first_home = root.path().join("first-home"); + let second_home = root.path().join("second-home"); + let shared_config = root.path().join("shared-config"); + fs::create_dir_all(&first_home).unwrap(); + fs::create_dir_all(&second_home).unwrap(); + + let _home_guard = EnvVarGuard::set("HOME", first_home.as_os_str()); + let _config_guard = EnvVarGuard::set("TOKSCALE_CONFIG_DIR", shared_config.as_os_str()); + let clients = make_filters(&[ClientId::Claude]); + let first_scope = CacheReportScope::for_request(None, None, None, None).unwrap(); + save_cached_data( + &UsageData::default(), + &clients, + &GroupBy::Model, + &first_scope, + test_signature(), + ) + .unwrap(); + + unsafe { + env::set_var("HOME", &second_home); + } + let second_scope = CacheReportScope::for_request(None, None, None, None).unwrap(); + assert_ne!(first_scope, second_scope); + assert!(first_scope.use_env_roots); + assert!(second_scope.use_env_roots); + assert!(matches!( + load_cache(&clients, &GroupBy::Model, &second_scope), + CacheResult::Miss + )); + + let explicit_second_scope = CacheReportScope::for_request( + Some(second_home.to_string_lossy().into_owned()), + None, + None, + None, + ) + .unwrap(); + assert_eq!( + second_scope.resolved_home_dir, + explicit_second_scope.resolved_home_dir + ); + assert!(second_scope.use_env_roots); + assert!(!explicit_second_scope.use_env_roots); + assert_ne!(second_scope, explicit_second_scope); + } + #[test] #[serial] fn test_load_cache_treats_future_timestamp_as_stale() { @@ -2285,7 +2392,7 @@ mod tests { r#"{ "schemaVersion": 24, "timestamp": 0, - "enabledClients": ["claude", "cursor"], + "enabledClients": ["claude", "gemini"], "groupBy": "model", "reportScope": { "since": null, @@ -2335,7 +2442,7 @@ mod tests { ]] } ], [ - "cursor", + "gemini", { "tokens": { "input": 20, @@ -2379,16 +2486,17 @@ mod tests { .unwrap(); cached["timestamp"] = serde_json::Value::from(fresh_timestamp_ms()); cached["schemaVersion"] = serde_json::Value::from(CACHE_SCHEMA_VERSION); + cached["reportScope"] = serde_json::to_value(CacheReportScope::default()).unwrap(); cached["sourceInventorySignature"] = serde_json::json!(vec![0x5a_u8; 32]); fs::write(&cache_path, serde_json::to_vec(&cached).unwrap()).unwrap(); - let clients = make_filters(&[ClientId::Claude, ClientId::Cursor]); + let clients = make_filters(&[ClientId::Claude, ClientId::Gemini]); match load_cache(&clients, &GroupBy::Model, &CacheReportScope::default()) { CacheResult::Fresh(data, signature) => { assert_eq!(signature, test_signature()); assert_eq!(data.daily[0].source_breakdown.len(), 2); - let cursor = data.daily[0].source_breakdown.get("cursor").unwrap(); - let model = cursor.models.get("claude-sonnet-4").unwrap(); + let gemini = data.daily[0].source_breakdown.get("gemini").unwrap(); + let model = gemini.models.get("claude-sonnet-4").unwrap(); assert_eq!(model.provider, "anthropic"); assert_eq!(model.tokens.total(), 30); } diff --git a/crates/tokscale-cli/src/tui/colors.rs b/crates/tokscale-cli/src/tui/colors.rs index ede1dccef..1900dd95d 100644 --- a/crates/tokscale-cli/src/tui/colors.rs +++ b/crates/tokscale-cli/src/tui/colors.rs @@ -44,7 +44,6 @@ pub fn get_provider_shade(provider: &str, rank: usize) -> Color { s if s.contains("deepseek") => &DEEPSEEK_SHADES, s if s.contains("xai") || s.contains("grok") => &XAI_SHADES, s if s.contains("meta") || s.contains("llama") => &META_SHADES, - s if s.contains("cursor") => &CURSOR_SHADES, _ => &UNKNOWN_SHADES, }; @@ -130,16 +129,6 @@ const META_SHADES: [(u8, u8, u8); 7] = [ (225, 226, 252), // #E1E2FC ]; -const CURSOR_SHADES: [(u8, u8, u8); 7] = [ - (139, 92, 246), // #8B5CF6 - (154, 114, 247), // #9A72F7 - (169, 135, 248), // #A987F8 - (184, 156, 250), // #B89CFA - (199, 177, 251), // #C7B1FB - (215, 199, 252), // #D7C7FC - (230, 220, 253), // #E6DCFD -]; - /// Neutral gray ramp for providers that don't match any known palette. /// Still produces distinct shades per rank instead of collapsing to white. const UNKNOWN_SHADES: [(u8, u8, u8); 7] = [ @@ -311,14 +300,6 @@ mod tests { assert_eq!(provider_color_key(" , anthropic"), "anthropic"); } - #[test] - fn cursor_provider_has_distinct_shades_per_rank() { - TokscaleConfig::initialize_default_for_tests(); - let rank_0 = get_provider_shade("cursor", 0); - let rank_6 = get_provider_shade("cursor", 6); - assert_ne!(rank_0, rank_6); - } - #[test] fn get_provider_shade_saturates_at_palette_end() { TokscaleConfig::initialize_default_for_tests(); diff --git a/crates/tokscale-cli/src/tui/data/mod.rs b/crates/tokscale-cli/src/tui/data/mod.rs index 93c1a75bb..30a7234a0 100644 --- a/crates/tokscale-cli/src/tui/data/mod.rs +++ b/crates/tokscale-cli/src/tui/data/mod.rs @@ -34,12 +34,19 @@ pub use tokscale_core::{ /// hermetic across developer machines; production builds still honor /// user-configured paths. #[cfg(not(test))] -fn data_loader_scanner_settings() -> Result { - crate::tui::settings::load_scanner_settings() +fn data_loader_scanner_settings( + home_dir: &Option, +) -> Result { + let home = home_dir + .as_ref() + .map(|path| path.to_string_lossy().into_owned()); + Ok(crate::tui::settings::load_scanner_settings_for_home(&home)?) } #[cfg(test)] -fn data_loader_scanner_settings() -> Result { +fn data_loader_scanner_settings( + _home_dir: &Option, +) -> Result { Ok(tokscale_core::scanner::ScannerSettings::default()) } @@ -54,7 +61,7 @@ pub(super) fn trim_allocator() { } pub struct DataLoader { - _sessions_path: Option, + pub home_dir: Option, pub since: Option, pub until: Option, pub year: Option, @@ -81,13 +88,13 @@ impl PreparedDataLoad { impl DataLoader { pub fn with_filters( - sessions_path: Option, + home_dir: Option, since: Option, until: Option, year: Option, ) -> Self { Self { - _sessions_path: sessions_path, + home_dir, since, until, year, @@ -110,10 +117,16 @@ impl DataLoader { } pub fn prepare(&self, enabled_clients: &[ClientId]) -> Result { - let home = dirs::home_dir() - .ok_or_else(|| anyhow::anyhow!("Could not find home directory"))? - .to_string_lossy() - .to_string(); + let (home, use_env_roots) = match &self.home_dir { + Some(home) => (home.to_string_lossy().into_owned(), false), + None => ( + dirs::home_dir() + .ok_or_else(|| anyhow::anyhow!("Could not find home directory"))? + .to_string_lossy() + .into_owned(), + true, + ), + }; let sources: Vec = enabled_clients .iter() @@ -122,17 +135,17 @@ impl DataLoader { let opts = LocalParseOptions { home_dir: Some(home), - use_env_roots: true, + use_env_roots, clients: Some(sources), since: self.since.clone(), until: self.until.clone(), year: self.year.clone(), - scanner_settings: data_loader_scanner_settings()?, + scanner_settings: data_loader_scanner_settings(&self.home_dir)?, }; prepare_local_sources(opts) .map(|sources| PreparedDataLoad { sources }) - .map_err(anyhow::Error::msg) + .map_err(anyhow::Error::new) } pub fn execute_with_diagnostics( @@ -142,34 +155,35 @@ impl DataLoader { ) -> Result { let group_by = group_by.clone(); - let usage_data = if Handle::try_current().is_ok() { + let usage_data: Result<_> = if Handle::try_current().is_ok() { std::thread::scope(|s| { - s.spawn(move || { - let rt = Runtime::new().map_err(|e| e.to_string())?; + s.spawn(move || -> Result<_> { + let rt = Runtime::new()?; rt.block_on(load_prepared_usage_data_with_diagnostics( prepared.sources, group_by, )) + .map_err(anyhow::Error::new) }) .join() - .unwrap_or_else(|_| Err("data loader thread panicked".to_string())) + .unwrap_or_else(|_| Err(anyhow::anyhow!("data loader thread panicked"))) }) } else { - Runtime::new()?.block_on(load_prepared_usage_data_with_diagnostics( - prepared.sources, - group_by, - )) + Runtime::new()? + .block_on(load_prepared_usage_data_with_diagnostics( + prepared.sources, + group_by, + )) + .map_err(anyhow::Error::new) }; trim_allocator(); - usage_data - .map(|result| DataLoadResult { - data: result.data, - pricing_diagnostics: result.pricing_diagnostics, - source_inventory_signature: result.source_inventory_signature, - source_digest: result.source_inventory_signature.process_digest(), - }) - .map_err(anyhow::Error::msg) + usage_data.map(|result| DataLoadResult { + data: result.data, + pricing_diagnostics: result.pricing_diagnostics, + source_inventory_signature: result.source_inventory_signature, + source_digest: result.source_inventory_signature.process_digest(), + }) } #[cfg(test)] @@ -180,10 +194,16 @@ impl DataLoader { group_by: &GroupBy, pricing: &tokscale_core::pricing::PricingService, ) -> Result { - let home = dirs::home_dir() - .ok_or_else(|| anyhow::anyhow!("Could not find home directory"))? - .to_string_lossy() - .to_string(); + let (home, use_env_roots) = match &self.home_dir { + Some(home) => (home.to_string_lossy().into_owned(), false), + None => ( + dirs::home_dir() + .ok_or_else(|| anyhow::anyhow!("Could not find home directory"))? + .to_string_lossy() + .into_owned(), + true, + ), + }; let sources: Vec = enabled_clients .iter() @@ -196,13 +216,13 @@ impl DataLoader { since: self.since.clone(), until: self.until.clone(), year: self.year.clone(), - use_env_roots: false, - scanner_settings: data_loader_scanner_settings()?, + use_env_roots, + scanner_settings: data_loader_scanner_settings(&self.home_dir)?, }; let usage_data = tokscale_core::load_usage_data_with_pricing(opts, group_by.clone(), Some(pricing)) - .map_err(anyhow::Error::msg)?; + .map_err(anyhow::Error::new)?; Ok(usage_data) } @@ -256,10 +276,16 @@ mod tests { group_by: &GroupBy, pricing: Option<&PricingService>, ) -> Result { - let home = dirs::home_dir() - .ok_or_else(|| anyhow::anyhow!("Could not find home directory"))? - .to_string_lossy() - .to_string(); + let (home, use_env_roots) = match &loader.home_dir { + Some(home) => (home.to_string_lossy().into_owned(), false), + None => ( + dirs::home_dir() + .ok_or_else(|| anyhow::anyhow!("Could not find home directory"))? + .to_string_lossy() + .into_owned(), + true, + ), + }; let sources: Vec = enabled_clients .iter() @@ -268,16 +294,16 @@ mod tests { let opts = LocalParseOptions { home_dir: Some(home), - use_env_roots: true, + use_env_roots, clients: Some(sources), since: loader.since.clone(), until: loader.until.clone(), year: loader.year.clone(), - scanner_settings: data_loader_scanner_settings()?, + scanner_settings: data_loader_scanner_settings(&loader.home_dir)?, }; tokscale_core::load_usage_data_with_pricing(opts, group_by.clone(), pricing) - .map_err(anyhow::Error::msg) + .map_err(anyhow::Error::new) } fn expected_message_cost( @@ -329,7 +355,6 @@ mod tests { assert_eq!(ClientId::short_name(ClientId::Claude), "Claude"); assert_eq!(ClientId::short_name(ClientId::Codex), "Codex"); assert_eq!(ClientId::short_name(ClientId::Copilot), "Copilot"); - assert_eq!(ClientId::short_name(ClientId::Cursor), "Cursor"); assert_eq!(ClientId::short_name(ClientId::Gemini), "Gemini"); assert_eq!(ClientId::short_name(ClientId::Amp), "Amp"); assert_eq!(ClientId::short_name(ClientId::Droid), "Droid"); @@ -349,7 +374,6 @@ mod tests { assert_eq!(ClientId::short_name(ClientId::Zed), "Zed Agent"); assert_eq!(ClientId::short_name(ClientId::Zcode), "ZCode"); assert_eq!(ClientId::short_name(ClientId::Kiro), "Kiro"); - assert_eq!(ClientId::short_name(ClientId::Trae), "Trae"); assert_eq!(ClientId::short_name(ClientId::Cline), "Cline"); } @@ -359,7 +383,6 @@ mod tests { assert_eq!(ClientId::hotkey(ClientId::Claude), Some('2')); assert_eq!(ClientId::hotkey(ClientId::Codex), Some('3')); assert_eq!(ClientId::hotkey(ClientId::Copilot), Some('c')); - assert_eq!(ClientId::hotkey(ClientId::Cursor), Some('4')); assert_eq!(ClientId::hotkey(ClientId::Gemini), Some('5')); assert_eq!(ClientId::hotkey(ClientId::Amp), Some('6')); assert_eq!(ClientId::hotkey(ClientId::Droid), Some('7')); @@ -379,7 +402,6 @@ mod tests { assert_eq!(ClientId::hotkey(ClientId::Zed), Some('z')); assert_eq!(ClientId::hotkey(ClientId::Zcode), Some('q')); assert_eq!(ClientId::hotkey(ClientId::Kiro), Some('i')); - assert_eq!(ClientId::hotkey(ClientId::Trae), Some('y')); assert_eq!(ClientId::hotkey(ClientId::Cline), Some('n')); } @@ -389,7 +411,7 @@ mod tests { assert_eq!(ClientId::from_hotkey('2'), Some(ClientId::Claude)); assert_eq!(ClientId::from_hotkey('3'), Some(ClientId::Codex)); assert_eq!(ClientId::from_hotkey('c'), Some(ClientId::Copilot)); - assert_eq!(ClientId::from_hotkey('4'), Some(ClientId::Cursor)); + assert_eq!(ClientId::from_hotkey('4'), None); assert_eq!(ClientId::from_hotkey('5'), Some(ClientId::Gemini)); assert_eq!(ClientId::from_hotkey('6'), Some(ClientId::Amp)); assert_eq!(ClientId::from_hotkey('7'), Some(ClientId::Droid)); @@ -409,7 +431,7 @@ mod tests { assert_eq!(ClientId::from_hotkey('z'), Some(ClientId::Zed)); assert_eq!(ClientId::from_hotkey('q'), Some(ClientId::Zcode)); assert_eq!(ClientId::from_hotkey('i'), Some(ClientId::Kiro)); - assert_eq!(ClientId::from_hotkey('y'), Some(ClientId::Trae)); + assert_eq!(ClientId::from_hotkey('y'), None); } #[test] @@ -451,7 +473,7 @@ mod tests { #[test] fn test_data_loader_new() { let loader = DataLoader::with_filters(None, None, None, None); - assert!(loader._sessions_path.is_none()); + assert!(loader.home_dir.is_none()); assert!(loader.since.is_none()); assert!(loader.until.is_none()); assert!(loader.year.is_none()); @@ -470,7 +492,7 @@ mod tests { // instead it asserts the cfg(test) helper returns a default // ScannerSettings regardless of what the real settings file // contains on the developer's machine. - let settings = super::data_loader_scanner_settings().unwrap(); + let settings = super::data_loader_scanner_settings(&None).unwrap(); assert!( settings.opencode_db_paths.is_empty(), "under #[cfg(test)] data_loader_scanner_settings must return \ @@ -489,7 +511,7 @@ mod tests { Some("2024".to_string()), ); - assert_eq!(loader._sessions_path, Some(PathBuf::from("/tmp/sessions"))); + assert_eq!(loader.home_dir, Some(PathBuf::from("/tmp/sessions"))); assert_eq!(loader.since, Some("2024-01-01".to_string())); assert_eq!(loader.until, Some("2024-12-31".to_string())); assert_eq!(loader.year, Some("2024".to_string())); diff --git a/crates/tokscale-cli/src/tui/mod.rs b/crates/tokscale-cli/src/tui/mod.rs index ca60f4cb0..22bd54014 100644 --- a/crates/tokscale-cli/src/tui/mod.rs +++ b/crates/tokscale-cli/src/tui/mod.rs @@ -10,12 +10,14 @@ pub mod settings; mod themes; mod ui; -pub use app::{App, Tab, TuiConfig}; +use app::KeyEventOutcome; +pub use app::{App, Tab, TuiConfig, TuiExit}; pub use cache::{ load_cache, save_cached_data, CacheReportScope, CacheResult, TUI_DEFAULT_GROUP_BY, }; pub use data::{DataLoader, UsageData}; pub use event::{Event, EventHandler}; +pub(crate) use themes::ThemeName; use std::collections::HashSet; use std::io; @@ -59,11 +61,12 @@ fn decide_initial_data(load_result: CacheResult) -> (Option, bool, Op } fn background_data_loader( + home_dir: Option, since: Option, until: Option, year: Option, ) -> DataLoader { - DataLoader::with_filters(None, since, until, year) + DataLoader::with_filters(home_dir.map(std::path::PathBuf::from), since, until, year) } fn should_force_source_reload( @@ -214,24 +217,27 @@ fn send_background_result( } fn background_cache_scope( + home_dir: &Option, since: &Option, until: &Option, year: &Option, -) -> CacheReportScope { - CacheReportScope::new(since.clone(), until.clone(), year.clone()) +) -> Result { + CacheReportScope::for_request(home_dir.clone(), since.clone(), until.clone(), year.clone()) } #[allow(clippy::too_many_arguments)] pub fn run( theme: Option<&str>, - refresh: u64, + refresh: Option, + no_refresh: bool, debug: bool, + home_dir: Option, clients: Option>, since: Option, until: Option, year: Option, initial_tab: Option, -) -> Result<()> { +) -> Result { if debug { let _ = tracing_subscriber::fmt() .with_env_filter("debug") @@ -241,8 +247,9 @@ pub fn run( let config = TuiConfig { theme: theme.map(str::to_string), - refresh, - sessions_path: None, + refresh: refresh.unwrap_or(0), + no_refresh, + home_dir: home_dir.clone(), clients: clients.clone(), since: since.clone(), until: until.clone(), @@ -266,7 +273,7 @@ pub fn run( // Single file read: load cache and check freshness in one pass. let initial_group_by = TUI_DEFAULT_GROUP_BY; - let initial_report_scope = background_cache_scope(&since, &until, &year); + let initial_report_scope = background_cache_scope(&home_dir, &since, &until, &year)?; let (cached_data, needs_background_load, initial_source_digest) = decide_initial_data( load_cache(&enabled_clients, &initial_group_by, &initial_report_scope), ); @@ -318,12 +325,13 @@ pub fn run( let bg_since = since.clone(); let bg_until = until.clone(); let bg_year = year.clone(); + let bg_home_dir = home_dir.clone(); let bg_enabled_clients = enabled_clients.clone(); let bg_group_by = app.group_by.borrow().clone(); - let bg_report_scope = background_cache_scope(&since, &until, &year); + let bg_report_scope = background_cache_scope(&home_dir, &since, &until, &year)?; thread::spawn(move || { - let loader = background_data_loader(bg_since, bg_until, bg_year); + let loader = background_data_loader(bg_home_dir, bg_since, bg_until, bg_year); let result = persist_background_load( load_background_data(&loader, &bg_clients, &bg_group_by, true, None), &bg_enabled_clients, @@ -391,7 +399,7 @@ fn run_loop_with_background( bg_tx: mpsc::Sender>, bg_rx: mpsc::Receiver>, #[cfg(unix)] sigcont_flag: &Arc, -) -> Result<()> { +) -> Result { loop { #[cfg(unix)] if sigcont_flag.swap(false, Ordering::Relaxed) { @@ -432,12 +440,17 @@ fn run_loop_with_background( let since = app.data_loader.since.clone(); let until = app.data_loader.until.clone(); let year = app.data_loader.year.clone(); + let home_dir = app + .data_loader + .home_dir + .as_ref() + .map(|path| path.to_string_lossy().into_owned()); let enabled_clients = app.enabled_clients.borrow().clone(); let group_by = app.group_by.borrow().clone(); - let report_scope = background_cache_scope(&since, &until, &year); + let report_scope = background_cache_scope(&home_dir, &since, &until, &year)?; thread::spawn(move || { - let loader = background_data_loader(since, until, year); + let loader = background_data_loader(home_dir, since, until, year); let result = persist_background_load( load_background_data(&loader, &clients, &group_by, force, last_digest), &enabled_clients, @@ -453,8 +466,8 @@ fn run_loop_with_background( app.on_tick(); } Event::Key(key) => { - if app.handle_key_event(key) { - break; + if let KeyEventOutcome::Exit(exit) = app.handle_key_event(key) { + return Ok(exit); } } Event::Mouse(mouse) => { @@ -464,12 +477,7 @@ fn run_loop_with_background( app.handle_resize(w, h); } } - - if app.should_quit { - break; - } } - Ok(()) } #[cfg(test)] @@ -591,6 +599,7 @@ mod tests { #[test] fn background_loader_preserves_filters() { let loader = background_data_loader( + None, Some("2026-05-01".to_string()), Some("2026-05-19".to_string()), Some("2026".to_string()), @@ -607,7 +616,7 @@ mod tests { let home = TempDir::new().unwrap(); let _guard = EnvGuard::set(home.path()); write_amp_source(home.path(), 10); - let loader = background_data_loader(None, None, None); + let loader = background_data_loader(None, None, None, None); let clients = [ClientId::Amp]; let signature_a = loader .load_with_diagnostics(&clients, &tokscale_core::GroupBy::Model) @@ -653,7 +662,7 @@ mod tests { let home = TempDir::new().unwrap(); let _guard = EnvGuard::set(home.path()); write_amp_source(home.path(), 10); - let loader = background_data_loader(None, None, None); + let loader = background_data_loader(None, None, None, None); let clients = [ClientId::Amp]; let mut prepared = loader.prepare(&clients).unwrap(); let baseline = Some( @@ -762,7 +771,8 @@ mod tests { TuiConfig { theme: Some("blue".to_string()), refresh: 0, - sessions_path: None, + no_refresh: false, + home_dir: None, clients: None, since: None, until: None, diff --git a/crates/tokscale-cli/src/tui/settings.rs b/crates/tokscale-cli/src/tui/settings.rs index ea5f9458b..258b327f9 100644 --- a/crates/tokscale-cli/src/tui/settings.rs +++ b/crates/tokscale-cli/src/tui/settings.rs @@ -4,7 +4,8 @@ use std::time::Duration; use anyhow::{Context, Result}; use serde::{Deserialize, Serialize}; -use tokscale_core::scanner::ScannerSettings; +use tokscale_core::paths::ConfigDirUnavailable; +use tokscale_core::scanner::{ScannerSettings, ScannerSettingsError}; use super::themes::ThemeName; @@ -16,6 +17,64 @@ const DEFAULT_NATIVE_TIMEOUT_MS: u64 = 300_000; const MIN_NATIVE_TIMEOUT_MS: u64 = 5_000; const MAX_NATIVE_TIMEOUT_MS: u64 = 3_600_000; +#[derive(Debug, thiserror::Error)] +pub(crate) enum SettingsLoadError { + #[error(transparent)] + ConfigDirectory(#[from] ConfigDirUnavailable), + #[error("failed to read settings file `{path}`: {source}")] + Read { + path: PathBuf, + #[source] + source: std::io::Error, + }, + #[error("failed to parse settings JSON `{path}`: {source}")] + Parse { + path: PathBuf, + #[source] + source: serde_json::Error, + }, + #[error("invalid settings in `{path}`: {source}")] + Invalid { + path: PathBuf, + #[source] + source: SettingsValidationError, + }, +} + +impl SettingsLoadError { + pub(crate) const fn is_invalid_environment(&self) -> bool { + !matches!(self, Self::Read { .. }) + } +} + +#[derive(Debug, thiserror::Error)] +pub(crate) enum SettingsValidationError { + #[error("invalid autoRefreshMs {value}; expected {min}..={max}")] + AutoRefreshRange { value: u64, min: u64, max: u64 }, + #[error("invalid nativeTimeoutMs {value}; expected {min}..={max}")] + NativeTimeoutRange { value: u64, min: u64, max: u64 }, + #[error("invalid colorPalette `{value}`; expected one of: {valid}")] + ColorPalette { value: String, valid: String }, + #[error("invalid scanner settings: {0}")] + Scanner(#[from] ScannerSettingsError), +} + +#[derive(Debug, thiserror::Error)] +pub(crate) enum NativeTimeoutError { + #[error("TOKSCALE_NATIVE_TIMEOUT_MS must be a positive integer: {source}")] + NotInteger { + #[source] + source: std::num::ParseIntError, + }, + #[error("failed to read TOKSCALE_NATIVE_TIMEOUT_MS: {source}")] + NotUnicode { + #[source] + source: std::env::VarError, + }, + #[error("invalid TOKSCALE_NATIVE_TIMEOUT_MS {value}; expected {min}..={max}")] + OutOfRange { value: u64, min: u64, max: u64 }, +} + #[derive(Debug, Clone, Copy)] enum ExplicitHomeConfigLayout { UnixDotConfig, @@ -32,16 +91,6 @@ impl ExplicitHomeConfigLayout { } } -#[derive(Debug, Clone, Default, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct LightSettings { - /// When true, every `tokscale --light` run atomically overwrites the - /// TUI cache (same semantics as `--light --write-cache`). The CLI - /// flags `--write-cache` / `--no-write-cache` override this per-invocation. - #[serde(default)] - pub write_cache: bool, -} - #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct Settings { @@ -75,8 +124,6 @@ pub struct Settings { /// override this list completely. #[serde(default)] pub default_clients: Vec, - #[serde(default)] - pub light: LightSettings, /// Opt-in toggle for the subscription quota Usage tab. Default is /// `false` so the tab strip stays focused on local token usage unless /// the user explicitly wants subscription usage lookups. @@ -88,7 +135,6 @@ pub struct Settings { /// such as `codex`, `zai`, and `minimax-token-plan-cn`. #[serde(default)] pub usage_providers: Vec, - #[cfg(test)] #[serde(skip)] pub save_path_override: Option, } @@ -115,26 +161,16 @@ impl Default for Settings { native_timeout_ms: DEFAULT_NATIVE_TIMEOUT_MS, scanner: ScannerSettings::default(), default_clients: Vec::new(), - light: LightSettings::default(), usage_tab_enabled: false, usage_providers: Vec::new(), - #[cfg(test)] save_path_override: None, } } } -/// Thin helper that loads settings and returns just the scanner portion. -/// -/// Every CLI entry point that builds `LocalParseOptions`/`ReportOptions` -/// calls this so user-configured scanner paths are honored on every -/// invocation. A missing file means the user has not configured scanner -/// overrides; malformed or unreadable files are reported to the command. -pub fn load_scanner_settings() -> Result { - Settings::load().map(|settings| settings.scanner) -} - -pub fn load_scanner_settings_for_home(home_dir: &Option) -> Result { +pub fn load_scanner_settings_for_home( + home_dir: &Option, +) -> std::result::Result { Settings::load_for_home_override(home_dir.as_deref().map(Path::new)) .map(|settings| settings.scanner) } @@ -145,44 +181,48 @@ pub fn load_scanner_settings_for_home(home_dir: &Option) -> Result Result> { - Settings::load().map(|settings| settings.default_clients) -} - -pub fn load_default_clients_for_home(home_dir: &Option) -> Result> { +pub fn load_default_clients_for_home( + home_dir: &Option, +) -> std::result::Result, SettingsLoadError> { Settings::load_for_home_override(home_dir.as_deref().map(Path::new)) .map(|settings| settings.default_clients) } impl Settings { - fn validate(self) -> Result { + fn validate(self) -> std::result::Result { if !(MIN_AUTO_REFRESH_MS..=MAX_AUTO_REFRESH_MS).contains(&self.auto_refresh_ms) { - anyhow::bail!( - "invalid autoRefreshMs {}; expected {}..={}", - self.auto_refresh_ms, - MIN_AUTO_REFRESH_MS, - MAX_AUTO_REFRESH_MS - ); + return Err(SettingsValidationError::AutoRefreshRange { + value: self.auto_refresh_ms, + min: MIN_AUTO_REFRESH_MS, + max: MAX_AUTO_REFRESH_MS, + }); } if !(MIN_NATIVE_TIMEOUT_MS..=MAX_NATIVE_TIMEOUT_MS).contains(&self.native_timeout_ms) { - anyhow::bail!( - "invalid nativeTimeoutMs {}; expected {}..={}", - self.native_timeout_ms, - MIN_NATIVE_TIMEOUT_MS, - MAX_NATIVE_TIMEOUT_MS - ); + return Err(SettingsValidationError::NativeTimeoutRange { + value: self.native_timeout_ms, + min: MIN_NATIVE_TIMEOUT_MS, + max: MAX_NATIVE_TIMEOUT_MS, + }); } - self.theme_name()?; - self.scanner - .validate() - .context("invalid scanner settings")?; + if self.color_palette.parse::().is_err() { + let valid = ThemeName::all() + .iter() + .map(ThemeName::as_str) + .collect::>() + .join(", "); + return Err(SettingsValidationError::ColorPalette { + value: self.color_palette.clone(), + valid, + }); + } + self.scanner.validate()?; Ok(self) } - fn config_path() -> Result { + fn config_path() -> std::result::Result { crate::paths::try_get_config_dir() .map(|directory| directory.join("settings.json")) - .map_err(anyhow::Error::new) + .map_err(SettingsLoadError::from) } fn writable_config_path() -> Result { @@ -217,40 +257,54 @@ impl Settings { Self::explicit_home_config_path_for_layout(home_dir, ExplicitHomeConfigLayout::current()) } - fn load_from_path(path: &Path) -> Result { - let content = match fs::read_to_string(path) { + fn load_from_path(path: &Path) -> std::result::Result { + let content = match fs::read(path) { Ok(content) => content, Err(error) if error.kind() == std::io::ErrorKind::NotFound => { return Ok(Self::default()); } Err(error) => { - return Err(error) - .with_context(|| format!("failed to read settings file `{}`", path.display())); + return Err(SettingsLoadError::Read { + path: path.to_path_buf(), + source: error, + }); } }; - serde_json::from_str::(&content) - .with_context(|| format!("failed to parse settings JSON `{}`", path.display()))? + let settings = serde_json::from_slice::(&content).map_err(|source| { + SettingsLoadError::Parse { + path: path.to_path_buf(), + source, + } + })?; + settings .validate() - .with_context(|| format!("invalid settings in `{}`", path.display())) + .map_err(|source| SettingsLoadError::Invalid { + path: path.to_path_buf(), + source, + }) } - pub fn load() -> Result { + pub fn load() -> std::result::Result { Self::load_from_path(&Self::config_path()?) } - pub fn load_for_home_override(home_dir: Option<&Path>) -> Result { + pub fn load_for_home_override( + home_dir: Option<&Path>, + ) -> std::result::Result { let Some(home_dir) = home_dir else { return Self::load(); }; - Self::load_from_path(&Self::explicit_home_config_path(home_dir)) + let path = Self::explicit_home_config_path(home_dir); + let mut settings = Self::load_from_path(&path)?; + settings.save_path_override = Some(path); + Ok(settings) } pub fn save(&self) -> Result<()> { self.clone().validate()?; - #[cfg(test)] let path = self.save_path_override.clone().map_or_else( Self::writable_config_path, |path| -> Result { @@ -261,9 +315,6 @@ impl Settings { }, )?; - #[cfg(not(test))] - let path = Self::writable_config_path()?; - let content = serde_json::to_string_pretty(self)?; tokscale_core::fs_atomic::write_atomic(&path, content.as_bytes())?; @@ -302,20 +353,22 @@ impl Settings { } } - pub fn get_native_timeout(&self) -> Result { + pub fn get_native_timeout(&self) -> std::result::Result { let timeout_ms = match std::env::var("TOKSCALE_NATIVE_TIMEOUT_MS") { Ok(value) => value .parse::() - .with_context(|| "TOKSCALE_NATIVE_TIMEOUT_MS must be a positive integer")?, + .map_err(|source| NativeTimeoutError::NotInteger { source })?, Err(std::env::VarError::NotPresent) => self.native_timeout_ms, Err(source) => { - return Err(source).context("failed to read TOKSCALE_NATIVE_TIMEOUT_MS"); + return Err(NativeTimeoutError::NotUnicode { source }); } }; if !(MIN_NATIVE_TIMEOUT_MS..=MAX_NATIVE_TIMEOUT_MS).contains(&timeout_ms) { - anyhow::bail!( - "invalid TOKSCALE_NATIVE_TIMEOUT_MS {timeout_ms}; expected {MIN_NATIVE_TIMEOUT_MS}..={MAX_NATIVE_TIMEOUT_MS}" - ); + return Err(NativeTimeoutError::OutOfRange { + value: timeout_ms, + min: MIN_NATIVE_TIMEOUT_MS, + max: MAX_NATIVE_TIMEOUT_MS, + }); } Ok(Duration::from_millis(timeout_ms)) } @@ -324,6 +377,7 @@ impl Settings { #[cfg(test)] mod tests { use super::*; + use std::error::Error as _; use std::path::PathBuf; #[test] @@ -374,6 +428,19 @@ mod tests { assert!(loaded.default_clients.is_empty()); } + #[test] + fn settings_loaded_for_explicit_home_save_back_to_that_home() { + let temp = tempfile::TempDir::new().unwrap(); + let path = Settings::explicit_home_config_path(temp.path()); + let mut loaded = Settings::load_for_home_override(Some(temp.path())).unwrap(); + loaded.color_palette = "halloween".to_string(); + + loaded.save().unwrap(); + + let saved: serde_json::Value = serde_json::from_slice(&fs::read(&path).unwrap()).unwrap(); + assert_eq!(saved["colorPalette"], "halloween"); + } + #[test] fn load_for_home_override_does_not_read_legacy_macos_path() { let temp = tempfile::TempDir::new().unwrap(); @@ -400,12 +467,32 @@ mod tests { assert!(message.contains("parse settings JSON"), "{message}"); assert!(message.contains(&path.display().to_string()), "{message}"); + assert!(error.is_invalid_environment()); assert!( error.source().is_some(), "parse error must remain in the chain" ); } + #[test] + fn load_for_home_override_reports_non_utf8_json_as_invalid_environment() { + let temp = tempfile::TempDir::new().unwrap(); + let path = Settings::explicit_home_config_path(temp.path()); + fs::create_dir_all(path.parent().unwrap()).unwrap(); + fs::write(&path, b"{\"colorPalette\":\"\xff\"}").unwrap(); + + let error = Settings::load_for_home_override(Some(temp.path())).unwrap_err(); + let message = format!("{error:#}"); + + assert!(message.contains("parse settings JSON"), "{message}"); + assert!(message.contains(&path.display().to_string()), "{message}"); + assert!(error.is_invalid_environment()); + assert!( + error.source().is_some(), + "UTF-8 decoding failure must remain in the parse error chain" + ); + } + #[test] fn load_for_home_override_reports_non_file_path_with_operation_and_source() { let temp = tempfile::TempDir::new().unwrap(); @@ -417,6 +504,7 @@ mod tests { assert!(message.contains("read settings file"), "{message}"); assert!(message.contains(&path.display().to_string()), "{message}"); + assert!(!error.is_invalid_environment()); assert!( error.source().is_some(), "I/O error must remain in the chain" @@ -435,6 +523,7 @@ mod tests { assert!(message.contains("invalid settings"), "{message}"); assert!(message.contains("autoRefreshMs 1"), "{message}"); assert!(message.contains(&path.display().to_string()), "{message}"); + assert!(error.is_invalid_environment()); } #[test] @@ -448,6 +537,7 @@ mod tests { let message = format!("{error:#}"); assert!(message.contains("colorPalette `ultraviolet`"), "{message}"); assert!(message.contains(&path.display().to_string()), "{message}"); + assert!(error.is_invalid_environment()); } #[test] @@ -624,7 +714,7 @@ mod tests { #[test] fn settings_default_clients_round_trips() { // User-configured list must survive load+save unchanged. This is - // what `tokscale --client opencode,claude` consults when no CLI + // what `tokscale models --client opencode,claude` consults when no CLI // flag is present. let json = r#"{ "colorPalette": "blue", @@ -661,27 +751,6 @@ mod tests { assert!(serde_json::from_str::(json).is_err()); } - #[test] - fn settings_load_accepts_legacy_json_without_light_section() { - let json = r#"{ - "colorPalette": "blue", - "autoRefreshEnabled": false, - "autoRefreshMs": 60000, - "includeUnusedModels": false, - "nativeTimeoutMs": 300000 - }"#; - let parsed: Settings = serde_json::from_str(json).unwrap(); - assert!(!parsed.light.write_cache); - } - - #[test] - fn light_settings_round_trip() { - let light = LightSettings { write_cache: true }; - let serialized = serde_json::to_string(&light).unwrap(); - let parsed: LightSettings = serde_json::from_str(&serialized).unwrap(); - assert!(parsed.write_cache); - } - #[test] fn settings_usage_tab_enabled_defaults_to_false() { let json = r#"{ "colorPalette": "blue" }"#; diff --git a/crates/tokscale-cli/src/tui/ui/agents.rs b/crates/tokscale-cli/src/tui/ui/agents.rs index 6123c344a..2fdfa297c 100644 --- a/crates/tokscale-cli/src/tui/ui/agents.rs +++ b/crates/tokscale-cli/src/tui/ui/agents.rs @@ -371,7 +371,8 @@ mod tests { TuiConfig { theme: Some("blue".to_string()), refresh: 0, - sessions_path: None, + no_refresh: false, + home_dir: None, clients: None, since: None, until: None, diff --git a/crates/tokscale-cli/src/tui/ui/daily.rs b/crates/tokscale-cli/src/tui/ui/daily.rs index 3c6e55857..96b4e2928 100644 --- a/crates/tokscale-cli/src/tui/ui/daily.rs +++ b/crates/tokscale-cli/src/tui/ui/daily.rs @@ -964,7 +964,8 @@ mod tests { let config = TuiConfig { theme: Some("blue".to_string()), refresh: 0, - sessions_path: None, + no_refresh: false, + home_dir: None, clients: None, since: None, until: None, diff --git a/crates/tokscale-cli/src/tui/ui/dialog/source_picker.rs b/crates/tokscale-cli/src/tui/ui/dialog/source_picker.rs index 566971714..4f8922d6d 100644 --- a/crates/tokscale-cli/src/tui/ui/dialog/source_picker.rs +++ b/crates/tokscale-cli/src/tui/ui/dialog/source_picker.rs @@ -383,7 +383,7 @@ mod tests { let expected = ClientId::iter().collect::>(); assert_eq!(dialog.sources, expected); - assert!(dialog.sources.contains(&ClientId::Cursor)); + assert!(dialog.sources.contains(&ClientId::Claude)); } #[test] diff --git a/crates/tokscale-cli/src/tui/ui/footer.rs b/crates/tokscale-cli/src/tui/ui/footer.rs index 359034011..3cbff97d4 100644 --- a/crates/tokscale-cli/src/tui/ui/footer.rs +++ b/crates/tokscale-cli/src/tui/ui/footer.rs @@ -521,7 +521,8 @@ mod tests { let config = TuiConfig { theme: Some("blue".to_string()), refresh: 0, - sessions_path: None, + no_refresh: false, + home_dir: None, clients: None, since: None, until: None, diff --git a/crates/tokscale-cli/src/tui/ui/header.rs b/crates/tokscale-cli/src/tui/ui/header.rs index 6bcf4b721..dec97f78b 100644 --- a/crates/tokscale-cli/src/tui/ui/header.rs +++ b/crates/tokscale-cli/src/tui/ui/header.rs @@ -226,7 +226,8 @@ mod tests { let config = TuiConfig { theme: Some("blue".to_string()), refresh: 0, - sessions_path: None, + no_refresh: false, + home_dir: None, clients: None, since: None, until: None, diff --git a/crates/tokscale-cli/src/tui/ui/hourly.rs b/crates/tokscale-cli/src/tui/ui/hourly.rs index 5e402f468..0f22f45d7 100644 --- a/crates/tokscale-cli/src/tui/ui/hourly.rs +++ b/crates/tokscale-cli/src/tui/ui/hourly.rs @@ -463,7 +463,8 @@ mod tests { let config = TuiConfig { theme: Some("blue".to_string()), refresh: 0, - sessions_path: None, + no_refresh: false, + home_dir: None, clients: None, since: None, until: None, diff --git a/crates/tokscale-cli/src/tui/ui/usage.rs b/crates/tokscale-cli/src/tui/ui/usage.rs index d0659a6d7..59c49835b 100644 --- a/crates/tokscale-cli/src/tui/ui/usage.rs +++ b/crates/tokscale-cli/src/tui/ui/usage.rs @@ -259,7 +259,8 @@ mod tests { let config = TuiConfig { theme: Some("blue".to_string()), refresh: 0, - sessions_path: None, + no_refresh: false, + home_dir: None, clients: None, since: None, until: None, diff --git a/crates/tokscale-cli/src/tui/ui/widgets.rs b/crates/tokscale-cli/src/tui/ui/widgets.rs index 6cad80cde..643895856 100644 --- a/crates/tokscale-cli/src/tui/ui/widgets.rs +++ b/crates/tokscale-cli/src/tui/ui/widgets.rs @@ -197,7 +197,6 @@ fn get_single_provider_display_name(provider: &str) -> String { "anthropic" => "Anthropic".to_string(), "openai" => "OpenAI".to_string(), "google" => "Google".to_string(), - "cursor" => "Cursor".to_string(), "deepseek" => "DeepSeek".to_string(), "zai" => "Z.AI".to_string(), "xiaomi" => "XiaoMi".to_string(), diff --git a/crates/tokscale-cli/tests/cli_tests.rs b/crates/tokscale-cli/tests/cli_tests.rs index cdeb4082c..d05a3d7db 100644 --- a/crates/tokscale-cli/tests/cli_tests.rs +++ b/crates/tokscale-cli/tests/cli_tests.rs @@ -205,6 +205,8 @@ fn headless_capture_command(fake_bin: &Path, output_path: &Path, mode: &str) -> output_path.to_str().unwrap(), "--no-auto-flags", "codex", + "--", + "codex", ]); cmd @@ -247,6 +249,28 @@ fn headless_capture_fast_nonzero_preserves_exit_code() { assert_eq!(fs::read_to_string(output_path).unwrap(), "captured fail"); } +#[test] +fn headless_rejects_invalid_native_timeout_before_starting_child() { + for value in ["bogus", "1"] { + let fake_bin = create_fake_codex_bin(); + let output_path = fake_bin + .path() + .join(format!("invalid-timeout-{value}.jsonl")); + + headless_capture_command(fake_bin.path(), &output_path, "success") + .env("TOKSCALE_NATIVE_TIMEOUT_MS", value) + .assert() + .code(2) + .stdout(predicate::str::is_empty()) + .stderr(predicate::str::contains("TOKSCALE_NATIVE_TIMEOUT_MS")); + + assert!( + !output_path.exists(), + "invalid execution environment must fail before creating output" + ); + } +} + #[test] fn headless_capture_slow_command_times_out() { let fake_bin = create_fake_codex_bin(); @@ -807,34 +831,6 @@ fn write_codex_token_session(dir: &Path, name: &str, model: &str, input: i64, ou .unwrap(); } -fn write_cursor_usage_cache(base: &Path) { - let cache_dir = base.join(".config/tokscale/cursor-cache"); - fs::create_dir_all(&cache_dir).unwrap(); - fs::write(cache_dir.join("usage.csv"), "Date,Model\n").unwrap(); -} - -fn write_cursor_credentials(base: &Path) { - let config_dir = base.join(".config/tokscale"); - fs::create_dir_all(&config_dir).unwrap(); - fs::write( - config_dir.join("cursor-credentials.json"), - serde_json::json!({ - "version": 1, - "activeAccountId": "active-account", - "accounts": { - "active-account": { - "sessionToken": "test-session-token", - "userId": "active-account", - "createdAt": "2026-01-01T00:00:00Z", - "label": "work" - } - } - }) - .to_string(), - ) - .unwrap(); -} - // ── Existing tests ───────────────────────────────────────────────────────── #[test] @@ -894,7 +890,9 @@ fn test_pricing_command_help() { .arg("--help") .assert() .success() - .stdout(predicate::str::contains("Show pricing for a model")); + .stdout(predicate::str::contains("Query model pricing")) + .stdout(predicate::str::contains("lookup")) + .stdout(predicate::str::contains("overrides")); } #[test] @@ -912,7 +910,7 @@ fn test_cache_prune_reports_empty_cache_stats() { let config_dir = TempDir::new().unwrap(); let mut cmd = cargo_bin_cmd!("tokscale"); cmd.env("TOKSCALE_CONFIG_DIR", config_dir.path()) - .args(["--no-spinner", "cache", "prune"]) + .args(["cache", "prune"]) .assert() .success() .stdout(predicate::str::contains( @@ -931,7 +929,7 @@ fn test_cache_prune_surfaces_unknown_shard_magic() { let mut cmd = cargo_bin_cmd!("tokscale"); cmd.env("TOKSCALE_CONFIG_DIR", config_dir.path()) - .args(["--no-spinner", "cache", "prune"]) + .args(["cache", "prune"]) .assert() .failure() .stderr(predicate::str::contains("has unrecognized magic")) @@ -939,15 +937,13 @@ fn test_cache_prune_surfaces_unknown_shard_magic() { } #[test] -fn test_codex_command_help() { - let mut cmd = cargo_bin_cmd!("tokscale"); - cmd.arg("codex") - .arg("--help") - .assert() - .success() - .stdout(predicate::str::contains( - "Codex account integration commands", - )); +fn test_removed_integration_namespaces_are_not_registered() { + for command in ["codex", "cursor", "trae"] { + cargo_bin_cmd!("tokscale") + .args([command, "--help"]) + .assert() + .code(2); + } } #[test] @@ -967,7 +963,86 @@ fn test_tui_command_help() { .arg("--help") .assert() .success() - .stdout(predicate::str::contains("Launch interactive TUI")); + .stdout(predicate::str::contains( + "Launch the interactive terminal interface", + )); +} + +#[test] +fn test_wrapped_ranking_rejects_ambiguous_and_irrelevant_options() { + cargo_bin_cmd!("tokscale") + .args(["wrapped", "--help"]) + .assert() + .success() + .stdout(predicate::str::contains("--ranking ")) + .stdout(predicate::str::contains("--agents").not()) + .stdout(predicate::str::contains("--clients").not()); + + cargo_bin_cmd!("tokscale") + .args(["wrapped", "--agents", "--clients", "--no-spinner"]) + .assert() + .code(2); + + cargo_bin_cmd!("tokscale") + .args([ + "wrapped", + "--ranking", + "agents", + "--client", + "claude", + "--no-spinner", + ]) + .assert() + .code(2) + .stderr(predicate::str::contains( + "--ranking agents requires `opencode`", + )); + + cargo_bin_cmd!("tokscale") + .args([ + "wrapped", + "--ranking", + "clients", + "--disable-pinned", + "--no-spinner", + ]) + .assert() + .code(2) + .stderr(predicate::str::contains("--disable-pinned does not apply")); +} + +#[test] +fn test_help_exposes_only_leaf_owned_options() { + cargo_bin_cmd!("tokscale") + .arg("--help") + .assert() + .success() + .stdout(predicate::str::contains("--json").not()) + .stdout(predicate::str::contains("--client").not()) + .stdout(predicate::str::contains("--group-by").not()); + + cargo_bin_cmd!("tokscale") + .args(["models", "--help"]) + .assert() + .success() + .stdout(predicate::str::contains("--json")) + .stdout(predicate::str::contains("--client")) + .stdout(predicate::str::contains("--group-by")); + + cargo_bin_cmd!("tokscale") + .args(["monthly", "--help"]) + .assert() + .success() + .stdout(predicate::str::contains("--json")) + .stdout(predicate::str::contains("--group-by").not()); + + cargo_bin_cmd!("tokscale") + .args(["tui", "--help"]) + .assert() + .success() + .stdout(predicate::str::contains("--tab")) + .stdout(predicate::str::contains("--theme")) + .stdout(predicate::str::contains("--json").not()); } #[test] @@ -1007,18 +1082,6 @@ fn test_invalid_subcommand() { cmd.arg("models").arg("invalid-flag").assert().failure(); } -#[test] -fn test_codex_accounts_empty_json() { - let tmp = TempDir::new().expect("failed to create temp home"); - let mut cmd = cargo_bin_cmd!("tokscale"); - cmd.env("HOME", tmp.path()) - .env_remove("CODEX_HOME") - .args(["codex", "accounts", "--json"]) - .assert() - .success() - .stdout(predicate::str::contains(r#""accounts": []"#)); -} - #[test] fn test_pricing_command_missing_model() { let mut cmd = cargo_bin_cmd!("tokscale"); @@ -1041,18 +1104,28 @@ fn test_headless_command_invalid_client() { .failure(); } +#[test] +fn test_headless_requires_explicit_child_command_separator() { + cargo_bin_cmd!("tokscale") + .args(["headless", "codex", "codex", "exec"]) + .assert() + .code(2) + .stderr(predicate::str::contains( + "separate Tokscale options from the child command with `--`", + )); +} + #[test] fn test_models_with_invalid_date_format() { let tmp = create_empty_fixture_dir(); cmd_with_home(tmp.path()) .arg("models") - .arg("--light") .args(["--client", "opencode"]) .arg("--no-spinner") .arg("--since") .arg("invalid-date") .assert() - .success(); + .code(2); } #[test] @@ -1060,29 +1133,101 @@ fn test_models_with_invalid_year() { let tmp = create_empty_fixture_dir(); cmd_with_home(tmp.path()) .arg("models") - .arg("--light") .args(["--client", "opencode"]) .arg("--no-spinner") .arg("--year") .arg("not-a-year") .assert() - .success(); + .code(2); +} + +#[test] +fn test_local_scope_rejects_nonexistent_home() { + cargo_bin_cmd!("tokscale") + .args([ + "models", + "--home", + "/definitely/not/a/tokscale/home", + "--no-spinner", + ]) + .assert() + .code(2) + .stderr(predicate::str::contains( + "--home must be an existing directory", + )); +} + +#[test] +fn test_date_presets_are_mutually_exclusive() { + cargo_bin_cmd!("tokscale") + .args(["models", "--week", "--month", "--no-spinner"]) + .assert() + .code(2) + .stderr(predicate::str::contains("cannot be used with")); +} + +#[test] +fn test_custom_date_range_must_be_ordered() { + cargo_bin_cmd!("tokscale") + .args([ + "models", + "--since", + "2026-07-15", + "--until", + "2026-07-14", + "--no-spinner", + ]) + .assert() + .code(2) + .stderr(predicate::str::contains("must not be later")); +} + +#[test] +fn json_report_suppresses_spinner_without_explicit_no_spinner() { + let tmp = create_empty_fixture_dir(); + let output = cmd_with_home(tmp.path()) + .args(["models", "--json", "--client", "opencode"]) + .output() + .unwrap(); + + assert!( + output.status.success(), + "stderr: {}", + String::from_utf8_lossy(&output.stderr) + ); + serde_json::from_slice::(&output.stdout).expect("valid JSON report"); + let stderr = String::from_utf8_lossy(&output.stderr); + assert!(!stderr.contains("Scanning session data..."), "{stderr}"); + assert!(!stderr.contains("\x1b[?25l"), "{stderr}"); } #[test] -fn test_global_theme_flag() { +fn test_theme_flag_is_owned_by_tui() { let mut cmd = cargo_bin_cmd!("tokscale"); - cmd.arg("--theme") - .arg("blue") - .arg("--help") + cmd.args(["tui", "--theme", "blue", "--help"]) .assert() .success(); + + let mut root = cargo_bin_cmd!("tokscale"); + root.args(["--theme", "blue"]).assert().code(2); } #[test] -fn test_global_debug_flag() { +fn test_debug_flag_is_owned_by_tui() { let mut cmd = cargo_bin_cmd!("tokscale"); - cmd.arg("--debug").arg("--help").assert().success(); + cmd.args(["tui", "--debug", "--help"]).assert().success(); + + let mut root = cargo_bin_cmd!("tokscale"); + root.arg("--debug").assert().code(2); +} + +#[test] +fn test_tui_refresh_modes_are_mutually_exclusive() { + cargo_bin_cmd!("tokscale") + .args(["tui", "--refresh", "30", "--no-refresh"]) + .assert() + .code(2) + .stderr(predicate::str::contains("cannot be used with")); } #[test] @@ -1233,9 +1378,9 @@ fn test_models_home_override_ignores_conflicting_xdg_env() { ); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - assert_eq!(json["totalMessages"].as_i64().unwrap(), 3); - assert_eq!(json["totalInput"].as_i64().unwrap(), 2400); - assert_eq!(json["totalOutput"].as_i64().unwrap(), 1000); + assert_eq!(json["data"]["totalMessages"].as_i64().unwrap(), 3); + assert_eq!(json["data"]["totalInput"].as_i64().unwrap(), 2400); + assert_eq!(json["data"]["totalOutput"].as_i64().unwrap(), 1000); assert!(!String::from_utf8_lossy(&output.stdout).contains("gemini-2.5-pro")); } @@ -1264,7 +1409,7 @@ fn test_monthly_home_override_ignores_conflicting_xdg_env() { ); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - let entries = json["entries"].as_array().unwrap(); + let entries = json["data"]["entries"].as_array().unwrap(); assert_eq!(entries.len(), 2); assert!(entries.iter().any(|entry| entry["month"] == "2024-06")); assert!(entries.iter().any(|entry| entry["month"] == "2025-01")); @@ -1295,7 +1440,7 @@ fn test_graph_home_override_ignores_conflicting_xdg_env() { ); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - let contributions = json["contributions"].as_array().unwrap(); + let contributions = json["data"]["contributions"].as_array().unwrap(); assert_eq!(contributions.len(), 2); assert!(!String::from_utf8_lossy(&output.stdout).contains("gemini-2.5-pro")); } @@ -1326,23 +1471,23 @@ fn test_models_home_override_ignores_conflicting_codex_home_env() { ); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - assert_eq!(json["totalMessages"].as_i64().unwrap(), 1); - assert_eq!(json["totalInput"].as_i64().unwrap(), 100); - assert_eq!(json["totalOutput"].as_i64().unwrap(), 30); - assert_eq!(json["totalCacheRead"].as_i64().unwrap(), 20); + assert_eq!(json["data"]["totalMessages"].as_i64().unwrap(), 1); + assert_eq!(json["data"]["totalInput"].as_i64().unwrap(), 100); + assert_eq!(json["data"]["totalOutput"].as_i64().unwrap(), 30); + assert_eq!(json["data"]["totalCacheRead"].as_i64().unwrap(), 20); assert!(!String::from_utf8_lossy(&output.stdout).contains("\"gpt-5\"")); } #[test] -fn test_tui_rejects_home_override() { +fn test_tui_accepts_home_but_requires_an_interactive_terminal() { let tmp = TempDir::new().unwrap(); cargo_bin_cmd!("tokscale") - .args(["--home", tmp.path().to_str().unwrap(), "tui"]) + .args(["tui", "--home", tmp.path().to_str().unwrap()]) .assert() - .failure() + .code(2) .stderr(predicate::str::contains( - "--home is currently supported for local report commands only", + "TUI requires an interactive terminal", )); } @@ -1361,9 +1506,9 @@ fn test_clients_home_override_uses_explicit_home_for_json() { let output = cmd_with_conflicting_env(conflicting_home.path()) .env("CODEX_HOME", conflicting_home.path().join(".codex")) .args([ + "clients", "--home", real_home.path().to_str().unwrap(), - "clients", "--json", ]) .output() @@ -1376,7 +1521,7 @@ fn test_clients_home_override_uses_explicit_home_for_json() { ); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - let codex = json["clients"] + let codex = json["data"]["clients"] .as_array() .unwrap() .iter() @@ -1399,9 +1544,9 @@ fn test_clients_home_override_ignores_copilot_exporter_env() { let output = cmd_with_conflicting_env(conflicting_home.path()) .env("COPILOT_OTEL_FILE_EXPORTER_PATH", &exporter_file) .args([ + "clients", "--home", real_home.path().to_str().unwrap(), - "clients", "--json", ]) .output() @@ -1414,7 +1559,7 @@ fn test_clients_home_override_ignores_copilot_exporter_env() { ); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - let copilot = json["clients"] + let copilot = json["data"]["clients"] .as_array() .unwrap() .iter() @@ -1460,7 +1605,7 @@ fn test_models_with_no_matching_date() { .unwrap(); assert!(output.status.success()); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - let entries = json["entries"].as_array().unwrap(); + let entries = json["data"]["entries"].as_array().unwrap(); assert!( entries.is_empty(), "No entries expected for future date range" @@ -1483,7 +1628,7 @@ fn test_graph_single_day_filter_uses_local_timezone_boundaries() { ); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - let contributions = json["contributions"].as_array().unwrap(); + let contributions = json["data"]["contributions"].as_array().unwrap(); assert_eq!( contributions.len(), 1, @@ -1504,7 +1649,7 @@ fn test_graph_with_year_filter() { .unwrap(); assert!(output.status.success()); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - let contributions = json["contributions"].as_array().unwrap(); + let contributions = json["data"]["contributions"].as_array().unwrap(); for c in contributions { let date = c["date"].as_str().unwrap(); assert!( @@ -1526,7 +1671,7 @@ fn test_models_with_client_filter_opencode() { .unwrap(); assert!(output.status.success()); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - let entries = json["entries"].as_array().unwrap(); + let entries = json["data"]["entries"].as_array().unwrap(); for entry in entries { assert_eq!(entry["client"].as_str().unwrap(), "opencode"); } @@ -1549,221 +1694,19 @@ fn test_models_with_client_filter_multiple() { .success(); } -fn assert_cursor_setup_warning(json: &serde_json::Value) { - let warnings = json["warnings"] - .as_array() - .expect("explicit Cursor report should expose setup warnings"); - assert!( - warnings - .iter() - .any(|warning| warning.as_str().is_some_and(|text| text - .contains("tokscale cursor login") - && text.contains("tokscale cursor sync --json") - && text.contains("cursor-cache/usage*.csv") - && text.contains("Tokscale does not parse local `~/.cursor`"))), - "warnings did not explain Cursor setup: {warnings:?}" - ); -} - -#[test] -fn test_models_cursor_explicit_missing_cache_reports_setup_warning_json() { - let tmp = create_empty_fixture_dir(); - let output = cmd_with_home(tmp.path()) - .args(["models", "--json", "--client", "cursor", "--no-spinner"]) - .output() - .unwrap(); - - assert!(output.status.success()); - let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - assert_cursor_setup_warning(&json); -} - -#[test] -fn test_models_cursor_explicit_local_cursor_state_still_reports_setup_warning_json() { - let tmp = create_empty_fixture_dir(); - fs::create_dir_all( - tmp.path() - .join(".cursor/projects/demo/agent-transcripts/session"), - ) - .unwrap(); - fs::write( - tmp.path() - .join(".cursor/projects/demo/agent-transcripts/session/session.jsonl"), - r#"{"role":"user","content":"hello"}"#, - ) - .unwrap(); - - let output = cmd_with_home(tmp.path()) - .args(["models", "--json", "--client", "cursor", "--no-spinner"]) - .output() - .unwrap(); - - assert!(output.status.success()); - let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - assert_cursor_setup_warning(&json); -} - -#[test] -fn test_monthly_cursor_explicit_missing_cache_reports_setup_warning_json() { - let tmp = create_empty_fixture_dir(); - let output = cmd_with_home(tmp.path()) - .args(["monthly", "--json", "--client", "cursor", "--no-spinner"]) - .output() - .unwrap(); - - assert!(output.status.success()); - let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - assert_cursor_setup_warning(&json); -} - #[test] -fn test_hourly_cursor_explicit_missing_cache_reports_setup_warning_json() { +fn test_reports_reject_removed_client_ids() { let tmp = create_empty_fixture_dir(); - let output = cmd_with_home(tmp.path()) - .args(["hourly", "--json", "--client", "cursor", "--no-spinner"]) - .output() - .unwrap(); - - assert!(output.status.success()); - let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - assert_cursor_setup_warning(&json); -} - -#[test] -fn test_models_cursor_explicit_home_override_reports_fixture_cache_path() { - let tmp = create_empty_fixture_dir(); - let output = cmd_with_home(tmp.path()) - .args([ - "--home", - tmp.path().to_str().unwrap(), - "models", - "--json", - "--client", - "cursor", - "--no-spinner", - ]) - .output() - .unwrap(); - - assert!(output.status.success()); - let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - let warnings = json["warnings"] - .as_array() - .expect("explicit Cursor --home report should expose setup warnings"); - assert!( - warnings - .iter() - .any(|warning| warning.as_str().is_some_and(|text| text - .contains(tmp.path().to_str().unwrap()) - && text.contains("tokscale cursor login") - && text.contains("tokscale cursor sync --json") - && text.contains("cursor-cache/usage*.csv"))), - "warnings did not explain Cursor --home setup: {warnings:?}" - ); -} - -#[test] -fn test_models_cursor_explicit_missing_cache_reports_setup_warning_text() { - let tmp = create_empty_fixture_dir(); - cmd_with_home(tmp.path()) - .args(["models", "--client", "cursor", "--no-spinner"]) - .assert() - .success() - .stderr(predicate::str::contains("Cursor usage requires")) - .stderr(predicate::str::contains("tokscale cursor login")) - .stderr(predicate::str::contains("tokscale cursor sync --json")) - .stderr(predicate::str::contains( - "Tokscale does not parse local `~/.cursor`", - )); -} - -#[test] -fn test_models_default_missing_cursor_cache_does_not_emit_setup_warning_json() { - let tmp = create_empty_fixture_dir(); - let output = cmd_with_home(tmp.path()) - .args(["models", "--json", "--no-spinner"]) - .output() - .unwrap(); - - assert!(output.status.success()); - let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - assert!( - json.get("warnings") - .and_then(serde_json::Value::as_array) - .is_none_or(Vec::is_empty), - "default all-client report should not warn about unrequested Cursor setup" - ); -} - -#[test] -fn test_models_cursor_explicit_existing_cache_suppresses_setup_warning_json() { - let tmp = create_empty_fixture_dir(); - write_cursor_usage_cache(tmp.path()); - - let output = cmd_with_home(tmp.path()) - .args(["models", "--json", "--client", "cursor", "--no-spinner"]) - .output() - .unwrap(); - - assert!(output.status.success()); - let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - assert!( - json.get("warnings") - .and_then(serde_json::Value::as_array) - .is_none_or(Vec::is_empty), - "existing Cursor cache should suppress setup warnings" - ); -} - -#[test] -fn test_models_cursor_logged_in_missing_cache_suggests_sync_only_json() { - let tmp = create_empty_fixture_dir(); - write_cursor_credentials(tmp.path()); - - let output = cmd_with_home(tmp.path()) - .env("HTTPS_PROXY", "http://127.0.0.1:9") - .env("HTTP_PROXY", "http://127.0.0.1:9") - .env("ALL_PROXY", "http://127.0.0.1:9") - .args(["models", "--json", "--client", "cursor", "--no-spinner"]) - .output() - .unwrap(); - - assert!( - output.status.success(), - "stderr: {}", - String::from_utf8_lossy(&output.stderr) - ); - let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - let warnings = json["warnings"].as_array().unwrap(); - let warning = warnings[0].as_str().unwrap(); - assert!(warning.contains("tokscale cursor sync --json")); - assert!( - !warning.contains("tokscale cursor login"), - "logged-in users with no cache should be told to sync, not log in again: {warning}" - ); -} - -#[test] -fn test_time_metrics_cursor_explicit_missing_cache_reports_setup_warning_json() { - let tmp = create_empty_fixture_dir(); - let output = cmd_with_home(tmp.path()) - .args([ - "time-metrics", - "--json", - "--client", - "cursor", - "--no-spinner", - ]) - .output() - .unwrap(); - - assert!( - output.status.success(), - "stderr: {}", - String::from_utf8_lossy(&output.stderr) - ); - let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - assert_cursor_setup_warning(&json); + for client in ["cursor", "trae"] { + let output = cmd_with_home(tmp.path()) + .args(["models", "--client", client, "--no-spinner"]) + .output() + .unwrap(); + assert_eq!(output.status.code(), Some(2)); + let stderr = String::from_utf8_lossy(&output.stderr); + assert!(stderr.contains("invalid value"), "stderr: {stderr}"); + assert!(stderr.contains(client), "stderr: {stderr}"); + } } #[test] @@ -1807,65 +1750,28 @@ fn test_time_metrics_reports_degraded_source_health_without_failing() { } #[test] -fn test_time_metrics_text_reports_degraded_source_health_without_failing() { - let tmp = create_empty_fixture_dir(); - let missing_db = tmp.path().join("missing/opencode.db"); - write_settings_json( - tmp.path(), - &format!( - r#"{{"scanner":{{"opencodeDbPaths":[{}]}}}}"#, - serde_json::to_string(missing_db.to_str().unwrap()).unwrap() - ), - ); - - cmd_with_home(tmp.path()) - .args(["time-metrics", "--client", "opencode", "--no-spinner"]) - .assert() - .success() - .stdout(predicate::str::contains("Session Time Metrics")) - .stderr(predicate::str::contains( - "Data health: 0 degraded source(s), 0 rejected record(s), 0 partial source(s), 1 failed source(s)", - )); -} - -#[test] -fn test_graph_cursor_explicit_missing_cache_reports_setup_warning_text() { - let tmp = create_empty_fixture_dir(); - cmd_with_home(tmp.path()) - .args(["graph", "--client", "cursor", "--no-spinner"]) - .assert() - .success() - .stderr(predicate::str::contains("Cursor usage requires")) - .stderr(predicate::str::contains("tokscale cursor login")); -} - -#[test] -fn test_graph_fresh_cursor_cache_skips_auto_sync_warning() { - let tmp = create_empty_fixture_dir(); - write_cursor_credentials(tmp.path()); - write_cursor_usage_cache(tmp.path()); - - let output = cmd_with_home(tmp.path()) - .env("HTTPS_PROXY", "http://127.0.0.1:9") - .env("HTTP_PROXY", "http://127.0.0.1:9") - .env("ALL_PROXY", "http://127.0.0.1:9") - .args(["graph", "--client", "cursor", "--no-spinner"]) - .output() - .unwrap(); - - assert!( - output.status.success(), - "stderr: {}", - String::from_utf8_lossy(&output.stderr) - ); - let stderr = String::from_utf8_lossy(&output.stderr); - assert!( - !stderr.contains("Cursor sync failed") && !stderr.contains("Cursor sync warning"), - "fresh Cursor cache should skip implicit graph sync; stderr: {stderr}" - ); -} - -#[test] +fn test_time_metrics_text_reports_degraded_source_health_without_failing() { + let tmp = create_empty_fixture_dir(); + let missing_db = tmp.path().join("missing/opencode.db"); + write_settings_json( + tmp.path(), + &format!( + r#"{{"scanner":{{"opencodeDbPaths":[{}]}}}}"#, + serde_json::to_string(missing_db.to_str().unwrap()).unwrap() + ), + ); + + cmd_with_home(tmp.path()) + .args(["time-metrics", "--client", "opencode", "--no-spinner"]) + .assert() + .success() + .stdout(predicate::str::contains("Session Time Metrics")) + .stderr(predicate::str::contains( + "Data health: 0 degraded source(s), 0 rejected record(s), 0 partial source(s), 1 failed source(s)", + )); +} + +#[test] fn test_models_with_repeated_client_filter() { let tmp = create_temp_fixture_dir(); cmd_with_home(tmp.path()) @@ -1882,8 +1788,6 @@ fn test_models_with_repeated_client_filter() { "--client", "gemini", "--client", - "cursor", - "--client", "amp", "--client", "droid", @@ -1920,31 +1824,49 @@ fn test_models_json_output() { assert!(output.status.success()); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - assert!(json.get("groupBy").is_some(), "Missing groupBy field"); - assert!(json.get("entries").is_some(), "Missing entries field"); - assert!(json.get("totalInput").is_some(), "Missing totalInput"); - assert!(json.get("totalOutput").is_some(), "Missing totalOutput"); assert!( - json.get("totalCacheRead").is_some(), + json["data"].get("groupBy").is_some(), + "Missing groupBy field" + ); + assert!( + json["data"].get("entries").is_some(), + "Missing entries field" + ); + assert!( + json["data"].get("totalInput").is_some(), + "Missing totalInput" + ); + assert!( + json["data"].get("totalOutput").is_some(), + "Missing totalOutput" + ); + assert!( + json["data"].get("totalCacheRead").is_some(), "Missing totalCacheRead" ); assert!( - json.get("totalCacheWrite").is_some(), + json["data"].get("totalCacheWrite").is_some(), "Missing totalCacheWrite" ); assert!( - json.get("totalReasoning").is_none(), + json["data"].get("totalReasoning").is_none(), "JSON report must fold reasoning into totalOutput" ); - assert!(json.get("totalTokens").is_some(), "Missing totalTokens"); - assert!(json.get("totalMessages").is_some(), "Missing totalMessages"); - assert!(json.get("totalCost").is_some(), "Missing totalCost"); assert!( - json.get("processingTimeMs").is_some(), + json["data"].get("totalTokens").is_some(), + "Missing totalTokens" + ); + assert!( + json["data"].get("totalMessages").is_some(), + "Missing totalMessages" + ); + assert!(json["data"].get("totalCost").is_some(), "Missing totalCost"); + assert!( + json["metadata"].get("processingTimeMs").is_some(), "Missing processingTimeMs" ); - let entries = json["entries"].as_array().unwrap(); + let entries = json["data"]["entries"].as_array().unwrap(); assert!(!entries.is_empty(), "Should have entries from fixture data"); let first = &entries[0]; assert!(first.get("client").is_some()); @@ -1972,6 +1894,56 @@ fn test_models_json_output() { assert!(performance["msPer1KTokens"].as_f64().unwrap() > 0.0); } +#[test] +fn test_every_local_json_command_uses_the_common_envelope() { + let tmp = create_empty_fixture_dir(); + let invocations: &[&[&str]] = &[ + &["models", "--json", "--client", "opencode", "--no-spinner"], + &["monthly", "--json", "--client", "opencode", "--no-spinner"], + &["hourly", "--json", "--client", "opencode", "--no-spinner"], + &[ + "time-metrics", + "--json", + "--client", + "opencode", + "--no-spinner", + ], + &["graph", "--client", "opencode", "--no-spinner"], + &["clients", "--json", "--client", "opencode"], + ]; + + for invocation in invocations { + let output = cmd_with_home(tmp.path()) + .args(*invocation) + .output() + .unwrap(); + assert!( + output.status.success(), + "{} failed: {}", + invocation.join(" "), + String::from_utf8_lossy(&output.stderr) + ); + let document: serde_json::Value = + serde_json::from_slice(&output.stdout).unwrap_or_else(|error| { + panic!("{} returned invalid JSON: {error}", invocation.join(" ")) + }); + let mut keys = document + .as_object() + .expect("report envelope must be an object") + .keys() + .map(String::as_str) + .collect::>(); + keys.sort_unstable(); + assert_eq!( + keys, + vec!["data", "health", "metadata"], + "{} returned a non-standard envelope", + invocation.join(" ") + ); + assert!(document["metadata"]["processingTimeMs"].is_number()); + } +} + #[test] fn test_models_json_offline_without_pricing_cache_still_succeeds() { let tmp = create_temp_fixture_dir_without_pricing_cache(); @@ -1986,11 +1958,11 @@ fn test_models_json_offline_without_pricing_cache_still_succeeds() { ); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - assert_eq!(json["totalInput"].as_i64().unwrap(), 2400); - assert_eq!(json["totalOutput"].as_i64().unwrap(), 1000); - assert_eq!(json["totalMessages"].as_i64().unwrap(), 3); - assert_eq!(json["entries"].as_array().unwrap().len(), 2); - let total_cost = json["totalCost"].as_f64().unwrap(); + assert_eq!(json["data"]["totalInput"].as_i64().unwrap(), 2400); + assert_eq!(json["data"]["totalOutput"].as_i64().unwrap(), 1000); + assert_eq!(json["data"]["totalMessages"].as_i64().unwrap(), 3); + assert_eq!(json["data"]["entries"].as_array().unwrap().len(), 2); + let total_cost = json["data"]["totalCost"].as_f64().unwrap(); assert_eq!(total_cost, 0.0); } @@ -2008,11 +1980,11 @@ fn test_monthly_json_offline_without_pricing_cache_still_succeeds() { ); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - let entries = json["entries"].as_array().unwrap(); + let entries = json["data"]["entries"].as_array().unwrap(); assert_eq!(entries.len(), 2); assert_eq!(entries[0]["month"].as_str().unwrap(), "2024-06"); assert_eq!(entries[1]["month"].as_str().unwrap(), "2025-01"); - let total_cost = json["totalCost"].as_f64().unwrap(); + let total_cost = json["data"]["totalCost"].as_f64().unwrap(); assert_eq!(total_cost, 0.0); } @@ -2030,10 +2002,13 @@ fn test_graph_offline_without_pricing_cache_still_succeeds() { ); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - assert_eq!(json["summary"]["totalTokens"].as_i64().unwrap(), 3950); - assert_eq!(json["summary"]["activeDays"].as_i64().unwrap(), 2); - assert_eq!(json["contributions"].as_array().unwrap().len(), 2); - let total_cost = json["summary"]["totalCost"].as_f64().unwrap(); + assert_eq!( + json["data"]["summary"]["totalTokens"].as_i64().unwrap(), + 3950 + ); + assert_eq!(json["data"]["summary"]["activeDays"].as_i64().unwrap(), 2); + assert_eq!(json["data"]["contributions"].as_array().unwrap().len(), 2); + let total_cost = json["data"]["summary"]["totalCost"].as_f64().unwrap(); assert_eq!(total_cost, 0.0); } @@ -2051,7 +2026,7 @@ fn test_hourly_json_offline_without_pricing_cache_still_succeeds() { ); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - let entries = json["entries"].as_array().unwrap(); + let entries = json["data"]["entries"].as_array().unwrap(); assert_eq!(entries.len(), 3); for entry in entries { let hour = entry["hour"].as_str().unwrap(); @@ -2083,7 +2058,7 @@ fn test_hourly_json_offline_without_pricing_cache_still_succeeds() { .sum::(), 1000 ); - let total_cost = json["totalCost"].as_f64().unwrap(); + let total_cost = json["data"]["totalCost"].as_f64().unwrap(); assert_eq!(total_cost, 0.0); } @@ -2129,7 +2104,7 @@ fn test_models_json_offline_uses_stale_pricing_cache_when_available() { ); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - let total_cost = json["totalCost"].as_f64().unwrap(); + let total_cost = json["data"]["totalCost"].as_f64().unwrap(); assert!( (total_cost - 0.0209).abs() < 1e-9, "unexpected totalCost: {total_cost}" @@ -2154,7 +2129,7 @@ fn test_monthly_json_offline_uses_stale_pricing_cache_when_available() { ); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - let total_cost = json["totalCost"].as_f64().unwrap(); + let total_cost = json["data"]["totalCost"].as_f64().unwrap(); assert!( (total_cost - 0.0209).abs() < 1e-9, "unexpected totalCost: {total_cost}" @@ -2179,7 +2154,7 @@ fn test_graph_offline_uses_stale_pricing_cache_when_available() { ); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - let total_cost = json["summary"]["totalCost"].as_f64().unwrap(); + let total_cost = json["data"]["summary"]["totalCost"].as_f64().unwrap(); assert!( (total_cost - 0.0209).abs() < 1e-9, "unexpected totalCost: {total_cost}" @@ -2204,7 +2179,7 @@ fn test_hourly_json_offline_uses_stale_pricing_cache_when_available() { ); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - let entries = json["entries"].as_array().unwrap(); + let entries = json["data"]["entries"].as_array().unwrap(); assert_eq!(entries.len(), 3); assert_eq!( entries @@ -2220,7 +2195,7 @@ fn test_hourly_json_offline_uses_stale_pricing_cache_when_available() { .sum::(), 1000 ); - let total_cost = json["totalCost"].as_f64().unwrap(); + let total_cost = json["data"]["totalCost"].as_f64().unwrap(); assert!( (total_cost - 0.0209).abs() < 1e-9, "unexpected totalCost: {total_cost}" @@ -2238,11 +2213,11 @@ fn test_models_json_total_consistency() { assert!(output.status.success()); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - let entries = json["entries"].as_array().unwrap(); + let entries = json["data"]["entries"].as_array().unwrap(); let sum_input: i64 = entries.iter().map(|e| e["input"].as_i64().unwrap()).sum(); let sum_output: i64 = entries.iter().map(|e| e["output"].as_i64().unwrap()).sum(); - let total_input = json["totalInput"].as_i64().unwrap(); - let total_output = json["totalOutput"].as_i64().unwrap(); + let total_input = json["data"]["totalInput"].as_i64().unwrap(); + let total_output = json["data"]["totalOutput"].as_i64().unwrap(); assert_eq!(json["health"]["complete"], true); @@ -2266,15 +2241,21 @@ fn test_monthly_json_output() { assert!(output.status.success()); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - assert!(json.get("entries").is_some(), "Missing entries field"); - assert!(json.get("totalCost").is_some(), "Missing totalCost field"); + assert!( + json["data"].get("entries").is_some(), + "Missing entries field" + ); + assert!( + json["data"].get("totalCost").is_some(), + "Missing totalCost field" + ); assert_eq!(json["health"]["complete"], true); assert!( - json.get("processingTimeMs").is_some(), + json["metadata"].get("processingTimeMs").is_some(), "Missing processingTimeMs" ); - let entries = json["entries"].as_array().unwrap(); + let entries = json["data"]["entries"].as_array().unwrap(); assert!( !entries.is_empty(), "Should have monthly entries from fixture data" @@ -2343,9 +2324,9 @@ fn test_hourly_home_override_uses_explicit_home_scanner_settings() { ); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - assert_eq!(json["entries"].as_array().unwrap().len(), 1); - assert_eq!(json["entries"][0]["input"].as_i64().unwrap(), 210); - assert_eq!(json["entries"][0]["output"].as_i64().unwrap(), 40); + assert_eq!(json["data"]["entries"].as_array().unwrap().len(), 1); + assert_eq!(json["data"]["entries"][0]["input"].as_i64().unwrap(), 210); + assert_eq!(json["data"]["entries"][0]["output"].as_i64().unwrap(), 40); assert!(!String::from_utf8_lossy(&output.stdout).contains("gpt-5")); } @@ -2359,7 +2340,7 @@ fn test_monthly_json_with_client_filter() { .unwrap(); assert!(output.status.success()); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - let entries = json["entries"].as_array().unwrap(); + let entries = json["data"]["entries"].as_array().unwrap(); for entry in entries { let month = entry["month"].as_str().unwrap(); assert!( @@ -2380,15 +2361,42 @@ fn test_graph_json_output() { assert!(output.status.success()); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - assert!(json.get("meta").is_some(), "Missing meta field"); - assert!(json.get("summary").is_some(), "Missing summary field"); - assert!(json.get("years").is_some(), "Missing years field"); + assert!(json["data"].get("meta").is_some(), "Missing meta field"); + assert!( + json["data"].get("summary").is_some(), + "Missing summary field" + ); + assert!(json["data"].get("years").is_some(), "Missing years field"); assert!( - json.get("contributions").is_some(), + json["data"].get("contributions").is_some(), "Missing contributions field" ); } +#[cfg(unix)] +#[test] +fn graph_output_preserves_non_utf8_path_bytes() { + use std::ffi::OsString; + use std::os::unix::ffi::OsStringExt; + + let tmp = create_empty_fixture_dir(); + let output_path = tmp + .path() + .join(OsString::from_vec(b"graph-\xff.json".to_vec())); + + cmd_with_home(tmp.path()) + .args(["graph", "--client", "opencode", "--output"]) + .arg(&output_path) + .arg("--no-spinner") + .assert() + .success(); + + assert!( + output_path.is_file(), + "Graph must write to the exact OS path supplied by the user" + ); +} + #[test] fn test_graph_json_has_meta() { let tmp = create_temp_fixture_dir(); @@ -2398,7 +2406,7 @@ fn test_graph_json_has_meta() { .unwrap(); assert!(output.status.success()); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - let meta = &json["meta"]; + let meta = &json["data"]["meta"]; assert!( meta.get("generatedAt").is_some(), "Missing meta.generatedAt" @@ -2416,7 +2424,7 @@ fn test_graph_json_has_summary() { .unwrap(); assert!(output.status.success()); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - let summary = &json["summary"]; + let summary = &json["data"]["summary"]; assert!( summary.get("totalTokens").is_some(), "Missing summary.totalTokens" @@ -2448,7 +2456,7 @@ fn test_models_group_by_default() { .unwrap(); assert!(output.status.success()); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - assert_eq!(json["groupBy"].as_str().unwrap(), "client,model"); + assert_eq!(json["data"]["groupBy"].as_str().unwrap(), "client,model"); } #[test] @@ -2478,11 +2486,11 @@ fn test_models_reports_project_reasoning_into_output() { "command failed: {json_output:?}" ); let json: serde_json::Value = serde_json::from_slice(&json_output.stdout).unwrap(); - assert_eq!(json["entries"][0]["output"], 50); - assert!(json["entries"][0].get("reasoning").is_none()); - assert_eq!(json["totalOutput"], 50); + assert_eq!(json["data"]["entries"][0]["output"], 50); + assert!(json["data"]["entries"][0].get("reasoning").is_none()); + assert_eq!(json["data"]["totalOutput"], 50); assert!(json.get("totalReasoning").is_none()); - assert_eq!(json["totalTokens"], 165); + assert_eq!(json["data"]["totalTokens"], 165); let table_output = cmd_with_home(base) .args(["models", "--client", "omp", "--no-spinner"]) @@ -2532,8 +2540,8 @@ fn test_monthly_reports_project_reasoning_into_output() { "command failed: {json_output:?}" ); let json: serde_json::Value = serde_json::from_slice(&json_output.stdout).unwrap(); - assert_eq!(json["entries"][0]["output"], 50); - assert!(json["entries"][0].get("reasoning").is_none()); + assert_eq!(json["data"]["entries"][0]["output"], 50); + assert!(json["data"]["entries"][0].get("reasoning").is_none()); let table_output = cmd_with_home(base) .args(["monthly", "--client", "omp", "--no-spinner"]) @@ -2580,8 +2588,8 @@ fn test_hourly_reports_project_reasoning_into_output() { "command failed: {json_output:?}" ); let json: serde_json::Value = serde_json::from_slice(&json_output.stdout).unwrap(); - assert_eq!(json["entries"][0]["output"], 50); - assert!(json["entries"][0].get("reasoning").is_none()); + assert_eq!(json["data"]["entries"][0]["output"], 50); + assert!(json["data"]["entries"][0].get("reasoning").is_none()); let table_output = cmd_with_home(base) .args(["hourly", "--client", "omp", "--no-spinner"]) @@ -2624,11 +2632,11 @@ fn test_models_report_clamps_reasoning_above_output() { assert!(output.status.success(), "command failed: {output:?}"); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - assert_eq!(json["entries"][0]["output"], 50); - assert!(json["entries"][0].get("reasoning").is_none()); - assert_eq!(json["totalOutput"], 50); + assert_eq!(json["data"]["entries"][0]["output"], 50); + assert!(json["data"]["entries"][0].get("reasoning").is_none()); + assert_eq!(json["data"]["totalOutput"], 50); assert!(json.get("totalReasoning").is_none()); - assert_eq!(json["totalTokens"], 165); + assert_eq!(json["data"]["totalTokens"], 165); assert!(json.get("warnings").is_none()); } @@ -2642,9 +2650,9 @@ fn test_models_group_by_model() { .unwrap(); assert!(output.status.success()); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - assert_eq!(json["groupBy"].as_str().unwrap(), "model"); + assert_eq!(json["data"]["groupBy"].as_str().unwrap(), "model"); - let entries = json["entries"].as_array().unwrap(); + let entries = json["data"]["entries"].as_array().unwrap(); let models: Vec<&str> = entries .iter() .map(|e| e["model"].as_str().unwrap()) @@ -2667,9 +2675,12 @@ fn test_models_group_by_client_provider_model() { .unwrap(); assert!(output.status.success()); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - assert_eq!(json["groupBy"].as_str().unwrap(), "client,provider,model"); + assert_eq!( + json["data"]["groupBy"].as_str().unwrap(), + "client,provider,model" + ); - let entries = json["entries"].as_array().unwrap(); + let entries = json["data"]["entries"].as_array().unwrap(); for entry in entries { assert!(entry.get("client").is_some(), "Entry must have client"); assert!(entry.get("provider").is_some(), "Entry must have provider"); @@ -2687,7 +2698,7 @@ fn test_models_json_with_group_by_model() { .unwrap(); assert!(output.status.success()); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - let entries = json["entries"].as_array().unwrap(); + let entries = json["data"]["entries"].as_array().unwrap(); for entry in entries { assert!( entry.get("mergedClients").is_some(), @@ -2718,9 +2729,9 @@ fn test_models_group_by_session_emits_session_id_per_entry() { .unwrap(); assert!(output.status.success(), "command failed: {:?}", output); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - assert_eq!(json["groupBy"].as_str().unwrap(), "session,model"); + assert_eq!(json["data"]["groupBy"].as_str().unwrap(), "session,model"); - let entries = json["entries"].as_array().unwrap(); + let entries = json["data"]["entries"].as_array().unwrap(); assert!(!entries.is_empty(), "expected at least one entry"); let mut session_ids: Vec<&str> = entries @@ -2761,9 +2772,12 @@ fn test_models_group_by_client_session_includes_client_and_session() { .unwrap(); assert!(output.status.success(), "command failed: {:?}", output); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - assert_eq!(json["groupBy"].as_str().unwrap(), "client,session,model"); + assert_eq!( + json["data"]["groupBy"].as_str().unwrap(), + "client,session,model" + ); - let entries = json["entries"].as_array().unwrap(); + let entries = json["data"]["entries"].as_array().unwrap(); assert!(!entries.is_empty()); for entry in entries { assert!(entry.get("sessionId").and_then(|v| v.as_str()).is_some()); @@ -2782,9 +2796,9 @@ fn test_models_group_by_workspace_model_uses_unknown_bucket_for_unsupported_clie .unwrap(); assert!(output.status.success()); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - assert_eq!(json["groupBy"].as_str().unwrap(), "workspace,model"); + assert_eq!(json["data"]["groupBy"].as_str().unwrap(), "workspace,model"); - let entries = json["entries"].as_array().unwrap(); + let entries = json["data"]["entries"].as_array().unwrap(); assert!(!entries.is_empty()); for entry in entries { assert!( @@ -2813,9 +2827,9 @@ fn test_models_group_by_workspace_model_surfaces_workspace_fields_for_qwen() { .unwrap(); assert!(output.status.success()); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - assert_eq!(json["groupBy"].as_str().unwrap(), "workspace,model"); + assert_eq!(json["data"]["groupBy"].as_str().unwrap(), "workspace,model"); - let entries = json["entries"].as_array().unwrap(); + let entries = json["data"]["entries"].as_array().unwrap(); assert_eq!(entries.len(), 1); assert_eq!( entries[0]["workspaceKey"].as_str().unwrap(), @@ -2838,9 +2852,9 @@ fn test_models_group_by_workspace_model_surfaces_workspace_fields_for_codex() { .unwrap(); assert!(output.status.success()); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - assert_eq!(json["groupBy"].as_str().unwrap(), "workspace,model"); + assert_eq!(json["data"]["groupBy"].as_str().unwrap(), "workspace,model"); - let entries = json["entries"].as_array().unwrap(); + let entries = json["data"]["entries"].as_array().unwrap(); assert_eq!(entries.len(), 1); assert_eq!( entries[0]["workspaceKey"].as_str().unwrap(), @@ -2869,9 +2883,9 @@ fn test_models_group_by_workspace_model_merges_claude_codex_pi_by_cwd() { .unwrap(); assert!(output.status.success(), "command failed: {:?}", output); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - assert_eq!(json["groupBy"].as_str().unwrap(), "workspace,model"); + assert_eq!(json["data"]["groupBy"].as_str().unwrap(), "workspace,model"); - let entries = json["entries"].as_array().unwrap(); + let entries = json["data"]["entries"].as_array().unwrap(); assert_eq!(entries.len(), 1); assert_eq!( entries[0]["workspaceKey"].as_str().unwrap(), @@ -2909,7 +2923,7 @@ fn test_models_client_filter_splits_pi_and_omp_sessions() { pi_output ); let pi_json: serde_json::Value = serde_json::from_slice(&pi_output.stdout).unwrap(); - let pi_entries = pi_json["entries"].as_array().unwrap(); + let pi_entries = pi_json["data"]["entries"].as_array().unwrap(); assert_eq!(pi_entries.len(), 1); assert_eq!(pi_entries[0]["client"].as_str().unwrap(), "pi"); assert_eq!(pi_entries[0]["input"].as_i64().unwrap(), 30); @@ -2925,7 +2939,7 @@ fn test_models_client_filter_splits_pi_and_omp_sessions() { omp_output ); let omp_json: serde_json::Value = serde_json::from_slice(&omp_output.stdout).unwrap(); - let omp_entries = omp_json["entries"].as_array().unwrap(); + let omp_entries = omp_json["data"]["entries"].as_array().unwrap(); assert_eq!(omp_entries.len(), 1); assert_eq!(omp_entries[0]["client"].as_str().unwrap(), "omp"); assert_eq!(omp_entries[0]["input"].as_i64().unwrap(), 40); @@ -2942,9 +2956,9 @@ fn test_models_group_by_workspace_model_surfaces_workspace_fields_for_opencode() .unwrap(); assert!(output.status.success()); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - assert_eq!(json["groupBy"].as_str().unwrap(), "workspace,model"); + assert_eq!(json["data"]["groupBy"].as_str().unwrap(), "workspace,model"); - let entries = json["entries"].as_array().unwrap(); + let entries = json["data"]["entries"].as_array().unwrap(); assert_eq!(entries.len(), 1); assert_eq!( entries[0]["workspaceKey"].as_str().unwrap(), @@ -2963,12 +2977,17 @@ fn test_models_group_by_workspace_model_surfaces_workspace_fields_for_opencode() fn test_pricing_command_success() { let tmp = create_pricing_fixture_dir(); let mut cmd = cmd_with_home(tmp.path()); - cmd.args(["pricing", "claude-sonnet-4-20250514", "--no-spinner"]) - .assert() - .success() - .stdout(predicate::str::contains("Pricing for")) - .stdout(predicate::str::contains("Input")) - .stdout(predicate::str::contains("Output")); + cmd.args([ + "pricing", + "lookup", + "claude-sonnet-4-20250514", + "--no-spinner", + ]) + .assert() + .success() + .stdout(predicate::str::contains("Pricing for")) + .stdout(predicate::str::contains("Input")) + .stdout(predicate::str::contains("Output")); } #[test] @@ -2977,6 +2996,7 @@ fn test_pricing_command_json() { let output = cmd_with_home(tmp.path()) .args([ "pricing", + "lookup", "claude-sonnet-4-20250514", "--json", "--no-spinner", @@ -2996,13 +3016,14 @@ fn test_pricing_command_json() { } #[test] -fn test_pricing_command_with_provider() { +fn test_pricing_command_with_source() { let tmp = create_pricing_fixture_dir(); let mut cmd = cmd_with_home(tmp.path()); cmd.args([ "pricing", + "lookup", "claude-sonnet-4-20250514", - "--provider", + "--source", "litellm", "--no-spinner", ]) @@ -3011,20 +3032,38 @@ fn test_pricing_command_with_provider() { } #[test] -fn test_pricing_command_invalid_provider() { +fn test_pricing_command_invalid_source() { let tmp = create_pricing_fixture_dir(); let mut cmd = cmd_with_home(tmp.path()); cmd.args([ "pricing", + "lookup", "claude-sonnet-4-20250514", - "--provider", - "invalid-provider", + "--source", + "invalid-source", "--no-spinner", ]) .assert() .failure(); } +#[test] +fn test_pricing_v4_spellings_are_rejected_with_exact_replacements() { + cargo_bin_cmd!("tokscale") + .args(["pricing", "list-overrides"]) + .assert() + .code(2) + .stderr(predicate::str::contains("use `tokscale pricing overrides`")); + + cargo_bin_cmd!("tokscale") + .args(["pricing", "gpt-5", "--json"]) + .assert() + .code(2) + .stderr(predicate::str::contains( + "use `tokscale pricing lookup gpt-5 --json`", + )); +} + #[test] fn test_pricing_command_does_not_fuzzy_match_provider_scoped_fireworks_model() { let tmp = TempDir::new().expect("failed to create temp dir"); @@ -3033,6 +3072,7 @@ fn test_pricing_command_does_not_fuzzy_match_provider_scoped_fireworks_model() { let output = cmd_with_home(tmp.path()) .args([ "pricing", + "lookup", "accounts/fireworks/models/deepseek-v4-pro", "--no-spinner", ]) @@ -3040,14 +3080,14 @@ fn test_pricing_command_does_not_fuzzy_match_provider_scoped_fireworks_model() { .unwrap(); assert!(!output.status.success()); - let stdout = String::from_utf8_lossy(&output.stdout); + let stderr = String::from_utf8_lossy(&output.stderr); assert!( - stdout.contains("Model not found: accounts/fireworks/models/deepseek-v4-pro"), - "expected a not-found message, got: {stdout}" + stderr.contains("Model not found: accounts/fireworks/models/deepseek-v4-pro"), + "expected a not-found message, got: {stderr}" ); assert!( - !stdout.contains("deepseek-r1-0528-distill-qwen3-8b"), - "provider-scoped pricing lookup must not report the wrong Fireworks match: {stdout}" + !stderr.contains("deepseek-r1-0528-distill-qwen3-8b"), + "provider-scoped pricing lookup must not report the wrong Fireworks match: {stderr}" ); } @@ -3072,7 +3112,7 @@ fn test_clients_command_reports_malformed_settings() { cmd_with_home(tmp.path()) .args(["clients", "--home", tmp.path().to_str().unwrap()]) .assert() - .failure() + .code(2) .stderr(predicate::str::contains("failed to parse settings JSON")) .stderr(predicate::str::contains( settings_json_path(tmp.path()).display().to_string(), @@ -3086,9 +3126,9 @@ fn excluded_crush_default_client_fails_before_report_output() { cmd_with_home(tmp.path()) .env("RUST_BACKTRACE", "1") - .args(["--light", "--no-spinner"]) + .args(["models", "--no-spinner"]) .assert() - .failure() + .code(2) .stdout(predicate::str::is_empty()) .stderr( predicate::str::contains("invalid client id(s) in settings.json defaultClients: crush") @@ -3098,6 +3138,67 @@ fn excluded_crush_default_client_fails_before_report_output() { ); } +#[test] +fn invalid_settings_range_is_invalid_execution_environment() { + let tmp = create_empty_fixture_dir(); + write_settings_json( + tmp.path(), + r#"{"autoRefreshMs":1,"nativeTimeoutMs":300000}"#, + ); + + cmd_with_home(tmp.path()) + .args(["clients", "--home", tmp.path().to_str().unwrap()]) + .assert() + .code(2) + .stdout(predicate::str::is_empty()) + .stderr(predicate::str::contains("invalid autoRefreshMs 1")); +} + +#[test] +fn non_utf8_settings_is_invalid_execution_environment() { + let tmp = create_empty_fixture_dir(); + let path = settings_json_path(tmp.path()); + fs::create_dir_all(path.parent().unwrap()).unwrap(); + fs::write(&path, b"{\"colorPalette\":\"\xff\"}").unwrap(); + + cmd_with_home(tmp.path()) + .args(["clients", "--home", tmp.path().to_str().unwrap()]) + .assert() + .code(2) + .stdout(predicate::str::is_empty()) + .stderr(predicate::str::contains("failed to parse settings JSON")); +} + +#[test] +fn malformed_scanner_environment_is_invalid_execution_environment() { + let tmp = create_empty_fixture_dir(); + + cmd_with_home(tmp.path()) + .env("TOKSCALE_EXTRA_DIRS", "broken") + .args(["models", "--client", "amp", "--no-spinner"]) + .assert() + .code(2) + .stdout(predicate::str::is_empty()) + .stderr( + predicate::str::contains("TOKSCALE_EXTRA_DIRS") + .and(predicate::str::contains("expected `client:path`")), + ); +} + +#[test] +fn unreadable_settings_path_remains_an_operational_error() { + let tmp = create_empty_fixture_dir(); + let path = settings_json_path(tmp.path()); + fs::create_dir_all(&path).unwrap(); + + cmd_with_home(tmp.path()) + .args(["clients", "--home", tmp.path().to_str().unwrap()]) + .assert() + .code(1) + .stdout(predicate::str::is_empty()) + .stderr(predicate::str::contains("failed to read settings file")); +} + #[test] fn test_clients_json() { let tmp = create_empty_fixture_dir(); @@ -3108,15 +3209,21 @@ fn test_clients_json() { assert!(output.status.success()); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); assert!(json.is_object(), "Clients JSON should be an object"); - assert!(json.get("clients").is_some(), "Should have 'clients' field"); assert!( - json.get("headlessRoots").is_some(), + json["data"].get("clients").is_some(), + "Should have 'clients' field" + ); + assert!( + json["data"].get("headlessRoots").is_some(), "Should have 'headlessRoots' field" ); - assert!(json.get("note").is_some(), "Should have 'note' field"); + assert!( + json["data"].get("note").is_some(), + "Should have 'note' field" + ); assert_eq!(json["health"]["complete"], true); - let arr = json["clients"].as_array().unwrap(); + let arr = json["data"]["clients"].as_array().unwrap(); assert!(!arr.is_empty(), "Should list at least one client"); let first = &arr[0]; @@ -3152,6 +3259,34 @@ fn test_clients_json() { && entry["exists"] == true)); } +#[test] +fn test_clients_filter_does_not_discover_unselected_opencode() { + let tmp = create_empty_fixture_dir(); + let opencode_data_root = tmp.path().join(".local/share/opencode"); + fs::remove_dir_all(&opencode_data_root).unwrap(); + fs::create_dir_all(opencode_data_root.parent().unwrap()).unwrap(); + fs::write(&opencode_data_root, "not a directory").unwrap(); + + let output = cmd_with_home(tmp.path()) + .args(["clients", "--json", "--client", "claude"]) + .output() + .unwrap(); + + assert!( + output.status.success(), + "stderr: {}", + String::from_utf8_lossy(&output.stderr) + ); + let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); + let clients = json["data"]["clients"].as_array().unwrap(); + assert_eq!(clients.len(), 1); + assert_eq!(clients[0]["client"], "claude"); + assert_eq!(json["health"]["failedSources"], 0); + assert!(!serde_json::to_string(&json["health"]) + .unwrap() + .contains("opencode")); +} + #[test] fn test_clients_json_reports_degraded_source_health_without_losing_payload() { let tmp = create_empty_fixture_dir(); @@ -3175,7 +3310,7 @@ fn test_clients_json_reports_degraded_source_health_without_losing_payload() { String::from_utf8_lossy(&output.stderr) ); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - assert!(json["clients"] + assert!(json["data"]["clients"] .as_array() .is_some_and(|rows| !rows.is_empty())); assert_eq!(json["health"]["complete"], false); @@ -3229,7 +3364,7 @@ fn test_clients_json_reports_broken_claude_mirror_without_losing_payload() { String::from_utf8_lossy(&output.stderr) ); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - assert!(json["clients"] + assert!(json["data"]["clients"] .as_array() .is_some_and(|rows| !rows.is_empty())); assert_eq!(json["health"]["failedSources"], 1); @@ -3290,7 +3425,7 @@ fn test_clients_json_opencode_diagnostics_match_adapter_for_non_utf8_xdg() { assert!(output.status.success()); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - let opencode = json["clients"] + let opencode = json["data"]["clients"] .as_array() .unwrap() .iter() @@ -3322,7 +3457,7 @@ fn test_clients_json_warp_sessions_path_exists_tracks_selected_root() { assert!(output.status.success()); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - let warp = json["clients"] + let warp = json["data"]["clients"] .as_array() .unwrap() .iter() @@ -3353,7 +3488,7 @@ fn test_clients_json_includes_claude_transcripts_path() { assert!(output.status.success()); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - let claude = json["clients"] + let claude = json["data"]["clients"] .as_array() .unwrap() .iter() @@ -3393,7 +3528,7 @@ fn test_clients_json_includes_claude_desktop_diagnostic() { assert!(output.status.success()); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - let claude = json["clients"] + let claude = json["data"]["clients"] .as_array() .unwrap() .iter() @@ -3429,7 +3564,7 @@ fn test_clients_command_includes_claude_desktop_diagnostic_text() { } #[test] -fn test_models_json_includes_claude_desktop_diagnostic_for_empty_explicit_claude_report() { +fn test_models_json_routes_claude_desktop_diagnostic_to_stderr() { let tmp = create_empty_fixture_dir(); fs::create_dir_all(tmp.path().join("Library/Application Support/Claude")).unwrap(); @@ -3440,15 +3575,9 @@ fn test_models_json_includes_claude_desktop_diagnostic_for_empty_explicit_claude assert!(output.status.success()); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - let diagnostics = json["diagnostics"].as_array().unwrap(); - - assert!(diagnostics.iter().any(|item| { - item["code"] == "claude_desktop_not_scanned" - && item["message"] - .as_str() - .unwrap() - .contains("Tokscale counts Claude Code JSONL transcripts") - })); + assert!(json.get("diagnostics").is_none()); + assert!(String::from_utf8_lossy(&output.stderr) + .contains("Tokscale counts Claude Code JSONL transcripts")); } #[test] @@ -3472,7 +3601,7 @@ fn test_clients_json_includes_settings_extra_paths() { assert!(output.status.success()); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - let codex = json["clients"] + let codex = json["data"]["clients"] .as_array() .unwrap() .iter() @@ -3515,7 +3644,7 @@ fn test_clients_json_includes_hermes_settings_extra_profile_path() { assert!(output.status.success()); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - let hermes = json["clients"] + let hermes = json["data"]["clients"] .as_array() .unwrap() .iter() @@ -3587,7 +3716,7 @@ fn test_clients_command_groups_opencode_database_paths_by_source() { .unwrap(); assert!(output.status.success()); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - let opencode = json["clients"] + let opencode = json["data"]["clients"] .as_array() .unwrap() .iter() @@ -3604,13 +3733,13 @@ fn test_clients_command_groups_opencode_database_paths_by_source() { })); } -// ── Light mode tests ─────────────────────────────────────────────────────── +// ── Table report tests ───────────────────────────────────────────────────── #[test] -fn test_models_light_output() { +fn test_models_table_output() { let tmp = create_temp_fixture_dir(); cmd_with_home(tmp.path()) - .args(["models", "--light", "--client", "opencode", "--no-spinner"]) + .args(["models", "--client", "opencode", "--no-spinner"]) .assert() .success() .stdout(predicate::str::contains("Token Usage Report by Model")) @@ -3618,20 +3747,57 @@ fn test_models_light_output() { } #[test] -fn test_monthly_light_output() { +fn test_monthly_table_output() { let tmp = create_temp_fixture_dir(); cmd_with_home(tmp.path()) - .args(["monthly", "--light", "--client", "opencode", "--no-spinner"]) + .args(["monthly", "--client", "opencode", "--no-spinner"]) .assert() .success() .stdout(predicate::str::contains("Monthly Token Usage Report")); } #[test] -fn test_models_light_with_client_filter() { +fn test_hourly_table_output() { + let tmp = create_temp_fixture_dir(); + cmd_with_home(tmp.path()) + .args(["hourly", "--client", "opencode", "--no-spinner"]) + .assert() + .success() + .stdout(predicate::str::contains("Hourly Usage")); +} + +#[test] +fn test_time_metrics_table_output() { + let tmp = create_temp_fixture_dir(); + cmd_with_home(tmp.path()) + .args(["time-metrics", "--client", "opencode", "--no-spinner"]) + .assert() + .success() + .stdout(predicate::str::contains("Session Time Metrics")); +} + +#[test] +fn test_time_metrics_benchmark_flag() { + let tmp = create_temp_fixture_dir(); + cmd_with_home(tmp.path()) + .args([ + "time-metrics", + "--client", + "opencode", + "--no-spinner", + "--benchmark", + ]) + .assert() + .success() + .stdout(predicate::str::contains("Processing time").not()) + .stderr(predicate::str::contains("Processing time")); +} + +#[test] +fn test_models_table_with_client_filter() { let tmp = create_temp_fixture_dir(); cmd_with_home(tmp.path()) - .args(["models", "--light", "--client", "opencode", "--no-spinner"]) + .args(["models", "--client", "opencode", "--no-spinner"]) .args(["--year", "2024"]) .assert() .success() @@ -3646,7 +3812,6 @@ fn test_models_benchmark_flag() { cmd_with_home(tmp.path()) .args([ "models", - "--light", "--client", "opencode", "--no-spinner", @@ -3654,7 +3819,8 @@ fn test_models_benchmark_flag() { ]) .assert() .success() - .stdout(predicate::str::contains("Processing time")); + .stdout(predicate::str::contains("Processing time").not()) + .stderr(predicate::str::contains("Processing time")); } #[test] @@ -3663,7 +3829,6 @@ fn test_monthly_benchmark_flag() { cmd_with_home(tmp.path()) .args([ "monthly", - "--light", "--client", "opencode", "--no-spinner", @@ -3671,7 +3836,8 @@ fn test_monthly_benchmark_flag() { ]) .assert() .success() - .stdout(predicate::str::contains("Processing time")); + .stdout(predicate::str::contains("Processing time").not()) + .stderr(predicate::str::contains("Processing time")); } // ── Empty fixture tests ──────────────────────────────────────────────────── @@ -3685,13 +3851,13 @@ fn test_models_empty_fixture() { .unwrap(); assert!(output.status.success()); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - let entries = json["entries"].as_array().unwrap(); + let entries = json["data"]["entries"].as_array().unwrap(); assert!( entries.is_empty(), "Empty fixture should produce no entries" ); - assert_eq!(json["totalInput"].as_i64().unwrap(), 0); - assert_eq!(json["totalOutput"].as_i64().unwrap(), 0); + assert_eq!(json["data"]["totalInput"].as_i64().unwrap(), 0); + assert_eq!(json["data"]["totalOutput"].as_i64().unwrap(), 0); } #[test] @@ -3703,7 +3869,7 @@ fn test_graph_empty_contributions() { .unwrap(); assert!(output.status.success()); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - let contributions = json["contributions"].as_array().unwrap(); + let contributions = json["data"]["contributions"].as_array().unwrap(); assert!( contributions.is_empty(), "Empty fixture should produce no contributions" @@ -3716,7 +3882,7 @@ fn test_graph_empty_contributions() { fn test_models_no_spinner_flag() { let tmp = create_temp_fixture_dir(); cmd_with_home(tmp.path()) - .args(["models", "--light", "--client", "opencode", "--no-spinner"]) + .args(["models", "--client", "opencode", "--no-spinner"]) .assert() .success(); } @@ -3741,7 +3907,7 @@ fn test_graph_with_client_filter() { .unwrap(); assert!(output.status.success()); let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - let contributions = json["contributions"].as_array().unwrap(); + let contributions = json["data"]["contributions"].as_array().unwrap(); for c in contributions { let clients = c["clients"].as_array().unwrap(); for cl in clients { @@ -3760,51 +3926,55 @@ fn test_graph_with_client_filter() { fn test_graph_output_to_file() { let tmp = create_temp_fixture_dir(); let output_file = tmp.path().join("graph-output.json"); - cmd_with_home(tmp.path()) + let output = cmd_with_home(tmp.path()) .args(["graph", "--client", "opencode", "--no-spinner"]) .args(["--output", output_file.to_str().unwrap()]) - .assert() - .success(); + .output() + .unwrap(); + assert!( + output.status.success(), + "stderr: {}", + String::from_utf8_lossy(&output.stderr) + ); + assert_eq!( + String::from_utf8_lossy(&output.stdout), + format!("{}\n", output_file.display()) + ); assert!(output_file.exists(), "Output file should be created"); let content = fs::read_to_string(&output_file).unwrap(); let json: serde_json::Value = serde_json::from_str(&content).unwrap(); - assert!(json.get("meta").is_some()); - assert!(json.get("contributions").is_some()); + assert!(json["data"].get("meta").is_some()); + assert!(json["data"].get("contributions").is_some()); } -// ── Root command tests (no subcommand) ───────────────────────────────────── +// ── Root command ownership tests ─────────────────────────────────────────── #[test] -fn test_root_json_output() { - let tmp = create_temp_fixture_dir(); - let output = cmd_with_home(tmp.path()) - .args(["--json", "--client", "opencode", "--no-spinner"]) - .output() - .unwrap(); - assert!(output.status.success()); - let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - assert!(json.get("entries").is_some()); - assert!(json.get("totalCost").is_some()); +fn test_root_rejects_json_report_options() { + cargo_bin_cmd!("tokscale") + .args(["--json", "models"]) + .assert() + .code(2) + .stderr(predicate::str::contains("use `tokscale models --json`")); } #[test] -fn test_root_light_output() { - let tmp = create_temp_fixture_dir(); - cmd_with_home(tmp.path()) - .args(["--light", "--client", "opencode", "--no-spinner"]) +fn test_root_rejects_removed_light_option() { + cargo_bin_cmd!("tokscale") + .arg("--light") .assert() - .success() - .stdout(predicate::str::contains("Token Usage Report by Model")); + .code(2) + .stderr(predicate::str::contains("use `tokscale models`")); } #[test] -fn light_report_surfaces_malformed_display_config_without_panicking() { +fn table_report_surfaces_malformed_display_config_without_panicking() { let tmp = create_temp_fixture_dir(); let config_path = tmp.path().join(".tokscale"); fs::write(&config_path, "[display_names.providers\n").unwrap(); cmd_with_home(tmp.path()) - .args(["--light", "--client", "opencode", "--no-spinner"]) + .args(["models", "--client", "opencode", "--no-spinner"]) .assert() .failure() .stdout(predicate::str::is_empty()) @@ -3816,16 +3986,16 @@ fn light_report_surfaces_malformed_display_config_without_panicking() { } #[test] -fn home_write_cache_conflict_fails_before_report_output() { +fn removed_write_cache_flag_fails_before_report_output() { let tmp = create_temp_fixture_dir(); let scoped_home = tmp.path().join("scoped-home"); fs::create_dir_all(&scoped_home).unwrap(); cmd_with_home(tmp.path()) .args([ + "models", "--home", scoped_home.to_str().unwrap(), - "--light", "--write-cache", "--client", "opencode", @@ -3834,13 +4004,11 @@ fn home_write_cache_conflict_fails_before_report_output() { .assert() .failure() .stdout(predicate::str::is_empty()) - .stderr(predicate::str::contains( - "--write-cache cannot be combined with --home", - )); + .stderr(predicate::str::contains("--write-cache")); } #[test] -fn home_settings_write_cache_conflict_fails_before_report_output() { +fn removed_light_setting_has_no_report_cache_side_effect() { let tmp = create_temp_fixture_dir(); let scoped_home = tmp.path().join("scoped-home"); fs::create_dir_all(&scoped_home).unwrap(); @@ -3848,66 +4016,56 @@ fn home_settings_write_cache_conflict_fails_before_report_output() { cmd_with_home(tmp.path()) .args([ + "models", "--home", scoped_home.to_str().unwrap(), - "--light", "--client", "opencode", "--no-spinner", ]) .assert() - .failure() - .stdout(predicate::str::is_empty()) - .stderr(predicate::str::contains( - "--write-cache cannot be combined with --home", - )); + .success(); + assert!(!tmp + .path() + .join(".cache/tokscale/tui-data-cache.json") + .exists()); } #[test] -fn light_with_write_cache_writes_to_canonical_path() { +fn cache_warm_writes_to_canonical_path() { let tmp = create_temp_fixture_dir(); let config_dir = tmp.path().join("custom-config-root"); prime_override_pricing_cache(&config_dir); cmd_with_home(tmp.path()) .env("TOKSCALE_CONFIG_DIR", &config_dir) - .args([ - "--light", - "--client", - "opencode", - "--write-cache", - "--no-spinner", - ]) + .args(["cache", "warm", "--client", "opencode"]) .assert() .success() - .stdout(predicate::str::contains("Token Usage Report by Model")); + .stdout(predicate::str::contains("TUI cache warmed")); assert!( config_dir.join("cache/tui-data-cache.json").exists(), - "--write-cache should populate the canonical cache path" + "cache warm should populate the canonical cache path" ); } #[test] -fn test_root_with_date_filter() { - let tmp = create_temp_fixture_dir(); - cmd_with_home(tmp.path()) - .args(["--json", "--client", "opencode", "--no-spinner"]) +fn test_root_rejects_date_filter() { + cargo_bin_cmd!("tokscale") .args(["--year", "2025"]) .assert() - .success() - .stdout(predicate::str::contains("gpt-4o")); + .code(2) + .stderr(predicate::str::contains("use `tokscale tui --year 2025`")); } #[test] -fn test_root_with_group_by() { - let tmp = create_temp_fixture_dir(); - let output = cmd_with_home(tmp.path()) - .args(["--json", "--client", "opencode", "--no-spinner"]) +fn test_root_rejects_group_by() { + cargo_bin_cmd!("tokscale") .args(["--group-by", "model"]) - .output() - .unwrap(); - assert!(output.status.success()); - let json: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); - assert_eq!(json["groupBy"].as_str().unwrap(), "model"); + .assert() + .code(2) + .stderr(predicate::str::contains( + "use `tokscale models --group-by model`", + )); } diff --git a/crates/tokscale-cli/tests/copilot_memory.rs b/crates/tokscale-cli/tests/copilot_memory.rs index b3a088781..37bd29e30 100644 --- a/crates/tokscale-cli/tests/copilot_memory.rs +++ b/crates/tokscale-cli/tests/copilot_memory.rs @@ -160,19 +160,18 @@ fn extract_max_rss_kb(stderr: &[u8]) -> u64 { .unwrap_or_else(|| panic!("missing /usr/bin/time max RSS output: {stderr}")) } -fn run_light_copilot(home: &Path) -> (Vec, u64) { +fn run_copilot_report(home: &Path) -> (Vec, u64) { let output = Command::new("/usr/bin/time") .arg("-f") .arg("MAXRSS_KB=%M") .arg(assert_cmd::cargo::cargo_bin!("tokscale")) .args([ + "models", "--home", home.to_str().unwrap(), - "--light", "--client", "copilot", "--no-spinner", - "--no-write-cache", ]) .env("HOME", home) .env("XDG_CONFIG_HOME", home.join(".config")) @@ -208,7 +207,7 @@ fn copilot_large_otel_cold_and_warm_source_cache_stay_below_memory_limit() { prime_pricing_cache(home.path()); write_large_copilot_fixture(home.path()); - let (cold_stdout, cold_rss_kb) = run_light_copilot(home.path()); + let (cold_stdout, cold_rss_kb) = run_copilot_report(home.path()); let cold_shards = source_cache_shards(home.path()); assert_eq!( cold_shards.len(), @@ -218,7 +217,7 @@ fn copilot_large_otel_cold_and_warm_source_cache_stay_below_memory_limit() { let shard_path = cold_shards[0].clone(); let cold_shard_identity = shard_identity(&shard_path); - let (warm_stdout, warm_rss_kb) = run_light_copilot(home.path()); + let (warm_stdout, warm_rss_kb) = run_copilot_report(home.path()); let warm_shards = source_cache_shards(home.path()); assert_eq!(warm_stdout, cold_stdout); diff --git a/crates/tokscale-cli/tests/tui_exit_tests.rs b/crates/tokscale-cli/tests/tui_exit_tests.rs new file mode 100644 index 000000000..6b6c5a6e0 --- /dev/null +++ b/crates/tokscale-cli/tests/tui_exit_tests.rs @@ -0,0 +1,161 @@ +#![cfg(unix)] + +use std::fs::File; +use std::io::{Read, Write}; +use std::os::fd::{AsRawFd, FromRawFd}; +use std::process::{Command, ExitStatus, Stdio}; +use std::sync::mpsc; +use std::thread; +use std::time::{Duration, Instant}; + +use tempfile::TempDir; + +const ENTER_ALTERNATE_SCREEN: &[u8] = b"\x1b[?1049h"; +const LEAVE_ALTERNATE_SCREEN: &[u8] = b"\x1b[?1049l"; + +fn terminal_mode(fd: i32) -> libc::termios { + // SAFETY: `tcgetattr` initializes the complete termios value on success, + // and every caller supplies a live PTY slave descriptor. + unsafe { + let mut mode = std::mem::zeroed(); + assert_eq!(libc::tcgetattr(fd, &mut mode), 0); + mode + } +} + +fn assert_terminal_mode_restored(before: &libc::termios, after: &libc::termios) { + let local_flags = libc::ECHO | libc::ICANON | libc::IEXTEN | libc::ISIG; + let input_flags = libc::BRKINT | libc::ICRNL | libc::INPCK | libc::ISTRIP | libc::IXON; + + assert_eq!(after.c_lflag & local_flags, before.c_lflag & local_flags); + assert_eq!(after.c_iflag & input_flags, before.c_iflag & input_flags); + assert_eq!(after.c_oflag & libc::OPOST, before.c_oflag & libc::OPOST); + assert_eq!(after.c_cc[libc::VMIN], before.c_cc[libc::VMIN]); + assert_eq!(after.c_cc[libc::VTIME], before.c_cc[libc::VTIME]); +} + +fn wait_for_exit(child: &mut std::process::Child) -> ExitStatus { + let deadline = Instant::now() + Duration::from_secs(10); + loop { + if let Some(status) = child.try_wait().expect("poll TUI child") { + return status; + } + if Instant::now() >= deadline { + let _ = child.kill(); + let _ = child.wait(); + panic!("TUI did not exit after input"); + } + thread::sleep(Duration::from_millis(10)); + } +} + +fn run_tui_with_input(input: &[u8]) -> (ExitStatus, Vec) { + let mut master_fd = -1; + let mut slave_fd = -1; + let size = libc::winsize { + ws_row: 30, + ws_col: 100, + ws_xpixel: 0, + ws_ypixel: 0, + }; + + // SAFETY: all output pointers are valid, and ownership of both returned + // descriptors is immediately transferred to `File` below. + let opened = unsafe { + libc::openpty( + &mut master_fd, + &mut slave_fd, + std::ptr::null_mut(), + std::ptr::null(), + &size, + ) + }; + assert_eq!( + opened, + 0, + "openpty failed: {}", + std::io::Error::last_os_error() + ); + + // SAFETY: `openpty` returned two new owned descriptors on success. + let mut master = unsafe { File::from_raw_fd(master_fd) }; + // SAFETY: same ownership transfer as the master descriptor above. + let slave = unsafe { File::from_raw_fd(slave_fd) }; + let original_mode = terminal_mode(slave.as_raw_fd()); + + let stdin = slave.try_clone().expect("clone PTY slave for stdin"); + let stdout = slave.try_clone().expect("clone PTY slave for stdout"); + let stderr = slave.try_clone().expect("clone PTY slave for stderr"); + let mut reader = master.try_clone().expect("clone PTY master for reader"); + let (ready_tx, ready_rx) = mpsc::sync_channel(1); + let reader_thread = thread::spawn(move || { + let mut output = Vec::new(); + let mut buffer = [0_u8; 4096]; + let mut ready_tx = Some(ready_tx); + loop { + match reader.read(&mut buffer) { + Ok(0) => break, + Ok(read) => { + output.extend_from_slice(&buffer[..read]); + if ready_tx.is_some() + && output + .windows(ENTER_ALTERNATE_SCREEN.len()) + .any(|window| window == ENTER_ALTERNATE_SCREEN) + { + let _ = ready_tx.take().unwrap().send(()); + } + } + Err(error) if error.raw_os_error() == Some(libc::EIO) => break, + Err(error) => panic!("read TUI PTY output: {error}"), + } + } + output + }); + + let home = TempDir::new().expect("create isolated TUI home"); + let mut child = Command::new(env!("CARGO_BIN_EXE_tokscale")) + .args(["tui", "--no-refresh"]) + .env("HOME", home.path()) + .env("TOKSCALE_CONFIG_DIR", home.path().join("tokscale-config")) + .env("TOKSCALE_PRICING_CACHE_ONLY", "1") + .env("TERM", "xterm-256color") + .stdin(Stdio::from(stdin)) + .stdout(Stdio::from(stdout)) + .stderr(Stdio::from(stderr)) + .spawn() + .expect("spawn TUI in PTY"); + + ready_rx + .recv_timeout(Duration::from_secs(10)) + .expect("TUI did not enter alternate screen"); + master.write_all(input).expect("write TUI key input"); + master.flush().expect("flush TUI key input"); + + let status = wait_for_exit(&mut child); + let restored_mode = terminal_mode(slave.as_raw_fd()); + assert_terminal_mode_restored(&original_mode, &restored_mode); + + drop(slave); + drop(master); + let output = reader_thread.join().expect("join PTY reader"); + assert!( + output + .windows(LEAVE_ALTERNATE_SCREEN.len()) + .any(|window| window == LEAVE_ALTERNATE_SCREEN), + "TUI did not leave the alternate screen" + ); + + (status, output) +} + +#[test] +fn q_exits_tui_successfully_after_restoring_terminal() { + let (status, _) = run_tui_with_input(b"q"); + assert_eq!(status.code(), Some(0)); +} + +#[test] +fn ctrl_c_exits_tui_with_130_after_restoring_terminal() { + let (status, _) = run_tui_with_input(b"\x03"); + assert_eq!(status.code(), Some(130)); +} diff --git a/crates/tokscale-core/client-catalog.json b/crates/tokscale-core/client-catalog.json index 3d563bd36..2e3cb7cef 100644 --- a/crates/tokscale-core/client-catalog.json +++ b/crates/tokscale-core/client-catalog.json @@ -26,15 +26,6 @@ "logo": "https://raw.githubusercontent.com/makoMakoGo/tokscale/personal/local-clients/.github/assets/client-openai.jpg", "color": "#3b82f6" }, - { - "variant": "Cursor", - "id": "cursor", - "displayName": "Cursor", - "shortName": "Cursor", - "hotkey": "4", - "logo": "https://raw.githubusercontent.com/makoMakoGo/tokscale/personal/local-clients/.github/assets/client-cursor.jpg", - "color": "#22c55e" - }, { "variant": "Gemini", "id": "gemini", @@ -235,15 +226,6 @@ "logo": "https://github.com/JetBrains.png", "color": "#7B61FF" }, - { - "variant": "Trae", - "id": "trae", - "displayName": "Trae", - "shortName": "Trae", - "hotkey": "y", - "logo": "https://raw.githubusercontent.com/makoMakoGo/tokscale/personal/local-clients/.github/assets/client-trae.png", - "color": "#00BFA5" - }, { "variant": "Warp", "id": "warp", diff --git a/crates/tokscale-core/src/adapters/file.rs b/crates/tokscale-core/src/adapters/file.rs index 4a269e568..7e9f5c230 100644 --- a/crates/tokscale-core/src/adapters/file.rs +++ b/crates/tokscale-core/src/adapters/file.rs @@ -1,4 +1,4 @@ -use std::path::Path; +use std::path::{Path, PathBuf}; use rayon::prelude::*; @@ -19,7 +19,6 @@ const GROK_TOTAL_ONLY_IMPUTATION_REVISION: u32 = MODEL_ID_CANONICALIZATION_REVIS const MUX_STABLE_DEDUP_REVISION: u32 = MODEL_ID_CANONICALIZATION_REVISION + 1; const QWEN_RECORD_REJECTION_REVISION: u32 = MODEL_ID_CANONICALIZATION_REVISION + 1; const ZCODE_OVERLAP_NORMALIZATION_REVISION: u32 = MODEL_ID_CANONICALIZATION_REVISION + 1; -const CURSOR_RECORD_REJECTION_REVISION: u32 = MODEL_ID_CANONICALIZATION_REVISION + 1; const KIMI_RECORD_REJECTION_REVISION: u32 = MODEL_ID_CANONICALIZATION_REVISION + 1; const COMMANDCODE_RECORD_REJECTION_REVISION: u32 = MODEL_ID_CANONICALIZATION_REVISION + 1; const ZCODE_RECORD_REJECTION_REVISION: u32 = ZCODE_OVERLAP_NORMALIZATION_REVISION + 1; @@ -30,6 +29,7 @@ const GROK_RECORD_REJECTION_REVISION: u32 = GROK_TOTAL_ONLY_IMPUTATION_REVISION const GROK_RELATED_METADATA_REVISION: u32 = GROK_RECORD_REJECTION_REVISION + 1; const GEMINI_RECORD_REJECTION_REVISION: u32 = MODEL_ID_CANONICALIZATION_REVISION + 1; const DROID_RECORD_REJECTION_REVISION: u32 = MODEL_ID_CANONICALIZATION_REVISION + 1; +const DROID_AGENT_ATTRIBUTION_REVISION: u32 = DROID_RECORD_REJECTION_REVISION + 1; const GROK_RELATED_METADATA_SIBLINGS: &[&str] = &["summary.json", "events.jsonl"]; pub(crate) struct CachedFileAdapter { @@ -37,6 +37,7 @@ pub(crate) struct CachedFileAdapter { parser_version: ParserVersion, fingerprint_policy: FingerprintPolicy, optional_related_inputs: bool, + dependency_path: Option Option>, parse: fn(&Path) -> SessionParseResult, } @@ -52,6 +53,24 @@ impl CachedFileAdapter { parser_version: ParserVersion::new(parser_id, revision), fingerprint_policy: FingerprintPolicy::PlainFile, optional_related_inputs: false, + dependency_path: None, + parse, + } + } + + pub(crate) const fn new_with_dependency( + client: ClientId, + parser_id: ParserId, + revision: u32, + dependency_path: fn(&Path) -> Option, + parse: fn(&Path) -> SessionParseResult, + ) -> Self { + Self { + client, + parser_version: ParserVersion::new(parser_id, revision), + fingerprint_policy: FingerprintPolicy::PlainFile, + optional_related_inputs: false, + dependency_path: Some(dependency_path), parse, } } @@ -68,6 +87,7 @@ impl CachedFileAdapter { parser_version: ParserVersion::new(parser_id, revision), fingerprint_policy: FingerprintPolicy::PrimaryWithSiblings { sibling_names }, optional_related_inputs: true, + dependency_path: None, parse, } } @@ -88,7 +108,16 @@ impl LocalSourceAdapter for CachedFileAdapter { self.fingerprint_policy.clone(), )? .into_iter() - .map(|unit| unit.with_parser_version(self.parser_version)) + .map(|unit| { + let dependency_path = self + .dependency_path + .and_then(|dependency_path| dependency_path(&unit.path)); + let unit = match dependency_path { + Some(dependency_path) => unit.with_dependency(dependency_path), + None => unit, + }; + unit.with_parser_version(self.parser_version) + }) .collect()) } @@ -200,12 +229,6 @@ impl LocalSourceAdapter for CopilotAdapter { } pub(crate) static COPILOT_ADAPTER: CopilotAdapter = CopilotAdapter; -pub(crate) static CURSOR_ADAPTER: CachedFileAdapter = CachedFileAdapter::new( - ClientId::Cursor, - ParserId::Cursor, - CURSOR_RECORD_REJECTION_REVISION, - sessions::cursor::parse_cursor_file, -); pub(crate) static GEMINI_ADAPTER: CachedFileAdapter = CachedFileAdapter::new( ClientId::Gemini, ParserId::Gemini, @@ -225,10 +248,11 @@ pub(crate) static AMP_ADAPTER: CachedFileAdapter = CachedFileAdapter::new( AMP_RECORD_REJECTION_REVISION, sessions::amp::parse_amp_file, ); -pub(crate) static DROID_ADAPTER: CachedFileAdapter = CachedFileAdapter::new( +pub(crate) static DROID_ADAPTER: CachedFileAdapter = CachedFileAdapter::new_with_dependency( ClientId::Droid, ParserId::Droid, - DROID_RECORD_REJECTION_REVISION, + DROID_AGENT_ATTRIBUTION_REVISION, + sessions::droid::droid_agent_dependency_path, sessions::droid::parse_droid_file, ); pub(crate) static KIMI_ADAPTER: CachedFileAdapter = CachedFileAdapter::new( @@ -570,14 +594,72 @@ not-json assert_eq!(fold_ctx.health.failed_sources(), 0); } + #[test] + fn droid_adapter_invalidates_cached_mission_worker_role_from_features() { + let home = tempfile::TempDir::new().unwrap(); + let session_dir = home.path().join(".factory/sessions/project"); + let settings_path = session_dir.join("mission-worker.settings.json"); + let features_path = home + .path() + .join(".factory/missions/mission-root/features.json"); + write_file( + &settings_path, + r#"{ + "model": "custom:gpt-5.6-sol-xhigh", + "providerLock": "openai", + "providerLockTimestamp": "2026-07-15T08:55:13.871Z", + "tokenUsage": {"inputTokens": 10, "outputTokens": 5}, + "tags": [ + {"name": "exec"}, + {"name": "mission-worker"}, + { + "name": "mission-session", + "metadata": {"role": "worker", "missionId": "mission-root"} + } + ] + }"#, + ); + write_file( + &features_path, + r#"{"features":[{"id":"implementation","skillName":"backend-worker","workerSessionIds":["mission-worker"]}]}"#, + ); + let settings = crate::scanner::ScannerSettings::default(); + let ctx = scan_context(home.path(), &settings); + let unit = DROID_ADAPTER.discover_checked(&ctx).unwrap().pop().unwrap(); + assert_eq!( + unit.fingerprint_policy, + FingerprintPolicy::PrimaryWithDependency { + dependency_path: features_path.clone() + } + ); + + let mut cache = message_cache::SourceMessageCache::default(); + let worker_messages = fold_with_adapter(&DROID_ADAPTER, vec![unit], &mut cache); + assert_eq!(worker_messages.len(), 1); + assert_eq!(worker_messages[0].agent.as_deref(), Some("Droid Worker")); + + write_file( + &features_path, + r#"{"features":[{"id":"scrutiny","skillName":"scrutiny-validator","workerSessionIds":["mission-worker"]}]}"#, + ); + let changed_unit = DROID_ADAPTER.discover_checked(&ctx).unwrap().pop().unwrap(); + let changed_unit = match DROID_ADAPTER.plan_cache_hit(changed_unit, &cache).unwrap() { + crate::adapters::CacheHitPlan::Miss(unit) => unit, + crate::adapters::CacheHitPlan::Hit(_) => { + panic!("changed Mission feature must invalidate the Droid source cache") + } + }; + let validator_messages = fold_with_adapter(&DROID_ADAPTER, vec![changed_unit], &mut cache); + assert_eq!(validator_messages.len(), 1); + assert_eq!( + validator_messages[0].agent.as_deref(), + Some("Droid Validator") + ); + } + #[test] fn cached_file_adapters_use_their_actual_record_rejection_revisions() { for (actual, parser_id, revision) in [ - ( - CURSOR_ADAPTER.parser_version, - ParserId::Cursor, - CURSOR_RECORD_REJECTION_REVISION, - ), ( GEMINI_ADAPTER.parser_version, ParserId::Gemini, @@ -596,7 +678,7 @@ not-json ( DROID_ADAPTER.parser_version, ParserId::Droid, - DROID_RECORD_REJECTION_REVISION, + DROID_AGENT_ATTRIBUTION_REVISION, ), ( KIMI_ADAPTER.parser_version, diff --git a/crates/tokscale-core/src/adapters/mod.rs b/crates/tokscale-core/src/adapters/mod.rs index af147c80d..1f2305ad5 100644 --- a/crates/tokscale-core/src/adapters/mod.rs +++ b/crates/tokscale-core/src/adapters/mod.rs @@ -17,7 +17,6 @@ mod opencode; mod omp; mod pi; -mod trae; mod vscode_tasks; mod warp; mod zed; @@ -514,7 +513,6 @@ fn default_parser_id(client: ClientId) -> ParserId { ClientId::OpenCode => ParserId::OpenCodeSqlite, ClientId::Claude => ParserId::Claude, ClientId::Codex => ParserId::Codex, - ClientId::Cursor => ParserId::Cursor, ClientId::Gemini => ParserId::Gemini, ClientId::Amp => ParserId::Amp, ClientId::Droid => ParserId::Droid, @@ -537,7 +535,6 @@ fn default_parser_id(client: ClientId) -> ParserId { ClientId::Zcode => ParserId::Zcode, ClientId::Kiro => ParserId::Kiro, ClientId::Junie => ParserId::Junie, - ClientId::Trae => ParserId::Trae, ClientId::Cline => ParserId::Cline, ClientId::CommandCode => ParserId::CommandCode, ClientId::Grok => ParserId::Grok, @@ -641,7 +638,7 @@ impl ParsedUnit { } } -static LOCAL_SOURCE_ADAPTERS: [&dyn LocalSourceAdapter; 31] = [ +static LOCAL_SOURCE_ADAPTERS: [&dyn LocalSourceAdapter; 29] = [ &zed::ZED_ADAPTER, &pi::PI_ADAPTER, &omp::OMP_ADAPTER, @@ -649,7 +646,6 @@ static LOCAL_SOURCE_ADAPTERS: [&dyn LocalSourceAdapter; 31] = [ &codex::CODEX_ADAPTER, &opencode::OPENCODE_ADAPTER, &file::COPILOT_ADAPTER, - &file::CURSOR_ADAPTER, &file::GEMINI_ADAPTER, &file::GROK_ADAPTER, &file::AMP_ADAPTER, @@ -664,7 +660,6 @@ static LOCAL_SOURCE_ADAPTERS: [&dyn LocalSourceAdapter; 31] = [ &vscode_tasks::KILOCODE_ADAPTER, &vscode_tasks::CLINE_ADAPTER, &antigravity::ANTIGRAVITY_ADAPTER, - &trae::TRAE_ADAPTER, &kilo::KILO_ADAPTER, &hermes::HERMES_ADAPTER, &goose::GOOSE_ADAPTER, diff --git a/crates/tokscale-core/src/adapters/trae.rs b/crates/tokscale-core/src/adapters/trae.rs deleted file mode 100644 index 8e9de8d15..000000000 --- a/crates/tokscale-core/src/adapters/trae.rs +++ /dev/null @@ -1,177 +0,0 @@ -use rayon::prelude::*; - -use crate::adapters::cache as adapter_cache; -use crate::adapters::discover as adapter_discover; -use crate::adapters::{ - AdapterScanContext, FingerprintPolicy, FoldContext, LocalSourceAdapter, MessageSink, - ParseContext, ParsedBatchSource, ParsedUnit, SourceDiscoveryError, SourceUnit, - UnitMessageSource, -}; -use crate::clients::ClientId; -use crate::sessions; - -pub(crate) struct TraeAdapter; - -impl LocalSourceAdapter for TraeAdapter { - fn client(&self) -> ClientId { - ClientId::Trae - } - - fn discover_checked( - &self, - ctx: &AdapterScanContext<'_>, - ) -> Result, SourceDiscoveryError> { - adapter_discover::discover_default_scanned_units( - ClientId::Trae, - ctx, - FingerprintPolicy::NoMessageCache, - ) - } - - fn parse_checked(&self, units: Vec, ctx: &ParseContext<'_>) -> Vec { - units - .into_par_iter() - .map(|unit| { - adapter_cache::parse_uncached_unit(unit, ctx, |path| { - sessions::trae::parse_trae_file("trae", path) - }) - }) - .collect() - } - - fn fold( - &self, - parsed: Vec, - ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, - ) -> Result<(), crate::adapters::SourcePipelineError> { - let mut messages = Vec::new(); - for unit in parsed { - ctx.health.record(unit.source_health()); - if let UnitMessageSource::Fresh(unit_messages) = unit.messages { - messages.extend(unit_messages); - } - } - sink.extend_messages(crate::dedupe_latest_trae_messages(messages)); - Ok(()) - } - - fn fold_batches( - &self, - batches: &mut ParsedBatchSource<'_>, - ctx: &mut FoldContext<'_>, - sink: &mut dyn MessageSink, - ) -> Result<(), crate::adapters::SourcePipelineError> { - let mut accumulator = crate::TraeMessageAccumulator::default(); - while let Some(parsed) = batches.next(ctx)? { - for unit in parsed { - ctx.health.record(unit.source_health()); - if let UnitMessageSource::Fresh(messages) = unit.messages { - accumulator.push_messages(messages); - } - } - } - sink.extend_messages(accumulator.finish()); - Ok(()) - } -} - -pub(crate) static TRAE_ADAPTER: TraeAdapter = TraeAdapter; - -#[cfg(test)] -mod tests { - use super::*; - use crate::adapters::FoldContext; - use crate::message_cache; - use crate::pricing::{ModelPricing, PricingService}; - - fn write_file(path: &std::path::Path, content: &str) { - std::fs::create_dir_all(path.parent().unwrap()).unwrap(); - std::fs::write(path, content).unwrap(); - } - - fn pricing_service() -> PricingService { - let mut litellm = std::collections::HashMap::new(); - litellm.insert( - "openai/gpt-5.4".to_string(), - ModelPricing { - input_cost_per_token: Some(10.0), - output_cost_per_token: Some(10.0), - ..Default::default() - }, - ); - PricingService::new(litellm, std::collections::HashMap::new()) - } - - #[test] - fn trae_adapter_dedupes_latest_session_and_applies_token_pricing() { - let dir = tempfile::TempDir::new().unwrap(); - let older = dir.path().join("older.json"); - let newer = dir.path().join("newer.json"); - write_file( - &older, - r#"[{"model_name":"GPT-5.4","session_id":"session-1","usage_time":1776000000,"dollar_float":0.1,"extra_info":{"input_token":10,"output_token":1,"cache_read_token":0,"cache_write_token":0}}]"#, - ); - write_file( - &newer, - r#"[{"model_name":"GPT-5.4","session_id":"session-1","usage_time":1776000001,"dollar_float":0.2,"extra_info":{"input_token":10,"output_token":1,"cache_read_token":0,"cache_write_token":0}}]"#, - ); - let mut cache = message_cache::SourceMessageCache::default(); - let pricing = pricing_service(); - - let units = vec![ - SourceUnit::no_message_cache(ClientId::Trae, older), - SourceUnit::no_message_cache(ClientId::Trae, newer), - ]; - let sink = rayon::ThreadPoolBuilder::new() - .num_threads(1) - .build() - .unwrap() - .install(|| { - let mut sink = Vec::new(); - let mut batches = crate::adapters::ParsedBatchSource::new(&TRAE_ADAPTER, units); - TRAE_ADAPTER - .fold_batches( - &mut batches, - &mut FoldContext::new(&mut cache, Some(&pricing)), - &mut sink, - ) - .unwrap(); - sink - }); - - assert_eq!(sink.len(), 1); - assert_eq!(sink[0].timestamp, 1_776_000_001_000); - assert_eq!(sink[0].cost, 110.0); - } - - #[test] - fn trae_adapter_keeps_valid_sessions_and_reports_rejections() { - let dir = tempfile::TempDir::new().unwrap(); - let source = dir.path().join("mixed.json"); - write_file( - &source, - r#"[ - {"model_name":"GPT-5.4","session_id":"good","usage_time":1776000000,"extra_info":{"input_token":10,"output_token":1}}, - {"model_name":"","session_id":"bad","usage_time":1776000001,"extra_info":{"input_token":10,"output_token":1}} - ]"#, - ); - let mut cache = message_cache::SourceMessageCache::default(); - let unit = SourceUnit::no_message_cache(ClientId::Trae, source); - let parsed = TRAE_ADAPTER - .parse_checked(vec![unit], &crate::adapters::ParseContext { pricing: None }); - let mut sink = Vec::new(); - let mut ctx = FoldContext::new(&mut cache, None); - - TRAE_ADAPTER.fold(parsed, &mut ctx, &mut sink).unwrap(); - - assert_eq!(sink.len(), 1); - assert_eq!(sink[0].session_id.as_ref(), "good"); - assert_eq!(ctx.health.rejected_records(), 1); - let source = &ctx.health.sources()[0]; - assert_eq!(source.client, ClientId::Trae); - let rejection = source.rejections.entries().next().unwrap(); - assert_eq!(rejection.key, "missing-model"); - assert_eq!(rejection.count, 1); - } -} diff --git a/crates/tokscale-core/src/aggregate/tui.rs b/crates/tokscale-core/src/aggregate/tui.rs index c00237fb8..b94c5cc6a 100644 --- a/crates/tokscale-core/src/aggregate/tui.rs +++ b/crates/tokscale-core/src/aggregate/tui.rs @@ -1677,7 +1677,7 @@ mod tests { 1.0, ), UnifiedMessage::new( - "cursor", + "gemini", "claude-sonnet-4.5", "anthropic", "session-2", @@ -1706,12 +1706,12 @@ mod tests { assert_eq!(claude_model.display_name, "claude-sonnet-4.5"); assert_eq!(claude_model.tokens.total(), 15); - let cursor = usage.daily[0].source_breakdown.get("cursor").unwrap(); - assert_eq!(cursor.cost, 2.0); - assert_eq!(cursor.models.len(), 1); - let cursor_model = cursor.models.get("v1|m|17:claude-sonnet-4.5").unwrap(); - assert_eq!(cursor_model.display_name, "claude-sonnet-4.5"); - assert_eq!(cursor_model.tokens.total(), 30); + let gemini = usage.daily[0].source_breakdown.get("gemini").unwrap(); + assert_eq!(gemini.cost, 2.0); + assert_eq!(gemini.models.len(), 1); + let gemini_model = gemini.models.get("v1|m|17:claude-sonnet-4.5").unwrap(); + assert_eq!(gemini_model.display_name, "claude-sonnet-4.5"); + assert_eq!(gemini_model.tokens.total(), 30); } #[test] diff --git a/crates/tokscale-core/src/clients.rs b/crates/tokscale-core/src/clients.rs index 151c6c8c0..d7648163a 100644 --- a/crates/tokscale-core/src/clients.rs +++ b/crates/tokscale-core/src/clients.rs @@ -72,6 +72,12 @@ mod tests { assert_eq!(ClientId::from_str("antigravity-cli"), None); } + #[test] + fn removed_clients_are_not_client_identities() { + assert_eq!(ClientId::from_str("cursor"), None); + assert_eq!(ClientId::from_str("trae"), None); + } + #[test] fn pi_and_omp_have_separate_identity_facts() { assert_eq!(ClientId::Pi.as_str(), "pi"); diff --git a/crates/tokscale-core/src/lib.rs b/crates/tokscale-core/src/lib.rs index f0b80e22e..d7c826b2b 100644 --- a/crates/tokscale-core/src/lib.rs +++ b/crates/tokscale-core/src/lib.rs @@ -6,6 +6,7 @@ mod client_catalog; pub mod clients; pub mod fs_atomic; mod local_clients; +mod local_report_error; mod message_cache; mod model_aliases; pub mod paths; @@ -31,6 +32,7 @@ pub use clients::{ warp_sqlite_roots_with_env_strategy, ClientCounts, ClientId, ClientIdentity, LocalClientDef, PathRoot, }; +pub use local_report_error::{LocalReportError, LocalReportErrorKind}; pub use message_cache::{prune_source_message_cache, SourceCachePruneError, SourceCachePruneStats}; pub use provider_identity::{inferred_provider_from_model, normalize_provider_for_grouping}; pub use sessionize::{ @@ -593,7 +595,7 @@ fn parse_all_messages_with_pricing( home_dir: &str, clients: &[String], pricing: Option<&pricing::PricingService>, -) -> Result, String> { +) -> Result, LocalReportError> { parse_all_messages_with_pricing_with_env_strategy( home_dir, clients, @@ -610,7 +612,7 @@ fn parse_all_messages_with_pricing_with_env_strategy( pricing: Option<&pricing::PricingService>, use_env_roots: bool, scanner_settings: &scanner::ScannerSettings, -) -> Result, String> { +) -> Result, LocalReportError> { parse_all_messages_with_health_with_env_strategy( home_dir, clients, @@ -626,7 +628,7 @@ fn parse_all_messages_with_health( home_dir: &str, clients: &[String], pricing: Option<&pricing::PricingService>, -) -> Result<(Vec, DataHealth), String> { +) -> Result<(Vec, DataHealth), LocalReportError> { parse_all_messages_with_health_with_env_strategy( home_dir, clients, @@ -643,7 +645,7 @@ fn parse_all_messages_with_health_with_env_strategy( pricing: Option<&pricing::PricingService>, use_env_roots: bool, scanner_settings: &scanner::ScannerSettings, -) -> Result<(Vec, DataHealth), String> { +) -> Result<(Vec, DataHealth), LocalReportError> { let prepared = prepare_local_sources(LocalParseOptions { home_dir: Some(home_dir.to_string()), use_env_roots, @@ -661,7 +663,7 @@ fn fold_prepared_local_sources_with_pricing( prepared: PreparedLocalSources, pricing: Option<&pricing::PricingService>, sink: &mut dyn adapters::MessageSink, -) -> Result<(SourceInventorySignature, DataHealth), String> { +) -> Result<(SourceInventorySignature, DataHealth), LocalReportError> { let PreparedLocalSources { clients, groups, @@ -670,7 +672,7 @@ fn fold_prepared_local_sources_with_pricing( } = prepared; let mut source_cache = message_cache::SourceMessageCache::load() .map_err(adapters::SourcePipelineError::from) - .map_err(|error| error.to_string())?; + .map_err(LocalReportError::operational)?; let parse_result = if clients.is_empty() { adapters::run_prepared_local_source_adapters( @@ -711,7 +713,7 @@ fn fold_prepared_local_sources_with_pricing( health, ) }) - .map_err(|error| error.to_string()) + .map_err(LocalReportError::operational) } struct RequestedClientFilterSink<'a> { @@ -777,18 +779,21 @@ fn stream_local_sources_into_engine( prepared: PreparedLocalSources, pricing: Option<&pricing::PricingService>, engine: &mut crate::aggregate::AggregationEngine, -) -> Result<(SourceInventorySignature, DataHealth), String> { +) -> Result<(SourceInventorySignature, DataHealth), LocalReportError> { let mut sink = AggregationSink(engine); fold_prepared_local_sources_with_pricing(prepared, pricing, &mut sink) } -pub fn prepare_local_sources(options: LocalParseOptions) -> Result { +pub fn prepare_local_sources( + options: LocalParseOptions, +) -> Result { let (home_dir, clients) = resolve_local_parse_request(&options)?; options .scanner_settings .validate() - .map_err(|error| error.to_string())?; - let selected_adapters = adapters::selected_adapters(&clients)?; + .map_err(LocalReportError::invalid_environment)?; + let selected_adapters = + adapters::selected_adapters(&clients).map_err(LocalReportError::invalid_request_message)?; let scan_ctx = adapters::AdapterScanContext { home_dir: &home_dir, use_env_roots: options.use_env_roots, @@ -797,7 +802,7 @@ pub fn prepare_local_sources(options: LocalParseOptions) -> Result = selected_adapters .into_iter() - .map(|adapter| -> Result<_, String> { + .map(|adapter| -> Result<_, LocalReportError> { #[cfg(test)] PREPARE_DISCOVERY_COUNT.with(|count| count.set(count.get() + 1)); // Third-party source and snapshot failures stay inside their @@ -808,7 +813,7 @@ pub fn prepare_local_sources(options: LocalParseOptions) -> Result { - return Err(error.to_string()); + return Err(LocalReportError::invalid_environment(error)); } Err(error) => { health.record(SourceHealth { @@ -933,53 +938,6 @@ fn confirmed_source_inventory_signature( SourceInventorySignature(hasher.finalize().into()) } -#[derive(Default)] -struct TraeMessageAccumulator { - latest_by_session: HashMap, UnifiedMessage>, -} - -impl TraeMessageAccumulator { - fn push_messages(&mut self, messages: Vec) { - for message in messages { - let session_id = std::sync::Arc::clone(&message.session_id); - match self.latest_by_session.get_mut(&session_id) { - Some(existing) => { - let should_replace = message.timestamp > existing.timestamp - || (message.timestamp == existing.timestamp - && message.dedup_key.as_ref().is_some_and(|key| { - existing - .dedup_key - .as_ref() - .is_none_or(|existing_key| key > existing_key) - })); - if should_replace { - *existing = message; - } - } - None => { - let _ = self.latest_by_session.insert(session_id, message); - } - } - } - } - - fn finish(self) -> Vec { - let mut deduped: Vec = self.latest_by_session.into_values().collect(); - deduped.sort_unstable_by(|a, b| { - a.session_id - .cmp(&b.session_id) - .then_with(|| a.timestamp.cmp(&b.timestamp)) - }); - deduped - } -} - -fn dedupe_latest_trae_messages(messages: Vec) -> Vec { - let mut accumulator = TraeMessageAccumulator::default(); - accumulator.push_messages(messages); - accumulator.finish() -} - /// Date-range retain shared by the report and local-parse filters. One /// `date_string()` per message, only when a date filter is active. fn retain_messages_in_date_range( @@ -1071,14 +1029,15 @@ fn normalize_token_breakdown(tokens: &mut TokenBreakdown) { tokens.reasoning = tokens.reasoning.max(0); } -fn resolve_report_request(options: &ReportOptions) -> Result<(String, Vec), String> { - let home_dir = get_home_dir_string(&options.home_dir)?; - let clients = options.clients.clone().unwrap_or_else(|| { - ClientId::iter() - .filter(|client| client.parse_local()) - .map(|c| c.as_str().to_string()) - .collect() - }); +fn resolve_report_request( + options: &ReportOptions, +) -> Result<(String, Vec), LocalReportError> { + let home_dir = get_home_dir_string(&options.home_dir) + .map_err(LocalReportError::invalid_environment_message)?; + let clients = options + .clients + .clone() + .unwrap_or_else(|| ClientId::iter().map(|c| c.as_str().to_string()).collect()); Ok((home_dir, clients)) } @@ -1095,7 +1054,7 @@ struct ResolvedAggregationRequest<'a> { fn load_aggregated_views_resolved( request: ResolvedAggregationRequest<'_>, -) -> Result { +) -> Result { let prepared = prepare_local_sources(LocalParseOptions { home_dir: Some(request.home_dir.to_string()), use_env_roots: request.use_env_roots, @@ -1121,7 +1080,7 @@ fn load_prepared_aggregated_views( date_range: DateRange, views: ViewSet, pricing: Option<&pricing::PricingService>, -) -> Result<(AggregatedViews, SourceInventorySignature), String> { +) -> Result<(AggregatedViews, SourceInventorySignature), LocalReportError> { let mut engine = crate::aggregate::AggregationEngine::new(AggregationConfig { group_by, date_range, @@ -1152,7 +1111,7 @@ fn load_aggregated_views_for_resolved_report( clients: &[String], views: ViewSet, pricing: Option<&pricing::PricingService>, -) -> Result { +) -> Result { load_aggregated_views_resolved(ResolvedAggregationRequest { home_dir, clients, @@ -1176,12 +1135,12 @@ pub fn load_aggregated_views_with_pricing( options: &ReportOptions, views: ViewSet, pricing: Option<&pricing::PricingService>, -) -> Result { +) -> Result { let (home_dir, clients) = resolve_report_request(options)?; load_aggregated_views_for_resolved_report(options, &home_dir, &clients, views, pricing) } -pub async fn get_model_report(options: ReportOptions) -> Result { +pub async fn get_model_report(options: ReportOptions) -> Result { let start = Instant::now(); let (home_dir, clients) = resolve_report_request(&options)?; let pricing = load_pricing_for_local_parse().await; @@ -1198,7 +1157,7 @@ pub async fn get_model_report(options: ReportOptions) -> Result Result { +pub async fn get_monthly_report(options: ReportOptions) -> Result { let start = Instant::now(); let (home_dir, clients) = resolve_report_request(&options)?; let pricing = load_pricing_for_local_parse().await; @@ -1219,7 +1178,7 @@ pub async fn get_monthly_report(options: ReportOptions) -> Result Result { +pub async fn get_hourly_report(options: ReportOptions) -> Result { let start = Instant::now(); let (home_dir, clients) = resolve_report_request(&options)?; let pricing = load_pricing_for_local_parse().await; @@ -1239,7 +1198,7 @@ pub async fn get_hourly_report(options: ReportOptions) -> Result, -) -> Result { +) -> Result { let start = Instant::now(); let (home_dir, clients) = resolve_report_request(&options)?; let mut views = load_aggregated_views_for_resolved_report( @@ -1262,7 +1221,9 @@ pub struct TimeMetricsReport { pub health: source_health::HealthReport, } -pub async fn get_time_metrics_report(options: ReportOptions) -> Result { +pub async fn get_time_metrics_report( + options: ReportOptions, +) -> Result { let start = Instant::now(); let (home_dir, clients) = resolve_report_request(&options)?; let views = load_aggregated_views_for_resolved_report( @@ -1278,12 +1239,14 @@ pub async fn get_time_metrics_report(options: ReportOptions) -> Result Result { +pub async fn generate_graph(options: ReportOptions) -> Result { let pricing = pricing::PricingService::get_or_init().await?; generate_graph_with_loaded_pricing(options, Some(&pricing)).await } -pub async fn generate_local_graph_report(options: ReportOptions) -> Result { +pub async fn generate_local_graph_report( + options: ReportOptions, +) -> Result { let pricing = load_pricing_for_local_parse().await; generate_graph_with_loaded_pricing(options, pricing.as_deref()).await } @@ -1494,16 +1457,17 @@ async fn load_pricing_for_local_parse_with_diagnostics( fn resolve_local_parse_request( options: &LocalParseOptions, -) -> Result<(String, Vec), String> { - let home_dir = get_home_dir_string(&options.home_dir)?; - let clients = options.clients.clone().unwrap_or_else(|| { - ClientId::iter() - .filter(|c| c.parse_local()) - .map(|c| c.as_str().to_string()) - .collect() - }); +) -> Result<(String, Vec), LocalReportError> { + let home_dir = get_home_dir_string(&options.home_dir) + .map_err(LocalReportError::invalid_environment_message)?; + let clients = options + .clients + .clone() + .unwrap_or_else(|| ClientId::iter().map(|c| c.as_str().to_string()).collect()); for client in &clients { - ClientId::from_str(client).ok_or_else(|| format!("unknown local client `{client}`"))?; + ClientId::from_str(client).ok_or_else(|| { + LocalReportError::invalid_request_message(format!("unknown local client `{client}`")) + })?; } Ok((home_dir, clients)) } @@ -1511,7 +1475,7 @@ fn resolve_local_parse_request( fn parse_prepared_local_unified_messages( prepared: PreparedLocalSources, pricing: Option<&pricing::PricingService>, -) -> Result>, String> { +) -> Result>, LocalReportError> { let filters = prepared.options.clone(); let mut messages = Vec::new(); let (source_inventory_signature, health) = @@ -1527,7 +1491,7 @@ fn parse_prepared_local_unified_messages( #[doc(hidden)] pub fn count_local_client_messages( options: LocalParseOptions, -) -> Result { +) -> Result { let start = Instant::now(); let prepared = prepare_local_sources(options)?; let mut sink = ClientCountSink::new(DateRange { @@ -1548,14 +1512,14 @@ pub fn count_local_client_messages( pub async fn parse_local_unified_messages_with_pricing( options: LocalParseOptions, pricing: Option<&pricing::PricingService>, -) -> Result>, String> { +) -> Result>, LocalReportError> { let prepared = prepare_local_sources(options)?; parse_prepared_local_unified_messages(prepared, pricing) } pub async fn parse_local_unified_messages( options: LocalParseOptions, -) -> Result>, String> { +) -> Result>, LocalReportError> { let prepared = prepare_local_sources(options)?; let pricing = load_pricing_for_local_parse().await; parse_prepared_local_unified_messages(prepared, pricing.as_deref()) @@ -1566,7 +1530,7 @@ pub fn load_usage_data_with_pricing( options: LocalParseOptions, group_by: GroupBy, pricing: Option<&pricing::PricingService>, -) -> Result { +) -> Result { let prepared = prepare_local_sources(options)?; load_prepared_usage_data_with_pricing(prepared, group_by, pricing) } @@ -1576,7 +1540,7 @@ pub fn load_prepared_usage_data_with_pricing( prepared: PreparedLocalSources, group_by: GroupBy, pricing: Option<&pricing::PricingService>, -) -> Result { +) -> Result { let date_range = DateRange { since: prepared.options.since.clone(), until: prepared.options.until.clone(), @@ -1600,7 +1564,7 @@ pub struct UsageDataWithDiagnostics { pub async fn load_usage_data_with_diagnostics( options: LocalParseOptions, group_by: GroupBy, -) -> Result { +) -> Result { let prepared = prepare_local_sources(options)?; load_prepared_usage_data_with_diagnostics(prepared, group_by).await } @@ -1608,7 +1572,7 @@ pub async fn load_usage_data_with_diagnostics( pub async fn load_prepared_usage_data_with_diagnostics( prepared: PreparedLocalSources, group_by: GroupBy, -) -> Result { +) -> Result { let mut pricing_diagnostics = pricing::PricingDiagnostics::new(); let pricing = load_pricing_for_local_parse_with_diagnostics(&mut pricing_diagnostics).await; let date_range = DateRange { @@ -1636,7 +1600,7 @@ pub async fn load_prepared_usage_data_with_diagnostics( pub async fn load_usage_data( options: LocalParseOptions, group_by: GroupBy, -) -> Result { +) -> Result { let prepared = prepare_local_sources(options)?; let pricing = load_pricing_for_local_parse().await; load_prepared_usage_data_with_pricing(prepared, group_by, pricing.as_deref()) diff --git a/crates/tokscale-core/src/lib_tests.rs b/crates/tokscale-core/src/lib_tests.rs index 9d6677fa3..5f962545c 100644 --- a/crates/tokscale-core/src/lib_tests.rs +++ b/crates/tokscale-core/src/lib_tests.rs @@ -1,15 +1,14 @@ use super::{ - aggregate_model_usage_entries, apply_token_pricing, dedupe_latest_trae_messages, - finalize_token_priced_messages, generate_graph_with_loaded_pricing, - load_aggregated_views_with_pricing, load_cache_only_pricing_with_diagnostics, - load_usage_data_with_pricing, message_cache, normalize_model_for_grouping, - parse_all_messages_with_health, parse_all_messages_with_health_with_env_strategy, - parse_all_messages_with_pricing, parse_all_messages_with_pricing_with_env_strategy, - positive_token_total, pricing, retain_for_requested_clients, scanner, - select_local_parse_pricing, AggregatedViews, AggregationConfig, ClientContribution, - ClientCounts, ClientId, DailyTotals, DateRange, GraphResult, GroupBy, LocalParseOptions, - ReportOptions, SessionContribution, TimeMetricsReport, TokenBreakdown, UnifiedMessage, ViewSet, - UNKNOWN_WORKSPACE_LABEL, + aggregate_model_usage_entries, apply_token_pricing, finalize_token_priced_messages, + generate_graph_with_loaded_pricing, load_aggregated_views_with_pricing, + load_cache_only_pricing_with_diagnostics, load_usage_data_with_pricing, message_cache, + normalize_model_for_grouping, parse_all_messages_with_health, + parse_all_messages_with_health_with_env_strategy, parse_all_messages_with_pricing, + parse_all_messages_with_pricing_with_env_strategy, positive_token_total, pricing, + retain_for_requested_clients, scanner, select_local_parse_pricing, AggregatedViews, + AggregationConfig, ClientContribution, ClientCounts, ClientId, DailyTotals, DateRange, + GraphResult, GroupBy, LocalParseOptions, ReportOptions, SessionContribution, TimeMetricsReport, + TokenBreakdown, UnifiedMessage, ViewSet, UNKNOWN_WORKSPACE_LABEL, }; use std::collections::{BTreeMap, HashMap, HashSet}; use std::ffi::OsString; @@ -27,7 +26,7 @@ struct LocalMessagesForTest { fn load_local_messages_for_test( options: LocalParseOptions, -) -> Result { +) -> Result { let counts = super::count_local_client_messages(options.clone())?.counts; let prepared = super::prepare_local_sources(options.clone())?; let mut messages = Vec::new(); @@ -165,30 +164,6 @@ fn test_session_contribution_serde_round_trip() { assert!(json.contains("\"session_id\":\"019e1e27")); } -fn make_trae_message( - session_id: &str, - timestamp: i64, - dedup_key: Option<&str>, - cost: f64, -) -> UnifiedMessage { - UnifiedMessage::new_with_dedup( - "trae", - "gpt-5.2", - "openai", - session_id, - timestamp, - TokenBreakdown { - input: 10, - output: 5, - cache_read: 0, - cache_write: 0, - reasoning: 0, - }, - cost, - dedup_key.map(crate::sessions::dedup_hash_str), - ) -} - #[allow(clippy::too_many_arguments)] fn make_message_with_tokens( client: &str, @@ -1941,35 +1916,6 @@ fn test_retain_for_requested_clients_preserves_kilo_split() { )); } -#[test] -#[serial_test::serial] -fn test_cursor_parse_path_keeps_zero_cost_for_unpriced_composer_rows() { - let temp_dir = tempfile::TempDir::new().unwrap(); - let cursor_cache_dir = temp_dir.path().join(".config/tokscale/cursor-cache"); - std::fs::create_dir_all(&cursor_cache_dir).unwrap(); - - let csv = r#"Date,Kind,Model,Max Mode,Input (w/ Cache Write),Input (w/o Cache Write),Cache Read,Output Tokens,Total Tokens,Cost -"2026-03-04T12:00:00.000Z","Included","Composer 1.5","No","1200","1000","5000","2000","8000","0""#; - std::fs::write(cursor_cache_dir.join("usage.csv"), csv).unwrap(); - - let pricing = pricing::PricingService::new(HashMap::new(), HashMap::new()); - let messages = parse_all_messages_with_pricing( - temp_dir.path().to_str().unwrap(), - &["cursor".to_string()], - Some(&pricing), - ) - .unwrap(); - - assert_eq!(messages.len(), 1); - assert_eq!(messages[0].client.as_ref(), "cursor"); - assert_eq!(messages[0].model_id.as_ref(), "composer 1.5"); - assert_eq!(messages[0].tokens.input, 1000); - assert_eq!(messages[0].tokens.output, 2000); - assert_eq!(messages[0].tokens.cache_read, 5000); - assert_eq!(messages[0].tokens.cache_write, 200); - assert_eq!(messages[0].cost, 0.0); -} - fn write_kimi_code_usage_fixture(source_home: &std::path::Path) { let kimi_home = source_home.join(".kimi-code"); std::fs::create_dir_all(&kimi_home).unwrap(); @@ -2172,8 +2118,13 @@ fn prepare_local_sources_rejects_invalid_extra_dirs_configuration() { .err() .expect("invalid extra-dir syntax must fail source preparation"); - assert!(error.contains("TOKSCALE_EXTRA_DIRS")); - assert!(error.contains("parse environment variable")); + assert_eq!( + error.kind(), + super::LocalReportErrorKind::InvalidEnvironment + ); + let message = error.to_string(); + assert!(message.contains("TOKSCALE_EXTRA_DIRS")); + assert!(message.contains("parse environment variable")); } #[cfg(unix)] @@ -2192,8 +2143,13 @@ fn prepare_local_sources_rejects_non_utf8_extra_dirs_configuration() { .err() .expect("non-UTF-8 extra-dir configuration must fail source preparation"); - assert!(error.contains("TOKSCALE_EXTRA_DIRS")); - assert!(error.contains("read environment variable")); + assert_eq!( + error.kind(), + super::LocalReportErrorKind::InvalidEnvironment + ); + let message = error.to_string(); + assert!(message.contains("TOKSCALE_EXTRA_DIRS")); + assert!(message.contains("read environment variable")); } #[test] @@ -4152,16 +4108,11 @@ fn test_source_cache_does_not_reuse_priced_cost_without_pricing_service() { let original_home = std::env::var("HOME").ok(); std::env::set_var("HOME", temp_home.path()); { - let cursor_cache_dir = source_home.path().join(".config/tokscale/cursor-cache"); - std::fs::create_dir_all(&cursor_cache_dir).unwrap(); - - let csv = r#"Date,Kind,Model,Max Mode,Input (w/ Cache Write),Input (w/o Cache Write),Cache Read,Output Tokens,Total Tokens,Cost -"2026-03-04T12:00:00.000Z","Included","Composer 1.5","No","1200","1000","5000","2000","8000","0""#; - std::fs::write(cursor_cache_dir.join("usage.csv"), csv).unwrap(); + write_kimi_code_usage_fixture(source_home.path()); let mut litellm = HashMap::new(); litellm.insert( - "composer 1.5".into(), + "gpt-5.5".into(), pricing::ModelPricing { input_cost_per_token: Some(0.001), output_cost_per_token: Some(0.002), @@ -4173,22 +4124,22 @@ fn test_source_cache_does_not_reuse_priced_cost_without_pricing_service() { let repriced_messages = parse_all_messages_with_pricing( source_home.path().to_str().unwrap(), - &["cursor".to_string()], + &["kimi".to_string()], Some(&pricing), ) .unwrap(); - assert_eq!(repriced_messages.len(), 1); - assert!(repriced_messages[0].cost > 0.0); + assert_eq!(repriced_messages.len(), 2); + assert!(repriced_messages.iter().all(|message| message.cost > 0.0)); let cached_messages = parse_all_messages_with_pricing( source_home.path().to_str().unwrap(), - &["cursor".to_string()], + &["kimi".to_string()], None, ) .unwrap(); - assert_eq!(cached_messages.len(), 1); - assert_eq!(cached_messages[0].cost, 0.0); + assert_eq!(cached_messages.len(), 2); + assert!(cached_messages.iter().all(|message| message.cost == 0.0)); } match original_home { @@ -5596,88 +5547,6 @@ fn test_select_local_parse_pricing_does_not_evaluate_stale_fallback_on_fresh_suc assert!(!stale_called); } -#[test] -fn test_dedupe_latest_trae_messages_keeps_latest_timestamp_for_session() { - let messages = vec![ - make_trae_message( - "session-stable", - 1_700_000_002_000, - Some("trae:session-stable:1_700_000_002"), - 0.2, - ), - make_trae_message( - "session-stable", - 1_700_000_003_000, - Some("trae:session-stable:1_700_000_003"), - 0.3, - ), - make_trae_message( - "session-other", - 1_700_000_001_000, - Some("trae:session-other:1_700_000_001"), - 0.1, - ), - ]; - - let deduped = dedupe_latest_trae_messages(messages); - - assert_eq!(deduped.len(), 2); - let stable = deduped - .iter() - .find(|msg| msg.session_id.as_ref() == "session-stable") - .expect("session-stable should remain after dedupe"); - assert_eq!(stable.timestamp, 1_700_000_003_000); - assert_eq!(stable.cost, 0.3); - assert_eq!( - stable.dedup_key, - Some(crate::sessions::dedup_hash_str( - "trae:session-stable:1_700_000_003" - )) - ); -} - -#[test] -fn test_dedupe_latest_trae_messages_tiebreaks_by_dedup_key() { - let messages = vec![ - make_trae_message( - "session-stable", - 1_700_000_010_000, - Some("dedupe-key-a"), - 0.2, - ), - make_trae_message( - "session-stable", - 1_700_000_010_000, - Some("dedupe-key-z"), - 0.4, - ), - make_trae_message( - "session-stable", - 1_700_000_009_000, - Some("dedupe-key-m"), - 0.1, - ), - ]; - - let deduped = dedupe_latest_trae_messages(messages); - - // Equal timestamps tiebreak on the greater dedup hash: arbitrary but - // stable across runs and machines (FNV-1a). Real trae keys embed - // usage_time = timestamp/1000, so production ties carry equal keys. - let key_a = crate::sessions::dedup_hash_str("dedupe-key-a"); - let key_z = crate::sessions::dedup_hash_str("dedupe-key-z"); - let (winning_key, winning_cost) = if key_z > key_a { - (key_z, 0.4) - } else { - (key_a, 0.2) - }; - - assert_eq!(deduped.len(), 1); - assert_eq!(deduped[0].timestamp, 1_700_000_010_000); - assert_eq!(deduped[0].dedup_key, Some(winning_key)); - assert_eq!(deduped[0].cost, winning_cost); -} - #[test] #[serial_test::serial] fn test_parse_all_messages_with_pricing_keeps_gateway_message_under_real_client_filter() { @@ -6685,37 +6554,3 @@ fn test_local_message_loader_amp_reads_current_thread_files() { assert_eq!(parsed.messages[0].tokens.input, 10); assert_eq!(parsed.messages[0].tokens.output, 2); } - -#[test] -#[serial_test::serial] -fn test_local_message_loader_default_keeps_cursor_out_of_local_count() { - let temp_dir = tempfile::TempDir::new().unwrap(); - let cursor_cache_dir = temp_dir.path().join(".config/tokscale/cursor-cache"); - std::fs::create_dir_all(&cursor_cache_dir).unwrap(); - std::fs::write( - cursor_cache_dir.join("usage.csv"), - r#"Date,Kind,Model,Max Mode,Input (w/ Cache Write),Input (w/o Cache Write),Cache Read,Output Tokens,Total Tokens,Cost -"2026-03-04T12:00:00.000Z","Included","Composer 1.5","No","1200","1000","5000","2000","8000","0""#, - ) - .unwrap(); - - let parsed = load_local_messages_for_test(LocalParseOptions { - home_dir: Some(temp_dir.path().to_str().unwrap().to_string()), - use_env_roots: false, - clients: None, - since: None, - until: None, - year: None, - scanner_settings: scanner::ScannerSettings::default(), - }) - .unwrap(); - - assert_eq!(parsed.counts.get(ClientId::Cursor), 0); - assert!( - parsed - .messages - .iter() - .all(|message| message.client.as_ref() != "cursor"), - "Cursor cache rows must not enter the default local message loading result" - ); -} diff --git a/crates/tokscale-core/src/local_clients.rs b/crates/tokscale-core/src/local_clients.rs index d2dfb3218..65ca73584 100644 --- a/crates/tokscale-core/src/local_clients.rs +++ b/crates/tokscale-core/src/local_clients.rs @@ -71,7 +71,6 @@ pub struct LocalClientDef { pub relative_path: &'static str, pub pattern: &'static str, pub headless: bool, - pub parse_local: bool, } impl LocalClientDef { @@ -100,7 +99,6 @@ pub const LOCAL_CLIENTS: &[LocalClientEntry] = &[ relative_path: "opencode", pattern: "*.db", headless: false, - parse_local: true, }, }, LocalClientEntry { @@ -110,7 +108,6 @@ pub const LOCAL_CLIENTS: &[LocalClientEntry] = &[ relative_path: ".claude/projects", pattern: "*.jsonl", headless: false, - parse_local: true, }, }, LocalClientEntry { @@ -123,17 +120,6 @@ pub const LOCAL_CLIENTS: &[LocalClientEntry] = &[ relative_path: "sessions", pattern: "*.jsonl", headless: true, - parse_local: true, - }, - }, - LocalClientEntry { - client: ClientId::Cursor, - def: LocalClientDef { - root: PathRoot::Home, - relative_path: ".config/tokscale/cursor-cache", - pattern: "usage*.csv", - headless: false, - parse_local: false, }, }, LocalClientEntry { @@ -146,7 +132,6 @@ pub const LOCAL_CLIENTS: &[LocalClientEntry] = &[ relative_path: "tmp", pattern: "*.json|*.jsonl", headless: false, - parse_local: true, }, }, LocalClientEntry { @@ -156,7 +141,6 @@ pub const LOCAL_CLIENTS: &[LocalClientEntry] = &[ relative_path: "amp/threads", pattern: "T-*.json", headless: false, - parse_local: true, }, }, LocalClientEntry { @@ -166,7 +150,6 @@ pub const LOCAL_CLIENTS: &[LocalClientEntry] = &[ relative_path: ".factory/sessions", pattern: "*.settings.json", headless: false, - parse_local: true, }, }, LocalClientEntry { @@ -176,7 +159,6 @@ pub const LOCAL_CLIENTS: &[LocalClientEntry] = &[ relative_path: ".openclaw/agents", pattern: "*.jsonl*", headless: false, - parse_local: true, }, }, LocalClientEntry { @@ -186,7 +168,6 @@ pub const LOCAL_CLIENTS: &[LocalClientEntry] = &[ relative_path: ".pi/agent/sessions", pattern: "*.jsonl", headless: false, - parse_local: true, }, }, LocalClientEntry { @@ -196,7 +177,6 @@ pub const LOCAL_CLIENTS: &[LocalClientEntry] = &[ relative_path: ".omp/agent/sessions", pattern: "*.jsonl", headless: false, - parse_local: true, }, }, LocalClientEntry { @@ -209,7 +189,6 @@ pub const LOCAL_CLIENTS: &[LocalClientEntry] = &[ relative_path: "sessions", pattern: "wire.jsonl", headless: false, - parse_local: true, }, }, LocalClientEntry { @@ -219,7 +198,6 @@ pub const LOCAL_CLIENTS: &[LocalClientEntry] = &[ relative_path: ".qwen/projects", pattern: "*.jsonl", headless: false, - parse_local: true, }, }, LocalClientEntry { @@ -229,7 +207,6 @@ pub const LOCAL_CLIENTS: &[LocalClientEntry] = &[ relative_path: ".config/Code/User/globalStorage/rooveterinaryinc.roo-cline/tasks", pattern: "ui_messages.json", headless: false, - parse_local: true, }, }, LocalClientEntry { @@ -239,7 +216,6 @@ pub const LOCAL_CLIENTS: &[LocalClientEntry] = &[ relative_path: ".config/Code/User/globalStorage/kilocode.kilo-code/tasks", pattern: "ui_messages.json", headless: false, - parse_local: true, }, }, LocalClientEntry { @@ -249,7 +225,6 @@ pub const LOCAL_CLIENTS: &[LocalClientEntry] = &[ relative_path: ".mux/sessions", pattern: "session-usage.json", headless: false, - parse_local: true, }, }, LocalClientEntry { @@ -259,7 +234,6 @@ pub const LOCAL_CLIENTS: &[LocalClientEntry] = &[ relative_path: "kilo/kilo.db", pattern: "kilo.db", headless: false, - parse_local: true, }, }, LocalClientEntry { @@ -272,7 +246,6 @@ pub const LOCAL_CLIENTS: &[LocalClientEntry] = &[ relative_path: "state.db", pattern: "state.db", headless: false, - parse_local: true, }, }, LocalClientEntry { @@ -282,7 +255,6 @@ pub const LOCAL_CLIENTS: &[LocalClientEntry] = &[ relative_path: ".copilot/otel", pattern: "*.jsonl", headless: false, - parse_local: true, }, }, LocalClientEntry { @@ -292,7 +264,6 @@ pub const LOCAL_CLIENTS: &[LocalClientEntry] = &[ relative_path: "goose/sessions/sessions.db", pattern: "sessions.db", headless: false, - parse_local: true, }, }, LocalClientEntry { @@ -305,7 +276,6 @@ pub const LOCAL_CLIENTS: &[LocalClientEntry] = &[ relative_path: "projects", pattern: "chat-messages.json", headless: false, - parse_local: true, }, }, LocalClientEntry { @@ -315,7 +285,6 @@ pub const LOCAL_CLIENTS: &[LocalClientEntry] = &[ relative_path: ".codebuddy/projects", pattern: "*.jsonl", headless: false, - parse_local: true, }, }, LocalClientEntry { @@ -325,7 +294,6 @@ pub const LOCAL_CLIENTS: &[LocalClientEntry] = &[ relative_path: "antigravity-cache/sessions", pattern: "*.jsonl", headless: false, - parse_local: true, }, }, LocalClientEntry { @@ -335,7 +303,6 @@ pub const LOCAL_CLIENTS: &[LocalClientEntry] = &[ relative_path: "zed/threads/threads.db", pattern: "threads.db", headless: false, - parse_local: true, }, }, LocalClientEntry { @@ -345,7 +312,6 @@ pub const LOCAL_CLIENTS: &[LocalClientEntry] = &[ relative_path: ".zcode/projects", pattern: "*.jsonl", headless: false, - parse_local: true, }, }, LocalClientEntry { @@ -355,7 +321,6 @@ pub const LOCAL_CLIENTS: &[LocalClientEntry] = &[ relative_path: ".kiro/sessions/cli", pattern: "*.json", headless: false, - parse_local: true, }, }, LocalClientEntry { @@ -365,17 +330,6 @@ pub const LOCAL_CLIENTS: &[LocalClientEntry] = &[ relative_path: ".junie/sessions", pattern: "events.jsonl", headless: false, - parse_local: true, - }, - }, - LocalClientEntry { - client: ClientId::Trae, - def: LocalClientDef { - root: PathRoot::Config, - relative_path: "trae-cache/sessions", - pattern: "*.json", - headless: false, - parse_local: true, }, }, LocalClientEntry { @@ -385,7 +339,6 @@ pub const LOCAL_CLIENTS: &[LocalClientEntry] = &[ relative_path: ".local/state/warp-terminal", pattern: "warp.sqlite", headless: false, - parse_local: true, }, }, LocalClientEntry { @@ -395,7 +348,6 @@ pub const LOCAL_CLIENTS: &[LocalClientEntry] = &[ relative_path: ".config/Code/User/globalStorage/saoudrizwan.claude-dev/tasks", pattern: "ui_messages.json", headless: false, - parse_local: true, }, }, LocalClientEntry { @@ -405,7 +357,6 @@ pub const LOCAL_CLIENTS: &[LocalClientEntry] = &[ relative_path: ".commandcode/projects", pattern: "*.jsonl", headless: false, - parse_local: true, }, }, LocalClientEntry { @@ -418,7 +369,6 @@ pub const LOCAL_CLIENTS: &[LocalClientEntry] = &[ relative_path: "sessions", pattern: "updates.jsonl", headless: false, - parse_local: true, }, }, ]; @@ -504,10 +454,6 @@ impl ClientId { pub fn supports_headless(self) -> bool { self.local_def().is_some_and(|def| def.headless) } - - pub fn parse_local(self) -> bool { - self.local_def().is_some_and(|def| def.parse_local) - } } #[cfg(test)] @@ -537,22 +483,10 @@ mod tests { assert_eq!(keyed, catalog); } - #[test] - fn cursor_is_registered_but_not_locally_parsed() { - let def = ClientId::Cursor - .local_def() - .expect("cursor has cache scan policy"); - assert_eq!(def.relative_path, ".config/tokscale/cursor-cache"); - assert_eq!(def.pattern, "usage*.csv"); - assert!(!ClientId::Cursor.parse_local()); - assert!(crate::adapters::adapter_for(ClientId::Cursor).is_some()); - } - #[test] fn warp_reads_local_sqlite_usage() { let warp = ClientId::Warp.local_def().expect("warp has scan policy"); assert_eq!(warp.pattern, "warp.sqlite"); - assert!(ClientId::Warp.parse_local()); } #[test] @@ -598,7 +532,6 @@ mod tests { let def = ClientId::Junie.local_def().expect("junie has scan policy"); assert_eq!(def.relative_path, ".junie/sessions"); assert_eq!(def.pattern, "events.jsonl"); - assert!(ClientId::Junie.parse_local()); } #[test] @@ -608,7 +541,6 @@ mod tests { .expect("zcode has local scan policy"); assert_eq!(def.relative_path, ".zcode/projects"); assert_eq!(def.pattern, "*.jsonl"); - assert!(ClientId::Zcode.parse_local()); } #[test] @@ -618,7 +550,6 @@ mod tests { .expect("omp has local scan policy"); assert_eq!(def.relative_path, ".omp/agent/sessions"); assert_eq!(def.pattern, "*.jsonl"); - assert!(ClientId::Omp.parse_local()); } #[test] @@ -714,7 +645,6 @@ mod tests { relative_path: ".test/sessions", pattern: "*.jsonl", headless: false, - parse_local: true, }; assert_eq!( diff --git a/crates/tokscale-core/src/local_report_error.rs b/crates/tokscale-core/src/local_report_error.rs new file mode 100644 index 000000000..3e3f19aa3 --- /dev/null +++ b/crates/tokscale-core/src/local_report_error.rs @@ -0,0 +1,93 @@ +use std::error::Error; +use std::fmt; + +type BoxError = Box; + +/// Stable failure category for callers that need deterministic exit behavior. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum LocalReportErrorKind { + /// The caller supplied an invalid report request. + InvalidRequest, + /// User-controlled process or scanner configuration is malformed. + InvalidEnvironment, + /// The request was valid but execution failed while reading or processing it. + Operational, +} + +/// Typed local-report failure that preserves its originating error chain. +#[derive(Debug)] +pub struct LocalReportError { + kind: LocalReportErrorKind, + source: BoxError, +} + +impl LocalReportError { + pub(crate) fn invalid_request(source: impl Error + Send + Sync + 'static) -> Self { + Self::new(LocalReportErrorKind::InvalidRequest, source) + } + + pub(crate) fn invalid_environment(source: impl Error + Send + Sync + 'static) -> Self { + Self::new(LocalReportErrorKind::InvalidEnvironment, source) + } + + pub(crate) fn operational(source: impl Error + Send + Sync + 'static) -> Self { + Self::new(LocalReportErrorKind::Operational, source) + } + + pub(crate) fn invalid_request_message(message: impl Into) -> Self { + Self::invalid_request(MessageError(message.into())) + } + + pub(crate) fn invalid_environment_message(message: impl Into) -> Self { + Self::invalid_environment(MessageError(message.into())) + } + + /// Return the stable category without parsing the display message. + pub const fn kind(&self) -> LocalReportErrorKind { + self.kind + } + + /// Whether a CLI should classify this as invalid invocation/environment. + pub const fn is_invalid_invocation(&self) -> bool { + matches!( + self.kind, + LocalReportErrorKind::InvalidRequest | LocalReportErrorKind::InvalidEnvironment + ) + } + + fn new(kind: LocalReportErrorKind, source: impl Error + Send + Sync + 'static) -> Self { + Self { + kind, + source: Box::new(source), + } + } +} + +impl fmt::Display for LocalReportError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + self.source.fmt(formatter) + } +} + +impl Error for LocalReportError { + fn source(&self) -> Option<&(dyn Error + 'static)> { + Some(self.source.as_ref()) + } +} + +impl From for LocalReportError { + fn from(message: String) -> Self { + Self::operational(MessageError(message)) + } +} + +#[derive(Debug)] +struct MessageError(String); + +impl fmt::Display for MessageError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str(&self.0) + } +} + +impl Error for MessageError {} diff --git a/crates/tokscale-core/src/message_cache.rs b/crates/tokscale-core/src/message_cache.rs index 48b7f3054..b92579e1c 100644 --- a/crates/tokscale-core/src/message_cache.rs +++ b/crates/tokscale-core/src/message_cache.rs @@ -252,7 +252,9 @@ pub(crate) enum ParserId { OpenCodeSqlite, Claude, Codex, - Cursor, + // Retired parser tags keep persisted bincode discriminants stable. No + // active adapter can construct or request these parser versions. + RetiredCursor, Gemini, Amp, Droid, @@ -278,7 +280,7 @@ pub(crate) enum ParserId { KiroSqlite, KiroGlobalStorage, Junie, - Trae, + RetiredTrae, Cline, CommandCode, Grok, @@ -309,7 +311,7 @@ impl ParserId { Self::OpenCodeSqlite => "opencode-sqlite", Self::Claude => "claude", Self::Codex => "codex", - Self::Cursor => "cursor", + Self::RetiredCursor => "cursor", Self::Gemini => "gemini", Self::Amp => "amp", Self::Droid => "droid", @@ -336,7 +338,7 @@ impl ParserId { Self::KiroSqlite => "kiro-sqlite", Self::KiroGlobalStorage => "kiro-global-storage", Self::Junie => "junie", - Self::Trae => "trae", + Self::RetiredTrae => "trae", Self::Cline => "cline", Self::CommandCode => "command-code", Self::Grok => "grok", @@ -4246,7 +4248,7 @@ mod tests { .unwrap() .is_some()); assert!(loaded - .get_meta(source.path(), ParserVersion::new(ParserId::Cursor, 1)) + .get_meta(source.path(), ParserVersion::new(ParserId::Gemini, 1)) .unwrap() .is_none()); @@ -4337,7 +4339,7 @@ mod tests { let source = write_temp_file(b"source\n"); let fingerprint = SourceFingerprint::from_path(source.path()).unwrap(); let copilot_version = ParserVersion::new(ParserId::Copilot, 1); - let cursor_version = ParserVersion::new(ParserId::Cursor, 1); + let gemini_version = ParserVersion::new(ParserId::Gemini, 1); let mut cache = SourceMessageCache::load().unwrap(); cache.insert(CachedSourceEntry::new_with_version( source.path(), @@ -4362,13 +4364,13 @@ mod tests { )); cache.insert(CachedSourceEntry::new_with_version( source.path(), - cursor_version, + gemini_version, fingerprint.clone(), vec![UnifiedMessage::new( - "cursor", + "gemini", "gpt-5", "openai", - "cursor-session", + "gemini-session", 1, TokenBreakdown { input: 2, @@ -4384,10 +4386,10 @@ mod tests { cache.save_if_dirty().unwrap(); let copilot_shard = shard_path(source.path(), copilot_version).unwrap(); - let cursor_shard = shard_path(source.path(), cursor_version).unwrap(); - assert_ne!(copilot_shard, cursor_shard); + let gemini_shard = shard_path(source.path(), gemini_version).unwrap(); + assert_ne!(copilot_shard, gemini_shard); assert!(copilot_shard.exists()); - assert!(cursor_shard.exists()); + assert!(gemini_shard.exists()); let mut loaded = SourceMessageCache::load().unwrap(); assert!(loaded @@ -4395,7 +4397,7 @@ mod tests { .unwrap() .is_some()); assert!(loaded - .get_meta(source.path(), cursor_version) + .get_meta(source.path(), gemini_version) .unwrap() .is_some()); let copilot_messages = loaded @@ -4405,15 +4407,15 @@ mod tests { fingerprint.clone(), )) .unwrap(); - let cursor_messages = loaded + let gemini_messages = loaded .take_messages(&CacheReadPlan::new( source.path(), - cursor_version, + gemini_version, fingerprint, )) .unwrap(); assert_eq!(copilot_messages[0].session_id.as_ref(), "copilot-session"); - assert_eq!(cursor_messages[0].session_id.as_ref(), "cursor-session"); + assert_eq!(gemini_messages[0].session_id.as_ref(), "gemini-session"); restore_cache_env(prev_env); } diff --git a/crates/tokscale-core/src/scanner.rs b/crates/tokscale-core/src/scanner.rs index 8857cae30..d7ae11600 100644 --- a/crates/tokscale-core/src/scanner.rs +++ b/crates/tokscale-core/src/scanner.rs @@ -925,7 +925,7 @@ fn scan_all_clients_with_env_strategy_inner( // Merge user-configured `scanner.opencodeDbPaths` here, INSIDE the // `enabled.contains(&ClientId::OpenCode)` guard, so a request like - // `tokscale --client claude` does not pull in OpenCode dbs the user + // `tokscale models --client claude` does not pull in OpenCode dbs the user // pinned for unrelated reasons. Inflated OpenCode `counts` and wasted // SQLite parsing work otherwise sneak past the message-level // client filter that runs much later in the pipeline. @@ -1373,9 +1373,7 @@ mod tests { result .get_mut(ClientId::Gemini) .push(PathBuf::from("d.json")); - result - .get_mut(ClientId::Cursor) - .push(PathBuf::from("e.csv")); + result.get_mut(ClientId::Amp).push(PathBuf::from("e.jsonl")); result.get_mut(ClientId::Pi).push(PathBuf::from("f.jsonl")); let all = result.all_files(); @@ -1383,8 +1381,8 @@ mod tests { assert_eq!(all[0], (ClientId::OpenCode, PathBuf::from("a.json"))); assert_eq!(all[1], (ClientId::Claude, PathBuf::from("b.jsonl"))); assert_eq!(all[2], (ClientId::Codex, PathBuf::from("c.jsonl"))); - assert_eq!(all[3], (ClientId::Cursor, PathBuf::from("e.csv"))); - assert_eq!(all[4], (ClientId::Gemini, PathBuf::from("d.json"))); + assert_eq!(all[3], (ClientId::Gemini, PathBuf::from("d.json"))); + assert_eq!(all[4], (ClientId::Amp, PathBuf::from("e.jsonl"))); assert_eq!(all[5], (ClientId::Pi, PathBuf::from("f.jsonl"))); } @@ -2423,7 +2421,7 @@ mod tests { // Regression guard: previously the scanner unconditionally // merged `scanner.opencodeDbPaths` after the inner scan, which // bypassed the existing `enabled.contains(&ClientId::OpenCode)` - // guard. A request like `tokscale --client claude` would still pull + // guard. A request like `tokscale models --client claude` would still pull // in user-pinned OpenCode dbs and inflate local message loading // counts plus waste SQLite parsing work. // @@ -2894,35 +2892,6 @@ mod tests { ); } - #[test] - fn test_scan_all_clients_keeps_cursor_cache_scanning() { - let dir = TempDir::new().unwrap(); - let home = dir.path(); - - let cursor_file = home - .join(".config") - .join("tokscale") - .join("cursor-cache") - .join("usage.csv"); - fs::create_dir_all(cursor_file.parent().unwrap()).unwrap(); - fs::write(&cursor_file, "Date,Model,Input Tokens,Output Tokens\n").unwrap(); - - let all_clients = - scan_all_clients_with_env_strategy(home.to_str().unwrap(), &[], false).unwrap(); - let explicit_cursor = scan_all_clients_with_env_strategy( - home.to_str().unwrap(), - &["cursor".to_string()], - false, - ) - .unwrap(); - - assert_eq!( - all_clients.get(ClientId::Cursor), - &vec![cursor_file.clone()] - ); - assert_eq!(explicit_cursor.get(ClientId::Cursor), &vec![cursor_file]); - } - #[test] #[serial] fn test_scan_all_clients_headless_paths() { diff --git a/crates/tokscale-core/src/sessionize.rs b/crates/tokscale-core/src/sessionize.rs index ff7f9008d..990d0b493 100644 --- a/crates/tokscale-core/src/sessionize.rs +++ b/crates/tokscale-core/src/sessionize.rs @@ -930,7 +930,7 @@ mod tests { #[test] fn test_compute_daily_active_time_matches_local_day_boundaries_for_fixed_offset() { let interval = time_interval( - "trae", + "droid", "session-local-boundary", FixedOffset::east_opt(9 * 3600) .unwrap() diff --git a/crates/tokscale-core/src/sessions/cursor.rs b/crates/tokscale-core/src/sessions/cursor.rs deleted file mode 100644 index f625a7241..000000000 --- a/crates/tokscale-core/src/sessions/cursor.rs +++ /dev/null @@ -1,520 +0,0 @@ -//! Cursor IDE session parser -//! -//! Parses CSV files from the Cursor usage export API. -//! CSV files are cached locally at ~/.config/tokscale/cursor-cache/*.csv -//! (legacy single-account cache uses usage.csv; additional accounts may use usage..csv) -//! -//! CSV Formats: -//! - v1 (old): Date,Model,Input (w/ Cache Write),Input (w/o Cache Write),Cache Read,Output Tokens,Total Tokens,Cost,Cost to you -//! - v2 (new): Date,Kind,Model,Max Mode,Input (w/ Cache Write),Input (w/o Cache Write),Cache Read,Output Tokens,Total Tokens,Cost -//! - v3 (latest): Date,Cloud Agent ID,Automation ID,Kind,Model,Max Mode,Input (w/ Cache Write),Input (w/o Cache Write),Cache Read,Output Tokens,Total Tokens,Cost - -use super::error::{SessionParseError, SessionParseResult}; -use super::UnifiedMessage; -use crate::source_health::{RecordRejectionReason, ScannedSource}; -use crate::{provider_identity, TokenBreakdown}; -use std::path::Path; - -fn account_id_from_cursor_cache_path(path: &Path) -> String { - let file_name = path - .file_name() - .and_then(|n| n.to_str()) - .unwrap_or("usage.csv"); - - if file_name == "usage.csv" { - return "active".to_string(); - } - - if let Some(stem) = file_name - .strip_prefix("usage.") - .and_then(|s| s.strip_suffix(".csv")) - { - // Keep it simple/ASCII. The CLI already sanitizes file names. - let cleaned = stem - .chars() - .map(|c| { - if c.is_ascii_alphanumeric() || c == '-' || c == '_' || c == '.' { - c - } else { - '-' - } - }) - .collect::(); - if cleaned.is_empty() { - return "unknown".to_string(); - } - return cleaned; - } - - "unknown".to_string() -} - -/// Provider inference from model name -fn infer_provider(model: &str) -> &'static str { - provider_identity::inferred_provider_from_model(model).unwrap_or("cursor") -} - -/// Parse a Cursor usage CSV file -/// -/// Handles both formats: -/// - New: Date,Kind,Model,Max Mode,Input (w/ Cache Write),Input (w/o Cache Write),Cache Read,Output Tokens,Total Tokens,Cost -/// - Old: Date,Model,Input (w/ Cache Write),Input (w/o Cache Write),Cache Read,Output Tokens,Total Tokens,Cost,Cost to you -pub fn parse_cursor_file(path: &Path) -> SessionParseResult { - let content = std::fs::read_to_string(path) - .map_err(|error| SessionParseError::new("read CSV file", error))?; - - let mut scanned = ScannedSource { - messages: Vec::with_capacity(128), - ..ScannedSource::default() - }; - let mut lines = content.lines(); - - // Parse header line to determine column indices - let header = match lines.next() { - Some(h) => h, - None => return Ok(scanned), - }; - - // Verify this is a valid Cursor CSV - if !header.contains("Date") || !header.contains("Model") { - return Err(SessionParseError::invalid( - "validate CSV header", - "missing required Date or Model column", - )); - } - - // Detect format by checking for "Kind" column and column count - let header_fields: Vec<&str> = parse_csv_line(header); - let has_kind_column = header_fields.iter().any(|f| f.trim() == "Kind"); - let column_count = header_fields.len(); - - // Column indices based on format - let (model_idx, input_cache_write_idx, input_no_cache_idx, cache_read_idx, output_idx) = - if has_kind_column && column_count >= 11 { - // v3 format: Date,Cloud Agent ID,Automation ID,Kind,Model,... - (4, 6, 7, 8, 9) - } else if has_kind_column { - // v2 format: Date,Kind,Model,Max Mode,Input (w/ Cache Write),... - (2, 4, 5, 6, 7) - } else { - // v1 format: Date,Model,Input (w/ Cache Write),... - (1, 2, 3, 4, 5) - }; - - let account_id = account_id_from_cursor_cache_path(path); - - for line in lines { - if line.trim().is_empty() { - continue; - } - - // Parse CSV line (simple parsing, handles quoted fields) - let fields: Vec<&str> = parse_csv_line(line); - - // Need at least enough columns for the format - let min_fields = output_idx + 1; - if fields.len() < min_fields { - scanned - .rejections - .record(RecordRejectionReason::MalformedRecord); - continue; - } - - let date_str = fields[0].trim().trim_matches('"'); - let model = fields[model_idx].trim().trim_matches('"'); - let parse_count = |value: &str| { - value - .trim() - .trim_matches('"') - .parse::() - .map_err(|error| SessionParseError::new("decode CSV token count", error)) - }; - let counts = [ - fields[input_cache_write_idx], - fields[input_no_cache_idx], - fields[cache_read_idx], - fields[output_idx], - ] - .map(parse_count); - let [Ok(input_with_cache_write), Ok(input_without_cache_write), Ok(cache_read), Ok(output_tokens)] = - counts - else { - scanned - .rejections - .record(RecordRejectionReason::MalformedRecord); - continue; - }; - if [ - input_with_cache_write, - input_without_cache_write, - cache_read, - output_tokens, - ] - .into_iter() - .any(|tokens| tokens < 0) - { - scanned - .rejections - .record(RecordRejectionReason::MalformedRecord); - continue; - } - if input_with_cache_write < input_without_cache_write { - scanned - .rejections - .record(RecordRejectionReason::MalformedRecord); - continue; - } - - // Cache write = input_with_cache_write - input_without_cache_write - let cache_write = input_with_cache_write - input_without_cache_write; - // Input tokens = input_without_cache_write - let input = input_without_cache_write; - let tokens = TokenBreakdown { - input: input.max(0), - output: output_tokens.max(0), - cache_read: cache_read.max(0), - cache_write, - reasoning: 0, - }; - let Some(token_total) = tokens.checked_total() else { - scanned - .rejections - .record(RecordRejectionReason::MalformedRecord); - continue; - }; - if token_total == 0 { - continue; - } - - if model.is_empty() { - scanned - .rejections - .record(RecordRejectionReason::MissingModel); - continue; - } - - // Parse timestamp from date string - let timestamp = parse_date_to_timestamp(date_str); - if timestamp == 0 { - scanned - .rejections - .record(RecordRejectionReason::MissingTimestamp); - continue; - } - - scanned.messages.push(UnifiedMessage::new( - "cursor", - model, - infer_provider(model), - format!("cursor-{}-{}", account_id, date_str), - timestamp, - tokens, - 0.0, - )); - } - - Ok(scanned) -} - -/// Simple CSV line parser that handles quoted fields -fn parse_csv_line(line: &str) -> Vec<&str> { - let mut fields = Vec::new(); - let mut start = 0; - let mut in_quotes = false; - let bytes = line.as_bytes(); - - for (i, &byte) in bytes.iter().enumerate() { - match byte { - b'"' => in_quotes = !in_quotes, - b',' if !in_quotes => { - fields.push(&line[start..i]); - start = i + 1; - } - _ => {} - } - } - - // Add the last field - if start <= line.len() { - fields.push(&line[start..]); - } - - fields -} - -/// Parse a date string to Unix milliseconds timestamp -fn parse_date_to_timestamp(date_str: &str) -> i64 { - use chrono::{NaiveDate, NaiveDateTime, TimeZone, Utc}; - - // Try ISO 8601 format with milliseconds: "2025-02-05T12:00:00.123Z" - if let Ok(dt) = NaiveDateTime::parse_from_str(date_str, "%Y-%m-%dT%H:%M:%S%.3fZ") { - return Utc.from_utc_datetime(&dt).timestamp_millis(); - } - - // Try ISO 8601 format with time: "2025-02-05T12:00:00Z" - if let Ok(dt) = NaiveDateTime::parse_from_str(date_str, "%Y-%m-%dT%H:%M:%SZ") { - return Utc.from_utc_datetime(&dt).timestamp_millis(); - } - - // Try ISO 8601 format with milliseconds without Z: "2025-02-05T12:00:00.123" - if let Ok(dt) = NaiveDateTime::parse_from_str(date_str, "%Y-%m-%dT%H:%M:%S%.3f") { - return Utc.from_utc_datetime(&dt).timestamp_millis(); - } - - // Try ISO 8601 format with time without Z: "2025-02-05T12:00:00" - if let Ok(dt) = NaiveDateTime::parse_from_str(date_str, "%Y-%m-%dT%H:%M:%S") { - return Utc.from_utc_datetime(&dt).timestamp_millis(); - } - - // Date-only format: "2025-02-05" - use noon UTC (12:00:00Z) - // Noon keeps the local date stable for all timezones from UTC-12 to UTC+14, - // so filtering by local day boundaries won't shift the record to an adjacent day. - if let Ok(date) = NaiveDate::parse_from_str(date_str, "%Y-%m-%d") { - let dt = date.and_hms_opt(12, 0, 0).unwrap(); - return Utc.from_utc_datetime(&dt).timestamp_millis(); - } - - 0 -} - -#[cfg(test)] -mod tests { - use super::*; - - fn parse_cursor_file(path: &Path) -> Vec { - super::parse_cursor_file(path).unwrap().messages - } - - #[test] - fn rejects_csv_without_required_header() { - let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("usage.csv"); - std::fs::write(&path, "Kind,Input,Output\nchat,1,2\n").unwrap(); - - let error = super::parse_cursor_file(&path).unwrap_err(); - assert_eq!(error.operation(), "validate CSV header"); - } - - #[test] - fn missing_csv_is_an_error() { - let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("missing.csv"); - - let error = super::parse_cursor_file(&path).unwrap_err(); - assert_eq!(error.operation(), "read CSV file"); - } - - #[test] - fn mixed_rows_reject_bad_record_and_keep_later_usage() { - let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("usage.csv"); - std::fs::write( - &path, - "Date,Model,Input (w/ Cache Write),Input (w/o Cache Write),Cache Read,Output Tokens,Total Tokens,Cost,Cost to you\n\ - 2025-02-01,gpt-4o,10,8,1,2,11,$0,$0\n\ - 2025-02-02,gpt-4o,not-a-number,8,1,2,11,$0,$0\n\ - 2025-02-03,gpt-4o,20,16,2,4,22,$0,$0\n", - ) - .unwrap(); - - let scanned = super::parse_cursor_file(&path).unwrap(); - - assert_eq!(scanned.messages.len(), 2); - assert_eq!(scanned.rejections.total(), 1); - assert!(scanned.interrupted.is_none()); - } - - #[test] - fn negative_row_tokens_are_malformed_instead_of_clamped() { - let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("usage.csv"); - std::fs::write( - &path, - "Date,Model,Input (w/ Cache Write),Input (w/o Cache Write),Cache Read,Output Tokens,Total Tokens,Cost,Cost to you\n\ - 2025-02-01,gpt-4o,10,8,1,2,11,$0,$0\n\ - 2025-02-02,gpt-4o,10,-8,1,2,11,$0,$0\n\ - 2025-02-03,gpt-4o,20,16,2,4,22,$0,$0\n", - ) - .unwrap(); - - let scanned = super::parse_cursor_file(&path).unwrap(); - - assert_eq!(scanned.messages.len(), 2); - assert_eq!(scanned.rejections.total(), 1); - assert_eq!( - scanned.rejections.entries().next().unwrap().key, - "malformed-record" - ); - } - - #[test] - fn inconsistent_cache_write_inputs_are_malformed_and_later_row_survives() { - let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("usage.csv"); - std::fs::write( - &path, - "Date,Model,Input (w/ Cache Write),Input (w/o Cache Write),Cache Read,Output Tokens,Total Tokens,Cost,Cost to you\n\ - 2025-02-01,gpt-4o,10,20,1,2,23,$0,$0\n\ - 2025-02-02,gpt-4o,20,16,2,4,22,$0,$0\n", - ) - .unwrap(); - - let scanned = super::parse_cursor_file(&path).unwrap(); - - assert_eq!(scanned.messages.len(), 1); - assert_eq!(scanned.messages[0].tokens.input, 16); - assert_eq!(scanned.messages[0].tokens.cache_write, 4); - assert_eq!(scanned.rejections.total(), 1); - assert_eq!( - scanned.rejections.entries().next().unwrap().key, - "malformed-record" - ); - } - - #[test] - fn overflowing_row_tokens_are_malformed_and_later_row_survives() { - let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("usage.csv"); - std::fs::write( - &path, - "Date,Model,Input (w/ Cache Write),Input (w/o Cache Write),Cache Read,Output Tokens,Total Tokens,Cost,Cost to you\n\ - 2025-02-01,gpt-4o,9223372036854775807,9223372036854775807,0,1,0,$0,$0\n\ - 2025-02-02,gpt-4o,20,16,2,4,22,$0,$0\n", - ) - .unwrap(); - - let scanned = super::parse_cursor_file(&path).unwrap(); - - assert_eq!(scanned.messages.len(), 1); - assert_eq!(scanned.messages[0].tokens.input, 16); - assert_eq!(scanned.rejections.total(), 1); - assert_eq!( - scanned.rejections.entries().next().unwrap().key, - "malformed-record" - ); - } - - #[test] - fn test_infer_provider() { - assert_eq!(infer_provider("claude-3-sonnet"), "anthropic"); - assert_eq!(infer_provider("gpt-4o"), "openai"); - assert_eq!(infer_provider("gemini-pro"), "google"); - assert_eq!(infer_provider("deepseek-coder"), "deepseek"); - assert_eq!(infer_provider("llama-3"), "meta"); - assert_eq!(infer_provider("unknown-model"), "cursor"); - } - - #[test] - fn test_parse_csv_line() { - let line = "2025-02-01,gpt-4o,10,5,0,15,30,$0.10,$0.10"; - let fields = parse_csv_line(line); - assert_eq!(fields.len(), 9); - assert_eq!(fields[0], "2025-02-01"); - assert_eq!(fields[1], "gpt-4o"); - assert_eq!(fields[8], "$0.10"); - } - - #[test] - fn test_parse_date_to_timestamp() { - // ISO with milliseconds and Z (new Cursor format) - let ts = parse_date_to_timestamp("2025-11-13T18:36:05.846Z"); - assert!(ts > 0); - - // ISO with Z - let ts = parse_date_to_timestamp("2025-02-05T12:00:00Z"); - assert!(ts > 0); - - // Date only - let ts = parse_date_to_timestamp("2025-02-05"); - assert!(ts > 0); - - // Invalid - let ts = parse_date_to_timestamp("invalid"); - assert_eq!(ts, 0); - } - - #[test] - fn test_parse_cursor_csv_sample_old_format() { - let csv = "Date,Model,Input (w/ Cache Write),Input (w/o Cache Write),Cache Read,Output Tokens,Total Tokens,Cost,Cost to you -2025-02-01,gpt-4o,10,5,0,15,30,$0.10,$0.10 -2025-02-02,gpt-4o-mini,0,0,0,5,5,$0.05,$0.05"; - - let temp_dir = tempfile::TempDir::new().unwrap(); - let file_path = temp_dir.path().join("usage.csv"); - std::fs::write(&file_path, csv).unwrap(); - - let messages = parse_cursor_file(&file_path); - assert_eq!(messages.len(), 2); - - assert_eq!(messages[0].client.as_ref(), "cursor"); - assert_eq!(messages[0].model_id.as_ref(), "gpt-4o"); - assert_eq!(messages[0].provider_id.as_ref(), "openai"); - assert_eq!(messages[0].tokens.input, 5); - assert_eq!(messages[0].tokens.output, 15); - assert_eq!(messages[0].tokens.cache_write, 5); // 10 - 5 - assert_eq!(messages[0].cost, 0.0); - - assert_eq!(messages[1].model_id.as_ref(), "gpt-4o-mini"); - } - - #[test] - fn test_parse_cursor_csv_sample_new_format() { - // Cursor API column layout with internally consistent cache-write totals. - let csv = r#"Date,Kind,Model,Max Mode,Input (w/ Cache Write),Input (w/o Cache Write),Cache Read,Output Tokens,Total Tokens,Cost -"2025-11-13T18:36:05.846Z","Included","auto","No","28342","775","105891","21282","156290","0.19" -"2025-11-13T13:35:04.658Z","On-Demand","gpt-5-codex","No","8263","8263","66964","1612","76839","0.03""#; - - let temp_dir = tempfile::TempDir::new().unwrap(); - let file_path = temp_dir.path().join("usage.csv"); - std::fs::write(&file_path, csv).unwrap(); - - let messages = parse_cursor_file(&file_path); - assert_eq!(messages.len(), 2); - - // First message: auto model - assert_eq!(messages[0].client.as_ref(), "cursor"); - assert_eq!(messages[0].model_id.as_ref(), "auto"); - assert_eq!(messages[0].provider_id.as_ref(), "cursor"); // unknown model -> cursor - assert_eq!(messages[0].tokens.input, 775); - assert_eq!(messages[0].tokens.output, 21282); - assert_eq!(messages[0].tokens.cache_read, 105891); - assert_eq!(messages[0].tokens.cache_write, 28342 - 775); // 27567 - assert_eq!(messages[0].cost, 0.0); - - // Second message: gpt-5-codex - assert_eq!(messages[1].model_id.as_ref(), "gpt-5-codex"); - assert_eq!(messages[1].provider_id.as_ref(), "openai"); // gpt -> openai - assert_eq!(messages[1].tokens.input, 8263); - assert_eq!(messages[1].tokens.cache_read, 66964); - } - - #[test] - fn test_parse_cursor_csv_sample_v3_format() { - // v3 format includes Cloud Agent ID and Automation ID columns - let csv = r#"Date,Cloud Agent ID,Automation ID,Kind,Model,Max Mode,Input (w/ Cache Write),Input (w/o Cache Write),Cache Read,Output Tokens,Total Tokens,Cost -"2026-04-09T20:01:10.528Z","bc-a380fb49-e1a5-414e-817d-6a85b6cdc51c","cc30782e-26cc-4359-bc22-7567efe282be","Included","composer-2","Yes","343446","343446","29045760","915201","30304407","Included" -"2026-04-09T18:02:13.576Z","bc-19a9b74b-2af3-46e2-9f61-3ba1cdac46c8","1a0df38f-1474-4dfe-896b-70b841d4a833","On-Demand","composer-2","Yes","43478","43478","420864","7957","472299","0.11" -"2026-04-09T07:39:09.091Z","bc-49262501-0ee0-49f9-b856-a5b0466deddb","","Errored, No Charge","composer-2","Yes","104504","104504","985600","3666","1093770","-""#; - - let temp_dir = tempfile::TempDir::new().unwrap(); - let file_path = temp_dir.path().join("usage.csv"); - std::fs::write(&file_path, csv).unwrap(); - - let messages = parse_cursor_file(&file_path); - assert_eq!(messages.len(), 3); - - // First message: "Included" cost should be 0 - assert_eq!(messages[0].client.as_ref(), "cursor"); - assert_eq!(messages[0].model_id.as_ref(), "composer-2"); - assert_eq!(messages[0].cost, 0.0); - assert_eq!(messages[0].tokens.cache_read, 29045760); - - // Second message: app-reported cost is ignored. - assert_eq!(messages[1].model_id.as_ref(), "composer-2"); - assert_eq!(messages[1].cost, 0.0); - - // Third message: "-" cost should be 0 (Errored, No Charge) - assert_eq!(messages[2].model_id.as_ref(), "composer-2"); - assert_eq!(messages[2].cost, 0.0); - } -} diff --git a/crates/tokscale-core/src/sessions/droid.rs b/crates/tokscale-core/src/sessions/droid.rs index 5ce76bb03..00e26b8bd 100644 --- a/crates/tokscale-core/src/sessions/droid.rs +++ b/crates/tokscale-core/src/sessions/droid.rs @@ -7,7 +7,20 @@ use super::UnifiedMessage; use crate::source_health::{RecordRejectionReason, ScannedSource}; use crate::{model_aliases, provider_identity, TokenBreakdown}; use serde::Deserialize; -use std::path::Path; +use std::io::{BufRead, BufReader}; +use std::path::{Path, PathBuf}; + +const DROID_EXPLORER_AGENT: &str = "Droid Explorer"; +const DROID_WORKER_AGENT: &str = "Droid Worker"; +const DROID_ORCHESTRATOR_AGENT: &str = "Droid Orchestrator"; +const DROID_VALIDATOR_AGENT: &str = "Droid Validator"; + +const MISSION_ORCHESTRATOR_TAG: &str = "mission-orchestrator"; +const MISSION_SESSION_TAG: &str = "mission-session"; +const MISSION_WORKER_TAG: &str = "mission-worker"; +const SUBAGENT_TAG: &str = "subagent"; +const SCRUTINY_VALIDATOR_SKILL: &str = "scrutiny-validator"; +const USER_TESTING_VALIDATOR_SKILL: &str = "user-testing-validator"; /// Droid settings.json structure #[derive(Debug, Deserialize)] @@ -19,6 +32,53 @@ pub struct DroidSettingsJson { pub provider_lock_timestamp: Option, #[serde(rename = "tokenUsage")] pub token_usage: Option, + #[serde(default)] + tags: Vec, +} + +#[derive(Debug, Deserialize)] +struct DroidTag { + name: String, + metadata: Option, +} + +#[derive(Debug, Deserialize)] +struct DroidTagMetadata { + role: Option, + #[serde(rename = "missionId")] + mission_id: Option, +} + +#[derive(Debug, Deserialize)] +struct DroidSessionStart { + #[serde(rename = "type")] + record_type: String, + title: Option, + #[serde(rename = "callingSessionId")] + calling_session_id: Option, +} + +impl DroidSessionStart { + fn parent_session_id(&self) -> Option<&str> { + self.calling_session_id + .as_deref() + .map(str::trim) + .filter(|session_id| !session_id.is_empty()) + } +} + +#[derive(Debug, Deserialize)] +struct DroidMissionFeatures { + #[serde(default)] + features: Vec, +} + +#[derive(Debug, Deserialize)] +struct DroidMissionFeature { + #[serde(rename = "skillName")] + skill_name: Option, + #[serde(default, rename = "workerSessionIds")] + worker_session_ids: Vec, } #[derive(Debug, Deserialize)] @@ -124,6 +184,179 @@ fn invalid_at_path( ) } +fn settings_session_id(path: &Path) -> Option<&str> { + path.file_name() + .and_then(|name| name.to_str()) + .and_then(|name| name.strip_suffix(".settings.json")) + .map(str::trim) + .filter(|session_id| !session_id.is_empty()) +} + +fn transcript_path(path: &Path) -> Option { + settings_session_id(path).map(|session_id| path.with_file_name(format!("{session_id}.jsonl"))) +} + +fn read_settings(path: &Path) -> Option { + let mut bytes = std::fs::read(path).ok()?; + simd_json::from_slice(&mut bytes).ok() +} + +fn read_session_start(path: &Path) -> Option { + let transcript = transcript_path(path)?; + let mut first_line = String::new(); + BufReader::new(std::fs::File::open(transcript).ok()?) + .read_line(&mut first_line) + .ok()?; + let start: DroidSessionStart = serde_json::from_str(&first_line).ok()?; + (start.record_type == "session_start").then_some(start) +} + +fn factory_root(path: &Path) -> Option<&Path> { + path.ancestors() + .find(|ancestor| ancestor.file_name().and_then(|name| name.to_str()) == Some("sessions")) + .and_then(Path::parent) +} + +fn mission_features_path(path: &Path, mission_id: &str) -> Option { + Some( + factory_root(path)? + .join("missions") + .join(mission_id) + .join("features.json"), + ) +} + +fn has_tag(settings: &DroidSettingsJson, name: &str) -> bool { + settings.tags.iter().any(|tag| tag.name == name) +} + +fn mission_session_metadata(settings: &DroidSettingsJson) -> Option<&DroidTagMetadata> { + settings + .tags + .iter() + .find(|tag| tag.name == MISSION_SESSION_TAG) + .and_then(|tag| tag.metadata.as_ref()) +} + +fn mission_session_role(settings: &DroidSettingsJson) -> Option<&str> { + mission_session_metadata(settings).and_then(|metadata| metadata.role.as_deref()) +} + +fn mission_id(settings: &DroidSettingsJson) -> Option<&str> { + mission_session_metadata(settings) + .and_then(|metadata| metadata.mission_id.as_deref()) + .map(str::trim) + .filter(|mission_id| !mission_id.is_empty()) +} + +fn mission_worker_is_validator( + path: &Path, + settings: &DroidSettingsJson, + session_id: &str, +) -> bool { + let Some(features_path) = + mission_id(settings).and_then(|mission_id| mission_features_path(path, mission_id)) + else { + return false; + }; + let Ok(mut bytes) = std::fs::read(features_path) else { + return false; + }; + let Ok(features) = simd_json::from_slice::(&mut bytes) else { + return false; + }; + + features.features.iter().any(|feature| { + feature + .worker_session_ids + .iter() + .any(|worker_id| worker_id == session_id) + && matches!( + feature.skill_name.as_deref(), + Some(SCRUTINY_VALIDATOR_SKILL | USER_TESTING_VALIDATOR_SKILL) + ) + }) +} + +fn agent_from_title(title: &str) -> Option<&'static str> { + let label = title + .split(':') + .next() + .unwrap_or(title) + .trim() + .to_ascii_lowercase() + .replace([' ', '_'], "-"); + + match label.as_str() { + "explorer" => Some(DROID_EXPLORER_AGENT), + "scrutiny-feature-reviewer" | "user-testing-flow-validator" => Some(DROID_VALIDATOR_AGENT), + _ => None, + } +} + +fn resolve_droid_agent( + path: &Path, + settings: &DroidSettingsJson, + inherit_validator_parent: bool, +) -> Option<&'static str> { + if has_tag(settings, MISSION_ORCHESTRATOR_TAG) + || mission_session_role(settings) == Some("orchestrator") + { + return Some(DROID_ORCHESTRATOR_AGENT); + } + // The built-in tag is the authoritative Mission Worker marker. + if has_tag(settings, MISSION_WORKER_TAG) { + let session_id = settings_session_id(path)?; + return Some(if mission_worker_is_validator(path, settings, session_id) { + DROID_VALIDATOR_AGENT + } else { + DROID_WORKER_AGENT + }); + } + if !has_tag(settings, SUBAGENT_TAG) { + return None; + } + + let start = read_session_start(path); + if let Some(agent) = start + .as_ref() + .and_then(|start| start.title.as_deref()) + .and_then(agent_from_title) + { + return Some(agent); + } + if inherit_validator_parent { + let parent_is_validator = start + .as_ref() + .and_then(DroidSessionStart::parent_session_id) + .map(|parent_id| path.with_file_name(format!("{parent_id}.settings.json"))) + .and_then(|parent_path| { + let parent_settings = read_settings(&parent_path)?; + resolve_droid_agent(&parent_path, &parent_settings, false) + }) + == Some(DROID_VALIDATOR_AGENT); + if parent_is_validator { + return Some(DROID_VALIDATOR_AGENT); + } + } + + Some(DROID_WORKER_AGENT) +} + +/// Return the role-bearing companion file that participates in Droid cache +/// invalidation. Task subagents derive their role from the session header; +/// Mission workers derive it from the Mission feature assigned to the worker. +pub(crate) fn droid_agent_dependency_path(path: &Path) -> Option { + let settings = read_settings(path)?; + if has_tag(&settings, MISSION_WORKER_TAG) { + return mission_id(&settings) + .and_then(|mission_id| mission_features_path(path, mission_id)); + } + has_tag(&settings, SUBAGENT_TAG) + .then(|| transcript_path(path)) + .flatten() +} + /// Parse a Droid settings.json file pub fn parse_droid_file(path: &Path) -> SessionParseResult { let data = std::fs::read(path) @@ -133,6 +366,8 @@ pub fn parse_droid_file(path: &Path) -> SessionParseResult { let settings: DroidSettingsJson = simd_json::from_slice(&mut bytes) .map_err(|error| SessionParseError::at_path(path, "decode JSON", error))?; + let agent = resolve_droid_agent(path, &settings, true).map(str::to_string); + // Skip if no token usage data let usage = match settings.token_usage { Some(u) => u, @@ -236,8 +471,8 @@ pub fn parse_droid_file(path: &Path) -> SessionParseResult { return Ok(scanned); } - scanned.messages.push(UnifiedMessage::new( - "droid", model, provider, session_id, timestamp, tokens, 0.0, + scanned.messages.push(UnifiedMessage::new_with_agent( + "droid", model, provider, session_id, timestamp, tokens, 0.0, agent, )); Ok(scanned) } @@ -245,11 +480,216 @@ pub fn parse_droid_file(path: &Path) -> SessionParseResult { #[cfg(test)] mod tests { use super::*; + use serde_json::json; fn parse_droid_file(path: &Path) -> Vec { super::parse_droid_file(path).unwrap().messages } + fn write_json(path: &Path, value: serde_json::Value) { + std::fs::create_dir_all(path.parent().unwrap()).unwrap(); + std::fs::write(path, serde_json::to_vec_pretty(&value).unwrap()).unwrap(); + } + + fn write_settings(path: &Path, tags: serde_json::Value) { + write_json( + path, + json!({ + "model": "custom:gpt-5.6-sol-xhigh", + "providerLock": "openai", + "providerLockTimestamp": "2026-07-15T08:55:13.871Z", + "tokenUsage": { + "inputTokens": 10, + "outputTokens": 5 + }, + "tags": tags + }), + ); + } + + fn write_session_start(path: &Path, title: &str, calling_session_id: Option<&str>) { + std::fs::create_dir_all(path.parent().unwrap()).unwrap(); + let value = json!({ + "type": "session_start", + "id": path.file_stem().and_then(|stem| stem.to_str()), + "title": title, + "callingSessionId": calling_session_id + }); + std::fs::write(path, serde_json::to_vec(&value).unwrap()).unwrap(); + } + + fn agent_for(path: &Path) -> Option { + parse_droid_file(path) + .into_iter() + .next() + .and_then(|message| message.agent.map(|agent| agent.to_string())) + } + + fn mission_worker_tags(mission_id: &str) -> serde_json::Value { + json!([ + {"name": "exec"}, + {"name": "mission-worker"}, + { + "name": "mission-session", + "metadata": {"role": "worker", "missionId": mission_id} + } + ]) + } + + #[test] + fn test_parse_droid_file_attributes_four_agent_roles() { + let temp_dir = tempfile::tempdir().unwrap(); + let factory = temp_dir.path().join(".factory"); + let sessions = factory.join("sessions/project"); + let mission_id = "mission-session"; + let features_path = factory + .join("missions") + .join(mission_id) + .join("features.json"); + + let orchestrator = sessions.join(format!("{mission_id}.settings.json")); + write_settings( + &orchestrator, + json!([{ + "name": "mission-session", + "metadata": {"role": "orchestrator", "missionId": mission_id} + }]), + ); + + let explorer = sessions.join("explorer.settings.json"); + write_settings(&explorer, json!([{"name": "subagent"}])); + write_session_start( + &sessions.join("explorer.jsonl"), + "Explorer: inspect parser flow", + Some(mission_id), + ); + + let worker = sessions.join("worker.settings.json"); + write_settings(&worker, json!([{"name": "subagent"}])); + write_session_start( + &sessions.join("worker.jsonl"), + "Worker: implement parser flow", + Some(mission_id), + ); + + let implementation_worker = sessions.join("implementation-worker.settings.json"); + write_settings(&implementation_worker, mission_worker_tags(mission_id)); + + let scrutiny_validator = sessions.join("scrutiny-validator.settings.json"); + write_settings(&scrutiny_validator, mission_worker_tags(mission_id)); + + let user_testing_validator = sessions.join("user-testing-validator.settings.json"); + write_settings(&user_testing_validator, mission_worker_tags(mission_id)); + + write_json( + &features_path, + json!({ + "features": [ + { + "id": "implementation", + "skillName": "backend-worker", + "workerSessionIds": ["implementation-worker"] + }, + { + "id": "scrutiny", + "skillName": "scrutiny-validator", + "workerSessionIds": ["scrutiny-validator"] + }, + { + "id": "user-testing", + "skillName": "user-testing-validator", + "workerSessionIds": ["user-testing-validator"] + } + ] + }), + ); + + let scrutiny_reviewer = sessions.join("scrutiny-reviewer.settings.json"); + write_settings(&scrutiny_reviewer, json!([{"name": "subagent"}])); + write_session_start( + &sessions.join("scrutiny-reviewer.jsonl"), + "Worker: review implementation", + Some("scrutiny-validator"), + ); + + let flow_validator = sessions.join("flow-validator.settings.json"); + write_settings(&flow_validator, json!([{"name": "subagent"}])); + write_session_start( + &sessions.join("flow-validator.jsonl"), + "Worker: validate user flow", + Some("user-testing-validator"), + ); + + assert_eq!( + agent_for(&orchestrator).as_deref(), + Some(DROID_ORCHESTRATOR_AGENT) + ); + assert_eq!(agent_for(&explorer).as_deref(), Some(DROID_EXPLORER_AGENT)); + assert_eq!(agent_for(&worker).as_deref(), Some(DROID_WORKER_AGENT)); + assert_eq!( + agent_for(&implementation_worker).as_deref(), + Some(DROID_WORKER_AGENT) + ); + assert_eq!( + agent_for(&scrutiny_validator).as_deref(), + Some(DROID_VALIDATOR_AGENT) + ); + assert_eq!( + agent_for(&user_testing_validator).as_deref(), + Some(DROID_VALIDATOR_AGENT) + ); + assert_eq!( + agent_for(&scrutiny_reviewer).as_deref(), + Some(DROID_VALIDATOR_AGENT) + ); + assert_eq!( + agent_for(&flow_validator).as_deref(), + Some(DROID_VALIDATOR_AGENT) + ); + + assert_eq!( + droid_agent_dependency_path(&explorer), + Some(sessions.join("explorer.jsonl")) + ); + assert_eq!( + droid_agent_dependency_path(&scrutiny_validator), + Some(features_path) + ); + } + + #[test] + fn test_parse_droid_file_requires_mission_worker_tag() { + let temp_dir = tempfile::tempdir().unwrap(); + let factory = temp_dir.path().join(".factory"); + let sessions = factory.join("sessions/project"); + let mission_id = "mission-session"; + let worker = sessions.join("metadata-only-worker.settings.json"); + + write_settings( + &worker, + json!([{ + "name": "mission-session", + "metadata": {"role": "worker", "missionId": mission_id} + }]), + ); + write_json( + &factory + .join("missions") + .join(mission_id) + .join("features.json"), + json!({ + "features": [{ + "id": "implementation", + "skillName": "backend-worker", + "workerSessionIds": ["metadata-only-worker"] + }] + }), + ); + + assert_eq!(agent_for(&worker), None); + assert_eq!(droid_agent_dependency_path(&worker), None); + } + #[test] fn test_normalize_model_name_custom_prefix() { assert_eq!( diff --git a/crates/tokscale-core/src/sessions/mod.rs b/crates/tokscale-core/src/sessions/mod.rs index e6d51bf4a..b8075c467 100644 --- a/crates/tokscale-core/src/sessions/mod.rs +++ b/crates/tokscale-core/src/sessions/mod.rs @@ -12,7 +12,6 @@ pub mod codebuff; pub mod codex; pub mod commandcode; pub mod copilot; -pub mod cursor; pub mod droid; pub mod error; pub mod gemini; @@ -31,7 +30,6 @@ pub mod opencode; pub mod pi; pub mod qwen; pub mod roocode; -pub mod trae; pub(crate) mod utils; pub mod warp; pub mod zcode; diff --git a/crates/tokscale-core/src/sessions/trae.rs b/crates/tokscale-core/src/sessions/trae.rs deleted file mode 100644 index 3f6e2b5ec..000000000 --- a/crates/tokscale-core/src/sessions/trae.rs +++ /dev/null @@ -1,418 +0,0 @@ -//! Trae usage API response parser. -//! -//! Reads `trae-cache/sessions/*.json` — the raw JSON dumped -//! from the usage API — and converts each entry into a `UnifiedMessage`. -//! -//! The API returns token counts and vendor spend. Tokscale ignores vendor spend -//! and derives report cost from token usage through its own pricing table. - -use super::error::{SessionParseError, SessionParseResult}; -use super::UnifiedMessage; -use crate::source_health::{RecordRejectionReason, ScannedSource}; -use crate::TokenBreakdown; - -/// Known mapping from Trae display names to tiktoken-style model ids. -/// Unknown names fall through to the raw `model_name` from the API -/// (mixed-case, space-separated). -fn normalize_trae_model(name: &str) -> String { - match name { - "GPT-5.4" => "gpt-5.4", - "GPT-5.3-Codex" | "GPT-5.3 Codex" => "gpt-5.3-codex", - "GPT-5.3" => "gpt-5.3", - "GPT-5.2-Codex" | "GPT-5.2 Codex" => "gpt-5.2-codex", - "GPT-5.2" => "gpt-5.2", - "GPT-5.1-Codex" | "GPT-5.1 Codex" => "gpt-5.1-codex", - "GPT-5.1" => "gpt-5.1", - "Gemini 3.1 Pro" => "gemini-3.1-pro", - "Gemini 3.1" => "gemini-3.1", - "GLM 5.1" | "GLM-5.1" => "glm-5.1", - "Claude Sonnet 4.6" | "Claude-Sonnet-4.6" => "claude-sonnet-4.6", - "Claude Sonnet 4.5" | "Claude-Sonnet-4.5" => "claude-sonnet-4.5", - other => other, - } - .to_string() -} - -/// Infer the provider from the display name. -fn provider_for_model(name: &str) -> &'static str { - if name.contains("GPT") || name.contains("gpt") { - "openai" - } else if name.contains("Claude") || name.contains("claude") { - "anthropic" - } else if name.contains("Gemini") || name.contains("gemini") { - "google" - } else if name.contains("GLM") || name.contains("glm") { - "zhipu" - } else { - "trae" - } -} - -/// Parse a single session JSON object into a `UnifiedMessage`. -struct TraeSessionRejection { - reason: RecordRejectionReason, -} - -impl TraeSessionRejection { - fn new(reason: RecordRejectionReason) -> Self { - Self { reason } - } -} - -fn parse_session( - client: &str, - session: &serde_json::Value, -) -> Result, TraeSessionRejection> { - if !session.is_object() { - return Err(TraeSessionRejection::new( - RecordRejectionReason::MalformedRecord, - )); - } - let extra = &session["extra_info"]; - if !extra.is_object() { - return Err(TraeSessionRejection::new( - RecordRejectionReason::MalformedRecord, - )); - } - let tokens = TokenBreakdown { - input: nonnegative_token(extra, "input_token")?, - output: nonnegative_token(extra, "output_token")?, - cache_read: nonnegative_token(extra, "cache_read_token")?, - cache_write: nonnegative_token(extra, "cache_write_token")?, - reasoning: 0, - }; - if crate::positive_token_total(&tokens) == 0 { - return Ok(None); - } - - let model_raw = session["model_name"] - .as_str() - .map(str::trim) - .filter(|model| !model.is_empty()) - .ok_or_else(|| TraeSessionRejection::new(RecordRejectionReason::MissingModel))?; - let model_id = normalize_trae_model(model_raw); - let provider = provider_for_model(&model_id); - // Records without a real `session_id` cannot be deduplicated correctly - // (every "missing-id" record would collide on the same key); records - // without a positive `usage_time` would land at epoch 0. Reject them - // rather than fabricating placeholders. - let session_id = session["session_id"] - .as_str() - .filter(|value| !value.trim().is_empty()) - .ok_or_else(|| TraeSessionRejection::new(RecordRejectionReason::MalformedRecord))?; - let usage_time = session["usage_time"] - .as_i64() - .ok_or_else(|| TraeSessionRejection::new(RecordRejectionReason::MissingTimestamp))?; - if usage_time <= 0 { - return Err(TraeSessionRejection::new( - RecordRejectionReason::MissingTimestamp, - )); - } - // API returns epoch seconds; UnifiedMessage expects milliseconds. Use - // `checked_mul` because the JSON cache is untrusted input — a crafted - // `usage_time` near `i64::MAX` would panic in debug builds and silently - // wrap to a negative timestamp in release builds. - let timestamp_ms = usage_time - .checked_mul(1000) - .ok_or_else(|| TraeSessionRejection::new(RecordRejectionReason::MissingTimestamp))?; - let dedup_key = Some(crate::sessions::dedup_hash_str(&format!( - "trae:{}:{}", - session_id, usage_time - ))); - - Ok(Some(UnifiedMessage::new_with_dedup( - client, - model_id, - provider, - session_id, - timestamp_ms, - tokens, - 0.0, - dedup_key, - ))) -} - -fn nonnegative_token(extra: &serde_json::Value, field: &str) -> Result { - let Some(value) = extra.get(field) else { - return Ok(0); - }; - let value = value - .as_i64() - .ok_or_else(|| TraeSessionRejection::new(RecordRejectionReason::MalformedRecord))?; - if value < 0 { - return Err(TraeSessionRejection::new( - RecordRejectionReason::MalformedRecord, - )); - } - Ok(value) -} - -/// Parse a cache file containing an array of sessions as returned by the API. -pub fn parse_trae_file(client: &str, path: &std::path::Path) -> SessionParseResult { - let content = std::fs::read_to_string(path) - .map_err(|error| SessionParseError::new("read Trae cache file", error))?; - let value: serde_json::Value = serde_json::from_str(&content) - .map_err(|error| SessionParseError::new("decode Trae cache file", error))?; - let sessions = value.as_array().ok_or_else(|| { - SessionParseError::invalid( - "validate Trae cache file", - "top-level value must be an array", - ) - })?; - let mut scanned = ScannedSource::default(); - for session in sessions { - match parse_session(client, session) { - Ok(Some(message)) => scanned.messages.push(message), - Ok(None) => {} - Err(rejection) => { - scanned.rejections.record(rejection.reason); - } - } - } - Ok(scanned) -} - -#[cfg(test)] -mod tests { - use super::*; - use std::io::Write; - - fn parse_trae_file(client: &str, path: &std::path::Path) -> Vec { - let scanned = super::parse_trae_file(client, path).unwrap(); - assert!(scanned.rejections.is_empty()); - assert!(scanned.interrupted.is_none()); - scanned.messages - } - - fn rejection_keys(client: &str, path: &std::path::Path) -> Vec { - super::parse_trae_file(client, path) - .unwrap() - .rejections - .entries() - .map(|entry| entry.key.to_string()) - .collect() - } - - fn write_fixture(data: &str) -> tempfile::NamedTempFile { - let mut f = tempfile::NamedTempFile::new().unwrap(); - f.write_all(data.as_bytes()).unwrap(); - f.flush().unwrap(); - f - } - - #[test] - fn test_parse_empty_file() { - let f = write_fixture("[]"); - let msgs = parse_trae_file("trae", f.path()); - assert!(msgs.is_empty()); - } - - #[test] - fn test_parse_single_session() { - let json = serde_json::json!([{ - "model_name": "GPT-5.4", - "session_id": "test-session-1", - "usage_time": 1776000000, - "dollar_float": 0.5, - "extra_info": { - "input_token": 1000, - "output_token": 500, - "cache_read_token": 200, - "cache_write_token": 100 - } - }]); - let f = write_fixture(&json.to_string()); - let msgs = parse_trae_file("trae", f.path()); - assert_eq!(msgs.len(), 1); - let m = &msgs[0]; - assert_eq!(m.client.as_ref(), "trae"); - assert_eq!(m.model_id.as_ref(), "gpt-5.4"); - assert_eq!(m.provider_id.as_ref(), "openai"); - assert_eq!(m.tokens.input, 1000); - assert_eq!(m.tokens.output, 500); - assert_eq!(m.tokens.cache_read, 200); - assert_eq!(m.tokens.cache_write, 100); - assert_eq!(m.cost, 0.0); - // timestamp: epoch seconds → ms - assert_eq!(m.timestamp, 1_776_000_000_000); - } - - #[test] - fn test_skip_zero_token_session() { - let json = serde_json::json!([{ - "model_name": "GPT-5.4", - "session_id": "empty-session", - "usage_time": 1776000000, - "dollar_float": 0.0, - "extra_info": { - "input_token": 0, - "output_token": 0, - "cache_read_token": 0, - "cache_write_token": 0 - } - }]); - let f = write_fixture(&json.to_string()); - let msgs = parse_trae_file("trae", f.path()); - assert!(msgs.is_empty()); - } - - #[test] - fn zero_usage_without_attribution_is_an_intentional_filter() { - let json = serde_json::json!([{ - "extra_info": { - "input_token": 0, - "output_token": 0, - "cache_read_token": 0, - "cache_write_token": 0 - } - }]); - let f = write_fixture(&json.to_string()); - - let scanned = super::parse_trae_file("trae", f.path()).unwrap(); - - assert!(scanned.messages.is_empty()); - assert!(scanned.rejections.is_empty()); - assert!(scanned.interrupted.is_none()); - } - - #[test] - fn test_normalize_model_names() { - assert_eq!(normalize_trae_model("GPT-5.4"), "gpt-5.4"); - assert_eq!(normalize_trae_model("GPT-5.3-Codex"), "gpt-5.3-codex"); - assert_eq!(normalize_trae_model("GPT-5.3 Codex"), "gpt-5.3-codex"); - assert_eq!(normalize_trae_model("Gemini 3.1 Pro"), "gemini-3.1-pro"); - assert_eq!(normalize_trae_model("GLM 5.1"), "glm-5.1"); - assert_eq!(normalize_trae_model("Unknown Model"), "Unknown Model"); - } - - #[test] - fn test_provider_mapping() { - assert_eq!(provider_for_model("GPT-5.4"), "openai"); - assert_eq!(provider_for_model("Claude Sonnet 4.6"), "anthropic"); - assert_eq!(provider_for_model("Gemini 3.1 Pro"), "google"); - assert_eq!(provider_for_model("GLM 5.1"), "zhipu"); - assert_eq!(provider_for_model("SomeOtherModel"), "trae"); - } - - #[test] - fn test_auto_mode_without_concrete_model_is_rejected() { - let json = serde_json::json!([{ - "model_name": "", - "mode": "Auto", - "session_id": "auto-session-1", - "usage_time": 1776000000, - "dollar_float": 0.27, - "extra_info": { - "input_token": 159213, - "output_token": 210, - "cache_read_token": 6144, - "cache_write_token": 0 - } - }]); - let f = write_fixture(&json.to_string()); - assert_eq!(rejection_keys("trae", f.path()), ["missing-model"]); - } - - #[test] - fn test_reject_session_without_session_id() { - // A record without `session_id` would otherwise dedup to the same - // key as every other malformed record. Reject it instead. - let json = serde_json::json!([{ - "model_name": "GPT-5.4", - "usage_time": 1776000000, - "dollar_float": 0.1, - "extra_info": { "input_token": 100, "output_token": 1, "cache_read_token": 0, "cache_write_token": 0 } - }]); - let f = write_fixture(&json.to_string()); - assert_eq!(rejection_keys("trae", f.path()), ["malformed-record"]); - } - - #[test] - fn test_reject_session_without_usage_time() { - // No `usage_time` → would land at epoch 0. Reject it. - let json = serde_json::json!([{ - "model_name": "GPT-5.4", - "session_id": "abc", - "dollar_float": 0.1, - "extra_info": { "input_token": 100, "output_token": 1, "cache_read_token": 0, "cache_write_token": 0 } - }]); - let f = write_fixture(&json.to_string()); - assert_eq!(rejection_keys("trae", f.path()), ["missing-timestamp"]); - } - - #[test] - fn test_reject_session_with_non_positive_usage_time() { - let json = serde_json::json!([{ - "model_name": "GPT-5.4", - "session_id": "abc", - "usage_time": 0, - "dollar_float": 0.1, - "extra_info": { "input_token": 100, "output_token": 1, "cache_read_token": 0, "cache_write_token": 0 } - }]); - let f = write_fixture(&json.to_string()); - assert_eq!(rejection_keys("trae", f.path()), ["missing-timestamp"]); - } - - #[test] - fn test_reject_session_with_overflowing_usage_time() { - // A maliciously crafted cache could contain a near-MAX `usage_time`. - // Multiplying by 1000 would overflow `i64` — debug-panic or wrap to - // a negative timestamp. Reject the record instead. - let json = serde_json::json!([{ - "model_name": "GPT-5.4", - "session_id": "evil", - "usage_time": i64::MAX, - "dollar_float": 0.1, - "extra_info": { "input_token": 100, "output_token": 1, "cache_read_token": 0, "cache_write_token": 0 } - }]); - let f = write_fixture(&json.to_string()); - assert_eq!(rejection_keys("trae", f.path()), ["missing-timestamp"]); - } - - #[test] - fn test_missing_model_and_mode_is_rejected() { - let json = serde_json::json!([{ - "session_id": "no-meta", - "usage_time": 1776000000, - "dollar_float": 0.01, - "extra_info": { "input_token": 100, "output_token": 1, "cache_read_token": 0, "cache_write_token": 0 } - }]); - let f = write_fixture(&json.to_string()); - assert_eq!(rejection_keys("trae", f.path()), ["missing-model"]); - } - - #[test] - fn malformed_session_does_not_discard_neighboring_sessions() { - let json = serde_json::json!([ - { - "model_name": "GPT-5.4", - "session_id": "good-before", - "usage_time": 1776000000_i64, - "extra_info": { "input_token": 100, "output_token": 1 } - }, - { - "model_name": "GPT-5.4", - "session_id": "bad", - "usage_time": 1776000001_i64, - "extra_info": { "input_token": "not-an-integer" } - }, - { - "model_name": "Claude Sonnet 4.6", - "session_id": "good-after", - "usage_time": 1776000002_i64, - "extra_info": { "input_token": 200, "output_token": 2 } - } - ]); - let f = write_fixture(&json.to_string()); - - let scanned = super::parse_trae_file("trae", f.path()).unwrap(); - - assert_eq!(scanned.messages.len(), 2); - assert_eq!(scanned.messages[0].session_id.as_ref(), "good-before"); - assert_eq!(scanned.messages[1].session_id.as_ref(), "good-after"); - let rejection = scanned.rejections.entries().next().unwrap(); - assert_eq!(rejection.key, "malformed-record"); - assert_eq!(rejection.count, 1); - assert!(scanned.interrupted.is_none()); - } -} diff --git a/docs/adr/0007-client-identity-catalog.md b/docs/adr/0007-client-identity-catalog.md index fb77ba76c..aac05836f 100644 --- a/docs/adr/0007-client-identity-catalog.md +++ b/docs/adr/0007-client-identity-catalog.md @@ -8,6 +8,10 @@ Superseded in part by ADR 0015: the hosted frontend registry and Narrowed by ADR 0012: excluded clients do not retain catalog-only identities in this local-only fork. +Narrowed by ADR 0024's Subscription Usage boundary: every remaining catalog +identity participates in ordinary local reports; the former `parse_local` +capability split has been removed. + ## Decision Use `crates/tokscale-core/client-catalog.json` as the canonical source for @@ -32,11 +36,6 @@ the same `ClientId` set without duplicates. Identity-only, remote-only, and display-placeholder catalog entries require a new explicit decision rather than a capability branch in callers. -`parse_local` only controls whether an integration participates in ordinary -reports without an explicit client filter. It does not indicate whether an -adapter exists; for example, Cursor remains an accepted adapter-backed client -while opting out of the ordinary no-filter report set. - `ClientId` is the only Rust client identity type. Do not add a second enum, hand-written base-client list, or hidden per-client CLI flag set. diff --git a/docs/adr/0014-explicit-subscription-usage-boundary.md b/docs/adr/0014-explicit-subscription-usage-boundary.md index 001b22efe..bdebd608d 100644 --- a/docs/adr/0014-explicit-subscription-usage-boundary.md +++ b/docs/adr/0014-explicit-subscription-usage-boundary.md @@ -41,6 +41,9 @@ Tokscale accepts purpose-specific subscription credentials for plan quota lookup General provider API keys are intentionally ignored for these quota lookups. Examples include `ZAI_API_KEY`, `GLM_API_KEY`, `KIMI_API_KEY`, `MINIMAX_API_KEY`, and `MINIMAX_API_TOKEN`. +Credential ownership and persistence follow ADR 0023. The broader Usage +subsystem boundary, including removal of Cursor and Trae, follows ADR 0024. + ## Configuration Policy Canonical TUI provider IDs are: diff --git a/docs/adr/0015-local-only-product-surface.md b/docs/adr/0015-local-only-product-surface.md index 9d1f63259..c1988077e 100644 --- a/docs/adr/0015-local-only-product-surface.md +++ b/docs/adr/0015-local-only-product-surface.md @@ -4,6 +4,9 @@ Accepted. +Superseded in part by ADR 0024's Subscription Usage boundary: Cursor and Trae +are no longer maintained integrations. + ## Context This fork is maintained for local CLI and TUI usage accounting on diff --git a/docs/adr/0018-bounded-source-fold-pipeline.md b/docs/adr/0018-bounded-source-fold-pipeline.md index a8bc46a7d..3ae6a6685 100644 --- a/docs/adr/0018-bounded-source-fold-pipeline.md +++ b/docs/adr/0018-bounded-source-fold-pipeline.md @@ -5,6 +5,9 @@ Status: Accepted OpenCode's retired JSON source-class and precedence details are superseded by ADR 0019; this document reflects the current SQLite-only fold contract. +Trae-specific fold clauses are superseded by ADR 0024's Subscription Usage +boundary; the integration and its fold state have been removed. + ## Context Source discovery already produced an ordered inventory, but execution parsed every @@ -16,9 +19,9 @@ finished. The fold also carries observable adapter-specific semantics. Codex, Claude, Hermes, Antigravity, OpenCode, and CodeBuddy deduplicate across source units; -Trae selects one latest message per session; and OMP emits cache hits before -misses while using one parent-task index for all misses. A bounded -implementation must preserve those rules across batch boundaries. +OMP emits cache hits before misses while using one parent-task index for all +misses. A bounded implementation must preserve those rules across batch +boundaries. ## Decision @@ -47,8 +50,7 @@ Execute each prepared adapter group as ordered, bounded batches. miss batch is parsed. Existing class-precedence rules are retained as described below. - Deduplication and merge state is created once per adapter group and survives - every batch. Trae retains only the current latest message per session until - its final sorted emission. + every batch. - OpenCode retains one deduplication set across all current-format SQLite databases and every batch. OMP retains its dedicated lightweight whole-group cache-hit/miss plan, builds one parent-task index from all miss paths, then diff --git a/docs/adr/0022-deterministic-cli-command-semantics.md b/docs/adr/0022-deterministic-cli-command-semantics.md new file mode 100644 index 000000000..9d82f4b66 --- /dev/null +++ b/docs/adr/0022-deterministic-cli-command-semantics.md @@ -0,0 +1,115 @@ +# ADR 0022: Deterministic CLI command semantics + +Status: Accepted + +## Context + +The v4 CLI mixed interactive navigation and report generation in the same +argument space. `tokscale models` could open a TUI on a terminal but print a +report in a pipe, root-level flags were copied into multiple execution paths, +and some successfully parsed options were ignored by the selected path. The +`--light` display flag also controlled whether a report wrote the TUI aggregate +cache. This made command meaning depend on TTY state, argument position, and +unrelated presentation choices. + +Those are not compatibility conveniences. They make automation impossible to +reason about and allow the parser to claim an option was accepted without a +single authoritative owner applying it. + +## Decision + +### Commands have one role + +The bare `tokscale` command is an exact shorthand for an unconfigured +`tokscale tui`. The root owns only help, version, and subcommand selection. Any +TUI option requires the explicit `tui` subcommand. + +`models`, `monthly`, `hourly`, and `time-metrics` are report commands. They +always emit a human-readable table by default and a JSON document with +`--json`; their function never changes with TTY state. Opening a report tab is +spelled `tokscale tui --tab `. + +TTY detection may control terminal presentation such as color and progress, +but it may not select a different command. A TUI requires interactive stdin +and stdout; otherwise it fails as invalid usage with a report-command hint. +Requesting a disabled optional tab also fails instead of silently opening a +different tab. + +### Every option has one owner + +Business options live on the narrowest command that applies them. Shared local +source scope consists of `--home` and repeatable or comma-separated +`--client`; shared date scope consists of one preset or an inclusive +`--since`/`--until` range. `--group-by` belongs only to `models`. `--json`, +`--benchmark`, and `--no-spinner` belong only to report commands that use +them. `--theme`, `--refresh`, `--no-refresh`, `--debug`, and `--tab` belong +only to `tui`. + +Parsing is followed by one resolve-and-validate step that produces a typed +`ExecutionPlan`. Execution consumes that plan and does not inspect Clap state +or TTY state again. The invariant is: + +> Every explicit argument accepted by the parser must change the execution +> plan; otherwise parsing must fail. + +An explicit `--home` must name an existing directory and is authoritative for +source discovery and settings. It never falls back to the process home or +client-specific environment roots. Client ids are canonicalized and +deduplicated. Date presets are mutually exclusive, dates use local-time +inclusive boundaries, and `since` may not be later than `until`. + +### Output and failures are stable + +Stdout contains only the command's primary product. Progress, benchmark +timing, health summaries, warnings, and errors use stderr. Local JSON commands +emit one common envelope: + +```json +{ + "data": {}, + "health": {}, + "metadata": { "processingTimeMs": 0 } +} +``` + +Third-party record or source damage remains in `health` under ADR 0021 and +does not change a successfully produced report's exit code. Invalid CLI usage +or environment is exit code `2`; internal, I/O, network, and authentication +failures are exit code `1`; user interruption remains `130` where the child or +terminal supplies it. + +Invalid environment includes malformed or out-of-range environment variables +and settings values. Failure to read or write an otherwise valid settings path +is an operational I/O failure, so it remains exit code `1`. + +Inside the TUI, `q` is an ordinary successful quit and returns `0`. `Ctrl-C` +is a typed user interruption and returns `130`, but only after terminal modes +and the alternate screen have been restored. + +### Explicit maintenance and leaf commands + +`graph` always produces JSON. Without `--output` it writes the document to +stdout; with `--output` it writes the file and prints only the final path to +stdout. + +Pricing is `pricing lookup ` or `pricing overrides`; the lookup's +catalog selector is named `--source`. Cache maintenance is `cache warm` or +`cache prune`. Reports never write the TUI aggregate cache, and the removed +`--write-cache`, `--no-write-cache`, and `light.writeCache` controls have no +replacement inside a report command. `headless` requires `--` between +Tokscale options and the child command. + +The old spellings are not aliases and are never rewritten into a successful +command. Known v4 invocations may receive one migration hint only after Clap +rejects them. + +## Consequences + +Scripts can determine output shape from argv alone, and help output exposes +only options the selected command will execute. TUI navigation is slightly +more verbose but unambiguous. The aggregate cache becomes an explicit product +boundary instead of a side effect of table rendering. + +This is a breaking CLI change and must ship in the next major release. The +version bump remains a separate release change so merging the implementation +does not implicitly publish packages before the release checks are complete. diff --git a/docs/adr/0023-provider-owned-credentials.md b/docs/adr/0023-provider-owned-credentials.md new file mode 100644 index 000000000..c65f2f326 --- /dev/null +++ b/docs/adr/0023-provider-owned-credentials.md @@ -0,0 +1,85 @@ +# ADR 0023: Provider-owned credentials + +Status: Accepted + +The Cursor subsection is superseded by ADR 0024: Cursor support has been +removed and the broader Subscription Usage redesign now owns that product +surface. The Codex credential boundary remains current. + +## Context + +Tokscale reports local token usage and optional subscription quota data. Some +upstream integrations expanded that role into account management: + +- Cursor asked users to paste a browser `WorkosCursorSessionToken`, copied it + into `~/.config/tokscale/cursor-credentials.json`, managed multiple accounts, + and made implicit Cursor API requests before local reports and the TUI. +- Codex copied the complete OAuth token set from the provider-owned + `auth.json` into `~/.config/tokscale/codex-credentials.json`, then exposed + import, switch, remove, and multi-account status commands. + +Neither credential store is required to parse local data. The extra copies +increase the secret-bearing surface, blur ownership, and let an analytics tool +modify account state. File mode `0600` reduces exposure but does not justify a +second plaintext copy of an access token, refresh token, ID token, or browser +session cookie. + +## Decision + +Tokscale is a credential consumer, not an account manager. + +- Authentication remains owned by the provider CLI, provider application, OS + keychain, or an explicitly named environment variable. +- Tokscale does not copy provider credentials into its configuration or cache + directories. +- Tokscale does not switch provider accounts, implement provider login/logout, + or refresh and rewrite provider OAuth credentials. +- A cached quota or usage result may be persisted only when it contains no + credential or raw authentication response. +- Missing, expired, or rejected credentials are reported explicitly. Users + repair authentication with the provider's own tool. + +### Cursor + +The `tokscale cursor` account-management and API-sync namespace is removed. +Local reports may still parse previously supplied Cursor usage CSV files under +`~/.config/tokscale/cursor-cache/usage*.csv`, because those files contain usage +data rather than credentials. Reports and the TUI never use a Tokscale-owned +Cursor credential file and never contact Cursor implicitly. + +Legacy `cursor-credentials.json` files are ignored. Tokscale does not silently +delete user files during startup; users may remove the obsolete file after +upgrading. + +### Codex and ChatGPT subscription usage + +Codex quota lookup reads only the current provider-owned authentication source: + +1. `$CODEX_HOME/auth.json` when `CODEX_HOME` is explicitly set; +2. `~/.config/codex/auth.json`; +3. `~/.codex/auth.json`; +4. the provider's macOS Keychain entry. + +Only the access token and account id required for the usage request are +deserialized. Tokscale does not deserialize the refresh token or ID token, +does not refresh OAuth, and never writes an auth file. The `tokscale codex` +multi-account namespace and `codex-credentials.json` store are removed. + +Legacy `codex-credentials.json` files are ignored and may be removed after the +provider-owned authentication has been verified. + +### Subscription Usage redesign + +ADR 0014 continues to govern when remote subscription requests may occur. +ADR 0024 defines the subsystem boundary. Multiple Z.ai Coding Plan keys and the +broader account/plan presentation model remain in design under issue #146; +that design must reference external secrets rather than store key values in +Tokscale. + +## Consequences + +Tokscale no longer offers Cursor or Codex account switching. Cursor API export +does not refresh automatically, and an expired Codex access token requires the +user to authenticate with Codex again. In exchange, local report commands no +longer create, retain, refresh, or mutate these providers' secrets, and command +behavior matches Tokscale's analytics role. diff --git a/docs/adr/0024-subscription-usage-redesign.md b/docs/adr/0024-subscription-usage-redesign.md new file mode 100644 index 000000000..153aee07b --- /dev/null +++ b/docs/adr/0024-subscription-usage-redesign.md @@ -0,0 +1,124 @@ +# ADR 0024: Subscription Usage product boundary and redesign + +## Status + +Accepted foundation. The provider, account, plan, cache, and presentation +details remain in design under issue #146. + +## Context + +Tokscale has two different data products that accumulated overlapping names +and responsibilities: + +- Local reports parse provider-owned transcripts, databases, and session + files to account for token usage. +- Subscription Usage contacts provider APIs to display account-level quota, + limits, reset windows, and plan state. + +Subscription Usage also accumulated credential copying, account switching, +provider-specific sync commands, inconsistent caches, and presentation models +that did not share one account or plan identity. Cursor and Trae expanded this +surface further by mixing local parsing with Tokscale-owned credentials and +remote synchronization. Treating each removal as an isolated integration +cleanup would leave the subsystem without a durable product boundary. + +This ADR records that boundary and the accepted direction for the whole Usage +subsystem. It does not pretend that the final provider/account/plan model is +already designed. + +## Scope + +Subscription Usage includes: + +- the `tokscale usage` report; +- the optional TUI Usage tab; +- remote subscription and coding-plan provider adapters; +- discovery of externally owned credentials; +- normalized quota-result caching; +- provider, account, and plan identity presented by the CLI and TUI. + +Local transcript and session reports remain a separate product surface. They +are affected only where an old integration blurred the boundary between local +parsing, remote quota lookup, and account management. + +## Decision + +### Stable product boundaries + +- Local report refresh never performs a remote Subscription Usage request. + Remote access remains explicit and bounded by ADR 0014. +- Tokscale is a credential consumer, not an account manager, as established by + ADR 0023. It does not copy provider credentials, log users in or out, switch + accounts, refresh OAuth credentials, or rewrite provider authentication. +- Subscription credentials remain in provider-owned files, OS credential + stores, or explicitly named environment variables. Tokscale settings and + caches must not contain secret values. +- A cache may persist only a normalized quota result and its freshness + metadata. It must not persist access tokens, refresh tokens, session cookies, + API keys, or raw authentication responses. +- Provider and account failures are isolated. One unavailable or malformed + provider result must not hide healthy Usage data from other providers or + accounts. +- CLI and TUI presentations must consume one provider/account/plan domain + model. Rendering code must not invent a second identity scheme. + +### Cursor and Trae removal + +Cursor and Trae are removed from the maintained product surface as the first +application of this boundary: + +- Remove them from the client catalog, local scan definitions, adapters, + parsers, reports, TUI, Wrapped, CLI namespaces, assets, tests, and user + documentation. +- `cursor` and `trae` are invalid client ids. Tokscale does not discover their + old cache directories, copy their credentials, refresh tokens, or contact + either service. +- Persisted parser discriminants remain only as retired cache-format tags + until the next cache-format break. No active adapter may request them. +- Legacy Tokscale-owned files are ignored rather than deleted automatically; + deleting user files is an explicit maintenance action. + +Codex and ChatGPT subscription lookup reads the active provider-owned Codex +authentication source under ADR 0023. It does not restore Tokscale account +switching or a second credential store. + +The existing `TOKSCALE_USAGE_ZAI_CODING_PLAN_API_KEY` contract remains valid. +Future support for multiple Z.ai plans must reference multiple external +secrets rather than copy their values into Tokscale. + +## Design still in progress + +Issue #146 owns the detailed redesign, including: + +- the provider/account/plan identity and ordering model; +- which existing providers remain supported; +- multiple external-secret references and deduplication; +- normalized cache schema and fresh, stale, and unavailable semantics; +- refresh scheduling and the optional Usage-tab lifecycle; +- CLI JSON, TUI rendering, redaction, and partial-failure presentation. + +Those decisions must preserve the boundaries above. They may refine or +supersede this ADR once the complete model is known; they must not reintroduce +Tokscale-owned secrets or account-management commands as incidental provider +features. + +## Relationship to earlier decisions + +This ADR extends ADR 0014's explicit remote-access boundary and ADR 0023's +provider-owned credential policy. It supersedes the Cursor and Trae clauses in +ADR 0015, ADR 0018, and ADR 0023, and removes the `parse_local` capability split +from ADR 0007 because every remaining catalog identity is an ordinary local +report source. + +## Consequences + +The fork currently has no Cursor or Trae local reporting, remote Usage, +account management, sync command, or TUI presence. Existing commands and +settings that name them fail as invalid usage instead of producing an empty +success. + +The Usage subsystem is intentionally an active redesign area rather than a +finished provider list frozen by this ADR. New Usage work must first identify +credential ownership, remote-request consent, cache contents, failure +isolation, and provider/account/plan identity before adding presentation or +convenience commands. diff --git a/docs/cli.md b/docs/cli.md index 6f7f8cc7d..466905e6f 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -1,70 +1,120 @@ # CLI usage -This page documents the command surface most users need. Run commands from this -fork with `bun run cli --` after building from source, or with `tokscale` when -using an installed binary. +Tokscale separates its interactive interface from report commands. Command +meaning is determined entirely by argv; piping or redirecting output never +selects a different feature. -Use `--no-spinner` in automation so output stays deterministic. +Run commands from a built checkout with `bun run cli --`, or use `tokscale` +with an installed fork package. Pass `--no-spinner` to report commands in +automation. -## Report commands +## Interactive TUI ```bash -# Interactive TUI tokscale tokscale tui -tokscale models -tokscale monthly -tokscale hourly +tokscale tui --tab models +tokscale tui --client opencode,claude --week +tokscale tui --theme blue --refresh 30 +tokscale tui --no-refresh +``` -# Table or JSON output -tokscale --no-spinner --light -tokscale models --no-spinner --json -tokscale monthly --no-spinner --json -tokscale hourly --no-spinner --json +Bare `tokscale` is exactly the default TUI shortcut. Any TUI option requires +the explicit `tui` subcommand. The TUI requires interactive stdin and stdout; +for example, `tokscale | jq` fails and points to `tokscale models --json` +instead of changing into a report command. -# Contribution graph export -tokscale graph --no-spinner --output graph.json +Available `--tab` values are `overview`, `models`, `monthly`, `weekly`, +`daily`, `hourly`, `stats`, `agents`, `issues`, and `usage`. Requesting the +Usage tab while `usageTabEnabled` is false is an error; Tokscale does not +silently open Overview. -# Session time metrics -tokscale time-metrics --no-spinner --json -``` +CLI options override settings for the current TUI process and do not rewrite +`settings.json`. The TUI captures normal mouse input; use the terminal's +modified selection gesture, usually `Shift+drag`, to select terminal text. -The root command defaults to the interactive TUI when stdin/stdout are terminals -and falls back to scriptable output otherwise. +## Local reports -The TUI captures normal mouse input for tabs, filters, and graph cells. Use your -terminal's native modified selection gesture, usually `Shift+drag`, when you -want to select and copy text from the TUI. +```bash +# Human-readable tables +tokscale models --no-spinner +tokscale monthly --no-spinner +tokscale hourly --no-spinner +tokscale time-metrics --no-spinner -## Filters +# Structured output +tokscale models --json +tokscale monthly --json +tokscale hourly --json +tokscale time-metrics --json +``` -Client filters accept comma-separated values or repeated flags: +These commands always produce reports, even when stdout is a terminal. Table +output is the default; the removed `--light` mode is not an alias. To open a +specific TUI view, use `tokscale tui --tab models` or the corresponding tab. -```bash -tokscale --client opencode -tokscale --client opencode,claude -tokscale -c opencode -c claude +All local JSON reports use the same top-level envelope: + +```json +{ + "data": {}, + "health": { + "complete": true, + "cleanSources": 0, + "degradedSources": 0, + "rejectedRecords": 0, + "partialSources": 0, + "failedSources": 0, + "sourceDataBytes": 0, + "issues": [] + }, + "metadata": { + "processingTimeMs": 0 + } +} ``` -Date filters are inclusive and use the local timezone: +Stdout contains only the table or JSON document. Progress, `--benchmark` +timing, health summaries, warnings, and errors go to stderr. A degraded report +still exits `0` when its payload was produced; inspect `health` when automation +must react to rejected records or unavailable sources. + +## Source and date scope + +Local commands that read usage share the same source scope: ```bash -tokscale --today -tokscale --week -tokscale --month -tokscale --since 2026-01-01 --until 2026-01-31 -tokscale --year 2026 +tokscale models --client opencode +tokscale models --client opencode,claude +tokscale models -c opencode -c claude +tokscale models --home /tmp/test-home --no-spinner +tokscale tui --client codex --home /tmp/test-home ``` -For testing or alternate home roots, local report commands accept: +Repeated client ids are deduplicated. Without a CLI filter, Tokscale uses +`defaultClients` when configured and otherwise scans all local clients. An +unknown client is an error. `--home` must be an existing directory and is +authoritative: source discovery does not silently fall back to the process +home or client-specific environment roots. + +Date boundaries are inclusive and use the local timezone: ```bash -tokscale --home /tmp/test-home --no-spinner --json +tokscale models --today +tokscale models --week +tokscale models --month +tokscale models --year 2026 +tokscale models --since 2026-01-01 +tokscale models --until 2026-01-31 +tokscale models --since 2026-01-01 --until 2026-01-31 ``` -## Grouping +Choose one preset or a custom range. Combining presets, combining `--year` +with `--since`/`--until`, or specifying `since > until` is invalid usage. + +## Model grouping -`models` output supports these `--group-by` values: +Only `models` owns `--group-by`: | Strategy | Effect | | --- | --- | @@ -75,134 +125,117 @@ tokscale --home /tmp/test-home --no-spinner --json | `session,model` | One row per session id and model. | | `client,session,model` | One row per client, session id, and model. | -Examples: - ```bash -tokscale models --no-spinner --json --group-by model -tokscale models --no-spinner --json --group-by client,provider,model -tokscale models --no-spinner --json --group-by session,model +tokscale models --json --group-by model +tokscale models --json --group-by client,provider,model ``` -## Inspecting local sources +## Graph and source inspection + +`graph` always produces JSON: ```bash -tokscale clients -tokscale clients --json +tokscale graph --no-spinner +tokscale graph --no-spinner --output graph.json ``` -This shows scan locations and session counts for local clients. +Without `--output`, the JSON document is stdout. With an output file, stdout +contains only the final path and operational details use stderr. -## Cache maintenance +Inspect source locations and counts with: ```bash -tokscale cache prune +tokscale clients +tokscale clients --json +tokscale clients --client codex --home /tmp/test-home ``` -This explicitly scans every source-message cache shard, removes shards whose -source file no longer exists, and removes older parser revisions when a newer -revision exists for the same source and parser. It prints the scanned, removed, -and retained shard counts. Normal reports and TUI loads do not run this full -cache traversal. If a shard cannot be read or decoded, or a selected shard -cannot be removed, the command fails instead of reporting a partial success. - -## Pricing lookup +## Wrapped ranking ```bash -tokscale pricing claude-sonnet-4-5 --no-spinner -tokscale pricing grok-code --provider openrouter --no-spinner -tokscale pricing list-overrides --json +tokscale wrapped +tokscale wrapped --ranking agents +tokscale wrapped --ranking clients ``` -Standalone pricing lookup is a catalog query. Local report pricing -canonicalizes parsed model ids before lookup, but standalone -`tokscale pricing ` does not replay arbitrary local source cleanup unless -the command itself is explicitly changed later. - -## Integration commands +Without `--ranking`, Wrapped automatically uses OpenCode agent rankings when +agent data exists and otherwise uses client rankings. An explicit +`--ranking agents` never changes into a client ranking: when no agent data is +available, the image keeps the requested panel and renders an explicit empty +state. `--ranking agents` requires OpenCode in an explicit `--client` scope. -Cursor: +## Cache maintenance ```bash -tokscale cursor login --name work -tokscale cursor status -tokscale cursor accounts --json -tokscale cursor sync --json -tokscale cursor switch work -tokscale cursor logout --name work -tokscale cursor logout --all --purge-cache +tokscale cache warm +tokscale cache warm --client codex +tokscale cache prune ``` -Codex account helpers: +`cache warm` explicitly builds the TUI aggregate cache for its source scope. +Report commands never modify that aggregate cache. Source-message shards remain +an internal derived cache and are written automatically while parsing. -```bash -tokscale codex import --name work -tokscale codex accounts --json -tokscale codex switch work -tokscale codex status --json -tokscale codex remove work -``` +`cache prune` traverses source-message shards, removes orphaned sources and +superseded parser revisions, and prints scanned, removed, and retained counts. +Unreadable or unclassifiable shards make the explicit maintenance command fail +instead of reporting partial success. -Antigravity: +## Pricing lookup ```bash -tokscale antigravity status --json -tokscale antigravity sync -tokscale antigravity purge-cache +tokscale pricing lookup claude-sonnet-4-5 --no-spinner +tokscale pricing lookup grok-code --source openrouter --no-spinner +tokscale pricing lookup claude-sonnet-4-5 --json +tokscale pricing overrides +tokscale pricing overrides --json ``` -Trae: +`--source` selects a pricing catalog and is distinct from a model's provider. +Standalone lookup is a catalog query; it does not replay arbitrary cleanup from +a local source parser. -```bash -tokscale trae login -tokscale trae login --manual --variant solo -tokscale trae status --json -tokscale trae sync --since 30 -tokscale trae sync --since 30 --include-aux -tokscale trae logout --variant solo -``` +## Integration and usage commands -Warp/Oz subscription aggregate data: +Tokscale does not provide `cursor`, `trae`, or `codex` account-management +command namespaces. Cursor and Trae are not supported local clients. +`tokscale usage` reads the currently authenticated Codex account from +provider-owned auth state without copying, switching, refreshing, or modifying +its credentials. ```bash -tokscale warp login -tokscale warp login --cookie +# Local integrations with explicit sync workflows +tokscale antigravity status --json +tokscale antigravity sync tokscale warp status --json tokscale warp sync --json -tokscale warp logout --purge-cache -``` - -## Subscription usage - -Subscription usage is separate from local token reports. -```bash +# Subscription quota, separate from local reports tokscale usage tokscale usage --json ``` -The TUI Usage tab is hidden unless `usageTabEnabled` is set in -`settings.json`. +Flags belong to the leaf command that executes them. They cannot be placed on +the root or before the owning subcommand. ## Headless capture -Headless capture currently supports Codex CLI: +Headless capture requires `--` between Tokscale options and the child command: ```bash -tokscale headless codex exec -m gpt-5 "review this change" +tokscale headless codex --format jsonl -- codex exec -m gpt-5 "review this change" ``` -Manual redirect is also possible: +This boundary prevents child flags such as `--json` or `--output` from being +claimed by Tokscale. Set `TOKSCALE_HEADLESS_DIR` to change the capture root. -```bash -mkdir -p ~/.config/tokscale/headless/codex -codex exec --json "review this change" \ - > ~/.config/tokscale/headless/codex/review.jsonl -``` +## Exit codes -Set `TOKSCALE_HEADLESS_DIR` to customize the headless log root. - -## Local-only surface +| Code | Meaning | +| --- | --- | +| `0` | The command produced its result, including an incomplete local report. | +| `1` | Internal, I/O, network, or authentication failure. | +| `2` | Invalid CLI arguments, option combinations, or runtime environment. | +| `130` | User interruption where supplied by the terminal or child process. | -This fork does not expose hosted Tokscale account or submission commands. The -CLI surface is limited to local reports, cache-backed local integrations, -subscription usage helpers, and source-build tooling for this fork. +This fork does not expose hosted login, submission, or leaderboard commands. diff --git a/docs/clients.md b/docs/clients.md index bd68ab715..43676c3ff 100644 --- a/docs/clients.md +++ b/docs/clients.md @@ -21,10 +21,9 @@ When using an installed binary, use `tokscale clients` instead. | `opencode` | OpenCode | `~/.local/share/opencode/opencode*.db` | Reads only current-format SQLite databases and combines multiple release channels when present. | | `claude` | Claude Code | `~/.claude/projects/**/*.jsonl`, `~/.claude/transcripts/**/*.jsonl` | Claude Desktop chat history is not treated as Claude Code token accounting. | | `codex` | Codex CLI | `$CODEX_HOME/sessions/**/*.jsonl`, fallback `~/.codex/sessions/` | Also supports `tokscale headless codex ...` capture. | -| `cursor` | Cursor | `~/.config/tokscale/cursor-cache/usage*.csv` | Reads a local API cache. Logged-in reports and the TUI may auto-refresh stale cache data; local `~/.cursor` state is not parsed. | | `gemini` | Gemini CLI | `$GEMINI_CLI_HOME/tmp/**/chats/*`, fallback `~/.gemini/tmp/` | Reads local chat files. | | `amp` | Amp | `~/.local/share/amp/threads/T-*.json` | Reads local thread files. | -| `droid` | Droid | `~/.factory/sessions/**/*.settings.json` | Reads Factory Droid sessions. | +| `droid` | Droid | `~/.factory/sessions/**/*.settings.json`, related session JSONL and Mission `features.json` | Reads Factory Droid sessions and attributes subagent usage to `Droid Explorer`, `Droid Worker`, `Droid Orchestrator`, or `Droid Validator`. | | `openclaw` | OpenClaw | `~/.openclaw/agents/` plus legacy `.clawdbot`, `.moltbot`, `.moldbot` roots | Reads agent session indexes and JSONL session files. | | `pi` | Pi | `~/.pi/agent/sessions/**/*.jsonl` | Separate from OMP by design. | | `omp` | OMP | `~/.omp/agent/sessions/**/*.jsonl` | Separate from Pi by design. | @@ -44,7 +43,6 @@ When using an installed binary, use `tokscale clients` instead. | `zcode` | ZCode | `~/.zcode/projects/**/*.jsonl` | Reads Z.ai ADE JSONL sessions. | | `kiro` | Kiro | `~/.kiro/sessions/cli/`, `~/.local/share/kiro-cli/data.sqlite3`, and Kiro IDE globalStorage snapshots | Combines CLI and IDE local sources when present. | | `junie` | Junie | `~/.junie/sessions/**/events.jsonl` | Reads JetBrains Junie session events. | -| `trae` | Trae | `~/.config/tokscale/trae-cache/sessions/*.json` | Requires `tokscale trae login` and `tokscale trae sync`. China variants are not supported. | | `cline` | Cline | VS Code globalStorage `saoudrizwan.claude-dev/tasks/**/ui_messages.json` | Same task-log family as Roo Code and KiloCode. | | `commandcode` | Command Code | `~/.commandcode/projects/**/*.jsonl` | Estimated from transcripts. | | `grok` | Grok Build | `$GROK_HOME/sessions/**/updates.jsonl`, fallback `~/.grok/sessions/` | Reads total-token deltas and applies the fixed total-only bucket allocation from ADR 0017. | @@ -108,39 +106,17 @@ Use `TOKSCALE_EXTRA_DIRS` for one-off runs: ```bash TOKSCALE_EXTRA_DIRS='codex:/abs/path/.codex/sessions,gemini:/abs/path/gemini/tmp' \ - tokscale --no-spinner --light + tokscale models --no-spinner ``` ## Cache-backed integrations -Cursor reads a local API cache, but it is not purely manual-sync-backed. Ordinary -local reports and the TUI may call the Cursor API before reading reports when -all of these are true: - -- no `--home` override is active; -- the client filter includes Cursor, including the default unfiltered report; -- saved Cursor credentials exist; -- the expected Cursor cache files are older than five minutes. - -Manual commands are still available: - -```bash -tokscale cursor login --name work -tokscale cursor sync --json -``` - -`tokscale cursor sync --json` forces a refresh. Filtering Cursor out with -`--client` or using `--home` prevents the implicit pre-report refresh. - -Antigravity and Trae are different: they do not refresh from the root report or -TUI command. Run their sync commands before reports when you need fresh data: +Antigravity does not refresh from the root report or TUI command. Run its sync +command before reports when you need fresh data: ```bash tokscale antigravity status tokscale antigravity sync - -tokscale trae login -tokscale trae sync --since 30 ``` `warp` has two separate surfaces. Local reports read `warp.sqlite` when it is diff --git a/docs/configuration.md b/docs/configuration.md index beef67568..648b94b69 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -6,15 +6,6 @@ Tokscale stores most local settings under the platform config directory: - Windows default: `%APPDATA%\tokscale\settings.json` - Override root: `TOKSCALE_CONFIG_DIR` -Known exceptions that are not moved by `TOKSCALE_CONFIG_DIR` today: - -- Cursor integration state: - `$HOME/.config/tokscale/cursor-credentials.json` and - `$HOME/.config/tokscale/cursor-cache/` - -Setting `TOKSCALE_CONFIG_DIR` does not isolate Cursor credentials or Cursor -usage cache today. - ## Example ```json @@ -56,7 +47,6 @@ usage cache today. | `autoRefreshMs` | number | TUI auto-refresh interval in milliseconds. | | `nativeTimeoutMs` | number | Maximum processing time for native subprocess work. | | `defaultClients` | string[] | Client filter used when no `--client/-c` flag is passed. | -| `light.writeCache` | boolean | Allow `tokscale --light` to refresh the TUI startup cache after rendering. | | `usageTabEnabled` | boolean | Show the subscription quota Usage tab in the TUI. | | `usageProviders` | string[] | Explicit allowlist of subscription providers the TUI may fetch. Empty means cache-display mode. | | `scanner.opencodeDbPaths` | string[] | Authoritative additional current-format OpenCode SQLite database files. Missing, unreadable, or obsolete entries fail explicitly. This is the only custom OpenCode scan setting. | @@ -74,7 +64,7 @@ reported explicitly. | Variable | Meaning | | --- | --- | -| `TOKSCALE_CONFIG_DIR` | Overrides the general config/cache root used by Tokscale. Non-empty values are used verbatim. Empty values are treated as unset. It does not currently move Cursor credentials or Cursor usage cache. | +| `TOKSCALE_CONFIG_DIR` | Overrides the general config/cache root used by Tokscale. Non-empty values are used verbatim. Empty values are treated as unset. | | `TOKSCALE_NATIVE_TIMEOUT_MS` | Overrides `nativeTimeoutMs`. | | `TOKSCALE_EXTRA_DIRS` | One-off extra scan roots as `client:/abs/path,client:/abs/path`. | | `TOKSCALE_HEADLESS_DIR` | Overrides the headless capture root. Surrounding whitespace is trimmed; blank values fall back to the default root. | @@ -121,23 +111,18 @@ explicitly want a full traversal that removes classified v2 shards; there is no automatic v2 migration. Retired `source-message-cache.bin` and `source-message-cache.lock` files are not current cache inputs. +The TUI aggregate cache is separate from source-message shards. Reports never +write it; use `tokscale cache warm` when you intentionally want to prebuild it. + Integration roots are mixed state, not all disposable caches: - `antigravity-cache/` contains synced Antigravity artifacts. Use `tokscale antigravity purge-cache` when you want to clear them. -- `trae-cache/` contains both synced Trae usage artifacts and credentials: - `credentials-solo.json` and `credentials-ide.json`. Deleting the directory - can log you out; preserve those files if you only want to clear synced usage. - `warp-cache/` contains both synced Warp aggregate usage and `credentials.json`. Deleting the directory can log you out; use `tokscale warp logout --purge-cache` when you intentionally want to remove credentials and cached usage together. -Cursor is separate from the `TOKSCALE_CONFIG_DIR` roots above: its credentials -and cache live at -`$HOME/.config/tokscale/cursor-credentials.json` and -`$HOME/.config/tokscale/cursor-cache/`, independent of `TOKSCALE_CONFIG_DIR`. - ## Subscription providers Canonical `usageProviders` ids: @@ -158,3 +143,9 @@ warp General-purpose provider API keys such as `ZAI_API_KEY`, `GLM_API_KEY`, `KIMI_API_KEY`, `MINIMAX_API_KEY`, and `MINIMAX_API_TOKEN` are not used for subscription quota lookups. + +Codex subscription usage reads the currently authenticated account from +provider-owned Codex auth state (`$CODEX_HOME/auth.json`, the standard Codex +config locations, or the official macOS keychain item). Tokscale does not copy, +refresh, switch, or modify those credentials. A legacy Tokscale +`codex-credentials.json` is obsolete and ignored by current versions. diff --git a/docs/development.md b/docs/development.md index 92643a083..a4b52577c 100644 --- a/docs/development.md +++ b/docs/development.md @@ -41,7 +41,7 @@ For quick local CLI runs: ```bash bun run cli -bun run cli -- --no-spinner --light +bun run cli -- models --no-spinner ``` ## Test diff --git a/docs/pricing.md b/docs/pricing.md index b166c9d48..97cf6d1d0 100644 --- a/docs/pricing.md +++ b/docs/pricing.md @@ -89,7 +89,7 @@ Overrides are exact-only and case-insensitive: necessarily the raw source label emitted by a client or parser. - For local report overrides, key the entry by that final canonical id unless a parser intentionally preserves the full route. -- `tokscale pricing ` matches the command argument as a catalog query. +- `tokscale pricing lookup ` matches the command argument as a catalog query. - Full gateway paths are only needed when you intentionally query or override that exact route as a catalog key. @@ -110,9 +110,9 @@ lookup or report that needs pricing. ## Standalone lookup ```bash -tokscale pricing claude-sonnet-4-5 --no-spinner -tokscale pricing grok-code --provider openrouter --no-spinner -tokscale pricing list-overrides --json +tokscale pricing lookup claude-sonnet-4-5 --no-spinner +tokscale pricing lookup grok-code --source openrouter --no-spinner +tokscale pricing overrides --json ``` Standalone lookup does not infer arbitrary source prefixes, route prefixes, diff --git a/packages/cli/package.json b/packages/cli/package.json index 0e9f70069..7c47a9fe5 100644 --- a/packages/cli/package.json +++ b/packages/cli/package.json @@ -41,7 +41,6 @@ "tokscale", "opencode", "claude-code", - "cursor", "gemini", "codex", "cli" diff --git a/packages/tokscale/package.json b/packages/tokscale/package.json index 19eb73816..bc2dc366f 100644 --- a/packages/tokscale/package.json +++ b/packages/tokscale/package.json @@ -23,7 +23,6 @@ "claude", "openai", "gemini", - "cursor", "cli" ], "dependencies": {