//! 请求处理器 //! //! 处理各种API端点的HTTP请求 use super::{ forwarder::RequestForwarder, providers::{get_adapter, transform, ProviderType}, server::ProxyState, session::ProxySession, 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; /// 记录请求使用量(带 ProxySession 支持) #[allow(dead_code, clippy::too_many_arguments)] async fn log_usage_with_session( state: &ProxyState, session: &ProxySession, provider_id: &str, app_type: &str, usage: TokenUsage, latency_ms: u64, first_token_ms: Option, status_code: u16, provider_type: Option<&ProviderType>, ) { 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 model = session .model .clone() .unwrap_or_else(|| "unknown".to_string()); let provider_type_str = provider_type.map(|pt| pt.as_str().to_string()); if let Err(e) = logger.log_with_calculation( session.session_id.clone(), provider_id.to_string(), app_type.to_string(), model, usage, multiplier, latency_ms, first_token_ms, status_code, Some(session.session_id.clone()), provider_type_str, session.is_streaming, ) { log::warn!("记录使用量失败: {e}"); } } /// 记录请求使用量(兼容旧接口) #[allow(clippy::too_many_arguments)] async fn log_usage( state: &ProxyState, provider_id: &str, app_type: &str, model: &str, usage: TokenUsage, latency_ms: u64, first_token_ms: Option, is_streaming: bool, 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, first_token_ms, status_code, None, None, // provider_type is_streaming, ) { log::warn!("记录使用量失败: {e}"); } } type UsageCallbackWithTiming = Arc, Option) + Send + Sync + 'static>; #[derive(Clone)] struct SseUsageCollector { inner: Arc, } struct SseUsageCollectorInner { events: Mutex>, first_event_time: Mutex>, start_time: std::time::Instant, on_complete: UsageCallbackWithTiming, finished: AtomicBool, } impl SseUsageCollector { fn new( start_time: std::time::Instant, callback: impl Fn(Vec, Option) + Send + Sync + 'static, ) -> Self { let on_complete: UsageCallbackWithTiming = Arc::new(callback); Self { inner: Arc::new(SseUsageCollectorInner { events: Mutex::new(Vec::new()), first_event_time: Mutex::new(None), start_time, on_complete, finished: AtomicBool::new(false), }), } } async fn push(&self, event: Value) { // 记录首个事件时间 { let mut first_time = self.inner.first_event_time.lock().await; if first_time.is_none() { *first_time = Some(std::time::Instant::now()); } } 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) }; let first_token_ms = { let first_time = self.inner.first_event_time.lock().await; first_time.map(|t| (t - self.inner.start_time).as_millis() as u64) }; (self.inner.on_complete)(events, first_token_ms); } } /// 创建带日志记录的透传流 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) { ( StatusCode::OK, Json(json!({ "status": "healthy", "timestamp": chrono::Utc::now().to_rfc3339(), })), ) } /// 获取服务状态 pub async fn get_status(State(state): State) -> Result, ProxyError> { let status = state.status.read().await.clone(); Ok(Json(status)) } /// 处理 /v1/messages 请求(Claude API) pub async fn handle_messages( State(state): State, headers: axum::http::HeaderMap, Json(body): Json, ) -> Result { let start_time = std::time::Instant::now(); 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::Claude, &failed_ids) .await?; // 检查是否需要转换(OpenRouter) let adapter = get_adapter(&AppType::Claude); let needs_transform = adapter.needs_transform(&provider); // 检查是否是流式请求 let is_stream = body .get("stream") .and_then(|s| s.as_bool()) .unwrap_or(false); log::info!( "[Claude] Provider: {}, needs_transform: {}, is_stream: {}", provider.name, needs_transform, is_stream ); let forwarder = RequestForwarder::new( state.db.clone(), config.request_timeout, config.max_retries, state.status.clone(), ); let response = forwarder .forward_with_retry(&AppType::Claude, "/v1/messages", body, headers) .await?; let status = response.status(); log::info!("[Claude] 上游响应状态: {status}"); // 如果需要转换 if needs_transform { if is_stream { // 流式响应转换 log::info!("[Claude] 开始流式响应转换 (OpenAI SSE → Anthropic SSE)"); 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(); let start_time_clone = start_time; SseUsageCollector::new(start_time, move |events, first_token_ms| { if let Some(usage) = TokenUsage::from_claude_stream_events(&events) { let latency_ms = start_time_clone.elapsed().as_millis() as u64; 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, latency_ms, first_token_ms, true, // is_streaming 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", axum::http::HeaderValue::from_static("text/event-stream"), ); headers.insert( "Cache-Control", axum::http::HeaderValue::from_static("no-cache"), ); headers.insert( "Connection", axum::http::HeaderValue::from_static("keep-alive"), ); let body = axum::body::Body::from_stream(logged_stream); log::info!("[Claude] ====== 请求结束 (流式转换) ======"); return Ok((headers, body).into_response()); } else { // 非流式响应转换 log::info!("[Claude] 开始转换响应 (OpenAI → Anthropic)"); let response_headers = response.headers().clone(); // 读取响应体 let body_bytes = response.bytes().await.map_err(|e| { log::error!("[Claude] 读取响应体失败: {e}"); ProxyError::ForwardFailed(format!("Failed to read response body: {e}")) })?; let body_str = String::from_utf8_lossy(&body_bytes); log::info!("[Claude] OpenAI 响应长度: {} bytes", body_bytes.len()); log::debug!("[Claude] OpenAI 原始响应: {body_str}"); // 解析并转换 let openai_response: Value = serde_json::from_slice(&body_bytes).map_err(|e| { log::error!("[Claude] 解析 OpenAI 响应失败: {e}, body: {body_str}"); ProxyError::TransformError(format!("Failed to parse OpenAI response: {e}")) })?; 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| { log::error!("[Claude] 转换响应失败: {e}"); e })?; log::info!("[Claude] 转换响应成功"); 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"); let latency_ms = start_time.elapsed().as_millis() as u64; 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, latency_ms, None, false, status.as_u16(), ) .await; } }); } log::info!("[Claude] ====== 请求结束 ======"); // 构建响应 let mut builder = axum::response::Response::builder().status(status); // 复制响应头(排除 content-length,因为内容已改变) for (key, value) in response_headers.iter() { if key.as_str().to_lowercase() != "content-length" && key.as_str().to_lowercase() != "transfer-encoding" { builder = builder.header(key, value); } } builder = builder.header("content-type", "application/json"); let response_body = serde_json::to_vec(&anthropic_response).map_err(|e| { log::error!("[Claude] 序列化响应失败: {e}"); ProxyError::TransformError(format!("Failed to serialize response: {e}")) })?; log::info!( "[Claude] 返回转换后的响应, 长度: {} bytes", response_body.len() ); let body = axum::body::Body::from(response_body); return Ok(builder.body(body).unwrap()); } } // 透传响应(直连 Anthropic) log::info!("[Claude] 透传响应模式"); // 检查是否流式响应 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(); let start_time_clone = start_time; SseUsageCollector::new(start_time, move |events, first_token_ms| { if let Some(usage) = TokenUsage::from_claude_stream_events(&events) { let latency_ms = start_time_clone.elapsed().as_millis() as u64; 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, latency_ms, first_token_ms, true, 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"); let latency_ms = start_time.elapsed().as_millis() as u64; 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, latency_ms, None, false, 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()) } } /// 处理 Gemini API 请求(透传,包括查询参数) pub async fn handle_gemini( State(state): State, uri: axum::http::Uri, headers: axum::http::HeaderMap, Json(body): Json, ) -> Result { let start_time = std::time::Instant::now(); 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, config.max_retries, state.status.clone(), ); // 提取完整的路径和查询参数 let endpoint = uri .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/")) .map(|s| s.split(':').next().unwrap_or(s)) .unwrap_or("unknown") .to_string(); log::info!("[Gemini] 请求端点: {endpoint}"); let response = forwarder .forward_with_retry(&AppType::Gemini, endpoint, body, headers) .await?; let status = response.status(); log::info!("[Gemini] 上游响应状态: {status}"); // 检查是否流式响应 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 fallback_model = gemini_model.clone(); let status_code = status.as_u16(); let start_time_clone = start_time; SseUsageCollector::new(start_time, move |events, first_token_ms| { if let Some(usage) = TokenUsage::from_gemini_stream_chunks(&events) { // 优先使用响应中的实际模型名称,否则使用从 URI 提取的模型名称 let model = usage .model .clone() .unwrap_or_else(|| fallback_model.clone()); let latency_ms = start_time_clone.elapsed().as_millis() as u64; let state = state.clone(); let provider_id = provider_id.clone(); tokio::spawn(async move { log_usage( &state, &provider_id, "gemini", &model, usage, latency_ms, first_token_ms, true, 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) { // 优先使用响应中的实际模型名称,否则使用从 URI 提取的模型名称 let model = usage.model.clone().unwrap_or_else(|| gemini_model.clone()); let latency_ms = start_time.elapsed().as_millis() as u64; tokio::spawn({ let state = state.clone(); let provider_id = provider.id.clone(); async move { log_usage( &state, &provider_id, "gemini", &model, usage, latency_ms, None, false, 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()) } } /// 处理 /v1/responses 请求(OpenAI Responses API - Codex CLI 透传) pub async fn handle_responses( State(state): State, headers: axum::http::HeaderMap, Json(body): Json, ) -> Result { let start_time = std::time::Instant::now(); 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, config.max_retries, state.status.clone(), ); let response = forwarder .forward_with_retry(&AppType::Codex, "/v1/responses", body, headers) .await?; let status = response.status(); log::info!("[Codex] 上游响应状态: {status}"); // 检查是否流式响应 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 request_model = request_model.clone(); let status_code = status.as_u16(); let start_time_clone = start_time; SseUsageCollector::new(start_time, move |events, first_token_ms| { if let Some(usage) = TokenUsage::from_codex_stream_events(&events) { // 尝试从事件中提取模型,回退到请求模型 let model = events .iter() .find_map(|e| { if e.get("type")?.as_str()? == "response.completed" { e.get("response")?.get("model")?.as_str() } else { None } }) .unwrap_or(&request_model) .to_string(); let latency_ms = start_time_clone.elapsed().as_millis() as u64; let state = state.clone(); let provider_id = provider_id.clone(); tokio::spawn(async move { log_usage( &state, &provider_id, "codex", &model, usage, latency_ms, first_token_ms, true, 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"); let latency_ms = start_time.elapsed().as_millis() as u64; log::info!( "[Codex] 解析到 usage: input={}, output={}", usage.input_tokens, usage.output_tokens ); 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, latency_ms, None, false, status.as_u16(), ) .await; } }); } else { log::warn!("[Codex] 未能解析 usage 信息,跳过记录"); } } 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()) } } /// 处理 /v1/chat/completions 请求(OpenAI Chat Completions API - Codex CLI) pub async fn handle_chat_completions( State(state): State, headers: axum::http::HeaderMap, Json(body): Json, ) -> Result { let start_time = std::time::Instant::now(); log::info!("[Codex] ====== /v1/chat/completions 请求开始 ======"); let config = state.config.read().await.clone(); let request_model = body .get("model") .and_then(|m| m.as_str()) .unwrap_or("unknown") .to_string(); let is_stream = body .get("stream") .and_then(|v| v.as_bool()) .unwrap_or(false); log::info!("[Codex] 请求模型: {request_model}, 流式: {is_stream}"); // 选择目标 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?; log::info!("[Codex] 选择 Provider: {}", provider.id); let forwarder = RequestForwarder::new( state.db.clone(), config.request_timeout, config.max_retries, state.status.clone(), ); let response = forwarder .forward_with_retry(&AppType::Codex, "/v1/chat/completions", body, headers) .await?; let status = response.status(); log::info!("[Codex] 上游响应状态: {status}"); // 检查是否流式响应 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 request_model = request_model.clone(); let status_code = status.as_u16(); let start_time_clone = start_time; SseUsageCollector::new(start_time, move |events, first_token_ms| { if let Some(usage) = TokenUsage::from_openai_stream_events(&events) { let model = events .iter() .find_map(|e| e.get("model")?.as_str()) .unwrap_or(&request_model) .to_string(); let latency_ms = start_time_clone.elapsed().as_millis() as u64; let state = state.clone(); let provider_id = provider_id.clone(); tokio::spawn(async move { log_usage( &state, &provider_id, "codex", &model, usage, latency_ms, first_token_ms, true, 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() ); // 记录使用量 (OpenAI 格式: prompt_tokens, completion_tokens) if let Some(usage) = TokenUsage::from_openai_response(&json_value) { let model = json_value .get("model") .and_then(|m| m.as_str()) .unwrap_or("unknown"); let latency_ms = start_time.elapsed().as_millis() as u64; log::info!( "[Codex] 解析到 usage: input={}, output={}", usage.input_tokens, usage.output_tokens ); 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, latency_ms, None, false, status.as_u16(), ) .await; } }); } else { log::warn!("[Codex] 未能解析 usage 信息,跳过记录"); } } 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()) } }