fix(keys): restart budget windows on use after reset (#335)

Co-authored-by: Codex <noreply@openai.com>

backend/crates/gateway-plugin/sdk/docs/capabilities.md

@@ -178,7 +178,7 @@ Key 明文仍通过宿主管理面查看，插件模型调用使用返回的 Key
 相同值再次赋值不产生新 revision 或审计；并发更新按事务顺序生效，同一字段由后提交的值覆盖，接口不提供调用去重或比较交换。
 将非零上限降低到已用金额及以下会拒绝后续准入，提高限额可恢复准入，在途请求仍按原生规则完成和结算
 
-重置只清零所选周期，保留额度上限、密钥、窗口到期时间和历史费用事件；未使用的 Key 不因此开启窗口。
+重置清零所选周期并清除到期时间，保留额度上限、密钥和历史费用事件；所选窗口下次使用时按原生规则重新开启。
 费用按完成时间归属，重置前完成但延迟落盘的费用不会重新扣入已重置周期，重置后完成的在途请求仍正常扣额。
 实例版本复验、清零和审计在同一事务提交，不推进配置 revision。每次调用均执行新重置，可能清掉两次调用之间的新消费
 

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::<Utc>::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::<Utc>::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::<Utc>::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<Utc>,
+    used_at: DateTime<Utc>,
 ) -> 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(())
 }
 

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),

docs/api.md

@@ -1128,8 +1128,8 @@ HTTP 请求头及新建 WS 的握手提示按当时的最终出站档位构造
 `dailyResetsAt`、`weeklyResetsAt`（RFC3339 或 `null`）。
 记账和限额比较保留完整精度。
 日窗口按北京时间零点重置；周窗口从首次准入当天零点起持续七天，到期后在下一次使用时重新开启。
-手动重置仅清零所选周期的已用金额，保留限额上限、原到期时间和历史费用，返回 `{ id }`。
-未使用或已过期的窗口不会因手动重置而重新开启。重置前完成但延迟结算的费用不再计入所选周期；
+手动重置清零所选周期的已用金额并清除到期时间，保留限额上限和历史费用，返回 `{ id }`。
+所选窗口的重置时间返回 `null`，下次使用时按新建 Key 的规则重新开启。重置前完成但延迟结算的费用不再计入所选周期，也不会开启窗口；
 重置后完成的请求继续计费，包括重置时仍在进行的请求。操作保留管理员审计，不改变账号上游额度。
 费用按请求完成时间归属窗口。并发按同一 Key 的执行中请求累计，包含 SSE 与每个 WebSocket
 `response.create`；空闲连接不占名额，内部重试不重复占用。

docs/architecture.md

@@ -556,7 +556,8 @@ Key 的 RPM 在成功准入时才计数，金额限制在入队前及成功准
 日金额、七天金额、并发和 RPM 按 Client Key 跨账号、跨 Provider 合计，零表示不限；修改限额不重置已用金额。
 Core 负责准入与结算时序，Store 持久化费用账本，Admin 负责限额配置。
 Admin 的手动重置复用同一账本与 Key 行锁，在一个事务中清零所选周期金额、推进计费起点并写入审计，
-保留窗口到期时间与费用事件，不推进配置 revision。结算仍按完成时间判断归属，重置前完成的费用不会重新扣入已重置周期。
+关闭所选窗口并保留费用事件，不推进配置 revision，下次使用时重新开启窗口。
+结算仍按完成时间判断归属，重置前完成的费用不会重新扣入已重置周期或开启窗口。
 插件预算回调通过 `PluginClientKeyAccess` 进入同一 `ClientKeyService` 和 `ClientKeyStore`，存储事务先复核插件实例版本，
 再执行原生预算操作；不建立另一份预算状态。查询复用账本投影，不触发准入。局部更新金额上限不改用量或窗口，
 仅在值变化时提交配置 revision、审计并发布；重置仍不推进配置版本。重复调用及结果未知的合同见

frontend/src/views/keys/components/ApiKeyBudgetResetModal.vue

@@ -46,7 +46,7 @@ async function resetBudget() {
   <BaseModal
     v-model="open"
     title="重置已用额度"
-    description="清零所选周期的已用金额，原重置时间保持不变"
+    description="清零所选周期的已用金额，下次使用时重新确定重置日期"
     size="sm"
     :dismissible="!loading"
   >