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
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:
@@ -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()
|
||||||
|
|||||||
@@ -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_quotas(tokens + 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);
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user