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(peer): include quota-backed plans in pool summaries

Browse Source 
Reuse account plan resolution when exporting peer pool summaries. Read saved quota observations for missing plans even when quota sharing is disabled, without exposing quota data or mutating account records.

Cover unknown and missing plans, subtype preservation, plan-only limits, and quota refresh from an unknown account plan.

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Zyy
2026-11-02 12:24:57 +08:00
co-authored by Codex
parent 7b05b14a2b
commit 6f71b27c3c
5 changed files with 149 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/crates/gateway-admin/src/use_case/peer_link/mod.rs
+14 -5
View File
@@ -15,7 +15,7 @@ use crate::{
model::{
AdminError,
peer_link::{PeerLinkExport, PeerLinkLimits},
provider_credentials::ProviderQuotaReadRequest,
provider_credentials::{ProviderQuotaReadRequest, explicit_plan_type},
},
ports::{
provider::ProviderAdminRegistry,
@@ -175,10 +175,16 @@ impl PeerLinkService for DefaultPeerLinkService {
.await
.map_err(|error| map_store_error(error, "peer account runtime"))?;
let mut quotas = BTreeMap::new();
if sharing.share_quota {
if sharing.share_quota || sharing.share_plans {
let mut window_count = 0usize;
let mut by_provider = BTreeMap::new();
for account in &accounts {
// 仅公开套餐时只补读缺失套餐的本地额度，不回源刷新或扩大公开范围
if !sharing.share_quota
&& explicit_plan_type(account.plan_type.as_deref()).is_some()
{
continue;
}
by_provider
.entry(&account.provider_kind)
.or_insert_with(Vec::new)
@@ -200,9 +206,11 @@ impl PeerLinkService for DefaultPeerLinkService {
.zip(provider.read_quotas(&requests).await)
{
if let Ok(quota) = quota {
window_count = window_count.saturating_add(quota.windows.len());
if window_count > limits.max_quota_windows {
return Err(PeerLinkError::SnapshotTooLarge);
if sharing.share_quota {
window_count = window_count.saturating_add(quota.windows.len());
if window_count > limits.max_quota_windows {
return Err(PeerLinkError::SnapshotTooLarge);
}
}
quotas.insert(account.id.clone(), quota);
}
@@ -215,6 +223,7 @@ impl PeerLinkService for DefaultPeerLinkService {
if sharing.share_pool {
result.pool = Some(projection::pool_summary(
&accounts,
&quotas,
&runtime,
&self.providers,
sharing.share_plans,
backend/crates/gateway-admin/src/use_case/peer_link/projection.rs
+6 -4
View File
@@ -25,6 +25,7 @@ use crate::{
pub(super) fn pool_summary(
accounts: &[AccountRecord],
quotas: &BTreeMap<String, ProviderQuota>,
runtime: &AccountRuntimeSnapshot,
providers: &ProviderAdminRegistry,
share_plans: bool,
@@ -43,12 +44,13 @@ pub(super) fn pool_summary(
AccountStatus::Error => summary.error += 1,
}
if let Some(plans) = &mut summary.plans
&& let (Some(plan), Ok(provider)) = (
&account.plan_type,
providers.require(&account.provider_kind),
&& let Some(plan) = providers.resolve_account_plan(
account.provider_kind.as_str(),
&mut account.plan_type.clone(),
quotas.get(&account.id),
)
{
*plans.entry(provider.plan_type_display(plan)).or_default() += 1;
*plans.entry(plan).or_default() += 1;
}
}
summary
backend/crates/gateway-api/tests/openai/peer_link/mod.rs
+108
View File
@@ -261,6 +261,114 @@ async fn exports_scoped_observations_even_when_the_key_budget_is_exhausted_witho
assert!(fixture.observations.lock().unwrap().summaries.is_empty());
}
#[tokio::test]
async fn pool_plans_use_account_plan_fallback_independently_of_quota_sharing() {
for share_quota in [true, false] {
let (fixture, app) = fixture(
ClientPeerSharing {
share_plans: true,
share_quota,
..sharing()
},
RAW_KEY,
)
.await;
for (account_plan, quota_plan, expected) in [
(None, Some("free"), Some("Free")),
(Some("unknown"), Some("free"), Some("Free")),
(Some(" "), Some("free"), Some("Free")),
(Some("edu_pro"), Some("edu"), Some("EduPro")),
(Some("free"), None, Some("Free")),
(None, None, None),
(None, Some("unknown"), None),
] {
fixture
.account
.lock()
.unwrap()
.as_mut()
.unwrap()
.account
.plan_type = account_plan.map(str::to_owned);
fixture
.provider_quota
.lock()
.unwrap()
.as_mut()
.unwrap()
.plan_type = quota_plan.map(str::to_owned);
let response = query(&app, PATH, Some(RAW_KEY)).await;
assert_eq!(response.status(), StatusCode::OK);
let body = response_json(response).await;
let plans =
expected.map_or_else(|| json!([]), |name| json!([{ "name": name, "count": 1 }]));
assert_eq!(
body["pool"]["plans"], plans,
"{account_plan:?} / {quota_plan:?}"
);
assert_eq!(body["pool"]["account_count"], 1);
assert_eq!(body.get("quota").is_some(), share_quota);
assert_eq!(
fixture
.account
.lock()
.unwrap()
.as_ref()
.unwrap()
.account
.plan_type
.as_deref(),
account_plan,
);
}
}
}
#[tokio::test]
async fn plan_only_summary_limits_exported_data_and_preserves_unknown_on_quota_read_failure() {
let (fixture, app) = fixture(
ClientPeerSharing {
share_plans: true,
share_quota: false,
..sharing()
},
RAW_KEY,
)
.await;
fixture
.account
.lock()
.unwrap()
.as_mut()
.unwrap()
.account
.plan_type = None;
{
let mut quota = fixture.provider_quota.lock().unwrap();
let quota = quota.as_mut().unwrap();
quota.plan_type = Some("free".into());
quota.windows =
vec![quota.windows[0].clone(); gateway_protocol::peer::MAX_QUOTA_WINDOWS + 1];
}
let response = query(&app, PATH, Some(RAW_KEY)).await;
assert_eq!(response.status(), StatusCode::OK);
let body = response_json(response).await;
assert_eq!(
body["pool"]["plans"],
json!([{ "name": "Free", "count": 1 }])
);
assert_eq!(body["capabilities"], json!(["key_usage", "pool_summary"]));
assert!(body.get("quota").is_none());
*fixture.provider_quota.lock().unwrap() = None;
let response = query(&app, PATH, Some(RAW_KEY)).await;
assert_eq!(response.status(), StatusCode::OK);
let body = response_json(response).await;
assert_eq!(body["pool"]["account_count"], 1);
assert_eq!(body["pool"]["plans"], json!([]));
}
#[tokio::test]
async fn requires_authentication_explicit_sharing_and_bound_groups_and_rejects_query_overrides() {
let (fixture, app) = fixture(sharing(), RAW_KEY).await;
backend/crates/providers/openai/tests/credential/quota/mod.rs
+18 -17
View File
@@ -410,42 +410,43 @@ async fn websocket_additional_rate_limit_resolves_by_name_without_touching_core(
#[tokio::test]
async fn quota_refresh_synchronizes_plan_changes_without_losing_subtypes() {
for (current_plan, observed_plan, expected) in [
("plus", Some("pro"), "pro"),
("pro", Some("plus"), "plus"),
("plus", Some("free"), "free"),
("unknown", Some("pro"), "pro"),
("plus", Some(" PRO "), "pro"),
("pro", None, "pro"),
("pro", Some("unknown"), "pro"),
("pro", Some(" "), "pro"),
(None, Some("free"), "free"),
(Some("plus"), Some("pro"), "pro"),
(Some("pro"), Some("plus"), "plus"),
(Some("plus"), Some("free"), "free"),
(Some("unknown"), Some("pro"), "pro"),
(Some("plus"), Some(" PRO "), "pro"),
(Some("pro"), None, "pro"),
(Some("pro"), Some("unknown"), "pro"),
(Some("pro"), Some(" "), "pro"),
(
"self_serve_business_prolite",
Some("self_serve_business_prolite"),
Some("team"),
"self_serve_business_prolite",
),
(
"self_serve_business_usage_based",
Some("self_serve_business_usage_based"),
Some("team"),
"self_serve_business_usage_based",
),
(
"enterprise_cbp_usage_based",
Some("enterprise_cbp_usage_based"),
Some("business"),
"enterprise_cbp_usage_based",
),
("edu_plus", Some("edu"), "edu_plus"),
("self_serve_business_prolite", Some("pro"), "pro"),
(Some("edu_plus"), Some("edu"), "edu_plus"),
(Some("self_serve_business_prolite"), Some("pro"), "pro"),
(
"team",
Some("team"),
Some("self_serve_business_prolite"),
"self_serve_business_prolite",
),
("plus", Some("future_plan"), "future_plan"),
(Some("plus"), Some("future_plan"), "future_plan"),
] {
let store = Arc::new(MemoryAccountStore::default());
let account_id = "acct_plan_refresh";
let mut verified_account = profile("chatgpt-plan-refresh");
verified_account.plan_type = Some(current_plan.to_owned());
verified_account.plan_type = current_plan.map(str::to_owned);
store
.seed_oauth_credential(ImportCodexOAuthCredential {
account_id: account_id.to_owned(),
@@ -480,7 +481,7 @@ async fn quota_refresh_synchronizes_plan_changes_without_losing_subtypes() {
assert_eq!(
after.plan_type(),
Some(expected),
"{current_plan} -> {observed_plan:?}"
"{current_plan:?} -> {observed_plan:?}"
);
assert_eq!(after.revision(), before.revision());
assert_eq!(after.email(), before.email());
docs/api.md
+3
View File
@@ -864,6 +864,9 @@ CPR 默认路径为 `/v1/cpr/link`，使用同一数据面入口认证，不接
| `sharePlans` | 已公开账号池中的 `pool.plans` |
| `shareQuota` | `quota_windows` 与 `model_availability` |
套餐分布沿用账号页口径：优先使用账号已有的明确套餐，缺失时从已保存的额度快照补全，无需开启 `shareQuota`。
仍未知的账号只计入账号总数，不推断为 Free
共享范围同时受当前 Key 的公开开关、分组及冻结权限限制，只包含 Provider 明确声明的直连 OAuth 账号。
互联 Key 的推理也使用该直连范围，不把 API Key 转接账号再次作为共享池，不递归查询其他节点。
这些开关和账号组织方式属于 CPR 内置实现；第三方通过自身授权规则生成协议数据