From 087622bf6dbe8d485d1bb0edf6617d9113d472e3 Mon Sep 17 00:00:00 2001 From: YoVinchen Date: Wed, 3 Dec 2025 21:09:51 +0800 Subject: [PATCH] feat(proxy/usage): enhance token parser and logger for multi-format support Parser enhancements: - Add OpenAI Chat Completions format parsing (prompt_tokens/completion_tokens) - Add model field to TokenUsage for actual model name extraction - Add from_codex_response_adjusted() for proper cache token handling - Add debug logging for better stream event tracing Logger enhancements: - Add first_token_ms, provider_type, is_streaming, cost_multiplier fields - Extend RequestLog struct with full metadata tracking - Update log_with_calculation() signature for new fields Calculator: Update tests with model field in TokenUsage. --- src-tauri/src/proxy/usage/calculator.rs | 4 + src-tauri/src/proxy/usage/logger.rs | 30 +++- src-tauri/src/proxy/usage/parser.rs | 195 +++++++++++++++++++++++- 3 files changed, 224 insertions(+), 5 deletions(-) diff --git a/src-tauri/src/proxy/usage/calculator.rs b/src-tauri/src/proxy/usage/calculator.rs index e7cd335fa..7295e13ae 100644 --- a/src-tauri/src/proxy/usage/calculator.rs +++ b/src-tauri/src/proxy/usage/calculator.rs @@ -105,6 +105,7 @@ mod tests { output_tokens: 500, cache_read_tokens: 200, cache_creation_tokens: 100, + model: None, }; let pricing = ModelPricing::from_strings("3.0", "15.0", "0.3", "3.75").unwrap(); @@ -134,6 +135,7 @@ mod tests { output_tokens: 0, cache_read_tokens: 0, cache_creation_tokens: 0, + model: None, }; let pricing = ModelPricing::from_strings("3.0", "15.0", "0", "0").unwrap(); @@ -153,6 +155,7 @@ mod tests { output_tokens: 500, cache_read_tokens: 0, cache_creation_tokens: 0, + model: None, }; let multiplier = Decimal::from_str("1.0").unwrap(); @@ -168,6 +171,7 @@ mod tests { output_tokens: 1, cache_read_tokens: 1, cache_creation_tokens: 1, + model: None, }; let pricing = ModelPricing::from_strings("0.075", "0.3", "0.01875", "0.075").unwrap(); diff --git a/src-tauri/src/proxy/usage/logger.rs b/src-tauri/src/proxy/usage/logger.rs index 8fb995e3f..6dccb8dcf 100644 --- a/src-tauri/src/proxy/usage/logger.rs +++ b/src-tauri/src/proxy/usage/logger.rs @@ -18,9 +18,16 @@ pub struct RequestLog { pub usage: TokenUsage, pub cost: Option, pub latency_ms: u64, + pub first_token_ms: Option, pub status_code: u16, pub error_message: Option, pub session_id: Option, + /// 供应商类型 (claude, claude_auth, codex, gemini, gemini_cli, openrouter) + pub provider_type: Option, + /// 是否为流式请求 + pub is_streaming: bool, + /// 成本倍数 + pub cost_multiplier: String, } /// 使用量记录器 @@ -66,8 +73,9 @@ impl<'a> UsageLogger<'a> { request_id, provider_id, app_type, model, input_tokens, output_tokens, cache_read_tokens, cache_creation_tokens, input_cost_usd, output_cost_usd, cache_read_cost_usd, cache_creation_cost_usd, total_cost_usd, - latency_ms, status_code, error_message, session_id, created_at - ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18)", + latency_ms, first_token_ms, status_code, error_message, session_id, + provider_type, is_streaming, cost_multiplier, created_at + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18, ?19, ?20, ?21, ?22)", rusqlite::params![ log.request_id, log.provider_id, @@ -83,9 +91,13 @@ impl<'a> UsageLogger<'a> { cache_creation_cost, total_cost, log.latency_ms as i64, + log.first_token_ms.map(|v| v as i64), log.status_code as i64, log.error_message, log.session_id, + log.provider_type, + log.is_streaming as i64, + log.cost_multiplier, created_at, ], ) @@ -114,9 +126,13 @@ impl<'a> UsageLogger<'a> { usage: TokenUsage::default(), cost: None, latency_ms, + first_token_ms: None, status_code, error_message: Some(error_message), session_id: None, + provider_type: None, + is_streaming: false, + cost_multiplier: "1.0".to_string(), }; self.log_request(&log) @@ -147,8 +163,11 @@ impl<'a> UsageLogger<'a> { usage: TokenUsage, cost_multiplier: Decimal, latency_ms: u64, + first_token_ms: Option, status_code: u16, session_id: Option, + provider_type: Option, + is_streaming: bool, ) -> Result<(), AppError> { let pricing = self.get_model_pricing(&model)?; @@ -166,9 +185,13 @@ impl<'a> UsageLogger<'a> { usage, cost, latency_ms, + first_token_ms, status_code, error_message: None, session_id, + provider_type, + is_streaming, + cost_multiplier: cost_multiplier.to_string(), }; self.log_request(&log) @@ -201,6 +224,7 @@ mod tests { output_tokens: 500, cache_read_tokens: 0, cache_creation_tokens: 0, + model: None, }; logger.log_with_calculation( @@ -213,6 +237,8 @@ mod tests { 100, 200, None, + Some("claude".to_string()), + false, )?; // 验证记录已插入 diff --git a/src-tauri/src/proxy/usage/parser.rs b/src-tauri/src/proxy/usage/parser.rs index 7945ff549..65fe61a69 100644 --- a/src-tauri/src/proxy/usage/parser.rs +++ b/src-tauri/src/proxy/usage/parser.rs @@ -16,6 +16,8 @@ pub struct TokenUsage { pub output_tokens: u32, pub cache_read_tokens: u32, pub cache_creation_tokens: u32, + /// 从响应中提取的实际模型名称（如果可用） + pub model: Option, } /// API 类型 @@ -43,6 +45,7 @@ impl TokenUsage { .get("cache_creation_input_tokens") .and_then(|v| v.as_u64()) .unwrap_or(0) as u32, + model: None, }) } @@ -96,15 +99,33 @@ impl TokenUsage { output_tokens: usage.get("completion_tokens")?.as_u64()? as u32, cache_read_tokens: 0, cache_creation_tokens: 0, + model: None, }) } /// 从 Codex API 非流式响应解析 pub fn from_codex_response(body: &Value) -> Option { - let usage = body.get("usage")?; + let usage = body.get("usage"); + if usage.is_none() { + log::debug!( + "[Codex] 响应中没有 usage 字段，body keys: {:?}", + body.as_object().map(|o| o.keys().collect::>()) + ); + return None; + } + let usage = usage?; + + let input_tokens = usage.get("input_tokens").and_then(|v| v.as_u64()); + let output_tokens = usage.get("output_tokens").and_then(|v| v.as_u64()); + + if input_tokens.is_none() || output_tokens.is_none() { + log::debug!("[Codex] usage 字段缺少 input_tokens 或 output_tokens，usage: {usage:?}"); + return None; + } + Some(Self { - input_tokens: usage.get("input_tokens")?.as_u64()? as u32, - output_tokens: usage.get("output_tokens")?.as_u64()? as u32, + input_tokens: input_tokens? as u32, + output_tokens: output_tokens? as u32, cache_read_tokens: usage .get("cache_read_input_tokens") .and_then(|v| v.as_u64()) @@ -113,27 +134,110 @@ impl TokenUsage { .get("cache_creation_input_tokens") .and_then(|v| v.as_u64()) .unwrap_or(0) as u32, + model: None, + }) + } + + /// 从 Codex API 响应解析并调整 input_tokens + /// + /// Codex 的 input_tokens 需要减去 cached_tokens 以获得实际计费的 token 数 + /// 公式: adjusted_input = max(input_tokens - cached_tokens, 0) + #[allow(dead_code)] + pub fn from_codex_response_adjusted(body: &Value) -> Option { + let usage = body.get("usage")?; + let input_tokens = usage.get("input_tokens")?.as_u64()? as u32; + let output_tokens = usage.get("output_tokens")?.as_u64()? as u32; + + // 获取 cached_tokens (可能在 input_tokens_details 中) + let cached_tokens = usage + .get("input_tokens_details") + .and_then(|d| d.get("cached_tokens")) + .and_then(|v| v.as_u64()) + .unwrap_or(0) as u32; + + // 调整 input_tokens: 减去 cached_tokens + let adjusted_input = input_tokens.saturating_sub(cached_tokens); + + Some(Self { + input_tokens: adjusted_input, + output_tokens, + cache_read_tokens: cached_tokens, + cache_creation_tokens: usage + .get("cache_creation_input_tokens") + .and_then(|v| v.as_u64()) + .unwrap_or(0) as u32, + model: None, }) } /// 从 Codex API 流式响应解析 #[allow(dead_code)] pub fn from_codex_stream_events(events: &[Value]) -> Option { + log::debug!("[Codex] 解析流式事件，共 {} 个事件", events.len()); for event in events { if let Some(event_type) = event.get("type").and_then(|v| v.as_str()) { + log::debug!("[Codex] 事件类型: {event_type}"); if event_type == "response.completed" { if let Some(response) = event.get("response") { + log::debug!("[Codex] 找到 response.completed 事件，解析 usage"); return Self::from_codex_response(response); } } } } + log::debug!("[Codex] 未找到 response.completed 事件"); + None + } + + /// 从 OpenAI Chat Completions API 响应解析 (prompt_tokens, completion_tokens) + pub fn from_openai_response(body: &Value) -> Option { + let usage = body.get("usage")?; + + // OpenAI 使用 prompt_tokens 和 completion_tokens + let prompt_tokens = usage.get("prompt_tokens").and_then(|v| v.as_u64())?; + let completion_tokens = usage.get("completion_tokens").and_then(|v| v.as_u64())?; + + // 获取 cached_tokens (可能在 prompt_tokens_details 中) + let cached_tokens = usage + .get("prompt_tokens_details") + .and_then(|d| d.get("cached_tokens")) + .and_then(|v| v.as_u64()) + .unwrap_or(0) as u32; + + Some(Self { + input_tokens: prompt_tokens as u32, + output_tokens: completion_tokens as u32, + cache_read_tokens: cached_tokens, + cache_creation_tokens: 0, + model: None, + }) + } + + /// 从 OpenAI Chat Completions API 流式响应解析 + pub fn from_openai_stream_events(events: &[Value]) -> Option { + log::debug!("[Codex] 解析 OpenAI 流式事件，共 {} 个事件", events.len()); + // OpenAI 流式响应在最后一个 chunk 中包含 usage + for event in events.iter().rev() { + if let Some(usage) = event.get("usage") { + if !usage.is_null() { + log::debug!("[Codex] 找到 usage: {usage:?}"); + return Self::from_openai_response(event); + } + } + } + log::debug!("[Codex] 未找到 usage 信息"); None } /// 从 Gemini API 非流式响应解析 pub fn from_gemini_response(body: &Value) -> Option { let usage = body.get("usageMetadata")?; + // 提取实际使用的模型名称（modelVersion 字段） + let model = body + .get("modelVersion") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + Some(Self { input_tokens: usage.get("promptTokenCount")?.as_u64()? as u32, output_tokens: usage.get("candidatesTokenCount")?.as_u64()? as u32, @@ -142,6 +246,7 @@ impl TokenUsage { .and_then(|v| v.as_u64()) .unwrap_or(0) as u32, cache_creation_tokens: 0, + model, }) } @@ -151,6 +256,7 @@ impl TokenUsage { let mut total_input = 0u32; let mut total_output = 0u32; let mut total_cache_read = 0u32; + let mut model: Option = None; for chunk in chunks { if let Some(usage) = chunk.get("usageMetadata") { @@ -167,6 +273,13 @@ impl TokenUsage { .and_then(|v| v.as_u64()) .unwrap_or(0) as u32; } + + // 提取实际使用的模型名称（modelVersion 字段） + if model.is_none() { + if let Some(model_version) = chunk.get("modelVersion").and_then(|v| v.as_str()) { + model = Some(model_version.to_string()); + } + } } if total_input > 0 || total_output > 0 { @@ -175,6 +288,7 @@ impl TokenUsage { output_tokens: total_output, cache_read_tokens: total_cache_read, cache_creation_tokens: 0, + model, }) } else { None @@ -252,6 +366,7 @@ mod tests { #[test] fn test_gemini_response_parsing() { let response = json!({ + "modelVersion": "gemini-3-pro-high", "usageMetadata": { "promptTokenCount": 100, "candidatesTokenCount": 50, @@ -264,5 +379,79 @@ mod tests { assert_eq!(usage.output_tokens, 50); assert_eq!(usage.cache_read_tokens, 20); assert_eq!(usage.cache_creation_tokens, 0); + assert_eq!(usage.model, Some("gemini-3-pro-high".to_string())); + } + + #[test] + fn test_gemini_response_parsing_no_model() { + // 测试没有 modelVersion 字段的情况 + let response = json!({ + "usageMetadata": { + "promptTokenCount": 100, + "candidatesTokenCount": 50, + "cachedContentTokenCount": 20 + } + }); + + let usage = TokenUsage::from_gemini_response(&response).unwrap(); + assert_eq!(usage.input_tokens, 100); + assert_eq!(usage.output_tokens, 50); + assert_eq!(usage.cache_read_tokens, 20); + assert_eq!(usage.cache_creation_tokens, 0); + assert_eq!(usage.model, None); + } + + #[test] + fn test_codex_response_adjusted() { + let response = json!({ + "usage": { + "input_tokens": 1000, + "output_tokens": 500, + "input_tokens_details": { + "cached_tokens": 300 + } + } + }); + + let usage = TokenUsage::from_codex_response_adjusted(&response).unwrap(); + // input_tokens 应该被调整: 1000 - 300 = 700 + assert_eq!(usage.input_tokens, 700); + assert_eq!(usage.output_tokens, 500); + assert_eq!(usage.cache_read_tokens, 300); + } + + #[test] + fn test_codex_response_adjusted_no_cache() { + let response = json!({ + "usage": { + "input_tokens": 1000, + "output_tokens": 500 + } + }); + + let usage = TokenUsage::from_codex_response_adjusted(&response).unwrap(); + // 没有 cached_tokens，input_tokens 保持不变 + assert_eq!(usage.input_tokens, 1000); + assert_eq!(usage.output_tokens, 500); + assert_eq!(usage.cache_read_tokens, 0); + } + + #[test] + fn test_codex_response_adjusted_saturating_sub() { + // 测试 cached_tokens > input_tokens 的边界情况 + let response = json!({ + "usage": { + "input_tokens": 100, + "output_tokens": 50, + "input_tokens_details": { + "cached_tokens": 200 + } + } + }); + + let usage = TokenUsage::from_codex_response_adjusted(&response).unwrap(); + // saturating_sub 确保不会下溢 + assert_eq!(usage.input_tokens, 0); + assert_eq!(usage.cache_read_tokens, 200); } }