starred/codex-proxy-rs
1
0
Fork 0
mirror of https://github.com/zyycn/codex-proxy-rs.git synced 2026-11-06 08:06:19 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(store): preserve client key ledgers after deletion (#602)

Browse Source 
Co-authored-by: Zyy <zyycn@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Zyy
2026-11-04 17:00:44 +08:00
committed by GitHub
co-authored by Zyy Codex
parent 5eaca14a64
commit 239def9a5b
6 changed files with 94 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/crates/gateway-store/src/postgres/client_budgets.rs
+18 -11
View File
@@ -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
backend/crates/gateway-store/tests/postgres/client_budgets.rs
+66 -34
View File
@@ -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<String> = 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<String> = 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;
}
backend/migrations/.frozen-sha256
+1
View File
@@ -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
backend/migrations/0038_preserve_client_key_ledger.sql
+5
View File
@@ -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;