mirror of
https://github.com/farion1231/cc-switch.git
synced 2026-07-27 16:26:16 +08:00
feat(proxy): integrate usage logging into request handlers
Add usage logging to forwarder and streaming handlers. Track token usage and costs for each proxy request.
This commit is contained in:
@@ -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()
|
||||
);
|
||||
|
||||
+583
-26
@@ -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<dyn Fn(Vec<Value>) + Send + Sync + 'static>;
|
||||
|
||||
#[derive(Clone)]
|
||||
struct SseUsageCollector {
|
||||
inner: Arc<SseUsageCollectorInner>,
|
||||
}
|
||||
|
||||
struct SseUsageCollectorInner {
|
||||
events: Mutex<Vec<Value>>,
|
||||
on_complete: UsageCallback,
|
||||
finished: AtomicBool,
|
||||
}
|
||||
|
||||
impl SseUsageCollector {
|
||||
fn new(callback: impl Fn(Vec<Value>) + 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<Item = Result<Bytes, std::io::Error>> + Send + 'static,
|
||||
tag: &'static str,
|
||||
usage_collector: Option<SseUsageCollector>,
|
||||
) -> impl Stream<Item = Result<Bytes, std::io::Error>> + 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::<Value>(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<Value>) {
|
||||
@@ -37,6 +204,11 @@ pub async fn handle_messages(
|
||||
Json(body): Json<Value>,
|
||||
) -> Result<axum::response::Response, ProxyError> {
|
||||
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::<Value>(&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<Value>,
|
||||
) -> Result<axum::response::Response, ProxyError> {
|
||||
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::<Value>(&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<Value>,
|
||||
) -> Result<axum::response::Response, ProxyError> {
|
||||
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::<Value>(&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())
|
||||
}
|
||||
|
||||
@@ -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)]
|
||||
|
||||
@@ -71,6 +71,8 @@ pub fn create_anthropic_sse_stream(
|
||||
let mut current_block_type: Option<String> = 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::<OpenAIStreamChunk>(data) {
|
||||
// 记录原始 OpenAI 事件(格式化显示)
|
||||
if let Ok(json_value) = serde_json::from_str::<serde_json::Value>(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());
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user