fix(saas): SSE 用量统计一致性修复 — 回写 usage_records + 消除 relay_requests 双重计数
Some checks failed
CI / Lint & TypeCheck (push) Has been cancelled
CI / Unit Tests (push) Has been cancelled
CI / Build Frontend (push) Has been cancelled
CI / Rust Check (push) Has been cancelled
CI / Security Scan (push) Has been cancelled
CI / E2E Tests (push) Has been cancelled

- service.rs: SSE 流结束后回写 usage_records 真实 token (status=success)
- service.rs: spawned task 中调用 increment_usage 统一递增 tokens + relay_requests
- handlers.rs: 移除 SSE 路径的 increment_dimension("relay_requests") 消除双重计数
- 从 request_body 提取 model_id 用于 usage_records 精准归因
This commit is contained in:
iven
2026-04-15 01:40:27 +08:00
parent 02a4ba5e75
commit 76cdfd0c00
2 changed files with 18 additions and 11 deletions

View File

@@ -333,14 +333,8 @@ pub async fn chat_completions(
} }
} }
// SSE: relay_requests 实时递增tokens 由 AggregateUsageWorker 对账修正)
if let Err(e) = crate::billing::service::increment_dimension(
&state.db, &account_id_usage, "relay_requests",
).await {
tracing::warn!("Failed to increment billing relay_requests for {}: {}", account_id_usage, e);
}
// SSE 流已返回,递减队列计数器(流式任务开始处理) // SSE 流已返回,递减队列计数器(流式任务开始处理)
// 注意: relay_requests 和 tokens 统一由 execute_relay spawned task 中的 increment_usage 递增
state.cache.relay_dequeue(&account_id_usage); state.cache.relay_dequeue(&account_id_usage);
let response = axum::response::Response::builder() let response = axum::response::Response::builder()

View File

@@ -333,6 +333,12 @@ pub async fn execute_relay(
let task_id_clone = task_id.to_string(); let task_id_clone = task_id.to_string();
let key_id_for_spawn = key_id.clone(); let key_id_for_spawn = key_id.clone();
let account_id_clone = account_id.to_string(); let account_id_clone = account_id.to_string();
let provider_id_clone = provider_id.to_string();
// 从 request_body 提取 model_id 用于 usage_records 归因
let model_id_clone = serde_json::from_str::<serde_json::Value>(request_body)
.ok()
.and_then(|v| v.get("model").and_then(|m| m.as_str()).map(String::from))
.unwrap_or_default();
// Bounded channel for backpressure: 128 chunks (~128KB) buffer. // Bounded channel for backpressure: 128 chunks (~128KB) buffer.
// If the client reads slowly, the upstream is signaled via // If the client reads slowly, the upstream is signaled via
@@ -434,16 +440,23 @@ pub async fn execute_relay(
let input_opt = if input > 0 { Some(input) } else { None }; let input_opt = if input > 0 { Some(input) } else { None };
let output_opt = if output > 0 { Some(output) } else { None }; let output_opt = if output > 0 { Some(output) } else { None };
// Record task status + billing usage + key usage // Record task status + billing usage + key usage + usage_records
let db_op = async { let db_op = async {
if let Err(e) = update_task_status(&db_clone, &task_id_clone, "completed", input_opt, output_opt, None).await { if let Err(e) = update_task_status(&db_clone, &task_id_clone, "completed", input_opt, output_opt, None).await {
tracing::warn!("Failed to update task status after SSE stream: {}", e); tracing::warn!("Failed to update task status after SSE stream: {}", e);
} }
// P2-9 修复: SSE 路径也更新 billing_usage_quotas // SSE 路径回写 usage_records + billing 配额
if input > 0 || output > 0 { if input > 0 || output > 0 {
// 回写 usage_records 真实 token补全 handlers.rs 中 token=0 的占位记录)
if let Err(e) = crate::model_config::service::record_usage(
&db_clone, &account_id_clone, &provider_id_clone, &model_id_clone,
input, output, None, "success", None,
).await {
tracing::warn!("Failed to record SSE usage for task {}: {}", task_id_clone, e);
}
// 更新 billing_usage_quotastokens + relay_requests 同步递增)
if let Err(e) = crate::billing::service::increment_usage( if let Err(e) = crate::billing::service::increment_usage(
&db_clone, &account_id_clone, &db_clone, &account_id_clone, input, output,
input, output,
).await { ).await {
tracing::warn!("Failed to increment billing usage for SSE task {}: {}", task_id_clone, e); tracing::warn!("Failed to increment billing usage for SSE task {}: {}", task_id_clone, e);
} }