From a5d3d9f12d2d531dc300ec5fd0e8304f35c67f88 Mon Sep 17 00:00:00 2001 From: Zyy Date: Sat, 31 Oct 2026 15:58:02 +0800 Subject: [PATCH] fix(peer): recover quota from explicit upstream budget state Export authoritative budget access flags and keep monetary fields display-only for peer consumers. Use the official usage_limit_reached error type for budget exhaustion while preserving ordinary rate-limit cooldowns. Apply fresh quota facts atomically with peer snapshots through the shared quota writer, retaining connection, claim, and observation-time guards. Reuse the existing sync worker and document the recovery contract. Validated with 234 selected tests, strict workspace Clippy, and Rustfmt. Co-authored-by: Codex --- .../gateway-admin/src/model/peer_link.rs | 3 +- .../src/use_case/peer_link/mod.rs | 4 + .../src/use_case/peer_link/upstream/mod.rs | 1 + .../use_case/peer_link/upstream/projection.rs | 7 +- .../src/use_case/peer_link/upstream/quota.rs | 79 +++ .../use_case/peer_link/upstream/scheduling.rs | 24 +- .../use_case/peer_link/upstream/service.rs | 11 +- .../use_case/peer_link/upstream/snapshot.rs | 53 -- .../tests/use_case/peer_link/upstream/mod.rs | 2 +- .../crates/gateway-api/src/openai/error.rs | 5 + .../gateway-api/src/openai/peer_link/wire.rs | 11 +- .../crates/gateway-api/tests/openai/error.rs | 5 +- .../gateway-api/tests/openai/peer_link/mod.rs | 2 + .../gateway-core/src/account/upstream.rs | 8 +- .../crates/gateway-core/src/engine/budget.rs | 23 + backend/crates/gateway-core/src/error/mod.rs | 3 + .../tests/engine/execution/mod.rs | 4 +- .../runtime/src/callback/error.rs | 1 + backend/crates/gateway-protocol/src/peer.rs | 5 + .../src/postgres/client_budgets.rs | 15 +- .../postgres/provider_accounts/repository.rs | 59 +- .../provider_accounts/upstream_info.rs | 25 +- .../tests/postgres/client_budgets.rs | 4 +- .../provider_accounts/upstream_info.rs | 156 +++++- .../openai/src/credential/selector.rs | 2 +- .../providers/openai/tests/admin/peer_info.rs | 524 ++++++++++++++---- deploy/README.md | 4 +- docs/api.md | 19 +- docs/architecture.md | 5 +- 29 files changed, 850 insertions(+), 214 deletions(-) create mode 100644 backend/crates/gateway-admin/src/use_case/peer_link/upstream/quota.rs diff --git a/backend/crates/gateway-admin/src/model/peer_link.rs b/backend/crates/gateway-admin/src/model/peer_link.rs index f9e47c922..88345a481 100644 --- a/backend/crates/gateway-admin/src/model/peer_link.rs +++ b/backend/crates/gateway-admin/src/model/peer_link.rs @@ -2,7 +2,7 @@ use super::provider_credentials::ProviderQuotaWindowRole; use chrono::{DateTime, Utc}; -use gateway_core::engine::budget::ClientBudgetStatus; +use gateway_core::engine::budget::{ClientBudgetAccess, ClientBudgetStatus}; /// 调用边界提供的响应资源预算，超限时整体拒绝而不截断授权范围 #[derive(Debug, Clone, Copy)] @@ -25,6 +25,7 @@ pub struct PeerLinkExport { pub observed_at: DateTime, pub expires_at: DateTime, pub budget: ClientBudgetStatus, + pub budget_access: ClientBudgetAccess, pub budget_observed_at: DateTime, pub pool: Option, pub quota: Option>, diff --git a/backend/crates/gateway-admin/src/use_case/peer_link/mod.rs b/backend/crates/gateway-admin/src/use_case/peer_link/mod.rs index 2d584a4bd..19525b1e8 100644 --- a/backend/crates/gateway-admin/src/use_case/peer_link/mod.rs +++ b/backend/crates/gateway-admin/src/use_case/peer_link/mod.rs @@ -119,6 +119,10 @@ impl PeerLinkService for DefaultPeerLinkService { node, observed_at: now, expires_at: now + TimeDelta::minutes(5), + budget_access: key + .budget + .limits + .access(key.budget.daily_used_usd, key.budget.weekly_used_usd), budget: key.budget, budget_observed_at, pool: None, diff --git a/backend/crates/gateway-admin/src/use_case/peer_link/upstream/mod.rs b/backend/crates/gateway-admin/src/use_case/peer_link/upstream/mod.rs index 4c7cefe88..17202f9d7 100644 --- a/backend/crates/gateway-admin/src/use_case/peer_link/upstream/mod.rs +++ b/backend/crates/gateway-admin/src/use_case/peer_link/upstream/mod.rs @@ -2,6 +2,7 @@ mod accounts; mod projection; +mod quota; mod scheduling; mod service; mod snapshot; diff --git a/backend/crates/gateway-admin/src/use_case/peer_link/upstream/projection.rs b/backend/crates/gateway-admin/src/use_case/peer_link/upstream/projection.rs index 972f617a6..4f5b22179 100644 --- a/backend/crates/gateway-admin/src/use_case/peer_link/upstream/projection.rs +++ b/backend/crates/gateway-admin/src/use_case/peer_link/upstream/projection.rs @@ -12,12 +12,12 @@ use crate::model::{ UpstreamSyncState, }, }; -use gateway_core::account::{UpstreamInfoPlan, UpstreamInfoRecord}; +use gateway_core::account::{QuotaAccessState, UpstreamInfoPlan, UpstreamInfoRecord}; use gateway_protocol::peer::{ PeerComponentStatus, PeerLink, PeerPoolDisclosure, PeerQuotaDisclosure, }; -use super::{PeerInfoService, snapshot}; +use super::{PeerInfoService, quota, snapshot}; impl PeerInfoService { pub(crate) fn project(record: &UpstreamInfoRecord, now: SystemTime) -> ProviderQuota { @@ -138,7 +138,8 @@ impl PeerInfoService { window_seconds: Some(seconds), used_percent, reset_at: window.resets_at.as_deref().and_then(snapshot::datetime), - limit_reached: !window_stale && snapshot::budget_exhausted(window, now), + limit_reached: !window_stale + && quota::window_access(window, now) == QuotaAccessState::Exhausted, local_usage: None, provider_data: None, }); diff --git a/backend/crates/gateway-admin/src/use_case/peer_link/upstream/quota.rs b/backend/crates/gateway-admin/src/use_case/peer_link/upstream/quota.rs new file mode 100644 index 000000000..1f3ba9472 --- /dev/null +++ b/backend/crates/gateway-admin/src/use_case/peer_link/upstream/quota.rs @@ -0,0 +1,79 @@ +//! 消费互联明确的额度访问结论，金额只由展示层解释 + +use std::time::SystemTime; + +use gateway_core::account::{QuotaAccessState, QuotaEvidence, QuotaState}; +use gateway_protocol::peer::{PeerBudgetWindow, PeerComponentStatus, PeerKeyUsage, PeerLink}; + +use super::snapshot::{component_fresh, fresh, timestamp}; + +/// 原始时间与请求开始时间共同限制观测，迟到响应不能覆盖请求期间的新失败 +pub(super) fn observe( + peer: &PeerLink, + started_at: SystemTime, + now: SystemTime, +) -> Option { + if !fresh(&peer.observed_at, now) + || !component_fresh(Some(&peer.observed_at), Some(&peer.expires_at), now) + || peer.usage.status != PeerComponentStatus::Available + || !component_fresh(peer.usage.observed_at.as_deref(), None, now) + { + return None; + } + let observed_at = started_at + .min(timestamp(&peer.observed_at)?) + .min(timestamp(peer.usage.observed_at.as_deref()?)?); + let usage = peer.usage.data.as_ref()?; + match access(usage, now) { + QuotaAccessState::Allowed => Some(QuotaState::allowed(observed_at)), + QuotaAccessState::Exhausted => { + let reset_at = [&usage.daily, &usage.weekly] + .into_iter() + .filter(|window| window_access(window, now) == QuotaAccessState::Exhausted) + .filter_map(|window| window.resets_at.as_deref().and_then(timestamp)) + .min(); + Some(QuotaState::exhausted( + if usage.allowed == Some(false) { + QuotaEvidence::ProviderDenied + } else { + QuotaEvidence::AccountLimitReached + }, + observed_at, + reset_at, + )) + } + QuotaAccessState::Unknown => None, + } +} + +/// 与 OAuth 相同，明确的 allowed 优先，窗口耗尽仅补充缺失的访问结论 +pub(super) fn access(usage: &PeerKeyUsage, now: SystemTime) -> QuotaAccessState { + match usage.allowed { + Some(true) => QuotaAccessState::Allowed, + Some(false) => QuotaAccessState::Exhausted, + None if [&usage.daily, &usage.weekly] + .into_iter() + .any(|window| window_access(window, now) == QuotaAccessState::Exhausted) => + { + QuotaAccessState::Exhausted + } + None => QuotaAccessState::Unknown, + } +} + +/// 过期窗口保持未知，不从重置时间或金额推断恢复 +pub(super) fn window_access(window: &PeerBudgetWindow, now: SystemTime) -> QuotaAccessState { + if window + .resets_at + .as_deref() + .and_then(timestamp) + .is_some_and(|reset| reset <= now) + { + return QuotaAccessState::Unknown; + } + match window.limit_reached { + Some(true) => QuotaAccessState::Exhausted, + Some(false) => QuotaAccessState::Allowed, + None => QuotaAccessState::Unknown, + } +} diff --git a/backend/crates/gateway-admin/src/use_case/peer_link/upstream/scheduling.rs b/backend/crates/gateway-admin/src/use_case/peer_link/upstream/scheduling.rs index cc98a8319..92b22e018 100644 --- a/backend/crates/gateway-admin/src/use_case/peer_link/upstream/scheduling.rs +++ b/backend/crates/gateway-admin/src/use_case/peer_link/upstream/scheduling.rs @@ -4,9 +4,9 @@ use std::{collections::HashMap, sync::Arc, time::SystemTime}; use gateway_protocol::peer::{PeerComponentStatus, PeerLink, PeerModelCoverage, PeerModelState}; -use super::snapshot; +use super::{quota, snapshot}; -use gateway_core::account::UpstreamSchedulingSignal; +use gateway_core::account::{QuotaAccessState, UpstreamSchedulingSignal}; #[derive(Clone, Copy)] struct Validity { @@ -75,14 +75,22 @@ impl SchedulingObservation { && let Some(usage) = &peer.usage.data && let Some(usage_validity) = Validity::component(peer.usage.observed_at.as_deref(), None) + && quota::access(usage, now) == QuotaAccessState::Exhausted { - for window in [&usage.daily, &usage.weekly] { - if snapshot::budget_exhausted(window, now) { - let mut validity = usage_validity; - if let Some(reset) = window.resets_at.as_deref().and_then(snapshot::timestamp) { - validity.expires_at = validity.expires_at.min(reset); + if usage.allowed.is_some() { + exhausted_budgets.push(usage_validity); + } else { + // 窗口提供的补充证据只在各自重置边界之前影响调度评分 + for window in [&usage.daily, &usage.weekly] { + if quota::window_access(window, now) == QuotaAccessState::Exhausted { + let mut validity = usage_validity; + if let Some(reset) = + window.resets_at.as_deref().and_then(snapshot::timestamp) + { + validity.expires_at = validity.expires_at.min(reset); + } + exhausted_budgets.push(validity); } - exhausted_budgets.push(validity); } } } diff --git a/backend/crates/gateway-admin/src/use_case/peer_link/upstream/service.rs b/backend/crates/gateway-admin/src/use_case/peer_link/upstream/service.rs index d719e1531..f6061fab1 100644 --- a/backend/crates/gateway-admin/src/use_case/peer_link/upstream/service.rs +++ b/backend/crates/gateway-admin/src/use_case/peer_link/upstream/service.rs @@ -1,4 +1,4 @@ -//! 后台计划协调与仅供 Smart 排序的内存观测 +//! 互联后台同步、额度恢复与 Smart 排序观测 use std::{ collections::{BTreeMap, HashMap}, @@ -27,7 +27,7 @@ use crate::ports::{ provider::{ProviderAdmin, ProviderAdminRegistry}, }; -use super::{scheduling::SchedulingObservation, snapshot}; +use super::{quota, scheduling::SchedulingObservation, snapshot}; const BATCH_SIZE: u32 = 100; const FETCH_BATCH: u32 = 20; @@ -334,12 +334,18 @@ impl PeerInfoService { let Some(authentication) = self.connection(&loaded) else { return self.reconcile(&loaded).await; }; + let started_at = SystemTime::now(); let (result, retry_after) = match self.transport.fetch(&loaded.account, &authentication).await { Ok(bytes) => (snapshot::parse(&bytes, SystemTime::now()), None), Err(failure) => (Err(failure.kind), failure.retry_after), }; let now = SystemTime::now(); + let quota = result + .as_ref() + .ok() + .and_then(snapshot::decode) + .and_then(|peer| quota::observe(&peer, started_at, now)); let delay = match result { Ok(_) => Duration::from_secs(240 + stable_jitter(loaded.account.id().as_str())), Err(UpstreamInfoFailure::Unauthorized | UpstreamInfoFailure::Unsupported) => { @@ -359,6 +365,7 @@ impl PeerInfoService { completed_at: now, next_attempt_at: now + delay, result, + quota, }; if self.store.complete(completion).await? { self.refresh_cache(std::slice::from_ref(loaded.account.id())) diff --git a/backend/crates/gateway-admin/src/use_case/peer_link/upstream/snapshot.rs b/backend/crates/gateway-admin/src/use_case/peer_link/upstream/snapshot.rs index d990b05e2..b051c11f9 100644 --- a/backend/crates/gateway-admin/src/use_case/peer_link/upstream/snapshot.rs +++ b/backend/crates/gateway-admin/src/use_case/peer_link/upstream/snapshot.rs @@ -237,56 +237,3 @@ pub(super) fn decimal(value: &str) -> Option { .ok() .filter(|value| value.is_finite() && *value >= 0.0) } - -pub(super) fn budget_exhausted( - window: &gateway_protocol::peer::PeerBudgetWindow, - now: SystemTime, -) -> bool { - let Some(total) = window.total.as_deref() else { - return false; - }; - if compare_decimal(total, "0") != std::cmp::Ordering::Greater - || window - .resets_at - .as_deref() - .and_then(timestamp) - .is_some_and(|reset| reset <= now) - { - return false; - } - window - .remaining - .as_deref() - .is_some_and(|remaining| compare_decimal(remaining, "0") == std::cmp::Ordering::Equal) - || compare_decimal(&window.used, total) != std::cmp::Ordering::Less -} - -// 金额触顶比较保留十进制精度，百分比展示的浮点近似不能成为耗尽依据 -fn compare_decimal(left: &str, right: &str) -> std::cmp::Ordering { - let parts = |value: &str| { - let (integer, fractional) = value.split_once('.').unwrap_or((value, "")); - ( - integer.trim_start_matches('0').to_owned(), - fractional.to_owned(), - ) - }; - let (left_integer, left_fractional) = parts(left); - let (right_integer, right_fractional) = parts(right); - left_integer - .len() - .cmp(&right_integer.len()) - .then_with(|| left_integer.cmp(&right_integer)) - .then_with(|| { - let width = left_fractional.len().max(right_fractional.len()); - left_fractional - .bytes() - .chain(std::iter::repeat(b'0')) - .take(width) - .cmp( - right_fractional - .bytes() - .chain(std::iter::repeat(b'0')) - .take(width), - ) - }) -} diff --git a/backend/crates/gateway-admin/tests/use_case/peer_link/upstream/mod.rs b/backend/crates/gateway-admin/tests/use_case/peer_link/upstream/mod.rs index 703e6e9a6..74f0e7d95 100644 --- a/backend/crates/gateway-admin/tests/use_case/peer_link/upstream/mod.rs +++ b/backend/crates/gateway-admin/tests/use_case/peer_link/upstream/mod.rs @@ -159,7 +159,7 @@ impl PeerTransport for TestTransport { "observed_at":now.to_rfc3339(), "expires_at":(now + chrono::Duration::minutes(5)).to_rfc3339(), "capabilities":["key_usage"], - "usage":{"status":"available","observed_at":now.to_rfc3339(),"data":{"unit":"USD", + "usage":{"status":"available","observed_at":now.to_rfc3339(),"data":{"unit":"USD","allowed":false, "daily":{"total":"1","used":"1","remaining":"0","resetsAt":(now + chrono::Duration::days(1)).to_rfc3339()}, "weekly":{"total":"0","used":"0","remaining":"0"}}}, "pool":{"disclosure":"none"}, "quota":{"disclosure":"none"}, "models":[] diff --git a/backend/crates/gateway-api/src/openai/error.rs b/backend/crates/gateway-api/src/openai/error.rs index 939505e63..d0e9e8288 100644 --- a/backend/crates/gateway-api/src/openai/error.rs +++ b/backend/crates/gateway-api/src/openai/error.rs @@ -300,6 +300,11 @@ pub const fn gateway_error_contract( "rate_limit_error", "rate_limit_exceeded", ), + GatewayErrorKind::BudgetExhausted => ( + StatusCode::TOO_MANY_REQUESTS, + "usage_limit_reached", + "usage_limit_reached", + ), GatewayErrorKind::UpstreamUnavailable => ( StatusCode::BAD_GATEWAY, "server_error", diff --git a/backend/crates/gateway-api/src/openai/peer_link/wire.rs b/backend/crates/gateway-api/src/openai/peer_link/wire.rs index 8d836735d..f1103dec3 100644 --- a/backend/crates/gateway-api/src/openai/peer_link/wire.rs +++ b/backend/crates/gateway-api/src/openai/peer_link/wire.rs @@ -89,15 +89,18 @@ pub(super) fn response(export: PeerLinkExport) -> PeerLink { observed_at: Some(export.budget_observed_at.to_rfc3339()), data: Some(PeerKeyUsage { unit: "USD".to_owned(), + allowed: Some(export.budget_access.allowed()), daily: budget_window( export.budget.limits.daily_usd, export.budget.daily_used_usd, export.budget.daily_resets_at, + export.budget_access.daily_limit_reached, ), weekly: budget_window( export.budget.limits.weekly_usd, export.budget.weekly_used_usd, export.budget.weekly_resets_at, + export.budget_access.weekly_limit_reached, ), }), }, @@ -125,9 +128,15 @@ pub(super) fn response(export: PeerLinkExport) -> PeerLink { } } -fn budget_window(total: Decimal, used: Decimal, resets_at: Option) -> PeerBudgetWindow { +fn budget_window( + total: Decimal, + used: Decimal, + resets_at: Option, + limit_reached: bool, +) -> PeerBudgetWindow { let limited = total != Decimal::ZERO; PeerBudgetWindow { + limit_reached: Some(limit_reached), total: limited.then(|| total.canonical()), used: used.canonical(), remaining: limited.then(|| total.saturating_sub(used).canonical()), diff --git a/backend/crates/gateway-api/tests/openai/error.rs b/backend/crates/gateway-api/tests/openai/error.rs index 89aa34656..d7f9d65a1 100644 --- a/backend/crates/gateway-api/tests/openai/error.rs +++ b/backend/crates/gateway-api/tests/openai/error.rs @@ -165,9 +165,9 @@ async fn classified_capacity_error_without_special_code_returns_retryable_http_s } #[tokio::test] -async fn key_budget_errors_preserve_limit_code_and_retry_after() { +async fn key_budget_errors_use_the_official_quota_type_and_preserve_limit_details() { for code in ["key_daily_budget_exceeded", "key_weekly_budget_exceeded"] { - let error = GatewayError::new(GatewayErrorKind::RateLimited, "key budget exhausted") + let error = GatewayError::new(GatewayErrorKind::BudgetExhausted, "key budget exhausted") .with_client_code(code) .with_retry_after(std::time::Duration::from_millis(1501)); let response = gateway_error_response(&error); @@ -176,6 +176,7 @@ async fn key_budget_errors_preserve_limit_code_and_retry_after() { let body: serde_json::Value = serde_json::from_slice(&to_bytes(response.into_body(), 4096).await.unwrap()).unwrap(); assert_eq!(body["error"]["code"], code); + assert_eq!(body["error"]["type"], "usage_limit_reached"); } } diff --git a/backend/crates/gateway-api/tests/openai/peer_link/mod.rs b/backend/crates/gateway-api/tests/openai/peer_link/mod.rs index aaa774e72..344fc54a9 100644 --- a/backend/crates/gateway-api/tests/openai/peer_link/mod.rs +++ b/backend/crates/gateway-api/tests/openai/peer_link/mod.rs @@ -244,6 +244,8 @@ async fn exports_scoped_observations_even_when_the_key_budget_is_exhausted_witho let body = response_json(response).await; assert_eq!(body["object"], "cpr.peer"); assert_eq!(body["usage"]["data"]["daily"]["remaining"], "0"); + assert_eq!(body["usage"]["data"]["allowed"], false); + assert_eq!(body["usage"]["data"]["daily"]["limit_reached"], true); assert_eq!(body["pool"]["account_count"], 1); assert!(body["pool"].get("plans").is_none()); assert_eq!( diff --git a/backend/crates/gateway-core/src/account/upstream.rs b/backend/crates/gateway-core/src/account/upstream.rs index 5b9223b4f..ca519227c 100644 --- a/backend/crates/gateway-core/src/account/upstream.rs +++ b/backend/crates/gateway-core/src/account/upstream.rs @@ -1,4 +1,4 @@ -//! 上游互联观测与后台计划；原始协议由控制面解释，不改变账号自身额度状态 +//! 上游互联观测、后台计划与额度恢复合同，原始协议由控制面解释 use std::{ sync::{Arc, OnceLock, Weak}, @@ -7,7 +7,9 @@ use std::{ use sha2::{Digest, Sha256}; -use super::{CredentialRevision, OpaqueProviderData, ProviderAccount, ProviderAccountId}; +use super::{ + CredentialRevision, OpaqueProviderData, ProviderAccount, ProviderAccountId, QuotaState, +}; /// 请求发出时的连接身份；出口摘要包含托管代理解析后的认证信息，但不持久化明文 #[derive(Debug, Clone, PartialEq, Eq)] @@ -110,6 +112,8 @@ pub struct UpstreamInfoCompletion { pub completed_at: SystemTime, pub next_attempt_at: SystemTime, pub result: Result, + /// 控制面确认的额度事实，随有效快照原子提交，未知或过期观测不得覆盖既有事实 + pub quota: Option, } /// 调度侧只消费已发布的中立观测，不读取协议、凭据或持久存储 diff --git a/backend/crates/gateway-core/src/engine/budget.rs b/backend/crates/gateway-core/src/engine/budget.rs index 3ab1f082d..527f580d0 100644 --- a/backend/crates/gateway-core/src/engine/budget.rs +++ b/backend/crates/gateway-core/src/engine/budget.rs @@ -19,6 +19,29 @@ impl ClientBudgetLimits { pub fn is_limited(self) -> bool { self.daily_usd != Decimal::ZERO || self.weekly_usd != Decimal::ZERO } + + /// 准入与对外预算查询共用真实账本口径，展示字符串不参与访问判断 + #[must_use] + pub fn access(self, daily_used: Decimal, weekly_used: Decimal) -> ClientBudgetAccess { + ClientBudgetAccess { + daily_limit_reached: self.daily_usd != Decimal::ZERO && daily_used >= self.daily_usd, + weekly_limit_reached: self.weekly_usd != Decimal::ZERO + && weekly_used >= self.weekly_usd, + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ClientBudgetAccess { + pub daily_limit_reached: bool, + pub weekly_limit_reached: bool, +} + +impl ClientBudgetAccess { + #[must_use] + pub const fn allowed(self) -> bool { + !self.daily_limit_reached && !self.weekly_limit_reached + } } #[derive(Debug, Clone, Default, PartialEq, Eq)] diff --git a/backend/crates/gateway-core/src/error/mod.rs b/backend/crates/gateway-core/src/error/mod.rs index af2fc0a04..9383d6fa8 100644 --- a/backend/crates/gateway-core/src/error/mod.rs +++ b/backend/crates/gateway-core/src/error/mod.rs @@ -1094,6 +1094,8 @@ pub enum GatewayErrorKind { ProviderInfrastructureUnavailable, /// 上游限流 RateLimited, + /// 当前客户端 Key 的费用预算耗尽 + BudgetExhausted, /// 上游暂不可用 UpstreamUnavailable, /// 上游拒收当前请求：消息过大（如 WebSocket close 1009） @@ -1124,6 +1126,7 @@ impl GatewayErrorKind { Self::ConcurrencyQueueTimeout => "concurrency_queue_timeout", Self::ProviderInfrastructureUnavailable => "provider_infrastructure_unavailable", Self::RateLimited => "rate_limited", + Self::BudgetExhausted => "budget_exhausted", Self::UpstreamUnavailable => "upstream_unavailable", Self::MessageTooBig => "message_too_big", Self::Timeout => "timeout", diff --git a/backend/crates/gateway-core/tests/engine/execution/mod.rs b/backend/crates/gateway-core/tests/engine/execution/mod.rs index dcffde6e4..6d2b96eff 100644 --- a/backend/crates/gateway-core/tests/engine/execution/mod.rs +++ b/backend/crates/gateway-core/tests/engine/execution/mod.rs @@ -82,7 +82,7 @@ impl ClientBudgetPort for Budget { assert!(self.active.load(Ordering::SeqCst)); if self.reject { Err( - GatewayError::new(GatewayErrorKind::RateLimited, "budget exhausted") + GatewayError::new(GatewayErrorKind::BudgetExhausted, "budget exhausted") .with_client_code("key_daily_budget_exceeded"), ) } else { @@ -2639,7 +2639,7 @@ fn final_observation_is_emitted_once_for_success_rejection_and_detached_cancella observation_service(ObservationProviderBehavior::CompleteWithUsage, true); assert!(matches!( service.start(request(&service, ClientTransport::HttpJson)).await, - Err(error) if error.kind() == GatewayErrorKind::RateLimited + Err(error) if error.kind() == GatewayErrorKind::BudgetExhausted )); { let observations = observer.observations.lock().unwrap(); diff --git a/backend/crates/gateway-plugin/runtime/src/callback/error.rs b/backend/crates/gateway-plugin/runtime/src/callback/error.rs index 8b20f84df..814884203 100644 --- a/backend/crates/gateway-plugin/runtime/src/callback/error.rs +++ b/backend/crates/gateway-plugin/runtime/src/callback/error.rs @@ -24,6 +24,7 @@ pub(super) fn gateway(error: GatewayError) -> PluginFault { | GatewayErrorKind::ConcurrencyQueueFull | GatewayErrorKind::ConcurrencyQueueTimeout | GatewayErrorKind::RateLimited + | GatewayErrorKind::BudgetExhausted | GatewayErrorKind::NoAvailableProvider => ErrorCode::Capacity, GatewayErrorKind::Timeout => ErrorCode::Timeout, GatewayErrorKind::Cancelled => ErrorCode::Cancelled, diff --git a/backend/crates/gateway-protocol/src/peer.rs b/backend/crates/gateway-protocol/src/peer.rs index d8e9ee988..776092032 100644 --- a/backend/crates/gateway-protocol/src/peer.rs +++ b/backend/crates/gateway-protocol/src/peer.rs @@ -54,6 +54,9 @@ pub struct PeerUsage { #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct PeerKeyUsage { pub unit: String, + /// 上游账本给出的访问结论，缺失时不能根据展示金额推断恢复 + #[serde(default, skip_serializing_if = "Option::is_none")] + pub allowed: Option, pub daily: PeerBudgetWindow, pub weekly: PeerBudgetWindow, } @@ -61,6 +64,8 @@ pub struct PeerKeyUsage { /// 金额保留十进制字符串，不限额时 total 与 remaining 为 null #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct PeerBudgetWindow { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub limit_reached: Option, pub total: Option, pub used: String, pub remaining: Option, diff --git a/backend/crates/gateway-store/src/postgres/client_budgets.rs b/backend/crates/gateway-store/src/postgres/client_budgets.rs index fa3c556d4..775d6eaea 100644 --- a/backend/crates/gateway-store/src/postgres/client_budgets.rs +++ b/backend/crates/gateway-store/src/postgres/client_budgets.rs @@ -235,10 +235,9 @@ impl PgClientBudgetStore { .ok_or_else(unavailable)?; } } - let daily_exceeded = limits.daily_usd != Decimal::ZERO && daily >= limits.daily_usd; - let weekly_exceeded = limits.weekly_usd != Decimal::ZERO && weekly >= limits.weekly_usd; - if daily_exceeded || weekly_exceeded { - let (reset, code, message) = if weekly_exceeded { + let access = limits.access(daily, weekly); + if !access.allowed() { + let (reset, code, message) = if access.weekly_limit_reached { ( weekly_end, "key_weekly_budget_exceeded", @@ -251,9 +250,11 @@ impl PgClientBudgetStore { "Daily spending limit reached.", ) }; - return Err(GatewayError::new(GatewayErrorKind::RateLimited, message) - .with_client_code(code) - .with_retry_after((reset - now).to_std().unwrap_or(Duration::from_secs(1)))); + return Err( + GatewayError::new(GatewayErrorKind::BudgetExhausted, message) + .with_client_code(code) + .with_retry_after((reset - now).to_std().unwrap_or(Duration::from_secs(1))), + ); } } // 首次使用或窗口到期才写入；与重置、结算保持相同锁序 diff --git a/backend/crates/gateway-store/src/postgres/provider_accounts/repository.rs b/backend/crates/gateway-store/src/postgres/provider_accounts/repository.rs index 41252da60..a47544ee9 100644 --- a/backend/crates/gateway-store/src/postgres/provider_accounts/repository.rs +++ b/backend/crates/gateway-store/src/postgres/provider_accounts/repository.rs @@ -424,29 +424,7 @@ impl ProviderAccountRepository for PgProviderAccountRepository { expected_revision: Revision, state: QuotaState, ) -> StoreResult { - require_nonempty(ENTITY, "account_id", account_id)?; - let observed_at = state - .observed_at() - .map(DateTime::::from) - .ok_or_else(|| invalid("quota access change requires observed_at"))?; - let result = sqlx::query( - "update provider_accounts - set quota_access_observed_at = $3, quota_access_state = $4, - quota_evidence = $5, quota_reset_at = $6, - updated_at = greatest(now(), updated_at, $3) - where id = $1 and credential_revision = $2 - and (quota_access_observed_at is null or quota_access_observed_at <= $3)", - ) - .bind(account_id) - .bind(to_i64(expected_revision.get())?) - .bind(observed_at) - .bind(state.access().as_str()) - .bind(state.evidence().map(QuotaEvidence::as_str)) - .bind(state.reset_at().map(DateTime::::from)) - .execute(&self.pool) - .await - .map_err(|source| postgres_unavailable("apply provider quota access", source))?; - Ok(result.rows_affected() == 1) + write_quota_access(&self.pool, account_id, expected_revision, state).await } async fn touch_provider_quota_observation( @@ -1249,3 +1227,38 @@ impl PgProviderAccountRepository { rows.into_iter().map(account_summary_from_row).collect() } } + +/// 真实请求反馈与互联同步共用额度事实写入，调用方可将其并入已有事务 +pub(super) async fn write_quota_access<'e, E>( + executor: E, + account_id: &str, + expected_revision: Revision, + state: QuotaState, +) -> StoreResult +where + E: sqlx::Executor<'e, Database = sqlx::Postgres>, +{ + require_nonempty(ENTITY, "account_id", account_id)?; + let observed_at = state + .observed_at() + .map(DateTime::::from) + .ok_or_else(|| invalid("quota access change requires observed_at"))?; + let result = sqlx::query( + "update provider_accounts + set quota_access_observed_at = $3, quota_access_state = $4, + quota_evidence = $5, quota_reset_at = $6, + updated_at = greatest(now(), updated_at, $3) + where id = $1 and credential_revision = $2 + and (quota_access_observed_at is null or quota_access_observed_at <= $3)", + ) + .bind(account_id) + .bind(to_i64(expected_revision.get())?) + .bind(observed_at) + .bind(state.access().as_str()) + .bind(state.evidence().map(QuotaEvidence::as_str)) + .bind(state.reset_at().map(DateTime::::from)) + .execute(executor) + .await + .map_err(|source| postgres_unavailable("apply provider quota access", source))?; + Ok(result.rows_affected() == 1) +} diff --git a/backend/crates/gateway-store/src/postgres/provider_accounts/upstream_info.rs b/backend/crates/gateway-store/src/postgres/provider_accounts/upstream_info.rs index d66c53c1b..8a8fc9edd 100644 --- a/backend/crates/gateway-store/src/postgres/provider_accounts/upstream_info.rs +++ b/backend/crates/gateway-store/src/postgres/provider_accounts/upstream_info.rs @@ -3,8 +3,8 @@ use std::time::{Duration, SystemTime}; use gateway_core::account::{ - UpstreamConnection, UpstreamInfoClaim, UpstreamInfoCompletion, UpstreamInfoFailure, - UpstreamInfoPlan, UpstreamInfoRecord, UpstreamInfoStore, + QuotaAccessState, UpstreamConnection, UpstreamInfoClaim, UpstreamInfoCompletion, + UpstreamInfoFailure, UpstreamInfoPlan, UpstreamInfoRecord, UpstreamInfoStore, }; use sqlx::postgres::PgRow; @@ -272,6 +272,27 @@ impl UpstreamInfoStore for PgProviderAccountRepository { .execute(&mut *transaction) .await .map_err(database_error)?; + if result.rows_affected() == 1 + && failure.is_none() + && let Some(state) = completion.quota + && state.access() != QuotaAccessState::Unknown + && let Some(observed_at) = state.observed_at() + && account + .quota() + .observed_at() + .is_none_or(|current| current < observed_at) + { + // 连接与领取围栏已核对，行锁内拒绝同刻及更早的事实，再复用通用写入 + super::repository::write_quota_access( + &mut *transaction, + completion.connection.account_id.as_str(), + Revision::new(completion.connection.credential_revision.get()) + .map_err(crate::core_store_error)?, + state, + ) + .await + .map_err(crate::core_store_error)?; + } transaction.commit().await.map_err(database_error)?; Ok(result.rows_affected() == 1) } diff --git a/backend/crates/gateway-store/tests/postgres/client_budgets.rs b/backend/crates/gateway-store/tests/postgres/client_budgets.rs index b48f1c472..68849f845 100644 --- a/backend/crates/gateway-store/tests/postgres/client_budgets.rs +++ b/backend/crates/gateway-store/tests/postgres/client_budgets.rs @@ -589,7 +589,7 @@ async fn budgets_settle_exactly_once_and_enforce_each_threshold_across_store_ins assert_eq!(day.daily_used_usd.canonical(), "0.3"); assert_eq!(day.weekly_used_usd.canonical(), "0.3"); let day_error = second.admit(key_id("day")).await.unwrap_err(); - assert_eq!(day_error.kind(), GatewayErrorKind::RateLimited); + assert_eq!(day_error.kind(), GatewayErrorKind::BudgetExhausted); assert_eq!( day_error.client_error_code(), Some("key_daily_budget_exceeded") @@ -1160,7 +1160,7 @@ async fn plugin_budget_limits_preserve_consumption_and_unrelated_configuration_a ); assert_eq!( budgets.admit(key_id("key")).await.unwrap_err().kind(), - GatewayErrorKind::RateLimited + GatewayErrorKind::BudgetExhausted ); assert!( store diff --git a/backend/crates/gateway-store/tests/postgres/provider_accounts/upstream_info.rs b/backend/crates/gateway-store/tests/postgres/provider_accounts/upstream_info.rs index 874c5b7f8..45ff26684 100644 --- a/backend/crates/gateway-store/tests/postgres/provider_accounts/upstream_info.rs +++ b/backend/crates/gateway-store/tests/postgres/provider_accounts/upstream_info.rs @@ -1,4 +1,4 @@ -//! 验证互联同步计划的并发领取、租约回收、快照保留与连接写入围栏 +//! 验证互联计划、快照与额度恢复的原子提交及连接写入围栏 use super::*; use gateway_core::account::{ @@ -55,6 +55,160 @@ fn completion( completed_at: now() + Duration::from_secs(at), next_attempt_at: now() + Duration::from_secs(at + 300), result, + quota: None, + } +} + +#[tokio::test] +async fn quota_recovery_commits_only_with_a_current_successful_claim_and_newer_evidence() { + for case in [ + "recovered", + "older", + "equal", + "new_failure", + "wrong_claim", + "expired_claim", + "disabled", + "rotated", + "proxy_changed", + "query_failed", + "no_evidence", + "credential_error", + "exhausted", + "unknown", + "old_exhaustion", + ] { + let Some((database, repository, connection)) = fixture("peer_recovery").await else { + return; + }; + let id = &connection.account_id; + let exhausted = QuotaState::exhausted( + QuotaEvidence::PaymentRequired, + now(), + Some(now() + Duration::from_secs(3600)), + ); + repository + .apply_provider_quota_access(id.as_str(), Revision::new(1).unwrap(), exhausted) + .await + .unwrap(); + let claim = repository + .claim_due( + &ProviderKind::new("openai").unwrap(), + now(), + 1, + Duration::from_secs(60), + ) + .await + .unwrap() + .remove(0); + let mut done = completion(&claim, 3, Ok(observation())); + let recovered_at = now() + Duration::from_secs(1); + done.quota = Some(QuotaState::allowed(recovered_at)); + match case { + "older" => done.quota = Some(QuotaState::allowed(now() - Duration::from_secs(1))), + "equal" => done.quota = Some(QuotaState::allowed(now())), + "new_failure" => { + repository + .apply_provider_quota_access( + id.as_str(), + Revision::new(1).unwrap(), + QuotaState::exhausted( + QuotaEvidence::UsageLimitReached, + now() + Duration::from_secs(2), + None, + ), + ) + .await + .unwrap(); + } + "wrong_claim" => done.claim_token = uuid::Uuid::new_v4(), + "expired_claim" => done.completed_at = now() + Duration::from_secs(60), + "disabled" => { + repository + .set_provider_account_enabled(id.as_str(), false) + .await + .unwrap(); + } + "rotated" => { + sqlx::query("update provider_accounts set credential_revision = credential_revision + 1 where id = $1") + .bind(id.as_str()).execute(&database.pool).await.unwrap(); + } + "proxy_changed" => { + sqlx::query("update provider_accounts set outbound_proxy_url = 'http://127.0.0.1:8080' where id = $1") + .bind(id.as_str()).execute(&database.pool).await.unwrap(); + } + "query_failed" => done.result = Err(UpstreamInfoFailure::Unavailable), + "no_evidence" => done.quota = None, + "unknown" => done.quota = Some(QuotaState::observed_unknown(recovered_at)), + "exhausted" | "old_exhaustion" => { + repository + .apply_provider_quota_access( + id.as_str(), + Revision::new(1).unwrap(), + QuotaState::allowed(if case == "old_exhaustion" { + now() + Duration::from_secs(2) + } else { + now() + }), + ) + .await + .unwrap(); + done.quota = Some(QuotaState::exhausted( + QuotaEvidence::AccountLimitReached, + recovered_at, + Some(now() + Duration::from_secs(3600)), + )); + } + "credential_error" => { + sqlx::query( + "update provider_accounts set credential_state = 'expired' where id = $1", + ) + .bind(id.as_str()) + .execute(&database.pool) + .await + .unwrap(); + } + _ => {} + } + let before = repository.get_account(id).await.unwrap().unwrap(); + let observed_quota = done.quota; + let accepted = repository.complete(done).await.unwrap(); + assert_eq!( + accepted, + !matches!( + case, + "wrong_claim" | "expired_claim" | "disabled" | "rotated" | "proxy_changed" + ), + "{case}" + ); + let after = repository.get_account(id).await.unwrap().unwrap(); + assert_eq!( + after.quota(), + if matches!(case, "recovered" | "credential_error") { + QuotaState::allowed(recovered_at) + } else if case == "exhausted" { + observed_quota.unwrap() + } else { + before.quota() + }, + "{case}" + ); + assert_eq!( + after.credential_state(), + before.credential_state(), + "{case}" + ); + let saved = repository + .read(std::slice::from_ref(id)) + .await + .unwrap() + .remove(0); + assert_eq!( + saved.snapshot.is_some(), + accepted && case != "query_failed", + "{case}" + ); + database.close().await; } } diff --git a/backend/crates/providers/openai/src/credential/selector.rs b/backend/crates/providers/openai/src/credential/selector.rs index cdab17674..ee4cc6843 100644 --- a/backend/crates/providers/openai/src/credential/selector.rs +++ b/backend/crates/providers/openai/src/credential/selector.rs @@ -345,7 +345,7 @@ impl CodexCredentialSelector { .remaining_percent .map(|percent| (percent.clamp(0.0, 100.0) * 100.0).round() as u64) }; - // 互联观测只提供 Smart 额度分，不修改真实错误、冷却或账号可用状态 + // 选号只读取 Smart 额度分，额度恢复由互联同步事务独立提交 Some(AccountQuotaSignals::new(None, rank)) } diff --git a/backend/crates/providers/openai/tests/admin/peer_info.rs b/backend/crates/providers/openai/tests/admin/peer_info.rs index c9172161b..bf89c31a6 100644 --- a/backend/crates/providers/openai/tests/admin/peer_info.rs +++ b/backend/crates/providers/openai/tests/admin/peer_info.rs @@ -1,6 +1,7 @@ //! OpenAI 凭据适配、公共互联服务、Host 传输与调度只读端口的集成验证 use super::*; +use crate::support::MemoryCooldownPort; use gateway_core::{ account::{ LoadedCredential, UpstreamConnection, UpstreamInfoClaim, UpstreamInfoCompletion, @@ -207,38 +208,64 @@ impl UpstreamInfoStore for MemoryUpstreamInfoStore { if self.current(&completion.connection).await.is_err() { return Ok(false); } - let mut claims = self.claims.lock().unwrap(); - let mut records = self.records.lock().unwrap(); - let Some(record) = records.get_mut(&completion.connection.account_id) else { - return Ok(false); - }; - if claims.get(&completion.connection.account_id) != Some(&completion.claim_token) - || record - .plan - .claimed_until - .is_none_or(|expiry| expiry <= completion.completed_at) - || record.plan.connection != completion.connection - || record.plan.request_generation != completion.request_generation + let recovery = completion.quota.filter(|_| completion.result.is_ok()); { - return Ok(false); - } - claims.remove(&completion.connection.account_id); - record.plan.completed_generation = completion.request_generation; - record.plan.claimed_until = None; - record.plan.next_attempt_at = completion.next_attempt_at; - match completion.result { - Ok(snapshot) => { - record.snapshot = Some(snapshot); - record.plan.last_success_at = Some(completion.completed_at); - record.plan.last_error = None; - record.plan.consecutive_failures = 0; + let mut claims = self.claims.lock().unwrap(); + let mut records = self.records.lock().unwrap(); + let Some(record) = records.get_mut(&completion.connection.account_id) else { + return Ok(false); + }; + if claims.get(&completion.connection.account_id) != Some(&completion.claim_token) + || record + .plan + .claimed_until + .is_none_or(|expiry| expiry <= completion.completed_at) + || record.plan.connection != completion.connection + || record.plan.request_generation != completion.request_generation + { + return Ok(false); } - Err(error) => { - record.plan.last_error = Some(error); - record.plan.consecutive_failures += 1; - if error == UpstreamInfoFailure::Unauthorized { - record.snapshot = None; + claims.remove(&completion.connection.account_id); + record.plan.completed_generation = completion.request_generation; + record.plan.claimed_until = None; + record.plan.next_attempt_at = completion.next_attempt_at; + match completion.result { + Ok(snapshot) => { + record.snapshot = Some(snapshot); + record.plan.last_success_at = Some(completion.completed_at); + record.plan.last_error = None; + record.plan.consecutive_failures = 0; } + Err(error) => { + record.plan.last_error = Some(error); + record.plan.consecutive_failures += 1; + if error == UpstreamInfoFailure::Unauthorized { + record.snapshot = None; + } + } + } + } + if let Some(state) = recovery + && state.access() != gateway_core::account::QuotaAccessState::Unknown + && let Some(observed_at) = state.observed_at() + { + let account = self + .accounts + .account(completion.connection.account_id.as_str()) + .unwrap(); + if account.enabled() + && account + .quota() + .observed_at() + .is_none_or(|time| time < observed_at) + { + self.accounts + .apply_quota_access(QuotaAccessChange { + account_id: account.id().clone(), + expected_revision: account.revision(), + state, + }) + .await?; } } Ok(true) @@ -252,12 +279,279 @@ fn peer_document() -> Value { "node":{"id":"synthetic-peer","name":"测试上游","software":"cpr"}, "observed_at":now.to_rfc3339(), "expires_at":(now + chrono::Duration::minutes(5)).to_rfc3339(), "capabilities":["key_usage", "pool_summary"], - "usage":{"status":"available","observed_at":now.to_rfc3339(),"data":{"unit":"USD","daily":{"total":"10","used":"2","remaining":"8","resetsAt":(now + chrono::Duration::days(1)).to_rfc3339()},"weekly":{"total":"50","used":"5","remaining":"45","resetsAt":(now + chrono::Duration::days(7)).to_rfc3339()}}}, + "usage":{"status":"available","observed_at":now.to_rfc3339(),"data":{"unit":"USD","allowed":true,"daily":{"limit_reached":false,"total":"10","used":"2","remaining":"8","resetsAt":(now + chrono::Duration::days(1)).to_rfc3339()},"weekly":{"limit_reached":false,"total":"50","used":"5","remaining":"45","resetsAt":(now + chrono::Duration::days(7)).to_rfc3339()}}}, "pool":{"disclosure":"summary","account_count":2,"states":{"ready":2,"cooldown":0,"exhausted":0,"disabled":0,"error":0,"unknown":0},"plans":[{"name":"plus","count":2}]}, "quota":{"disclosure":"none"},"models":[] }) } +async fn exhaust_account(accounts: &MemoryAccountStore, id: &ProviderAccountId, at: SystemTime) { + let account = accounts.account(id.as_str()).unwrap(); + accounts + .apply_quota_access(QuotaAccessChange { + account_id: id.clone(), + expected_revision: account.revision(), + state: QuotaState::exhausted(QuotaEvidence::UsageLimitReached, at, None), + }) + .await + .unwrap(); +} + +#[tokio::test] +async fn peer_sync_restores_exhausted_accounts_only_from_current_key_budget_facts() { + for case in [ + "available", + "unlimited", + "budget_exhausted", + "weekly_exhausted", + "expired_budget", + "missing_remaining", + "missing_access", + "missing_allowed", + "undisclosed_access", + "amounts_disagree", + "usage_unknown", + "budget_only", + "old_usage", + "old_snapshot", + "unauthorized", + "unavailable", + ] { + let server = MockServer::start().await; + let (mut bundle, accounts, _, id, _config) = fixture(&server).await; + let exhausted_at = SystemTime::now() - Duration::from_secs(60); + exhaust_account(&accounts, &id, exhausted_at).await; + let mut document = peer_document(); + let older = DateTime::::from(exhausted_at - Duration::from_secs(1)).to_rfc3339(); + match case { + "unlimited" => { + for window in ["daily", "weekly"] { + document["usage"]["data"][window] = json!({ + "total":null, "remaining":null, "used":"7", "resetsAt":null + }); + } + } + "budget_exhausted" => { + document["usage"]["data"]["allowed"] = json!(false); + document["usage"]["data"]["daily"]["limit_reached"] = json!(true); + } + "weekly_exhausted" => { + document["usage"]["data"] + .as_object_mut() + .unwrap() + .remove("allowed"); + document["usage"]["data"]["weekly"]["limit_reached"] = json!(true); + } + "expired_budget" => { + document["usage"]["data"] + .as_object_mut() + .unwrap() + .remove("allowed"); + document["usage"]["data"]["daily"]["limit_reached"] = json!(true); + document["usage"]["data"]["daily"]["resetsAt"] = json!(older); + } + "missing_allowed" => { + document["usage"]["data"] + .as_object_mut() + .unwrap() + .remove("allowed"); + } + "undisclosed_access" => document["capabilities"] = json!(["pool_summary"]), + "missing_access" => { + document["usage"]["data"] + .as_object_mut() + .unwrap() + .remove("allowed"); + for window in ["daily", "weekly"] { + document["usage"]["data"][window] + .as_object_mut() + .unwrap() + .remove("limit_reached"); + } + } + "amounts_disagree" => { + document["usage"]["data"]["daily"]["remaining"] = json!("0"); + document["usage"]["data"]["daily"]["used"] = json!("10"); + } + "missing_remaining" => document["usage"]["data"]["daily"]["remaining"] = Value::Null, + "usage_unknown" => document["usage"]["status"] = json!("unknown"), + "budget_only" => document["capabilities"] = json!(["key_usage"]), + "old_usage" => document["usage"]["observed_at"] = json!(older), + "old_snapshot" => document["observed_at"] = json!(older), + _ => {} + } + let status = match case { + "unauthorized" => 403, + "unavailable" => 503, + _ => 200, + }; + Mock::given(method("GET")) + .and(path("/custom/v1/cpr/link")) + .respond_with(ResponseTemplate::new(status).set_body_json(document)) + .expect(1) + .mount(&server) + .await; + cycle(&worker(&mut bundle)).await; + let account = accounts.account(id.as_str()).unwrap(); + assert_eq!( + account.quota().is_exhausted(), + !matches!( + case, + "available" + | "unlimited" + | "budget_only" + | "missing_remaining" + | "amounts_disagree" + ), + "{case}" + ); + assert_eq!(account.credential_state(), CredentialState::Ready); + assert!( + server + .received_requests() + .await + .unwrap() + .iter() + .all(|request| request.method == "GET") + ); + } +} + +#[tokio::test] +async fn peer_sync_cannot_clear_an_exhaustion_recorded_while_the_request_is_in_flight() { + let server = MockServer::start().await; + let (mut bundle, accounts, _, id, _config) = fixture(&server).await; + exhaust_account(&accounts, &id, SystemTime::now() - Duration::from_secs(60)).await; + let mut document = peer_document(); + // 上游时钟略快时也必须按本地请求开始时间保护并发失败 + let future = (Utc::now() + chrono::Duration::seconds(10)).to_rfc3339(); + document["observed_at"] = json!(future); + document["usage"]["observed_at"] = json!(future); + Mock::given(method("GET")) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(document) + .set_delay(Duration::from_millis(200)), + ) + .mount(&server) + .await; + let worker = Arc::new(worker(&mut bundle)); + let running = tokio::spawn({ + let worker = worker.clone(); + async move { cycle(&worker).await } + }); + wait_for_request(&server).await; + let failed_at = SystemTime::now(); + exhaust_account(&accounts, &id, failed_at).await; + running.await.unwrap(); + let account = accounts.account(id.as_str()).unwrap(); + assert_eq!( + account.quota(), + QuotaState::exhausted(QuotaEvidence::UsageLimitReached, failed_at, None) + ); +} + +#[tokio::test] +async fn key_budget_rejections_recover_through_sync_while_ordinary_429_keeps_its_cooldown() { + use gateway_core::error::ProviderErrorKind; + + for code in [ + "key_daily_budget_exceeded", + "key_weekly_budget_exceeded", + "rate_limit_exceeded", + ] { + let budget = code != "rate_limit_exceeded"; + let server = MockServer::start().await; + let (mut bundle, accounts, _, id, _config) = fixture(&server).await; + Mock::given(method("POST")) + .and(path("/custom/v1/responses")) + .respond_with(ResponseTemplate::new(429).insert_header("retry-after", "86400") + .set_body_json(json!({"error":{"code":code,"type":if budget { "usage_limit_reached" } else { "rate_limit_error" },"message":"limit reached"}}))) + .expect(1).mount(&server).await; + let mut response = execute_peer_request(&bundle, "round_robin").await.unwrap(); + let failure = loop { + match response.next().await { + Some(Err(error)) => break error, + Some(Ok(_)) => {} + None => panic!("expected rejection"), + } + }; + drop(response); + assert_eq!( + failure.kind(), + if budget { + ProviderErrorKind::QuotaExhausted + } else { + ProviderErrorKind::RateLimited + }, + "{code}" + ); + assert_eq!( + accounts + .account(id.as_str()) + .unwrap() + .quota() + .is_exhausted(), + budget, + "{code}" + ); + let cooldown = bundle.cooldowns.read(&id).await.unwrap(); + assert_eq!(cooldown.is_some(), !budget, "{code}"); + server.verify().await; + server.reset().await; + + Mock::given(method("GET")) + .and(path("/custom/v1/cpr/link")) + .respond_with(ResponseTemplate::new(200).set_body_json(peer_document())) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/custom/v1/responses")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(COMPLETED_SESSION_SSE), + ) + .expect(if budget { 1 } else { 0 }) + .mount(&server) + .await; + cycle(&worker(&mut bundle)).await; + assert!( + !accounts + .account(id.as_str()) + .unwrap() + .quota() + .is_exhausted(), + "{code}" + ); + assert_eq!( + bundle.cooldowns.read(&id).await.unwrap(), + cooldown, + "{code}" + ); + if budget { + let mut response = execute_peer_request(&bundle, "round_robin").await.unwrap(); + while let Some(event) = response.next().await { + event.unwrap(); + } + } else { + let error = match execute_peer_request(&bundle, "round_robin").await { + Err(error) => error, + Ok(mut response) => loop { + match response.next().await { + Some(Err(error)) => break error, + Some(Ok(_)) => {} + None => panic!("cooldown must prevent sending"), + } + }, + }; + assert_eq!(error.kind(), ProviderErrorKind::NoEligibleAccount); + } + server.verify().await; + } +} + async fn fixture( server: &MockServer, ) -> ( @@ -680,11 +974,13 @@ async fn invalid_or_redirected_responses_never_replace_the_last_valid_snapshot() } #[tokio::test] -async fn smart_prefers_another_account_when_key_budget_is_exhausted_without_pool_disclosure() { +async fn peer_sync_excludes_budget_exhausted_accounts_without_pool_disclosure() { let server = MockServer::start().await; let mut peer = peer_document(); peer["capabilities"] = json!(["key_usage"]); peer["pool"] = json!({"disclosure":"none"}); + peer["usage"]["data"]["allowed"] = json!(false); + peer["usage"]["data"]["daily"]["limit_reached"] = json!(true); peer["usage"]["data"]["daily"]["used"] = json!("10.0000"); peer["usage"]["data"]["daily"]["remaining"] = json!("0.0000"); Mock::given(method("GET")) @@ -718,72 +1014,14 @@ async fn smart_prefers_another_account_when_key_budget_is_exhausted_without_pool ) .await; cycle(&worker(&mut bundle)).await; - let provider = ProviderKind::new("openai").unwrap(); - let model = UpstreamModelId::new("gpt-5.4").unwrap(); - let scope = Arc::new(FrozenAccountScope::new( - Arc::new(RuntimeAccountDirectory::new(BTreeMap::from([ - ( - ProviderAccountId::new("acct_peer").unwrap(), - RuntimeAccount::new(provider.clone(), BTreeSet::new()), - ), - ( - ProviderAccountId::new("acct_available").unwrap(), - RuntimeAccount::new(provider.clone(), BTreeSet::new()), - ), - ]))), - ClientRoutingScope::all_accounts(), - )); - let operation = Operation::Generate(GenerateRequest::from_protocol_payload( - ProtocolPayload::json_object( - "openai", - json!({"model":"gpt-5.4", "input":"budget routing"}) - .as_object() - .unwrap() - .clone(), - ) - .unwrap(), - )); - let snapshot = RuntimeSnapshot::new( - ConfigRevision::new(1).unwrap(), - gateway_core::settings::SettingsValues::new(2, "smart", Default::default(), None, None), - vec![provider.clone()], - vec![ProviderModel::new( - provider, - model, - ModelCapabilities::new(BTreeSet::from([operation.kind()]), Some(32_000)) - .with_upstream_feature_validation(), - )], - Vec::new(), - ) - .unwrap(); - let plan = snapshot - .plan( - &PublicModelId::new("gpt-5.4").unwrap(), - &operation, - scope.clone(), - &RoutingContext::default(), - ) - .unwrap(); - let context = AttemptContext::new( - RequestAttemptContext::new( - ModelRequestId::new("req_peer_smart").unwrap(), - ClientApiKeyId::new("key_peer_smart").unwrap(), - ), - NonZeroU32::new(1).unwrap(), - SystemTime::now() + Duration::from_secs(5), - account_policy(), - AccountAttemptContext::new(BTreeSet::new(), None, None).with_account_scope(scope), - None, - CancellationToken::new(), + assert!( + accounts + .account("acct_peer") + .unwrap() + .quota() + .is_exhausted() ); - let mut response = bundle - .core_provider() - .execute( - ProviderRequest::new(operation, plan.candidates()[0].clone()), - context, - ) - .await - .unwrap(); + let mut response = execute_peer_request(&bundle, "smart").await.unwrap(); assert_eq!( response.metadata().provider_account_id().as_str(), "acct_available" @@ -1046,9 +1284,84 @@ async fn cold_lists_bound_full_snapshot_batches_and_reuse_every_hydrated_summary ); } +async fn execute_peer_request( + bundle: &PeerFixture, + strategy: &str, +) -> Result { + let provider = ProviderKind::new("openai").unwrap(); + let model = UpstreamModelId::new("gpt-5.4").unwrap(); + let scope = Arc::new(FrozenAccountScope::new( + Arc::new(RuntimeAccountDirectory::new(BTreeMap::from([ + ( + ProviderAccountId::new("acct_peer").unwrap(), + RuntimeAccount::new(provider.clone(), BTreeSet::new()), + ), + ( + ProviderAccountId::new("acct_available").unwrap(), + RuntimeAccount::new(provider.clone(), BTreeSet::new()), + ), + ]))), + ClientRoutingScope::all_accounts(), + )); + let operation = Operation::Generate(GenerateRequest::from_protocol_payload( + ProtocolPayload::json_object( + "openai", + json!({"model":"gpt-5.4", "input":"budget routing"}) + .as_object() + .unwrap() + .clone(), + ) + .unwrap(), + )); + let snapshot = RuntimeSnapshot::new( + ConfigRevision::new(1).unwrap(), + gateway_core::settings::SettingsValues::new(2, strategy, Default::default(), None, None), + vec![provider.clone()], + vec![ProviderModel::new( + provider, + model, + ModelCapabilities::new(BTreeSet::from([operation.kind()]), Some(32_000)) + .with_upstream_feature_validation(), + )], + Vec::new(), + ) + .unwrap(); + let plan = snapshot + .plan( + &PublicModelId::new("gpt-5.4").unwrap(), + &operation, + scope.clone(), + &RoutingContext::default(), + ) + .unwrap(); + let context = AttemptContext::new( + RequestAttemptContext::new( + ModelRequestId::new("req_peer_smart").unwrap(), + ClientApiKeyId::new("key_peer_smart").unwrap(), + ), + NonZeroU32::new(1).unwrap(), + SystemTime::now() + Duration::from_secs(5), + gateway_core::account::AccountSelectionPolicy::new( + gateway_core::account::RotationStrategy::parse(strategy).unwrap(), + NonZeroU32::new(2).unwrap(), + ), + AccountAttemptContext::new(BTreeSet::new(), None, None).with_account_scope(scope), + None, + CancellationToken::new(), + ); + bundle + .core_provider() + .execute( + ProviderRequest::new(operation, plan.candidates()[0].clone()), + context, + ) + .await +} + struct PeerFixture { provider: provider_openai::ProviderBundle, peer: gateway_admin::PeerUpstreamBundle, + cooldowns: Arc, } impl PeerFixture { @@ -1071,12 +1384,21 @@ async fn peer_fixture( store: Arc, ) -> PeerFixture { let reader = Arc::new(gateway_core::account::UpstreamInfoReader::default()); + let cooldowns = Arc::new(MemoryCooldownPort::new()); let provider = provider_openai::initialize( config, - provider_ports_with_catalog( + ProviderStorePorts::new( accounts.clone(), - Arc::new(TestOAuthPending::default()), + Arc::new(TestLeaseCoordinator::default()), + Arc::new(MemorySessionAffinity::default()), + Arc::new(MemorySessionExclusions::default()), Arc::new(TestCatalogCache::default()), + Arc::new(TestArtifactProfiles), + Arc::new(TestCredentialState), + cooldowns.clone(), + Arc::new(TestRuntimePolicy), + Arc::new(TestOAuthPending::default()), + Arc::new(RecordingDiagnostics::default()), ), reader.clone(), ) @@ -1095,5 +1417,9 @@ async fn peer_fixture( ) .unwrap(); reader.bind(&peer.scheduling_source()).unwrap(); - PeerFixture { provider, peer } + PeerFixture { + provider, + peer, + cooldowns, + } } diff --git a/deploy/README.md b/deploy/README.md index b4867a6c9..b8c4c352f 100644 --- a/deploy/README.md +++ b/deploy/README.md @@ -247,7 +247,9 @@ PostgreSQL/Redis 启动密码。日常校验使用 `config --quiet` 列表与详情显示已保存的观测，手动刷新提交后台任务；旧观测和查询失败分别提示，不要求上游实时可用 第三方网关通过插件实现相同协议也可被识别，不按软件名称判定兼容。 -额度观测仅作为 Smart 选择的参考，不能保证下一次请求成功；协议和公开边界见 [CPR 互联 API](../docs/api.md#cpr-互联只读协议) +上游明确返回 Key 预算耗尽时，下游暂停该账号；后续同步确认允许使用后自动恢复，普通限流冷却独立保留。 +金额只供展示，缺少访问状态时不根据余额推断恢复。账号池额度观测作为 Smart 选择的参考，不能保证下一次请求成功； +协议和公开边界见 [CPR 互联 API](../docs/api.md#cpr-互联只读协议) ### 批量维护代理 diff --git a/docs/api.md b/docs/api.md index c118c254e..e845bd9a4 100644 --- a/docs/api.md +++ b/docs/api.md @@ -680,6 +680,10 @@ SSE/WS 的 `response.failed` 保留原消息、响应 ID 与其他业务字段 明确额度耗尽触发账号隔离与安全换号，包括 WebSocket 握手返回的 429；不会因其长 `Retry-After` 而转入同账号传输恢复等待 +Client Key 日周预算超限返回 `429`，`error.type` 为官方额度耗尽类型 `usage_limit_reached`， +`error.code` 分别保留 `key_daily_budget_exceeded` 与 `key_weekly_budget_exceeded`，`Retry-After` 指向预算重置边界。 +下游按官方错误类型确认额度耗尽，由互联预算同步确认恢复，不写普通限流冷却；普通 429 限流仍独立处理 + OpenAI 选号阶段确认本次可选账号全部额度耗尽时，HTTP 返回 `429`，WebSocket 错误帧返回 `status: 429`，两者的 `error.type` 与 `error.code` 均为 `usage_limit_reached`，提示客户端停止 本轮自动重试，等待额度恢复或补充可用账号。空账号池、认证失效和临时容量不足不按额度耗尽处理； @@ -742,8 +746,9 @@ OpenAI 选号阶段确认本次可选账号全部额度耗尽时，HTTP 返回 ` "observed_at": "2026-10-29T12:00:00Z", "data": { "unit": "USD", - "daily": { "total": "1", "used": "0.6", "remaining": "0.4", "resetsAt": "2026-10-29T16:00:00Z" }, - "weekly": { "total": null, "used": "2", "remaining": null, "resetsAt": null } + "allowed": true, + "daily": { "limit_reached": false, "total": "1", "used": "0.6", "remaining": "0.4", "resetsAt": "2026-10-29T16:00:00Z" }, + "weekly": { "limit_reached": false, "total": null, "used": "2", "remaining": null, "resetsAt": null } } }, "pool": { "disclosure": "none" }, @@ -754,6 +759,10 @@ OpenAI 选号阶段确认本次可选账号全部额度耗尽时，HTTP 返回 ` `usage.data` 沿用 [API Key 额度查询](#api-key-额度查询) 的十进制金额、窗口及不限额语义。 `usage.status` 可为 `available`、`unavailable` 或 `unknown`；没有可用数据时不把缺失值解释为零余额 +`usage.data.allowed` 是上游根据本地预算账本给出的访问结论，日周窗口的 `limit_reached` 表示对应预算是否耗尽。 +接入端以明确的 `allowed` 为准；缺失时，仍有效的 `limit_reached: true` 只能证明耗尽，不能据此推断恢复。 +缺少这些状态字段时金额仅用于展示，接入端不从 `total`、`used` 或 `remaining` 计算账号资格 + | 能力 | 数据合同 | 公开条件 | | --- | --- | --- | | `key_usage` | `usage` 为当前认证 Key 的本地预算账本 | `enabled` | @@ -1089,6 +1098,12 @@ Key 消费窗口的 `localUsage` 按各自 `[resetAt - windowSeconds, resetAt)` 网络错误、限流或无效响应保留同一连接的旧观测；401/403 或连接身份变化清除旧观测。 单个账号停用时停止同步并清理互联观测，重新启用后建立计划 +成功同步后，明确的预算访问结论决定账号是否进入 `quota_exhausted`； +有效的 `usage.data.allowed: true` 自动解除更早的耗尽，恢复参与调度。 +Key 预算及其访问状态随互联默认公开，不要求开启 `shareQuota`。 +快照和预算观测必须仍有效且晚于耗尽事实，重置时间到期、金额变化、缺失状态或查询失败均不代表恢复。 +恢复只改变额度状态，凭据错误和独立的 `rate_limited` 冷却仍按各自规则处理 + ### 账号模型限制 账号列表和详情返回 `modelAccess: { mode, models }`，模型 ID 区分大小写并精确匹配： diff --git a/docs/architecture.md b/docs/architecture.md index 52ae7c84e..e9d18bf12 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -1009,8 +1009,11 @@ Admin 为账号聚合设置独立并发上限，满额立即返回可重试的 Admin 的 `peer_link` 分别拥有共享导出与上游消费用例；上游消费统一处理发现、协议校验、同步、展示和调度投影。 Provider 声明可连接的认证种类，并从本轮凭据中提取连接材料；当前 OpenAI API Key 提供该适配，xAI 原生账号继续参与共享导出。 -Host 执行有界 HTTP 请求，复用组合根注入的证书策略和账号出口；Store 只持久化不透明观测及同步计划，不按厂商或认证种类硬编码候选。 +Host 执行有界 HTTP 请求，复用组合根注入的证书策略和账号出口；Store 持久化不透明观测、同步计划及控制面确认的额度事实，不按厂商或认证种类硬编码候选。 计划查询与领取只读元数据，每次领取带独立 token，删除重建计划后的旧执行者不能提交观测。 +共享端按 Core 的预算规则从本地账本生成明确的访问状态，与本地准入共用口径；金额字段只供接入端展示。 +接入端 Admin 的独立额度模块解释默认公开的访问状态，展示、调度投影与同步共用，不要求公开账号池额度；未知或过期事实保留已有状态，Provider 不新增 CPR 错误码分类。 +Store 在同一事务内核对连接身份、启用状态、领取代次与 token，原子提交快照并复用通用额度写入；观测时间受本地请求开始时间限制，迟到响应不能覆盖并发失败，额度同步不改变凭据或独立冷却。 Admin 按连接身份及完成版本缓存管理摘要与调度投影，完整观测分批读取，解析和投影在锁外完成。 Core 定义中立调度信号与只读端口，组合根绑定观测来源；Provider 选号不持有控制面服务，弱引用避免互相保活。 账号和凭据提交后的共用通知同时失效原生 Provider 与互联派生状态，插件账号写入复用同一流程。