From 3a199884fcfd815244da0d87aea85f38dd0e4dad Mon Sep 17 00:00:00 2001 From: makoMakoGo <48956204+makoMakoGo@users.noreply.github.com> Date: Thu, 21 May 2026 14:08:03 +0800 Subject: [PATCH 1/2] fix(amp): parse exported thread usage --- crates/tokscale-cli/src/main.rs | 20 +- crates/tokscale-core/src/clients.rs | 6 +- crates/tokscale-core/src/lib.rs | 248 ++++---- crates/tokscale-core/src/scanner.rs | 32 +- crates/tokscale-core/src/sessions/amp.rs | 747 +++++++---------------- 5 files changed, 397 insertions(+), 656 deletions(-) diff --git a/crates/tokscale-cli/src/main.rs b/crates/tokscale-cli/src/main.rs index 22f5c62a3..e6967033c 100644 --- a/crates/tokscale-cli/src/main.rs +++ b/crates/tokscale-cli/src/main.rs @@ -2916,8 +2916,16 @@ fn run_clients_command(json: bool) -> Result<()> { let clients: Vec = ClientId::iter() .map(|client| { - let sessions_path = client.data().resolve_path(&home_dir.to_string_lossy()); - let sessions_path_exists = Path::new(&sessions_path).exists(); + let (sessions_path, sessions_path_exists) = if client == ClientId::Amp { + ( + tokscale_core::sessions::amp::amp_source_label().to_string(), + tokscale_core::sessions::amp::amp_cli_available(), + ) + } else { + let sessions_path = client.data().resolve_path(&home_dir.to_string_lossy()); + let sessions_path_exists = Path::new(&sessions_path).exists(); + (sessions_path, sessions_path_exists) + }; let additional_paths: Vec = built_in_extra_paths .iter() .filter(|(c, _)| *c == client) @@ -3059,10 +3067,16 @@ fn run_clients_command(json: bool) -> Result<()> { for row in clients { println!(" {}", row.label.white()); + let source_label = if row.client == "amp" { + "source" + } else { + "sessions" + }; println!( " {}", format!( - "sessions: {}", + "{}: {}", + source_label, describe_path(&row.sessions_path, row.sessions_path_exists) ) .bright_black() diff --git a/crates/tokscale-core/src/clients.rs b/crates/tokscale-core/src/clients.rs index 805689731..286209896 100644 --- a/crates/tokscale-core/src/clients.rs +++ b/crates/tokscale-core/src/clients.rs @@ -215,9 +215,9 @@ define_clients!( }, Amp = 5 => { id: "amp", - root: PathRoot::XdgData, - relative: "amp/threads", - pattern: "T-*.json", + root: PathRoot::Home, + relative: ".local/bin/amp", + pattern: "", headless: false, parse_local: true, submit_default: true diff --git a/crates/tokscale-core/src/lib.rs b/crates/tokscale-core/src/lib.rs index 14c3ed42f..bae6c80a6 100644 --- a/crates/tokscale-core/src/lib.rs +++ b/crates/tokscale-core/src/lib.rs @@ -643,7 +643,7 @@ fn parse_all_messages_with_pricing( home_dir: &str, clients: &[String], pricing: Option<&pricing::PricingService>, -) -> Vec { +) -> Result, String> { parse_all_messages_with_pricing_with_env_strategy( home_dir, clients, @@ -659,7 +659,7 @@ fn parse_all_messages_with_pricing_with_env_strategy( pricing: Option<&pricing::PricingService>, use_env_roots: bool, scanner_settings: &scanner::ScannerSettings, -) -> Vec { +) -> Result, String> { #[derive(Debug)] struct CachedParseOutcome { messages: Vec, @@ -986,6 +986,7 @@ fn parse_all_messages_with_pricing_with_env_strategy( let mut all_messages: Vec = Vec::new(); let include_all = clients.is_empty(); let include_synthetic = include_all || clients.iter().any(|c| c == "synthetic"); + let include_amp = include_all || clients.iter().any(|c| c == ClientId::Amp.as_str()); // Parse OpenCode: prefer SQLite, collapse forked SQLite history there, then // suppress legacy JSON overlap by message identity. @@ -1154,20 +1155,11 @@ fn parse_all_messages_with_pricing_with_env_strategy( } } - let amp_outcomes: Vec = scan_result - .get(ClientId::Amp) - .par_iter() - .map(|path| { - load_or_parse_source(path, &source_cache, pricing, |path| { - sessions::amp::parse_amp_file(path) - }) - }) - .collect(); - for outcome in amp_outcomes { - all_messages.extend(outcome.messages); - if let Some(entry) = outcome.cache_entry { - source_cache.insert(entry); - } + if include_amp && should_collect_amp_cli_source(home_dir, use_env_roots) { + let mut amp_messages = + sessions::amp::parse_amp_threads_from_cli().map_err(|err| err.to_string())?; + apply_pricing_to_messages(&mut amp_messages, pricing); + all_messages.extend(amp_messages); } let codebuff_outcomes: Vec = scan_result @@ -1449,7 +1441,7 @@ fn parse_all_messages_with_pricing_with_env_strategy( source_cache.save_if_dirty(); - all_messages + Ok(all_messages) } fn filter_unified_messages( @@ -1620,7 +1612,7 @@ pub async fn get_model_report(options: ReportOptions) -> Result Result Result bool { + let home_path = Path::new(home_dir); + use_env_roots + && dirs::home_dir() + .map(|current_home| current_home == home_path) + .unwrap_or(false) + && home_path.join(".local/share/amp/session.json").is_file() +} + fn pricing_multiplier(message: &UnifiedMessage) -> f64 { // Zed bills hosted models at provider list price + 10%. // Source: https://zed.dev/docs/ai/plans-and-usage and https://zed.dev/docs/ai/models @@ -2005,7 +2006,7 @@ fn parse_local_unified_messages_resolved( pricing, options.use_env_roots, &options.scanner_settings, - ); + )?; Ok(filter_unified_messages(messages, &options)) } pub fn parse_local_clients(options: LocalParseOptions) -> Result { @@ -2023,6 +2024,7 @@ pub fn parse_local_clients(options: LocalParseOptions) -> Result Result = scan_result - .get(ClientId::Amp) - .par_iter() - .flat_map(|path| { - sessions::amp::parse_amp_file(path) + let amp_msgs: Vec = + if include_amp && should_collect_amp_cli_source(&home_dir, options.use_env_roots) { + sessions::amp::parse_amp_threads_from_cli() + .map_err(|err| err.to_string())? .into_iter() .map(|msg| unified_to_parsed(&msg)) - .collect::>() - }) - .collect(); + .collect() + } else { + Vec::new() + }; let amp_count = amp_msgs.len() as i32; counts.set(ClientId::Amp, amp_count); messages.extend(amp_msgs); @@ -3337,7 +3339,8 @@ mod tests { temp_dir.path().to_str().unwrap(), &["cursor".to_string()], Some(&pricing), - ); + ) + .unwrap(); assert_eq!(messages.len(), 1); assert_eq!(messages[0].client, "cursor"); @@ -3397,7 +3400,8 @@ mod tests { source_home.path().to_str().unwrap(), &["opencode".to_string()], None, - ); + ) + .unwrap(); assert_eq!(messages.len(), 1); assert_ne!(messages[0].date, "1900-01-01"); @@ -3459,7 +3463,8 @@ mod tests { source_home.path().to_str().unwrap(), &["opencode".to_string()], None, - ); + ) + .unwrap(); assert!(first_messages.is_empty()); let cache = message_cache::SourceMessageCache::load(); @@ -3473,7 +3478,8 @@ mod tests { source_home.path().to_str().unwrap(), &["opencode".to_string()], None, - ); + ) + .unwrap(); assert_eq!(second_messages.len(), 1); } @@ -3518,7 +3524,8 @@ mod tests { source_home.path().to_str().unwrap(), &["opencode".to_string()], None, - ); + ) + .unwrap(); assert_eq!(messages.len(), 1); let loaded = message_cache::SourceMessageCache::load(); @@ -3585,7 +3592,8 @@ mod tests { source_home.path().to_str().unwrap(), &["opencode".to_string()], None, - ); + ) + .unwrap(); assert_eq!(first_messages.len(), 1); conn.execute( @@ -3599,7 +3607,8 @@ mod tests { source_home.path().to_str().unwrap(), &["opencode".to_string()], None, - ); + ) + .unwrap(); assert_eq!(refreshed_messages.len(), 2); } @@ -3693,7 +3702,8 @@ mod tests { source_home.path().to_str().unwrap(), &["opencode".to_string()], None, - ); + ) + .unwrap(); assert_eq!( messages.len(), 3, @@ -3711,7 +3721,8 @@ mod tests { source_home.path().to_str().unwrap(), &["opencode".to_string()], None, - ); + ) + .unwrap(); assert_eq!( messages_warm.len(), 3, @@ -3789,7 +3800,8 @@ mod tests { source_home.path().to_str().unwrap(), &["opencode".to_string()], None, - ); + ) + .unwrap(); assert_eq!(messages.len(), 3); assert_eq!(messages.iter().map(|m| m.tokens.input).sum::(), 600); @@ -3932,7 +3944,8 @@ mod tests { source_home.path().to_str().unwrap(), &["codex".to_string()], None, - ); + ) + .unwrap(); assert_eq!(messages.len(), 2); assert_eq!( @@ -4003,7 +4016,8 @@ mod tests { source_home.path().to_str().unwrap(), &["codex".to_string()], None, - ); + ) + .unwrap(); assert_eq!( messages.len(), @@ -4124,7 +4138,8 @@ mod tests { source_home.path().to_str().unwrap(), &["codex".to_string()], None, - ); + ) + .unwrap(); assert_eq!(initial_messages.len(), 1); assert_eq!(initial_messages[0].model_id, "gpt-5.4"); assert!(message_cache::SourceMessageCache::load() @@ -4150,13 +4165,15 @@ mod tests { source_home.path().to_str().unwrap(), &["codex".to_string()], None, - ); + ) + .unwrap(); std::env::set_var("HOME", fresh_cache_home.path()); let fresh_messages = parse_all_messages_with_pricing( source_home.path().to_str().unwrap(), &["codex".to_string()], None, - ); + ) + .unwrap(); assert_eq!(warm_messages, fresh_messages); assert_eq!(warm_messages.len(), 2); @@ -4199,7 +4216,8 @@ mod tests { source_home.path().to_str().unwrap(), &["codex".to_string()], None, - ); + ) + .unwrap(); assert_eq!(first_messages.len(), 1); std::thread::sleep(std::time::Duration::from_millis(5)); @@ -4221,13 +4239,15 @@ mod tests { source_home.path().to_str().unwrap(), &["codex".to_string()], None, - ); + ) + .unwrap(); std::env::set_var("HOME", fresh_cache_home.path()); let fresh_messages = parse_all_messages_with_pricing( source_home.path().to_str().unwrap(), &["codex".to_string()], None, - ); + ) + .unwrap(); assert_eq!(warm_messages, fresh_messages); } @@ -4268,7 +4288,8 @@ mod tests { source_home.path().to_str().unwrap(), &["codex".to_string()], None, - ); + ) + .unwrap(); assert_eq!(initial_messages.len(), 1); std::thread::sleep(std::time::Duration::from_millis(5)); @@ -4290,7 +4311,8 @@ mod tests { source_home.path().to_str().unwrap(), &["codex".to_string()], None, - ); + ) + .unwrap(); assert!(message_cache::SourceMessageCache::load() .get(&path) .is_none()); @@ -4300,7 +4322,8 @@ mod tests { source_home.path().to_str().unwrap(), &["codex".to_string()], None, - ); + ) + .unwrap(); assert_eq!(warm_messages, fresh_messages); } @@ -4356,7 +4379,8 @@ mod tests { source_home.path().to_str().unwrap(), &["codex".to_string()], None, - ); + ) + .unwrap(); assert_eq!(messages, expected); } @@ -4392,7 +4416,8 @@ mod tests { source_home.path().to_str().unwrap(), &["codex".to_string()], None, - ); + ) + .unwrap(); assert_eq!(initial_messages.len(), 1); std::thread::sleep(std::time::Duration::from_millis(20)); @@ -4402,14 +4427,16 @@ mod tests { source_home.path().to_str().unwrap(), &["codex".to_string()], None, - ); + ) + .unwrap(); std::env::set_var("HOME", fresh_cache_home.path()); let fresh_messages = parse_all_messages_with_pricing( source_home.path().to_str().unwrap(), &["codex".to_string()], None, - ); + ) + .unwrap(); assert_eq!(warm_messages, fresh_messages); assert_ne!(warm_messages[0].timestamp, initial_messages[0].timestamp); @@ -4452,7 +4479,8 @@ mod tests { source_home.path().to_str().unwrap(), &["codex".to_string()], None, - ); + ) + .unwrap(); assert_eq!(messages.len(), 1); assert_eq!(messages[0].model_id, "gpt-5.4"); @@ -4494,7 +4522,8 @@ mod tests { source_home.path().to_str().unwrap(), &["codex".to_string()], None, - ); + ) + .unwrap(); assert_eq!(initial_messages.len(), 1); assert_eq!(initial_messages[0].model_id, "unknown"); assert!(message_cache::SourceMessageCache::load() @@ -4520,14 +4549,16 @@ mod tests { source_home.path().to_str().unwrap(), &["codex".to_string()], None, - ); + ) + .unwrap(); std::env::set_var("HOME", fresh_cache_home.path()); let fresh_messages = parse_all_messages_with_pricing( source_home.path().to_str().unwrap(), &["codex".to_string()], None, - ); + ) + .unwrap(); assert_eq!(resumed_messages, fresh_messages); assert_eq!(resumed_messages.len(), 1); @@ -4572,7 +4603,8 @@ mod tests { source_home.path().to_str().unwrap(), &["codex".to_string()], None, - ); + ) + .unwrap(); assert_eq!(initial_messages.len(), 1); assert!(message_cache::SourceMessageCache::load() .get(&path) @@ -4598,14 +4630,16 @@ mod tests { source_home.path().to_str().unwrap(), &["codex".to_string()], None, - ); + ) + .unwrap(); std::env::set_var("HOME", fresh_cache_home.path()); let fresh_messages = parse_all_messages_with_pricing( source_home.path().to_str().unwrap(), &["codex".to_string()], None, - ); + ) + .unwrap(); assert_eq!(warm_messages, fresh_messages); assert_eq!(warm_messages.len(), 2); @@ -4648,7 +4682,8 @@ mod tests { source_home.path().to_str().unwrap(), &["cursor".to_string()], Some(&pricing), - ); + ) + .unwrap(); assert_eq!(repriced_messages.len(), 1); assert!(repriced_messages[0].cost > 0.0); @@ -4656,7 +4691,8 @@ mod tests { source_home.path().to_str().unwrap(), &["cursor".to_string()], None, - ); + ) + .unwrap(); assert_eq!(cached_messages.len(), 1); assert_eq!(cached_messages[0].cost, 0.0); @@ -5198,7 +5234,8 @@ mod tests { temp_dir.path().to_str().unwrap(), &["synthetic".to_string()], Some(&pricing), - ); + ) + .unwrap(); assert_eq!(messages.len(), 1); assert_eq!(messages[0].client, "opencode"); @@ -5255,7 +5292,8 @@ mod tests { temp_dir.path().to_str().unwrap(), &["synthetic".to_string()], Some(&pricing), - ); + ) + .unwrap(); assert_eq!( messages.len(), @@ -5521,59 +5559,29 @@ mod tests { } #[test] - fn test_parse_local_clients_amp_partial_ledger_recovers_message_fallback_day() { - use chrono::TimeZone; - + fn test_parse_local_clients_amp_ignores_legacy_thread_files() { let temp_dir = tempfile::TempDir::new().unwrap(); let amp_dir = temp_dir.path().join(".local/share/amp/threads"); std::fs::create_dir_all(&_dir).unwrap(); - - let thread_created = chrono::DateTime::parse_from_rfc3339("2026-04-04T12:00:00Z") - .unwrap() - .timestamp_millis(); - let ledger_timestamp = chrono::DateTime::parse_from_rfc3339("2026-04-08T12:00:00Z") - .unwrap() - .timestamp_millis(); - - let thread = format!( - r#"{{ - "id": "thread-amp-gap", - "created": {thread_created}, - "usageLedger": {{ - "events": [ - {{ - "timestamp": "2026-04-08T12:00:00Z", - "model": "claude-sonnet-4-0", - "credits": 0.75, - "tokens": {{ "input": 100, "output": 20 }} - }} - ] - }}, + std::fs::write( + amp_dir.join("T-legacy.json"), + r#"{ + "id": "legacy-thread", "messages": [ - {{ + { "role": "assistant", "messageId": 1, - "usage": {{ - "model": "claude-sonnet-4-0", - "inputTokens": 100, - "outputTokens": 20, - "credits": 0.75 - }} - }}, - {{ - "role": "assistant", - "messageId": 2, - "usage": {{ - "model": "claude-sonnet-4-0", - "inputTokens": 50, - "outputTokens": 10, - "credits": 0.40 - }} - }} + "usage": { + "timestamp": "2026-05-21T04:00:00Z", + "model": "claude-opus-4-7", + "inputTokens": 10, + "outputTokens": 2 + } + } ] - }}"# - ); - std::fs::write(amp_dir.join("T-thread-amp-gap.json"), thread).unwrap(); + }"#, + ) + .unwrap(); let parsed = parse_local_clients(LocalParseOptions { home_dir: Some(temp_dir.path().to_str().unwrap().to_string()), @@ -5586,19 +5594,7 @@ mod tests { }) .unwrap(); - assert_eq!(parsed.counts.get(ClientId::Amp), 2); - assert_eq!(parsed.messages.len(), 2); - - let dates: HashSet = parsed.messages.iter().map(|msg| msg.date.clone()).collect(); - let local_date = |timestamp_ms: i64| { - chrono::Local - .timestamp_millis_opt(timestamp_ms) - .single() - .unwrap() - .format("%Y-%m-%d") - .to_string() - }; - assert!(dates.contains(&local_date(thread_created + 2000))); - assert!(dates.contains(&local_date(ledger_timestamp))); + assert_eq!(parsed.counts.get(ClientId::Amp), 0); + assert!(parsed.messages.is_empty()); } } diff --git a/crates/tokscale-core/src/scanner.rs b/crates/tokscale-core/src/scanner.rs index be98a2419..53f1e6129 100644 --- a/crates/tokscale-core/src/scanner.rs +++ b/crates/tokscale-core/src/scanner.rs @@ -458,7 +458,12 @@ fn supports_extra_dir_scanning(client_id: ClientId) -> bool { // registry rather than scanned file paths. !matches!( client_id, - ClientId::Kilo | ClientId::Crush | ClientId::Hermes | ClientId::Goose | ClientId::Zed + ClientId::Kilo + | ClientId::Crush + | ClientId::Hermes + | ClientId::Goose + | ClientId::Zed + | ClientId::Amp ) } @@ -598,6 +603,7 @@ fn scan_all_clients_with_env_strategy_inner( | ClientId::Goose | ClientId::Zed | ClientId::Crush + | ClientId::Amp | ClientId::Codebuff ) { continue; @@ -2476,9 +2482,14 @@ mod tests { #[test] fn test_parse_extra_dirs_skips_unsupported_clients() { - let enabled: HashSet = - [ClientId::Claude, ClientId::Kilo].iter().copied().collect(); - let dirs = parse_extra_dirs("claude:/tmp/mac-sessions,kilo:/tmp/kilo", &enabled); + let enabled: HashSet = [ClientId::Claude, ClientId::Kilo, ClientId::Amp] + .iter() + .copied() + .collect(); + let dirs = parse_extra_dirs( + "claude:/tmp/mac-sessions,kilo:/tmp/kilo,amp:/tmp/amp", + &enabled, + ); assert_eq!(dirs.len(), 1); assert_eq!(dirs[0].0, ClientId::Claude); assert_eq!(dirs[0].1, "/tmp/mac-sessions"); @@ -2529,6 +2540,19 @@ mod tests { restore_env("TOKSCALE_EXTRA_DIRS", previous); } + #[test] + fn test_scan_all_clients_does_not_scan_legacy_amp_threads() { + let dir = TempDir::new().unwrap(); + let home = dir.path(); + let amp_threads = home.join(".local/share/amp/threads"); + fs::create_dir_all(&_threads).unwrap(); + File::create(amp_threads.join("T-legacy.json")).unwrap(); + + let result = scan_all_clients(home.to_str().unwrap(), &["amp".to_string()]); + + assert!(result.get(ClientId::Amp).is_empty()); + } + fn setup_mock_codebuff_chat(base: &Path, channel: &str, chat_id: &str) -> PathBuf { let chat_dir = base .join(".config") diff --git a/crates/tokscale-core/src/sessions/amp.rs b/crates/tokscale-core/src/sessions/amp.rs index 49afbb3c7..d07a8120e 100644 --- a/crates/tokscale-core/src/sessions/amp.rs +++ b/crates/tokscale-core/src/sessions/amp.rs @@ -1,42 +1,50 @@ -//! Amp (Sourcegraph) session parser +//! Amp (Sourcegraph) session parser. //! -//! Parses JSON files from ~/.local/share/amp/threads/ +//! Amp no longer keeps the token-bearing thread JSON under a stable local +//! `~/.local/share/amp/threads` directory. Tokscale reads the current Amp +//! server-backed source through `amp threads list` and `amp threads export`. -use super::utils::read_file_or_none; use super::UnifiedMessage; use crate::{provider_identity, TokenBreakdown}; use serde::Deserialize; -use std::path::Path; +use std::fmt; +use std::process::Command; -/// Amp usage event from usageLedger -#[derive(Debug, Deserialize)] -pub struct AmpUsageEvent { - pub timestamp: Option, - pub model: Option, - pub credits: Option, - pub tokens: Option, - #[serde(rename = "operationType")] - pub _operation_type: Option, - #[serde(rename = "fromMessageId")] - pub _from_message_id: Option, - #[serde(rename = "toMessageId")] - pub to_message_id: Option, +const AMP_COMMAND: &str = "amp"; +const AMP_SOURCE_LABEL: &str = "amp threads list/export"; + +#[derive(Debug)] +pub struct AmpCliError { + message: String, +} + +impl AmpCliError { + fn new(message: impl Into) -> Self { + Self { + message: message.into(), + } + } +} + +impl fmt::Display for AmpCliError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}", self.message) + } } +impl std::error::Error for AmpCliError {} + #[derive(Debug, Deserialize)] -pub struct AmpTokens { - pub input: Option, - pub output: Option, - #[serde(rename = "cacheReadInputTokens")] - pub cache_read_input_tokens: Option, - #[serde(rename = "cacheCreationInputTokens")] - pub cache_creation_input_tokens: Option, +struct AmpThreadExport { + id: Option, + messages: Option>, } -/// Amp message usage (per-message, more detailed) +/// Amp exported message usage. #[derive(Debug, Deserialize)] pub struct AmpMessageUsage { pub model: Option, + pub timestamp: Option, #[serde(rename = "inputTokens")] pub input_tokens: Option, #[serde(rename = "outputTokens")] @@ -45,7 +53,6 @@ pub struct AmpMessageUsage { pub cache_read_input_tokens: Option, #[serde(rename = "cacheCreationInputTokens")] pub cache_creation_input_tokens: Option, - pub credits: Option, } #[derive(Debug, Deserialize)] @@ -56,133 +63,95 @@ pub struct AmpMessage { pub usage: Option, } -#[derive(Debug, Deserialize)] -pub struct AmpUsageLedger { - pub events: Option>, +pub fn amp_source_label() -> &'static str { + AMP_SOURCE_LABEL } -#[derive(Debug, Deserialize)] -pub struct AmpThread { - pub id: Option, - pub created: Option, - pub messages: Option>, - #[serde(rename = "usageLedger")] - pub usage_ledger: Option, +pub fn amp_cli_available() -> bool { + run_amp(["--version"]).is_ok() } -/// Get provider from model name +/// Get provider from model name. fn get_provider_from_model(model: &str) -> &'static str { provider_identity::inferred_provider_from_model(model).unwrap_or("anthropic") } -#[derive(Debug, Clone)] -struct AmpUsageRecord { - model: String, - timestamp: i64, - has_explicit_timestamp: bool, - message_id: Option, - ledger_to_message_id: Option, - tokens: TokenBreakdown, - cost: f64, +fn parse_amp_timestamp(timestamp: Option<&str>) -> Option { + timestamp + .and_then(|ts| chrono::DateTime::parse_from_rfc3339(ts).ok()) + .map(|dt| dt.timestamp_millis()) + .filter(|timestamp| *timestamp != 0) } -impl AmpUsageRecord { - fn matches_message_usage(&self, other: &Self) -> bool { - self.model == other.model && self.tokens == other.tokens +fn run_amp(args: [&str; N]) -> Result, AmpCliError> { + let log_file = std::env::temp_dir().join(format!("tokscale-amp-{}.log", std::process::id())); + let output = Command::new(AMP_COMMAND) + .env("TERM", "xterm-256color") + .arg("--no-color") + .arg("--no-notifications") + .arg("--no-ide") + .arg("--log-file") + .arg(log_file) + .args(args) + .output() + .map_err(|err| AmpCliError::new(format!("failed to run amp: {err}")))?; + + if output.status.success() { + return Ok(output.stdout); } - fn into_unified(self, thread_id: &str) -> UnifiedMessage { - UnifiedMessage::new( - "amp", - &self.model, - get_provider_from_model(&self.model), - thread_id.to_string(), - self.timestamp, - self.tokens, - self.cost, - ) - } + let stderr = String::from_utf8_lossy(&output.stderr); + let stdout = String::from_utf8_lossy(&output.stdout); + let detail = stderr.trim(); + let detail = if detail.is_empty() { + stdout.trim() + } else { + detail + }; + Err(AmpCliError::new(format!( + "amp {} failed{}", + args.join(" "), + if detail.is_empty() { + String::new() + } else { + format!(": {detail}") + } + ))) } -fn parse_amp_timestamp(timestamp: Option) -> Option { - timestamp - .and_then(|ts| chrono::DateTime::parse_from_rfc3339(&ts).ok()) - .map(|dt| dt.timestamp_millis()) - .filter(|timestamp| *timestamp != 0) +fn is_amp_thread_id(value: &str) -> bool { + value.strip_prefix("T-").is_some_and(|rest| { + !rest.is_empty() && rest.chars().all(|c| c.is_ascii_alphanumeric() || c == '-') + }) } -fn fallback_amp_timestamp( - explicit: Option, - thread_created_ms: i64, - file_mtime_ms: i64, -) -> i64 { - explicit - .filter(|timestamp| *timestamp != 0) - .or_else(|| (thread_created_ms != 0).then_some(thread_created_ms)) - .unwrap_or(file_mtime_ms) +pub fn parse_amp_thread_ids(list_output: &str) -> Vec { + let mut ids = Vec::new(); + let mut seen = std::collections::HashSet::new(); + for line in list_output.lines() { + let Some(candidate) = line.split_whitespace().last() else { + continue; + }; + if is_amp_thread_id(candidate) && seen.insert(candidate.to_string()) { + ids.push(candidate.to_string()); + } + } + ids } -fn parse_amp_ledger_records( - usage_ledger: Option, - thread_created_ms: i64, - file_mtime_ms: i64, -) -> Vec { - let Some(ledger) = usage_ledger else { - return Vec::new(); - }; - let Some(events) = ledger.events else { - return Vec::new(); - }; - - events - .into_iter() - .filter_map(|event| { - let model = event.model?; - let explicit_timestamp = parse_amp_timestamp(event.timestamp); - let timestamp = - fallback_amp_timestamp(explicit_timestamp, thread_created_ms, file_mtime_ms); - let tokens = event.tokens.unwrap_or(AmpTokens { - input: Some(0), - output: Some(0), - cache_read_input_tokens: Some(0), - cache_creation_input_tokens: Some(0), - }); - - Some(AmpUsageRecord { - model, - timestamp, - has_explicit_timestamp: explicit_timestamp.is_some(), - message_id: None, - ledger_to_message_id: event.to_message_id.filter(|id| *id > 0), - tokens: TokenBreakdown { - input: tokens.input.unwrap_or(0).max(0), - output: tokens.output.unwrap_or(0).max(0), - cache_read: tokens.cache_read_input_tokens.unwrap_or(0).max(0), - cache_write: tokens.cache_creation_input_tokens.unwrap_or(0).max(0), - reasoning: 0, - }, - cost: event.credits.unwrap_or(0.0).max(0.0), - }) - }) - .collect() +pub fn parse_amp_export_bytes(bytes: &mut [u8]) -> Result, AmpCliError> { + let thread: AmpThreadExport = simd_json::from_slice(bytes) + .map_err(|err| AmpCliError::new(format!("failed to parse amp thread export: {err}")))?; + Ok(parse_amp_export(thread)) } -fn parse_amp_message_records( - thread_messages: Option>, - thread_created_ms: i64, - file_mtime_ms: i64, -) -> Vec { - let Some(thread_messages) = thread_messages else { +fn parse_amp_export(thread: AmpThreadExport) -> Vec { + let thread_id = thread.id.unwrap_or_else(|| "unknown".to_string()); + let Some(messages) = thread.messages else { return Vec::new(); }; - let base_timestamp = if thread_created_ms != 0 { - thread_created_ms - } else { - file_mtime_ms - }; - - thread_messages + let mut parsed: Vec = messages .into_iter() .filter_map(|msg| { if msg.role.as_deref() != Some("assistant") { @@ -191,159 +160,68 @@ fn parse_amp_message_records( let usage = msg.usage?; let model = usage.model?; - let message_id = msg.message_id.unwrap_or(0).max(0); - let timestamp = base_timestamp.saturating_add(message_id.saturating_mul(1000)); + let model = model.trim(); + if model.is_empty() { + return None; + } - Some(AmpUsageRecord { + let timestamp = parse_amp_timestamp(usage.timestamp.as_deref())?; + let tokens = TokenBreakdown { + input: usage.input_tokens.unwrap_or(0).max(0), + output: usage.output_tokens.unwrap_or(0).max(0), + cache_read: usage.cache_read_input_tokens.unwrap_or(0).max(0), + cache_write: usage.cache_creation_input_tokens.unwrap_or(0).max(0), + reasoning: 0, + }; + let dedup_key = msg + .message_id + .filter(|id| *id > 0) + .map(|id| format!("amp:{thread_id}:{id}")); + + Some(UnifiedMessage::new_with_dedup( + "amp", model, + get_provider_from_model(model), + thread_id.clone(), timestamp, - has_explicit_timestamp: false, - message_id: Some(message_id).filter(|id| *id > 0), - ledger_to_message_id: None, - tokens: TokenBreakdown { - input: usage.input_tokens.unwrap_or(0).max(0), - output: usage.output_tokens.unwrap_or(0).max(0), - cache_read: usage.cache_read_input_tokens.unwrap_or(0).max(0), - cache_write: usage.cache_creation_input_tokens.unwrap_or(0).max(0), - reasoning: 0, - }, - cost: usage.credits.unwrap_or(0.0).max(0.0), - }) + tokens, + 0.0, + dedup_key, + )) }) - .collect() -} - -fn find_matching_ledger_record( - ledger_records: &[AmpUsageRecord], - consumed: &[bool], - search_start: usize, - message_record: &AmpUsageRecord, -) -> Option { - let find_match = |predicate: &dyn Fn(usize) -> bool| { - (search_start..ledger_records.len()) - .find(|&index| predicate(index)) - .or_else(|| (0..search_start).find(|&index| predicate(index))) - }; + .collect(); - if let Some(message_id) = message_record.message_id { - if let Some(index) = find_match(&|index| { - !consumed[index] && ledger_records[index].ledger_to_message_id == Some(message_id) - }) { - return Some(index); - } - } - - find_match(&|index| { - !consumed[index] && ledger_records[index].matches_message_usage(message_record) - }) + parsed.sort_by_key(|message| message.timestamp); + parsed } -fn merge_amp_records( - ledger_record: AmpUsageRecord, - message_record: &AmpUsageRecord, -) -> AmpUsageRecord { - if ledger_record.has_explicit_timestamp { - if ledger_record.cost > 0.0 || message_record.cost <= 0.0 { - ledger_record - } else { - AmpUsageRecord { - cost: message_record.cost, - message_id: message_record.message_id, - ..ledger_record - } - } - } else { - AmpUsageRecord { - model: ledger_record.model, - timestamp: message_record.timestamp, - has_explicit_timestamp: false, - message_id: message_record.message_id, - ledger_to_message_id: ledger_record.ledger_to_message_id, - tokens: ledger_record.tokens, - cost: if ledger_record.cost > 0.0 { - ledger_record.cost - } else { - message_record.cost - }, - } - } -} - -/// Parse an Amp thread JSON file -pub fn parse_amp_file(path: &Path) -> Vec { - let Some(content) = read_file_or_none(path) else { - return Vec::new(); - }; - - // Get file mtime as last-resort timestamp fallback - let file_mtime_ms = std::fs::metadata(path) - .ok() - .and_then(|m| m.modified().ok()) - .and_then(|t| t.duration_since(std::time::UNIX_EPOCH).ok()) - .map(|d| d.as_millis() as i64) - .unwrap_or(0); - - let mut bytes = content; - let thread: AmpThread = match simd_json::from_slice(&mut bytes) { - Ok(t) => t, - Err(_) => return Vec::new(), - }; - - let thread_id = thread.id.clone().unwrap_or_else(|| { - path.file_stem() - .and_then(|s| s.to_str()) - .unwrap_or("unknown") - .to_string() - }); - - let thread_created_ms = thread.created.unwrap_or(0); - let mut ledger_records = - parse_amp_ledger_records(thread.usage_ledger, thread_created_ms, file_mtime_ms); - let message_records = - parse_amp_message_records(thread.messages, thread_created_ms, file_mtime_ms); - - if ledger_records.is_empty() { - let mut message_records = message_records; - message_records.sort_by_key(|record| record.timestamp); - return message_records - .into_iter() - .map(|record| record.into_unified(&thread_id)) - .collect(); +pub fn parse_amp_threads_from_cli() -> Result, AmpCliError> { + let list_output = run_amp(["threads", "--include-archived", "list"])?; + let list_text = String::from_utf8(list_output).map_err(|err| { + AmpCliError::new(format!("amp threads list returned invalid UTF-8: {err}")) + })?; + let thread_ids = parse_amp_thread_ids(&list_text); + + let mut messages = Vec::new(); + let mut seen_keys = std::collections::HashSet::new(); + for thread_id in thread_ids { + let mut export = run_amp(["threads", "export", thread_id.as_str()])?; + let thread_messages = parse_amp_export_bytes(&mut export)?; + messages.extend(thread_messages.into_iter().filter(|message| { + message + .dedup_key + .as_ref() + .is_none_or(|key| seen_keys.insert(key.clone())) + })); } - let mut consumed = vec![false; ledger_records.len()]; - let mut search_start = 0usize; - let mut unmatched_message_records = Vec::new(); - - for message_record in &message_records { - if let Some(index) = - find_matching_ledger_record(&ledger_records, &consumed, search_start, message_record) - { - consumed[index] = true; - search_start = index.saturating_add(1); - let merged = merge_amp_records(ledger_records[index].clone(), message_record); - ledger_records[index] = merged; - } else { - unmatched_message_records.push(message_record.clone()); - } - } - - ledger_records.extend(unmatched_message_records); - ledger_records.sort_by_key(|record| record.timestamp); - ledger_records - .into_iter() - .map(|record| record.into_unified(&thread_id)) - .collect() + messages.sort_by_key(|message| message.timestamp); + Ok(messages) } #[cfg(test)] mod tests { - use super::parse_amp_file; - use std::path::Path; - - fn write_amp_thread(path: &Path, content: &str) { - std::fs::write(path, content).unwrap(); - } + use super::{parse_amp_export_bytes, parse_amp_thread_ids}; fn timestamp_ms(value: &str) -> i64 { chrono::DateTime::parse_from_rfc3339(value) @@ -351,272 +229,101 @@ mod tests { .timestamp_millis() } - fn local_date(timestamp_ms: i64) -> String { - use chrono::TimeZone; - - chrono::Local - .timestamp_millis_opt(timestamp_ms) - .single() - .unwrap() - .format("%Y-%m-%d") - .to_string() - } - #[test] - fn test_parse_amp_reconciles_partial_ledger_with_message_usage() { - let temp_dir = tempfile::TempDir::new().unwrap(); - let path = temp_dir.path().join("T-partial.json"); - let thread_created = timestamp_ms("2026-04-04T12:00:00Z"); - let ledger_timestamp = "2026-04-08T12:00:00Z"; - - write_amp_thread( - &path, - &serde_json::json!({ - "id": "thread-partial", - "created": thread_created, - "usageLedger": { - "events": [ - { - "timestamp": ledger_timestamp, - "model": "claude-sonnet-4-0", - "credits": 0.75, - "tokens": { "input": 100, "output": 20 } - } - ] - }, - "messages": [ - { - "role": "assistant", - "messageId": 1, - "usage": { - "model": "claude-sonnet-4-0", - "inputTokens": 100, - "outputTokens": 20, - "credits": 0.75 - } - }, - { - "role": "assistant", - "messageId": 2, - "usage": { - "model": "claude-sonnet-4-0", - "inputTokens": 50, - "outputTokens": 10, - "credits": 0.40 - } - } - ] - }) - .to_string(), - ); - - let messages = parse_amp_file(&path); - assert_eq!(messages.len(), 2); - assert_eq!(messages[0].date, local_date(thread_created + 2000)); - assert_eq!(messages[1].date, local_date(timestamp_ms(ledger_timestamp))); - assert_eq!(messages[0].tokens.input, 50); - assert_eq!(messages[1].tokens.input, 100); - } - - #[test] - fn test_parse_amp_does_not_double_count_full_ledger() { - let temp_dir = tempfile::TempDir::new().unwrap(); - let path = temp_dir.path().join("T-full.json"); - let thread_created = timestamp_ms("2026-04-04T12:00:00Z"); - let first_ledger_timestamp = "2026-04-04T12:00:00Z"; - let second_ledger_timestamp = "2026-04-05T12:00:00Z"; - - write_amp_thread( - &path, - &serde_json::json!({ - "id": "thread-full", - "created": thread_created, - "usageLedger": { - "events": [ - { - "timestamp": first_ledger_timestamp, - "model": "claude-sonnet-4-0", - "credits": 0.20, - "tokens": { "input": 20, "output": 5 } - }, - { - "timestamp": second_ledger_timestamp, - "model": "claude-sonnet-4-0", - "credits": 0.25, - "tokens": { "input": 25, "output": 5 } - } - ] - }, - "messages": [ - { - "role": "assistant", - "messageId": 1, - "usage": { - "model": "claude-sonnet-4-0", - "inputTokens": 20, - "outputTokens": 5, - "credits": 0.20 - } - }, - { - "role": "assistant", - "messageId": 2, - "usage": { - "model": "claude-sonnet-4-0", - "inputTokens": 25, - "outputTokens": 5, - "credits": 0.25 - } - } - ] - }) - .to_string(), - ); + fn test_parse_amp_thread_ids_from_table_output() { + let output = r#" +Title Last Updated Visibility Messages Thread ID +──────────────────────────────────────────── ──────────── ────────── ──────── ────────────────────────────────────── +Hi 6m ago Private 13 T-019e48e2-c44e-73fb-8eaf-3c09017ef567 +Reply only with OK 46m ago Private 1 T-019e48eb-861c-75ca-9974-826334df1dea +Reply only with OK 46m ago Private 1 T-019e48eb-861c-75ca-9974-826334df1dea +"#; - let messages = parse_amp_file(&path); - assert_eq!(messages.len(), 2); - assert_eq!( - messages[0].date, - local_date(timestamp_ms(first_ledger_timestamp)) - ); assert_eq!( - messages[1].date, - local_date(timestamp_ms(second_ledger_timestamp)) + parse_amp_thread_ids(output), + vec![ + "T-019e48e2-c44e-73fb-8eaf-3c09017ef567", + "T-019e48eb-861c-75ca-9974-826334df1dea" + ] ); } #[test] - fn test_parse_amp_prefers_message_id_match_over_token_heuristic() { - let temp_dir = tempfile::TempDir::new().unwrap(); - let path = temp_dir.path().join("T-message-id-match.json"); - let thread_created = timestamp_ms("2026-04-04T12:00:00Z"); - let first_ledger_timestamp = "2026-04-10T12:00:00Z"; - let second_ledger_timestamp = "2026-04-05T12:00:00Z"; - - write_amp_thread( - &path, - &serde_json::json!({ - "id": "thread-message-id-match", - "created": thread_created, - "usageLedger": { - "events": [ - { - "timestamp": first_ledger_timestamp, - "model": "claude-sonnet-4-0", - "credits": 0.20, - "tokens": { "input": 20, "output": 5 }, - "toMessageId": 2 - }, - { - "timestamp": second_ledger_timestamp, - "model": "claude-sonnet-4-0", - "credits": 0.20, - "tokens": { "input": 20, "output": 5 }, - "toMessageId": 1 - } - ] + fn test_parse_amp_export_uses_message_usage_timestamp_and_tokens() { + let mut export = serde_json::json!({ + "id": "T-export", + "messages": [ + { + "role": "user", + "messageId": 1, + "content": "hi" }, - "messages": [ - { - "role": "assistant", - "messageId": 1, - "usage": { - "model": "claude-sonnet-4-0", - "inputTokens": 20, - "outputTokens": 5, - "credits": 0.20 - } - }, - { - "role": "assistant", - "messageId": 2, - "usage": { - "model": "claude-sonnet-4-0", - "inputTokens": 20, - "outputTokens": 5, - "credits": 0.20 - } + { + "role": "assistant", + "messageId": 2, + "usage": { + "timestamp": "2026-05-21T04:00:00Z", + "model": "claude-opus-4-7", + "inputTokens": 18232, + "outputTokens": 29, + "totalInputTokens": 18232, + "cacheReadInputTokens": null, + "cacheCreationInputTokens": null } - ] - }) - .to_string(), - ); - - let messages = parse_amp_file(&path); - assert_eq!(messages.len(), 2); - assert_eq!(messages[0].timestamp, timestamp_ms(second_ledger_timestamp)); - assert_eq!(messages[1].timestamp, timestamp_ms(first_ledger_timestamp)); - } - - #[test] - fn test_parse_amp_prefers_message_timestamp_when_ledger_timestamp_missing() { - let temp_dir = tempfile::TempDir::new().unwrap(); - let path = temp_dir.path().join("T-missing-ledger-ts.json"); - let thread_created = timestamp_ms("2026-04-04T12:00:00Z"); - - write_amp_thread( - &path, - &serde_json::json!({ - "id": "thread-missing-ts", - "created": thread_created, - "usageLedger": { - "events": [ - { - "model": "claude-sonnet-4-0", - "credits": 0.20, - "tokens": { "input": 20, "output": 5 } - } - ] }, - "messages": [ - { - "role": "assistant", - "messageId": 7, - "usage": { - "model": "claude-sonnet-4-0", - "inputTokens": 20, - "outputTokens": 5, - "credits": 0.20 - } + { + "role": "assistant", + "messageId": 3, + "usage": { + "timestamp": "2026-05-21T04:01:00Z", + "model": "claude-opus-4-7", + "inputTokens": 100, + "outputTokens": 20, + "cacheReadInputTokens": 5, + "cacheCreationInputTokens": 7 } - ] - }) - .to_string(), - ); + } + ] + }) + .to_string() + .into_bytes(); - let messages = parse_amp_file(&path); - assert_eq!(messages.len(), 1); - assert_eq!(messages[0].timestamp, thread_created + 7000); + let messages = parse_amp_export_bytes(&mut export).unwrap(); + assert_eq!(messages.len(), 2); + assert_eq!(messages[0].client, "amp"); + assert_eq!(messages[0].model_id, "claude-opus-4-7"); + assert_eq!(messages[0].provider_id, "anthropic"); + assert_eq!(messages[0].session_id, "T-export"); + assert_eq!(messages[0].timestamp, timestamp_ms("2026-05-21T04:00:00Z")); + assert_eq!(messages[0].tokens.input, 18_232); + assert_eq!(messages[0].tokens.output, 29); + assert_eq!(messages[0].tokens.cache_read, 0); + assert_eq!(messages[0].tokens.cache_write, 0); + assert_eq!(messages[0].cost, 0.0); + assert_eq!(messages[0].dedup_key.as_deref(), Some("amp:T-export:2")); + assert_eq!(messages[1].tokens.cache_read, 5); + assert_eq!(messages[1].tokens.cache_write, 7); } #[test] - fn test_parse_amp_uses_file_mtime_when_thread_created_missing() { - let temp_dir = tempfile::TempDir::new().unwrap(); - let path = temp_dir.path().join("T-no-created.json"); - - write_amp_thread( - &path, - r#"{ - "id": "thread-no-created", - "messages": [ - { - "role": "assistant", - "messageId": 5, - "usage": { - "model": "claude-sonnet-4-0", - "inputTokens": 10, - "outputTokens": 2, - "credits": 0.11 - } + fn test_parse_amp_export_requires_usage_timestamp() { + let mut export = br#"{ + "id": "T-no-timestamp", + "messages": [ + { + "role": "assistant", + "messageId": 1, + "usage": { + "model": "claude-opus-4-7", + "inputTokens": 10, + "outputTokens": 2 } - ] - }"#, - ); + } + ] + }"# + .to_vec(); - let file_mtime_ms = crate::sessions::utils::file_modified_timestamp_ms(&path); - let messages = parse_amp_file(&path); - assert_eq!(messages.len(), 1); - assert!(messages[0].timestamp >= file_mtime_ms); - assert_ne!(messages[0].date, "1970-01-01"); + let messages = parse_amp_export_bytes(&mut export).unwrap(); + assert!(messages.is_empty()); } } From 7214aab513762280f92df093474a9db2f98b42bd Mon Sep 17 00:00:00 2001 From: makoMakoGo <48956204+makoMakoGo@users.noreply.github.com> Date: Thu, 21 May 2026 17:41:40 +0800 Subject: [PATCH 2/2] fix(amp): parallelize thread exports --- crates/tokscale-core/src/clients.rs | 2 + crates/tokscale-core/src/sessions/amp.rs | 72 ++++++++++++++++++++---- 2 files changed, 63 insertions(+), 11 deletions(-) diff --git a/crates/tokscale-core/src/clients.rs b/crates/tokscale-core/src/clients.rs index 286209896..aca6df363 100644 --- a/crates/tokscale-core/src/clients.rs +++ b/crates/tokscale-core/src/clients.rs @@ -215,6 +215,8 @@ define_clients!( }, Amp = 5 => { id: "amp", + // Amp usage is collected through `amp threads list/export`; this path + // is only the local CLI marker used to decide whether Amp is present. root: PathRoot::Home, relative: ".local/bin/amp", pattern: "", diff --git a/crates/tokscale-core/src/sessions/amp.rs b/crates/tokscale-core/src/sessions/amp.rs index d07a8120e..2366c7ee3 100644 --- a/crates/tokscale-core/src/sessions/amp.rs +++ b/crates/tokscale-core/src/sessions/amp.rs @@ -6,12 +6,17 @@ use super::UnifiedMessage; use crate::{provider_identity, TokenBreakdown}; +use rayon::prelude::*; use serde::Deserialize; use std::fmt; +use std::io::ErrorKind; +use std::path::{Path, PathBuf}; use std::process::Command; +use std::sync::atomic::{AtomicU64, Ordering}; const AMP_COMMAND: &str = "amp"; const AMP_SOURCE_LABEL: &str = "amp threads list/export"; +static AMP_LOG_SEQUENCE: AtomicU64 = AtomicU64::new(0); #[derive(Debug)] pub struct AmpCliError { @@ -68,7 +73,7 @@ pub fn amp_source_label() -> &'static str { } pub fn amp_cli_available() -> bool { - run_amp(["--version"]).is_ok() + run_amp(&["--version"]).is_ok() } /// Get provider from model name. @@ -83,20 +88,53 @@ fn parse_amp_timestamp(timestamp: Option<&str>) -> Option { .filter(|timestamp| *timestamp != 0) } -fn run_amp(args: [&str; N]) -> Result, AmpCliError> { - let log_file = std::env::temp_dir().join(format!("tokscale-amp-{}.log", std::process::id())); +fn amp_log_file_path(args: &[&str]) -> PathBuf { + let sequence = AMP_LOG_SEQUENCE.fetch_add(1, Ordering::Relaxed); + let command = args + .iter() + .map(|arg| { + arg.chars() + .filter(|c| c.is_ascii_alphanumeric()) + .collect::() + }) + .filter(|arg| !arg.is_empty()) + .collect::>() + .join("-"); + std::env::temp_dir().join(format!( + "tokscale-amp-{}-{}-{}.log", + std::process::id(), + sequence, + command + )) +} + +fn remove_amp_log_file(path: &Path) -> Result<(), AmpCliError> { + match std::fs::remove_file(path) { + Ok(()) => Ok(()), + Err(err) if err.kind() == ErrorKind::NotFound => Ok(()), + Err(err) => Err(AmpCliError::new(format!( + "failed to remove amp log file {}: {err}", + path.display() + ))), + } +} + +fn run_amp(args: &[&str]) -> Result, AmpCliError> { + let log_file = amp_log_file_path(args); let output = Command::new(AMP_COMMAND) .env("TERM", "xterm-256color") .arg("--no-color") .arg("--no-notifications") .arg("--no-ide") .arg("--log-file") - .arg(log_file) + .arg(&log_file) .args(args) .output() .map_err(|err| AmpCliError::new(format!("failed to run amp: {err}")))?; + let cleanup_result = remove_amp_log_file(&log_file); if output.status.success() { + cleanup_result?; return Ok(output.stdout); } @@ -108,7 +146,7 @@ fn run_amp(args: [&str; N]) -> Result, AmpCliError> { } else { detail }; - Err(AmpCliError::new(format!( + let command_error = format!( "amp {} failed{}", args.join(" "), if detail.is_empty() { @@ -116,7 +154,13 @@ fn run_amp(args: [&str; N]) -> Result, AmpCliError> { } else { format!(": {detail}") } - ))) + ); + if let Err(cleanup_err) = cleanup_result { + return Err(AmpCliError::new(format!( + "{command_error}; additionally, {cleanup_err}" + ))); + } + Err(AmpCliError::new(command_error)) } fn is_amp_thread_id(value: &str) -> bool { @@ -196,18 +240,24 @@ fn parse_amp_export(thread: AmpThreadExport) -> Vec { } pub fn parse_amp_threads_from_cli() -> Result, AmpCliError> { - let list_output = run_amp(["threads", "--include-archived", "list"])?; + let list_output = run_amp(&["threads", "--include-archived", "list"])?; let list_text = String::from_utf8(list_output).map_err(|err| { AmpCliError::new(format!("amp threads list returned invalid UTF-8: {err}")) })?; let thread_ids = parse_amp_thread_ids(&list_text); + let thread_messages: Vec> = thread_ids + .par_iter() + .map(|thread_id| { + let mut export = run_amp(&["threads", "export", thread_id.as_str()])?; + parse_amp_export_bytes(&mut export) + }) + .collect::, AmpCliError>>()?; + let mut messages = Vec::new(); let mut seen_keys = std::collections::HashSet::new(); - for thread_id in thread_ids { - let mut export = run_amp(["threads", "export", thread_id.as_str()])?; - let thread_messages = parse_amp_export_bytes(&mut export)?; - messages.extend(thread_messages.into_iter().filter(|message| { + for thread in thread_messages { + messages.extend(thread.into_iter().filter(|message| { message .dedup_key .as_ref()