From 239def9a5b0c60080f037ee6dadc19cf86cefb5d Mon Sep 17 00:00:00 2001 From: Zyy <0x@ainz.cc> Date: Wed, 4 Nov 2026 17:00:44 +0800 Subject: [PATCH] fix(store): preserve client key ledgers after deletion (#602) Co-authored-by: Zyy Co-authored-by: Codex --- .../src/postgres/client_budgets.rs | 29 +++-- .../tests/postgres/client_budgets.rs | 100 ++++++++++++------ backend/migrations/.frozen-sha256 | 1 + .../0038_preserve_client_key_ledger.sql | 5 + docs/api.md | 4 +- docs/architecture.md | 3 +- 6 files changed, 94 insertions(+), 48 deletions(-) create mode 100644 backend/migrations/0038_preserve_client_key_ledger.sql diff --git a/backend/crates/gateway-store/src/postgres/client_budgets.rs b/backend/crates/gateway-store/src/postgres/client_budgets.rs index 775d6eaea..867a936bb 100644 --- a/backend/crates/gateway-store/src/postgres/client_budgets.rs +++ b/backend/crates/gateway-store/src/postgres/client_budgets.rs @@ -297,16 +297,7 @@ impl PgClientBudgetStore { .begin() .await .map_err(|source| ClientBudgetError(Some(source.into())))?; - // 与窗口开启和重置统一先锁 Key，再写窗口和费用 - let key = sqlx::query_scalar::<_, String>( - "select id from client_api_keys where id = $1 for update", - ) - .bind(first.key_id.as_str()) - .fetch_optional(&mut *tx) - .await - .map_err(|source| ClientBudgetError(Some(source.into())))?; - let Some(key) = key else { return Ok(()) }; // 删除 Key 时也会删除其费用记录 - settle_in_transaction(&mut tx, &key, charges, self.timezone) + settle_in_transaction(&mut tx, first.key_id.as_str(), charges, self.timezone) .await .map_err(|source| ClientBudgetError(Some(source.into())))?; tx.commit() @@ -321,6 +312,22 @@ async fn settle_in_transaction( charges: &[ClientBudgetCharge], timezone: gateway_core::time::DeploymentTimeZone, ) -> Result<(), sqlx::Error> { + // 凭据存在时与准入和重置共用 Key 行锁；删除后由准入已开启的窗口串行化迟到结算 + let active = sqlx::query("select 1 from client_api_keys where id = $1 for update") + .bind(key) + .fetch_optional(&mut **tx) + .await?; + if active.is_none() + && sqlx::query( + "select 1 from client_api_key_budget_windows where client_api_key_id = $1 for update", + ) + .bind(key) + .fetch_optional(&mut **tx) + .await? + .is_none() + { + return Ok(()); + } let ids = charges .iter() .map(|charge| charge.request_id.as_str()) @@ -397,7 +404,7 @@ async fn advance_windows( let weekly_end = timezone .days_after(now, 7) .ok_or_else(|| sqlx::Error::Protocol("invalid budget week end".to_owned()))?; - // 调用方持有 Key 行锁；未滚动的窗口无需产生新元组与 WAL + // 未滚动的窗口无需产生新元组与 WAL sqlx::query("insert into client_api_key_budget_windows (client_api_key_id, daily_start, daily_end, weekly_start, weekly_end) select $1, $4, $5, $4, $6 diff --git a/backend/crates/gateway-store/tests/postgres/client_budgets.rs b/backend/crates/gateway-store/tests/postgres/client_budgets.rs index 68849f845..4e9a03de5 100644 --- a/backend/crates/gateway-store/tests/postgres/client_budgets.rs +++ b/backend/crates/gateway-store/tests/postgres/client_budgets.rs @@ -7,8 +7,8 @@ use gateway_admin::{ model::{ MutationActor, MutationContext, Revision, client_keys::{ - ClientKeyBudgetMutationOrigin, ClientKeyBudgetPeriod, ResetClientKeyBudget, - UpdateClientKey, UpdateClientKeyBudgetLimits, + ClientKeyBudgetMutationOrigin, ClientKeyBudgetPeriod, DeleteClientKey, NewClientKey, + ResetClientKeyBudget, UpdateClientKey, UpdateClientKeyBudgetLimits, }, plugin_resources::PluginResourceOwner, plugins::{PluginSource, instances::PluginInstance}, @@ -129,51 +129,83 @@ async fn open_budget_windows_only_write_for_new_charges() { } #[tokio::test] -async fn indexed_charge_cleanup_keeps_other_keys_and_rolls_back_with_key_deletion() { - let Some(database) = TestDatabase::create("budget_key_cleanup").await else { +async fn deleting_a_key_preserves_its_ledger_and_settles_inflight_requests_once() { + let Some(database) = TestDatabase::create("deleted_key_ledger").await else { return; }; - let store = PgClientBudgetStore::new(database.pool.clone()); - for key in ["deleted", "retained"] { - seed(&database, key, "0", "0").await; - store.settle(charge(key, key, "1.25")).await.unwrap(); - } - let mut tx = database.pool.begin().await.unwrap(); - sqlx::query("delete from client_api_keys where id = 'deleted'") - .execute(&mut *tx) + let budgets = PgClientBudgetStore::new(database.pool.clone()); + let keys = PgAdminClientKeyStore::new(database.pool.clone()); + seed(&database, "original", "0", "0").await; + budgets.admit(key_id("original")).await.unwrap(); + budgets + .settle(charge("original", "before", "1.25")) .await .unwrap(); - let remaining: Vec = sqlx::query_scalar( - "select client_api_key_id from client_api_key_charge_events order by client_api_key_id", + budgets.admit(key_id("original")).await.unwrap(); + keys.delete_client_key( + DeleteClientKey { + id: key_id("original"), + }, + &context(), ) - .fetch_all(&mut *tx) .await .unwrap(); - assert_eq!(remaining, ["retained"]); - tx.rollback().await.unwrap(); - let restored: i64 = sqlx::query_scalar("select count(*) from client_api_key_charge_events") - .fetch_one(&database.pool) - .await - .unwrap(); - assert_eq!(restored, 2); - sqlx::query("delete from client_api_keys where id = 'deleted'") - .execute(&database.pool) - .await - .unwrap(); - assert_eq!( - status(&database, "retained") + assert!( + keys.reveal_client_key(&key_id("original")) .await - .daily_used_usd - .canonical(), - "1.25" + .unwrap() + .is_none() ); - let remaining: Vec = sqlx::query_scalar( - "select client_api_key_id from client_api_key_charge_events order by client_api_key_id", + assert_eq!( + budgets.admit(key_id("original")).await.unwrap_err().kind(), + GatewayErrorKind::Unauthorized + ); + + keys.create_client_key( + NewClientKey { + id: key_id("replacement"), + name: "original".into(), + label: None, + group_ids: Vec::new(), + peer_sharing: Default::default(), + request_profile_overrides: Default::default(), + limits: RateLimits::unlimited(), + budget: Default::default(), + plaintext: "synthetic-replacement-key".into(), + }, + &context(), + ) + .await + .unwrap(); + let late = charge("original", "late", "2.50"); + let (first, retry) = tokio::join!(budgets.settle(late.clone()), budgets.settle(late)); + first.unwrap(); + retry.unwrap(); + let events: Vec<(String, String)> = sqlx::query_as( + "select request_id, trim_scale(amount_usd)::text from client_api_key_charge_events + where client_api_key_id='original' order by request_id", ) .fetch_all(&database.pool) .await .unwrap(); - assert_eq!(remaining, ["retained"]); + assert_eq!( + events, + [ + ("req_before".into(), "1.25".into()), + ("req_late".into(), "2.5".into()) + ] + ); + let totals: (String, String) = sqlx::query_as( + "select trim_scale(daily_used_usd)::text, trim_scale(weekly_used_usd)::text + from client_api_key_budget_windows where client_api_key_id='original'", + ) + .fetch_one(&database.pool) + .await + .unwrap(); + assert_eq!(totals, ("3.75".into(), "3.75".into())); + let replacement = status(&database, "replacement").await; + assert_eq!(replacement.daily_used_usd.canonical(), "0"); + assert_eq!(replacement.weekly_used_usd.canonical(), "0"); database.close().await; } diff --git a/backend/migrations/.frozen-sha256 b/backend/migrations/.frozen-sha256 index 1b9910757..40dd3d391 100644 --- a/backend/migrations/.frozen-sha256 +++ b/backend/migrations/.frozen-sha256 @@ -35,3 +35,4 @@ a0e409402611dbc76bb71393b8eda7c717756c18974dbabca305fdcb4b115f31 0032_usage_and bd73379456839e8a12d7591ad2a839901af0f5ba27871248569abdfa26bc83ef 0035_cpr_link_and_concurrency_defaults.sql d05442ee4cb85db8131e015dbed9c6ad988f93306abcda34647317674857ee44 0036_account_display_names.sql b215ea90d5c8753aeb8192b17ffe65f7a95c1d620aaeaf9fc4337d8cf6f2634e 0037_retry_after_wait_limit.sql +321f4a09367717eb13c3cf5dbfd1f4b014e31c615a74bf6e3b09b8e810d87198 0038_preserve_client_key_ledger.sql diff --git a/backend/migrations/0038_preserve_client_key_ledger.sql b/backend/migrations/0038_preserve_client_key_ledger.sql new file mode 100644 index 000000000..f442f7a90 --- /dev/null +++ b/backend/migrations/0038_preserve_client_key_ledger.sql @@ -0,0 +1,5 @@ +-- 预算窗口和费用事件归属请求使用的 Key ID，凭据删除后仍可结算在途请求 +alter table client_api_key_budget_windows + drop constraint client_api_key_budget_windows_client_api_key_id_fkey; +alter table client_api_key_charge_events + drop constraint client_api_key_charge_events_client_api_key_id_fkey; diff --git a/docs/api.md b/docs/api.md index 8c360daff..a5d3d382a 100644 --- a/docs/api.md +++ b/docs/api.md @@ -1898,8 +1898,8 @@ HTTP 返回 `429`，`error.code` 为 `key_daily_budget_exceeded` 或 `key_weekly 预算检查超时、存储不可用或待结算队列无法接受新请求时返回 `503`、`key_budget_unavailable` 自动结算按网关请求 ID 幂等执行，响应交付不等待费用落库，额度查询中的已用金额在后台写入后可见。 -账本独立于使用统计日志，记录保留至删除 Key， -不受 `usageRetentionDays` 影响 +删除 Key 会撤销访问，保留请求记录和原 Key 的费用账本。删除前已准入的请求仍按原 Key ID 结算， +同名新 Key 单独记账。请求记录按 `usageRetentionDays` 清理，费用账本不受该设置影响 ## 8. 运行设置 diff --git a/docs/architecture.md b/docs/architecture.md index 783a534f0..f91700d86 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -890,7 +890,8 @@ Admin 的手动重置复用同一账本与 Key 行锁，在一个事务中清零 仅首次使用或窗口到期时开启窗口，仍与重置和结算使用相同锁序。预算准入总等待上限为一秒。 待结算费用达到 4096 笔后拒绝新准入，已准入请求的费用继续接收；缓存上限因此还包含在途请求的收尾费用。 进程退出后无法恢复的费用不会形成欠账。PostgreSQL 不可用或费用队列不可用时拒绝新的计费请求，Redis 继续管理并发/RPM 租约。 -账本独立于可丢弃的请求观测日志，日志清理不重置金额；费用事件保留至删除 Key，已有日志不会回填为账本费用。 +预算窗口与费用事件按请求准入时的 Key ID 保存，独立于凭据和请求观测日志。删除 Key 后，在途结算锁定已开启的预算窗口，沿用请求 ID 去重。 +日志清理不重置金额，已有日志不会回填为账本费用。 字段与错误合同见 [Client Key API](api.md#7-client-key) ### 用量维护服务