From 4a413f5bb7991479c2a7f4f90d55dadc686ff10e Mon Sep 17 00:00:00 2001 From: Zyy Date: Fri, 11 Sep 2026 05:13:06 +0800 Subject: [PATCH] fix(budgets): remove manual charge reconciliation Settle only observed USD costs and treat missing costs as zero without blocking client keys. Preserve known retry costs and idempotent settlement. Remove reconciliation controls, endpoints, pending records, and obsolete schema fields; migrate legacy unknown charges to zero without changing recorded usage. Display daily and weekly amounts with two decimal places while retaining full precision for accounting and tooltips. Distinguish missing WebSocket closing handshakes from TCP resets in connection diagnostics. Validation: workspace tests with PostgreSQL/Redis, targeted regressions, strict Clippy, rustfmt, frontend lint/build, and migration checks passed. Co-authored-by: Codex --- .../gateway-admin/src/model/client_keys.rs | 16 -- .../crates/gateway-admin/src/ports/store.rs | 23 -- .../gateway-admin/src/use_case/client_keys.rs | 41 ---- .../gateway-api/src/admin/client_keys.rs | 70 ------ .../gateway-api/tests/admin/client_keys.rs | 40 +-- .../crates/gateway-core/src/engine/budget.rs | 21 +- .../gateway-core/src/engine/coordinator.rs | 35 ++- .../gateway-core/src/engine/execution.rs | 15 +- .../crates/gateway-core/src/metering/mod.rs | 2 + .../gateway-core/tests/engine/coordinator.rs | 91 ++++++- .../gateway-core/tests/engine/execution.rs | 8 +- .../src/postgres/client_budgets.rs | 207 +++------------- .../gateway-store/src/postgres/client_keys.rs | 16 -- .../tests/postgres/client_budgets.rs | 228 +++++++----------- .../openai/src/transport/websocket/pump.rs | 3 +- .../openai/tests/transport/latency.rs | 2 +- backend/migrations/.frozen-sha256 | 1 + ...04_remove_client_charge_reconciliation.sql | 17 ++ docs/api.md | 13 +- docs/architecture.md | 18 +- frontend/src/api/modules/api-keys.ts | 16 -- .../api-keys/components/ApiKeyBudgetCell.vue | 12 +- .../components/ApiKeyReconcileModal.vue | 124 ---------- frontend/src/views/api-keys/index.vue | 13 +- 24 files changed, 275 insertions(+), 757 deletions(-) create mode 100644 backend/migrations/0004_remove_client_charge_reconciliation.sql delete mode 100644 frontend/src/views/api-keys/components/ApiKeyReconcileModal.vue diff --git a/backend/crates/gateway-admin/src/model/client_keys.rs b/backend/crates/gateway-admin/src/model/client_keys.rs index 59e5eb131..6532773df 100644 --- a/backend/crates/gateway-admin/src/model/client_keys.rs +++ b/backend/crates/gateway-admin/src/model/client_keys.rs @@ -223,19 +223,3 @@ pub struct ClientKeyMutation { pub record: Option, pub id: ClientApiKeyId, } - -#[derive(Debug, Clone)] -pub struct UnresolvedClientCharge { - pub request_id: String, - pub started_at: DateTime, - pub completed_at: Option>, - pub state: String, -} - -#[derive(Debug, Clone)] -pub struct ReconcileClientCharge { - pub key_id: ClientApiKeyId, - pub request_id: String, - pub amount_usd: gateway_core::metering::Decimal, - pub reason: String, -} diff --git a/backend/crates/gateway-admin/src/ports/store.rs b/backend/crates/gateway-admin/src/ports/store.rs index 9cc4cc02a..181cfe95f 100644 --- a/backend/crates/gateway-admin/src/ports/store.rs +++ b/backend/crates/gateway-admin/src/ports/store.rs @@ -220,29 +220,6 @@ pub trait AuthStore: Send + Sync { /// Client API Key 管理写入。 #[async_trait] pub trait ClientKeyStore: Send + Sync { - async fn unresolved_charges( - &self, - _id: &gateway_core::policy::ClientApiKeyId, - ) -> AdminStoreResult> { - Err(AdminStoreError::new( - AdminStoreErrorKind::Unavailable, - "client budget", - "unsupported", - )) - } - - async fn reconcile_charge( - &self, - _command: crate::model::client_keys::ReconcileClientCharge, - _context: &MutationContext, - ) -> AdminStoreResult<()> { - Err(AdminStoreError::new( - AdminStoreErrorKind::Unavailable, - "client budget", - "unsupported", - )) - } - async fn list_client_keys(&self, query: ClientKeyListQuery) -> AdminStoreResult; async fn reveal_client_key( diff --git a/backend/crates/gateway-admin/src/use_case/client_keys.rs b/backend/crates/gateway-admin/src/use_case/client_keys.rs index 5a8f5b3fe..4a8c4bbfe 100644 --- a/backend/crates/gateway-admin/src/use_case/client_keys.rs +++ b/backend/crates/gateway-admin/src/use_case/client_keys.rs @@ -26,19 +26,6 @@ use super::{map_store_error, publish_committed}; /// API 消费的 Client Key 管理服务。 #[async_trait] pub trait ClientKeyService: Send + Sync { - async fn unresolved_charges( - &self, - _id: &ClientApiKeyId, - ) -> Result, AdminError> { - Err(AdminError::unavailable("client budget unavailable")) - } - async fn reconcile_charge( - &self, - _context: &MutationContext, - _command: crate::model::client_keys::ReconcileClientCharge, - ) -> Result<(), AdminError> { - Err(AdminError::unavailable("client budget unavailable")) - } async fn list(&self, query: ClientKeyListQuery) -> Result; async fn reveal(&self, id: &ClientApiKeyId) -> Result; async fn create( @@ -77,34 +64,6 @@ impl DefaultClientKeyService { #[async_trait] impl ClientKeyService for DefaultClientKeyService { - async fn unresolved_charges( - &self, - id: &ClientApiKeyId, - ) -> Result, AdminError> { - self.store - .unresolved_charges(id) - .await - .map_err(|error| map_store_error(error, "client budget")) - } - - async fn reconcile_charge( - &self, - context: &MutationContext, - command: crate::model::client_keys::ReconcileClientCharge, - ) -> Result<(), AdminError> { - if command.request_id.is_empty() - || command.request_id.len() > 128 - || command.reason.trim().is_empty() - || command.reason.len() > 1024 - || command.reason.chars().any(char::is_control) - { - return Err(AdminError::invalid("invalid charge reconciliation")); - } - self.store - .reconcile_charge(command, context) - .await - .map_err(|error| map_store_error(error, "client budget")) - } async fn list(&self, query: ClientKeyListQuery) -> Result { validate_cursor(&query)?; self.store diff --git a/backend/crates/gateway-api/src/admin/client_keys.rs b/backend/crates/gateway-api/src/admin/client_keys.rs index 9e0952949..8370d0a51 100644 --- a/backend/crates/gateway-api/src/admin/client_keys.rs +++ b/backend/crates/gateway-api/src/admin/client_keys.rs @@ -275,7 +275,6 @@ pub struct ClientKeyView { weekly_used_usd: String, daily_resets_at: Option>, weekly_resets_at: Option>, - unresolved_requests: u64, created_at: DateTime, updated_at: DateTime, last_used_at: Option>, @@ -327,7 +326,6 @@ impl From for ClientKeyView { weekly_used_usd: record.budget.weekly_used_usd.canonical(), daily_resets_at: record.budget.daily_resets_at.map(DateTime::from), weekly_resets_at: record.budget.weekly_resets_at.map(DateTime::from), - unresolved_requests: record.budget.unresolved_requests, created_at: record.created_at, updated_at: record.updated_at, last_used_at: record.last_used_at, @@ -679,14 +677,6 @@ where { Router::new() .route("/api/admin/client-keys", get(list_client_keys::)) - .route( - "/api/admin/client-keys/unresolved-charges", - get(unresolved_client_charges::), - ) - .route( - "/api/admin/client-keys/reconcile-charge", - post(reconcile_client_charge::), - ) .route( "/api/admin/client-keys/create", post(create_client_key::), @@ -710,66 +700,6 @@ where ) } -#[derive(Deserialize)] -#[serde(rename_all = "camelCase", deny_unknown_fields)] -struct ReconcileClientChargeRequest { - id: String, - request_id: String, - amount_usd: String, - reason: String, -} - -async fn unresolved_client_charges( - _auth: AdminAuth, - State(state): State, - AdminQuery(query): AdminQuery, -) -> Result -where - S: AdminSessionState + Send + Sync, -{ - let id = query.into_domain_id().map_err(map_wire_error)?; - let items = state.admin_services().client_keys().unresolved_charges(&id).await.map_err(map_service_error)? - .into_iter().map(|event| serde_json::json!({ - "requestId": event.request_id, "startedAt": event.started_at.to_rfc3339(), - "completedAt": event.completed_at.map(|time| time.to_rfc3339()), "state": event.state, - })).collect::>(); - Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(serde_json::json!({"items": items})), - )) -} - -async fn reconcile_client_charge( - auth: AdminAuth, - State(state): State, - AdminJson(payload): AdminJson, -) -> Result -where - S: AdminSessionState + Send + Sync, -{ - let key_id = client_key_id(payload.id, "id").map_err(map_wire_error)?; - let amount_usd = payload - .amount_usd - .parse::() - .map_err(|_| map_wire_error(WireValidationError::new("amountUsd")))?; - let command = gateway_admin::model::client_keys::ReconcileClientCharge { - key_id, - request_id: payload.request_id, - amount_usd, - reason: payload.reason, - }; - state - .admin_services() - .client_keys() - .reconcile_charge(&auth.context().mutation_context(), command) - .await - .map_err(map_service_error)?; - Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(serde_json::json!({"reconciled": true})), - )) -} - async fn list_client_keys( _auth: AdminAuth, State(state): State, diff --git a/backend/crates/gateway-api/tests/admin/client_keys.rs b/backend/crates/gateway-api/tests/admin/client_keys.rs index f53cf6014..50e02e511 100644 --- a/backend/crates/gateway-api/tests/admin/client_keys.rs +++ b/backend/crates/gateway-api/tests/admin/client_keys.rs @@ -50,44 +50,6 @@ fn budget_inputs_preserve_decimal_precision_and_omitted_updates() { ); } -#[tokio::test] -async fn budget_reconciliation_routes_require_an_admin_session() { - use axum::{ - body::Body, - http::{Request, StatusCode}, - }; - use tower::ServiceExt as _; - let fixture = AdminTestFixture::new().await; - let router = client_keys::router::().with_state(fixture.state()); - for (method, path, body) in [ - ( - "GET", - "/api/admin/client-keys/unresolved-charges?id=key_1", - "", - ), - ( - "POST", - "/api/admin/client-keys/reconcile-charge", - r#"{"id":"key_1","requestId":"req_charge","amountUsd":"0","reason":"verified"}"#, - ), - ] { - let response = router - .clone() - .oneshot( - Request::builder() - .method(method) - .uri(path) - .header("content-type", "application/json") - .header("x-request-id", "req_auth_budget") - .body(Body::from(body)) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(response.status(), StatusCode::UNAUTHORIZED); - } -} - #[test] fn client_key_queries_should_reject_unknown_zero_and_oversized_fields() { let unknown = serde_json::from_value::(json!({ "other": true })); @@ -320,7 +282,7 @@ fn client_key_responses_should_keep_shape_and_redact_creation_debug() { assert_eq!(list["items"][0]["dailyLimitUsd"], "0"); assert_eq!(list["items"][0]["weeklyLimitUsd"], "0"); assert_eq!(list["items"][0]["dailyUsedUsd"], "0"); - assert_eq!(list["items"][0]["unresolvedRequests"], 0); + assert!(list["items"][0].get("unresolvedRequests").is_none()); assert_eq!( DateTime::parse_from_rfc3339( list["items"][0]["lastUsedAt"] diff --git a/backend/crates/gateway-core/src/engine/budget.rs b/backend/crates/gateway-core/src/engine/budget.rs index 92126d6cb..1d5d30cea 100644 --- a/backend/crates/gateway-core/src/engine/budget.rs +++ b/backend/crates/gateway-core/src/engine/budget.rs @@ -1,4 +1,4 @@ -//! Durable per-key USD budgets, independent of best-effort request observations. +//! 按 Key 持久化 USD 限额与已取得费用，不依赖尽力写入的请求观测。 use std::time::SystemTime; @@ -28,21 +28,14 @@ pub struct ClientBudgetStatus { pub weekly_used_usd: Decimal, pub daily_resets_at: Option, pub weekly_resets_at: Option, - pub unresolved_requests: u64, -} - -#[derive(Debug, Clone)] -pub struct ClientBudgetAdmission { - pub key_id: ClientApiKeyId, - pub request_id: ModelRequestId, - pub deadline_at: SystemTime, } #[derive(Debug, Clone)] pub struct ClientBudgetCharge { + pub key_id: ClientApiKeyId, pub request_id: ModelRequestId, - /// None means the request may have incurred an unknown charge, never zero. - pub amount_usd: Option, + /// 已取得的 USD 费用，包含重试；缺少费用的尝试按零累计。 + pub amount_usd: Decimal, pub completed_at: SystemTime, } @@ -51,9 +44,9 @@ pub struct ClientBudgetCharge { pub struct ClientBudgetError; pub trait ClientBudgetPort: Send + Sync { - /// Persist a pending charge before any upstream work; read the current policy atomically. - fn admit(&self, request: ClientBudgetAdmission) -> BoxFuture<'_, Result<(), GatewayError>>; + /// 原子检查当前限额与已用金额，不创建预扣费或待结算记录。 + fn admit(&self, key_id: ClientApiKeyId) -> BoxFuture<'_, Result<(), GatewayError>>; - /// Idempotent by gateway request ID. Failed writes leave the durable pending charge intact. + /// 按网关请求 ID 幂等累计已取得费用；写入失败由 Store 重试。 fn settle(&self, charge: ClientBudgetCharge) -> BoxFuture<'_, Result<(), ClientBudgetError>>; } diff --git a/backend/crates/gateway-core/src/engine/coordinator.rs b/backend/crates/gateway-core/src/engine/coordinator.rs index bdd9250bf..2bfa24523 100644 --- a/backend/crates/gateway-core/src/engine/coordinator.rs +++ b/backend/crates/gateway-core/src/engine/coordinator.rs @@ -167,8 +167,7 @@ where request_id, client_api_key_ref, observation: ResponseObservation::new(timing_started_at), - budget_prior_attempts_usd: Some(Decimal::ZERO), - budget_attempt_send_state: UpstreamSendState::NotSent, + budget_prior_attempts_usd: Decimal::ZERO, budget_attempt_already_counted: false, trace, deadline, @@ -258,8 +257,7 @@ pub struct ResponseExecutionSession { request_id: ModelRequestId, client_api_key_ref: crate::policy::ClientApiKeyId, observation: ResponseObservation, - budget_prior_attempts_usd: Option, - budget_attempt_send_state: UpstreamSendState, + budget_prior_attempts_usd: Decimal, budget_attempt_already_counted: bool, trace: TraceContext, deadline: SystemTime, @@ -486,30 +484,31 @@ where #[must_use] pub fn budget_charge(&self) -> super::budget::ClientBudgetCharge { + // 超出数据库可表示范围时保留最大金额，避免溢出后误记为零。 let amount_usd = self .budget_prior_attempts_usd - .zip(self.budget_attempt_usd()) - .and_then(|(prior, current)| prior.checked_add(current)); + .checked_add(self.budget_attempt_usd()) + .unwrap_or(Decimal::MAX); super::budget::ClientBudgetCharge { + key_id: self.client_api_key_ref.clone(), request_id: self.request_id.clone(), amount_usd, completed_at: self.finalized_at.unwrap_or_else(SystemTime::now), } } - fn budget_attempt_usd(&self) -> Option { + fn budget_attempt_usd(&self) -> Decimal { if self.budget_attempt_already_counted { - return Some(Decimal::ZERO); + return Decimal::ZERO; } + // Key 只累计已取得的 USD 费用；缺少费用的请求按零结算。 + // 发送状态仍用于判断重放是否安全并保留在诊断中，不能据此推断存在欠费。 self.observation .cost .total() .filter(|money| money.currency().as_str() == "USD") .map(crate::metering::Money::amount) - .or_else(|| { - (self.budget_attempt_send_state == UpstreamSendState::NotSent) - .then_some(Decimal::ZERO) - }) + .unwrap_or(Decimal::ZERO) } /// 返回截至当前已完成的实际上游调用结果。 @@ -786,8 +785,6 @@ where }), ); let provider_request = ProviderRequest::new(self.operation.clone(), candidate.clone()); - // Cancellation while a provider is opening cannot prove the request was never sent. - self.budget_attempt_send_state = UpstreamSendState::Ambiguous; let stream = match poll_provider( provider, provider_request, @@ -811,7 +808,6 @@ where ProviderBoundary::Result(result) => match *result { Ok(stream) => stream, Err(error) => { - self.budget_attempt_send_state = error.send_state(); record_trace_error(&attempt_trace, &error); if self.prepare_unavailable_native_continuation_replay(&error) { return Ok(Some(PullOutcome::AttemptDiscarded)); @@ -1068,7 +1064,6 @@ where .await; } current.send_observed = true; - self.budget_attempt_send_state = UpstreamSendState::Sent; self.send_state_watermark = UpstreamSendState::Sent; } } @@ -1108,7 +1103,6 @@ where error.send_state() }; let send_state = self.raise_send_watermark(attempt_send_state); - self.budget_attempt_send_state = attempt_send_state; if self.request_persisted { best_effort_store_write( "mark_send_state", @@ -1365,12 +1359,11 @@ where } fn reset_uncommitted_observations(&mut self) { - // Response facts are attempt-local; the key must also pay for discarded attempts. + // 响应观测按尝试隔离；被丢弃尝试中已取得的费用仍计入 Key。 self.budget_prior_attempts_usd = self .budget_prior_attempts_usd - .zip(self.budget_attempt_usd()) - .and_then(|(prior, current)| prior.checked_add(current)); - self.budget_attempt_send_state = UpstreamSendState::NotSent; + .checked_add(self.budget_attempt_usd()) + .unwrap_or(Decimal::MAX); self.budget_attempt_already_counted = false; self.observation.reset_for_attempt(); self.upstream_complete = false; diff --git a/backend/crates/gateway-core/src/engine/execution.rs b/backend/crates/gateway-core/src/engine/execution.rs index 4a6647231..4ad8bb939 100644 --- a/backend/crates/gateway-core/src/engine/execution.rs +++ b/backend/crates/gateway-core/src/engine/execution.rs @@ -14,7 +14,7 @@ use uuid::Uuid; use crate::engine::admission::{ ClientAdmissionDecision, ClientAdmissionPort, ClientAdmissionRejection, ClientAdmissionRequest, }; -use crate::engine::budget::{ClientBudgetAdmission, ClientBudgetCharge, ClientBudgetPort}; +use crate::engine::budget::{ClientBudgetCharge, ClientBudgetPort}; use crate::engine::continuation::{ ContinuationBinding, NativeContinuationPin, NativeContinuationPort, NativeContinuationStoreErrorKind, PreviousResponseId, @@ -523,13 +523,7 @@ impl DefaultExecutionService { model_request_id: request_id.clone(), }; if let Some(budget) = &self.budget - && let Err(error) = budget - .admit(ClientBudgetAdmission { - key_id: client.policy.key_id().clone(), - request_id: request_id.clone(), - deadline_at, - }) - .await + && let Err(error) = budget.admit(client.policy.key_id().clone()).await { admission.release().await; return Err(error); @@ -586,8 +580,9 @@ impl DefaultExecutionService { settle_budget( budget.as_ref(), ClientBudgetCharge { + key_id: client.policy.key_id().clone(), request_id: request_id.clone(), - amount_usd: Some(crate::metering::Decimal::ZERO), + amount_usd: crate::metering::Decimal::ZERO, completed_at: SystemTime::now(), }, ) @@ -966,7 +961,7 @@ struct AdmissionLease { async fn settle_budget(port: &dyn ClientBudgetPort, charge: ClientBudgetCharge) { if let Err(error) = port.settle(charge).await { - tracing::error!(%error, "Client budget settlement pending; durable admission requires reconciliation"); + tracing::error!(%error, "Client budget settlement failed; storage will retry on the next request"); } } diff --git a/backend/crates/gateway-core/src/metering/mod.rs b/backend/crates/gateway-core/src/metering/mod.rs index 242b465b2..d3e0245ed 100644 --- a/backend/crates/gateway-core/src/metering/mod.rs +++ b/backend/crates/gateway-core/src/metering/mod.rs @@ -14,6 +14,8 @@ pub struct Decimal(u128); impl Decimal { pub const ZERO: Self = Self(0); + /// 数据库可表示的最大非负金额。 + pub const MAX: Self = Self(MAX_SCALED_DECIMAL); /// 从按十位小数缩放的整数创建。 /// diff --git a/backend/crates/gateway-core/tests/engine/coordinator.rs b/backend/crates/gateway-core/tests/engine/coordinator.rs index 579ccbaa3..9838c8f8c 100644 --- a/backend/crates/gateway-core/tests/engine/coordinator.rs +++ b/backend/crates/gateway-core/tests/engine/coordinator.rs @@ -2159,18 +2159,23 @@ fn discarded_attempt_cost_never_leaks_into_retry_result() { assert_eq!(state.finalizations[0].cost_source, CostSource::Unavailable); assert_eq!(state.finalizations[0].cost_ticks, None); assert_eq!( - session.budget_charge().amount_usd, - None, - "unknown final usage cannot erase a charged retry" + session.budget_charge().amount_usd.scaled(), + 999, + "missing final usage cannot erase a charged retry" ); } #[test] -fn key_budget_accumulates_retries_and_preserves_unknown_attempts() { +fn key_budget_accumulates_only_observed_retry_costs() { for (first_cost, first_send, expected) in [ - (Some(999), UpstreamSendState::Sent, Some(1122)), - (None, UpstreamSendState::Sent, None), - (None, UpstreamSendState::NotSent, Some(123)), + (Some(999), UpstreamSendState::Sent, 1122), + ( + Some(gateway_core::metering::Decimal::MAX.scaled()), + UpstreamSendState::Sent, + gateway_core::metering::Decimal::MAX.scaled(), + ), + (None, UpstreamSendState::Sent, 123), + (None, UpstreamSendState::NotSent, 123), ] { let operation = generate_operation(); let route_plan = plan(&operation); @@ -2225,7 +2230,7 @@ fn key_budget_accumulates_retries_and_preserves_unknown_attempts() { block_on(session.collect_uncommitted()).unwrap(); block_on(session.commit_downstream(Some(200))).unwrap(); let charge = session.budget_charge(); - assert_eq!(charge.amount_usd.map(|amount| amount.scaled()), expected); + assert_eq!(charge.amount_usd.scaled(), expected); assert_eq!( session.budget_charge().completed_at, charge.completed_at, @@ -2234,6 +2239,67 @@ fn key_budget_accumulates_retries_and_preserves_unknown_attempts() { } } +#[test] +fn key_budget_settles_post_send_failures_from_observed_cost_without_replaying() { + for send_state in [UpstreamSendState::Sent, UpstreamSendState::Ambiguous] { + for cost in [None, Some(123)] { + for committed in [false, true] { + let operation = generate_operation(); + let route_plan = plan(&operation); + let mut items = vec![Ok(GatewayEvent::Started(ResponseMeta::new( + "response-interrupted", + "gpt-5", + )))]; + if let Some(cost) = cost { + items.push(Ok(GatewayEvent::ProviderCost( + ProviderReportedCost::from_usd_ticks(cost).unwrap(), + ))); + } + items.push(Err(ProviderError::new( + ProviderErrorKind::Transport, + send_state, + ))); + let (coordinator, _, provider) = coordinator(vec![Script::Stream { + account_id: "acct_first", + items, + }]); + let mut session = block_on(coordinator.start( + model_request(&operation, SystemTime::now() + Duration::from_secs(30)), + operation, + route_plan, + None, + None, + CancellationToken::new(), + )) + .unwrap(); + if committed { + block_on(session.next_event()).unwrap().unwrap(); + block_on(session.commit_downstream(Some(200))).unwrap(); + loop { + match block_on(session.next_event()) { + Ok(Some(_)) => {} + Err(EngineError::Provider(_)) => break, + other => panic!("expected terminal transport error, got {other:?}"), + } + } + } else { + assert!(matches!( + block_on(session.collect_uncommitted()), + Err(EngineError::Provider(_)) + )); + } + assert!(session.is_finalized()); + assert_eq!( + session.budget_charge().amount_usd.scaled(), + cost.unwrap_or(0), + "send={send_state:?}, committed={committed}, cost={cost:?}" + ); + assert_eq!(provider.contexts.lock().unwrap().len(), 1); + } + } + } +} + #[test] fn pre_commit_failure_excludes_account_and_retries_same_target() { let operation = generate_operation(); @@ -2812,11 +2878,8 @@ fn empty_selection_after_atomic_failure_should_deliver_the_last_upstream_wire() assert_eq!(provider.contexts.lock().expect("contexts lock").len(), 2); assert_eq!( - session - .budget_charge() - .amount_usd - .map(|amount| amount.scaled()), - Some(500), + session.budget_charge().amount_usd.scaled(), + 500, "restoring a discarded failure must not charge its usage twice" ); let state = store.state.lock().expect("store lock"); @@ -2956,6 +3019,7 @@ fn ambiguous_send_state_stops_retry() { let finalization = &state.finalizations[0]; assert_eq!(finalization.outcome, ExecutionOutcome::Failed); assert_eq!(finalization.send_state, UpstreamSendState::Ambiguous); + assert_eq!(session.budget_charge().amount_usd.scaled(), 0); assert_eq!(finalization.attempt_count, 1); assert!(!finalization.committed); assert!(finalization.latency_ms.is_some()); @@ -3430,6 +3494,7 @@ fn cancellation_before_pending_delivery_commit_reaches_terminal_state() { assert_eq!(state.commits, 0); assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Cancelled); assert!(!state.finalizations[0].committed); + assert_eq!(session.budget_charge().amount_usd.scaled(), 0); } #[test] diff --git a/backend/crates/gateway-core/tests/engine/execution.rs b/backend/crates/gateway-core/tests/engine/execution.rs index 783d890a6..6362421ed 100644 --- a/backend/crates/gateway-core/tests/engine/execution.rs +++ b/backend/crates/gateway-core/tests/engine/execution.rs @@ -7,9 +7,7 @@ use std::{ time::{Duration, Instant, SystemTime}, }; -use gateway_core::engine::budget::{ - ClientBudgetAdmission, ClientBudgetCharge, ClientBudgetError, ClientBudgetPort, -}; +use gateway_core::engine::budget::{ClientBudgetCharge, ClientBudgetError, ClientBudgetPort}; use gateway_core::error::GatewayError; #[derive(Default)] @@ -51,7 +49,7 @@ struct Budget { } impl ClientBudgetPort for Budget { - fn admit(&self, _: ClientBudgetAdmission) -> BoxFuture<'_, Result<(), GatewayError>> { + fn admit(&self, _: ClientApiKeyId) -> BoxFuture<'_, Result<(), GatewayError>> { Box::pin(async { assert!(self.active.load(Ordering::SeqCst)); if self.reject { @@ -230,7 +228,7 @@ fn cancellation_and_pre_send_failure_settle_zero_and_release_concurrency_for_all assert!( charges .iter() - .all(|charge| charge.amount_usd == Some(gateway_core::metering::Decimal::ZERO)) + .all(|charge| charge.amount_usd == gateway_core::metering::Decimal::ZERO) ); } diff --git a/backend/crates/gateway-store/src/postgres/client_budgets.rs b/backend/crates/gateway-store/src/postgres/client_budgets.rs index 7bb762360..cf734f169 100644 --- a/backend/crates/gateway-store/src/postgres/client_budgets.rs +++ b/backend/crates/gateway-store/src/postgres/client_budgets.rs @@ -1,4 +1,4 @@ -//! Durable admission and idempotent USD settlement, serialized per client key. +//! 按 Key 串行检查限额，并幂等累计已取得的 USD 费用。 use std::{collections::BTreeMap, sync::Mutex, time::Duration}; @@ -6,11 +6,12 @@ use chrono::{DateTime, Utc}; use futures::future::BoxFuture; use gateway_core::{ engine::budget::{ - ClientBudgetAdmission, ClientBudgetCharge, ClientBudgetError, ClientBudgetLimits, - ClientBudgetPort, ClientBudgetStatus, + ClientBudgetCharge, ClientBudgetError, ClientBudgetLimits, ClientBudgetPort, + ClientBudgetStatus, }, error::{GatewayError, GatewayErrorKind}, metering::Decimal, + policy::ClientApiKeyId, }; use sqlx::{PgPool, Postgres, Row, Transaction}; @@ -30,42 +31,25 @@ impl PgClientBudgetStore { } } - async fn admit_inner(&self, request: ClientBudgetAdmission) -> Result<(), GatewayError> { - // Preserve exact charges across a transient outage. On process loss the durable - // pending event still blocks a limited key after its deadline, requiring reconciliation. + async fn admit_inner(&self, key_id: ClientApiKeyId) -> Result<(), GatewayError> { + // 短暂存储故障后按原金额重试；进程退出丢失的费用不转成人工核账或阻断 Key。 let retries = self .retry .lock() .map_err(|_| unavailable())? .values() + .filter(|charge| charge.key_id == key_id) .cloned() .collect::>(); - if !retries.is_empty() { - let ids = retries - .iter() - .map(|charge| charge.request_id.as_str()) - .collect::>(); - let owned = sqlx::query_scalar::<_, String>( - "select request_id from client_key_charge_events - where client_api_key_id = $1 and request_id = any($2)", - ) - .bind(request.key_id.as_str()) - .bind(ids) - .fetch_all(&self.pool) - .await - .map_err(|_| unavailable())?; - for charge in retries { - if owned.iter().any(|id| id == charge.request_id.as_str()) { - self.settle(charge).await.map_err(|_| unavailable())?; - } - } + for charge in retries { + self.settle(charge).await.map_err(|_| unavailable())?; } let mut tx = self.pool.begin().await.map_err(|_| unavailable())?; let row = sqlx::query( "select daily_limit_usd::text, weekly_limit_usd::text, enabled from client_api_keys where id = $1 for update", ) - .bind(request.key_id.as_str()) + .bind(key_id.as_str()) .fetch_optional(&mut *tx) .await .map_err(|_| unavailable())? @@ -92,25 +76,15 @@ impl PgClientBudgetStore { .map_err(|_| unavailable())?, }; let now = Utc::now(); - advance_windows(&mut tx, request.key_id.as_str(), now) + advance_windows(&mut tx, key_id.as_str(), now) .await .map_err(|_| unavailable())?; if limits.is_limited() { - let unresolved = sqlx::query_scalar::<_, bool>("select exists(select 1 from client_key_charge_events - where client_api_key_id = $1 and (state = 'unknown' or (state = 'pending' and deadline_at <= $2)))") - .bind(request.key_id.as_str()).bind(now).fetch_one(&mut *tx).await.map_err(|_| unavailable())?; - if unresolved { - return Err(GatewayError::new( - GatewayErrorKind::RateLimited, - "client API key has unresolved usage; administrator reconciliation is required", - ) - .with_client_code("key_budget_unresolved")); - } let window = sqlx::query( "select daily_used_usd::text, weekly_used_usd::text, daily_end, weekly_end from client_key_budget_windows where client_api_key_id = $1", ) - .bind(request.key_id.as_str()) + .bind(key_id.as_str()) .fetch_one(&mut *tx) .await .map_err(|_| unavailable())?; @@ -145,32 +119,20 @@ impl PgClientBudgetStore { .with_retry_after(retry)); } } - sqlx::query( - "insert into client_key_charge_events (request_id, client_api_key_id, deadline_at) - values ($1, $2, $3)", - ) - .bind(request.request_id.as_str()) - .bind(request.key_id.as_str()) - .bind(DateTime::::from(request.deadline_at)) - .execute(&mut *tx) - .await - .map_err(|_| unavailable())?; tx.commit().await.map_err(|_| unavailable()) } async fn settle_inner(&self, charge: &ClientBudgetCharge) -> Result<(), ClientBudgetError> { let mut tx = self.pool.begin().await.map_err(|_| ClientBudgetError)?; - // Match admission's key -> window -> event lock order, including concurrent settlements. + // 与准入统一先锁 Key，再写窗口和费用，串行化同一 Key 的并发结算。 let key = sqlx::query_scalar::<_, String>( - "select k.id from client_api_keys k - join client_key_charge_events e on e.client_api_key_id = k.id - where e.request_id = $1 for update of k", + "select id from client_api_keys where id = $1 for update", ) - .bind(charge.request_id.as_str()) + .bind(charge.key_id.as_str()) .fetch_optional(&mut *tx) .await .map_err(|_| ClientBudgetError)?; - let Some(key) = key else { return Ok(()) }; // Deleting a key also deletes its ledger. + let Some(key) = key else { return Ok(()) }; // 删除 Key 时也会删除其费用记录。 settle_in_transaction(&mut tx, &key, charge) .await .map_err(|_| ClientBudgetError)?; @@ -184,38 +146,33 @@ async fn settle_in_transaction( charge: &ClientBudgetCharge, ) -> Result<(), sqlx::Error> { advance_windows(tx, key, Utc::now()).await?; + // 仅在请求结束时写入费用；请求 ID 冲突时不重复累计。 let changed = sqlx::query( - "update client_key_charge_events - set state = $2, amount_usd = $3::text::numeric, completed_at = $4 - where request_id = $1 and state <> 'settled'", + "insert into client_key_charge_events (request_id, client_api_key_id, amount_usd, completed_at) + values ($1, $2, $3::text::numeric, $4) + on conflict (request_id) do nothing", ) .bind(charge.request_id.as_str()) - .bind(if charge.amount_usd.is_some() { - "settled" - } else { - "unknown" - }) - .bind(charge.amount_usd.map(|amount| amount.canonical())) + .bind(key) + .bind(charge.amount_usd.canonical()) .bind(DateTime::::from(charge.completed_at)) .execute(&mut **tx) .await? .rows_affected(); - if changed == 1 - && let Some(amount) = charge.amount_usd - { + if changed == 1 { sqlx::query("update client_key_budget_windows set daily_used_usd = daily_used_usd + case when $3 >= daily_start and $3 < daily_end then $2::text::numeric else 0 end, weekly_used_usd = weekly_used_usd + case when $3 >= weekly_start and $3 < weekly_end then $2::text::numeric else 0 end where client_api_key_id = $1") - .bind(key).bind(amount.canonical()).bind(DateTime::::from(charge.completed_at)) + .bind(key).bind(charge.amount_usd.canonical()).bind(DateTime::::from(charge.completed_at)) .execute(&mut **tx).await?; } Ok(()) } impl ClientBudgetPort for PgClientBudgetStore { - fn admit(&self, request: ClientBudgetAdmission) -> BoxFuture<'_, Result<(), GatewayError>> { - Box::pin(async move { self.admit_inner(request).await }) + fn admit(&self, key_id: ClientApiKeyId) -> BoxFuture<'_, Result<(), GatewayError>> { + Box::pin(async move { self.admit_inner(key_id).await }) } fn settle(&self, charge: ClientBudgetCharge) -> BoxFuture<'_, Result<(), ClientBudgetError>> { @@ -263,16 +220,19 @@ pub(super) async fn load_client_key_budgets( .iter() .map(|record| record.id.as_str()) .collect::>(); - let rows = sqlx::query("select k.id, k.daily_limit_usd::text, k.weekly_limit_usd::text, + let rows = sqlx::query( + "select k.id, k.daily_limit_usd::text, k.weekly_limit_usd::text, (case when w.daily_end > now() then w.daily_used_usd else 0 end)::text as daily_used, (case when w.weekly_end > now() then w.weekly_used_usd else 0 end)::text as weekly_used, case when w.daily_end > now() then w.daily_end end as daily_end, - case when w.weekly_end > now() then w.weekly_end end as weekly_end, - (select count(*) from client_key_charge_events e where e.client_api_key_id = k.id - and (e.state = 'unknown' or (e.state = 'pending' and e.deadline_at <= now()))) as unresolved + case when w.weekly_end > now() then w.weekly_end end as weekly_end from client_api_keys k left join client_key_budget_windows w on w.client_api_key_id = k.id - where k.id = any($1)") - .bind(ids).fetch_all(pool).await.map_err(|_| postgres_unavailable("load client budgets"))?; + where k.id = any($1)", + ) + .bind(ids) + .fetch_all(pool) + .await + .map_err(|_| postgres_unavailable("load client budgets"))?; let mut budgets = BTreeMap::new(); for row in rows { let parse = |field| -> StoreResult { @@ -295,7 +255,6 @@ pub(super) async fn load_client_key_budgets( weekly_resets_at: row .get::>, _>("weekly_end") .map(Into::into), - unresolved_requests: row.get::("unresolved").unsigned_abs(), }, ); } @@ -314,99 +273,3 @@ fn unavailable() -> GatewayError { ) .with_client_code("key_budget_unavailable") } - -pub(super) async fn unresolved_charges( - pool: &PgPool, - key: &gateway_core::policy::ClientApiKeyId, -) -> gateway_admin::ports::store::AdminStoreResult< - Vec, -> { - let rows = sqlx::query("select request_id, started_at, completed_at, state from client_key_charge_events - where client_api_key_id = $1 and (state = 'unknown' or (state = 'pending' and deadline_at <= now())) - order by started_at, request_id limit 200") - .bind(key.as_str()).fetch_all(pool).await.map_err(|_| budget_admin_unavailable())?; - Ok(rows - .into_iter() - .map( - |row| gateway_admin::model::client_keys::UnresolvedClientCharge { - request_id: row.get("request_id"), - started_at: row.get("started_at"), - completed_at: row.get("completed_at"), - state: row.get("state"), - }, - ) - .collect()) -} - -pub(super) async fn reconcile_charge( - pool: &PgPool, - command: gateway_admin::model::client_keys::ReconcileClientCharge, - context: &gateway_admin::model::MutationContext, - revision: gateway_admin::model::Revision, -) -> gateway_admin::ports::store::AdminStoreResult<()> { - use gateway_admin::ports::store::{AdminStoreError, AdminStoreErrorKind}; - let conflict = || { - AdminStoreError::new( - AdminStoreErrorKind::Conflict, - "client budget", - "charge is missing, active, or already settled with another amount", - ) - }; - let mut tx = pool.begin().await.map_err(|_| budget_admin_unavailable())?; - sqlx::query("select id from client_api_keys where id = $1 for update") - .bind(command.key_id.as_str()) - .fetch_optional(&mut *tx) - .await - .map_err(|_| budget_admin_unavailable())? - .ok_or_else(conflict)?; - let row = sqlx::query("select state, amount_usd::text, coalesce(completed_at, deadline_at) as completed_at, deadline_at <= now() as expired - from client_key_charge_events where request_id = $1 and client_api_key_id = $2 for update") - .bind(&command.request_id).bind(command.key_id.as_str()).fetch_optional(&mut *tx).await.map_err(|_| budget_admin_unavailable())?.ok_or_else(conflict)?; - let state: &str = row.get("state"); - if state == "settled" { - return if row.get::("amount_usd").parse::().ok() - == Some(command.amount_usd) - { - Ok(()) - } else { - Err(conflict()) - }; - } - if state == "pending" && !row.get::("expired") { - return Err(conflict()); - } - let charge = ClientBudgetCharge { - request_id: gateway_core::engine::ModelRequestId::new(command.request_id.clone()) - .map_err(|_| conflict())?, - amount_usd: Some(command.amount_usd), - completed_at: row.get::, _>("completed_at").into(), - }; - settle_in_transaction(&mut tx, command.key_id.as_str(), &charge) - .await - .map_err(|_| budget_admin_unavailable())?; - sqlx::query("update client_key_charge_events set reconciled_at = now(), reconciliation_reason = $2 where request_id = $1") - .bind(&command.request_id).bind(command.reason).execute(&mut *tx).await.map_err(|_| budget_admin_unavailable())?; - let audit = crate::mutation_audit( - context, - "reconcile", - "client_key_charge", - &command.request_id, - vec!["amount_usd".to_owned(), "reconciliation_reason".to_owned()], - ); - super::admin_security_audit::append_admin_audit_event_in_transaction( - &mut tx, - audit, - crate::store_revision(revision)?, - ) - .await - .map_err(|_| budget_admin_unavailable())?; - tx.commit().await.map_err(|_| budget_admin_unavailable()) -} - -fn budget_admin_unavailable() -> gateway_admin::ports::store::AdminStoreError { - gateway_admin::ports::store::AdminStoreError::new( - gateway_admin::ports::store::AdminStoreErrorKind::Unavailable, - "client budget", - "storage unavailable", - ) -} diff --git a/backend/crates/gateway-store/src/postgres/client_keys.rs b/backend/crates/gateway-store/src/postgres/client_keys.rs index 26f7ec873..f279ca220 100644 --- a/backend/crates/gateway-store/src/postgres/client_keys.rs +++ b/backend/crates/gateway-store/src/postgres/client_keys.rs @@ -678,22 +678,6 @@ impl PgAdminClientKeyStore { #[async_trait] impl ClientKeyStore for PgAdminClientKeyStore { - async fn unresolved_charges( - &self, - id: &ClientApiKeyId, - ) -> AdminStoreResult> { - self.required_record(id).await?; - super::client_budgets::unresolved_charges(&self.keys.pool, id).await - } - - async fn reconcile_charge( - &self, - command: gateway_admin::model::client_keys::ReconcileClientCharge, - context: &MutationContext, - ) -> AdminStoreResult<()> { - let revision = self.revision().await?; - super::client_budgets::reconcile_charge(&self.keys.pool, command, context, revision).await - } async fn list_client_keys( &self, query: AdminClientKeyListQuery, diff --git a/backend/crates/gateway-store/tests/postgres/client_budgets.rs b/backend/crates/gateway-store/tests/postgres/client_budgets.rs index 82fdc8a17..f95f5e4d7 100644 --- a/backend/crates/gateway-store/tests/postgres/client_budgets.rs +++ b/backend/crates/gateway-store/tests/postgres/client_budgets.rs @@ -1,17 +1,14 @@ -use std::time::{Duration, SystemTime}; +use std::time::SystemTime; use chrono::{DateTime, Utc}; use gateway_admin::{ - model::{ - MutationActor, MutationContext, - client_keys::{ReconcileClientCharge, UpdateClientKey}, - }, + model::{MutationActor, MutationContext, client_keys::UpdateClientKey}, ports::store::ClientKeyStore as _, }; use gateway_core::{ engine::{ ModelRequestId, - budget::{ClientBudgetAdmission, ClientBudgetCharge, ClientBudgetPort, ClientBudgetStatus}, + budget::{ClientBudgetCharge, ClientBudgetPort, ClientBudgetStatus}, }, error::GatewayErrorKind, policy::{ClientApiKeyId, RateLimits}, @@ -23,18 +20,15 @@ use gateway_store::postgres::{ use super::TestDatabase; -fn admission(key: &str, request: &str) -> ClientBudgetAdmission { - ClientBudgetAdmission { - key_id: ClientApiKeyId::new(key).unwrap(), - request_id: ModelRequestId::new(format!("req_{request}")).unwrap(), - deadline_at: SystemTime::now() + Duration::from_secs(600), - } +fn key_id(key: &str) -> ClientApiKeyId { + ClientApiKeyId::new(key).unwrap() } -fn charge(request: &str, amount: Option<&str>) -> ClientBudgetCharge { +fn charge(key: &str, request: &str, amount: &str) -> ClientBudgetCharge { ClientBudgetCharge { + key_id: key_id(key), request_id: ModelRequestId::new(format!("req_{request}")).unwrap(), - amount_usd: amount.map(|value| value.parse().unwrap()), + amount_usd: amount.parse().unwrap(), completed_at: SystemTime::now(), } } @@ -65,7 +59,7 @@ async fn status(database: &TestDatabase, key: &str) -> ClientBudgetStatus { fn context() -> MutationContext { MutationContext { actor: MutationActor::System, - request_id: "reconcile-test".to_owned(), + request_id: "budget-test".to_owned(), } } @@ -79,16 +73,15 @@ async fn budgets_settle_exactly_once_and_enforce_each_threshold_across_store_ins let first = PgClientBudgetStore::new(database.pool.clone()); let second = PgClientBudgetStore::new(database.pool.clone()); for (key, prefix, amount) in [("day", "d", "0.1"), ("week", "w", "0.1")] { - for index in 0..3 { - let id = format!("{prefix}-{index}"); - // Admitted work can finish above the threshold; no estimated cost reservation. - first.admit(admission(key, &id)).await.unwrap(); + for _ in 0..3 { + // 已准入请求可完成并超过限额，不预占估算费用。 + first.admit(key_id(key)).await.unwrap(); } for index in 0..3 { let id = format!("{prefix}-{index}"); let (a, b) = tokio::join!( - first.settle(charge(&id, Some(amount))), - second.settle(charge(&id, Some(amount))) + first.settle(charge(key, &id, amount)), + second.settle(charge(key, &id, amount)) ); a.unwrap(); b.unwrap(); @@ -97,31 +90,23 @@ async fn budgets_settle_exactly_once_and_enforce_each_threshold_across_store_ins let day = status(&database, "day").await; assert_eq!(day.daily_used_usd.canonical(), "0.3"); assert_eq!(day.weekly_used_usd.canonical(), "0.3"); - let day_error = second - .admit(admission("day", "d-rejected")) - .await - .unwrap_err(); + let day_error = second.admit(key_id("day")).await.unwrap_err(); assert_eq!(day_error.kind(), GatewayErrorKind::RateLimited); assert_eq!( day_error.client_error_code(), Some("key_daily_budget_exceeded") ); assert!(day_error.retry_after().is_some()); - let week_error = first - .admit(admission("week", "w-rejected")) - .await - .unwrap_err(); + let week_error = first.admit(key_id("week")).await.unwrap_err(); assert_eq!( week_error.client_error_code(), Some("key_weekly_budget_exceeded") ); - let rejected_count: i64 = sqlx::query_scalar( - "select count(*) from client_key_charge_events where request_id like '%rejected'", - ) - .fetch_one(&database.pool) - .await - .unwrap(); - assert_eq!(rejected_count, 0); + let event_count: i64 = sqlx::query_scalar("select count(*) from client_key_charge_events") + .fetch_one(&database.pool) + .await + .unwrap(); + assert_eq!(event_count, 6, "rejected admission must not create charges"); database.close().await; } @@ -132,17 +117,17 @@ async fn window_rollover_is_shanghai_midnight_and_seven_days_with_late_settlemen }; seed(&database, "key", "1", "2").await; let store = PgClientBudgetStore::new(database.pool.clone()); - store.admit(admission("key", "first")).await.unwrap(); + store.admit(key_id("key")).await.unwrap(); let (day_start, day_end, week_end): (DateTime, DateTime, DateTime) = sqlx::query_as("select daily_start, daily_end, weekly_end from client_key_budget_windows where client_api_key_id = 'key'") .fetch_one(&database.pool).await.unwrap(); assert_eq!(day_start.timestamp().rem_euclid(86400), 16 * 3600); assert_eq!((day_end - day_start).num_hours(), 24); assert_eq!((week_end - day_start).num_hours(), 168); - store.settle(charge("first", Some("1"))).await.unwrap(); + store.settle(charge("key", "first", "1")).await.unwrap(); sqlx::query("update client_key_budget_windows set daily_start = daily_start - interval '1 day', daily_end = daily_start") .execute(&database.pool).await.unwrap(); - store.admit(admission("key", "second")).await.unwrap(); + store.admit(key_id("key")).await.unwrap(); assert_eq!( status(&database, "key").await.daily_used_usd.canonical(), "0" @@ -151,16 +136,16 @@ async fn window_rollover_is_shanghai_midnight_and_seven_days_with_late_settlemen status(&database, "key").await.weekly_used_usd.canonical(), "1" ); - store.settle(charge("second", Some("1"))).await.unwrap(); + store.settle(charge("key", "second", "1")).await.unwrap(); sqlx::query("update client_key_budget_windows set daily_end = now() - interval '1 second', weekly_end = now() - interval '1 second'") .execute(&database.pool).await.unwrap(); let virtual_reset = status(&database, "key").await; assert_eq!(virtual_reset.daily_used_usd.canonical(), "0"); assert_eq!(virtual_reset.weekly_used_usd.canonical(), "0"); - store.admit(admission("key", "after-reset")).await.unwrap(); + store.admit(key_id("key")).await.unwrap(); let old = ClientBudgetCharge { completed_at: (day_start - chrono::Duration::seconds(1)).into(), - ..charge("after-reset", Some("0.9")) + ..charge("key", "after-reset", "0.9") }; store.settle(old).await.unwrap(); let reset = status(&database, "key").await; @@ -171,103 +156,80 @@ async fn window_rollover_is_shanghai_midnight_and_seven_days_with_late_settlemen } #[tokio::test] -async fn unknown_and_orphan_charges_block_limited_keys_until_audited_reconciliation() { - let Some(database) = TestDatabase::create("budgets_unknown").await else { +async fn zero_cost_and_interrupted_requests_never_block_limited_keys() { + let Some(database) = TestDatabase::create("budgets_zero").await else { return; }; seed(&database, "key", "1", "5").await; - seed(&database, "other", "0", "0").await; let store = PgClientBudgetStore::new(database.pool.clone()); - let admin = PgAdminClientKeyStore::new(database.pool.clone()); - store.admit(admission("key", "unknown")).await.unwrap(); - store.settle(charge("unknown", None)).await.unwrap(); + store.admit(key_id("key")).await.unwrap(); + store.settle(charge("key", "no-cost", "0")).await.unwrap(); + // 模拟准入后进程退出，没有费用可结算；重启后仍应允许同一 Key 使用。 + store.admit(key_id("key")).await.unwrap(); + let restarted = PgClientBudgetStore::new(database.pool.clone()); + restarted.admit(key_id("key")).await.unwrap(); + let events: i64 = sqlx::query_scalar("select count(*) from client_key_charge_events") + .fetch_one(&database.pool) + .await + .unwrap(); + assert_eq!(events, 1, "admission must not leave pending charges"); assert_eq!( - store - .admit(admission("key", "blocked")) - .await - .unwrap_err() - .client_error_code(), - Some("key_budget_unresolved") + status(&database, "key").await.daily_used_usd.canonical(), + "0" ); - let key_id = ClientApiKeyId::new("key").unwrap(); - assert_eq!(admin.unresolved_charges(&key_id).await.unwrap().len(), 1); - assert_eq!(status(&database, "key").await.unresolved_requests, 1); - let reconcile = |key: &str, request: &str, amount: &str| ReconcileClientCharge { - key_id: ClientApiKeyId::new(key).unwrap(), - request_id: format!("req_{request}"), - amount_usd: amount.parse().unwrap(), - reason: "verified against upstream usage".to_owned(), - }; - assert!( - admin - .reconcile_charge(reconcile("other", "unknown", "0.4"), &context()) - .await - .is_err() - ); - admin - .reconcile_charge(reconcile("key", "unknown", "0.4"), &context()) + restarted + .settle(charge("key", "known", "0.4")) .await .unwrap(); - admin - .reconcile_charge(reconcile("key", "unknown", "0.4"), &context()) - .await - .unwrap(); - assert!( - admin - .reconcile_charge(reconcile("key", "unknown", "0.5"), &context()) - .await - .is_err() - ); assert_eq!( status(&database, "key").await.daily_used_usd.canonical(), "0.4" ); - let audits: i64 = sqlx::query_scalar( - "select count(*) from admin_audit_events where entity_kind = 'client_key_charge'", + database.close().await; +} + +#[tokio::test] +async fn transient_settlement_failure_rolls_back_and_retries_exact_cost_before_admission() { + let Some(database) = TestDatabase::create("budgets_retry").await else { + return; + }; + seed(&database, "key", "1", "5").await; + let store = PgClientBudgetStore::new(database.pool.clone()); + store.admit(key_id("key")).await.unwrap(); + // 在费用事件插入后让窗口写入失败，验证整个事务回滚。 + sqlx::raw_sql( + "create function reject_budget_update() returns trigger language plpgsql as $$ + begin + if new.daily_used_usd > 0 then raise exception 'temporary test failure'; end if; + return new; + end $$; + create trigger reject_budget_update before update on client_key_budget_windows + for each row execute function reject_budget_update();", ) - .fetch_one(&database.pool) + .execute(&database.pool) .await .unwrap(); - assert_eq!(audits, 1); - store.admit(admission("key", "orphan")).await.unwrap(); - assert!( - admin - .reconcile_charge(reconcile("key", "orphan", "0"), &context()) - .await - .is_err() - ); - sqlx::query("update client_key_charge_events set deadline_at = now() - interval '1 second' where request_id = 'req_orphan'") - .execute(&database.pool).await.unwrap(); - let restarted = PgClientBudgetStore::new(database.pool.clone()); + assert!(store.settle(charge("key", "retry", "1.25")).await.is_err()); + let events: i64 = sqlx::query_scalar("select count(*) from client_key_charge_events") + .fetch_one(&database.pool) + .await + .unwrap(); + assert_eq!(events, 0); assert_eq!( - restarted - .admit(admission("key", "restart-blocked")) - .await - .unwrap_err() - .client_error_code(), - Some("key_budget_unresolved") + status(&database, "key").await.daily_used_usd.canonical(), + "0" ); - admin - .reconcile_charge(reconcile("key", "orphan", "0"), &context()) - .await - .unwrap(); - restarted - .admit(admission("key", "recovered")) - .await - .unwrap(); - assert_eq!(status(&database, "key").await.unresolved_requests, 0); - store - .admit(admission("other", "unlimited-unknown")) - .await - .unwrap(); - store - .settle(charge("unlimited-unknown", None)) - .await - .unwrap(); - store - .admit(admission("other", "unlimited-next")) + sqlx::query("drop trigger reject_budget_update on client_key_budget_windows") + .execute(&database.pool) .await .unwrap(); + let error = store.admit(key_id("key")).await.unwrap_err(); + assert_eq!(error.client_error_code(), Some("key_daily_budget_exceeded")); + store.settle(charge("key", "retry", "1.25")).await.unwrap(); + assert_eq!( + status(&database, "key").await.daily_used_usd.canonical(), + "1.25" + ); database.close().await; } @@ -279,9 +241,9 @@ async fn budget_updates_preserve_omitted_limits_and_do_not_clear_usage() { seed(&database, "key", "0", "0").await; let store = PgClientBudgetStore::new(database.pool.clone()); let admin = PgAdminClientKeyStore::new(database.pool.clone()); - store.admit(admission("key", "unlimited")).await.unwrap(); + store.admit(key_id("key")).await.unwrap(); store - .settle(charge("unlimited", Some("2.75"))) + .settle(charge("key", "unlimited", "2.75")) .await .unwrap(); let update = UpdateClientKey { @@ -317,7 +279,7 @@ async fn budget_updates_preserve_omitted_limits_and_do_not_clear_usage() { assert_eq!(current.daily_used_usd.canonical(), "2.75"); assert_eq!( store - .admit(admission("key", "blocked")) + .admit(key_id("key")) .await .unwrap_err() .client_error_code(), @@ -334,17 +296,13 @@ async fn budget_updates_preserve_omitted_limits_and_do_not_clear_usage() { ) .await .unwrap(); - store.admit(admission("key", "allowed")).await.unwrap(); + store.admit(key_id("key")).await.unwrap(); sqlx::query("update client_api_keys set enabled = false where id = 'key'") .execute(&database.pool) .await .unwrap(); assert_eq!( - store - .admit(admission("key", "disabled")) - .await - .unwrap_err() - .kind(), + store.admit(key_id("key")).await.unwrap_err().kind(), GatewayErrorKind::PolicyDenied ); sqlx::query("delete from client_api_keys where id = 'key'") @@ -352,14 +310,10 @@ async fn budget_updates_preserve_omitted_limits_and_do_not_clear_usage() { .await .unwrap(); assert_eq!( - store - .admit(admission("key", "deleted")) - .await - .unwrap_err() - .kind(), + store.admit(key_id("key")).await.unwrap_err().kind(), GatewayErrorKind::Unauthorized ); - store.settle(charge("allowed", Some("1"))).await.unwrap(); + store.settle(charge("key", "allowed", "1")).await.unwrap(); database.close().await; } @@ -372,12 +326,12 @@ async fn budget_database_outage_fails_closed() { database.pool.close().await; assert_eq!( store - .admit(admission("key", "offline")) + .admit(key_id("key")) .await .unwrap_err() .client_error_code(), Some("key_budget_unavailable") ); - assert!(store.settle(charge("offline", Some("1"))).await.is_err()); + assert!(store.settle(charge("key", "offline", "1")).await.is_err()); database.close().await; } diff --git a/backend/crates/providers/openai/src/transport/websocket/pump.rs b/backend/crates/providers/openai/src/transport/websocket/pump.rs index 7952a8b1e..a16e40dce 100644 --- a/backend/crates/providers/openai/src/transport/websocket/pump.rs +++ b/backend/crates/providers/openai/src/transport/websocket/pump.rs @@ -670,9 +670,10 @@ pub(crate) fn transport_metric_reason(error: &tungstenite::Error) -> &'static st std::io::ErrorKind::TimedOut => "transport_timeout", _ => "transport_error", }, + // 未收到关闭握手只能证明连接异常结束，不能据此认定收到 TCP RST。 tungstenite::Error::Protocol( tungstenite::error::ProtocolError::ResetWithoutClosingHandshake, - ) => "tcp_reset", + ) => "reset_without_closing_handshake", _ => "transport_error", } } diff --git a/backend/crates/providers/openai/tests/transport/latency.rs b/backend/crates/providers/openai/tests/transport/latency.rs index 5419f96a3..86f87a41f 100644 --- a/backend/crates/providers/openai/tests/transport/latency.rs +++ b/backend/crates/providers/openai/tests/transport/latency.rs @@ -1000,7 +1000,7 @@ async fn connection_local_continuation_should_fail_after_its_live_socket_disappe let observation = error .connection_observation() .expect("disappeared connection should retain its lifecycle observation"); - assert_eq!(observation.exit_reason(), "tcp_reset"); + assert_eq!(observation.exit_reason(), "reset_without_closing_handshake"); assert!(observation.age_ms() >= observation.idle_ms()); } diff --git a/backend/migrations/.frozen-sha256 b/backend/migrations/.frozen-sha256 index 80e287176..ce50996c3 100644 --- a/backend/migrations/.frozen-sha256 +++ b/backend/migrations/.frozen-sha256 @@ -1,3 +1,4 @@ b7d3df7b48d835a00a3b74cbefc958a039d38fea5529cdad04bfbc703dca72aa 0001_initial.sql 78c427c17d8b5759405f1629a80ceb624012cb6b31371b762078b000266d86c6 0002_client_key_budgets.sql 46f493ff410cf77550114542e0418475876ee2b2c975436290f0f72440e098c4 0003_account_outbound_proxy.sql +5e4bfd0b47a4de5e99b7ae1d89e0294de61856a40fba3480572db765147c3b3a 0004_remove_client_charge_reconciliation.sql diff --git a/backend/migrations/0004_remove_client_charge_reconciliation.sql b/backend/migrations/0004_remove_client_charge_reconciliation.sql new file mode 100644 index 000000000..17a31a72b --- /dev/null +++ b/backend/migrations/0004_remove_client_charge_reconciliation.sql @@ -0,0 +1,17 @@ +-- 移除人工核账：历史未取得费用的请求按零结算，已记录金额保持不变。 +-- 迁移在监听请求前执行；遗留未完成记录不再阻断 Key。 +update client_key_charge_events +set state = 'settled', + amount_usd = coalesce(amount_usd, 0), + completed_at = coalesce(completed_at, least(deadline_at, now())); + +-- 费用表只保存已完成的幂等结算，不再维护预扣费和待核账状态。 +drop index client_key_charge_events_unresolved; +alter table client_key_charge_events + drop column state, + drop column started_at, + drop column deadline_at, + drop column reconciled_at, + drop column reconciliation_reason, + alter column amount_usd set not null, + alter column completed_at set not null; diff --git a/docs/api.md b/docs/api.md index 4a88f7405..80b44539d 100644 --- a/docs/api.md +++ b/docs/api.md @@ -437,8 +437,6 @@ PostgreSQL 或 Redis。管理端只在用户打开弹窗或点击刷新时调用 | `POST` | `/api/admin/client-keys/enable` | `{ id }` | 启用 | | `POST` | `/api/admin/client-keys/disable` | `{ id }` | 禁用 | | `POST` | `/api/admin/client-keys/delete` | `{ id }` | 删除 | -| `GET` | `/api/admin/client-keys/unresolved-charges` | `id` | 按时间读取最早的 200 笔待核账请求 | -| `POST` | `/api/admin/client-keys/reconcile-charge` | `{ id, requestId, amountUsd, reason }` | 管理员确认未知费用并记录审计 | 创建字段为 `name`、可选 `label`、`groupIds`、`maxConcurrency`、`requestsPerMinute`、可选 `dailyLimitUsd` 和 `weeklyLimitUsd`，更新请求再增加 @@ -452,7 +450,8 @@ PostgreSQL 或 Redis。管理端只在用户打开弹窗或点击刷新时调用 `maxConcurrency` 和 `requestsPerMinute` 是非负整数，零表示不限。 列表增加 `dailyLimitUsd`、`weeklyLimitUsd`、`dailyUsedUsd`、`weeklyUsedUsd`（均为字符串）、 -`dailyResetsAt`、`weeklyResetsAt`（RFC3339 或 `null`）与 `unresolvedRequests`。 +`dailyResetsAt`、`weeklyResetsAt`（RFC3339 或 `null`）。 +管理端日／周金额显示两位小数，悬停可查看原始值；记账和限额比较保留完整精度。 日窗口按北京时间零点重置；周窗口从首次准入当天零点起持续七天，到期后在下一次使用时重新开启。 费用按请求完成时间归属窗口。并发按同一 Key 的执行中请求累计，包含 SSE 与每个 WebSocket `response.create`；空闲连接不占名额，内部重试不重复占用。 @@ -461,12 +460,12 @@ PostgreSQL 或 Redis。管理端只在用户打开弹窗或点击刷新时调用 任一已结算金额达到限额后拒绝新请求，已准入请求可完成并使金额超过阈值。 HTTP 返回 `429`，`error.code` 为 `key_daily_budget_exceeded` 或 `key_weekly_budget_exceeded`， 并附 `Retry-After`；WebSocket 每次 `response.create` 执行相同检查并返回协议错误事件。 -已发送但费用不明、或超过请求期限仍未结算的请求使受限 Key 返回 `key_budget_unresolved`。 +只累计上游上报或按用量与模型价格计算出的 USD 费用；无法取得费用的尝试按零累计， +保留错误和用量诊断，不产生待核账记录或阻断。内部重试中已经取得的费用仍会累计。 预算存储不可用时返回 `503`、`key_budget_unavailable`。 -核账的 `amountUsd` 使用相同金额格式，`reason` 为 1–1024 字节的非空原因。只能处理该 Key 下未知或 -已超过期限的请求；相同金额重复提交不重复计费，修改已结算金额返回 `409`。账本独立于使用统计日志， -记录保留至删除 Key，不受 `usageRetentionDays` 影响。 +自动结算按网关请求 ID 幂等执行。账本独立于使用统计日志，记录保留至删除 Key， +不受 `usageRetentionDays` 影响。 ## 8. 运行设置 diff --git a/docs/architecture.md b/docs/architecture.md index c5fd8083d..4b6e63c3b 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -135,7 +135,7 @@ sequenceDiagram A->>E: Operation + client context E->>E: freeze snapshot and compile routing plan E->>S: enqueue request / attempt observations - E->>S: check Key budget and persist pending charge + E->>S: check Key budget against recorded usage E->>P: one candidate, one credential, one attempt P-->>E: cold canonical stream + raw wire E->>E: enforce send and downstream commit barriers @@ -250,23 +250,23 @@ Continuation 仍受原请求的 Client Key、账号范围、Provider 和发送/ ### Client Key 限额与结算 日金额、七天金额、并发和 RPM 按 Client Key 跨账号、跨 Provider 合计，零表示不限；修改限额不重置已用金额。 -Core 负责准入与结算时序，Store 持久化费用账本，Admin 负责限额配置及带审计的费用核账。 +Core 负责准入与结算时序，Store 持久化费用账本，Admin 负责限额配置。 - 日窗口按北京时间零点划分；七天窗口从首次准入当天零点开始，到期后由下一次使用重新开启，不固定为周一。 - 金额优先使用 Provider 上报的 USD，否则按现有模型价格估算；订阅账号的估算费用不代表上游订阅账单。 - 费用归属请求完成时间，跨日请求计入完成日，延迟写入或核账仍保留原完成时间。 + 费用归属请求完成时间，跨日请求计入完成日，延迟写入仍保留原完成时间。 - 准入检查已结算金额，不预占未来费用。达到任一金额阈值后拒绝新请求，已准入请求仍可完成并使总额超过阈值。 - SSE 持有并发名额直到终态；WebSocket 每个 `response.create` 独立准入，空闲连接不占名额。 Core 在每次请求开始时重新鉴权并冻结当前策略，既有连接也应用已发布的限额与授权范围变更。 换号和内部重试复用同一名额，Key 与账号并发上限同时生效；预算拒绝或启动失败释放名额。 -每次计费请求在发送上游前持久化待结算记录。完成、失败、取消和断连使用同一结算路径，以网关请求 ID -保证幂等；明确未发送的失败按零结算，已发送或可能已发送但没有可信费用的请求保留未知状态。 -未知费用或超过请求期限仍未结算的记录阻止受限 Key 接受新的计费请求，需管理员确认金额并填写核账原因； -零金额也必须显式确认。限额或核账阻断不会修改 Key 的管理员启停状态。 +完成、失败、取消和断连使用同一结算路径，在结束时以网关请求 ID 幂等累计已取得费用。 +缺少用量或价格、无法取得 USD 费用的尝试按零累计，不创建待核账记录，也不会阻断 Key； +内部重试中已经取得的费用仍须累计。发送状态只决定重放是否安全，不能据此推断费用。 +请求日志继续保留真实错误、用量与费用来源，账本按零累计不表示上游实际免费。 -结算写入失败时，进程内保留精确费用，在同一 Key 下次请求前重试；进程崩溃后依靠持久化待结算记录要求核账， -不能把遗失费用按零放行。PostgreSQL 不可用时拒绝所有新的计费请求，Redis 继续管理并发/RPM 租约。 +结算写入失败时，进程内保留精确费用，在同一 Key 下次请求前重试；进程退出后无法恢复的费用不会形成欠账。 +PostgreSQL 不可用时拒绝所有新的计费请求，Redis 继续管理并发/RPM 租约。 账本独立于可丢弃的请求观测日志，日志清理不重置金额；费用事件保留至删除 Key，已有日志不会回填为账本费用。 字段与错误合同见 [Client Key API](api.md#7-client-key)。 diff --git a/frontend/src/api/modules/api-keys.ts b/frontend/src/api/modules/api-keys.ts index 43c782412..60e55d98f 100644 --- a/frontend/src/api/modules/api-keys.ts +++ b/frontend/src/api/modules/api-keys.ts @@ -17,7 +17,6 @@ export interface ApiKey { weeklyUsedUsd: string dailyResetsAt: string | null weeklyResetsAt: string | null - unresolvedRequests: number createdAt: string updatedAt: string lastUsedAt: string | null @@ -47,21 +46,6 @@ export interface ApiKeyMutationResponse { id: string } -export interface UnresolvedClientCharge { - requestId: string - startedAt: string - completedAt: string | null - state: string -} - -export function getUnresolvedClientCharges(id: string) { - return request<{ items: UnresolvedClientCharge[] }>({ url: '/api/admin/client-keys/unresolved-charges', method: 'GET', params: { id } }) -} - -export function reconcileClientCharge(data: { id: string, requestId: string, amountUsd: string, reason: string }) { - return request<{ reconciled: boolean }>({ url: '/api/admin/client-keys/reconcile-charge', method: 'POST', data }) -} - // 请求参数类型：仅定义 API 边界的形状，调用方不依赖显式声明。 interface ApiKeyListParams { cursor?: string diff --git a/frontend/src/views/api-keys/components/ApiKeyBudgetCell.vue b/frontend/src/views/api-keys/components/ApiKeyBudgetCell.vue index 69107a6d1..df6798ede 100644 --- a/frontend/src/views/api-keys/components/ApiKeyBudgetCell.vue +++ b/frontend/src/views/api-keys/components/ApiKeyBudgetCell.vue @@ -1,11 +1,8 @@ @@ -26,11 +24,5 @@ function amount(value: string) { ${{ amount(window.used) }} / {{ Number(window.limit) === 0 ? '∞' : `$${amount(window.limit)}` }} -
- 待核账 {{ apiKey.unresolvedRequests }} 笔 - - - -
diff --git a/frontend/src/views/api-keys/components/ApiKeyReconcileModal.vue b/frontend/src/views/api-keys/components/ApiKeyReconcileModal.vue deleted file mode 100644 index 1acae34a5..000000000 --- a/frontend/src/views/api-keys/components/ApiKeyReconcileModal.vue +++ /dev/null @@ -1,124 +0,0 @@ - - - diff --git a/frontend/src/views/api-keys/index.vue b/frontend/src/views/api-keys/index.vue index 1ef082366..ffddfcfac 100644 --- a/frontend/src/views/api-keys/index.vue +++ b/frontend/src/views/api-keys/index.vue @@ -1,5 +1,4 @@