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

feat(store): materialize quota_limit_reached for single-source limit status

Browse Source 
Dashboard reads the materialized column instead of parsing raw quota
JSON (design §3.3/§5.3): Provider computes snapshot-level limit_reached
on quota writes, store persists it, Admin account_status and Dashboard
both consume the same value.

- QuotaObservation gains limit_reached; openai writes fact.exhausted()
  (429 always true), xai writes used_percent>=100
- compare_and_swap_provider_quota updates the column; get_quotas and
  ACCOUNT_SELECT read it back
- ProviderAccountSummary carries quota_limit_reached; admin_account_status
  derives RateLimited from it (Unknown -> Invalid)
- Dashboard provider_account_metrics reads quota_limit_reached instead of
  provider_quota_json->'rate_limit'
- Migration 0003 adds the column with backfill (coalesce for NULL when
  rate_limit key is absent); 0004 fixes NULL backfill for already-applied
  environments

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Zyy
2026-08-03 03:32:17 +08:00
co-authored by Claude Fable 5
parent df83ffdafb
commit d199141a91
21 changed files with 84 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/crates/gateway-core/src/engine/credential.rs
+2
View File
@@ -578,6 +578,8 @@ pub struct QuotaObservation {
pub expected_revision: CredentialRevision,
pub quota: Option<OpaqueProviderData>,
pub observed_at: Option<SystemTime>,
/// 快照级限流事实（Provider 解析后物化）；None 表示不更新物化列。
pub limit_reached: Option<bool>,
}
impl fmt::Debug for QuotaObservation {
backend/crates/gateway-store/src/postgres/observability/queries/dashboard.rs
+2 -8
View File
@@ -348,7 +348,7 @@ pub(crate) async fn provider_account_metrics(
) -> StoreResult<ProviderAccountMetrics> {
let row = sqlx::query(
"with normalized as (
select enabled, availability, provider_quota_json,
select enabled, availability, quota_limit_reached,
case
when not enabled then 'disabled'
when availability = 'banned' then 'banned'
@@ -357,13 +357,7 @@ pub(crate) async fn provider_account_metrics(
when availability = 'expired'
or (access_token_expires_at is not null and access_token_expires_at <= $1)
then 'expired'
when provider_quota_json is not null
and provider_quota_json->'rate_limit'->>'limit_reached' = 'true'
and (
provider_quota_json->'rate_limit'->'primary_window'->>'reset_at' is null
or (provider_quota_json->'rate_limit'->'primary_window'->>'reset_at')::bigint > extract(epoch from $1)::bigint
)
then 'rate_limited'
when quota_limit_reached then 'rate_limited'
when availability = 'unknown' then 'invalid'
else 'active'
end as status
backend/crates/gateway-store/src/postgres/provider_accounts/core_adapter.rs
+7 -1
View File
@@ -163,7 +163,8 @@ impl ProviderAccountStore for PgProviderAccountRepository {
.map(|account| account.as_str().to_owned())
.collect::<Vec<_>>();
let rows = sqlx::query(
"select id, credential_revision, provider_quota_json, quota_observed_at \
"select id, credential_revision, provider_quota_json, quota_observed_at, \
quota_limit_reached \
from provider_accounts \
where id = any($1) and provider_quota_json is not null",
)
@@ -191,12 +192,16 @@ impl ProviderAccountStore for PgProviderAccountRepository {
let observed_at = row
.try_get::<DateTime<Utc>, _>("quota_observed_at")
.map_err(|_| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?;
let limit_reached = row
.try_get::<bool, _>("quota_limit_reached")
.map_err(|_| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?;
Ok(QuotaObservation {
account_id: CoreProviderAccountId::new(account_id)
.map_err(|_| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?,
expected_revision: revision,
quota: Some(quota),
observed_at: Some(observed_at.into()),
limit_reached: Some(limit_reached),
})
})
.collect()
@@ -223,6 +228,7 @@ impl ProviderAccountStore for PgProviderAccountRepository {
Revision::new(observation.expected_revision.get()).map_err(core_store_error)?,
quota,
observation.observed_at.map(DateTime::<Utc>::from),
observation.limit_reached,
)
.await
.map_err(core_store_error)?;
backend/crates/gateway-store/src/postgres/provider_accounts/mapping.rs
+4
View File
@@ -72,6 +72,10 @@ pub(crate) fn admin_account_status(
.is_some_and(|expires_at| expires_at <= now)
{
AdminAccountStatus::Expired
} else if account.quota_limit_reached {
AdminAccountStatus::RateLimited
} else if account.availability == AccountAvailability::Unknown {
AdminAccountStatus::Invalid
} else {
AdminAccountStatus::Active
}
backend/crates/gateway-store/src/postgres/provider_accounts/repository.rs
+5 -1
View File
@@ -27,6 +27,7 @@ pub trait ProviderAccountRepository: Send + Sync {
expected_revision: Revision,
quota: Option<JsonObject>,
observed_at: Option<DateTime<Utc>>,
limit_reached: Option<bool>,
) -> StoreResult<bool>;
async fn delete_provider_account(&self, id: &str) -> StoreResult<bool>;
}
@@ -94,7 +95,7 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
upstream_account_id, plan_type, authentication_kind, credential_revision, has_refresh_token,
access_token_expires_at, next_refresh_at, enabled, availability,
availability_observed_at,
quota_observed_at, created_at, updated_at
quota_observed_at, quota_limit_reached, created_at, updated_at
from provider_accounts
where ($1::text is null or provider_kind = $1) and ($2 or enabled)
order by provider_kind, name, id",
@@ -246,6 +247,7 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
expected_revision: Revision,
quota: Option<JsonObject>,
observed_at: Option<DateTime<Utc>>,
limit_reached: Option<bool>,
) -> StoreResult<bool> {
require_nonempty(ENTITY, "account_id", account_id)?;
if quota.is_some() != observed_at.is_some() {
@@ -257,6 +259,7 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
let result = sqlx::query(
"update provider_accounts
set provider_quota_json = $3, quota_observed_at = $4,
quota_limit_reached = coalesce($5, quota_limit_reached),
updated_at = greatest(now(), coalesce($4, now()))
where id = $1 and credential_revision = $2
and ($4 is null or quota_observed_at is null or quota_observed_at <= $4)",
@@ -265,6 +268,7 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
.bind(to_i64(expected_revision.get())?)
.bind(quota.map(|value| value.as_value()))
.bind(observed_at)
.bind(limit_reached)
.execute(&self.pool)
.await
.map_err(|_| postgres_unavailable("compare and swap provider quota"))?;
backend/crates/gateway-store/src/postgres/provider_accounts/rows.rs
+7 -2
View File
@@ -88,6 +88,8 @@ pub struct ProviderAccountSummary {
pub availability: AccountAvailability,
pub availability_observed_at: DateTime<Utc>,
pub quota_observed_at: Option<DateTime<Utc>>,
/// 快照级限流事实（Provider 物化）。
pub quota_limit_reached: bool,
pub created_at: DateTime<Utc>,
pub updated_at: DateTime<Utc>,
}
@@ -298,6 +300,8 @@ pub struct ProviderAccountObservation {
pub provider_quota_json: Option<JsonObject>,
pub availability_observed_at: DateTime<Utc>,
pub quota_observed_at: Option<DateTime<Utc>>,
/// 快照级限流事实（Provider 物化）。
pub quota_limit_reached: bool,
}
#[derive(Debug, Clone, PartialEq, Eq)]
@@ -331,14 +335,14 @@ impl ProviderAccountObservation {
pub(crate) const ACCOUNT_SELECT: &str = "select id, provider_kind, name, email, upstream_user_id,
upstream_account_id, plan_type, authentication_kind, provider_credentials_json, credential_revision,
has_refresh_token, access_token_expires_at, next_refresh_at, enabled, availability,
provider_quota_json,
provider_quota_json, quota_limit_reached,
availability_observed_at, quota_observed_at, created_at, updated_at
from provider_accounts where id = $1";
pub(crate) const ACCOUNT_SELECT_BY_IDS: &str = "select id, provider_kind, name, email, upstream_user_id,
upstream_account_id, plan_type, authentication_kind, provider_credentials_json, credential_revision,
has_refresh_token, access_token_expires_at, next_refresh_at, enabled, availability,
provider_quota_json,
provider_quota_json, quota_limit_reached,
availability_observed_at, quota_observed_at, created_at, updated_at
from provider_accounts
where id = any($1::text[]) and provider_kind = $2
@@ -440,6 +444,7 @@ pub(crate) fn account_summary_from_row(
availability: parse_availability(&availability)?,
availability_observed_at: get(&row, "availability_observed_at")?,
quota_observed_at: get(&row, "quota_observed_at")?,
quota_limit_reached: get(&row, "quota_limit_reached")?,
created_at: get(&row, "created_at")?,
updated_at: get(&row, "updated_at")?,
})
backend/crates/gateway-store/tests/postgres/mod.rs
+1 -1
View File
@@ -156,7 +156,7 @@ async fn connect_and_migrate_should_apply_all_migrations_once_and_reopen_cleanly
.expect("drop migration test database");
admin.close().await;
assert_eq!((first_table_count, migration_count), (10, 2));
assert_eq!((first_table_count, migration_count), (10, 4));
assert_eq!(response_id_types, ["bytea", "bytea"]);
assert!(!raw_response_id_index_exists);
}
backend/crates/gateway-store/tests/postgres/provider_accounts.rs
+2
View File
@@ -74,6 +74,7 @@ async fn core_quota_batch_reads_only_observed_accounts_in_one_contract_call() {
.clone(),
)),
observed_at: Some(SystemTime::now()),
limit_reached: None,
})
.await
.expect("persist quota");
@@ -140,6 +141,7 @@ async fn delayed_provider_observations_cannot_overwrite_newer_state_or_quota() {
.clone(),
)),
observed_at: Some(observed_at),
limit_reached: None,
};
assert_eq!(
repository
backend/crates/providers/openai/src/credential/quota/mod.rs
+4
View File
@@ -554,6 +554,7 @@ impl CodexCredentialQuotaService {
expected_revision: account.revision(),
quota: Some(OpaqueProviderData::new(object)),
observed_at: Some(observed_at),
limit_reached: Some(fact.exhausted()),
})
.await?;
if outcome == QuotaWriteOutcome::Conflict {
@@ -626,6 +627,7 @@ impl CodexCredentialQuotaService {
expected_revision: account.revision(),
quota: Some(OpaqueProviderData::new(quota)),
observed_at: Some(observed_at),
limit_reached: Some(true),
})
.await?;
if outcome == QuotaWriteOutcome::Conflict {
@@ -728,6 +730,7 @@ impl CodexCredentialQuotaService {
expected_revision: account.revision(),
quota: Some(OpaqueProviderData::new(object)),
observed_at: Some(observed_at),
limit_reached: Some(true),
})
.await?
== QuotaWriteOutcome::Conflict
@@ -870,6 +873,7 @@ impl CodexCredentialQuotaService {
expected_revision: account.revision(),
quota: Some(OpaqueProviderData::new(object)),
observed_at: Some(observed_at),
limit_reached: Some(snapshot.fact().exhausted()),
})
.await?
== QuotaWriteOutcome::Conflict
backend/crates/providers/openai/tests/admin.rs
+2
View File
@@ -471,6 +471,7 @@ async fn openai_admin_provider_projects_codex_additional_limit_as_the_primary_qu
raw.as_object().expect("quota object").clone(),
)),
observed_at: Some(SystemTime::now()),
limit_reached: None,
})
.await
.expect("persist quota");
@@ -537,6 +538,7 @@ async fn openai_admin_projects_confirmed_quota_exhaustion_as_full_without_mutati
.clone(),
)),
observed_at: Some(SystemTime::now()),
limit_reached: None,
})
.await
.expect("persist raw quota");
backend/crates/providers/openai/tests/credential/quota/mod.rs
+4
View File
@@ -358,6 +358,7 @@ async fn passive_rate_limit_headers_replace_stale_secondary_window() {
baseline.as_object().expect("quota object").clone(),
)),
observed_at: Some(SystemTime::now()),
limit_reached: None,
})
.await
.expect("persist baseline quota");
@@ -436,6 +437,7 @@ async fn periodic_quota_synchronization_does_not_scan_stale_ready_accounts() {
)),
// 正常账号即使缓存已旧，也只等待真实请求被动同步，不能被定时任务主动拉取。
observed_at: Some(SystemTime::UNIX_EPOCH),
limit_reached: None,
})
.await
.expect("persist stale quota");
@@ -481,6 +483,7 @@ async fn periodic_quota_synchronization_preserves_availability_when_usage_peaks(
.clone(),
)),
observed_at: Some(SystemTime::now()),
limit_reached: None,
})
.await
.expect("persist exhausted quota");
@@ -718,6 +721,7 @@ async fn periodic_quota_synchronization_does_not_recover_from_percent_drop_befor
.clone(),
)),
observed_at: Some(SystemTime::now()),
limit_reached: None,
})
.await
.expect("seed post-failure quota snapshot");
backend/crates/providers/openai/tests/credential/quota/recovery.rs
+1
View File
@@ -155,6 +155,7 @@ async fn expired_old_reset_with_fresh_allowed_snapshot_recovers_ready() {
.clone(),
)),
observed_at: Some(SystemTime::now()),
limit_reached: None,
})
.await
.expect("seed old exhausted snapshot");
backend/crates/providers/openai/tests/credential/quota/snapshot.rs
+5
View File
@@ -196,6 +196,7 @@ async fn persisted_provider_quota_projects_dynamic_windows_without_network_io()
raw.as_object().expect("quota object").clone(),
)),
observed_at: Some(SystemTime::now()),
limit_reached: None,
})
.await
.expect("persist quota");
@@ -243,6 +244,7 @@ async fn roll_expired_windows_is_o1_for_extremely_old_reset() {
raw.as_object().expect("quota object").clone(),
)),
observed_at: Some(SystemTime::now()),
limit_reached: None,
})
.await
.expect("persist quota");
@@ -284,6 +286,7 @@ async fn scheduling_signal_survives_limit_reached_without_percent_or_reset() {
raw.as_object().expect("quota object").clone(),
)),
observed_at: Some(SystemTime::now()),
limit_reached: None,
})
.await
.expect("persist quota");
@@ -335,6 +338,7 @@ async fn persisted_codex_additional_limit_replaces_the_top_level_rate_limit() {
raw.as_object().expect("quota object").clone(),
)),
observed_at: Some(SystemTime::now()),
limit_reached: None,
})
.await
.expect("persist quota");
@@ -404,6 +408,7 @@ async fn code_review_limit_projects_as_independent_window_with_limit_name() {
raw.as_object().expect("quota object").clone(),
)),
observed_at: Some(SystemTime::now()),
limit_reached: None,
})
.await
.expect("persist quota");
backend/crates/providers/openai/tests/provider/contract.rs
+1
View File
@@ -1293,6 +1293,7 @@ async fn usage_limit_failure_returns_promptly_and_refreshes_the_authoritative_qu
.clone(),
)),
observed_at: Some(SystemTime::now()),
limit_reached: None,
})
.await
.expect("seed stale passive quota");
backend/crates/providers/xai/src/credential/catalog.rs
+3
View File
@@ -451,6 +451,9 @@ impl GrokCredentialQuotaService {
loaded.account.revision(),
document,
observed_at.into(),
presentation
.used_percent()
.is_some_and(|used| used.is_finite() && used >= 100.0),
)
.await
.map_err(map_quota_repository_error)?;
backend/crates/providers/xai/src/credential/repository.rs
+2
View File
@@ -519,6 +519,7 @@ impl GrokCredentialRepository {
expected_revision: CredentialRevision,
document: serde_json::Map<String, Value>,
observed_at: SystemTime,
limit_reached: bool,
) -> Result<(), GrokCredentialRepositoryError> {
let outcome = self
.store
@@ -527,6 +528,7 @@ impl GrokCredentialRepository {
expected_revision,
quota: Some(OpaqueProviderData::new(document)),
observed_at: Some(observed_at),
limit_reached: Some(limit_reached),
})
.await
.map_err(map_store_error)?;
backend/crates/providers/xai/tests/credential/catalog.rs
+1
View File
@@ -919,6 +919,7 @@ async fn quota_read_rejects_corrupt_provider_document() {
expected_revision: CredentialRevision::new(1).expect("revision"),
quota: Some(OpaqueProviderData::new(document)),
observed_at: Some(SystemTime::now()),
limit_reached: None,
})
.await
.expect("seed corrupt quota");
backend/crates/providers/xai/tests/credential/selector.rs
+1
View File
@@ -245,6 +245,7 @@ impl SelectorFixture {
document.as_object().expect("quota object").clone(),
)),
observed_at: Some(SystemTime::now()),
limit_reached: None,
})
.await
.expect("persist quota");
backend/migrations/.frozen-sha256
+2
View File
@@ -1,2 +1,4 @@
54be07512951a561391b6e611355a1bfc119d82363253fb9bbdb225c9e3e39bb 0001_initial.sql
f480ac4cb70f6963eee71c0b8a67777ebb492f85501c90ebe7544bee0b2c0525 0002_remove_account_cooldown.sql
5cf5b5387c66b4d437d79c002cec46cb85196789060fa87c39e97a382487cac9 0004_fix_quota_limit_reached_null.sql
a3f236d558e7598f69277520458c801ee43130f7704ca25455df468cf8899ff2 0003_materialize_quota_limit_reached.sql
backend/migrations/0003_materialize_quota_limit_reached.sql
+20
View File
@@ -0,0 +1,20 @@
-- 物化快照级限流事实：Dashboard 复用 Provider 解析结果，不再 SQL 解析 raw quota JSON。
--
-- `quota_limit_reached` 由 Provider 在写 quota 时计算（快照级：顶层或任一窗口触顶，
-- 已过 reset 的窗口不计数），Admin/Dashboard 直接读该列。
alter table provider_accounts
add column quota_limit_reached boolean not null default false;
-- 存量回填：从现有 raw JSON 推导（仅 primary_window 的 limit_reached + reset 未过期）。
update provider_accounts
set quota_limit_reached = coalesce(
provider_quota_json is not null
and provider_quota_json->'rate_limit'->>'limit_reached' = 'true'
and (
provider_quota_json->'rate_limit'->'primary_window'->>'reset_at' is null
or (provider_quota_json->'rate_limit'->'primary_window'->>'reset_at')::bigint > extract(epoch from now())::bigint
),
false
)
where provider_quota_json is not null;
backend/migrations/0004_fix_quota_limit_reached_null.sql
+8
View File
@@ -0,0 +1,8 @@
-- 修正 0003 存量回填：raw JSON 缺 rate_limit 键时表达式为 NULL，
-- 赋 NULL 违反 not-null 约束。0004 把仍为 NULL 的行兜底为 false
-- （0003 已应用的环境，add column 用 default false 填充，通常无 NULL；
-- 本迁移保证任何环境下该列都非 NULL）。
update provider_accounts
set quota_limit_reached = coalesce(quota_limit_reached, false)
where quota_limit_reached is null;