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:
YoVinchen
2025-12-03 00:18:42 +08:00
parent 05973bf959
commit 6423e2e8ab
4 changed files with 607 additions and 28 deletions
+9 -2
View File
@@ -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
View File
@@ -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())
}
+1
View File
@@ -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());
}