mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-06 17:37:47 +08:00
2007 lines
71 KiB
Rust
2007 lines
71 KiB
Rust
mod capability;
|
|
mod plan;
|
|
mod presets;
|
|
mod provider;
|
|
mod quota;
|
|
mod quota_refresh;
|
|
mod service;
|
|
|
|
pub mod providers;
|
|
|
|
pub use capability::{ProviderPoolCapabilities, ProviderPoolCapability};
|
|
pub use plan::{derive_oauth_plan_type, derive_plan_tier, normalize_provider_plan_tier};
|
|
pub use presets::{
|
|
build_admin_pool_scheduling_presets_payload, normalize_provider_scheduling_presets,
|
|
};
|
|
pub use provider::{ProviderPoolAdapter, ProviderPoolMemberInput};
|
|
pub use providers::{
|
|
build_antigravity_pool_quota_request, build_antigravity_pool_quota_summary_request,
|
|
build_chatgpt_web_pool_quota_request, build_codex_pool_quota_request,
|
|
build_codex_pool_reset_credit_consume_request, build_codex_pool_reset_credits_request,
|
|
build_gemini_cli_pool_quota_request, build_kiro_pool_quota_request,
|
|
build_windsurf_pool_model_configs_request,
|
|
build_windsurf_pool_model_configs_request_with_base_url, build_windsurf_pool_quota_request,
|
|
build_windsurf_pool_quota_request_with_base_url, build_windsurf_pool_rate_limit_request,
|
|
build_windsurf_pool_rate_limit_request_with_base_url, build_xai_pool_billing_request,
|
|
build_xai_pool_user_request, enrich_chatgpt_web_quota_metadata, grok_mode_id_for_model,
|
|
grok_pool_tier_from_quota_bucket, grok_quota_window_key_for_model,
|
|
grok_supported_quota_windows_for_tier, normalize_chatgpt_web_image_quota_limit,
|
|
AntigravityProviderPoolAdapter, ChatGptWebProviderPoolAdapter, CodexProviderPoolAdapter,
|
|
DefaultProviderPoolAdapter, GeminiCliProviderPoolAdapter, GrokProviderPoolAdapter,
|
|
KiroPoolQuotaAuthInput, KiroProviderPoolAdapter, UnsupportedQuotaProviderPoolAdapter,
|
|
XaiProviderPoolAdapter, ANTIGRAVITY_FETCH_AVAILABLE_MODELS_PATH,
|
|
ANTIGRAVITY_RETRIEVE_USER_QUOTA_SUMMARY_PATH, CHATGPT_WEB_CONVERSATION_INIT_PATH,
|
|
CHATGPT_WEB_DEFAULT_BASE_URL, CODEX_WHAM_RESET_CREDITS_CONSUME_URL,
|
|
CODEX_WHAM_RESET_CREDITS_URL, CODEX_WHAM_USAGE_URL, GEMINI_CLI_RETRIEVE_USER_QUOTA_PATH,
|
|
GEMINI_CLI_USER_AGENT, KIRO_USAGE_LIMITS_PATH, KIRO_USAGE_SDK_VERSION,
|
|
WINDSURF_MODEL_CONFIGS_PATH, WINDSURF_RATE_LIMIT_PATH, WINDSURF_USER_STATUS_PATH,
|
|
XAI_BILLING_PATH, XAI_USER_PATH,
|
|
};
|
|
pub use quota::{
|
|
provider_pool_codex_metadata_has_account_quota, provider_pool_key_account_quota_exhausted,
|
|
provider_pool_key_minimum_quota_reached, provider_pool_key_model_quota_exhausted,
|
|
provider_pool_key_model_quota_hard_blocked, provider_pool_key_quota_hard_blocked,
|
|
provider_pool_key_scheduling_label, provider_pool_member_quota_snapshot,
|
|
provider_pool_quota_metadata_provider_type, provider_pool_quota_metadata_updated_at,
|
|
provider_pool_quota_snapshot_updated_at,
|
|
};
|
|
pub use quota_refresh::ProviderPoolQuotaRequestSpec;
|
|
pub use service::ProviderPoolService;
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
|
|
use aether_pool_core::PoolSchedulingPreset;
|
|
use serde_json::{json, Value};
|
|
|
|
fn sample_key(upstream_metadata: Option<Value>) -> StoredProviderCatalogKey {
|
|
let mut key = StoredProviderCatalogKey::new(
|
|
"key-1".to_string(),
|
|
"provider-1".to_string(),
|
|
"key-1".to_string(),
|
|
"oauth".to_string(),
|
|
None,
|
|
true,
|
|
)
|
|
.expect("key should build");
|
|
key.upstream_metadata = upstream_metadata;
|
|
key
|
|
}
|
|
|
|
#[test]
|
|
fn builtin_service_registers_provider_pool_adapters() {
|
|
let service = ProviderPoolService::with_builtin_adapters();
|
|
|
|
assert_eq!(
|
|
service.provider_types().collect::<Vec<_>>(),
|
|
[
|
|
"antigravity",
|
|
"chatgpt_web",
|
|
"claude_code",
|
|
"codex",
|
|
"gemini_cli",
|
|
"grok",
|
|
"kiro",
|
|
"vertex_ai",
|
|
"windsurf",
|
|
"xai"
|
|
]
|
|
);
|
|
assert!(service
|
|
.adapter("codex")
|
|
.capabilities()
|
|
.supports(ProviderPoolCapability::PlanTier));
|
|
assert_eq!(service.adapter("unknown").provider_type(), "default");
|
|
}
|
|
|
|
#[test]
|
|
fn builtin_service_owns_quota_refresh_support_and_endpoint_selection() {
|
|
let service = ProviderPoolService::with_builtin_adapters();
|
|
|
|
assert_eq!(
|
|
service.provider_types_for_capability(ProviderPoolCapability::QuotaRefresh),
|
|
[
|
|
"antigravity",
|
|
"chatgpt_web",
|
|
"codex",
|
|
"gemini_cli",
|
|
"grok",
|
|
"kiro",
|
|
"windsurf",
|
|
"xai"
|
|
]
|
|
);
|
|
assert!(service.supports_quota_refresh("codex"));
|
|
assert!(service.supports_quota_refresh("antigravity"));
|
|
assert!(service.supports_quota_refresh("grok"));
|
|
assert!(service.supports_quota_refresh("gemini_cli"));
|
|
assert!(service.supports_quota_refresh("windsurf"));
|
|
assert!(service.supports_quota_refresh("xai"));
|
|
assert_eq!(
|
|
service.quota_refresh_unsupported_message("claude_code"),
|
|
"Claude Code 暂不支持自动刷新额度：上游没有稳定可用的账号额度查询接口"
|
|
);
|
|
assert_eq!(
|
|
service.quota_refresh_unsupported_message("vertex_ai"),
|
|
"Vertex AI 暂不支持自动刷新额度：额度属于 Google Cloud 项目/区域配额"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn codex_quota_request_adds_account_header_for_paid_accounts() {
|
|
let spec = build_codex_pool_quota_request(
|
|
"key-1",
|
|
Some(("authorization".to_string(), "Bearer access".to_string())),
|
|
None,
|
|
Some(&json!({
|
|
"plan_type": "plus",
|
|
"account_id": "acct-1"
|
|
})),
|
|
)
|
|
.expect("spec should build");
|
|
|
|
assert_eq!(
|
|
spec.headers.get("chatgpt-account-id").map(String::as_str),
|
|
Some("acct-1")
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn codex_quota_request_prefers_imported_authorization_header() {
|
|
let spec = build_codex_pool_quota_request(
|
|
"key-1",
|
|
Some((
|
|
"authorization".to_string(),
|
|
"Bearer jwt-access-token".to_string(),
|
|
)),
|
|
None,
|
|
Some(&json!({
|
|
"headers": {
|
|
"authorization": "Bearer imported-session"
|
|
}
|
|
})),
|
|
)
|
|
.expect("spec should build");
|
|
|
|
assert_eq!(
|
|
spec.headers.get("authorization").map(String::as_str),
|
|
Some("Bearer imported-session")
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn codex_agent_identity_quota_request_prefers_dynamic_assertion() {
|
|
let spec = build_codex_pool_quota_request(
|
|
"key-1",
|
|
Some((
|
|
"authorization".to_string(),
|
|
"AgentAssertion signed-at-request-time".to_string(),
|
|
)),
|
|
None,
|
|
Some(&json!({
|
|
"auth_mode": "agentIdentity",
|
|
"headers": {
|
|
"authorization": "Bearer stale-imported-session"
|
|
}
|
|
})),
|
|
)
|
|
.expect("spec should build");
|
|
|
|
assert_eq!(
|
|
spec.headers.get("authorization").map(String::as_str),
|
|
Some("AgentAssertion signed-at-request-time")
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn codex_nested_agent_identity_quota_request_prefers_dynamic_assertion() {
|
|
let spec = build_codex_pool_quota_request(
|
|
"key-1",
|
|
Some((
|
|
"authorization".to_string(),
|
|
"AgentAssertion signed-at-request-time".to_string(),
|
|
)),
|
|
None,
|
|
Some(&json!({
|
|
"agent_identity": {
|
|
"agent_runtime_id": "runtime-1",
|
|
"agent_private_key": "private-key"
|
|
},
|
|
"headers": {
|
|
"authorization": "Bearer stale-imported-session"
|
|
}
|
|
})),
|
|
)
|
|
.expect("spec should build");
|
|
|
|
assert_eq!(
|
|
spec.headers.get("authorization").map(String::as_str),
|
|
Some("AgentAssertion signed-at-request-time")
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn gemini_cli_quota_request_uses_v1internal_retrieve_user_quota() {
|
|
let spec = build_gemini_cli_pool_quota_request(
|
|
"key-1",
|
|
"https://cloudcode-pa.googleapis.com/",
|
|
("authorization".to_string(), "Bearer access".to_string()),
|
|
"project-1",
|
|
);
|
|
|
|
assert_eq!(spec.method, "POST");
|
|
assert_eq!(
|
|
spec.url,
|
|
"https://cloudcode-pa.googleapis.com/v1internal:retrieveUserQuota"
|
|
);
|
|
assert_eq!(
|
|
spec.headers.get("authorization").map(String::as_str),
|
|
Some("Bearer access")
|
|
);
|
|
assert_eq!(
|
|
spec.json_body.as_ref().and_then(|body| body.get("project")),
|
|
Some(&json!("project-1"))
|
|
);
|
|
assert_eq!(
|
|
spec.headers.get("user-agent").map(String::as_str),
|
|
Some(GEMINI_CLI_USER_AGENT)
|
|
);
|
|
assert!(spec
|
|
.json_body
|
|
.as_ref()
|
|
.is_some_and(|body| body.get("userAgent").is_none()));
|
|
assert_eq!(spec.client_api_format, "gemini:generate_content");
|
|
assert_eq!(spec.provider_api_format, "gemini_cli:retrieve_user_quota");
|
|
}
|
|
|
|
#[test]
|
|
fn codex_quota_request_uses_wham_usage_endpoint() {
|
|
let spec = build_codex_pool_quota_request(
|
|
"key-1",
|
|
Some(("authorization".to_string(), "Bearer access".to_string())),
|
|
None,
|
|
None,
|
|
)
|
|
.expect("spec should build");
|
|
|
|
assert_eq!(spec.method, "GET");
|
|
assert_eq!(spec.url, "https://chatgpt.com/backend-api/wham/usage");
|
|
assert_eq!(
|
|
spec.headers.get("authorization").map(String::as_str),
|
|
Some("Bearer access")
|
|
);
|
|
assert_eq!(
|
|
spec.headers.get("accept").map(String::as_str),
|
|
Some("application/json")
|
|
);
|
|
assert_eq!(spec.model_name.as_deref(), Some("codex-wham-usage"));
|
|
}
|
|
|
|
#[test]
|
|
fn codex_reset_credits_request_uses_wham_detail_endpoint() {
|
|
let spec = build_codex_pool_reset_credits_request(
|
|
"key-1",
|
|
Some(("authorization".to_string(), "Bearer access".to_string())),
|
|
None,
|
|
Some(&json!({
|
|
"plan_type": "plus",
|
|
"account_id": "acct-1"
|
|
})),
|
|
)
|
|
.expect("spec should build");
|
|
|
|
assert_eq!(spec.method, "GET");
|
|
assert_eq!(spec.url, CODEX_WHAM_RESET_CREDITS_URL);
|
|
assert_eq!(
|
|
spec.headers.get("authorization").map(String::as_str),
|
|
Some("Bearer access")
|
|
);
|
|
assert_eq!(
|
|
spec.headers.get("chatgpt-account-id").map(String::as_str),
|
|
Some("acct-1")
|
|
);
|
|
assert_eq!(spec.model_name.as_deref(), Some("codex-wham-reset-credits"));
|
|
}
|
|
|
|
#[test]
|
|
fn codex_reset_credit_consume_request_posts_redeem_request_id() {
|
|
let spec = build_codex_pool_reset_credit_consume_request(
|
|
"key-1",
|
|
Some(("authorization".to_string(), "Bearer access".to_string())),
|
|
None,
|
|
None,
|
|
"8ae6f1c7-7e9e-4f5d-9b8a-000000000000",
|
|
)
|
|
.expect("spec should build");
|
|
|
|
assert_eq!(spec.method, "POST");
|
|
assert_eq!(spec.url, CODEX_WHAM_RESET_CREDITS_CONSUME_URL);
|
|
assert_eq!(spec.content_type.as_deref(), Some("application/json"));
|
|
assert_eq!(
|
|
spec.json_body
|
|
.as_ref()
|
|
.and_then(|body| body.get("redeem_request_id")),
|
|
Some(&json!("8ae6f1c7-7e9e-4f5d-9b8a-000000000000"))
|
|
);
|
|
assert_eq!(
|
|
spec.json_body
|
|
.as_ref()
|
|
.and_then(|body| body.as_object())
|
|
.map(|body| body.len()),
|
|
Some(1)
|
|
);
|
|
assert_eq!(
|
|
spec.model_name.as_deref(),
|
|
Some("codex-wham-reset-credit-consume")
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn codex_quota_request_skips_account_header_for_free_accounts() {
|
|
let spec = build_codex_pool_quota_request(
|
|
"key-1",
|
|
Some(("authorization".to_string(), "Bearer access".to_string())),
|
|
None,
|
|
Some(&json!({
|
|
"plan_type": "codex:free",
|
|
"account_id": "acct-1"
|
|
})),
|
|
)
|
|
.expect("spec should build");
|
|
|
|
assert!(!spec.headers.contains_key("chatgpt-account-id"));
|
|
}
|
|
|
|
#[test]
|
|
fn kiro_quota_request_includes_profile_arn_when_present() {
|
|
let spec = build_kiro_pool_quota_request(
|
|
"key-1",
|
|
&KiroPoolQuotaAuthInput {
|
|
authorization_value: "Bearer access".to_string(),
|
|
api_region: "us-west-2".to_string(),
|
|
kiro_version: "0.3.210".to_string(),
|
|
machine_id: "machine".to_string(),
|
|
profile_arn: Some("arn:aws:sso:::profile/p-1".to_string()),
|
|
},
|
|
);
|
|
|
|
assert!(spec.url.contains("q.us-west-2.amazonaws.com"));
|
|
assert!(spec
|
|
.url
|
|
.contains("profileArn=arn%3Aaws%3Asso%3A%3A%3Aprofile%2Fp-1"));
|
|
}
|
|
|
|
#[test]
|
|
fn kiro_quota_request_rejects_region_url_injection() {
|
|
let spec = build_kiro_pool_quota_request(
|
|
"key-1",
|
|
&KiroPoolQuotaAuthInput {
|
|
authorization_value: "Bearer sensitive-access-token".to_string(),
|
|
api_region: "attacker.example/evil".to_string(),
|
|
kiro_version: "0.3.210".to_string(),
|
|
machine_id: "machine".to_string(),
|
|
profile_arn: None,
|
|
},
|
|
);
|
|
|
|
assert_eq!(
|
|
spec.url,
|
|
"https://q.us-east-1.amazonaws.com/getUsageLimits?origin=AI_EDITOR&resourceType=AGENTIC_REQUEST&isEmailRequired=true"
|
|
);
|
|
assert_eq!(
|
|
spec.headers.get("host").map(String::as_str),
|
|
Some("q.us-east-1.amazonaws.com")
|
|
);
|
|
assert!(!spec.url.contains("attacker.example"));
|
|
assert!(!spec
|
|
.headers
|
|
.get("host")
|
|
.is_some_and(|value| value.contains("attacker.example")));
|
|
}
|
|
|
|
#[test]
|
|
fn chatgpt_web_quota_request_uses_default_base_url_when_empty() {
|
|
let spec = build_chatgpt_web_pool_quota_request(
|
|
"key-1",
|
|
"",
|
|
("authorization".to_string(), "Bearer access".to_string()),
|
|
);
|
|
|
|
assert_eq!(
|
|
spec.url,
|
|
"https://chatgpt.com/backend-api/conversation/init"
|
|
);
|
|
assert_eq!(
|
|
spec.headers.get("origin").map(String::as_str),
|
|
Some("https://chatgpt.com")
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn chatgpt_web_quota_metadata_enriches_auth_and_uses_first_remaining_as_limit() {
|
|
let mut metadata = json!({
|
|
"image_quota_remaining": 12,
|
|
});
|
|
enrich_chatgpt_web_quota_metadata(
|
|
&mut metadata,
|
|
Some(&json!({
|
|
"plan": "free",
|
|
"email": "[email protected]",
|
|
"accountId": "acct-1"
|
|
})),
|
|
);
|
|
normalize_chatgpt_web_image_quota_limit(&mut metadata, None);
|
|
|
|
assert_eq!(metadata["plan_type"], json!("free"));
|
|
assert_eq!(metadata["email"], json!("[email protected]"));
|
|
assert_eq!(metadata["account_id"], json!("acct-1"));
|
|
assert_eq!(metadata["image_quota_total"], json!(12.0));
|
|
assert_eq!(metadata["image_quota_used"], json!(0.0));
|
|
}
|
|
|
|
#[test]
|
|
fn chatgpt_web_quota_metadata_preserves_existing_paid_limit() {
|
|
let mut metadata = json!({
|
|
"plan_type": "plus",
|
|
"image_quota_remaining": 7,
|
|
});
|
|
normalize_chatgpt_web_image_quota_limit(
|
|
&mut metadata,
|
|
Some(&json!({
|
|
"chatgpt_web": {
|
|
"image_quota_total": 40
|
|
}
|
|
})),
|
|
);
|
|
|
|
assert_eq!(metadata["image_quota_total"], json!(40.0));
|
|
assert_eq!(metadata["image_quota_used"], json!(33.0));
|
|
}
|
|
|
|
#[test]
|
|
fn chatgpt_web_quota_metadata_does_not_preserve_legacy_free_25_limit() {
|
|
let mut metadata = json!({
|
|
"plan_type": "free",
|
|
"image_quota_remaining": 19,
|
|
});
|
|
normalize_chatgpt_web_image_quota_limit(
|
|
&mut metadata,
|
|
Some(&json!({
|
|
"chatgpt_web": {
|
|
"plan_type": "free",
|
|
"image_quota_total": 25
|
|
}
|
|
})),
|
|
);
|
|
|
|
assert_eq!(metadata["image_quota_total"], json!(19.0));
|
|
assert_eq!(metadata["image_quota_used"], json!(0.0));
|
|
assert_eq!(
|
|
metadata["image_quota_limit_source"],
|
|
json!("first_remaining")
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn chatgpt_web_quota_metadata_ignores_upstream_free_25_default() {
|
|
let mut metadata = json!({
|
|
"plan_type": "free",
|
|
"image_quota_remaining": 19,
|
|
"image_quota_total": 25,
|
|
});
|
|
normalize_chatgpt_web_image_quota_limit(&mut metadata, None);
|
|
|
|
assert_eq!(metadata["image_quota_total"], json!(19.0));
|
|
assert_eq!(metadata["image_quota_used"], json!(0.0));
|
|
assert_eq!(
|
|
metadata["image_quota_limit_source"],
|
|
json!("first_remaining")
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn chatgpt_web_quota_metadata_preserves_marked_free_first_limit() {
|
|
let mut metadata = json!({
|
|
"plan_type": "free",
|
|
"image_quota_remaining": 18,
|
|
});
|
|
normalize_chatgpt_web_image_quota_limit(
|
|
&mut metadata,
|
|
Some(&json!({
|
|
"chatgpt_web": {
|
|
"plan_type": "free",
|
|
"image_quota_total": 19,
|
|
"image_quota_limit_source": "first_remaining"
|
|
}
|
|
})),
|
|
);
|
|
|
|
assert_eq!(metadata["image_quota_total"], json!(19.0));
|
|
assert_eq!(metadata["image_quota_used"], json!(1.0));
|
|
assert_eq!(
|
|
metadata["image_quota_limit_source"],
|
|
json!("first_remaining")
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn windsurf_quota_request_uses_user_status_connect_rpc() {
|
|
let spec = build_windsurf_pool_quota_request("key-ws", "session-token-123");
|
|
|
|
assert_eq!(spec.request_id, "windsurf-quota:key-ws");
|
|
assert_eq!(spec.method, "POST");
|
|
assert_eq!(
|
|
spec.url,
|
|
format!("https://server.codeium.com{WINDSURF_USER_STATUS_PATH}")
|
|
);
|
|
assert_eq!(spec.content_type.as_deref(), Some("application/json"));
|
|
assert_eq!(
|
|
spec.headers
|
|
.get("connect-protocol-version")
|
|
.map(String::as_str),
|
|
Some("1")
|
|
);
|
|
assert_eq!(
|
|
spec.json_body
|
|
.as_ref()
|
|
.and_then(|body| body.pointer("/metadata/apiKey"))
|
|
.and_then(Value::as_str),
|
|
Some("session-token-123")
|
|
);
|
|
assert_eq!(spec.provider_api_format, "windsurf:user_status");
|
|
}
|
|
|
|
#[test]
|
|
fn windsurf_model_and_rate_limit_requests_use_connect_rpc_metadata() {
|
|
let models = build_windsurf_pool_model_configs_request("key-ws", "api-key-123");
|
|
let rate_limit = build_windsurf_pool_rate_limit_request("key-ws", "api-key-123");
|
|
|
|
assert_eq!(
|
|
models.url,
|
|
format!("https://server.codeium.com{WINDSURF_MODEL_CONFIGS_PATH}")
|
|
);
|
|
assert_eq!(
|
|
rate_limit.url,
|
|
format!("https://server.codeium.com{WINDSURF_RATE_LIMIT_PATH}")
|
|
);
|
|
for spec in [models, rate_limit] {
|
|
assert_eq!(spec.method, "POST");
|
|
assert_eq!(
|
|
spec.headers
|
|
.get("connect-protocol-version")
|
|
.map(String::as_str),
|
|
Some("1")
|
|
);
|
|
assert_eq!(
|
|
spec.json_body
|
|
.as_ref()
|
|
.and_then(|body| body.pointer("/metadata/apiKey"))
|
|
.and_then(Value::as_str),
|
|
Some("api-key-123")
|
|
);
|
|
assert_eq!(spec.client_api_format, "openai:chat");
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn windsurf_rate_limit_metadata_keeps_member_schedulable() {
|
|
let service = ProviderPoolService::with_builtin_adapters();
|
|
let key = sample_key(Some(json!({
|
|
"windsurf": {
|
|
"updated_at": 1_700_000_000u64,
|
|
"rate_limit": {
|
|
"limited": true,
|
|
"retry_after_ms": 60_000
|
|
}
|
|
}
|
|
})));
|
|
|
|
let signals = service.member_signals("windsurf", &key, None, None);
|
|
|
|
assert!(!signals.quota_exhausted);
|
|
}
|
|
|
|
#[test]
|
|
fn windsurf_status_snapshot_ban_marks_member_exhausted() {
|
|
let service = ProviderPoolService::with_builtin_adapters();
|
|
let mut key = sample_key(Some(json!({
|
|
"windsurf": {
|
|
"updated_at": 1_700_000_000u64,
|
|
"daily_remaining_percent": 100.0
|
|
}
|
|
})));
|
|
key.status_snapshot = Some(json!({
|
|
"quota": {
|
|
"provider_type": "windsurf",
|
|
"code": "banned",
|
|
"exhausted": false,
|
|
"windows": [{
|
|
"code": "daily",
|
|
"used_ratio": 0.0,
|
|
"remaining_ratio": 1.0
|
|
}]
|
|
}
|
|
}));
|
|
|
|
let signals = service.member_signals("windsurf", &key, None, None);
|
|
|
|
assert!(signals.quota_exhausted);
|
|
}
|
|
|
|
#[test]
|
|
fn preset_payload_derives_provider_support_from_capabilities() {
|
|
let payload = build_admin_pool_scheduling_presets_payload();
|
|
let items = payload.as_array().expect("payload should be array");
|
|
let free_first = items
|
|
.iter()
|
|
.find(|item| item["name"] == "free_first")
|
|
.expect("free_first should exist");
|
|
let recent_refresh = items
|
|
.iter()
|
|
.find(|item| item["name"] == "recent_refresh")
|
|
.expect("recent_refresh should exist");
|
|
let legacy_free_team = items
|
|
.iter()
|
|
.find(|item| item["name"] == "free_team_first")
|
|
.expect("legacy free_team_first should remain configurable");
|
|
|
|
assert_eq!(
|
|
free_first["providers"],
|
|
json!(["codex", "grok", "kiro", "windsurf", "xai"])
|
|
);
|
|
assert_eq!(
|
|
recent_refresh["providers"],
|
|
json!(["codex", "grok", "kiro", "windsurf", "xai"])
|
|
);
|
|
assert_eq!(free_first["default_enabled"], json!(false));
|
|
assert_eq!(recent_refresh["default_enabled"], json!(false));
|
|
assert_eq!(
|
|
recent_refresh["default_enabled_providers"],
|
|
json!(["codex", "windsurf"])
|
|
);
|
|
assert_eq!(legacy_free_team["default_mode"], json!("both"));
|
|
assert_eq!(legacy_free_team["modes"].as_array().map(Vec::len), Some(3));
|
|
}
|
|
|
|
#[test]
|
|
fn quota_metadata_provider_type_comes_from_pool_registry() {
|
|
assert_eq!(
|
|
provider_pool_quota_metadata_provider_type(&json!({
|
|
"gemini_cli": {
|
|
"updated_at": 1_700_000_000u64
|
|
}
|
|
}))
|
|
.as_deref(),
|
|
Some("gemini_cli")
|
|
);
|
|
assert_eq!(
|
|
provider_pool_quota_metadata_provider_type(&json!({
|
|
"custom_provider": {
|
|
"updated_at": 1_700_000_000u64
|
|
}
|
|
}))
|
|
.as_deref(),
|
|
Some("custom_provider")
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn codex_adapter_injects_recent_refresh_and_filters_by_capability() {
|
|
let service = ProviderPoolService::with_builtin_adapters();
|
|
let normalized = service.normalize_scheduling_presets(
|
|
"codex",
|
|
&[PoolSchedulingPreset {
|
|
preset: "cache_affinity".to_string(),
|
|
enabled: true,
|
|
mode: None,
|
|
}],
|
|
);
|
|
|
|
assert_eq!(
|
|
normalized
|
|
.iter()
|
|
.map(|preset| preset.preset.as_str())
|
|
.collect::<Vec<_>>(),
|
|
["cache_affinity", "recent_refresh"]
|
|
);
|
|
|
|
let legacy_free_team = service.normalize_scheduling_presets(
|
|
"codex",
|
|
&[PoolSchedulingPreset {
|
|
preset: "free_team_first".to_string(),
|
|
enabled: true,
|
|
mode: Some("team_only".to_string()),
|
|
}],
|
|
);
|
|
assert_eq!(legacy_free_team[0].preset, "free_team_first");
|
|
assert_eq!(legacy_free_team[0].mode.as_deref(), Some("team_only"));
|
|
|
|
let unsupported = service.normalize_scheduling_presets(
|
|
"chatgpt_web",
|
|
&[PoolSchedulingPreset {
|
|
preset: "plus_first".to_string(),
|
|
enabled: true,
|
|
mode: None,
|
|
}],
|
|
);
|
|
assert!(unsupported.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn codex_minimum_quota_respects_boundary_and_provider() {
|
|
for (used_percent, expected) in [
|
|
(98.9, false),
|
|
(98.99999, false),
|
|
(99.0, true),
|
|
(99.5, true),
|
|
(100.0, true),
|
|
] {
|
|
let key = sample_key(Some(json!({
|
|
"codex": { "primary_used_percent": used_percent }
|
|
})));
|
|
assert_eq!(
|
|
provider_pool_key_minimum_quota_reached(&key, "codex", None),
|
|
expected,
|
|
"used_percent={used_percent}"
|
|
);
|
|
assert!(!provider_pool_key_minimum_quota_reached(&key, "kiro", None));
|
|
}
|
|
assert!(!provider_pool_key_minimum_quota_reached(
|
|
&sample_key(None),
|
|
"codex",
|
|
None
|
|
));
|
|
}
|
|
|
|
#[test]
|
|
fn codex_minimum_quota_short_model_names_still_use_account_windows() {
|
|
for (used_percent, expected) in [(98.0, false), (99.0, true)] {
|
|
let key = sample_key(Some(json!({
|
|
"codex": { "primary_used_percent": used_percent }
|
|
})));
|
|
for model in ["o1", "o3", "", " "] {
|
|
assert_eq!(
|
|
provider_pool_key_minimum_quota_reached(&key, "codex", Some(model)),
|
|
expected,
|
|
"model={model:?}, used_percent={used_percent}"
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn codex_minimum_quota_respects_windows_and_reset() {
|
|
let now = std::time::SystemTime::now()
|
|
.duration_since(std::time::UNIX_EPOCH)
|
|
.expect("system time should be after unix epoch")
|
|
.as_secs();
|
|
for (reset_at, expected) in [(now + 3600, true), (now - 60, false)] {
|
|
let mut key = sample_key(None);
|
|
key.status_snapshot = Some(json!({
|
|
"quota": {
|
|
"provider_type": "codex",
|
|
"updated_at": now - 600,
|
|
"windows": [
|
|
{ "code": "5h", "used_ratio": 0.2 },
|
|
{ "code": "weekly", "used_ratio": 0.99, "reset_at": reset_at }
|
|
]
|
|
}
|
|
}));
|
|
for model in [None, Some("gpt-5.4")] {
|
|
assert_eq!(
|
|
provider_pool_key_minimum_quota_reached(&key, "codex", model),
|
|
expected
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn codex_minimum_quota_checks_either_window_and_relative_resets() {
|
|
let now = std::time::SystemTime::now()
|
|
.duration_since(std::time::UNIX_EPOCH)
|
|
.expect("system time should be after unix epoch")
|
|
.as_secs();
|
|
for prefix in ["primary", "secondary"] {
|
|
for (reset_seconds, expected) in [(3600, true), (60, false)] {
|
|
let key = sample_key(Some(json!({
|
|
"codex": {
|
|
"updated_at": now - 600,
|
|
"allowed": true,
|
|
"limit_reached": false,
|
|
format!("{prefix}_used_percent"): 99.0,
|
|
format!("{prefix}_reset_after_seconds"): reset_seconds
|
|
}
|
|
})));
|
|
for model in [None, Some("gpt-5.4")] {
|
|
assert_eq!(
|
|
provider_pool_key_minimum_quota_reached(&key, "codex", model),
|
|
expected,
|
|
"prefix={prefix}, reset_seconds={reset_seconds}, model={model:?}"
|
|
);
|
|
}
|
|
assert!(!provider_pool_key_account_quota_exhausted(&key, "codex"));
|
|
}
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn codex_minimum_quota_uses_newest_source_with_or_without_model() {
|
|
for (snapshot_at, metadata_at, metadata_wins) in [
|
|
(Some(200), Some(100), false),
|
|
(Some(100), Some(200), true),
|
|
(None, None, false),
|
|
(Some(100), None, false),
|
|
(None, Some(100), true),
|
|
] {
|
|
for snapshot_reached in [false, true] {
|
|
let mut key = sample_key(Some(json!({
|
|
"codex": {
|
|
"updated_at": metadata_at,
|
|
"primary_used_percent": if snapshot_reached { 20.0 } else { 99.0 }
|
|
}
|
|
})));
|
|
key.status_snapshot = Some(json!({
|
|
"quota": {
|
|
"provider_type": "codex",
|
|
"observed_at": snapshot_at,
|
|
"windows": [{
|
|
"code": "5h",
|
|
"used_ratio": if snapshot_reached { 0.99 } else { 0.2 },
|
|
"is_exhausted": false
|
|
}]
|
|
}
|
|
}));
|
|
for model in [None, Some("gpt-5.4")] {
|
|
assert_eq!(
|
|
provider_pool_key_minimum_quota_reached(&key, "codex", model),
|
|
snapshot_reached != metadata_wins,
|
|
"snapshot_at={snapshot_at:?}, metadata_at={metadata_at:?}, model={model:?}"
|
|
);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn codex_minimum_quota_isolates_model_buckets_and_checks_each_model_window() {
|
|
for explicit_model in [false, true] {
|
|
let mut key = sample_key(None);
|
|
key.status_snapshot = Some(json!({
|
|
"quota": {
|
|
"provider_type": "codex",
|
|
"windows": [
|
|
{ "code": "5h", "used_ratio": 0.2 },
|
|
{
|
|
"code": "spark_5h", "used_ratio": 0.99,
|
|
"model": if explicit_model { Some("spark") } else { None },
|
|
"is_exhausted": false
|
|
},
|
|
{
|
|
"code": "spark_weekly", "used_ratio": 0.2,
|
|
"model": if explicit_model { Some("spark") } else { None },
|
|
"is_exhausted": false
|
|
}
|
|
]
|
|
}
|
|
}));
|
|
assert!(provider_pool_key_minimum_quota_reached(
|
|
&key,
|
|
"codex",
|
|
Some("gpt-5.3-codex-spark")
|
|
));
|
|
for model in [None, Some("gpt-5.4")] {
|
|
assert!(!provider_pool_key_minimum_quota_reached(
|
|
&key, "codex", model
|
|
));
|
|
}
|
|
assert_eq!(
|
|
provider_pool_key_model_quota_exhausted(&key, "codex", "gpt-5.3-codex-spark"),
|
|
Some(false)
|
|
);
|
|
}
|
|
for (account_percent, spark_percent) in [(99.0, 20.0), (20.0, 99.0)] {
|
|
let key = sample_key(Some(json!({
|
|
"codex": {
|
|
"primary_used_percent": account_percent,
|
|
"spark_primary_used_percent": spark_percent,
|
|
"spark_secondary_used_percent": 20.0
|
|
}
|
|
})));
|
|
assert_eq!(
|
|
provider_pool_key_minimum_quota_reached(&key, "codex", Some("gpt-5.3-codex-spark")),
|
|
spark_percent == 99.0
|
|
);
|
|
assert_eq!(
|
|
provider_pool_key_minimum_quota_reached(&key, "codex", Some("gpt-5.4")),
|
|
account_percent == 99.0
|
|
);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn codex_minimum_quota_keeps_model_bucket_when_account_metadata_is_newer() {
|
|
for (snapshot_ratio, account_percent) in [(0.2, 99.0), (0.99, 20.0)] {
|
|
let mut key = sample_key(Some(json!({
|
|
"codex": { "updated_at": 200, "primary_used_percent": account_percent }
|
|
})));
|
|
key.status_snapshot = Some(json!({
|
|
"quota": {
|
|
"provider_type": "codex",
|
|
"observed_at": 100,
|
|
"windows": [{ "code": "spark_5h", "used_ratio": snapshot_ratio }]
|
|
}
|
|
}));
|
|
assert_eq!(
|
|
provider_pool_key_minimum_quota_reached(&key, "codex", Some("gpt-5.3-codex-spark")),
|
|
snapshot_ratio == 0.99
|
|
);
|
|
assert_eq!(
|
|
provider_pool_key_minimum_quota_reached(&key, "codex", Some("gpt-5.4")),
|
|
account_percent == 99.0
|
|
);
|
|
|
|
key.upstream_metadata.as_mut().unwrap()["codex"]["spark_primary_used_percent"] =
|
|
json!(account_percent);
|
|
assert_eq!(
|
|
provider_pool_key_minimum_quota_reached(&key, "codex", Some("gpt-5.3-codex-spark")),
|
|
account_percent == 99.0,
|
|
"the newer observation for the same model bucket must win"
|
|
);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn codex_minimum_quota_supports_remaining_values_and_ignores_unknown_data() {
|
|
for (metrics, expected) in [
|
|
(json!({ "remaining_ratio": 0.01 }), true),
|
|
(json!({ "remaining_percent": "1" }), true),
|
|
(json!({ "remaining": 1, "limit": 100 }), true),
|
|
(json!({ "used_percent": "99" }), true),
|
|
(json!({ "used_ratio": null }), false),
|
|
(json!({ "used_ratio": "NaN" }), false),
|
|
(json!({ "remaining": 0, "limit": 0 }), false),
|
|
(json!({ "remaining_ratio": 0.01001 }), false),
|
|
] {
|
|
let key = sample_key(Some(json!({
|
|
"codex": { "quota_by_model": { "spark": metrics } }
|
|
})));
|
|
assert_eq!(
|
|
provider_pool_key_minimum_quota_reached(&key, "codex", Some("gpt-5.3-codex-spark")),
|
|
expected,
|
|
"metrics={metrics}"
|
|
);
|
|
assert!(!provider_pool_key_minimum_quota_reached(
|
|
&key, "codex", None
|
|
));
|
|
}
|
|
let disabled = sample_key(Some(json!({
|
|
"codex": { "primary_used_percent": 99.0, "primary_window_minutes": 0 }
|
|
})));
|
|
assert!(!provider_pool_key_minimum_quota_reached(
|
|
&disabled, "codex", None
|
|
));
|
|
|
|
let mut mismatched = sample_key(None);
|
|
mismatched.status_snapshot = Some(json!({
|
|
"quota": {
|
|
"provider_type": "kiro",
|
|
"windows": [{ "code": "weekly", "used_ratio": 0.99 }]
|
|
}
|
|
}));
|
|
assert!(!provider_pool_key_minimum_quota_reached(
|
|
&mismatched,
|
|
"codex",
|
|
None
|
|
));
|
|
}
|
|
|
|
#[test]
|
|
fn provider_quota_exhaustion_is_adapter_owned() {
|
|
assert!(provider_pool_key_account_quota_exhausted(
|
|
&sample_key(Some(json!({
|
|
"codex": {
|
|
"allowed": false,
|
|
"limit_reached": true
|
|
}
|
|
}))),
|
|
"codex",
|
|
));
|
|
assert!(provider_pool_key_account_quota_exhausted(
|
|
&sample_key(Some(json!({
|
|
"codex": {
|
|
"has_credits": false,
|
|
"credits_unlimited": false
|
|
}
|
|
}))),
|
|
"codex",
|
|
));
|
|
assert!(provider_pool_key_account_quota_exhausted(
|
|
&sample_key(Some(json!({
|
|
"kiro": {
|
|
"remaining": 0
|
|
}
|
|
}))),
|
|
"kiro",
|
|
));
|
|
assert!(provider_pool_key_account_quota_exhausted(
|
|
&sample_key(Some(json!({
|
|
"chatgpt_web": {
|
|
"image_quota_blocked": true
|
|
}
|
|
}))),
|
|
"chatgpt_web",
|
|
));
|
|
assert!(provider_pool_key_account_quota_exhausted(
|
|
&sample_key(Some(json!({
|
|
"grok": {
|
|
"quota_by_model": {
|
|
"quota_fast": {
|
|
"is_exhausted": true,
|
|
"remaining": 0.0
|
|
}
|
|
}
|
|
}
|
|
}))),
|
|
"grok",
|
|
));
|
|
assert!(!provider_pool_key_account_quota_exhausted(
|
|
&sample_key(Some(json!({
|
|
"grok": {
|
|
"pool_tier": "basic",
|
|
"quota_by_model": {
|
|
"quota_fast": {
|
|
"is_exhausted": false,
|
|
"remaining": 1.0
|
|
},
|
|
"quota_heavy": {
|
|
"is_exhausted": true,
|
|
"remaining": 0.0
|
|
}
|
|
}
|
|
}
|
|
}))),
|
|
"grok",
|
|
));
|
|
assert!(!provider_pool_key_account_quota_exhausted(
|
|
&sample_key(Some(json!({
|
|
"codex": {
|
|
"has_credits": false,
|
|
"credits_unlimited": true
|
|
}
|
|
}))),
|
|
"codex",
|
|
));
|
|
assert!(!provider_pool_key_account_quota_exhausted(
|
|
&sample_key(Some(json!({
|
|
"codex": {
|
|
"allowed": true,
|
|
"primary_used_percent": 100.0
|
|
}
|
|
}))),
|
|
"codex",
|
|
));
|
|
|
|
let mut explicit_codex_limit = sample_key(None);
|
|
explicit_codex_limit.status_snapshot = Some(json!({
|
|
"quota": {
|
|
"version": 2,
|
|
"provider_type": "codex",
|
|
"exhausted": true,
|
|
"allowed": false,
|
|
"limit_reached": true,
|
|
"usage_ratio": 0.91,
|
|
"windows": [{
|
|
"code": "weekly",
|
|
"used_ratio": 0.91
|
|
}]
|
|
}
|
|
}));
|
|
assert!(provider_pool_key_account_quota_exhausted(
|
|
&explicit_codex_limit,
|
|
"codex",
|
|
));
|
|
}
|
|
|
|
#[test]
|
|
fn provider_quota_exhaustion_snapshot_expires_after_reset_at() {
|
|
let now = std::time::SystemTime::now()
|
|
.duration_since(std::time::UNIX_EPOCH)
|
|
.expect("system time should be after unix epoch")
|
|
.as_secs();
|
|
|
|
let mut expired = sample_key(None);
|
|
expired.status_snapshot = Some(json!({
|
|
"quota": {
|
|
"version": 2,
|
|
"provider_type": "codex",
|
|
"code": "exhausted",
|
|
"exhausted": true,
|
|
"updated_at": now.saturating_sub(600),
|
|
"windows": [{
|
|
"code": "5h",
|
|
"used_ratio": 1.0,
|
|
"reset_at": now.saturating_sub(60),
|
|
"is_exhausted": true
|
|
}]
|
|
}
|
|
}));
|
|
assert!(!provider_pool_key_account_quota_exhausted(
|
|
&expired, "codex"
|
|
));
|
|
|
|
let mut active = sample_key(None);
|
|
active.status_snapshot = Some(json!({
|
|
"quota": {
|
|
"version": 2,
|
|
"provider_type": "codex",
|
|
"code": "exhausted",
|
|
"exhausted": true,
|
|
"updated_at": now,
|
|
"windows": [{
|
|
"code": "5h",
|
|
"used_ratio": 1.0,
|
|
"reset_at": now.saturating_add(3600),
|
|
"is_exhausted": true
|
|
}]
|
|
}
|
|
}));
|
|
assert!(provider_pool_key_account_quota_exhausted(&active, "codex"));
|
|
}
|
|
|
|
#[test]
|
|
fn antigravity_model_quota_exhaustion_does_not_block_other_models() {
|
|
let service = ProviderPoolService::with_builtin_adapters();
|
|
let mut key = sample_key(None);
|
|
key.status_snapshot = Some(json!({
|
|
"quota": {
|
|
"version": 2,
|
|
"provider_type": "antigravity",
|
|
"exhausted": false,
|
|
"windows": [
|
|
{
|
|
"code": "model:gemini-3.1-pro-high",
|
|
"scope": "model",
|
|
"model": "gemini-3.1-pro-high",
|
|
"used_ratio": 1.0,
|
|
"is_exhausted": true
|
|
},
|
|
{
|
|
"code": "model:gemini-3-flash-agent",
|
|
"scope": "model",
|
|
"model": "gemini-3-flash-agent",
|
|
"used_ratio": 0.1,
|
|
"is_exhausted": false
|
|
}
|
|
]
|
|
}
|
|
}));
|
|
|
|
let exhausted =
|
|
service.member_signals("antigravity", &key, None, Some("gemini-3.1-pro-high"));
|
|
let available =
|
|
service.member_signals("antigravity", &key, None, Some("gemini-3-flash-agent"));
|
|
|
|
assert!(exhausted.quota_exhausted);
|
|
assert!(!available.quota_exhausted);
|
|
}
|
|
|
|
#[test]
|
|
fn antigravity_tiered_model_quota_wins_over_display_only_group_windows() {
|
|
let service = ProviderPoolService::with_builtin_adapters();
|
|
let mut key = sample_key(None);
|
|
key.status_snapshot = Some(json!({
|
|
"quota": {
|
|
"version": 2,
|
|
"provider_type": "antigravity",
|
|
"exhausted": true,
|
|
"windows": [{
|
|
"code": "model:gemini-3.7-flash-tiered",
|
|
"scope": "model",
|
|
"model": "gemini-3.7-flash-tiered",
|
|
"remaining_ratio": 0.906,
|
|
"used_ratio": 0.094,
|
|
"is_exhausted": false
|
|
}, {
|
|
"code": "group:0:3p-5h",
|
|
"scope": "quota_group",
|
|
"quota_group": "group:0",
|
|
"bucket_id": "3p-5h",
|
|
"used_ratio": 1.0,
|
|
"is_exhausted": true
|
|
}]
|
|
}
|
|
}));
|
|
|
|
let signals =
|
|
service.member_signals("antigravity", &key, None, Some("gemini-3.7-flash-tiered"));
|
|
|
|
assert!(!signals.quota_exhausted);
|
|
assert!(!signals.quota_hard_blocked);
|
|
}
|
|
|
|
#[test]
|
|
fn antigravity_tiered_variant_does_not_match_another_model_family() {
|
|
let service = ProviderPoolService::with_builtin_adapters();
|
|
let mut key = sample_key(None);
|
|
key.status_snapshot = Some(json!({
|
|
"quota": {
|
|
"version": 2,
|
|
"provider_type": "antigravity",
|
|
"exhausted": false,
|
|
"windows": [{
|
|
"code": "model:gemini-3.7-pro-tiered",
|
|
"scope": "model",
|
|
"model": "gemini-3.7-pro-tiered",
|
|
"used_ratio": 1.0,
|
|
"is_exhausted": true
|
|
}]
|
|
}
|
|
}));
|
|
|
|
let signals =
|
|
service.member_signals("antigravity", &key, None, Some("gemini-3.7-flash-tiered"));
|
|
|
|
assert!(!signals.quota_exhausted);
|
|
assert!(!signals.quota_hard_blocked);
|
|
}
|
|
|
|
#[test]
|
|
fn codex_standard_and_spark_quota_families_are_independent() {
|
|
let service = ProviderPoolService::with_builtin_adapters();
|
|
let mut standard_exhausted = sample_key(None);
|
|
standard_exhausted.status_snapshot = Some(json!({
|
|
"quota": {
|
|
"version": 2,
|
|
"provider_type": "codex",
|
|
"exhausted": true,
|
|
"allowed": false,
|
|
"limit_reached": true,
|
|
"windows": [
|
|
{ "code": "weekly", "used_ratio": 1.0, "is_exhausted": true },
|
|
{ "code": "5h", "used_ratio": 0.5, "is_exhausted": false },
|
|
{ "code": "spark_weekly", "used_ratio": 0.2, "is_exhausted": false },
|
|
{ "code": "spark_5h", "used_ratio": 0.1, "is_exhausted": false }
|
|
]
|
|
}
|
|
}));
|
|
|
|
let standard =
|
|
service.member_signals("codex", &standard_exhausted, None, Some("gpt-5.3-codex"));
|
|
let spark = service.member_signals(
|
|
"codex",
|
|
&standard_exhausted,
|
|
None,
|
|
Some("gpt-5.3-codex-spark"),
|
|
);
|
|
assert!(standard.quota_exhausted);
|
|
assert!(!standard.quota_hard_blocked);
|
|
assert!(!spark.quota_exhausted);
|
|
assert!(!spark.quota_hard_blocked);
|
|
|
|
let mut spark_exhausted = sample_key(None);
|
|
spark_exhausted.status_snapshot = Some(json!({
|
|
"quota": {
|
|
"version": 2,
|
|
"provider_type": "codex",
|
|
"exhausted": false,
|
|
"windows": [
|
|
{ "code": "weekly", "used_ratio": 0.2, "is_exhausted": false },
|
|
{ "code": "5h", "used_ratio": 0.1, "is_exhausted": false },
|
|
{ "code": "spark_weekly", "used_ratio": 1.0, "is_exhausted": true },
|
|
{ "code": "spark_5h", "used_ratio": 0.4, "is_exhausted": false }
|
|
]
|
|
}
|
|
}));
|
|
|
|
let standard =
|
|
service.member_signals("codex", &spark_exhausted, None, Some("gpt-5.3-codex"));
|
|
let spark =
|
|
service.member_signals("codex", &spark_exhausted, None, Some("gpt-5.3-codex-spark"));
|
|
assert!(!standard.quota_exhausted);
|
|
assert!(spark.quota_exhausted);
|
|
}
|
|
|
|
#[test]
|
|
fn provider_quota_exhaustion_metadata_expires_after_reset_at() {
|
|
let now = std::time::SystemTime::now()
|
|
.duration_since(std::time::UNIX_EPOCH)
|
|
.expect("system time should be after unix epoch")
|
|
.as_secs();
|
|
|
|
assert!(!provider_pool_key_account_quota_exhausted(
|
|
&sample_key(Some(json!({
|
|
"codex": {
|
|
"updated_at": now.saturating_sub(600),
|
|
"allowed": false,
|
|
"limit_reached": true,
|
|
"primary_used_percent": 100.0,
|
|
"primary_reset_at": now.saturating_sub(60)
|
|
}
|
|
}))),
|
|
"codex",
|
|
));
|
|
assert!(provider_pool_key_account_quota_exhausted(
|
|
&sample_key(Some(json!({
|
|
"codex": {
|
|
"updated_at": now,
|
|
"primary_used_percent": 100.0,
|
|
"primary_reset_at": now.saturating_add(3600)
|
|
}
|
|
}))),
|
|
"codex",
|
|
));
|
|
}
|
|
|
|
#[test]
|
|
fn model_quota_windows_are_isolated_without_provider_specific_names() {
|
|
let service = ProviderPoolService::with_builtin_adapters();
|
|
let mut key = sample_key(None);
|
|
key.status_snapshot = Some(json!({
|
|
"quota": {
|
|
"version": 2,
|
|
"provider_type": "codex",
|
|
"exhausted": true,
|
|
"windows": [
|
|
{
|
|
"code": "alpha_short",
|
|
"quota_group": "alpha",
|
|
"model": "vendor-alpha-model",
|
|
"used_ratio": 1.0,
|
|
"is_exhausted": true
|
|
},
|
|
{
|
|
"code": "alpha_long",
|
|
"quota_group": "alpha",
|
|
"model": "vendor-alpha-model",
|
|
"used_ratio": 0.2,
|
|
"is_exhausted": false
|
|
},
|
|
{
|
|
"code": "beta_short",
|
|
"quota_group": "beta",
|
|
"model": "vendor-beta-model",
|
|
"used_ratio": 1.0,
|
|
"is_exhausted": true
|
|
}
|
|
]
|
|
}
|
|
}));
|
|
|
|
let alpha = service.member_signals("codex", &key, None, Some("vendor-alpha-model"));
|
|
let beta = service.member_signals("codex", &key, None, Some("vendor-beta-model"));
|
|
assert!(
|
|
!alpha.quota_exhausted,
|
|
"one available alpha window must keep it usable"
|
|
);
|
|
assert!(beta.quota_exhausted);
|
|
assert!(!beta.quota_hard_blocked);
|
|
}
|
|
|
|
#[test]
|
|
fn legacy_family_prefix_is_matched_by_model_tokens() {
|
|
let service = ProviderPoolService::with_builtin_adapters();
|
|
let mut key = sample_key(None);
|
|
key.status_snapshot = Some(json!({
|
|
"quota": {
|
|
"version": 2,
|
|
"provider_type": "codex",
|
|
"exhausted": true,
|
|
"windows": [
|
|
{ "code": "alpha_weekly", "used_ratio": 1.0, "is_exhausted": true },
|
|
{ "code": "default_weekly", "used_ratio": 0.1, "is_exhausted": false }
|
|
]
|
|
}
|
|
}));
|
|
|
|
let alpha = service.member_signals("codex", &key, None, Some("vendor-alpha-v2"));
|
|
assert!(alpha.quota_exhausted);
|
|
// An unrelated model has no identifiable bucket and therefore falls
|
|
// back to the account-level snapshot rather than guessing a family.
|
|
let unknown = service.member_signals("codex", &key, None, Some("vendor-gamma-v2"));
|
|
assert!(unknown.quota_exhausted);
|
|
}
|
|
|
|
#[test]
|
|
fn compact_model_bucket_identity_matches_request_tokens() {
|
|
let service = ProviderPoolService::with_builtin_adapters();
|
|
let mut key = sample_key(None);
|
|
key.status_snapshot = Some(json!({
|
|
"quota": {
|
|
"version": 2,
|
|
"provider_type": "codex",
|
|
"exhausted": true,
|
|
"windows": [
|
|
{
|
|
"code": "additional_0_primary_window",
|
|
"scope": "model",
|
|
"model": "spark",
|
|
"used_ratio": 0.0,
|
|
"is_exhausted": false
|
|
}
|
|
]
|
|
}
|
|
}));
|
|
|
|
let signals = service.member_signals("codex", &key, None, Some("gpt-5.3-codex-spark"));
|
|
assert!(
|
|
!signals.quota_exhausted,
|
|
"a compact bucket name should match a token in the selected model"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn model_family_token_matches_versioned_alias_without_hardcoding_name() {
|
|
let service = ProviderPoolService::with_builtin_adapters();
|
|
let mut key = sample_key(None);
|
|
key.status_snapshot = Some(json!({
|
|
"quota": {
|
|
"version": 2,
|
|
"provider_type": "codex",
|
|
"exhausted": true,
|
|
"windows": [{
|
|
"code": "additional_0_primary_window",
|
|
"scope": "model",
|
|
"model": "gpt-5.3-codex-spark",
|
|
"used_ratio": 1.0,
|
|
"is_exhausted": true
|
|
}]
|
|
}
|
|
}));
|
|
|
|
let signals = service.member_signals("codex", &key, None, Some("gpt-5.4-codex-spark"));
|
|
assert!(signals.quota_exhausted);
|
|
}
|
|
|
|
#[test]
|
|
fn model_only_snapshot_does_not_poison_account_fallback() {
|
|
let service = ProviderPoolService::with_builtin_adapters();
|
|
let mut key = sample_key(None);
|
|
key.status_snapshot = Some(json!({
|
|
"quota": {
|
|
"version": 2,
|
|
"provider_type": "codex",
|
|
"exhausted": true,
|
|
"windows": [{
|
|
"code": "model:vendor-alpha",
|
|
"scope": "model",
|
|
"model": "vendor-alpha",
|
|
"used_ratio": 1.0,
|
|
"is_exhausted": true
|
|
}]
|
|
}
|
|
}));
|
|
|
|
let signals = service.member_signals("codex", &key, None, None);
|
|
assert!(
|
|
!signals.quota_exhausted,
|
|
"account-level inspection must ignore model-only buckets"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn model_map_quota_is_resolved_without_materialized_windows() {
|
|
let service = ProviderPoolService::with_builtin_adapters();
|
|
let mut key = sample_key(None);
|
|
key.status_snapshot = Some(json!({
|
|
"quota": {
|
|
"version": 2,
|
|
"provider_type": "grok",
|
|
"exhausted": false,
|
|
"quota_by_model": {
|
|
"vendor-alpha": {
|
|
"remaining": 0.0,
|
|
"total": 10.0
|
|
},
|
|
"vendor-beta": {
|
|
"remaining": 5.0,
|
|
"total": 10.0
|
|
}
|
|
}
|
|
}
|
|
}));
|
|
|
|
let alpha = service.member_signals("grok", &key, None, Some("vendor-alpha"));
|
|
let beta = service.member_signals("grok", &key, None, Some("vendor-beta"));
|
|
assert!(alpha.quota_exhausted);
|
|
assert!(!beta.quota_exhausted);
|
|
}
|
|
|
|
#[test]
|
|
fn raw_provider_metadata_model_bucket_overrides_account_snapshot() {
|
|
let service = ProviderPoolService::with_builtin_adapters();
|
|
let mut key = sample_key(Some(json!({
|
|
"codex": {
|
|
"updated_at": 1_700_000_001u64,
|
|
"additional_quota_windows": [{
|
|
"scope": "model",
|
|
"model": "future-spark",
|
|
"used_ratio": 0.0,
|
|
"is_exhausted": false
|
|
}]
|
|
}
|
|
})));
|
|
key.status_snapshot = Some(json!({
|
|
"quota": {
|
|
"version": 2,
|
|
"provider_type": "codex",
|
|
"exhausted": true,
|
|
"windows": [{
|
|
"code": "primary",
|
|
"scope": "account",
|
|
"used_ratio": 1.0,
|
|
"is_exhausted": true
|
|
}]
|
|
}
|
|
}));
|
|
|
|
let signals = service.member_signals("codex", &key, None, Some("future-spark"));
|
|
assert!(
|
|
!signals.quota_exhausted,
|
|
"a newer model bucket in provider metadata must not inherit an exhausted account bucket"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn model_quota_availability_suppresses_account_hard_block() {
|
|
let service = ProviderPoolService::with_builtin_adapters();
|
|
let mut key = sample_key(None);
|
|
key.status_snapshot = Some(json!({
|
|
"quota": {
|
|
"version": 2,
|
|
"provider_type": "codex",
|
|
"allowed": false,
|
|
"limit_reached": true,
|
|
"exhausted": true,
|
|
"windows": [{
|
|
"scope": "model",
|
|
"model": "future-model",
|
|
"used_ratio": 0.0,
|
|
"is_exhausted": false
|
|
}]
|
|
}
|
|
}));
|
|
|
|
let signals = service.member_signals("codex", &key, None, Some("future-model"));
|
|
assert!(!signals.quota_exhausted);
|
|
assert!(!signals.quota_hard_blocked);
|
|
}
|
|
|
|
#[test]
|
|
fn newer_model_quota_observation_wins_over_stale_snapshot() {
|
|
let service = ProviderPoolService::with_builtin_adapters();
|
|
let mut key = sample_key(Some(json!({
|
|
"codex": {
|
|
"updated_at": 1_700_000_200u64,
|
|
"additional_quota_windows": [{
|
|
"scope": "model",
|
|
"model": "future-model",
|
|
"used_ratio": 0.0,
|
|
"is_exhausted": false
|
|
}]
|
|
}
|
|
})));
|
|
key.status_snapshot = Some(json!({
|
|
"quota": {
|
|
"version": 2,
|
|
"provider_type": "codex",
|
|
"observed_at": 1_700_000_100u64,
|
|
"exhausted": true,
|
|
"windows": [{
|
|
"scope": "model",
|
|
"model": "future-model",
|
|
"used_ratio": 1.0,
|
|
"is_exhausted": true
|
|
}]
|
|
}
|
|
}));
|
|
|
|
let signals = service.member_signals("codex", &key, None, Some("future-model"));
|
|
assert!(!signals.quota_exhausted);
|
|
}
|
|
|
|
#[test]
|
|
fn codex_newer_flat_quota_metadata_clears_stale_exhausted_snapshot() {
|
|
let now = std::time::SystemTime::now()
|
|
.duration_since(std::time::UNIX_EPOCH)
|
|
.expect("system time should be after unix epoch")
|
|
.as_secs();
|
|
let service = ProviderPoolService::with_builtin_adapters();
|
|
// Headers and quota refreshes persist the flat metadata shape while the
|
|
// derived snapshot may still describe the previous quota observation.
|
|
for used_percent in [83.0, 99.0, 100.0] {
|
|
let mut key = sample_key(Some(json!({
|
|
"codex": {
|
|
"updated_at": now,
|
|
"primary_used_percent": used_percent,
|
|
"primary_reset_at": now + 3600
|
|
}
|
|
})));
|
|
key.status_snapshot = Some(json!({
|
|
"quota": {
|
|
"provider_type": "codex",
|
|
"observed_at": now - 60,
|
|
"updated_at": now - 60,
|
|
"code": "exhausted",
|
|
"exhausted": true,
|
|
"allowed": false,
|
|
"limit_reached": true,
|
|
"reset_at": now + 3600,
|
|
"windows": [{
|
|
"code": "weekly",
|
|
"scope": "account",
|
|
"used_ratio": 1.0,
|
|
"remaining_ratio": 0.0,
|
|
"is_exhausted": true,
|
|
"reset_at": now + 3600
|
|
}]
|
|
}
|
|
}));
|
|
for model in [None, Some("gpt-5.4"), Some("o3")] {
|
|
let signals = service.member_signals("codex", &key, None, model);
|
|
assert_eq!(signals.quota_exhausted, used_percent >= 100.0);
|
|
assert!(!signals.quota_hard_blocked);
|
|
assert_eq!(
|
|
provider_pool_key_minimum_quota_reached(&key, "codex", model),
|
|
used_percent >= 99.0
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn codex_account_and_model_quota_agree_on_explicit_allow_and_deny_flags() {
|
|
let service = ProviderPoolService::with_builtin_adapters();
|
|
for (flags, exhausted) in [
|
|
(json!({ "allowed": true }), false),
|
|
(json!({ "limit_reached": false }), false),
|
|
(json!({ "allowed": true, "limit_reached": true }), true),
|
|
(json!({ "allowed": false, "limit_reached": false }), true),
|
|
] {
|
|
for use_windows in [false, true] {
|
|
let mut metadata = flags.clone();
|
|
metadata["updated_at"] = json!(300);
|
|
if use_windows {
|
|
metadata["windows"] = json!([{ "code": "weekly", "used_ratio": 1.0 }]);
|
|
} else {
|
|
metadata["primary_used_percent"] = json!(100.0);
|
|
}
|
|
let key = sample_key(Some(json!({ "codex": metadata })));
|
|
for model in [None, Some("gpt-5.4"), Some("o3")] {
|
|
assert_eq!(
|
|
service
|
|
.member_signals("codex", &key, None, model)
|
|
.quota_exhausted,
|
|
exhausted,
|
|
"flags={flags}, use_windows={use_windows}, model={model:?}"
|
|
);
|
|
// The opt-in reserve still protects a numerically full
|
|
// window even when the upstream reports it as allowed.
|
|
assert!(provider_pool_key_minimum_quota_reached(
|
|
&key, "codex", model
|
|
));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn codex_model_quota_honors_latest_account_refusal_without_blocking_spark() {
|
|
let service = ProviderPoolService::with_builtin_adapters();
|
|
for include_window in [false, true] {
|
|
let mut metadata = json!({
|
|
"updated_at": 300,
|
|
"allowed": false,
|
|
"limit_reached": true,
|
|
"spark_primary_used_percent": 17.0
|
|
});
|
|
if include_window {
|
|
metadata["primary_used_percent"] = json!(83.0);
|
|
}
|
|
let mut key = sample_key(Some(json!({ "codex": metadata })));
|
|
for include_snapshot in [false, true] {
|
|
if include_snapshot {
|
|
key.status_snapshot = Some(json!({
|
|
"quota": {
|
|
"provider_type": "codex",
|
|
"observed_at": 200,
|
|
"exhausted": false,
|
|
"windows": [{ "code": "weekly", "used_ratio": 0.83 }]
|
|
}
|
|
}));
|
|
}
|
|
assert!(
|
|
service
|
|
.member_signals("codex", &key, None, Some("gpt-5.4"))
|
|
.quota_exhausted
|
|
);
|
|
assert!(
|
|
!service
|
|
.member_signals("codex", &key, None, Some("gpt-5.3-codex-spark"))
|
|
.quota_exhausted
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn codex_quota_source_freshness_supports_all_timestamp_formats() {
|
|
let service = ProviderPoolService::with_builtin_adapters();
|
|
for observed_at in [
|
|
json!(1_700_000_200_u64),
|
|
json!(1_700_000_200_000_u64),
|
|
json!("1700000200000"),
|
|
json!("2023-11-14T22:16:40Z"),
|
|
] {
|
|
for use_windows in [false, true] {
|
|
let metadata = if use_windows {
|
|
json!({
|
|
"updated_at": observed_at,
|
|
"windows": [{ "code": "weekly", "used_ratio": 0.83 }]
|
|
})
|
|
} else {
|
|
json!({ "updated_at": observed_at, "primary_used_percent": 83.0 })
|
|
};
|
|
let mut key = sample_key(Some(json!({ "codex": metadata })));
|
|
key.status_snapshot = Some(json!({
|
|
"quota": {
|
|
"provider_type": "codex",
|
|
"observed_at": "2023-11-14T22:15:00Z",
|
|
"exhausted": true,
|
|
"allowed": false,
|
|
"windows": [{ "code": "weekly", "used_ratio": 1.0 }]
|
|
}
|
|
}));
|
|
for model in [None, Some("gpt-5.4")] {
|
|
let signals = service.member_signals("codex", &key, None, model);
|
|
assert!(!signals.quota_exhausted, "timestamp={observed_at}");
|
|
assert!(!signals.quota_hard_blocked, "timestamp={observed_at}");
|
|
assert!(!provider_pool_key_minimum_quota_reached(
|
|
&key, "codex", model
|
|
));
|
|
}
|
|
|
|
// Swap freshness while retaining the same observations. An
|
|
// older usable bucket cannot erase a later exhausted snapshot.
|
|
key.status_snapshot.as_mut().unwrap()["quota"]["observed_at"] =
|
|
json!("2023-11-14T22:18:20Z");
|
|
assert!(provider_pool_key_account_quota_exhausted(&key, "codex"));
|
|
assert!(provider_pool_key_quota_hard_blocked(&key, "codex"));
|
|
assert_eq!(
|
|
provider_pool_key_model_quota_exhausted(&key, "codex", "gpt-5.4"),
|
|
Some(true)
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn codex_account_quota_recovery_requires_new_account_observation() {
|
|
for metadata in [
|
|
json!({ "updated_at": 100, "primary_used_percent": 83.0 }),
|
|
json!({ "primary_used_percent": 83.0 }),
|
|
json!({ "updated_at": 300, "plan_type": "plus" }),
|
|
json!({ "updated_at": 300, "credits_unlimited": false }),
|
|
json!({ "updated_at": 300, "spark_primary_used_percent": 83.0 }),
|
|
json!({ "updated_at": 300, "windows": [{ "code": "weekly" }] }),
|
|
] {
|
|
let mut key = sample_key(Some(json!({ "codex": metadata })));
|
|
key.status_snapshot = Some(json!({
|
|
"quota": {
|
|
"provider_type": "codex",
|
|
"observed_at": 200,
|
|
"exhausted": true,
|
|
"allowed": false,
|
|
"limit_reached": true,
|
|
"windows": [{ "code": "weekly", "used_ratio": 1.0 }]
|
|
}
|
|
}));
|
|
assert!(provider_pool_key_account_quota_exhausted(&key, "codex"));
|
|
assert!(provider_pool_key_quota_hard_blocked(&key, "codex"));
|
|
assert_eq!(
|
|
provider_pool_key_model_quota_exhausted(&key, "codex", "gpt-5.4"),
|
|
Some(true)
|
|
);
|
|
}
|
|
|
|
let mut key = sample_key(Some(json!({
|
|
"codex": {
|
|
"updated_at": 300,
|
|
"primary_used_percent": 83.0,
|
|
"allowed": false,
|
|
"limit_reached": true
|
|
}
|
|
})));
|
|
key.status_snapshot = Some(json!({
|
|
"quota": { "provider_type": "codex", "updated_at": 200, "exhausted": false }
|
|
}));
|
|
assert!(provider_pool_key_account_quota_exhausted(&key, "codex"));
|
|
assert!(provider_pool_key_quota_hard_blocked(&key, "codex"));
|
|
}
|
|
|
|
#[test]
|
|
fn codex_account_updates_do_not_override_independent_model_quota() {
|
|
for (spark_ratio, account_percent) in [(0.83, 100.0), (1.0, 83.0)] {
|
|
let mut key = sample_key(Some(json!({
|
|
"codex": { "updated_at": 300, "primary_used_percent": account_percent }
|
|
})));
|
|
key.status_snapshot = Some(json!({
|
|
"quota": {
|
|
"provider_type": "codex",
|
|
"updated_at": 200,
|
|
"windows": [{ "code": "spark_5h", "used_ratio": spark_ratio }]
|
|
}
|
|
}));
|
|
assert_eq!(
|
|
provider_pool_key_model_quota_exhausted(&key, "codex", "gpt-5.3-codex-spark"),
|
|
Some(spark_ratio >= 1.0)
|
|
);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn codex_explicit_quota_block_is_hard_until_reset() {
|
|
let now = std::time::SystemTime::now()
|
|
.duration_since(std::time::UNIX_EPOCH)
|
|
.expect("system time should be after unix epoch")
|
|
.as_secs();
|
|
|
|
assert!(provider_pool_key_quota_hard_blocked(
|
|
&sample_key(Some(json!({
|
|
"codex": {
|
|
"updated_at": now,
|
|
"allowed": false,
|
|
"limit_reached": true,
|
|
"primary_reset_at": now.saturating_add(3600)
|
|
}
|
|
}))),
|
|
"codex",
|
|
));
|
|
assert!(!provider_pool_key_quota_hard_blocked(
|
|
&sample_key(Some(json!({
|
|
"codex": {
|
|
"updated_at": now,
|
|
"primary_used_percent": 100.0,
|
|
"primary_reset_at": now.saturating_add(3600)
|
|
}
|
|
}))),
|
|
"codex",
|
|
));
|
|
assert!(!provider_pool_key_quota_hard_blocked(
|
|
&sample_key(Some(json!({
|
|
"codex": {
|
|
"updated_at": now.saturating_sub(600),
|
|
"allowed": false,
|
|
"limit_reached": true,
|
|
"primary_reset_at": now.saturating_sub(60)
|
|
}
|
|
}))),
|
|
"codex",
|
|
));
|
|
|
|
let mut snapshot_blocked = sample_key(None);
|
|
snapshot_blocked.status_snapshot = Some(json!({
|
|
"quota": {
|
|
"version": 2,
|
|
"provider_type": "codex",
|
|
"exhausted": true,
|
|
"allowed": false,
|
|
"limit_reached": true,
|
|
"usage_ratio": 0.91,
|
|
"reset_at": now.saturating_add(3600)
|
|
}
|
|
}));
|
|
assert!(provider_pool_key_quota_hard_blocked(
|
|
&snapshot_blocked,
|
|
"codex",
|
|
));
|
|
snapshot_blocked.status_snapshot = Some(json!({
|
|
"quota": {
|
|
"version": 2,
|
|
"provider_type": "codex",
|
|
"exhausted": true,
|
|
"allowed": false,
|
|
"limit_reached": true,
|
|
"usage_ratio": 0.91,
|
|
"reset_at": now.saturating_sub(60)
|
|
}
|
|
}));
|
|
assert!(!provider_pool_key_quota_hard_blocked(
|
|
&snapshot_blocked,
|
|
"codex",
|
|
));
|
|
}
|
|
|
|
#[test]
|
|
fn grok_quota_tier_boundaries_match_pool_modes() {
|
|
assert_eq!(
|
|
grok_supported_quota_windows_for_tier(Some("basic")),
|
|
[("quota_fast", "fast")]
|
|
);
|
|
assert_eq!(
|
|
grok_supported_quota_windows_for_tier(Some("super")),
|
|
[
|
|
("quota_auto", "auto"),
|
|
("quota_fast", "fast"),
|
|
("quota_expert", "expert"),
|
|
("quota_grok_4_3", "grok-420-computer-use-sa")
|
|
]
|
|
);
|
|
assert_eq!(
|
|
grok_supported_quota_windows_for_tier(Some("heavy")),
|
|
[
|
|
("quota_auto", "auto"),
|
|
("quota_fast", "fast"),
|
|
("quota_expert", "expert"),
|
|
("quota_heavy", "heavy"),
|
|
("quota_grok_4_3", "grok-420-computer-use-sa")
|
|
]
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn grok_pool_tier_infers_from_live_quota_totals() {
|
|
let bucket = json!({
|
|
"quota_by_model": {
|
|
"quota_fast": {
|
|
"remaining": 20.0,
|
|
"total": 30.0
|
|
},
|
|
"quota_auto": {
|
|
"remaining": 7.0,
|
|
"total": 7.0
|
|
}
|
|
}
|
|
});
|
|
let bucket = bucket.as_object().expect("bucket should be object");
|
|
|
|
assert_eq!(grok_pool_tier_from_quota_bucket(bucket), Some("basic"));
|
|
}
|
|
|
|
#[test]
|
|
fn grok_model_name_maps_to_quota_window() {
|
|
assert_eq!(
|
|
grok_quota_window_key_for_model(Some("grok-4.20-fast")),
|
|
Some("quota_fast")
|
|
);
|
|
assert_eq!(
|
|
grok_quota_window_key_for_model(Some("grok-4.20-multi-agent-0309")),
|
|
Some("quota_heavy")
|
|
);
|
|
assert_eq!(
|
|
grok_quota_window_key_for_model(Some("grok-4.3-beta")),
|
|
Some("quota_grok_4_3")
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn plan_tier_derivation_normalizes_provider_prefix() {
|
|
let key = sample_key(Some(json!({
|
|
"codex": {
|
|
"plan_type": "codex:Plus"
|
|
}
|
|
})));
|
|
|
|
assert_eq!(
|
|
derive_oauth_plan_type("codex", &key, None).as_deref(),
|
|
Some("plus")
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn plan_tier_derivation_reads_quota_snapshot() {
|
|
let mut key = sample_key(None);
|
|
key.status_snapshot = Some(json!({
|
|
"quota": {
|
|
"plan_type": "team"
|
|
}
|
|
}));
|
|
|
|
assert_eq!(
|
|
derive_oauth_plan_type("codex", &key, None).as_deref(),
|
|
Some("team")
|
|
);
|
|
}
|
|
}
|