From d97c6e064b3714fa496f575aa6d7d64fb405601e Mon Sep 17 00:00:00 2001 From: Zyy <39424122+zyycn@users.noreply.github.com> Date: Wed, 30 Sep 2026 17:54:47 +0800 Subject: [PATCH] fix(keys): restart budget windows on use after reset (#335) Co-authored-by: Codex --- .../gateway-plugin/sdk/docs/capabilities.md | 2 +- .../src/postgres/client_budgets.rs | 45 +++--- .../tests/postgres/client_budgets.rs | 136 +++++++++++++++++- docs/api.md | 4 +- docs/architecture.md | 3 +- .../components/ApiKeyBudgetResetModal.vue | 2 +- 6 files changed, 165 insertions(+), 27 deletions(-) diff --git a/backend/crates/gateway-plugin/sdk/docs/capabilities.md b/backend/crates/gateway-plugin/sdk/docs/capabilities.md index 0a899bbfa..58e9117ad 100644 --- a/backend/crates/gateway-plugin/sdk/docs/capabilities.md +++ b/backend/crates/gateway-plugin/sdk/docs/capabilities.md @@ -178,7 +178,7 @@ Key 明文仍通过宿主管理面查看，插件模型调用使用返回的 Key 相同值再次赋值不产生新 revision 或审计；并发更新按事务顺序生效，同一字段由后提交的值覆盖，接口不提供调用去重或比较交换。 将非零上限降低到已用金额及以下会拒绝后续准入，提高限额可恢复准入，在途请求仍按原生规则完成和结算 -重置只清零所选周期，保留额度上限、密钥、窗口到期时间和历史费用事件；未使用的 Key 不因此开启窗口。 +重置清零所选周期并清除到期时间，保留额度上限、密钥和历史费用事件；所选窗口下次使用时按原生规则重新开启。 费用按完成时间归属，重置前完成但延迟落盘的费用不会重新扣入已重置周期，重置后完成的在途请求仍正常扣额。 实例版本复验、清零和审计在同一事务提交，不推进配置 revision。每次调用均执行新重置，可能清掉两次调用之间的新消费 diff --git a/backend/crates/gateway-store/src/postgres/client_budgets.rs b/backend/crates/gateway-store/src/postgres/client_budgets.rs index f2bd70b69..1d2aca365 100644 --- a/backend/crates/gateway-store/src/postgres/client_budgets.rs +++ b/backend/crates/gateway-store/src/postgres/client_budgets.rs @@ -78,13 +78,15 @@ async fn reset_client_key_budget_in_transaction( ClientKeyBudgetPeriod::Weekly | ClientKeyBudgetPeriod::All ); let reset_at = Utc::now(); - // 推进计费起点，避免重置前完成、稍后落盘的费用重新扣额；未使用的 Key 不开启窗口。 + // 起止时间收拢到重置边界，窗口保持未开启，同时排除重置前完成的迟到费用。 sqlx::query( "update client_key_budget_windows set daily_used_usd = case when $2 then 0 else daily_used_usd end, - daily_start = case when $2 and daily_end > $4 then $4 else daily_start end, + daily_start = case when $2 then $4 else daily_start end, + daily_end = case when $2 then $4 else daily_end end, weekly_used_usd = case when $3 then 0 else weekly_used_usd end, - weekly_start = case when $3 and weekly_end > $4 then $4 else weekly_start end + weekly_start = case when $3 then $4 else weekly_start end, + weekly_end = case when $3 then $4 else weekly_end end where client_api_key_id = $1", ) .bind(command.id.as_str()) @@ -96,10 +98,18 @@ async fn reset_client_key_budget_in_transaction( .map_err(|_| postgres_unavailable("reset client budget"))?; let mut fields = Vec::new(); if daily { - fields.extend(["daily_used_usd".to_owned(), "daily_start".to_owned()]); + fields.extend([ + "daily_used_usd".to_owned(), + "daily_start".to_owned(), + "daily_end".to_owned(), + ]); } if weekly { - fields.extend(["weekly_used_usd".to_owned(), "weekly_start".to_owned()]); + fields.extend([ + "weekly_used_usd".to_owned(), + "weekly_start".to_owned(), + "weekly_end".to_owned(), + ]); } super::append_admin_audit_event_in_transaction( tx, @@ -174,7 +184,7 @@ impl PgClientBudgetStore { .map_err(|_| unavailable())?, }; let now = Utc::now(); - advance_windows(&mut tx, key_id.as_str(), now) + advance_windows(&mut tx, key_id.as_str(), now, now) .await .map_err(|_| unavailable())?; if limits.is_limited() { @@ -243,7 +253,7 @@ async fn settle_in_transaction( key: &str, charge: &ClientBudgetCharge, ) -> Result<(), sqlx::Error> { - advance_windows(tx, key, Utc::now()).await?; + let completed_at = DateTime::::from(charge.completed_at); // 仅在请求结束时写入费用；请求 ID 冲突时不重复累计。 let changed = sqlx::query( "insert into client_key_charge_events (request_id, client_api_key_id, amount_usd, completed_at) @@ -253,16 +263,17 @@ async fn settle_in_transaction( .bind(charge.request_id.as_str()) .bind(key) .bind(charge.amount_usd.canonical()) - .bind(DateTime::::from(charge.completed_at)) + .bind(completed_at) .execute(&mut **tx) .await? .rows_affected(); if changed == 1 { + advance_windows(tx, key, Utc::now(), completed_at).await?; 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(charge.amount_usd.canonical()).bind(DateTime::::from(charge.completed_at)) + .bind(key).bind(charge.amount_usd.canonical()).bind(completed_at) .execute(&mut **tx).await?; } Ok(()) @@ -291,19 +302,21 @@ async fn advance_windows( tx: &mut Transaction<'_, Postgres>, key: &str, now: DateTime, + used_at: DateTime, ) -> Result<(), sqlx::Error> { + // 到期日仍按北京时间零点计算，计费起点不能退回重置前；迟到费用不启动新窗口。 sqlx::query("insert into client_key_budget_windows (client_api_key_id, daily_start, daily_end, weekly_start, weekly_end) select $1, day, day + interval '24 hours', day, day + interval '168 hours' from (select date_trunc('day', $2::timestamptz at time zone 'Asia/Shanghai') at time zone 'Asia/Shanghai' as day) d on conflict (client_api_key_id) do update set - daily_start = case when client_key_budget_windows.daily_end <= $2 then excluded.daily_start else client_key_budget_windows.daily_start end, - daily_end = case when client_key_budget_windows.daily_end <= $2 then excluded.daily_end else client_key_budget_windows.daily_end end, - daily_used_usd = case when client_key_budget_windows.daily_end <= $2 then 0 else client_key_budget_windows.daily_used_usd end, - weekly_start = case when client_key_budget_windows.weekly_end <= $2 then excluded.weekly_start else client_key_budget_windows.weekly_start end, - weekly_end = case when client_key_budget_windows.weekly_end <= $2 then excluded.weekly_end else client_key_budget_windows.weekly_end end, - weekly_used_usd = case when client_key_budget_windows.weekly_end <= $2 then 0 else client_key_budget_windows.weekly_used_usd end") - .bind(key).bind(now).execute(&mut **tx).await?; + daily_start = case when client_key_budget_windows.daily_end <= $2 and client_key_budget_windows.daily_start <= $3 then greatest(excluded.daily_start, client_key_budget_windows.daily_start) else client_key_budget_windows.daily_start end, + daily_end = case when client_key_budget_windows.daily_end <= $2 and client_key_budget_windows.daily_start <= $3 then excluded.daily_end else client_key_budget_windows.daily_end end, + daily_used_usd = case when client_key_budget_windows.daily_end <= $2 and client_key_budget_windows.daily_start <= $3 then 0 else client_key_budget_windows.daily_used_usd end, + weekly_start = case when client_key_budget_windows.weekly_end <= $2 and client_key_budget_windows.weekly_start <= $3 then greatest(excluded.weekly_start, client_key_budget_windows.weekly_start) else client_key_budget_windows.weekly_start end, + weekly_end = case when client_key_budget_windows.weekly_end <= $2 and client_key_budget_windows.weekly_start <= $3 then excluded.weekly_end else client_key_budget_windows.weekly_end end, + weekly_used_usd = case when client_key_budget_windows.weekly_end <= $2 and client_key_budget_windows.weekly_start <= $3 then 0 else client_key_budget_windows.weekly_used_usd end") + .bind(key).bind(now).bind(used_at).execute(&mut **tx).await?; Ok(()) } diff --git a/backend/crates/gateway-store/tests/postgres/client_budgets.rs b/backend/crates/gateway-store/tests/postgres/client_budgets.rs index 771eddaad..be8b2c07f 100644 --- a/backend/crates/gateway-store/tests/postgres/client_budgets.rs +++ b/backend/crates/gateway-store/tests/postgres/client_budgets.rs @@ -75,7 +75,7 @@ fn context() -> MutationContext { } #[tokio::test] -async fn manual_reset_clears_only_selected_budget_and_preserves_policy_history_and_expiry() { +async fn manual_reset_clears_selected_windows_and_preserves_policy_and_history() { let Some(database) = TestDatabase::create("budget_manual_reset").await else { return; }; @@ -108,14 +108,25 @@ async fn manual_reset_clears_only_selected_budget_and_preserves_policy_history_a ) .await .unwrap(); - // 同一费用重试仍由事件 ID 去重，不因清零而被再次累计。 - store.settle(billed).await.unwrap(); let after = status(&database, key).await; assert_eq!(after.daily_used_usd.canonical(), daily); assert_eq!(after.weekly_used_usd.canonical(), weekly); assert_eq!(after.limits, before.limits); - assert_eq!(after.daily_resets_at, before.daily_resets_at); - assert_eq!(after.weekly_resets_at, before.weekly_resets_at); + assert_eq!( + after.daily_resets_at, + (period == ClientKeyBudgetPeriod::Weekly) + .then_some(before.daily_resets_at) + .flatten() + ); + assert_eq!( + after.weekly_resets_at, + (period == ClientKeyBudgetPeriod::Daily) + .then_some(before.weekly_resets_at) + .flatten() + ); + // 同一费用重试既不重复扣额，也不能重新开启已重置的窗口。 + store.settle(billed).await.unwrap(); + assert_eq!(status(&database, key).await, after); assert_eq!( store.admit(key_id(key)).await.is_ok(), period == ClientKeyBudgetPeriod::All @@ -130,6 +141,118 @@ async fn manual_reset_clears_only_selected_budget_and_preserves_policy_history_a database.close().await; } +#[tokio::test] +async fn manual_reset_restarts_selected_windows_only_after_new_usage() { + let Some(database) = TestDatabase::create("budget_reset_restart").await else { + return; + }; + let store = PgClientBudgetStore::new(database.pool.clone()); + let admin = PgAdminClientKeyStore::new(database.pool.clone()); + for (key, period) in [ + ("daily", ClientKeyBudgetPeriod::Daily), + ("weekly", ClientKeyBudgetPeriod::Weekly), + ("all", ClientKeyBudgetPeriod::All), + ] { + seed(&database, key, "10", "20").await; + store.admit(key_id(key)).await.unwrap(); + store.settle(charge(key, key, "1")).await.unwrap(); + sqlx::query("update client_key_budget_windows set weekly_start = weekly_start - interval '3 days', weekly_end = weekly_end - interval '3 days' where client_api_key_id = $1") + .bind(key).execute(&database.pool).await.unwrap(); + let before = status(&database, key).await; + let delayed = charge(key, &format!("{key}-delayed"), "0.5"); + let old = charge(key, &format!("{key}-old"), "0.7"); + admin + .reset_client_key_budget( + ResetClientKeyBudget { + id: key_id(key), + period, + }, + ClientKeyBudgetMutationOrigin::Admin, + &context(), + ) + .await + .unwrap(); + let reset = status(&database, key).await; + // 重置前完成的迟到费用保留历史，但不能开启窗口或重新扣入所选周期。 + store.settle(delayed).await.unwrap(); + let delayed = status(&database, key).await; + assert_eq!(delayed.daily_resets_at, reset.daily_resets_at); + assert_eq!(delayed.weekly_resets_at, reset.weekly_resets_at); + assert_eq!( + delayed.daily_used_usd.canonical(), + if period == ClientKeyBudgetPeriod::Weekly { + "1.5" + } else { + "0" + } + ); + assert_eq!( + delayed.weekly_used_usd.canonical(), + if period == ClientKeyBudgetPeriod::Daily { + "1.5" + } else { + "0" + } + ); + + store.admit(key_id(key)).await.unwrap(); + let restarted = status(&database, key).await; + assert_eq!(restarted.daily_resets_at, before.daily_resets_at); + if period == ClientKeyBudgetPeriod::Daily { + assert_eq!(restarted.weekly_resets_at, before.weekly_resets_at); + } else { + assert_eq!( + restarted.weekly_resets_at.unwrap(), + restarted.daily_resets_at.unwrap() + std::time::Duration::from_secs(6 * 86400) + ); + assert_ne!(restarted.weekly_resets_at, before.weekly_resets_at); + } + // 开启窗口后仍保留重置边界，另一笔重置前的迟到费用也不能重新扣额。 + store.settle(old).await.unwrap(); + let after_old = status(&database, key).await; + assert_eq!(after_old.daily_resets_at, restarted.daily_resets_at); + assert_eq!(after_old.weekly_resets_at, restarted.weekly_resets_at); + assert_eq!( + after_old.daily_used_usd.canonical(), + if period == ClientKeyBudgetPeriod::Weekly { + "2.2" + } else { + "0" + } + ); + assert_eq!( + after_old.weekly_used_usd.canonical(), + if period == ClientKeyBudgetPeriod::Daily { + "2.2" + } else { + "0" + } + ); + store + .settle(charge(key, &format!("{key}-new"), "0.25")) + .await + .unwrap(); + let used = status(&database, key).await; + assert_eq!( + used.daily_used_usd.canonical(), + if period == ClientKeyBudgetPeriod::Weekly { + "2.45" + } else { + "0.25" + } + ); + assert_eq!( + used.weekly_used_usd.canonical(), + if period == ClientKeyBudgetPeriod::Daily { + "2.45" + } else { + "0.25" + } + ); + } + database.close().await; +} + #[tokio::test] async fn manual_reset_excludes_old_completions_but_counts_inflight_requests_after_reset() { let Some(database) = TestDatabase::create("budget_reset_late").await else { @@ -648,7 +771,7 @@ async fn plugin_resets_share_the_native_ledger_and_each_call_resets_again() { assert_eq!(after.weekly_used_usd.canonical(), "0"); assert_eq!(after.limits, before.limits); assert_eq!(after.daily_resets_at, before.daily_resets_at); - assert_eq!(after.weekly_resets_at, before.weekly_resets_at); + assert!(after.weekly_resets_at.is_none()); budgets.settle(charge("key", "after", "2")).await.unwrap(); assert_eq!( @@ -666,6 +789,7 @@ async fn plugin_resets_share_the_native_ledger_and_each_call_resets_again() { let after = status(&database, "key").await; assert_eq!(after.daily_used_usd.canonical(), "5"); assert_eq!(after.weekly_used_usd.canonical(), "0"); + assert!(after.weekly_resets_at.is_none()); let (audits, charges, current_revision): (i64, i64, i64) = sqlx::query_as( "select (select count(*) from admin_audit_events where action='reset_budget' and config_revision is null), (select count(*) from client_key_charge_events), diff --git a/docs/api.md b/docs/api.md index 4b5f4bcbd..b86fd2bd0 100644 --- a/docs/api.md +++ b/docs/api.md @@ -1128,8 +1128,8 @@ HTTP 请求头及新建 WS 的握手提示按当时的最终出站档位构造 `dailyResetsAt`、`weeklyResetsAt`（RFC3339 或 `null`）。 记账和限额比较保留完整精度。 日窗口按北京时间零点重置；周窗口从首次准入当天零点起持续七天，到期后在下一次使用时重新开启。 -手动重置仅清零所选周期的已用金额，保留限额上限、原到期时间和历史费用，返回 `{ id }`。 -未使用或已过期的窗口不会因手动重置而重新开启。重置前完成但延迟结算的费用不再计入所选周期； +手动重置清零所选周期的已用金额并清除到期时间，保留限额上限和历史费用，返回 `{ id }`。 +所选窗口的重置时间返回 `null`，下次使用时按新建 Key 的规则重新开启。重置前完成但延迟结算的费用不再计入所选周期，也不会开启窗口； 重置后完成的请求继续计费，包括重置时仍在进行的请求。操作保留管理员审计，不改变账号上游额度。 费用按请求完成时间归属窗口。并发按同一 Key 的执行中请求累计，包含 SSE 与每个 WebSocket `response.create`；空闲连接不占名额，内部重试不重复占用。 diff --git a/docs/architecture.md b/docs/architecture.md index 89a4ad9e9..a4e0f08de 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -556,7 +556,8 @@ Key 的 RPM 在成功准入时才计数，金额限制在入队前及成功准 日金额、七天金额、并发和 RPM 按 Client Key 跨账号、跨 Provider 合计，零表示不限；修改限额不重置已用金额。 Core 负责准入与结算时序，Store 持久化费用账本，Admin 负责限额配置。 Admin 的手动重置复用同一账本与 Key 行锁，在一个事务中清零所选周期金额、推进计费起点并写入审计， -保留窗口到期时间与费用事件，不推进配置 revision。结算仍按完成时间判断归属，重置前完成的费用不会重新扣入已重置周期。 +关闭所选窗口并保留费用事件，不推进配置 revision，下次使用时重新开启窗口。 +结算仍按完成时间判断归属，重置前完成的费用不会重新扣入已重置周期或开启窗口。 插件预算回调通过 `PluginClientKeyAccess` 进入同一 `ClientKeyService` 和 `ClientKeyStore`，存储事务先复核插件实例版本， 再执行原生预算操作；不建立另一份预算状态。查询复用账本投影，不触发准入。局部更新金额上限不改用量或窗口， 仅在值变化时提交配置 revision、审计并发布；重置仍不推进配置版本。重复调用及结果未知的合同见 diff --git a/frontend/src/views/keys/components/ApiKeyBudgetResetModal.vue b/frontend/src/views/keys/components/ApiKeyBudgetResetModal.vue index 4765c1710..eddc96b2c 100644 --- a/frontend/src/views/keys/components/ApiKeyBudgetResetModal.vue +++ b/frontend/src/views/keys/components/ApiKeyBudgetResetModal.vue @@ -46,7 +46,7 @@ async function resetBudget() {