diff --git a/src-tauri/src/proxy/forwarder.rs b/src-tauri/src/proxy/forwarder.rs index 9f9fcf6fc..245759d04 100644 --- a/src-tauri/src/proxy/forwarder.rs +++ b/src-tauri/src/proxy/forwarder.rs @@ -199,12 +199,19 @@ impl RequestForwarder { // 检查是否需要格式转换 let needs_transform = adapter.needs_transform(provider); + // 记录原始请求 JSON + log::info!( + "[{}] ====== 请求开始 ======\n>>> 原始请求 JSON:\n{}", + adapter.name(), + serde_json::to_string_pretty(body).unwrap_or_else(|_| body.to_string()) + ); + // 转换请求体（如果需要） let request_body = if needs_transform { log::info!("[{}] 转换请求格式 (Anthropic → OpenAI)", adapter.name()); let transformed = adapter.transform_request(body.clone(), provider)?; - log::debug!( - "[{}] 转换后的请求: {}", + log::info!( + "[{}] >>> 转换后的请求 JSON:\n{}", adapter.name(), serde_json::to_string_pretty(&transformed).unwrap_or_default() ); diff --git a/src-tauri/src/proxy/handlers.rs b/src-tauri/src/proxy/handlers.rs index 6d24449c4..00c3f71f0 100644 --- a/src-tauri/src/proxy/handlers.rs +++ b/src-tauri/src/proxy/handlers.rs @@ -7,11 +7,178 @@ use super::{ providers::{get_adapter, transform}, server::ProxyState, types::*, + usage::{logger::UsageLogger, parser::TokenUsage}, ProxyError, }; use crate::app_config::AppType; use axum::{extract::State, http::StatusCode, response::IntoResponse, Json}; +use bytes::Bytes; +use futures::stream::{Stream, StreamExt}; +use rust_decimal::Decimal; use serde_json::{json, Value}; +use std::{ + str::FromStr, + sync::{ + atomic::{AtomicBool, Ordering}, + Arc, + }, +}; +use tokio::sync::Mutex; + +/// 记录请求使用量 +async fn log_usage( + state: &ProxyState, + provider_id: &str, + app_type: &str, + model: &str, + usage: TokenUsage, + latency_ms: u64, + status_code: u16, +) { + let logger = UsageLogger::new(&state.db); + + // 获取 provider 的 cost_multiplier + let multiplier = match state.db.get_provider_by_id(provider_id, app_type) { + Ok(Some(p)) => { + if let Some(meta) = p.meta { + if let Some(cm) = meta.cost_multiplier { + Decimal::from_str(&cm).unwrap_or(Decimal::from(1)) + } else { + Decimal::from(1) + } + } else { + Decimal::from(1) + } + } + _ => Decimal::from(1), + }; + + let request_id = uuid::Uuid::new_v4().to_string(); + + if let Err(e) = logger.log_with_calculation( + request_id, + provider_id.to_string(), + app_type.to_string(), + model.to_string(), + usage, + multiplier, + latency_ms, + status_code, + None, + ) { + log::warn!("记录使用量失败: {e}"); + } +} + +type UsageCallback = Arc) + Send + Sync + 'static>; + +#[derive(Clone)] +struct SseUsageCollector { + inner: Arc, +} + +struct SseUsageCollectorInner { + events: Mutex>, + on_complete: UsageCallback, + finished: AtomicBool, +} + +impl SseUsageCollector { + fn new(callback: impl Fn(Vec) + Send + Sync + 'static) -> Self { + let on_complete: UsageCallback = Arc::new(callback); + Self { + inner: Arc::new(SseUsageCollectorInner { + events: Mutex::new(Vec::new()), + on_complete, + finished: AtomicBool::new(false), + }), + } + } + + async fn push(&self, event: Value) { + let mut events = self.inner.events.lock().await; + events.push(event); + } + + async fn finish(&self) { + if self.inner.finished.swap(true, Ordering::SeqCst) { + return; + } + + let events = { + let mut guard = self.inner.events.lock().await; + std::mem::take(&mut *guard) + }; + + (self.inner.on_complete)(events); + } +} + +/// 创建带日志记录的透传流 +fn create_logged_passthrough_stream( + stream: impl Stream> + Send + 'static, + tag: &'static str, + usage_collector: Option, +) -> impl Stream> + Send { + async_stream::stream! { + let mut buffer = String::new(); + let mut collector = usage_collector; + + tokio::pin!(stream); + + while let Some(chunk) = stream.next().await { + match chunk { + Ok(bytes) => { + let text = String::from_utf8_lossy(&bytes); + buffer.push_str(&text); + + // 尝试解析并记录完整的 SSE 事件 + while let Some(pos) = buffer.find("\n\n") { + let event_text = buffer[..pos].to_string(); + buffer = buffer[pos + 2..].to_string(); + + if !event_text.trim().is_empty() { + // 提取 data 部分并尝试解析为 JSON + for line in event_text.lines() { + if let Some(data) = line.strip_prefix("data: ") { + if data.trim() != "[DONE]" { + if let Ok(json_value) = serde_json::from_str::(data) { + if let Some(c) = &collector { + c.push(json_value.clone()).await; + } + log::info!( + "[{}] <<< SSE 事件:\n{}", + tag, + serde_json::to_string_pretty(&json_value).unwrap_or_else(|_| data.to_string()) + ); + } else { + log::info!("[{tag}] <<< SSE 数据: {data}"); + } + } else { + log::info!("[{tag}] <<< SSE: [DONE]"); + } + } + } + } + } + + yield Ok(bytes); + } + Err(e) => { + log::error!("[{tag}] 流错误: {e}"); + yield Err(std::io::Error::other(e.to_string())); + break; + } + } + } + + log::info!("[{}] ====== 流结束 ======", tag); + + if let Some(c) = collector.take() { + c.finish().await; + } + } +} /// 健康检查 pub async fn health_check() -> (StatusCode, Json) { @@ -37,6 +204,11 @@ pub async fn handle_messages( Json(body): Json, ) -> Result { let config = state.config.read().await.clone(); + let request_model = body + .get("model") + .and_then(|m| m.as_str()) + .unwrap_or("unknown") + .to_string(); // 选择目标 Provider let router = super::router::ProviderRouter::new(state.db.clone()); @@ -85,6 +257,40 @@ pub async fn handle_messages( let stream = response.bytes_stream(); let sse_stream = super::providers::streaming::create_anthropic_sse_stream(stream); + let usage_collector = { + let state = state.clone(); + let provider_id = provider.id.clone(); + let model = request_model.clone(); + let status_code = status.as_u16(); + SseUsageCollector::new(move |events| { + if let Some(usage) = TokenUsage::from_claude_stream_events(&events) { + let state = state.clone(); + let provider_id = provider_id.clone(); + let model = model.clone(); + tokio::spawn(async move { + log_usage( + &state, + &provider_id, + "claude", + &model, + usage, + 0, + status_code, + ) + .await; + }); + } else { + log::debug!("[Claude] OpenRouter 流式响应缺少 usage 统计，跳过消费记录"); + } + }) + }; + + let logged_stream = create_logged_passthrough_stream( + sse_stream, + "Claude/OpenRouter", + Some(usage_collector), + ); + let mut headers = axum::http::HeaderMap::new(); headers.insert( "Content-Type", @@ -99,7 +305,8 @@ pub async fn handle_messages( axum::http::HeaderValue::from_static("keep-alive"), ); - let body = axum::body::Body::from_stream(sse_stream); + let body = axum::body::Body::from_stream(logged_stream); + log::info!("[Claude] ====== 请求结束 (流式转换) ======"); return Ok((headers, body).into_response()); } else { // 非流式响应转换 @@ -124,6 +331,10 @@ pub async fn handle_messages( })?; log::info!("[Claude] 解析 OpenAI 响应成功"); + log::info!( + "[Claude] <<< OpenAI 响应 JSON:\n{}", + serde_json::to_string_pretty(&openai_response).unwrap_or_default() + ); let anthropic_response = transform::openai_to_anthropic(openai_response).map_err(|e| { @@ -132,11 +343,39 @@ pub async fn handle_messages( })?; log::info!("[Claude] 转换响应成功"); - log::debug!( - "[Claude] Anthropic 响应: {}", - serde_json::to_string(&anthropic_response).unwrap_or_default() + log::info!( + "[Claude] <<< Anthropic 响应 JSON:\n{}", + serde_json::to_string_pretty(&anthropic_response).unwrap_or_default() ); + // 记录使用量 + if let Some(usage) = TokenUsage::from_claude_response(&anthropic_response) { + let model = anthropic_response + .get("model") + .and_then(|m| m.as_str()) + .unwrap_or("unknown"); + + tokio::spawn({ + let state = state.clone(); + let provider_id = provider.id.clone(); + let model = model.to_string(); + async move { + log_usage( + &state, + &provider_id, + "claude", + &model, + usage, + 0, + status.as_u16(), + ) + .await; + } + }); + } + + log::info!("[Claude] ====== 请求结束 ======"); + // 构建响应 let mut builder = axum::response::Response::builder().status(status); @@ -167,15 +406,119 @@ pub async fn handle_messages( } // 透传响应（直连 Anthropic） - log::info!("[Claude] 透传响应"); - let mut builder = axum::response::Response::builder().status(response.status()); + log::info!("[Claude] 透传响应模式"); - for (key, value) in response.headers() { - builder = builder.header(key, value); + // 检查是否流式响应 + let content_type = response + .headers() + .get("content-type") + .and_then(|v| v.to_str().ok()) + .unwrap_or(""); + let is_sse = content_type.contains("text/event-stream"); + + if is_sse { + // 流式透传：使用包装流记录 SSE 事件 + log::info!("[Claude] 流式透传响应 (SSE)"); + let mut builder = axum::response::Response::builder().status(status); + + for (key, value) in response.headers() { + builder = builder.header(key, value); + } + + let stream = response + .bytes_stream() + .map(|chunk| chunk.map_err(|e| std::io::Error::other(e.to_string()))); + let usage_collector = { + let state = state.clone(); + let provider_id = provider.id.clone(); + let model = request_model.clone(); + let status_code = status.as_u16(); + SseUsageCollector::new(move |events| { + if let Some(usage) = TokenUsage::from_claude_stream_events(&events) { + let state = state.clone(); + let provider_id = provider_id.clone(); + let model = model.clone(); + tokio::spawn(async move { + log_usage( + &state, + &provider_id, + "claude", + &model, + usage, + 0, + status_code, + ) + .await; + }); + } else { + log::debug!("[Claude] 流式响应缺少 usage 统计，跳过消费记录"); + } + }) + }; + let logged_stream = + create_logged_passthrough_stream(stream, "Claude", Some(usage_collector)); + + let body = axum::body::Body::from_stream(logged_stream); + log::info!("[Claude] ====== 请求结束 (流式) ======"); + Ok(builder.body(body).unwrap()) + } else { + // 非流式透传：读取完整响应并记录 + let response_headers = response.headers().clone(); + let status = response.status(); + + let body_bytes = response.bytes().await.map_err(|e| { + log::error!("[Claude] 读取透传响应失败: {e}"); + ProxyError::ForwardFailed(format!("Failed to read response body: {e}")) + })?; + + // 记录响应 JSON + if let Ok(json_value) = serde_json::from_slice::(&body_bytes) { + log::info!( + "[Claude] <<< Anthropic 透传响应 JSON:\n{}", + serde_json::to_string_pretty(&json_value).unwrap_or_default() + ); + + // 记录使用量 + if let Some(usage) = TokenUsage::from_claude_response(&json_value) { + let model = json_value + .get("model") + .and_then(|m| m.as_str()) + .unwrap_or("unknown"); + + tokio::spawn({ + let state = state.clone(); + let provider_id = provider.id.clone(); + let model = model.to_string(); + async move { + log_usage( + &state, + &provider_id, + "claude", + &model, + usage, + 0, + status.as_u16(), + ) + .await; + } + }); + } + } else { + log::info!( + "[Claude] <<< 透传响应 (非 JSON): {} bytes", + body_bytes.len() + ); + } + log::info!("[Claude] ====== 请求结束 ======"); + + let mut builder = axum::response::Response::builder().status(status); + for (key, value) in response_headers.iter() { + builder = builder.header(key, value); + } + + let body = axum::body::Body::from(body_bytes); + Ok(builder.body(body).unwrap()) } - - let body = axum::body::Body::from_stream(response.bytes_stream()); - Ok(builder.body(body).unwrap()) } /// 处理 Gemini API 请求（透传，包括查询参数） @@ -186,6 +529,14 @@ pub async fn handle_gemini( Json(body): Json, ) -> Result { let config = state.config.read().await.clone(); + + // 选择目标 Provider + let router = super::router::ProviderRouter::new(state.db.clone()); + let failed_ids = Vec::new(); + let provider = router + .select_provider(&AppType::Gemini, &failed_ids) + .await?; + let forwarder = RequestForwarder::new( state.db.clone(), config.request_timeout, @@ -198,22 +549,124 @@ pub async fn handle_gemini( .path_and_query() .map(|pq| pq.as_str()) .unwrap_or(uri.path()); + let gemini_model = endpoint + .split('/') + .find(|s| s.starts_with("models/")) + .and_then(|s| s.strip_prefix("models/")) + .unwrap_or("unknown") + .to_string(); - log::debug!("Gemini request endpoint (with query): {endpoint}"); + log::info!("[Gemini] 请求端点: {endpoint}"); let response = forwarder .forward_with_retry(&AppType::Gemini, endpoint, body, headers) .await?; - // 透传响应 - let mut builder = axum::response::Response::builder().status(response.status()); + let status = response.status(); + log::info!("[Gemini] 上游响应状态: {status}"); - for (key, value) in response.headers() { - builder = builder.header(key, value); + // 检查是否流式响应 + let content_type = response + .headers() + .get("content-type") + .and_then(|v| v.to_str().ok()) + .unwrap_or(""); + let is_sse = content_type.contains("text/event-stream"); + + if is_sse { + // 流式透传 + log::info!("[Gemini] 流式透传响应 (SSE)"); + let mut builder = axum::response::Response::builder().status(status); + + for (key, value) in response.headers() { + builder = builder.header(key, value); + } + + let stream = response + .bytes_stream() + .map(|chunk| chunk.map_err(|e| std::io::Error::other(e.to_string()))); + let usage_collector = { + let state = state.clone(); + let provider_id = provider.id.clone(); + let model = gemini_model.clone(); + let status_code = status.as_u16(); + SseUsageCollector::new(move |events| { + if let Some(usage) = TokenUsage::from_gemini_stream_chunks(&events) { + let state = state.clone(); + let provider_id = provider_id.clone(); + let model = model.clone(); + tokio::spawn(async move { + log_usage( + &state, + &provider_id, + "gemini", + &model, + usage, + 0, + status_code, + ) + .await; + }); + } else { + log::debug!("[Gemini] 流式响应缺少 usage 统计，跳过消费记录"); + } + }) + }; + let logged_stream = + create_logged_passthrough_stream(stream, "Gemini", Some(usage_collector)); + + let body = axum::body::Body::from_stream(logged_stream); + Ok(builder.body(body).unwrap()) + } else { + // 非流式透传 + let response_headers = response.headers().clone(); + let status = response.status(); + + let body_bytes = response.bytes().await.map_err(|e| { + log::error!("[Gemini] 读取响应失败: {e}"); + ProxyError::ForwardFailed(format!("Failed to read response body: {e}")) + })?; + + // 记录响应 JSON + if let Ok(json_value) = serde_json::from_slice::(&body_bytes) { + log::info!( + "[Gemini] <<< 响应 JSON:\n{}", + serde_json::to_string_pretty(&json_value).unwrap_or_default() + ); + + // 记录使用量 + if let Some(usage) = TokenUsage::from_gemini_response(&json_value) { + let model = gemini_model.clone(); + tokio::spawn({ + let state = state.clone(); + let provider_id = provider.id.clone(); + async move { + log_usage( + &state, + &provider_id, + "gemini", + &model, + usage, + 0, + status.as_u16(), + ) + .await; + } + }); + } + } else { + log::info!("[Gemini] <<< 响应 (非 JSON): {} bytes", body_bytes.len()); + } + log::info!("[Gemini] ====== 请求结束 ======"); + + let mut builder = axum::response::Response::builder().status(status); + for (key, value) in response_headers.iter() { + builder = builder.header(key, value); + } + + let body = axum::body::Body::from(body_bytes); + Ok(builder.body(body).unwrap()) } - - let body = axum::body::Body::from_stream(response.bytes_stream()); - Ok(builder.body(body).unwrap()) } /// 处理 /v1/responses 请求（OpenAI Responses API - Codex CLI 透传） @@ -223,6 +676,17 @@ pub async fn handle_responses( Json(body): Json, ) -> Result { let config = state.config.read().await.clone(); + let request_model = body + .get("model") + .and_then(|m| m.as_str()) + .unwrap_or("unknown") + .to_string(); + + // 选择目标 Provider + let router = super::router::ProviderRouter::new(state.db.clone()); + let failed_ids = Vec::new(); + let provider = router.select_provider(&AppType::Codex, &failed_ids).await?; + let forwarder = RequestForwarder::new( state.db.clone(), config.request_timeout, @@ -234,13 +698,106 @@ pub async fn handle_responses( .forward_with_retry(&AppType::Codex, "/v1/responses", body, headers) .await?; - // 透传响应（包括流式和非流式） - let mut builder = axum::response::Response::builder().status(response.status()); + let status = response.status(); + log::info!("[Codex] 上游响应状态: {status}"); - for (key, value) in response.headers() { - builder = builder.header(key, value); + // 检查是否流式响应 + let content_type = response + .headers() + .get("content-type") + .and_then(|v| v.to_str().ok()) + .unwrap_or(""); + let is_sse = content_type.contains("text/event-stream"); + + if is_sse { + // 流式透传 + log::info!("[Codex] 流式透传响应 (SSE)"); + let mut builder = axum::response::Response::builder().status(status); + + for (key, value) in response.headers() { + builder = builder.header(key, value); + } + + let stream = response + .bytes_stream() + .map(|chunk| chunk.map_err(|e| std::io::Error::other(e.to_string()))); + let usage_collector = { + let state = state.clone(); + let provider_id = provider.id.clone(); + let model = request_model.clone(); + let status_code = status.as_u16(); + SseUsageCollector::new(move |events| { + if let Some(usage) = TokenUsage::from_codex_stream_events(&events) { + let state = state.clone(); + let provider_id = provider_id.clone(); + let model = model.clone(); + tokio::spawn(async move { + log_usage(&state, &provider_id, "codex", &model, usage, 0, status_code) + .await; + }); + } else { + log::debug!("[Codex] 流式响应缺少 usage 统计，跳过消费记录"); + } + }) + }; + let logged_stream = + create_logged_passthrough_stream(stream, "Codex", Some(usage_collector)); + + let body = axum::body::Body::from_stream(logged_stream); + Ok(builder.body(body).unwrap()) + } else { + // 非流式透传 + let response_headers = response.headers().clone(); + let status = response.status(); + + let body_bytes = response.bytes().await.map_err(|e| { + log::error!("[Codex] 读取响应失败: {e}"); + ProxyError::ForwardFailed(format!("Failed to read response body: {e}")) + })?; + + // 记录响应 JSON + if let Ok(json_value) = serde_json::from_slice::(&body_bytes) { + log::info!( + "[Codex] <<< 响应 JSON:\n{}", + serde_json::to_string_pretty(&json_value).unwrap_or_default() + ); + + // 记录使用量 + if let Some(usage) = TokenUsage::from_codex_response(&json_value) { + let model = json_value + .get("model") + .and_then(|m| m.as_str()) + .unwrap_or("unknown"); + + tokio::spawn({ + let state = state.clone(); + let provider_id = provider.id.clone(); + let model = model.to_string(); + async move { + log_usage( + &state, + &provider_id, + "codex", + &model, + usage, + 0, + status.as_u16(), + ) + .await; + } + }); + } + } else { + log::info!("[Codex] <<< 响应 (非 JSON): {} bytes", body_bytes.len()); + } + log::info!("[Codex] ====== 请求结束 ======"); + + let mut builder = axum::response::Response::builder().status(status); + for (key, value) in response_headers.iter() { + builder = builder.header(key, value); + } + + let body = axum::body::Body::from(body_bytes); + Ok(builder.body(body).unwrap()) } - - let body = axum::body::Body::from_stream(response.bytes_stream()); - Ok(builder.body(body).unwrap()) } diff --git a/src-tauri/src/proxy/mod.rs b/src-tauri/src/proxy/mod.rs index 8149aba85..b61756dc0 100644 --- a/src-tauri/src/proxy/mod.rs +++ b/src-tauri/src/proxy/mod.rs @@ -10,6 +10,7 @@ pub mod providers; mod router; pub(crate) mod server; pub(crate) mod types; +pub mod usage; // 公开导出给外部使用（commands, services等模块需要） #[allow(unused_imports)] diff --git a/src-tauri/src/proxy/providers/streaming.rs b/src-tauri/src/proxy/providers/streaming.rs index 5e6f2bf52..48ecefa6c 100644 --- a/src-tauri/src/proxy/providers/streaming.rs +++ b/src-tauri/src/proxy/providers/streaming.rs @@ -71,6 +71,8 @@ pub fn create_anthropic_sse_stream( let mut current_block_type: Option = None; let mut tool_call_id = None; + log::info!("[Claude/OpenRouter] ====== 开始流式响应转换 ======"); + tokio::pin!(stream); while let Some(chunk) = stream.next().await { @@ -90,14 +92,26 @@ pub fn create_anthropic_sse_stream( for l in line.lines() { if let Some(data) = l.strip_prefix("data: ") { if data.trim() == "[DONE]" { + log::info!("[Claude/OpenRouter] <<< OpenAI SSE: [DONE]"); let event = json!({"type": "message_stop"}); let sse_data = format!("event: message_stop\ndata: {}\n\n", serde_json::to_string(&event).unwrap_or_default()); + log::info!("[Claude/OpenRouter] >>> Anthropic SSE: message_stop"); yield Ok(Bytes::from(sse_data)); continue; } if let Ok(chunk) = serde_json::from_str::(data) { + // 记录原始 OpenAI 事件（格式化显示） + if let Ok(json_value) = serde_json::from_str::(data) { + log::info!( + "[Claude/OpenRouter] <<< OpenAI SSE 事件:\n{}", + serde_json::to_string_pretty(&json_value).unwrap_or_else(|_| data.to_string()) + ); + } else { + log::info!("[Claude/OpenRouter] <<< OpenAI SSE 数据: {data}"); + } + if message_id.is_none() { message_id = Some(chunk.id.clone()); }