mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-06 08:06:19 +00:00
fix(accounts): recover snapshot races and support unlimited concurrency (#225)
Retry account snapshot conflicts before reserving scheduling capacity. Allow zero as an unlimited default account concurrency limit, preserving per-account overrides, lease accounting, and request intervals. Keep connection-test diagnostics expanded within the modal scroll area. Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
@@ -38,7 +38,8 @@ pub struct AccountGroupAccountSummary {
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct AccountGroupCapacity {
|
||||
pub used_slots: Option<u64>,
|
||||
pub total_slots: u64,
|
||||
/// `None` 表示可用账号中存在不限制并发的账号。
|
||||
pub total_slots: Option<u64>,
|
||||
}
|
||||
|
||||
/// Successful, downstream-committed USD request costs for group accounts.
|
||||
@@ -55,7 +56,7 @@ pub struct AccountGroupMemberFact {
|
||||
pub group_id: AccountGroupId,
|
||||
pub account_id: String,
|
||||
pub status: AccountStatusFacts,
|
||||
pub total_slots: u64,
|
||||
pub total_slots: Option<u64>,
|
||||
}
|
||||
|
||||
/// Lightweight group reference embedded in account and client-key views.
|
||||
|
||||
@@ -992,7 +992,8 @@ pub struct DashboardWireAttribute {
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct DashboardCapacity {
|
||||
pub max_concurrent_per_account: u64,
|
||||
pub total_slots: u64,
|
||||
/// `None` 表示可用账号池不限制并发。
|
||||
pub total_slots: Option<u64>,
|
||||
pub used_slots: Option<u64>,
|
||||
pub available_slots: Option<u64>,
|
||||
}
|
||||
|
||||
@@ -231,8 +231,8 @@ fn project_group_runtime(
|
||||
})
|
||||
})
|
||||
},
|
||||
total_slots: available_ids.iter().fold(0_u64, |sum, account_id| {
|
||||
sum.saturating_add(slots_by_account.get(*account_id).copied().unwrap_or(0))
|
||||
total_slots: available_ids.iter().try_fold(0_u64, |sum, account_id| {
|
||||
Some(sum.saturating_add(slots_by_account.get(*account_id).copied().flatten()?))
|
||||
}),
|
||||
};
|
||||
}
|
||||
|
||||
@@ -207,27 +207,25 @@ impl DefaultObservabilityService {
|
||||
}
|
||||
let wire_profiles = self.providers.dashboard_wire_profiles(&configurations);
|
||||
let max_concurrent_per_account = u64::from(settings.max_concurrent_per_account);
|
||||
let total_slots = runtime_slots.as_ref().map_or_else(
|
||||
|| {
|
||||
observation
|
||||
.provider_accounts
|
||||
.normal
|
||||
.saturating_mul(max_concurrent_per_account)
|
||||
},
|
||||
|slots| {
|
||||
slots
|
||||
.inherited_accounts
|
||||
.saturating_mul(max_concurrent_per_account)
|
||||
.saturating_add(slots.overridden_slots)
|
||||
},
|
||||
);
|
||||
let (inherited_accounts, overridden_slots) = runtime_slots
|
||||
.as_ref()
|
||||
.map_or((observation.provider_accounts.normal, 0), |slots| {
|
||||
(slots.inherited_accounts, slots.overridden_slots)
|
||||
});
|
||||
let total_slots = (max_concurrent_per_account > 0 || inherited_accounts == 0).then(|| {
|
||||
inherited_accounts
|
||||
.saturating_mul(max_concurrent_per_account)
|
||||
.saturating_add(overridden_slots)
|
||||
});
|
||||
let used_slots = runtime_slots.and_then(|slots| slots.used_slots);
|
||||
Ok(DashboardResult {
|
||||
capacity: DashboardCapacity {
|
||||
max_concurrent_per_account,
|
||||
total_slots,
|
||||
used_slots,
|
||||
available_slots: used_slots.map(|used| total_slots.saturating_sub(used)),
|
||||
available_slots: used_slots
|
||||
.zip(total_slots)
|
||||
.map(|(used, total)| total.saturating_sub(used)),
|
||||
},
|
||||
rotation_strategy: settings.rotation_strategy,
|
||||
observation,
|
||||
|
||||
@@ -381,7 +381,6 @@ fn validate_settings(command: &ReplaceRuntimeSettings) -> Result<(), AdminError>
|
||||
&& isize::try_from(command.responses_max_decompressed_body_bytes).is_ok()
|
||||
&& command.refresh_margin_seconds > 0
|
||||
&& command.refresh_concurrency > 0
|
||||
&& command.max_concurrent_per_account > 0
|
||||
&& command.max_waiting_per_key <= 1_000
|
||||
&& command.max_waiting_per_account <= 1_000
|
||||
&& (1..=120).contains(&command.concurrency_wait_timeout_seconds)
|
||||
|
||||
@@ -82,7 +82,7 @@ async fn group_query_service_enriches_only_current_page_members_with_runtime_fac
|
||||
group.capacity,
|
||||
AccountGroupCapacity {
|
||||
used_slots: Some(2),
|
||||
total_slots: 4,
|
||||
total_slots: Some(4),
|
||||
}
|
||||
);
|
||||
}
|
||||
@@ -90,6 +90,40 @@ async fn group_query_service_enriches_only_current_page_members_with_runtime_fac
|
||||
#[derive(Default)]
|
||||
struct FakeGroupStore {
|
||||
requested_groups: Mutex<Vec<String>>,
|
||||
members: Option<Vec<AccountGroupMemberFact>>,
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn group_capacity_only_becomes_unlimited_for_available_unlimited_members() {
|
||||
for (available_slots, unavailable_slots, expected) in
|
||||
[(None, Some(3), None), (Some(4), None, Some(4))]
|
||||
{
|
||||
let mut available = member("acct_available", 4);
|
||||
available.total_slots = available_slots;
|
||||
let mut unavailable = member("acct_limited", 3);
|
||||
unavailable.total_slots = unavailable_slots;
|
||||
let groups = Arc::new(FakeGroupStore {
|
||||
members: Some(vec![available, unavailable]),
|
||||
..Default::default()
|
||||
});
|
||||
let service = AdminHarness::new()
|
||||
.account_groups(groups)
|
||||
.account_runtime(Arc::new(FakeRuntimeStore::default()))
|
||||
.build()
|
||||
.await;
|
||||
let page = service
|
||||
.account_groups()
|
||||
.list(AccountGroupListQuery {
|
||||
page: 1,
|
||||
page_size: PageSize::new(20).expect("page size"),
|
||||
search: None,
|
||||
enabled: None,
|
||||
})
|
||||
.await
|
||||
.expect("group capacity");
|
||||
assert_eq!(page.items[0].capacity.total_slots, expected);
|
||||
assert_eq!(page.items[0].capacity.used_slots, Some(2));
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -115,7 +149,10 @@ impl AccountGroupStore for FakeGroupStore {
|
||||
.iter()
|
||||
.map(|group_id| group_id.as_str().to_owned())
|
||||
.collect();
|
||||
Ok(vec![member("acct_available", 4), member("acct_limited", 3)])
|
||||
Ok(self
|
||||
.members
|
||||
.clone()
|
||||
.unwrap_or_else(|| vec![member("acct_available", 4), member("acct_limited", 3)]))
|
||||
}
|
||||
|
||||
async fn create_account_group(
|
||||
@@ -218,7 +255,7 @@ fn group_record() -> AccountGroupRecord {
|
||||
},
|
||||
capacity: AccountGroupCapacity {
|
||||
used_slots: None,
|
||||
total_slots: 0,
|
||||
total_slots: Some(0),
|
||||
},
|
||||
usage: AccountGroupUsage {
|
||||
today_usd: DecimalAmount::from_str("1").expect("today usage"),
|
||||
@@ -242,7 +279,7 @@ fn member(account_id: &str, total_slots: u64) -> AccountGroupMemberFact {
|
||||
last_error_reason: None,
|
||||
last_error_message: None,
|
||||
},
|
||||
total_slots,
|
||||
total_slots: Some(total_slots),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -237,11 +237,45 @@ async fn dashboard_summary_should_project_rebuildable_runtime_slots() {
|
||||
.expect("dashboard summary")
|
||||
.capacity;
|
||||
|
||||
assert_eq!(capacity.total_slots, 7);
|
||||
assert_eq!(capacity.total_slots, Some(7));
|
||||
assert_eq!(capacity.used_slots, Some(2));
|
||||
assert_eq!(capacity.available_slots, Some(5));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dashboard_capacity_distinguishes_unlimited_inheritance_finite_overrides_and_empty_pool() {
|
||||
for (inherited_accounts, overridden_slots, expected_total, expected_available) in [
|
||||
(2, 5, None, None),
|
||||
(0, 5, Some(5), Some(3)),
|
||||
(0, 0, Some(0), Some(0)),
|
||||
] {
|
||||
let now = Utc::now();
|
||||
let store = Arc::new(FixtureObservabilityStore::new(observation_range(now)));
|
||||
store.replace_runtime_slots(Some(DashboardRuntimeSlots {
|
||||
inherited_accounts,
|
||||
overridden_slots,
|
||||
used_slots: Some(2),
|
||||
}));
|
||||
let services = super::AdminHarness::new()
|
||||
.observability(store)
|
||||
.settings(Arc::new(FixtureSettingsStore {
|
||||
max_concurrent_per_account: 0,
|
||||
}))
|
||||
.provider(super::dashboard_profile_provider())
|
||||
.build()
|
||||
.await;
|
||||
let capacity = services
|
||||
.observability()
|
||||
.dashboard_summary(observation_range(now), TrendKind::Usage)
|
||||
.await
|
||||
.expect("dashboard")
|
||||
.capacity;
|
||||
assert_eq!(capacity.total_slots, expected_total);
|
||||
assert_eq!(capacity.available_slots, expected_available);
|
||||
assert_eq!(capacity.used_slots, Some(2));
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dashboard_summary_should_share_one_current_observation_time_across_runtime_facts() {
|
||||
let historical_end = Utc::now() - Duration::hours(1);
|
||||
@@ -749,7 +783,9 @@ impl ObservabilityStore for FixtureObservabilityStore {
|
||||
}
|
||||
}
|
||||
|
||||
struct FixtureSettingsStore;
|
||||
struct FixtureSettingsStore {
|
||||
max_concurrent_per_account: u32,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl SettingsStore for FixtureSettingsStore {
|
||||
@@ -780,7 +816,7 @@ impl SettingsStore for FixtureSettingsStore {
|
||||
model_mappings: Default::default(),
|
||||
refresh_margin_seconds: 300,
|
||||
refresh_concurrency: 2,
|
||||
max_concurrent_per_account: 1,
|
||||
max_concurrent_per_account: self.max_concurrent_per_account,
|
||||
request_interval_ms: 0,
|
||||
max_waiting_per_key: 0,
|
||||
max_waiting_per_account: 0,
|
||||
@@ -834,7 +870,9 @@ impl SettingsStore for FixtureSettingsStore {
|
||||
async fn observability_services(store: Arc<FixtureObservabilityStore>) -> AdminServices {
|
||||
super::AdminHarness::new()
|
||||
.observability(store)
|
||||
.settings(Arc::new(FixtureSettingsStore))
|
||||
.settings(Arc::new(FixtureSettingsStore {
|
||||
max_concurrent_per_account: 1,
|
||||
}))
|
||||
.provider(super::dashboard_profile_provider())
|
||||
.build()
|
||||
.await
|
||||
@@ -845,7 +883,9 @@ async fn observability_services_with_calculated_billing(
|
||||
) -> AdminServices {
|
||||
super::AdminHarness::new()
|
||||
.observability(store)
|
||||
.settings(Arc::new(FixtureSettingsStore))
|
||||
.settings(Arc::new(FixtureSettingsStore {
|
||||
max_concurrent_per_account: 1,
|
||||
}))
|
||||
.provider(super::calculated_billing_provider())
|
||||
.build()
|
||||
.await
|
||||
|
||||
@@ -126,7 +126,7 @@ struct AccountGroupAccountSummaryView {
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct AccountGroupCapacityView {
|
||||
used_slots: Option<u64>,
|
||||
total_slots: u64,
|
||||
total_slots: Option<u64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
|
||||
@@ -425,7 +425,7 @@ pub struct DashboardPoolSummaryView {
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct DashboardCapacityInfoView {
|
||||
pub max_concurrent_per_account: u64,
|
||||
pub total_slots: u64,
|
||||
pub total_slots: Option<u64>,
|
||||
pub used_slots: Option<u64>,
|
||||
pub available_slots: Option<u64>,
|
||||
}
|
||||
|
||||
@@ -125,7 +125,6 @@ impl UpdateRuntimeSettingsRequest {
|
||||
for (value, field) in [
|
||||
(self.refresh_margin_seconds, "refreshMarginSeconds"),
|
||||
(self.refresh_concurrency, "refreshConcurrency"),
|
||||
(self.max_concurrent_per_account, "maxConcurrentPerAccount"),
|
||||
(self.usage_retention_days, "usageRetentionDays"),
|
||||
(self.ops_event_retention_days, "opsEventRetentionDays"),
|
||||
(self.audit_retention_days, "auditRetentionDays"),
|
||||
|
||||
@@ -723,7 +723,7 @@ impl AccountGroupStore for MemoryAccountGroupStore {
|
||||
last_error_reason: None,
|
||||
last_error_message: None,
|
||||
},
|
||||
total_slots: 1,
|
||||
total_slots: Some(1),
|
||||
},
|
||||
AccountGroupMemberFact {
|
||||
group_id: group_id(PRIMARY_GROUP_ID),
|
||||
@@ -737,7 +737,7 @@ impl AccountGroupStore for MemoryAccountGroupStore {
|
||||
last_error_reason: None,
|
||||
last_error_message: None,
|
||||
},
|
||||
total_slots: 1,
|
||||
total_slots: Some(1),
|
||||
},
|
||||
])
|
||||
}
|
||||
@@ -1490,7 +1490,7 @@ fn group_color(value: &str) -> AccountGroupColor {
|
||||
fn capacity(used_slots: Option<u64>, total_slots: u64) -> AccountGroupCapacity {
|
||||
AccountGroupCapacity {
|
||||
used_slots,
|
||||
total_slots,
|
||||
total_slots: Some(total_slots),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -77,6 +77,15 @@ fn settings_request_should_reject_unknown_rotation_strategy() {
|
||||
assert_eq!(request.validate().unwrap_err().field(), "rotationStrategy");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn settings_request_accepts_unlimited_default_account_concurrency() {
|
||||
let mut body = update_body();
|
||||
body["maxConcurrentPerAccount"] = json!(0);
|
||||
let request: UpdateRuntimeSettingsRequest =
|
||||
serde_json::from_value(body).expect("decode settings");
|
||||
request.validate().expect("zero means unlimited");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn settings_request_should_reject_non_semver_client_min() {
|
||||
let mut body = update_body();
|
||||
|
||||
@@ -58,6 +58,46 @@ impl AccountConcurrencyLimit {
|
||||
}
|
||||
}
|
||||
|
||||
/// 账号继承默认值或应用独立覆盖后的实际并发约束。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum AccountConcurrency {
|
||||
Unlimited,
|
||||
Limited(NonZeroU32),
|
||||
}
|
||||
|
||||
impl AccountConcurrency {
|
||||
/// 配置值零表示不限制;账号独立覆盖仍只接受正数。
|
||||
#[must_use]
|
||||
pub const fn new(value: u32) -> Self {
|
||||
match NonZeroU32::new(value) {
|
||||
Some(limit) => Self::Limited(limit),
|
||||
None => Self::Unlimited,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn get(self) -> u32 {
|
||||
match self {
|
||||
Self::Unlimited => 0,
|
||||
Self::Limited(limit) => limit.get(),
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn limit(self) -> Option<NonZeroU32> {
|
||||
match self {
|
||||
Self::Unlimited => None,
|
||||
Self::Limited(limit) => Some(limit),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<NonZeroU32> for AccountConcurrency {
|
||||
fn from(limit: NonZeroU32) -> Self {
|
||||
Self::Limited(limit)
|
||||
}
|
||||
}
|
||||
|
||||
/// Relative scheduling priority for an account.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
|
||||
pub struct AccountWeight(NonZeroU16);
|
||||
@@ -958,9 +998,9 @@ impl ProviderAccount {
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn effective_concurrency(&self, default: NonZeroU32) -> NonZeroU32 {
|
||||
pub const fn effective_concurrency(&self, default: AccountConcurrency) -> AccountConcurrency {
|
||||
match self.concurrency_limit {
|
||||
Some(limit) => limit.into_non_zero(),
|
||||
Some(limit) => AccountConcurrency::Limited(limit.into_non_zero()),
|
||||
None => default,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2,7 +2,6 @@
|
||||
|
||||
use std::cmp::Reverse;
|
||||
use std::collections::{BTreeSet, HashMap};
|
||||
use std::num::NonZeroU32;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::sync::{Mutex, RwLock};
|
||||
use std::time::{Duration, Instant, SystemTime};
|
||||
@@ -10,7 +9,7 @@ use std::time::{Duration, Instant, SystemTime};
|
||||
use crate::concurrency::ConcurrencyQueuePolicy;
|
||||
use crate::identity::ProviderKind;
|
||||
|
||||
use super::{AccountStatus, ProviderAccount, ProviderAccountId};
|
||||
use super::{AccountConcurrency, AccountStatus, ProviderAccount, ProviderAccountId};
|
||||
|
||||
/// `runtime_settings.rotation_strategy` 的稳定值。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
|
||||
@@ -48,21 +47,21 @@ impl RotationStrategy {
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct AccountSelectionPolicy {
|
||||
strategy: RotationStrategy,
|
||||
max_concurrent_per_account: NonZeroU32,
|
||||
max_concurrent_per_account: AccountConcurrency,
|
||||
request_interval: Duration,
|
||||
queue_policy: ConcurrencyQueuePolicy,
|
||||
}
|
||||
|
||||
impl AccountSelectionPolicy {
|
||||
#[must_use]
|
||||
pub const fn new(
|
||||
pub fn new(
|
||||
strategy: RotationStrategy,
|
||||
max_concurrent_per_account: NonZeroU32,
|
||||
max_concurrent_per_account: impl Into<AccountConcurrency>,
|
||||
request_interval: Duration,
|
||||
) -> Self {
|
||||
Self {
|
||||
strategy,
|
||||
max_concurrent_per_account,
|
||||
max_concurrent_per_account: max_concurrent_per_account.into(),
|
||||
request_interval,
|
||||
queue_policy: ConcurrencyQueuePolicy {
|
||||
max_waiting: 0,
|
||||
@@ -88,7 +87,7 @@ impl AccountSelectionPolicy {
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn max_concurrent_per_account(self) -> NonZeroU32 {
|
||||
pub const fn max_concurrent_per_account(self) -> AccountConcurrency {
|
||||
self.max_concurrent_per_account
|
||||
}
|
||||
|
||||
@@ -502,18 +501,19 @@ impl AccountSelector {
|
||||
)
|
||||
)
|
||||
})
|
||||
.fold((0_u64, 0_u64), |(used, total), candidate| {
|
||||
.try_fold((0_u64, 0_u64), |(used, total), candidate| {
|
||||
let capacity = u64::from(
|
||||
candidate
|
||||
.account
|
||||
.effective_concurrency(context.policy.max_concurrent_per_account())
|
||||
.limit()?
|
||||
.get(),
|
||||
);
|
||||
(
|
||||
Some((
|
||||
used.saturating_add(u64::from(candidate.signals.in_flight)),
|
||||
total.saturating_add(capacity),
|
||||
)
|
||||
});
|
||||
))
|
||||
})?;
|
||||
(total_slots > 0).then_some(AccountCapacitySnapshot {
|
||||
used_slots: used_slots.min(total_slots),
|
||||
total_slots,
|
||||
@@ -671,11 +671,11 @@ impl AccountSelector {
|
||||
if context.excluded_accounts.contains(candidate.account.id()) {
|
||||
return Some(AccountSchedulingBlocker::Excluded);
|
||||
}
|
||||
if candidate.signals.in_flight
|
||||
>= candidate
|
||||
.account
|
||||
.effective_concurrency(context.policy.max_concurrent_per_account())
|
||||
.get()
|
||||
if candidate
|
||||
.account
|
||||
.effective_concurrency(context.policy.max_concurrent_per_account())
|
||||
.limit()
|
||||
.is_some_and(|limit| candidate.signals.in_flight >= limit.get())
|
||||
{
|
||||
return Some(AccountSchedulingBlocker::ConcurrencyLimit);
|
||||
}
|
||||
@@ -704,19 +704,22 @@ pub(crate) const SMART_SCORE_TOLERANCE: f64 = 0.05;
|
||||
// 首输出 10 秒时延迟得分减半;固定尺度不随其他候选账号变化。
|
||||
const SMART_LATENCY_HALF_SCORE_MS: f64 = 10_000.0;
|
||||
|
||||
fn capacity_utilization(candidate: &AccountCandidate, default_concurrency: NonZeroU32) -> f64 {
|
||||
f64::from(candidate.signals.in_flight)
|
||||
/ f64::from(
|
||||
candidate
|
||||
.account
|
||||
.effective_concurrency(default_concurrency)
|
||||
.get(),
|
||||
)
|
||||
fn capacity_utilization(
|
||||
candidate: &AccountCandidate,
|
||||
default_concurrency: AccountConcurrency,
|
||||
) -> f64 {
|
||||
candidate
|
||||
.account
|
||||
.effective_concurrency(default_concurrency)
|
||||
.limit()
|
||||
.map_or(0.0, |limit| {
|
||||
f64::from(candidate.signals.in_flight) / f64::from(limit.get())
|
||||
})
|
||||
}
|
||||
|
||||
fn select_smart_candidate<'a>(
|
||||
candidates: &[&'a AccountCandidate],
|
||||
default_concurrency: NonZeroU32,
|
||||
default_concurrency: AccountConcurrency,
|
||||
cursor: u64,
|
||||
) -> Option<&'a AccountCandidate> {
|
||||
let mut ranked = candidates
|
||||
@@ -734,7 +737,10 @@ fn select_smart_candidate<'a>(
|
||||
Some(ranked[index].0)
|
||||
}
|
||||
|
||||
pub(crate) fn smart_score(candidate: &AccountCandidate, default_concurrency: NonZeroU32) -> f64 {
|
||||
pub(crate) fn smart_score(
|
||||
candidate: &AccountCandidate,
|
||||
default_concurrency: AccountConcurrency,
|
||||
) -> f64 {
|
||||
let load = 1.0 - capacity_utilization(candidate, default_concurrency).clamp(0.0, 1.0);
|
||||
let quota = candidate
|
||||
.signals
|
||||
|
||||
@@ -9,8 +9,8 @@ use std::time::{Duration, SystemTime};
|
||||
use futures::future::BoxFuture;
|
||||
|
||||
use crate::account::{
|
||||
AccountFeedbackStats, AccountRuntimeSignals, CredentialRevision, CredentialState,
|
||||
OpaqueProviderData, ProviderAccountId, ProviderAccountStore,
|
||||
AccountConcurrency, AccountFeedbackStats, AccountRuntimeSignals, CredentialRevision,
|
||||
CredentialState, OpaqueProviderData, ProviderAccountId, ProviderAccountStore,
|
||||
};
|
||||
use crate::identity::ProviderKind;
|
||||
use crate::policy::ClientApiKeyId;
|
||||
@@ -83,18 +83,18 @@ pub struct ProviderSchedulingLeaseRequest {
|
||||
provider_kind: ProviderKind,
|
||||
account_id: ProviderAccountId,
|
||||
credential_revision: CredentialRevision,
|
||||
max_concurrent: NonZeroU32,
|
||||
max_concurrent: AccountConcurrency,
|
||||
request_interval: Duration,
|
||||
deadline: SystemTime,
|
||||
}
|
||||
|
||||
impl ProviderSchedulingLeaseRequest {
|
||||
#[must_use]
|
||||
pub const fn new(
|
||||
pub fn new(
|
||||
provider_kind: ProviderKind,
|
||||
account_id: ProviderAccountId,
|
||||
credential_revision: CredentialRevision,
|
||||
max_concurrent: NonZeroU32,
|
||||
max_concurrent: impl Into<AccountConcurrency>,
|
||||
request_interval: Duration,
|
||||
deadline: SystemTime,
|
||||
) -> Self {
|
||||
@@ -102,7 +102,7 @@ impl ProviderSchedulingLeaseRequest {
|
||||
provider_kind,
|
||||
account_id,
|
||||
credential_revision,
|
||||
max_concurrent,
|
||||
max_concurrent: max_concurrent.into(),
|
||||
request_interval,
|
||||
deadline,
|
||||
}
|
||||
@@ -124,7 +124,7 @@ impl ProviderSchedulingLeaseRequest {
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn max_concurrent(&self) -> NonZeroU32 {
|
||||
pub const fn max_concurrent(&self) -> AccountConcurrency {
|
||||
self.max_concurrent
|
||||
}
|
||||
|
||||
|
||||
@@ -7,7 +7,9 @@ use std::time::Duration;
|
||||
|
||||
use futures::future::BoxFuture;
|
||||
|
||||
use crate::account::{AccountSelectionPolicy, ProviderAccountId, RotationStrategy};
|
||||
use crate::account::{
|
||||
AccountConcurrency, AccountSelectionPolicy, ProviderAccountId, RotationStrategy,
|
||||
};
|
||||
use crate::concurrency::ConcurrencyQueuePolicy;
|
||||
use crate::operation::Operation;
|
||||
use crate::policy::{
|
||||
@@ -474,8 +476,7 @@ async fn compile_runtime_snapshot(
|
||||
.ok_or(RuntimeSnapshotCompileError::InvalidData)?;
|
||||
let selection_policy = AccountSelectionPolicy::new(
|
||||
rotation_strategy,
|
||||
NonZeroU32::new(facts.settings.max_concurrent_per_account)
|
||||
.ok_or(RuntimeSnapshotCompileError::InvalidData)?,
|
||||
AccountConcurrency::new(facts.settings.max_concurrent_per_account),
|
||||
Duration::from_millis(facts.settings.request_interval_ms),
|
||||
)
|
||||
.with_queue(ConcurrencyQueuePolicy {
|
||||
|
||||
@@ -77,7 +77,7 @@ fn weighted_candidate(id: &str, weight: u16, in_flight: u32) -> AccountCandidate
|
||||
|
||||
#[test]
|
||||
fn account_scheduling_should_default_and_resolve_concurrency_override() {
|
||||
let default = NonZeroU32::new(12).expect("default concurrency");
|
||||
let default = gateway_core::account::AccountConcurrency::new(12);
|
||||
let account = account("acct_scheduling_default");
|
||||
assert_eq!(account.concurrency_limit(), None);
|
||||
assert_eq!(account.weight(), AccountWeight::DEFAULT);
|
||||
|
||||
@@ -10,6 +10,91 @@ use super::{candidate, candidate_with_concurrency, context};
|
||||
|
||||
const FAILURE_RATE_HALF_LIFE: Duration = Duration::from_secs(15 * 60);
|
||||
|
||||
#[test]
|
||||
fn unlimited_account_concurrency_preserves_overrides_intervals_and_finite_scores() {
|
||||
use gateway_core::account::{AccountConcurrency, AccountSelectionPolicy};
|
||||
|
||||
for strategy in [
|
||||
RotationStrategy::Smart,
|
||||
RotationStrategy::RoundRobin,
|
||||
RotationStrategy::Sticky,
|
||||
RotationStrategy::QuotaResetPriority,
|
||||
] {
|
||||
let mut context = context(strategy);
|
||||
context.policy =
|
||||
AccountSelectionPolicy::new(strategy, AccountConcurrency::Unlimited, Duration::ZERO);
|
||||
let candidates = [
|
||||
candidate("acct_unlimited", u32::MAX, None),
|
||||
candidate_with_concurrency("acct_limited", 2, 2),
|
||||
];
|
||||
assert_eq!(
|
||||
AccountSelector
|
||||
.select(&candidates, &context)
|
||||
.expect("unlimited remains eligible")
|
||||
.candidate()
|
||||
.account
|
||||
.id()
|
||||
.as_str(),
|
||||
"acct_unlimited"
|
||||
);
|
||||
assert!(AccountSelector.select(&candidates[1..], &context).is_none());
|
||||
assert!(
|
||||
AccountSelector
|
||||
.capacity_snapshot(&candidates, &context)
|
||||
.is_none()
|
||||
);
|
||||
let trace = gateway_core::diagnostics::TraceContext::new("req_unlimited_concurrency");
|
||||
trace.account_selection(
|
||||
&candidates,
|
||||
&context,
|
||||
AccountSelector.select(&candidates, &context).as_ref(),
|
||||
);
|
||||
if strategy == RotationStrategy::Smart {
|
||||
let snapshot = trace.snapshot().expect("trace");
|
||||
assert!(
|
||||
snapshot["events"][0]["data"]["candidates"][0]["smartScore"]
|
||||
.as_f64()
|
||||
.expect("finite score")
|
||||
.is_finite()
|
||||
);
|
||||
}
|
||||
context.policy = AccountSelectionPolicy::new(
|
||||
strategy,
|
||||
AccountConcurrency::Unlimited,
|
||||
Duration::from_secs(1),
|
||||
);
|
||||
let mut recent = candidate("acct_unlimited", u32::MAX, None);
|
||||
recent.signals.last_started_at = Some(context.now);
|
||||
assert!(AccountSelector.select(&[recent], &context).is_none());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unavailable_unlimited_account_does_not_hide_the_finite_pool_capacity() {
|
||||
use gateway_core::account::{
|
||||
AccountConcurrency, AccountSelectionPolicy, CredentialState, QuotaState,
|
||||
};
|
||||
let mut context = context(RotationStrategy::Smart);
|
||||
context.policy = AccountSelectionPolicy::new(
|
||||
RotationStrategy::Smart,
|
||||
AccountConcurrency::Unlimited,
|
||||
Duration::ZERO,
|
||||
);
|
||||
let mut unavailable = candidate("acct_disabled", 0, None);
|
||||
unavailable.account = unavailable.account.with_account_facts(
|
||||
false,
|
||||
CredentialState::Ready,
|
||||
QuotaState::unknown(),
|
||||
None,
|
||||
None,
|
||||
);
|
||||
let finite = candidate_with_concurrency("acct_limited", 1, 2);
|
||||
let capacity = AccountSelector
|
||||
.capacity_snapshot(&[unavailable, finite], &context)
|
||||
.expect("finite available pool");
|
||||
assert_eq!((capacity.used_slots(), capacity.total_slots()), (1, 2));
|
||||
}
|
||||
|
||||
fn feedback_subject() -> (AccountFeedbackStats, ProviderKind, ProviderAccountId) {
|
||||
(
|
||||
AccountFeedbackStats::default(),
|
||||
|
||||
@@ -135,6 +135,43 @@ fn compiler_should_reject_revision_changed_during_consistent_read() {
|
||||
assert_eq!(error, RuntimeSnapshotCompileError::RevisionChanged);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compiler_accepts_unlimited_default_account_concurrency() {
|
||||
let facts = SnapshotFacts::new(
|
||||
revision(1),
|
||||
revision(1),
|
||||
SnapshotSettingsFacts::new(0, 50, "smart", BTreeMap::new(), None, None),
|
||||
Vec::new(),
|
||||
Vec::new(),
|
||||
vec![SnapshotProviderAccountFacts::new(
|
||||
gateway_core::account::ProviderAccountId::new("acct_unlimited").expect("account"),
|
||||
"alpha",
|
||||
)],
|
||||
Vec::new(),
|
||||
);
|
||||
let compiler = RuntimeSnapshotCompiler::new(
|
||||
Arc::new(TestSnapshotStore::new(Ok(facts))),
|
||||
Arc::new(TestCatalog::Unavailable),
|
||||
);
|
||||
let snapshot = block_on(compiler.compile()).expect("compile unlimited default");
|
||||
let plan = snapshot
|
||||
.plan(
|
||||
&PublicModelId::new("any-model").expect("model"),
|
||||
&super::operation(),
|
||||
snapshot.all_account_scope(),
|
||||
&Default::default(),
|
||||
)
|
||||
.expect("plan request");
|
||||
assert_eq!(
|
||||
plan.account_selection_policy().max_concurrent_per_account(),
|
||||
gateway_core::account::AccountConcurrency::Unlimited
|
||||
);
|
||||
assert_eq!(
|
||||
plan.account_selection_policy().request_interval(),
|
||||
std::time::Duration::from_millis(50)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn routing_plans_share_frozen_pricing_after_a_new_snapshot_is_published() {
|
||||
use gateway_core::{metering::PricingOverrides, runtime::RuntimeSnapshotHandle};
|
||||
|
||||
@@ -213,9 +213,10 @@ impl AccountGroupStore for PgAccountGroupRepository {
|
||||
last_error_reason: account.last_error_reason,
|
||||
last_error_message: account.last_error_message,
|
||||
},
|
||||
total_slots: account
|
||||
.concurrency_limit
|
||||
.map_or(default_slots, |limit| u64::from(limit.get())),
|
||||
total_slots: account.concurrency_limit.map_or_else(
|
||||
|| (default_slots > 0).then_some(default_slots),
|
||||
|limit| Some(u64::from(limit.get())),
|
||||
),
|
||||
})
|
||||
})
|
||||
.collect::<StoreResult<Vec<_>>>()
|
||||
@@ -504,7 +505,7 @@ fn group_record(row: &sqlx::postgres::PgRow) -> StoreResult<AccountGroupRecord>
|
||||
},
|
||||
capacity: AccountGroupCapacity {
|
||||
used_slots: None,
|
||||
total_slots: 0,
|
||||
total_slots: Some(0),
|
||||
},
|
||||
usage: AccountGroupUsage {
|
||||
today_usd: DecimalAmount::from_str("0").map_err(|_| invalid("invalid zero cost"))?,
|
||||
|
||||
@@ -755,7 +755,8 @@ impl AccountStore for PgAdminAccountStore {
|
||||
.map_err(|_| postgres_unavailable("lock default concurrency"))?;
|
||||
let changed = sqlx::query_scalar::<_, String>(
|
||||
"update provider_accounts set concurrency_limit = $2, updated_at = greatest(now(), updated_at)
|
||||
where id = $1 and enabled = true and coalesce(concurrency_limit, $3) > $2
|
||||
where id = $1 and enabled = true
|
||||
and (coalesce(concurrency_limit, $3) = 0 or coalesce(concurrency_limit, $3) > $2)
|
||||
returning id",
|
||||
)
|
||||
.bind(account_id.as_str())
|
||||
|
||||
@@ -162,7 +162,6 @@ impl RuntimeSettingsUpdate {
|
||||
|| isize::try_from(self.responses_max_decompressed_body_bytes).is_err()
|
||||
|| self.refresh_margin_seconds == 0
|
||||
|| self.refresh_concurrency == 0
|
||||
|| self.max_concurrent_per_account == 0
|
||||
|| self.max_waiting_per_key > 1_000
|
||||
|| self.max_waiting_per_account > 1_000
|
||||
|| !(1..=120).contains(&self.concurrency_wait_timeout_seconds)
|
||||
|
||||
@@ -38,7 +38,7 @@ local max_concurrent = tonumber(ARGV[4])
|
||||
local interval_ms = tonumber(ARGV[5])
|
||||
local retry_ms = 0
|
||||
|
||||
if in_flight >= max_concurrent then
|
||||
if max_concurrent > 0 and in_flight >= max_concurrent then
|
||||
local earliest = redis.call('ZRANGE', KEYS[1], 0, 0, 'WITHSCORES')
|
||||
if #earliest == 2 then
|
||||
retry_ms = math.max(retry_ms, math.ceil(tonumber(earliest[2]) - now_ms))
|
||||
@@ -177,7 +177,7 @@ impl CredentialBoundedLeaseRequest {
|
||||
pub fn validate(&self) -> StoreResult<()> {
|
||||
require_nonempty("credential bounded lease", "resource_id", &self.resource_id)?;
|
||||
require_nonempty("credential bounded lease", "owner_id", &self.owner_id)?;
|
||||
if self.max_concurrent == 0 {
|
||||
if self.max_concurrent == 0 && self.scope != CredentialLeaseScope::ProviderAccount {
|
||||
return Err(invalid("max_concurrent must be positive"));
|
||||
}
|
||||
supported_duration(self.request_interval, true, "request interval")?;
|
||||
@@ -395,7 +395,7 @@ impl RedisCredentialLeaseRepository {
|
||||
request_interval: Duration,
|
||||
) -> StoreResult<LeaseAttempt> {
|
||||
request.validate()?;
|
||||
if max_concurrent == 0 {
|
||||
if max_concurrent == 0 && request.scope != CredentialLeaseScope::ProviderAccount {
|
||||
return Err(invalid("max_concurrent must be positive"));
|
||||
}
|
||||
let keys = self.keys(request)?;
|
||||
|
||||
@@ -113,8 +113,11 @@ async fn groups_aggregate_cross_provider_members_and_key_bindings_without_multip
|
||||
.expect("load current-page group members");
|
||||
assert_eq!(members.len(), 2);
|
||||
assert_eq!(
|
||||
members.iter().map(|member| member.total_slots).sum::<u64>(),
|
||||
7
|
||||
members
|
||||
.iter()
|
||||
.map(|member| member.total_slots)
|
||||
.sum::<Option<u64>>(),
|
||||
Some(7)
|
||||
);
|
||||
assert_eq!(page.total, 2);
|
||||
let by_id = page
|
||||
@@ -134,7 +137,7 @@ async fn groups_aggregate_cross_provider_members_and_key_bindings_without_multip
|
||||
assert_eq!(mixed.account_summary.limited, 0);
|
||||
assert_eq!(mixed.account_summary.total, 0);
|
||||
assert_eq!(mixed.capacity.used_slots, None);
|
||||
assert_eq!(mixed.capacity.total_slots, 0);
|
||||
assert_eq!(mixed.capacity.total_slots, Some(0));
|
||||
assert_eq!(mixed.usage.today_usd.as_str(), "0");
|
||||
assert_eq!(mixed.usage.retained_total_usd.as_str(), "0");
|
||||
let empty = by_id.get(EMPTY_GROUP).expect("empty group");
|
||||
|
||||
@@ -3445,7 +3445,8 @@ async fn model_access_only_batch_update_preserves_other_settings_and_survives_re
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn adaptive_concurrency_uses_latest_locked_settings_without_overwriting_admin_fields() {
|
||||
async fn adaptive_concurrency_handles_unlimited_and_latest_locked_settings_without_overwriting_admin_fields()
|
||||
{
|
||||
let Some(database) = TestDatabase::create("adaptive_concurrency").await else {
|
||||
return;
|
||||
};
|
||||
@@ -3466,6 +3467,9 @@ async fn adaptive_concurrency_uses_latest_locked_settings_without_overwriting_ad
|
||||
(true, Some(2), 10, Some(2), false),
|
||||
(true, None, 2, None, false),
|
||||
(true, None, 10, Some(3), true),
|
||||
(true, None, 0, Some(3), true),
|
||||
(true, Some(2), 0, Some(2), false),
|
||||
(false, None, 0, None, false),
|
||||
] {
|
||||
let mut admin = database.pool.begin().await.expect("admin transaction");
|
||||
sqlx::query("update runtime_settings set max_concurrent_per_account = $1 where id = 1")
|
||||
@@ -3520,6 +3524,9 @@ async fn adaptive_concurrency_uses_latest_locked_settings_without_overwriting_ad
|
||||
.fetch_all(&database.pool)
|
||||
.await
|
||||
.expect("audit");
|
||||
assert_eq!(audited_fields, vec![vec!["concurrency_limit".to_owned()]]);
|
||||
assert_eq!(
|
||||
audited_fields,
|
||||
vec![vec!["concurrency_limit".to_owned()]; 2]
|
||||
);
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
@@ -48,6 +48,45 @@ fn runtime_settings_keep_account_rotation_global() {
|
||||
assert!(settings.validate().is_ok());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unlimited_default_account_concurrency_round_trips_without_relaxing_other_limits() {
|
||||
let Some(database) = TestDatabase::create("unlimited_default_concurrency").await else {
|
||||
return;
|
||||
};
|
||||
let repository = PgRuntimeSettingsRepository::new(database.pool.clone());
|
||||
let mut update = settings_with_margin(3_600);
|
||||
update.max_concurrent_per_account = 0;
|
||||
repository
|
||||
.update_runtime_settings(update)
|
||||
.await
|
||||
.expect("persist unlimited default");
|
||||
let settings = repository
|
||||
.load_runtime_settings()
|
||||
.await
|
||||
.expect("read unlimited default");
|
||||
assert_eq!(settings.max_concurrent_per_account, 0);
|
||||
for statement in [
|
||||
"update runtime_settings set max_concurrent_per_account = -1 where id = 1",
|
||||
"update runtime_settings set refresh_concurrency = 0 where id = 1",
|
||||
"update runtime_settings set refresh_margin_seconds = 0 where id = 1",
|
||||
"update runtime_settings set request_interval_ms = -1 where id = 1",
|
||||
] {
|
||||
let error = sqlx::query(statement)
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.expect_err("constraint must reject invalid setting");
|
||||
assert_eq!(
|
||||
error
|
||||
.as_database_error()
|
||||
.and_then(|error| error.code())
|
||||
.as_deref(),
|
||||
Some("23514"),
|
||||
"{statement}"
|
||||
);
|
||||
}
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn runtime_settings_reject_invalid_model_mapping() {
|
||||
let settings = RuntimeSettingsUpdate {
|
||||
|
||||
@@ -19,9 +19,101 @@ fn credential_lease_rejects_zero_ttl() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn scheduling_lease_rejects_zero_concurrency() {
|
||||
let request = scheduling_request("acct_invalid", "worker-invalid", 0, Duration::ZERO);
|
||||
assert!(request.validate().is_err());
|
||||
fn only_account_scheduling_leases_allow_unlimited_concurrency() {
|
||||
let mut request = scheduling_request("acct_unlimited", "worker-unlimited", 0, Duration::ZERO);
|
||||
request.validate().expect("unlimited account concurrency");
|
||||
for scope in [
|
||||
CredentialLeaseScope::OAuthRefresh,
|
||||
CredentialLeaseScope::OAuthRefreshCapacity,
|
||||
] {
|
||||
request.scope = scope;
|
||||
assert!(request.validate().is_err());
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unlimited_scheduling_leases_still_count_release_and_enforce_request_interval() {
|
||||
let Some((repository, mut connection, namespace)) = repository().await else {
|
||||
return;
|
||||
};
|
||||
let request = scheduling_request("acct_unlimited", "worker-unlimited", 0, Duration::ZERO);
|
||||
let mut leases = Vec::new();
|
||||
for _ in 0..5 {
|
||||
leases.push(acquired(
|
||||
repository
|
||||
.try_acquire_bounded_lease(&request)
|
||||
.await
|
||||
.expect("unlimited slot"),
|
||||
));
|
||||
}
|
||||
assert_eq!(
|
||||
repository
|
||||
.credential_runtime_signals(std::slice::from_ref(&request.resource_id))
|
||||
.await
|
||||
.expect("signals")[0]
|
||||
.in_flight,
|
||||
5
|
||||
);
|
||||
let limited = CredentialBoundedLeaseRequest {
|
||||
max_concurrent: 3,
|
||||
..request.clone()
|
||||
};
|
||||
assert!(matches!(
|
||||
repository
|
||||
.try_acquire_bounded_lease(&limited)
|
||||
.await
|
||||
.expect("positive limit still applies"),
|
||||
CredentialBoundedLeaseAcquisition::Busy { .. }
|
||||
));
|
||||
for lease in leases {
|
||||
assert!(lease.release().await.expect("release unlimited slot"));
|
||||
}
|
||||
assert_eq!(
|
||||
repository
|
||||
.credential_runtime_signals(std::slice::from_ref(&request.resource_id))
|
||||
.await
|
||||
.expect("released signals")[0]
|
||||
.in_flight,
|
||||
0
|
||||
);
|
||||
let interval = scheduling_request(
|
||||
"acct_unlimited_interval",
|
||||
"worker-interval",
|
||||
0,
|
||||
Duration::from_secs(30),
|
||||
);
|
||||
acquired(
|
||||
repository
|
||||
.try_acquire_bounded_lease(&interval)
|
||||
.await
|
||||
.expect("first start"),
|
||||
)
|
||||
.release()
|
||||
.await
|
||||
.expect("release interval slot");
|
||||
assert!(
|
||||
matches!(repository.try_acquire_bounded_lease(&interval).await.expect("interval decision"), CredentialBoundedLeaseAcquisition::Busy { retry_after: Some(value) } if value > Duration::ZERO)
|
||||
);
|
||||
let keys = redis::cmd("KEYS")
|
||||
.arg(format!("{namespace}:*"))
|
||||
.query_async::<Vec<String>>(&mut connection)
|
||||
.await
|
||||
.expect("isolated keys");
|
||||
for key in &keys {
|
||||
assert!(
|
||||
redis::cmd("PTTL")
|
||||
.arg(key)
|
||||
.query_async::<i64>(&mut connection)
|
||||
.await
|
||||
.expect("TTL")
|
||||
> 0
|
||||
);
|
||||
}
|
||||
redis::cmd("DEL")
|
||||
.arg(keys)
|
||||
.query_async::<i64>(&mut connection)
|
||||
.await
|
||||
.expect("cleanup isolated keys");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
@@ -43,6 +43,7 @@ const CLOUDFLARE_CHALLENGE_BACKOFF: [Duration; 4] = [
|
||||
const CLOUDFLARE_PATH_BLOCK_THRESHOLD: u32 = 3;
|
||||
const SESSION_AFFINITY_TIMEOUT: Duration = Duration::from_millis(100);
|
||||
const CYBER_POLICY_SESSION_TTL: Duration = Duration::from_secs(60 * 60);
|
||||
const MAX_ACCOUNT_SNAPSHOT_RETRIES: u32 = 3;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum CodexAccountFailure {
|
||||
@@ -351,6 +352,7 @@ impl CodexCredentialSelector {
|
||||
request.attempt.deadline(),
|
||||
request.attempt.concurrency_wait_budget(),
|
||||
);
|
||||
let mut snapshot_retries = 0;
|
||||
'capacity: loop {
|
||||
let diagnostic = request.attempt.is_diagnostic_required_account();
|
||||
let mut accounts = self.repository.list_for_provider().await?;
|
||||
@@ -396,7 +398,18 @@ impl CodexCredentialSelector {
|
||||
if request.requires_websocket
|
||||
&& account.authentication_kind() == super::CODEX_AUTHENTICATION_KIND_API_KEY
|
||||
{
|
||||
let runtime = self.repository.load_runtime_credential(&account).await?;
|
||||
let runtime = match self.repository.load_runtime_credential(&account).await {
|
||||
Ok(runtime) => runtime,
|
||||
Err(CredentialRepositoryError::RevisionConflict) => {
|
||||
retry_account_snapshot(
|
||||
request.attempt,
|
||||
&account,
|
||||
&mut snapshot_retries,
|
||||
)?;
|
||||
continue 'capacity;
|
||||
}
|
||||
Err(error) => return Err(error.into()),
|
||||
};
|
||||
if !matches!(runtime.authentication, CodexRuntimeAuthentication::ApiKey(ref auth)
|
||||
if auth.configuration.transport == super::ApiKeyTransport::PreferWebsocket)
|
||||
{
|
||||
@@ -609,6 +622,16 @@ impl CodexCredentialSelector {
|
||||
.find(|candidate| candidate.account.id() == selected.account.id())
|
||||
.map(|candidate| candidate.account.clone())
|
||||
.ok_or(CredentialSelectionError::InvalidCredential)?;
|
||||
// 额度观测等并发更新会使整个账号快照失效,必须重新选号并校验资格。
|
||||
// 在占用租约和请求间隔前完成校验,避免重读被自己的异步释放挡住。
|
||||
let runtime = match self.repository.load_runtime_credential(&account).await {
|
||||
Ok(runtime) => runtime,
|
||||
Err(CredentialRepositoryError::RevisionConflict) => {
|
||||
retry_account_snapshot(request.attempt, &account, &mut snapshot_retries)?;
|
||||
continue 'capacity;
|
||||
}
|
||||
Err(error) => return Err(error.into()),
|
||||
};
|
||||
let allows_account_state_mutation = !diagnostic || account.enabled();
|
||||
match self
|
||||
.leases
|
||||
@@ -690,7 +713,6 @@ impl CodexCredentialSelector {
|
||||
.is_some_and(CodexSessionAffinity::session_id_present),
|
||||
"OpenAI account selected"
|
||||
);
|
||||
let runtime = self.repository.load_runtime_credential(&account).await?;
|
||||
let cookies = runtime
|
||||
.cookies
|
||||
.into_iter()
|
||||
@@ -1459,6 +1481,32 @@ impl fmt::Debug for CodexCredentialLease {
|
||||
}
|
||||
}
|
||||
|
||||
fn retry_account_snapshot(
|
||||
attempt: &AttemptContext,
|
||||
account: &ProviderAccount,
|
||||
retries: &mut u32,
|
||||
) -> Result<(), CredentialSelectionError> {
|
||||
let retry = *retries < MAX_ACCOUNT_SNAPSHOT_RETRIES;
|
||||
if retry {
|
||||
*retries += 1;
|
||||
}
|
||||
attempt.trace().record(
|
||||
"account.snapshot_conflict",
|
||||
serde_json::json!({
|
||||
"accountId": account.id().as_str(),
|
||||
"credentialRevision": account.revision().get(),
|
||||
"retry": retry,
|
||||
"retryCount": *retries,
|
||||
"maxRetries": MAX_ACCOUNT_SNAPSHOT_RETRIES,
|
||||
}),
|
||||
);
|
||||
if retry {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(CredentialSelectionError::AccountSnapshotChanged)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
pub enum CredentialSelectionError {
|
||||
#[error(transparent)]
|
||||
@@ -1471,6 +1519,8 @@ pub enum CredentialSelectionError {
|
||||
CapacityUnavailable { retry_after: Option<Duration> },
|
||||
#[error("Codex account data is invalid")]
|
||||
InvalidCredential,
|
||||
#[error("Codex account changed repeatedly during selection")]
|
||||
AccountSnapshotChanged,
|
||||
#[error("Codex account store is unavailable")]
|
||||
Store,
|
||||
#[error("Codex account lease runtime is unavailable")]
|
||||
|
||||
@@ -797,6 +797,14 @@ pub(super) fn map_selection_error(error: CredentialSelectionError) -> ProviderEr
|
||||
Some("usage_limit_reached".to_owned()),
|
||||
Some("usage_limit_reached".to_owned()),
|
||||
)),
|
||||
CredentialSelectionError::AccountSnapshotChanged => provider_error(
|
||||
ProviderErrorKind::ProviderInfrastructureUnavailable,
|
||||
UpstreamSendState::NotSent,
|
||||
)
|
||||
.with_diagnostic(
|
||||
ProviderDiagnostic::new("OpenAI account changed repeatedly during selection")
|
||||
.with_classification("account_selection", "account_snapshot_conflict"),
|
||||
),
|
||||
CredentialSelectionError::InvalidCredential
|
||||
| CredentialSelectionError::Store
|
||||
| CredentialSelectionError::Coordinator
|
||||
|
||||
@@ -411,6 +411,226 @@ fn selector_uses_the_account_concurrency_override_for_the_redis_lease() {
|
||||
assert_eq!(requests[0].max_concurrent().get(), 7);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn selector_reloads_quota_snapshot_before_acquiring_the_only_lease() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_primary", "at-primary");
|
||||
store.set_scheduling(
|
||||
"acct_primary",
|
||||
Some(AccountConcurrencyLimit::new(1).expect("single slot")),
|
||||
AccountWeight::DEFAULT,
|
||||
);
|
||||
let original = store.account("acct_primary").expect("account");
|
||||
let observed_at = SystemTime::now();
|
||||
block_on(store.apply_quota_access(QuotaAccessChange {
|
||||
account_id: original.id().clone(),
|
||||
expected_revision: original.revision(),
|
||||
state: QuotaState::allowed(observed_at),
|
||||
}))
|
||||
.expect("initial quota observation");
|
||||
store.on_credential_load(Arc::new(move |store, id, count| {
|
||||
Box::pin(async move {
|
||||
if count == 1 {
|
||||
let account = store.account(id.as_str()).expect("account");
|
||||
store
|
||||
.apply_quota_access(QuotaAccessChange {
|
||||
account_id: id.clone(),
|
||||
expected_revision: account.revision(),
|
||||
state: QuotaState::allowed(observed_at + Duration::from_millis(1)),
|
||||
})
|
||||
.await?;
|
||||
}
|
||||
Ok(())
|
||||
})
|
||||
}));
|
||||
let leases = Arc::new(TestLeaseCoordinator::default());
|
||||
let selector = selector(&store, Arc::clone(&leases));
|
||||
let attempt = attempt(BTreeSet::new());
|
||||
let selected = block_on(selector.select(&SelectCodexCredential {
|
||||
upstream_model: "gpt-5.4",
|
||||
request_url: &Url::parse("https://chatgpt.com/backend-api/codex/responses").expect("URL"),
|
||||
attempt: &attempt,
|
||||
session_affinity_key: None,
|
||||
}))
|
||||
.expect("a newer allowed observation must not fail selection");
|
||||
|
||||
assert_eq!(selected.account().revision(), original.revision());
|
||||
assert_eq!(
|
||||
selected.account(),
|
||||
&store.account("acct_primary").expect("current account")
|
||||
);
|
||||
assert_eq!(store.credential_loads(), 2);
|
||||
let requests = leases.requests.lock().expect("lease requests");
|
||||
assert_eq!(
|
||||
requests.len(),
|
||||
1,
|
||||
"conflicted snapshots must never consume capacity or request interval"
|
||||
);
|
||||
assert_eq!(requests[0].max_concurrent().get(), 1);
|
||||
assert_eq!(requests[0].request_interval(), Duration::from_millis(10));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn selector_uses_rotated_credentials_after_a_snapshot_conflict() {
|
||||
use gateway_core::account::{CredentialCasOutcome, CredentialCasUpdate, ProviderAccountUpdate};
|
||||
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_primary", "at-old");
|
||||
let old_revision = store.account("acct_primary").expect("account").revision();
|
||||
store.on_credential_load(Arc::new(|store, id, count| {
|
||||
Box::pin(async move {
|
||||
if count == 1 {
|
||||
let account = store.account(id.as_str()).expect("account");
|
||||
let update = CredentialCasUpdate::new(
|
||||
id.clone(),
|
||||
account.revision(),
|
||||
ProviderAccountUpdate {
|
||||
account_id: id.clone(),
|
||||
name: account.name().to_owned(),
|
||||
email: account.email().map(str::to_owned),
|
||||
plan_type: account.plan_type().map(str::to_owned),
|
||||
},
|
||||
CodexCredentialCodec::encode_new(
|
||||
&secret("at-rotated"),
|
||||
&profile("chatgpt-acct_primary"),
|
||||
Vec::new(),
|
||||
)
|
||||
.expect("rotated credential"),
|
||||
account.has_refresh_token(),
|
||||
account.access_token_expires_at(),
|
||||
account.next_refresh_at(),
|
||||
)
|
||||
.expect("credential update")
|
||||
.preserving_profile();
|
||||
assert!(matches!(
|
||||
store.compare_and_swap_credential(update).await?,
|
||||
CredentialCasOutcome::Updated(_)
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
})
|
||||
}));
|
||||
let leases = Arc::new(TestLeaseCoordinator::default());
|
||||
let selector = selector(&store, Arc::clone(&leases));
|
||||
let attempt = attempt(BTreeSet::new());
|
||||
let selected = block_on(selector.select(&SelectCodexCredential {
|
||||
upstream_model: "gpt-5.4",
|
||||
request_url: &Url::parse("https://chatgpt.com/backend-api/codex/responses").expect("URL"),
|
||||
attempt: &attempt,
|
||||
session_affinity_key: None,
|
||||
}))
|
||||
.expect("reload rotated credentials");
|
||||
|
||||
assert_eq!(
|
||||
selected.account().revision(),
|
||||
old_revision.next().expect("next revision")
|
||||
);
|
||||
assert_eq!(
|
||||
selected
|
||||
.authentication()
|
||||
.oauth()
|
||||
.expect("OAuth")
|
||||
.access_token
|
||||
.expose_secret(),
|
||||
"at-rotated"
|
||||
);
|
||||
assert_eq!(store.credential_loads(), 2);
|
||||
let requests = leases.requests.lock().expect("lease requests");
|
||||
assert_eq!(requests.len(), 1);
|
||||
assert_eq!(
|
||||
requests[0].credential_revision(),
|
||||
selected.account().revision()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn selector_snapshot_retry_preserves_frozen_account_scope_and_model_permissions() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
for id in ["acct_primary", "acct_other", "acct_outside_scope"] {
|
||||
create_account(&store, id, "at-test");
|
||||
}
|
||||
store.on_credential_load(Arc::new(|store, id, count| {
|
||||
Box::pin(async move {
|
||||
assert_eq!(id.as_str(), "acct_primary");
|
||||
assert_eq!(count, 1);
|
||||
store.set_enabled(id, false).await
|
||||
})
|
||||
}));
|
||||
let leases = Arc::new(TestLeaseCoordinator::default());
|
||||
let selector = selector(&store, Arc::clone(&leases));
|
||||
let attempt = model_restricted_attempt(None, BTreeSet::new());
|
||||
let error = block_on(selector.select(&SelectCodexCredential {
|
||||
upstream_model: "test-luna",
|
||||
request_url: &Url::parse("https://chatgpt.com/backend-api/codex/responses").expect("URL"),
|
||||
attempt: &attempt,
|
||||
session_affinity_key: None,
|
||||
}))
|
||||
.expect_err("retry cannot escape the frozen scope or use a denied model");
|
||||
|
||||
assert!(matches!(
|
||||
error,
|
||||
CredentialSelectionError::NoEligibleCredential
|
||||
));
|
||||
assert!(leases.requests.lock().expect("lease requests").is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn selector_does_not_retry_credential_store_unavailability() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_primary", "at-primary");
|
||||
store.on_credential_load(Arc::new(|_, _, _| {
|
||||
Box::pin(async {
|
||||
Err(gateway_core::error::StoreError::new(
|
||||
gateway_core::error::StoreErrorKind::Unavailable,
|
||||
))
|
||||
})
|
||||
}));
|
||||
let leases = Arc::new(TestLeaseCoordinator::default());
|
||||
let selector = selector(&store, Arc::clone(&leases));
|
||||
let attempt = attempt(BTreeSet::new());
|
||||
let error = block_on(selector.select(&SelectCodexCredential {
|
||||
upstream_model: "gpt-5.4",
|
||||
request_url: &Url::parse("https://chatgpt.com/backend-api/codex/responses").expect("URL"),
|
||||
attempt: &attempt,
|
||||
session_affinity_key: None,
|
||||
}))
|
||||
.expect_err("unavailability is not an account snapshot conflict");
|
||||
|
||||
assert!(matches!(error, CredentialSelectionError::Store));
|
||||
assert_eq!(store.credential_loads(), 1);
|
||||
assert!(leases.requests.lock().expect("lease requests").is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn selector_does_not_retry_invalid_credential_data() {
|
||||
use gateway_core::account::{NewProviderAccount, PlaintextCredential};
|
||||
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_primary", "at-primary");
|
||||
let account = store.account("acct_primary").expect("account");
|
||||
block_on(store.delete_account(account.id())).expect("replace stored credential");
|
||||
block_on(store.create_account(NewProviderAccount {
|
||||
account,
|
||||
credential: PlaintextCredential::new(serde_json::Map::new()),
|
||||
model_access: None,
|
||||
}))
|
||||
.expect("persist invalid provider data");
|
||||
let leases = Arc::new(TestLeaseCoordinator::default());
|
||||
let selector = selector(&store, Arc::clone(&leases));
|
||||
let attempt = attempt(BTreeSet::new());
|
||||
let error = block_on(selector.select(&SelectCodexCredential {
|
||||
upstream_model: "gpt-5.4",
|
||||
request_url: &Url::parse("https://chatgpt.com/backend-api/codex/responses").expect("URL"),
|
||||
attempt: &attempt,
|
||||
session_affinity_key: None,
|
||||
}))
|
||||
.expect_err("invalid data is not an account snapshot conflict");
|
||||
|
||||
assert!(matches!(error, CredentialSelectionError::InvalidCredential));
|
||||
assert_eq!(store.credential_loads(), 1);
|
||||
assert!(leases.requests.lock().expect("lease requests").is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn selector_round_robin_cursor_advances_across_requests() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
|
||||
@@ -1504,6 +1504,340 @@ async fn generate_without_an_eligible_openai_account_fails_before_network_io() {
|
||||
assert_eq!(error.send_state(), UpstreamSendState::NotSent);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn quota_snapshot_race_reloads_and_sends_the_request_once() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_provider_contract").await;
|
||||
let original = store.account("acct_provider_contract").expect("account");
|
||||
let observed_at = SystemTime::now();
|
||||
store
|
||||
.apply_quota_access(QuotaAccessChange {
|
||||
account_id: original.id().clone(),
|
||||
expected_revision: original.revision(),
|
||||
state: QuotaState::allowed(observed_at),
|
||||
})
|
||||
.await
|
||||
.expect("initial quota");
|
||||
store.on_credential_load(Arc::new(move |store, id, count| {
|
||||
Box::pin(async move {
|
||||
if count == 1 {
|
||||
let account = store.account(id.as_str()).expect("account");
|
||||
store
|
||||
.apply_quota_access(QuotaAccessChange {
|
||||
account_id: id.clone(),
|
||||
expected_revision: account.revision(),
|
||||
state: QuotaState::allowed(observed_at + Duration::from_millis(1)),
|
||||
})
|
||||
.await?;
|
||||
}
|
||||
Ok(())
|
||||
})
|
||||
}));
|
||||
let upstream = MockServer::start().await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/codex/responses"))
|
||||
.and(header("authorization", "Bearer at-acct_provider_contract"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200).set_body_raw(CAPTURE_COMPLETED_SSE, "text/event-stream"),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&upstream)
|
||||
.await;
|
||||
let leases = Arc::new(TestLeaseCoordinator::default());
|
||||
let provider = provider_with_affinity_and_base_url_and_leases(
|
||||
&store,
|
||||
Arc::new(MemorySessionAffinity::default()),
|
||||
upstream.uri(),
|
||||
Arc::clone(&leases),
|
||||
);
|
||||
let mut stream = provider
|
||||
.execute(
|
||||
planned_request("openai", http_generate_operation()),
|
||||
context("req_quota_snapshot_race", CancellationToken::new()),
|
||||
)
|
||||
.await
|
||||
.expect("concurrent quota observation must not fail prepare");
|
||||
assert_eq!(stream.metadata().provider_account_id(), original.id());
|
||||
assert_eq!(store.credential_loads(), 2);
|
||||
assert_eq!(leases.requests.lock().expect("leases").len(), 1);
|
||||
let mut completed = false;
|
||||
while let Some(event) = stream.next().await {
|
||||
let event = event.expect("upstream event");
|
||||
completed |= event
|
||||
.wire_event()
|
||||
.is_some_and(|wire| wire.data()["type"] == "response.completed");
|
||||
}
|
||||
assert!(
|
||||
completed,
|
||||
"the successful upstream response must reach the client stream"
|
||||
);
|
||||
upstream.verify().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn snapshot_retry_rechecks_account_availability_without_sending() {
|
||||
for change in ["disabled", "revoked", "exhausted"] {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_provider_contract").await;
|
||||
store.on_credential_load(Arc::new(move |store, id, count| {
|
||||
Box::pin(async move {
|
||||
assert_eq!(count, 1, "an unavailable account cannot be selected again");
|
||||
let account = store.account(id.as_str()).expect("account");
|
||||
match change {
|
||||
"disabled" => store.set_enabled(id, false).await?,
|
||||
"revoked" => {
|
||||
store
|
||||
.apply_state_change(gateway_core::account::AccountStateChange {
|
||||
account_id: id.clone(),
|
||||
expected_revision: account.revision(),
|
||||
credential_state: CredentialState::Invalid,
|
||||
observed_at: SystemTime::now(),
|
||||
error_reason: None,
|
||||
message: None,
|
||||
})
|
||||
.await?
|
||||
}
|
||||
"exhausted" => {
|
||||
store
|
||||
.apply_quota_access(QuotaAccessChange {
|
||||
account_id: id.clone(),
|
||||
expected_revision: account.revision(),
|
||||
state: QuotaState::exhausted(
|
||||
QuotaEvidence::UsageLimitReached,
|
||||
SystemTime::now(),
|
||||
None,
|
||||
),
|
||||
})
|
||||
.await?;
|
||||
}
|
||||
_ => unreachable!("test case"),
|
||||
}
|
||||
Ok(())
|
||||
})
|
||||
}));
|
||||
let upstream = MockServer::start().await;
|
||||
let leases = Arc::new(TestLeaseCoordinator::default());
|
||||
let provider = provider_with_affinity_and_base_url_and_leases(
|
||||
&store,
|
||||
Arc::new(MemorySessionAffinity::default()),
|
||||
upstream.uri(),
|
||||
Arc::clone(&leases),
|
||||
);
|
||||
let Err(error) = provider
|
||||
.execute(
|
||||
planned_request("openai", http_generate_operation()),
|
||||
context("req_snapshot_unavailable", CancellationToken::new()),
|
||||
)
|
||||
.await
|
||||
else {
|
||||
panic!("{change} account must not produce an upstream stream");
|
||||
};
|
||||
let expected = if change == "exhausted" {
|
||||
ProviderErrorKind::QuotaExhausted
|
||||
} else {
|
||||
ProviderErrorKind::NoEligibleAccount
|
||||
};
|
||||
assert_eq!(error.kind(), expected, "{change}");
|
||||
assert_eq!(error.send_state(), UpstreamSendState::NotSent);
|
||||
assert!(leases.requests.lock().expect("leases").is_empty());
|
||||
assert!(
|
||||
upstream
|
||||
.received_requests()
|
||||
.await
|
||||
.expect("requests")
|
||||
.is_empty()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn snapshot_retry_cannot_switch_a_disabled_native_continuation_account() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_provider_contract").await;
|
||||
create_account(&store, "acct_scope_new").await;
|
||||
store.on_credential_load(Arc::new(|store, id, count| {
|
||||
Box::pin(async move {
|
||||
assert_eq!(id.as_str(), "acct_provider_contract");
|
||||
assert_eq!(count, 1);
|
||||
store.set_enabled(id, false).await
|
||||
})
|
||||
}));
|
||||
let upstream = MockServer::start().await;
|
||||
let leases = Arc::new(TestLeaseCoordinator::default());
|
||||
let provider = provider_with_affinity_and_base_url_and_leases(
|
||||
&store,
|
||||
Arc::new(MemorySessionAffinity::default()),
|
||||
upstream.uri(),
|
||||
Arc::clone(&leases),
|
||||
);
|
||||
let Err(error) = provider
|
||||
.execute(
|
||||
planned_request("openai", http_generate_operation()),
|
||||
pinned_continuation_context(
|
||||
"req_snapshot_pinned",
|
||||
"acct_provider_contract",
|
||||
"client-previous",
|
||||
"upstream-previous",
|
||||
1,
|
||||
ContinuationAttempt::Native,
|
||||
),
|
||||
)
|
||||
.await
|
||||
else {
|
||||
panic!("native continuation cannot move to the healthy second account");
|
||||
};
|
||||
assert_eq!(error.kind(), ProviderErrorKind::NoEligibleAccount);
|
||||
assert_eq!(error.send_state(), UpstreamSendState::NotSent);
|
||||
assert!(leases.requests.lock().expect("leases").is_empty());
|
||||
assert!(
|
||||
upstream
|
||||
.received_requests()
|
||||
.await
|
||||
.expect("requests")
|
||||
.is_empty()
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn repeated_snapshot_conflicts_are_bounded_and_report_the_selection_stage() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_provider_contract").await;
|
||||
store.on_credential_load(Arc::new(|store, id, count| {
|
||||
Box::pin(async move {
|
||||
assert!(count <= 4, "snapshot conflict retries must be bounded");
|
||||
let account = store.account(id.as_str()).expect("account");
|
||||
let observed_at = account
|
||||
.quota()
|
||||
.observed_at()
|
||||
.unwrap_or_else(SystemTime::now)
|
||||
+ Duration::from_millis(1);
|
||||
store
|
||||
.apply_quota_access(QuotaAccessChange {
|
||||
account_id: id.clone(),
|
||||
expected_revision: account.revision(),
|
||||
state: QuotaState::allowed(observed_at),
|
||||
})
|
||||
.await?;
|
||||
Ok(())
|
||||
})
|
||||
}));
|
||||
let upstream = MockServer::start().await;
|
||||
let leases = Arc::new(TestLeaseCoordinator::default());
|
||||
let provider = provider_with_affinity_and_base_url_and_leases(
|
||||
&store,
|
||||
Arc::new(MemorySessionAffinity::default()),
|
||||
upstream.uri(),
|
||||
Arc::clone(&leases),
|
||||
);
|
||||
let Err(error) = provider
|
||||
.execute(
|
||||
planned_request("openai", http_generate_operation()),
|
||||
context("req_snapshot_conflicts", CancellationToken::new()),
|
||||
)
|
||||
.await
|
||||
else {
|
||||
panic!("continuously changing snapshots must not be sent");
|
||||
};
|
||||
assert_eq!(
|
||||
error.kind(),
|
||||
ProviderErrorKind::ProviderInfrastructureUnavailable
|
||||
);
|
||||
assert_eq!(error.send_state(), UpstreamSendState::NotSent);
|
||||
let diagnostic = error.diagnostic().expect("selection diagnostic");
|
||||
assert_eq!(diagnostic.stage(), Some("account_selection"));
|
||||
assert_eq!(diagnostic.code(), Some("account_snapshot_conflict"));
|
||||
assert_eq!(store.credential_loads(), 4);
|
||||
assert!(leases.requests.lock().expect("leases").is_empty());
|
||||
assert!(
|
||||
upstream
|
||||
.received_requests()
|
||||
.await
|
||||
.expect("requests")
|
||||
.is_empty()
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn api_websocket_precheck_and_selection_share_the_snapshot_retry_budget() {
|
||||
let upstream = MockServer::start().await;
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
store
|
||||
.seed_api_key(
|
||||
"acct_provider_contract",
|
||||
upstream.uri(),
|
||||
provider_openai::credential::ApiKeyTransport::PreferWebsocket,
|
||||
)
|
||||
.await;
|
||||
store.on_credential_load(Arc::new(|store, id, count| {
|
||||
Box::pin(async move {
|
||||
assert!(
|
||||
count <= 6,
|
||||
"both credential checks must share one retry budget"
|
||||
);
|
||||
// 第 1、4 次在传输预检冲突,第 3、6 次在最终候选校验冲突。
|
||||
if matches!(count, 1 | 3 | 4 | 6) {
|
||||
let account = store.account(id.as_str()).expect("account");
|
||||
let observed_at = account
|
||||
.quota()
|
||||
.observed_at()
|
||||
.unwrap_or_else(SystemTime::now)
|
||||
+ Duration::from_millis(1);
|
||||
store
|
||||
.apply_quota_access(QuotaAccessChange {
|
||||
account_id: id.clone(),
|
||||
expected_revision: account.revision(),
|
||||
state: QuotaState::allowed(observed_at),
|
||||
})
|
||||
.await?;
|
||||
}
|
||||
Ok(())
|
||||
})
|
||||
}));
|
||||
let leases = Arc::new(TestLeaseCoordinator::default());
|
||||
let provider = provider_with_affinity_and_base_url_and_leases(
|
||||
&store,
|
||||
Arc::new(MemorySessionAffinity::default()),
|
||||
upstream.uri(),
|
||||
Arc::clone(&leases),
|
||||
);
|
||||
let warmup = Operation::Generate(GenerateRequest::from_protocol_payload(
|
||||
ProtocolPayload::json_object(
|
||||
"openai",
|
||||
json!({"model":"gpt-5.4","input":[],"store":false,"generate":false})
|
||||
.as_object()
|
||||
.expect("object")
|
||||
.clone(),
|
||||
)
|
||||
.expect("warmup"),
|
||||
));
|
||||
let Err(error) = provider
|
||||
.execute(
|
||||
planned_request("openai", warmup),
|
||||
context("req_api_snapshot_conflicts", CancellationToken::new()),
|
||||
)
|
||||
.await
|
||||
else {
|
||||
panic!("mixed snapshot conflicts must stop before WebSocket connection");
|
||||
};
|
||||
assert_eq!(
|
||||
error.kind(),
|
||||
ProviderErrorKind::ProviderInfrastructureUnavailable
|
||||
);
|
||||
assert_eq!(
|
||||
error.diagnostic().and_then(|diagnostic| diagnostic.code()),
|
||||
Some("account_snapshot_conflict")
|
||||
);
|
||||
assert_eq!(store.credential_loads(), 6);
|
||||
assert!(leases.requests.lock().expect("leases").is_empty());
|
||||
assert!(
|
||||
upstream
|
||||
.received_requests()
|
||||
.await
|
||||
.expect("requests")
|
||||
.is_empty()
|
||||
);
|
||||
}
|
||||
|
||||
async fn exhaust_account_quota(store: &Arc<MemoryAccountStore>, account_id: &str) {
|
||||
let account = store.account(account_id).expect("test account");
|
||||
store
|
||||
|
||||
@@ -41,11 +41,21 @@ struct StoredAccount {
|
||||
state_observed_at: Option<SystemTime>,
|
||||
}
|
||||
|
||||
type CredentialLoadHook = dyn for<'a> Fn(
|
||||
&'a MemoryAccountStore,
|
||||
&'a ProviderAccountId,
|
||||
usize,
|
||||
) -> BoxFuture<'a, Result<(), StoreError>>
|
||||
+ Send
|
||||
+ Sync;
|
||||
|
||||
#[derive(Default)]
|
||||
pub(crate) struct MemoryAccountStore {
|
||||
accounts: Mutex<BTreeMap<ProviderAccountId, StoredAccount>>,
|
||||
quota_reads: AtomicUsize,
|
||||
fail_provider_listing: AtomicBool,
|
||||
credential_load_hook: Mutex<Option<Arc<CredentialLoadHook>>>,
|
||||
credential_loads: AtomicUsize,
|
||||
}
|
||||
|
||||
impl MemoryAccountStore {
|
||||
@@ -175,6 +185,17 @@ impl MemoryAccountStore {
|
||||
pub(crate) fn fail_provider_listing(&self) {
|
||||
self.fail_provider_listing.store(true, Ordering::SeqCst);
|
||||
}
|
||||
|
||||
pub(crate) fn on_credential_load(&self, hook: Arc<CredentialLoadHook>) {
|
||||
*self
|
||||
.credential_load_hook
|
||||
.lock()
|
||||
.expect("credential hook lock") = Some(hook);
|
||||
}
|
||||
|
||||
pub(crate) fn credential_loads(&self) -> usize {
|
||||
self.credential_loads.load(Ordering::SeqCst)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -266,6 +287,16 @@ impl ProviderAccountStore for MemoryAccountStore {
|
||||
account: &ProviderAccountId,
|
||||
expected_revision: CredentialRevision,
|
||||
) -> Result<LoadedCredential, StoreError> {
|
||||
let count = self.credential_loads.fetch_add(1, Ordering::SeqCst) + 1;
|
||||
let hook = self
|
||||
.credential_load_hook
|
||||
.lock()
|
||||
.expect("credential hook lock")
|
||||
.clone();
|
||||
// 在列表快照与凭据读取之间确定性提交并发变更,不依赖线程调度或延时。
|
||||
if let Some(hook) = hook {
|
||||
hook(self, account, count).await?;
|
||||
}
|
||||
let loaded = self.load_current_credential(account).await?;
|
||||
if loaded.account.revision() != expected_revision {
|
||||
return Err(store_error(StoreErrorKind::Conflict));
|
||||
|
||||
@@ -14,3 +14,4 @@ b20e425c0b062b72d35fdebe294e023eb17ceb44ac93083ff3d48400715bf430 0012_upstream_
|
||||
12a540be31d65c4925cbe63dab42ea2d4ec2864fecf71a6dd4ee942e71726857 0014_disable_fast.sql
|
||||
9bd37bd314b6864816c43cb4fd1869fb415920d4f10f766976e344334a1f60fc 0015_proxy_dual_stack_exit.sql
|
||||
a1e23b8535ed41de27861833cd2d5ee7498a9bb0e523156ba3f011f9c6b3966b 0016_client_profiles_and_model_pricing.sql
|
||||
fcfae912ddefd54551ebf8e7e733442519d2862d2977134d548ebc1e75c1dc34 0017_unlimited_default_account_concurrency.sql
|
||||
|
||||
@@ -0,0 +1,8 @@
|
||||
alter table runtime_settings
|
||||
drop constraint runtime_settings_refresh_ck,
|
||||
add constraint runtime_settings_refresh_ck check (
|
||||
refresh_margin_seconds > 0
|
||||
and refresh_concurrency > 0
|
||||
and max_concurrent_per_account >= 0
|
||||
and request_interval_ms >= 0
|
||||
);
|
||||
+10
@@ -934,6 +934,8 @@ HTTP 请求头及新建 WS 的握手提示按当时的最终出站档位构造
|
||||
列表数据为 `{ items, page, configRevision }`,其中 item 返回 `memberCount`、按 Provider 聚合的
|
||||
`providerCounts` 和 `clientKeyCount`。查询分组成员使用账号列表的 `groupId` 筛选,
|
||||
不提供独立的分组成员路由;账号的 Provider 不代表整个分组的 Provider。
|
||||
`capacity.totalSlots` 为 `number | null`:`null` 表示可用成员中存在继承无限并发的账号,`0` 表示没有可用槽位。
|
||||
`capacity.usedSlots` 继续返回实际在途数;Redis 不可用时为 `null`。
|
||||
|
||||
## 7. Client Key
|
||||
|
||||
@@ -1071,6 +1073,10 @@ accountAutoFreezeAdaptiveConcurrency
|
||||
关联代理配置了自定义位置时优先使用代理值。保存后通过现有配置发布机制对新请求生效,
|
||||
已开始请求及其重试保持同一份全局值;普通文本、绝对时间戳和数据驻留要求不受影响。
|
||||
|
||||
`maxConcurrentPerAccount` 是默认账号并发上限,取值 0~4294967295;`0` 表示不限制。
|
||||
账号的 `concurrencyLimit: null` 继承该默认值,单独设置的正数上限仍优先生效。
|
||||
无限并发仍统计在途请求,并遵守最小请求间隔、账号可用性与 Client Key 限制。
|
||||
|
||||
`maxWaitingPerKey` 与 `maxWaitingPerAccount` 是全局统一的排队容量,取值 0~1,000,默认 0(关闭);
|
||||
每个 Key、每个账号各自独立计数,没有单对象覆盖字段。执行并发为 5、最大排队数为 5 时,
|
||||
该对象最多容纳 5 个执行请求与 5 个等待请求。Key 并发为 0(不限)时跳过 Key 排队。
|
||||
@@ -1325,6 +1331,10 @@ errorCode, errorMessage, startedAt, completedAt, expiresAt, createdAt, updatedAt
|
||||
| `GET` | `/api/admin/usage/insights/diagnostics` | 按维度聚合诊断 |
|
||||
| `GET` | `/api/admin/operations/errors` | 运维错误分页列表 |
|
||||
|
||||
Dashboard 的 `capacityInfo.maxConcurrentPerAccount` 为默认账号并发上限,`0` 表示不限制。
|
||||
`capacityInfo.totalSlots` 为 `number | null`;可用账号池含无限并发账号时为 `null`,此时 `availableSlots` 也为 `null`。
|
||||
`usedSlots` 仍表示实际在途数,Redis 不可用时为 `null`;没有可用账号时 `totalSlots` 为 `0`。
|
||||
|
||||
用量查询可组合页码/游标、时间范围、Provider、Client Key、账号、模型、route、transport、状态码、
|
||||
request/response/upstream ID、outcome 与搜索文本。诊断 `dimension` 可取 `model`、`account`、
|
||||
`apiKey`、`provider`、`transport`、`failureClass`、`status`。
|
||||
|
||||
@@ -355,6 +355,8 @@ Client Key 与账号分组形成授权范围:
|
||||
- 分组可以包含多个 Provider,账号也可以属于多个分组。
|
||||
|
||||
账号选择综合启停状态、credential/quota 事实、Redis cooldown、并发上限、权重、请求间隔和会话亲和。
|
||||
默认账号并发上限为 `0` 时表示无限;账号独立正数上限仍优先,未设置则继承默认值。
|
||||
无限并发只跳过并发上限判断,保留在途租约计数、请求间隔及其他准入约束;含无限账号的容量不投影有限占用比例。
|
||||
`account::AccountModelAccess` 拥有管理员模型政策的校验与精确匹配语义,存入账号行的 `model_access_json`,
|
||||
由 `RuntimeAccountDirectory` / `FrozenAccountScope` 随配置快照冻结。Provider 在额度、亲和与租约之前
|
||||
按映射后的上游模型筛选账号;重试和 fallback 使用同一冻结政策。上游目录和凭据不承载或改写该政策。
|
||||
|
||||
@@ -16,7 +16,7 @@ export interface AccountGroupAccountSummary {
|
||||
|
||||
export interface AccountGroupCapacity {
|
||||
usedSlots: number | null
|
||||
totalSlots: number
|
||||
totalSlots: number | null
|
||||
}
|
||||
|
||||
export interface AccountGroupUsage {
|
||||
|
||||
@@ -165,7 +165,7 @@ export interface DashboardPoolSummary {
|
||||
|
||||
export interface DashboardCapacityInfo {
|
||||
maxConcurrentPerAccount: number
|
||||
totalSlots: number
|
||||
totalSlots: number | null
|
||||
usedSlots: number | null
|
||||
availableSlots: number | null
|
||||
}
|
||||
|
||||
@@ -5,7 +5,6 @@ import { RefreshCw } from '@lucide/vue'
|
||||
import BaseButton from '@/components/base/BaseButton.vue'
|
||||
import BaseIconButton from '@/components/base/BaseIconButton.vue'
|
||||
import BaseModal from '@/components/base/BaseModal/index.vue'
|
||||
import BaseScrollbar from '@/components/base/BaseScrollbar.vue'
|
||||
import BaseSelect from '@/components/base/BaseSelect.vue'
|
||||
import AccountIdentityCell from './AccountIdentityCell.vue'
|
||||
import AccountStatusBadge from './AccountStatusBadge/index.vue'
|
||||
@@ -176,46 +175,40 @@ function connectionLogClass(tone: string) {
|
||||
<p class="m-0 text-cp-xs font-heavy text-cp-text-quaternary">
|
||||
事件轨迹
|
||||
</p>
|
||||
<BaseScrollbar max-height="260px">
|
||||
<div class="pt-2">
|
||||
<div v-if="logs.length === 0" class="text-cp-sm font-emphasis text-cp-text-quaternary">
|
||||
-
|
||||
</div>
|
||||
<div v-else class="flex flex-col gap-1.5">
|
||||
<div
|
||||
v-for="item in logs"
|
||||
:key="item.key"
|
||||
class="grid grid-cols-[54px_minmax(0,1fr)] gap-2 text-cp-sm leading-[1.45] font-emphasis"
|
||||
>
|
||||
<span class="font-mono text-cp-text-quaternary">{{ item.time }}</span>
|
||||
<div class="min-w-0">
|
||||
<div class="pt-2">
|
||||
<div v-if="logs.length === 0" class="text-cp-sm font-emphasis text-cp-text-quaternary">
|
||||
-
|
||||
</div>
|
||||
<div v-else class="flex flex-col gap-1.5">
|
||||
<div
|
||||
v-for="item in logs"
|
||||
:key="item.key"
|
||||
class="grid grid-cols-[54px_minmax(0,1fr)] gap-2 text-cp-sm leading-[1.45] font-emphasis"
|
||||
>
|
||||
<span class="font-mono text-cp-text-quaternary">{{ item.time }}</span>
|
||||
<div class="min-w-0">
|
||||
<p
|
||||
class="m-0 wrap-break-word"
|
||||
:class="connectionLogClass(item.tone)"
|
||||
>
|
||||
{{ item.text }}
|
||||
</p>
|
||||
<div v-if="item.detail" class="mt-2 rounded-lg bg-cp-fill-quaternary px-3 py-2">
|
||||
<p
|
||||
class="m-0 wrap-break-word"
|
||||
:class="connectionLogClass(item.tone)"
|
||||
v-if="item.tone === 'danger'"
|
||||
class="mt-0 mb-2 text-cp-xs font-heavy text-cp-text-quaternary"
|
||||
>
|
||||
{{ item.text }}
|
||||
原始诊断
|
||||
</p>
|
||||
<div v-if="item.detail" class="mt-2 rounded-lg bg-cp-fill-quaternary px-3 py-2">
|
||||
<p
|
||||
v-if="item.tone === 'danger'"
|
||||
class="mt-0 mb-2 text-cp-xs font-heavy text-cp-text-quaternary"
|
||||
>
|
||||
原始诊断
|
||||
</p>
|
||||
<BaseScrollbar max-height="138px">
|
||||
<div>
|
||||
<pre
|
||||
class="m-0 whitespace-pre-wrap wrap-break-word font-mono text-cp-xs leading-[1.6] font-emphasis text-cp-text"
|
||||
v-text="item.detail"
|
||||
/>
|
||||
</div>
|
||||
</BaseScrollbar>
|
||||
</div>
|
||||
<pre
|
||||
class="m-0 whitespace-pre-wrap wrap-break-word font-mono text-cp-xs leading-[1.6] font-emphasis text-cp-text"
|
||||
v-text="item.detail"
|
||||
/>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</BaseScrollbar>
|
||||
</div>
|
||||
</div>
|
||||
</section>
|
||||
</div>
|
||||
|
||||
@@ -19,15 +19,17 @@ const props = defineProps<{
|
||||
rotationStrategy?: string | null
|
||||
}>()
|
||||
|
||||
const unlimitedCapacity = computed(() => props.capacity?.maxConcurrentPerAccount === 0 && props.capacity.totalSlots === null)
|
||||
|
||||
const scheduleStats = computed(() => {
|
||||
const cap = props.capacity
|
||||
const display = (value: number | null | undefined) => value === null || value === undefined
|
||||
? '—'
|
||||
: formatCompactNumber(value)
|
||||
return [
|
||||
{ label: '默认并发', value: display(cap?.maxConcurrentPerAccount) },
|
||||
{ label: '总槽位', value: display(cap?.totalSlots) },
|
||||
{ label: '空闲槽位', value: display(cap?.availableSlots) },
|
||||
{ label: '默认并发', value: cap?.maxConcurrentPerAccount === 0 ? '不限制' : display(cap?.maxConcurrentPerAccount) },
|
||||
{ label: '总槽位', value: unlimitedCapacity.value ? '不限制' : display(cap?.totalSlots) },
|
||||
{ label: '空闲槽位', value: unlimitedCapacity.value ? '不限制' : display(cap?.availableSlots) },
|
||||
]
|
||||
})
|
||||
|
||||
@@ -44,6 +46,8 @@ const usedProgressStyle = computed(() => {
|
||||
|
||||
const usedRatio = computed(() => {
|
||||
const cap = props.capacity
|
||||
if (unlimitedCapacity.value)
|
||||
return `${cap?.usedSlots === null || cap?.usedSlots === undefined ? '—' : formatCompactNumber(cap.usedSlots)} / ∞`
|
||||
if (!cap || cap.totalSlots === null || cap.totalSlots === undefined)
|
||||
return '— / —'
|
||||
if (cap.usedSlots === null || cap.usedSlots === undefined)
|
||||
@@ -180,7 +184,7 @@ const statusBars = computed(() => {
|
||||
{{ usedRatio }}
|
||||
</strong>
|
||||
</div>
|
||||
<div class="mt-4 h-2.5 w-full overflow-hidden rounded-full bg-cp-progress-remaining">
|
||||
<div v-if="!unlimitedCapacity" class="mt-4 h-2.5 w-full overflow-hidden rounded-full bg-cp-progress-remaining">
|
||||
<i
|
||||
class="block h-2.5 rounded-full bg-cp-success"
|
||||
:style="usedProgressStyle"
|
||||
|
||||
@@ -57,7 +57,7 @@ const metrics = computed<MetricItem[]>(() => {
|
||||
return [
|
||||
{
|
||||
label: '并发',
|
||||
value: `${props.group.capacity.usedSlots ?? '—'} / ${props.group.capacity.totalSlots}`,
|
||||
value: `${props.group.capacity.usedSlots ?? '—'} / ${props.group.capacity.totalSlots ?? '∞'}`,
|
||||
tone: (props.group.capacity.usedSlots ?? 0) > 0 ? 'active' : 'primary',
|
||||
icon: Grid2X2,
|
||||
},
|
||||
|
||||
@@ -15,12 +15,15 @@ const requestIntervalMs = defineModel<string>('requestIntervalMs', { required: t
|
||||
<BaseForm class="max-w-6xl sm:grid-cols-2">
|
||||
<BaseFormItem
|
||||
label="默认账号并发上限"
|
||||
description="账号未单独设置时使用的并发上限"
|
||||
description="账号未单独设置时使用的并发上限,0 表示不限制"
|
||||
>
|
||||
<BaseInput
|
||||
v-model="maxConcurrentPerAccount"
|
||||
aria-label="默认账号并发上限"
|
||||
type="number"
|
||||
min="0"
|
||||
max="4294967295"
|
||||
step="1"
|
||||
>
|
||||
<template #prefix>
|
||||
<Gauge class="size-4" />
|
||||
|
||||
@@ -192,6 +192,10 @@ export function useSettingsForm() {
|
||||
toast.warning('请完整填写并发、队列、凭据刷新参数和调度策略')
|
||||
return
|
||||
}
|
||||
if (!Number.isInteger(maxConcurrentPerAccount) || maxConcurrentPerAccount < 0 || maxConcurrentPerAccount > 4294967295) {
|
||||
toast.warning('默认账号并发上限应为 0~4294967295 的整数,0 表示不限制')
|
||||
return
|
||||
}
|
||||
if (responsesMaxDecompressedBodyMiB === null || !Number.isInteger(responsesMaxDecompressedBodyMiB) || responsesMaxDecompressedBodyMiB < 1
|
||||
|| !Number.isSafeInteger(responsesMaxDecompressedBodyMiB * MIB)) {
|
||||
toast.warning('Responses 解压上限应为有效的正整数(MiB)')
|
||||
|
||||
Reference in New Issue
Block a user