diff --git a/src-tauri/src/usage.rs b/src-tauri/src/usage.rs index 05cd334..385865d 100644 --- a/src-tauri/src/usage.rs +++ b/src-tauri/src/usage.rs @@ -1,4 +1,5 @@ mod resp; +mod token; #[cfg(target_os = "macos")] use super::executable_dir; @@ -27,7 +28,12 @@ use std::{ use tauri::{Emitter, Manager}; use tokio_util::sync::CancellationToken; -use self::resp::UsageSubscription; +use self::{ + resp::{pull_usage_queue, UsageSubscription}, + token::{ + normalize as normalize_keeper_tokens, ClampedTokenFields, TokenValues as KeeperTokenValues, + }, +}; const USAGE_DIR_NAME: &str = "usage-records"; const USAGE_DATABASE_FILE: &str = "usage.db"; @@ -42,10 +48,13 @@ const USAGE_UPDATED_EVENT: &str = "usage-records-updated"; const USAGE_SCHEMA_VERSION: u8 = 1; const USAGE_DATABASE_SCHEMA_VERSION: i64 = 5; const MAX_USAGE_FAILURE_BODY_CHARS: usize = 2_000; -const USAGE_QUEUE_BATCH_SIZE: usize = 500; +const USAGE_QUEUE_BATCH_SIZE: usize = 10_000; const USAGE_INBOX_PROCESS_LIMIT: usize = 500; const USAGE_INBOX_MAX_ATTEMPTS: i64 = 5; const USAGE_SUBSCRIBE_RETRY_SECONDS: u64 = 30; +const USAGE_QUEUE_KEY: &str = "usage"; +const LEGACY_USAGE_QUEUE_KEY: &str = "queue"; +const HTTP_USAGE_QUEUE_SOURCE: &str = "http_pull:usage_queue"; const SQLITE_BUSY_TIMEOUT_SECONDS: u64 = 5; const TOKENS_PER_PRICE_UNIT: f64 = 1_000_000.0; const LONG_CONTEXT_INPUT_TOKEN_THRESHOLD: u64 = 272_000; @@ -1390,6 +1399,65 @@ fn validate_legacy_schema(version: u8, path: &Path) -> Result<(), String> { } } +#[derive(Default)] +struct RedisUsageQueueSource { + selected_key: Option<&'static str>, +} + +struct RawUsageQueueBatch { + source: String, + messages: Vec, +} + +impl RedisUsageQueueSource { + async fn pull(&mut self, config: &GuiConfigFile) -> Result { + if let Some(queue_key) = self.selected_key { + return Self::pull_key(config, queue_key).await; + } + + match Self::pull_key(config, USAGE_QUEUE_KEY).await { + Ok(batch) => { + self.selected_key = Some(USAGE_QUEUE_KEY); + Ok(batch) + } + Err(usage_error) if redis_pull_can_try_legacy_queue(&usage_error) => { + match Self::pull_key(config, LEGACY_USAGE_QUEUE_KEY).await { + Ok(batch) => { + self.selected_key = Some(LEGACY_USAGE_QUEUE_KEY); + Ok(batch) + } + Err(legacy_error) => Err(format!( + "Redis usage 队列拉取失败（usage: {usage_error}; queue: {legacy_error}）" + )), + } + } + Err(error) => Err(error), + } + } + + async fn pull_key( + config: &GuiConfigFile, + queue_key: &'static str, + ) -> Result { + let messages = pull_usage_queue( + config.port, + &config.management_secret_key, + queue_key, + USAGE_QUEUE_BATCH_SIZE, + ) + .await?; + Ok(RawUsageQueueBatch { + source: format!("redis_pull:{queue_key}"), + messages, + }) + } +} + +fn redis_pull_can_try_legacy_queue(error: &str) -> bool { + let error = error.to_ascii_lowercase(); + error.contains("unsupported channel") || error.contains("unsupported queue") +} + pub(crate) fn start_usage_collector(app: tauri::AppHandle) { let state = app.state::(); if let Ok(total_records) = total_usage_records() { @@ -1418,6 +1486,7 @@ async fn usage_collector_loop(app: tauri::AppHandle, token: CancellationToken) { let mut retry_seconds = 1_u64; let mut subscription: Option = None; + let mut redis_queue = RedisUsageQueueSource::default(); let mut subscribe_retry_at = tokio::time::Instant::now(); let mut next_inbox_cleanup_at = tokio::time::Instant::now() + Duration::from_secs(60 * 60); let mut next_inbox_recovery_at = tokio::time::Instant::now(); @@ -1482,6 +1551,8 @@ async fn usage_collector_loop(app: tauri::AppHandle, token: CancellationToken) { } if !core_running { subscription = None; + redis_queue = RedisUsageQueueSource::default(); + subscribe_retry_at = tokio::time::Instant::now(); set_collector_status(&app, "waiting-core", "等待内核启动", None); retry_seconds = 1; wait_or_cancel(&token, 1).await; @@ -1493,7 +1564,7 @@ async fn usage_collector_loop(app: tauri::AppHandle, token: CancellationToken) { Ok(next_subscription) => { subscription = Some(next_subscription); set_collector_status(&app, "collecting", "已连接 CPA usage 实时订阅", None); - match backfill_usage_queue(&root, &config).await { + match backfill_usage_queue(&root, &config, &mut redis_queue).await { Ok(saved) => { publish_collected_records( &app, @@ -1577,27 +1648,43 @@ async fn usage_collector_loop(app: tauri::AppHandle, token: CancellationToken) { } } - match fetch_usage_queue(&config).await { - Ok(items) if items.is_empty() => { - set_collector_status(&app, "collecting", "使用记录采集中", None); + match pull_usage_queue_with_fallback(&mut redis_queue, &config).await { + Ok(batch) if batch.messages.is_empty() => { + set_collector_status( + &app, + "collecting", + &format!( + "使用记录采集中（{}）", + collector_source_label(&batch.source) + ), + None, + ); retry_seconds = 1; wait_or_cancel(&token, 1).await; } - Ok(items) => match persist_queue_items(&root, items, &config) { - Ok(saved) => { - publish_collected_records( - &app, - saved, - &format!("HTTP 兼容模式已保存 {saved} 条新记录"), - ); - retry_seconds = 1; + Ok(batch) => { + let source_label = collector_source_label(&batch.source); + match persist_raw_usage_messages_from_source( + &root, + &batch.source, + batch.messages, + &config, + ) { + Ok(saved) => { + publish_collected_records( + &app, + saved, + &format!("{source_label}已保存 {saved} 条新记录"), + ); + retry_seconds = 1; + } + Err(error) => { + set_collector_error(&app, error); + wait_or_cancel(&token, retry_seconds).await; + retry_seconds = (retry_seconds * 2).min(10); + } } - Err(error) => { - set_collector_error(&app, error); - wait_or_cancel(&token, retry_seconds).await; - retry_seconds = (retry_seconds * 2).min(10); - } - }, + } Err(error) => { set_collector_error(&app, error); wait_or_cancel(&token, retry_seconds).await; @@ -1607,18 +1694,55 @@ async fn usage_collector_loop(app: tauri::AppHandle, token: CancellationToken) { } } -async fn backfill_usage_queue(root: &Path, config: &GuiConfigFile) -> Result { +async fn backfill_usage_queue( + root: &Path, + config: &GuiConfigFile, + redis_queue: &mut RedisUsageQueueSource, +) -> Result { let mut saved_total = 0_usize; loop { - let items = fetch_usage_queue(config).await?; - let fetched = items.len(); + match redis_queue.pull(config).await { + Ok(batch) => { + let fetched = batch.messages.len(); + if fetched == 0 { + return Ok(saved_total); + } + saved_total = saved_total.saturating_add(persist_raw_usage_messages_from_source( + root, + &batch.source, + batch.messages, + config, + )?); + if fetched < USAGE_QUEUE_BATCH_SIZE { + return Ok(saved_total); + } + } + Err(redis_error) => { + return backfill_http_usage_queue(root, config, saved_total) + .await + .map_err(|http_error| { + format!("补录队列失败（Redis: {redis_error}; HTTP: {http_error}）") + }); + } + } + } +} + +async fn backfill_http_usage_queue( + root: &Path, + config: &GuiConfigFile, + mut saved_total: usize, +) -> Result { + loop { + let messages = fetch_usage_queue_raw(config).await?; + let fetched = messages.len(); if fetched == 0 { return Ok(saved_total); } - saved_total = saved_total.saturating_add(persist_queue_items_from_source( + saved_total = saved_total.saturating_add(persist_raw_usage_messages_from_source( root, - "http_backfill", - items, + HTTP_USAGE_QUEUE_SOURCE, + messages, config, )?); if fetched < USAGE_QUEUE_BATCH_SIZE { @@ -1696,6 +1820,43 @@ async fn fetch_usage_queue(config: &GuiConfigFile) -> Result, String> .map_err(|error| format!("解析 CPA 使用记录失败: {error}")) } +async fn fetch_usage_queue_raw(config: &GuiConfigFile) -> Result, String> { + fetch_usage_queue(config) + .await? + .into_iter() + .map(|item| { + serde_json::to_string(&item) + .map_err(|error| format!("序列化 CPA 使用记录失败: {error}")) + }) + .collect() +} + +async fn pull_usage_queue_with_fallback( + redis_queue: &mut RedisUsageQueueSource, + config: &GuiConfigFile, +) -> Result { + match redis_queue.pull(config).await { + Ok(batch) => Ok(batch), + Err(redis_error) => fetch_usage_queue_raw(config) + .await + .map(|messages| RawUsageQueueBatch { + source: HTTP_USAGE_QUEUE_SOURCE.to_string(), + messages, + }) + .map_err(|http_error| { + format!("读取 CPA 使用记录失败（Redis: {redis_error}; HTTP: {http_error}）") + }), + } +} + +fn collector_source_label(source: &str) -> &'static str { + if source.starts_with("redis_pull:") { + "Redis 队列" + } else { + "HTTP 兼容模式" + } +} + fn set_collector_error(app: &tauri::AppHandle, error: String) { set_collector_status(app, "error", &error, None); } @@ -1726,7 +1887,7 @@ fn persist_queue_items( items: Vec, config: &GuiConfigFile, ) -> Result { - persist_queue_items_from_source(root, "http_pull", items, config) + persist_queue_items_from_source(root, HTTP_USAGE_QUEUE_SOURCE, items, config) } fn persist_queue_items_from_source( @@ -1834,9 +1995,18 @@ fn persist_raw_usage_message_from_source( source: &str, raw_message: String, config: &GuiConfigFile, +) -> Result { + persist_raw_usage_messages_from_source(root, source, vec![raw_message], config) +} + +fn persist_raw_usage_messages_from_source( + root: &Path, + source: &str, + raw_messages: Vec, + config: &GuiConfigFile, ) -> Result { let mut connection = open_usage_database_at(root)?; - enqueue_usage_raw_messages(&mut connection, source, vec![raw_message])?; + enqueue_usage_raw_messages(&mut connection, source, raw_messages)?; process_usage_inbox(&mut connection, config) } @@ -2077,25 +2247,6 @@ fn insert_usage_records_in_transaction( } else { "unknown" }; - let cache_components = record - .tokens - .cache_read_tokens - .saturating_add(record.tokens.cache_creation_tokens); - let input_before_invariant = record.tokens.input_tokens; - let input_tokens = if cache_components > input_before_invariant { - input_before_invariant.saturating_add(cache_components) - } else { - input_before_invariant - }; - let total_tokens = if record.tokens.total_tokens == 0 - || record.tokens.total_tokens - == input_before_invariant.saturating_add(record.tokens.output_tokens) - { - input_tokens.saturating_add(record.tokens.output_tokens) - } else { - record.tokens.total_tokens - }; - let cached_tokens = record.cached_tokens.max(cache_components); let collector_source = if record.collector_source.trim().is_empty() { "legacy_json" } else { @@ -2136,14 +2287,14 @@ fn insert_usage_records_in_transaction( record.x_forwarded_for, record.user_agent, record.generate, - to_sql_i64(cached_tokens), + to_sql_i64(record.cached_tokens), collector_source, - to_sql_i64(input_tokens), + to_sql_i64(record.tokens.input_tokens), to_sql_i64(record.tokens.output_tokens), to_sql_i64(record.tokens.reasoning_tokens), to_sql_i64(record.tokens.cache_read_tokens), to_sql_i64(record.tokens.cache_creation_tokens), - to_sql_i64(total_tokens), + to_sql_i64(record.tokens.total_tokens), record.canceled, i64::from(record.failure_status), record.failure_body, @@ -2174,6 +2325,7 @@ fn normalize_usage_record(value: Value, config: &GuiConfigFile) -> Result Result 0 - && (raw_input_tokens < cache_components || tokens.total_tokens == raw_total_with_cache); - if claude_excludes_cache { - tokens.input_tokens = raw_input_tokens - .saturating_add(tokens.cache_read_tokens) - .saturating_add(tokens.cache_creation_tokens); - } - let input_before_invariant = tokens.input_tokens; - if cache_components > input_before_invariant { - tokens.input_tokens = input_before_invariant.saturating_add(cache_components); - if tokens.total_tokens == 0 - || tokens.total_tokens == input_before_invariant.saturating_add(tokens.output_tokens) - { - tokens.total_tokens = tokens.input_tokens.saturating_add(tokens.output_tokens); - } - } - if tokens.total_tokens == 0 - || (claude_excludes_cache && tokens.total_tokens == raw_total_without_cache) - { - tokens.total_tokens = tokens.input_tokens.saturating_add(tokens.output_tokens); - } let id = request_id.clone(); let endpoint = string_field(object, "endpoint").unwrap_or_default(); let failed = object @@ -2284,18 +2423,14 @@ fn normalize_usage_record(value: Value, config: &GuiConfigFile) -> Result 0 - AND ttft_ms IS NOT NULL - AND ttft_ms > 0 - AND latency_ms > ttft_ms + AND latency_ms > 0 THEN output_tokens ELSE 0 END) * 1000.0 @@ -2450,10 +2583,8 @@ fn load_usage_overview( AND failed = 0 AND canceled = 0 AND output_tokens > 0 - AND ttft_ms IS NOT NULL - AND ttft_ms > 0 - AND latency_ms > ttft_ms - THEN latency_ms - ttft_ms + AND latency_ms > 0 + THEN latency_ms ELSE 0 END), 0), 0.0 @@ -2463,9 +2594,7 @@ fn load_usage_overview( AND failed = 0 AND canceled = 0 AND output_tokens > 0 - AND ttft_ms IS NOT NULL - AND ttft_ms > 0 - AND latency_ms > ttft_ms + AND latency_ms > 0 THEN 1 ELSE 0 END), 0), @@ -3617,6 +3746,17 @@ fn token_u64(object: Option<&serde_json::Map>, key: &str) -> u64 object.map(|object| u64_field(object, key)).unwrap_or(0) } +fn token_was_negative(object: Option<&serde_json::Map>, key: &str) -> bool { + object + .and_then(|object| object.get(key)) + .and_then(|value| { + value + .as_i64() + .or_else(|| value.as_str().and_then(|text| text.parse::().ok())) + }) + .is_some_and(|value| value < 0) +} + fn hash_text(value: &str) -> String { if value.is_empty() { return String::new(); @@ -3680,6 +3820,76 @@ fn usage_source_display(config: &GuiConfigFile, provider: &str, source: &str) -> mod tests { use super::*; + #[tokio::test] + async fn redis_usage_queue_falls_back_to_legacy_key_and_reuses_selection() { + use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::TcpListener, + }; + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + let legacy_message = r#"{"request_id":"legacy"}"#; + let expectations = vec![ + ( + USAGE_QUEUE_KEY.to_string(), + "-ERR unsupported channel 'usage'\r\n".to_string(), + ), + (LEGACY_USAGE_QUEUE_KEY.to_string(), "*0\r\n".to_string()), + ( + LEGACY_USAGE_QUEUE_KEY.to_string(), + format!("*1\r\n${}\r\n{legacy_message}\r\n", legacy_message.len()), + ), + ]; + let server = tokio::spawn(async move { + for (queue_key, response) in expectations { + let (mut stream, _) = listener.accept().await.unwrap(); + let expected_auth = b"*2\r\n$4\r\nAUTH\r\n$6\r\nsecret\r\n"; + let mut auth = vec![0_u8; expected_auth.len()]; + stream.read_exact(&mut auth).await.unwrap(); + assert_eq!(auth, expected_auth); + stream.write_all(b"+OK\r\n").await.unwrap(); + + let batch_size = USAGE_QUEUE_BATCH_SIZE.to_string(); + let expected_pop = format!( + "*3\r\n$4\r\nLPOP\r\n${}\r\n{queue_key}\r\n${}\r\n{batch_size}\r\n", + queue_key.len(), + batch_size.len() + ); + let mut pop = vec![0_u8; expected_pop.len()]; + stream.read_exact(&mut pop).await.unwrap(); + assert_eq!(pop, expected_pop.as_bytes()); + stream.write_all(response.as_bytes()).await.unwrap(); + } + }); + + let mut config = GuiConfigFile::default(); + config.port = port; + config.management_secret_key = "secret".to_string(); + let mut source = RedisUsageQueueSource::default(); + + let first = source.pull(&config).await.unwrap(); + assert_eq!(first.source, "redis_pull:queue"); + assert!(first.messages.is_empty()); + + let second = source.pull(&config).await.unwrap(); + assert_eq!(second.source, "redis_pull:queue"); + assert_eq!(second.messages, vec![legacy_message]); + server.await.unwrap(); + } + + #[test] + fn redis_usage_queue_only_falls_back_for_unsupported_key_errors() { + assert!(redis_pull_can_try_legacy_queue( + "ERR unsupported channel 'usage'" + )); + assert!(redis_pull_can_try_legacy_queue( + "ERR unsupported queue 'usage'" + )); + assert!(!redis_pull_can_try_legacy_queue("CPA usage 订阅认证失败")); + assert!(!redis_pull_can_try_legacy_queue("connection refused")); + } + #[tokio::test] async fn background_usage_jobs_limit_database_concurrency() { use std::sync::{ @@ -4206,10 +4416,11 @@ mod tests { } #[test] - fn keeps_already_inclusive_claude_input_unchanged() { + fn legacy_claude_identity_uses_keeper_missing_cache_contract() { let record = normalize_usage_record( serde_json::json!({ "provider": "anthropic", + "auth_type": "oauth", "request_id": "claude-inclusive", "tokens": { "input_tokens": 720, @@ -4222,12 +4433,246 @@ mod tests { &GuiConfigFile::default(), ) .unwrap(); - assert_eq!(record.tokens.input_tokens, 720); - assert_eq!(record.tokens.total_tokens, 740); + assert_eq!(record.tokens.input_tokens, 1_340); + assert_eq!(record.tokens.total_tokens, 1_360); } #[test] - fn enforces_cache_input_invariant_for_unknown_producers() { + fn unresolved_executors_only_use_provider_for_oauth() { + for executor in ["", "unknown", "FutureExecutor"] { + for auth_type in ["", "apikey", "api_key", "unknown", " OAUTH "] { + let record = normalize_usage_record( + serde_json::json!({ + "request_id": "unresolved-executor", + "executor_type": executor, + "provider": "anthropic", + "auth_type": auth_type, + "auth_index": "missing-identity", + "tokens": { + "input_tokens": 100, + "output_tokens": 20, + "cache_read_tokens": 30, + "total_tokens": 120 + } + }), + &GuiConfigFile::default(), + ) + .unwrap(); + let expected_input = if auth_type.trim().eq_ignore_ascii_case("oauth") { + 130 + } else { + 100 + }; + assert_eq!( + record.tokens.input_tokens, expected_input, + "{executor}/{auth_type}" + ); + assert_eq!( + record.tokens.total_tokens, + expected_input + 20, + "{executor}/{auth_type}" + ); + } + } + } + + #[test] + fn codex_cached_only_records_preserve_total_with_missing_or_zero_read() { + for explicit_read in [false, true] { + let mut tokens = serde_json::json!({ + "input_tokens": 0, + "output_tokens": 0, + "reasoning_tokens": 0, + "cached_tokens": 30, + "cache_creation_tokens": 0, + "total_tokens": 30 + }); + if explicit_read { + tokens["cache_read_tokens"] = serde_json::json!(0); + } + let record = normalize_usage_record( + serde_json::json!({ + "request_id": "codex-cached-only", + "executor_type": "CodexExecutor", + "tokens": tokens + }), + &GuiConfigFile::default(), + ) + .unwrap(); + assert_eq!( + record.tokens.total_tokens, 30, + "explicit_read={explicit_read}" + ); + assert_eq!(record.tokens.input_tokens, 0); + assert_eq!(record.tokens.output_tokens, 0); + assert_eq!( + record.tokens.cache_read_tokens, + if explicit_read { 0 } else { 30 } + ); + } + } + + #[test] + fn clamped_contract_fields_do_not_rewrite_nonzero_total() { + for executor in ["ClaudeExecutor", "GeminiExecutor", "CodexExecutor"] { + for field in [ + "input_tokens", + "output_tokens", + "reasoning_tokens", + "cache_read_tokens", + "cache_creation_tokens", + ] { + let mut tokens = serde_json::json!({ + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 99 + }); + tokens[field] = serde_json::json!(-1); + let record = normalize_usage_record( + serde_json::json!({ + "request_id": "clamped-contract", + "executor_type": executor, + "tokens": tokens + }), + &GuiConfigFile::default(), + ) + .unwrap(); + assert_eq!(record.tokens.total_tokens, 99, "{executor}/{field}"); + } + } + } + + #[test] + fn zero_total_reconciliation_respects_keeper_clamped_field_dependencies() { + for executor in [ + "ClaudeExecutor", + "GeminiExecutor", + "CodexExecutor", + "OpenAICompatExecutor", + "KimiExecutor", + "unknown", + ] { + for field in [ + "input_tokens", + "output_tokens", + "total_tokens", + "reasoning_tokens", + "cache_read_tokens", + "cache_creation_tokens", + "cached_tokens", + ] { + let mut tokens = serde_json::json!({ + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 0 + }); + tokens[field] = serde_json::json!(-1); + let record = normalize_usage_record( + serde_json::json!({ + "request_id": "clamped-zero-fallback", + "executor_type": executor, + "tokens": tokens + }), + &GuiConfigFile::default(), + ) + .unwrap(); + let expected = if matches!(field, "input_tokens" | "output_tokens" | "total_tokens") + || (executor == "ClaudeExecutor" && field == "reasoning_tokens") + { + 0 + } else { + 15 + }; + assert_eq!(record.tokens.total_tokens, expected, "{executor}/{field}"); + } + } + } + + #[test] + fn clamped_fields_cannot_trigger_dependent_token_folds() { + let cases = [ + ( + "ClaudeExecutor", + "anthropic", + serde_json::json!({ + "input_tokens": 10, "output_tokens": 5, "cache_read_tokens": -1, + "cache_creation_tokens": 30, "total_tokens": 99 + }), + 10, + 5, + ), + ( + "GeminiExecutor", + "gemini", + serde_json::json!({ + "input_tokens": 10, "output_tokens": -1, "reasoning_tokens": 5, + "total_tokens": 99 + }), + 10, + 0, + ), + ( + "", + "gemini", + serde_json::json!({ + "input_tokens": -1, "output_tokens": 5, "reasoning_tokens": 10, + "total_tokens": 15 + }), + 0, + 5, + ), + ( + "OpenAICompatExecutor", + "openai", + serde_json::json!({ + "input_tokens": -1, "output_tokens": 5, "reasoning_tokens": 10, + "total_tokens": 15 + }), + 0, + 5, + ), + ]; + for (executor, provider, tokens, input, output) in cases { + let record = normalize_usage_record( + serde_json::json!({ + "request_id": "clamped-fold", + "executor_type": executor, + "provider": provider, + "auth_type": "oauth", + "tokens": tokens + }), + &GuiConfigFile::default(), + ) + .unwrap(); + assert_eq!(record.tokens.input_tokens, input, "{executor}/{provider}"); + assert_eq!(record.tokens.output_tokens, output, "{executor}/{provider}"); + } + } + + #[test] + fn clamped_cache_read_is_not_backfilled_from_cached_alias() { + for executor in ["CodexExecutor", "KimiExecutor", "unknown"] { + let record = normalize_usage_record( + serde_json::json!({ + "request_id": "clamped-cache-alias", + "executor_type": executor, + "tokens": { + "input_tokens": 100, + "output_tokens": 20, + "cache_read_tokens": -1, + "cached_tokens": 30, + "total_tokens": 120 + } + }), + &GuiConfigFile::default(), + ) + .unwrap(); + assert_eq!(record.tokens.cache_read_tokens, 0, "{executor}"); + } + } + + #[test] + fn unknown_producers_preserve_nonzero_parent_token_fields() { let record = normalize_usage_record( serde_json::json!({ "provider": "custom", @@ -4242,9 +4687,9 @@ mod tests { &GuiConfigFile::default(), ) .unwrap(); - assert_eq!(record.tokens.input_tokens, 700); - assert_eq!(record.tokens.total_tokens, 720); - assert!(record.tokens.cache_read_tokens <= record.tokens.input_tokens); + assert_eq!(record.tokens.input_tokens, 100); + assert_eq!(record.tokens.total_tokens, 120); + assert_eq!(record.tokens.cache_read_tokens, 600); } #[test] @@ -5039,7 +5484,7 @@ mod tests { } #[test] - fn overview_tps_uses_weighted_generation_time_and_ignores_invalid_records() { + fn overview_tps_uses_weighted_total_latency_without_requiring_ttft() { let root = test_root("tps-overview"); let mut connection = open_test_database(&root); @@ -5075,6 +5520,13 @@ mod tests { non_generation.generate = false; non_generation.tokens.output_tokens = 100; + let mut zero_latency = sample_record("tps-8", "2026-07-17T20:37:00+08:00", "gpt-a"); + zero_latency.latency_ms = 0; + zero_latency.tokens.output_tokens = 100; + + let mut no_output = sample_record("tps-9", "2026-07-17T20:38:00+08:00", "gpt-a"); + no_output.tokens.output_tokens = 0; + insert_usage_records( &mut connection, &[ @@ -5085,14 +5537,16 @@ mod tests { failed, canceled, non_generation, + zero_latency, + no_output, ], ) .unwrap(); let overview = load_usage_overview(&connection, &UsageQuery::default()).unwrap(); - assert!((overview.tps - (100.0 * 1_000.0 / 1_800.0)).abs() < f64::EPSILON); - assert_eq!(overview.tps_sample_count, 2); + assert!((overview.tps - (300.0 * 1_000.0 / 4_000.0)).abs() < f64::EPSILON); + assert_eq!(overview.tps_sample_count, 4); drop(connection); fs::remove_dir_all(root).unwrap(); } diff --git a/src-tauri/src/usage/resp.rs b/src-tauri/src/usage/resp.rs index 83b39bc..7a4edbb 100644 --- a/src-tauri/src/usage/resp.rs +++ b/src-tauri/src/usage/resp.rs @@ -9,9 +9,37 @@ use tokio::{ const USAGE_CHANNEL: &str = "usage"; const CONNECT_TIMEOUT: Duration = Duration::from_secs(5); const IO_TIMEOUT: Duration = Duration::from_secs(10); -const MAX_FRAME_BYTES: usize = 16 * 1024 * 1024; -const MAX_ARRAY_LENGTH: usize = 10_000; -const MAX_NESTING_DEPTH: usize = 8; +const MAX_BULK_BYTES: usize = 4 * 1024 * 1024; +const MAX_RESP_LINE_BYTES: usize = 4096; +const MAX_NESTING_DEPTH: usize = 4; +const SUBSCRIPTION_MAX_FRAME_BYTES: usize = MAX_BULK_BYTES + MAX_RESP_LINE_BYTES; +const SUBSCRIPTION_MAX_ARRAY_LENGTH: usize = 16; +const SUBSCRIPTION_MAX_TOTAL_BULK_BYTES: usize = MAX_BULK_BYTES; +const QUEUE_MAX_ARRAY_LENGTH: usize = 10_000; +const QUEUE_MAX_TOTAL_BULK_BYTES: usize = 16 * 1024 * 1024; + +#[derive(Clone, Copy)] +struct RespLimits { + max_frame_bytes: usize, + max_array_length: usize, + max_total_bulk_bytes: usize, +} + +const SUBSCRIPTION_LIMITS: RespLimits = RespLimits { + max_frame_bytes: SUBSCRIPTION_MAX_FRAME_BYTES, + max_array_length: SUBSCRIPTION_MAX_ARRAY_LENGTH, + max_total_bulk_bytes: SUBSCRIPTION_MAX_TOTAL_BULK_BYTES, +}; + +fn queue_limits(max_array_length: usize) -> RespLimits { + let max_array_length = max_array_length.min(QUEUE_MAX_ARRAY_LENGTH); + RespLimits { + max_frame_bytes: QUEUE_MAX_TOTAL_BULK_BYTES + .saturating_add((max_array_length + 1).saturating_mul(MAX_RESP_LINE_BYTES + 3)), + max_array_length, + max_total_bulk_bytes: QUEUE_MAX_TOTAL_BULK_BYTES, + } +} pub(super) struct UsageSubscription { stream: TcpStream, @@ -20,6 +48,18 @@ pub(super) struct UsageSubscription { impl UsageSubscription { pub(super) async fn connect(port: u16, management_key: &str) -> Result { + let mut subscription = Self::connect_authenticated(port, management_key).await?; + subscription + .send_command(&["SUBSCRIBE", USAGE_CHANNEL]) + .await?; + let acknowledgement = subscription.read_frame().await?; + if !is_subscription_ack(&acknowledgement) { + return Err("CPA 未确认 usage 订阅".to_string()); + } + Ok(subscription) + } + + async fn connect_authenticated(port: u16, management_key: &str) -> Result { let address = format!("127.0.0.1:{port}"); let stream = timeout(CONNECT_TIMEOUT, TcpStream::connect(&address)) .await @@ -35,13 +75,6 @@ impl UsageSubscription { RespValue::Error(error) => return Err(format!("CPA usage 订阅认证失败: {error}")), value => return Err(format!("CPA usage 订阅认证响应无效: {}", value.kind())), } - subscription - .send_command(&["SUBSCRIBE", USAGE_CHANNEL]) - .await?; - let acknowledgement = subscription.read_frame().await?; - if !is_subscription_ack(&acknowledgement) { - return Err("CPA 未确认 usage 订阅".to_string()); - } Ok(subscription) } @@ -71,16 +104,30 @@ impl UsageSubscription { } async fn read_frame(&mut self) -> Result { + self.read_frame_with_limits(SUBSCRIPTION_LIMITS).await + } + + async fn read_frame_with_limits(&mut self, limits: RespLimits) -> Result { loop { - match parse_resp_frame(&self.read_buffer, 0, 0)? { + let mut remaining_bulk_bytes = limits.max_total_bulk_bytes; + match parse_resp_frame_with_limits( + &self.read_buffer, + 0, + 0, + limits, + &mut remaining_bulk_bytes, + )? { ParseResult::Complete(value, consumed) => { + if consumed > limits.max_frame_bytes { + return Err("CPA usage RESP 响应超过大小限制".to_string()); + } self.read_buffer.drain(..consumed); return Ok(value); } ParseResult::Incomplete => {} } - if self.read_buffer.len() >= MAX_FRAME_BYTES { - return Err("CPA usage 订阅响应超过大小限制".to_string()); + if self.read_buffer.len() >= limits.max_frame_bytes { + return Err("CPA usage RESP 响应超过大小限制".to_string()); } let mut chunk = [0_u8; 8192]; let read = timeout(IO_TIMEOUT, self.stream.read(&mut chunk)) @@ -95,6 +142,43 @@ impl UsageSubscription { } } +pub(super) async fn pull_usage_queue( + port: u16, + management_key: &str, + queue_key: &str, + count: usize, +) -> Result, String> { + if count == 0 || count > QUEUE_MAX_ARRAY_LENGTH { + return Err(format!( + "CPA usage 队列批量大小必须在 1..={QUEUE_MAX_ARRAY_LENGTH} 之间" + )); + } + let mut connection = UsageSubscription::connect_authenticated(port, management_key).await?; + let count_argument = count.to_string(); + connection + .send_command(&["LPOP", queue_key, count_argument.as_str()]) + .await?; + match connection + .read_frame_with_limits(queue_limits(count)) + .await? + { + RespValue::Array(Some(values)) => values + .into_iter() + .map(|value| { + value + .text() + .ok_or_else(|| "CPA usage 队列包含非文本消息".to_string()) + }) + .collect(), + RespValue::Bulk(Some(value)) => String::from_utf8(value) + .map(|value| vec![value]) + .map_err(|_| "CPA usage 队列消息不是 UTF-8".to_string()), + RespValue::Array(None) | RespValue::Bulk(None) => Ok(Vec::new()), + RespValue::Error(error) => Err(format!("CPA usage 队列 LPOP 失败: {error}")), + value => Err(format!("CPA usage 队列 LPOP 响应无效: {}", value.kind())), + } +} + #[derive(Debug)] enum RespValue { Simple(String), @@ -131,31 +215,48 @@ enum ParseResult { } fn parse_resp_frame(input: &[u8], offset: usize, depth: usize) -> Result { + let mut remaining_bulk_bytes = SUBSCRIPTION_LIMITS.max_total_bulk_bytes; + parse_resp_frame_with_limits( + input, + offset, + depth, + SUBSCRIPTION_LIMITS, + &mut remaining_bulk_bytes, + ) +} + +fn parse_resp_frame_with_limits( + input: &[u8], + offset: usize, + depth: usize, + limits: RespLimits, + remaining_bulk_bytes: &mut usize, +) -> Result { if depth > MAX_NESTING_DEPTH { - return Err("CPA usage 订阅响应嵌套过深".to_string()); + return Err("CPA usage RESP 响应嵌套过深".to_string()); } let Some(prefix) = input.get(offset).copied() else { return Ok(ParseResult::Incomplete); }; match prefix { b'+' | b'-' | b':' => { - let Some((line, next)) = resp_line(input, offset + 1) else { + let Some((line, next)) = resp_line(input, offset + 1)? else { return Ok(ParseResult::Incomplete); }; let text = String::from_utf8(line.to_vec()) - .map_err(|_| "CPA usage 订阅文本响应不是 UTF-8".to_string())?; + .map_err(|_| "CPA usage RESP 文本响应不是 UTF-8".to_string())?; let value = match prefix { b'+' => RespValue::Simple(text), b'-' => RespValue::Error(text), _ => RespValue::Integer( text.parse::() - .map_err(|_| "CPA usage 订阅整数响应无效".to_string())?, + .map_err(|_| "CPA usage RESP 整数响应无效".to_string())?, ), }; Ok(ParseResult::Complete(value, next - offset)) } b'$' => { - let Some((line, data_start)) = resp_line(input, offset + 1) else { + let Some((line, data_start)) = resp_line(input, offset + 1)? else { return Ok(ParseResult::Incomplete); }; let length = parse_resp_length(line)?; @@ -166,9 +267,12 @@ fn parse_resp_frame(input: &[u8], offset: usize, depth: usize) -> Result MAX_FRAME_BYTES { - return Err("CPA usage 订阅字符串超过大小限制".to_string()); + usize::try_from(length).map_err(|_| "CPA usage RESP 字符串长度无效".to_string())?; + if length > MAX_BULK_BYTES { + return Err("CPA usage RESP 字符串超过大小限制".to_string()); + } + if length > *remaining_bulk_bytes { + return Err("CPA usage RESP 响应超过总字符串大小限制".to_string()); } let data_end = data_start.saturating_add(length); let frame_end = data_end.saturating_add(2); @@ -176,15 +280,16 @@ fn parse_resp_frame(input: &[u8], offset: usize, depth: usize) -> Result { - let Some((line, mut next)) = resp_line(input, offset + 1) else { + let Some((line, mut next)) = resp_line(input, offset + 1)? else { return Ok(ParseResult::Incomplete); }; let length = parse_resp_length(line)?; @@ -192,13 +297,19 @@ fn parse_resp_frame(input: &[u8], offset: usize, depth: usize) -> Result MAX_ARRAY_LENGTH { - return Err("CPA usage 订阅数组超过大小限制".to_string()); + usize::try_from(length).map_err(|_| "CPA usage RESP 数组长度无效".to_string())?; + if length > limits.max_array_length { + return Err("CPA usage RESP 数组超过大小限制".to_string()); } let mut values = Vec::with_capacity(length); for _ in 0..length { - match parse_resp_frame(input, next, depth + 1)? { + match parse_resp_frame_with_limits( + input, + next, + depth + 1, + limits, + remaining_bulk_bytes, + )? { ParseResult::Complete(value, consumed) => { values.push(value); next = next.saturating_add(consumed); @@ -211,24 +322,32 @@ fn parse_resp_frame(input: &[u8], offset: usize, depth: usize) -> Result Err("CPA usage 订阅响应类型无效".to_string()), + _ => Err("CPA usage RESP 响应类型无效".to_string()), } } -fn resp_line(input: &[u8], start: usize) -> Option<(&[u8], usize)> { - let relative_end = input - .get(start..)? - .windows(2) - .position(|pair| pair == b"\r\n")?; +fn resp_line(input: &[u8], start: usize) -> Result, String> { + let Some(remaining) = input.get(start..) else { + return Ok(None); + }; + let Some(relative_end) = remaining.windows(2).position(|pair| pair == b"\r\n") else { + if remaining.len() > MAX_RESP_LINE_BYTES + 1 { + return Err("CPA usage RESP 行超过大小限制".to_string()); + } + return Ok(None); + }; + if relative_end > MAX_RESP_LINE_BYTES { + return Err("CPA usage RESP 行超过大小限制".to_string()); + } let end = start + relative_end; - Some((&input[start..end], end + 2)) + Ok(Some((&input[start..end], end + 2))) } fn parse_resp_length(value: &[u8]) -> Result { std::str::from_utf8(value) - .map_err(|_| "CPA usage 订阅长度不是 UTF-8".to_string())? + .map_err(|_| "CPA usage RESP 长度不是 UTF-8".to_string())? .parse::() - .map_err(|_| "CPA usage 订阅长度无效".to_string()) + .map_err(|_| "CPA usage RESP 长度无效".to_string()) } fn is_subscription_ack(value: &RespValue) -> bool { @@ -277,7 +396,7 @@ mod tests { #[test] fn rejects_oversized_arrays() { - let error = match parse_resp_frame(b"*10001\r\n", 0, 0) { + let error = match parse_resp_frame(b"*17\r\n", 0, 0) { Err(error) => error, _ => panic!("expected oversized RESP array to fail"), }; @@ -317,4 +436,41 @@ mod tests { ); server.await.unwrap(); } + + #[tokio::test] + async fn authenticates_and_pulls_usage_queue_batch() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + let server = tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + let expected_auth = b"*2\r\n$4\r\nAUTH\r\n$6\r\nsecret\r\n"; + let mut auth = vec![0_u8; expected_auth.len()]; + stream.read_exact(&mut auth).await.unwrap(); + assert_eq!(auth, expected_auth); + stream.write_all(b"+OK\r\n").await.unwrap(); + + let expected_pop = b"*3\r\n$4\r\nLPOP\r\n$5\r\nusage\r\n$2\r\n17\r\n"; + let mut pop = vec![0_u8; expected_pop.len()]; + stream.read_exact(&mut pop).await.unwrap(); + assert_eq!(pop, expected_pop); + let mut response = "*17\r\n".to_string(); + for request_id in 0..17 { + let message = format!("{{\"request_id\":{request_id}}}"); + response.push_str(&format!("${}\r\n{message}\r\n", message.len())); + } + stream.write_all(response.as_bytes()).await.unwrap(); + }); + + let messages = pull_usage_queue(port, "secret", "usage", 17).await.unwrap(); + assert_eq!(messages.len(), 17); + assert_eq!( + messages.first().map(String::as_str), + Some("{\"request_id\":0}") + ); + assert_eq!( + messages.last().map(String::as_str), + Some("{\"request_id\":16}") + ); + server.await.unwrap(); + } } diff --git a/src-tauri/src/usage/token.rs b/src-tauri/src/usage/token.rs new file mode 100644 index 0000000..c5e4450 --- /dev/null +++ b/src-tauri/src/usage/token.rs @@ -0,0 +1,398 @@ +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub(super) struct TokenValues { + pub(super) input: u64, + pub(super) output: u64, + pub(super) reasoning: u64, + pub(super) cached: u64, + pub(super) cache_read: u64, + pub(super) cache_read_present: bool, + pub(super) cache_creation: u64, + pub(super) total: u64, + pub(super) clamped: ClampedTokenFields, +} + +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub(super) struct ClampedTokenFields { + pub(super) input: bool, + pub(super) output: bool, + pub(super) reasoning: bool, + pub(super) cached: bool, + pub(super) cache_read: bool, + pub(super) cache_creation: bool, + pub(super) total: bool, +} + +impl ClampedTokenFields { + fn blocks_parent_contract(self) -> bool { + self.input || self.output || self.reasoning || self.cache_read || self.cache_creation + } + + fn blocks_zero_total(self) -> bool { + self.input || self.output || self.total + } + + fn blocks_reasoning_evidence(self) -> bool { + self.blocks_zero_total() || self.reasoning + } + + fn any(self) -> bool { + self.blocks_parent_contract() || self.cached || self.total + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum TokenHandler { + Claude, + Gemini, + ResponsesInclusive, + Strict, + OpenAiCompatibility, +} + +pub(super) fn normalize( + executor_type: &str, + provider: &str, + auth_type: &str, + mut tokens: TokenValues, +) -> TokenValues { + let (handler, parser_contract) = resolve_handler(executor_type, provider, auth_type); + let raw_tokens = tokens; + if handler != TokenHandler::Claude + && !tokens.clamped.cached + && !tokens.clamped.cache_read + && tokens.cache_read == 0 + && tokens.cached > 0 + && (!tokens.cache_read_present || !parser_contract) + { + tokens.cache_read = tokens.cached; + } + + match handler { + TokenHandler::Claude => normalize_claude(tokens, parser_contract), + TokenHandler::Gemini => normalize_gemini(tokens, parser_contract), + TokenHandler::ResponsesInclusive => { + normalize_responses(tokens, parser_contract, executor_type, raw_tokens) + } + TokenHandler::OpenAiCompatibility => normalize_openai_compatibility(tokens), + TokenHandler::Strict => reconcile_zero_total(tokens), + } +} + +fn resolve_handler(executor_type: &str, provider: &str, auth_type: &str) -> (TokenHandler, bool) { + let executor = executor_type.trim().to_ascii_lowercase(); + let handler = match executor.as_str() { + "claudeexecutor" => Some(TokenHandler::Claude), + "geminiexecutor" + | "geminivertexexecutor" + | "geminicliexecutor" + | "aistudioexecutor" + | "antigravityexecutor" => Some(TokenHandler::Gemini), + "codexexecutor" + | "codexwebsocketsexecutor" + | "codexautoexecutor" + | "xaiexecutor" + | "xaiwebsocketsexecutor" + | "xaiautoexecutor" => Some(TokenHandler::ResponsesInclusive), + "kimiexecutor" => Some(TokenHandler::Strict), + "openaicompatexecutor" => Some(TokenHandler::OpenAiCompatibility), + _ => None, + }; + if let Some(handler) = handler { + return (handler, true); + } + + if !auth_type.trim().eq_ignore_ascii_case("oauth") { + return (TokenHandler::Strict, false); + } + let identity = provider.trim().to_ascii_lowercase(); + let handler = match identity.as_str() { + "claude" | "anthropic" => TokenHandler::Claude, + "gemini" + | "vertex" + | "gemini-cli" + | "gemini-cli-code-assist" + | "gemini-interactions" + | "aistudio" + | "ai-studio" + | "antigravity" => TokenHandler::Gemini, + "codex" | "xai" => TokenHandler::ResponsesInclusive, + "kimi" | "moonshot" => TokenHandler::Strict, + "openai" | "openai-compatible" | "openai_compatibility" | "openai-compatibility" => { + TokenHandler::OpenAiCompatibility + } + value if value.starts_with("openai-compatible-") => TokenHandler::OpenAiCompatibility, + _ => TokenHandler::Strict, + }; + (handler, false) +} + +fn normalize_claude(mut tokens: TokenValues, parser_contract: bool) -> TokenValues { + tokens.cached = tokens.cache_read; + if tokens.clamped.input || tokens.clamped.cache_read || tokens.clamped.cache_creation { + return reconcile_zero_total(tokens); + } + let raw_input = tokens.input; + let raw_total = tokens.total; + let Some(canonical_input) = + checked_sum(&[tokens.input, tokens.cache_read, tokens.cache_creation]) + else { + return tokens; + }; + tokens.input = canonical_input; + + let cache_total = checked_sum(&[tokens.cache_read, tokens.cache_creation]); + let raw_expected = checked_sum(&[raw_input, tokens.output]); + let canonical_expected = checked_sum(&[tokens.input, tokens.output]); + let legacy_missing_cache = !tokens.clamped.blocks_zero_total() + && cache_total.is_some_and(|value| value > 0) + && raw_expected == Some(raw_total) + && canonical_expected.is_some_and(|value| value != raw_total); + if parser_contract || legacy_missing_cache { + reconcile_canonical_total(tokens) + } else { + reconcile_zero_total(tokens) + } +} + +fn normalize_gemini(mut tokens: TokenValues, parser_contract: bool) -> TokenValues { + let should_fold = if tokens.clamped.output || tokens.clamped.reasoning { + false + } else if parser_contract { + tokens.reasoning > 0 + } else if tokens.clamped.blocks_reasoning_evidence() || tokens.reasoning == 0 { + false + } else if tokens.total == 0 { + true + } else if checked_sum(&[tokens.input, tokens.output]) == Some(tokens.total) { + false + } else { + checked_sum(&[tokens.input, tokens.output, tokens.reasoning]) == Some(tokens.total) + }; + if should_fold { + let Some(output) = checked_sum(&[tokens.output, tokens.reasoning]) else { + return tokens; + }; + tokens.output = output; + } + if parser_contract && !tokens.clamped.blocks_parent_contract() { + reconcile_canonical_total(tokens) + } else { + reconcile_zero_total(tokens) + } +} + +fn normalize_responses( + tokens: TokenValues, + parser_contract: bool, + executor_type: &str, + raw_tokens: TokenValues, +) -> TokenValues { + let legacy_codex_cached_only = parser_contract + && executor_type.trim().eq_ignore_ascii_case("CodexExecutor") + && !raw_tokens.clamped.any() + && raw_tokens.input == 0 + && raw_tokens.output == 0 + && raw_tokens.reasoning == 0 + && raw_tokens.cached > 0 + && raw_tokens.cache_read == 0 + && raw_tokens.cache_creation == 0 + && raw_tokens.total == raw_tokens.cached; + if legacy_codex_cached_only { + return tokens; + } + if parser_contract && !tokens.clamped.blocks_parent_contract() { + reconcile_canonical_total(tokens) + } else { + reconcile_zero_total(tokens) + } +} + +fn normalize_openai_compatibility(mut tokens: TokenValues) -> TokenValues { + if !tokens.clamped.blocks_reasoning_evidence() + && tokens.reasoning > 0 + && tokens.total > 0 + && checked_sum(&[tokens.input, tokens.output]) != Some(tokens.total) + && checked_sum(&[tokens.input, tokens.output, tokens.reasoning]) == Some(tokens.total) + { + if let Some(output) = checked_sum(&[tokens.output, tokens.reasoning]) { + tokens.output = output; + } + } + reconcile_zero_total(tokens) +} + +fn reconcile_canonical_total(mut tokens: TokenValues) -> TokenValues { + if tokens.clamped.blocks_parent_contract() || tokens.clamped.total { + return tokens; + } + let Some(cache_total) = checked_sum(&[tokens.cache_read, tokens.cache_creation]) else { + return reconcile_zero_total(tokens); + }; + if tokens.input < cache_total || tokens.output < tokens.reasoning { + return reconcile_zero_total(tokens); + } + if let Some(total) = checked_sum(&[tokens.input, tokens.output]) { + tokens.total = total; + } + tokens +} + +fn reconcile_zero_total(mut tokens: TokenValues) -> TokenValues { + if tokens.total != 0 || tokens.clamped.blocks_zero_total() { + return tokens; + } + if let Some(total) = checked_sum(&[tokens.input, tokens.output]) { + if total > 0 { + tokens.total = total; + return tokens; + } + } + if tokens.cache_read > 0 { + tokens.total = tokens.cache_read; + } + tokens +} + +fn checked_sum(values: &[u64]) -> Option { + values + .iter() + .try_fold(0_u64, |total, value| total.checked_add(*value)) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn claude_executor_folds_cache_into_input_and_reconciles_total() { + let tokens = normalize( + "ClaudeExecutor", + "anthropic", + "oauth", + TokenValues { + input: 100, + output: 20, + cached: 10, + cache_read: 10, + cache_read_present: true, + cache_creation: 5, + total: 120, + ..TokenValues::default() + }, + ); + assert_eq!(tokens.input, 115); + assert_eq!(tokens.cached, 10); + assert_eq!(tokens.total, 135); + } + + #[test] + fn gemini_executor_folds_reasoning_into_output() { + let tokens = normalize( + "GeminiExecutor", + "gemini", + "oauth", + TokenValues { + input: 11, + output: 7, + reasoning: 3, + total: 21, + ..TokenValues::default() + }, + ); + assert_eq!(tokens.output, 10); + assert_eq!(tokens.total, 21); + } + + #[test] + fn responses_executor_preserves_inclusive_output() { + let tokens = normalize( + "CodexExecutor", + "codex", + "oauth", + TokenValues { + input: 100, + output: 20, + reasoning: 5, + cache_read: 30, + cache_read_present: true, + total: 120, + ..TokenValues::default() + }, + ); + assert_eq!(tokens.output, 20); + assert_eq!(tokens.total, 120); + } + + #[test] + fn openai_compatibility_only_folds_proven_separated_reasoning() { + let tokens = normalize( + "OpenAICompatExecutor", + "openai", + "apikey", + TokenValues { + input: 1_000, + output: 20, + reasoning: 50, + total: 1_070, + ..TokenValues::default() + }, + ); + assert_eq!(tokens.output, 70); + assert_eq!(tokens.total, 1_070); + } + + #[test] + fn explicit_zero_cache_read_is_not_backfilled_from_legacy_cached() { + let tokens = normalize( + "CodexExecutor", + "codex", + "oauth", + TokenValues { + input: 100, + output: 20, + cached: 30, + cache_read_present: true, + total: 120, + ..TokenValues::default() + }, + ); + assert_eq!(tokens.cache_read, 0); + } + + #[test] + fn identity_hint_keeps_keeper_legacy_cache_read_fallback() { + let tokens = normalize( + "", + "codex", + "oauth", + TokenValues { + input: 100, + output: 20, + cached: 30, + cache_read_present: true, + total: 120, + ..TokenValues::default() + }, + ); + assert_eq!(tokens.cache_read, 30); + } + + #[test] + fn unknown_identity_does_not_use_fuzzy_provider_matching() { + let tokens = normalize( + "", + "custom-anthropic-proxy", + "oauth", + TokenValues { + input: 100, + output: 20, + cache_read: 30, + cache_read_present: true, + total: 120, + ..TokenValues::default() + }, + ); + assert_eq!(tokens.input, 100); + assert_eq!(tokens.total, 120); + } +} diff --git a/src/i18n/ja.ts b/src/i18n/ja.ts index 881ef65..1aad1fe 100644 --- a/src/i18n/ja.ts +++ b/src/i18n/ja.ts @@ -414,7 +414,7 @@ export const jaOverrides = { 'usage.stat.averageLatency': '平均レイテンシ', 'usage.stat.tps': 'TPS', 'usage.stat.performanceMeta': '有効サンプル {samples} 件 · RPM {rpm}', - 'usage.stat.performanceMetaTitle': 'TPS {tps}、有効サンプル {samples} 件、RPM {rpm}、平均レイテンシ {latency} ms', + 'usage.stat.performanceMetaTitle': 'TPS {tps}（総レイテンシで加重）、有効サンプル {samples} 件、RPM {rpm}、平均レイテンシ {latency} ms', 'usage.stat.cacheHitRate': 'キャッシュヒット率', 'usage.stat.cacheHitMeta': 'キャッシュ読取 {hit} · 入力 {input}', 'usage.stat.cacheHitMetaTitle': 'キャッシュヒット率 {rate}%、読取 {hit} Token、入力 {input} Token', @@ -498,7 +498,7 @@ export const jaOverrides = { 'usage.column.cache': 'キャッシュ', 'usage.column.cacheRate': 'キャッシュ率', 'usage.column.total': '合計', - 'usage.column.speed': '生成速度', + 'usage.column.speed': '出力速度', 'usage.key.noRemark': 'メモなし', 'usage.previous': '前へ', 'usage.next': '次へ', diff --git a/src/i18n/locales/en.ts b/src/i18n/locales/en.ts index 12b7499..1cfcfa4 100644 --- a/src/i18n/locales/en.ts +++ b/src/i18n/locales/en.ts @@ -401,7 +401,7 @@ export const en: Record = { 'usage.stat.averageLatency': 'Average Latency', 'usage.stat.tps': 'TPS', 'usage.stat.performanceMeta': '{samples} samples · RPM {rpm}', - 'usage.stat.performanceMetaTitle': 'TPS {tps}, {samples} valid samples, RPM {rpm}, average latency {latency} ms', + 'usage.stat.performanceMetaTitle': 'TPS {tps} (weighted by total latency), {samples} valid samples, RPM {rpm}, average latency {latency} ms', 'usage.stat.cacheHitRate': 'Cache Hit Rate', 'usage.stat.cacheHitMeta': 'Cache read {hit} · Input {input}', 'usage.stat.cacheHitMetaTitle': 'Cache hit rate {rate}%, cache read {hit} tokens, input {input} tokens', @@ -486,7 +486,7 @@ export const en: Record = { 'usage.column.cache': 'Cache', 'usage.column.cacheRate': 'Cache Rate', 'usage.column.total': 'Total', - 'usage.column.speed': 'Generation Speed', + 'usage.column.speed': 'Output Speed', 'usage.key.noRemark': 'No remark', 'usage.previous': 'Previous', 'usage.next': 'Next', diff --git a/src/i18n/locales/zh-CN.ts b/src/i18n/locales/zh-CN.ts index f048e05..e34d03c 100644 --- a/src/i18n/locales/zh-CN.ts +++ b/src/i18n/locales/zh-CN.ts @@ -400,7 +400,7 @@ export const zhCN = { 'usage.stat.averageLatency': '平均延迟', 'usage.stat.tps': 'TPS', 'usage.stat.performanceMeta': '{samples} 条有效样本 · RPM {rpm}', - 'usage.stat.performanceMetaTitle': 'TPS {tps}，{samples} 条有效样本，RPM {rpm}，平均延迟 {latency} ms', + 'usage.stat.performanceMetaTitle': 'TPS {tps}（按总延迟加权），{samples} 条有效样本，RPM {rpm}，平均延迟 {latency} ms', 'usage.stat.cacheHitRate': '缓存命中率', 'usage.stat.cacheHitMeta': '缓存读取 {hit} · 输入 {input}', 'usage.stat.cacheHitMetaTitle': '缓存命中率 {rate}%，缓存读取 {hit} Token，输入 {input} Token', @@ -485,7 +485,7 @@ export const zhCN = { 'usage.column.cache': '缓存', 'usage.column.cacheRate': '缓存率', 'usage.column.total': '总计', - 'usage.column.speed': '生成速度', + 'usage.column.speed': '输出速度', 'usage.key.noRemark': '未备注', 'usage.previous': '上一页', 'usage.next': '下一页', diff --git a/src/pages/UsageRecordsPage.tsx b/src/pages/UsageRecordsPage.tsx index 536b71a..7cbe337 100644 --- a/src/pages/UsageRecordsPage.tsx +++ b/src/pages/UsageRecordsPage.tsx @@ -400,9 +400,9 @@ export function UsageRecordsPage() { ); const loadData = useCallback( - (quiet = false) => { + (quiet = false, immediate = !quiet) => { if (!quiet) setLoading(true); - return schedulerRef.current!.schedule(() => executeLoadData(quiet), !quiet); + return schedulerRef.current!.schedule(() => executeLoadData(quiet), immediate); }, [executeLoadData], ); @@ -419,7 +419,7 @@ export function UsageRecordsPage() { let disposed = false; let unlisten: (() => void) | null = null; const refresh = () => { - if (!disposed && !document.hidden) void loadData(true); + if (!disposed && !document.hidden) void loadData(true, true); }; listen('usage-records-updated', refresh) .then((stop) => { @@ -1440,7 +1440,6 @@ function UsageEventCell({ const value = formatGenerationSpeed({ outputTokens: record.tokens.output_tokens, latencyMs: record.latency_ms, - ttftMs: record.ttft_ms, }); return {value}; } diff --git a/src/services/usageMetrics.ts b/src/services/usageMetrics.ts index 1556834..aa7af60 100644 --- a/src/services/usageMetrics.ts +++ b/src/services/usageMetrics.ts @@ -1,7 +1,6 @@ type GenerationSpeedInput = { outputTokens: number; latencyMs: number; - ttftMs: number | null; }; type CacheReadRateInput = { @@ -12,21 +11,17 @@ type CacheReadRateInput = { export const calculateGenerationSpeed = ({ outputTokens, latencyMs, - ttftMs, }: GenerationSpeedInput): number | null => { if ( !Number.isFinite(outputTokens) || !Number.isFinite(latencyMs) || - ttftMs === null || - !Number.isFinite(ttftMs) || outputTokens <= 0 || - ttftMs <= 0 || - latencyMs <= ttftMs + latencyMs <= 0 ) { return null; } - const speed = outputTokens / ((latencyMs - ttftMs) / 1_000); + const speed = outputTokens / (latencyMs / 1_000); return Number.isFinite(speed) && speed > 0 ? speed : null; }; diff --git a/tests/usageMetrics.test.ts b/tests/usageMetrics.test.ts index 1645103..c2001a6 100644 --- a/tests/usageMetrics.test.ts +++ b/tests/usageMetrics.test.ts @@ -7,18 +7,19 @@ import { } from '../src/services/usageMetrics'; describe('generation speed', () => { - test('uses only the generation interval after the first token', () => { - const input = { outputTokens: 344, latencyMs: 10_600, ttftMs: 2_770 }; - expect(calculateGenerationSpeed(input)).toBeCloseTo(43.9336, 4); - expect(formatGenerationSpeed(input)).toBe('43.9 t/s'); + test('uses the total request latency without requiring TTFT', () => { + const input = { outputTokens: 344, latencyMs: 10_600 }; + expect(calculateGenerationSpeed(input)).toBeCloseTo(32.4528, 4); + expect(formatGenerationSpeed(input)).toBe('32.5 t/s'); }); test.each([ - { outputTokens: 344, latencyMs: 10_600, ttftMs: 0 }, - { outputTokens: 344, latencyMs: 10_600, ttftMs: null }, - { outputTokens: 344, latencyMs: 2_770, ttftMs: 2_770 }, - { outputTokens: 344, latencyMs: 2_000, ttftMs: 2_770 }, - { outputTokens: 0, latencyMs: 10_600, ttftMs: 2_770 }, + { outputTokens: 344, latencyMs: 0 }, + { outputTokens: 344, latencyMs: -1 }, + { outputTokens: 0, latencyMs: 10_600 }, + { outputTokens: -1, latencyMs: 10_600 }, + { outputTokens: Number.NaN, latencyMs: 10_600 }, + { outputTokens: 344, latencyMs: Number.POSITIVE_INFINITY }, ])('returns an em dash when generation speed cannot be calculated', (input) => { expect(calculateGenerationSpeed(input)).toBeNull(); expect(formatGenerationSpeed(input)).toBe('—');