mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-04 16:37:46 +08:00
refactor(workspace): enforce layered crate boundaries
This commit is contained in:
@@ -0,0 +1,77 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogEndpoint;
|
||||
use serde_json::json;
|
||||
|
||||
use crate::capability::ProviderPoolCapabilities;
|
||||
use crate::provider::{
|
||||
provider_pool_endpoint_format_matches, provider_pool_matching_endpoint, ProviderPoolAdapter,
|
||||
};
|
||||
use crate::quota_refresh::ProviderPoolQuotaRequestSpec;
|
||||
|
||||
pub const ANTIGRAVITY_FETCH_AVAILABLE_MODELS_PATH: &str = "/v1internal:fetchAvailableModels";
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct AntigravityProviderPoolAdapter;
|
||||
|
||||
impl ProviderPoolAdapter for AntigravityProviderPoolAdapter {
|
||||
fn provider_type(&self) -> &'static str {
|
||||
"antigravity"
|
||||
}
|
||||
|
||||
fn capabilities(&self) -> ProviderPoolCapabilities {
|
||||
ProviderPoolCapabilities {
|
||||
quota_refresh: true,
|
||||
..ProviderPoolCapabilities::default()
|
||||
}
|
||||
}
|
||||
|
||||
fn quota_refresh_endpoint(
|
||||
&self,
|
||||
endpoints: &[StoredProviderCatalogEndpoint],
|
||||
include_inactive: bool,
|
||||
) -> Option<StoredProviderCatalogEndpoint> {
|
||||
provider_pool_matching_endpoint(endpoints, include_inactive, |endpoint| {
|
||||
provider_pool_endpoint_format_matches(endpoint, "gemini:generate_content")
|
||||
})
|
||||
}
|
||||
|
||||
fn quota_refresh_missing_endpoint_message(&self) -> String {
|
||||
"找不到有效的 gemini:generate_content 端点".to_string()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn build_antigravity_pool_quota_request(
|
||||
key_id: &str,
|
||||
endpoint_base_url: &str,
|
||||
authorization: (String, String),
|
||||
project_id: &str,
|
||||
mut identity_headers: BTreeMap<String, String>,
|
||||
) -> ProviderPoolQuotaRequestSpec {
|
||||
let mut headers = std::mem::take(&mut identity_headers);
|
||||
headers.insert("authorization".to_string(), authorization.1);
|
||||
headers.insert("content-type".to_string(), "application/json".to_string());
|
||||
headers.insert("accept".to_string(), "application/json".to_string());
|
||||
headers
|
||||
.entry("user-agent".to_string())
|
||||
.or_insert_with(|| "antigravity".to_string());
|
||||
|
||||
ProviderPoolQuotaRequestSpec {
|
||||
request_id: format!("antigravity-quota:{key_id}"),
|
||||
provider_name: "antigravity".to_string(),
|
||||
quota_kind: "antigravity".to_string(),
|
||||
method: "POST".to_string(),
|
||||
url: format!(
|
||||
"{}{}",
|
||||
endpoint_base_url.trim_end_matches('/'),
|
||||
ANTIGRAVITY_FETCH_AVAILABLE_MODELS_PATH
|
||||
),
|
||||
headers,
|
||||
content_type: Some("application/json".to_string()),
|
||||
json_body: Some(json!({ "project": project_id })),
|
||||
client_api_format: "gemini:generate_content".to_string(),
|
||||
provider_api_format: "antigravity:fetch_available_models".to_string(),
|
||||
model_name: Some("fetchAvailableModels".to_string()),
|
||||
accept_invalid_certs: false,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,359 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogEndpoint;
|
||||
use serde_json::{json, Map, Value};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::capability::ProviderPoolCapabilities;
|
||||
use crate::provider::{
|
||||
provider_pool_endpoint_format_matches, provider_pool_matching_endpoint, ProviderPoolAdapter,
|
||||
ProviderPoolMemberInput,
|
||||
};
|
||||
use crate::quota::{
|
||||
provider_pool_current_unix_secs, provider_pool_json_bool, provider_pool_json_f64,
|
||||
provider_pool_metadata_bucket, provider_pool_quota_snapshot_exhausted_decision,
|
||||
provider_pool_reset_deadline_elapsed, provider_pool_timestamp_unix_secs,
|
||||
};
|
||||
use crate::quota_refresh::ProviderPoolQuotaRequestSpec;
|
||||
|
||||
pub const CHATGPT_WEB_DEFAULT_BASE_URL: &str = "https://chatgpt.com";
|
||||
pub const CHATGPT_WEB_CONVERSATION_INIT_PATH: &str = "/backend-api/conversation/init";
|
||||
|
||||
const CHATGPT_WEB_USER_AGENT: &str = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/143.0.0.0 Safari/537.36 Edg/143.0.0.0";
|
||||
const CHATGPT_WEB_CLIENT_VERSION: &str = "prod-be885abbfcfe7b1f511e88b3003d9ee44757fbad";
|
||||
const CHATGPT_WEB_BUILD_NUMBER: &str = "5955942";
|
||||
const CHATGPT_WEB_SEC_CH_UA: &str =
|
||||
r#""Microsoft Edge";v="143", "Chromium";v="143", "Not A(Brand";v="24""#;
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct ChatGptWebProviderPoolAdapter;
|
||||
|
||||
impl ProviderPoolAdapter for ChatGptWebProviderPoolAdapter {
|
||||
fn provider_type(&self) -> &'static str {
|
||||
"chatgpt_web"
|
||||
}
|
||||
|
||||
fn capabilities(&self) -> ProviderPoolCapabilities {
|
||||
ProviderPoolCapabilities {
|
||||
quota_refresh: true,
|
||||
..ProviderPoolCapabilities::default()
|
||||
}
|
||||
}
|
||||
|
||||
fn quota_exhausted(&self, input: &ProviderPoolMemberInput<'_>) -> bool {
|
||||
if let Some(exhausted) =
|
||||
provider_pool_quota_snapshot_exhausted_decision(input.key, input.provider_type)
|
||||
{
|
||||
return exhausted;
|
||||
}
|
||||
provider_pool_metadata_bucket(input.key.upstream_metadata.as_ref(), input.provider_type)
|
||||
.is_some_and(quota_exhausted_from_bucket)
|
||||
}
|
||||
|
||||
fn quota_refresh_endpoint(
|
||||
&self,
|
||||
endpoints: &[StoredProviderCatalogEndpoint],
|
||||
include_inactive: bool,
|
||||
) -> Option<StoredProviderCatalogEndpoint> {
|
||||
provider_pool_matching_endpoint(endpoints, include_inactive, |endpoint| {
|
||||
provider_pool_endpoint_format_matches(endpoint, "openai:image")
|
||||
})
|
||||
}
|
||||
|
||||
fn quota_refresh_missing_endpoint_message(&self) -> String {
|
||||
"找不到有效的 openai:image 端点".to_string()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn build_chatgpt_web_pool_quota_request(
|
||||
key_id: &str,
|
||||
endpoint_base_url: &str,
|
||||
authorization: (String, String),
|
||||
) -> ProviderPoolQuotaRequestSpec {
|
||||
let base_url = chatgpt_web_base_url(endpoint_base_url);
|
||||
let device_id = Uuid::new_v4().to_string();
|
||||
let session_id = Uuid::new_v4().to_string();
|
||||
let mut headers = BTreeMap::from([
|
||||
("accept".to_string(), "application/json".to_string()),
|
||||
("content-type".to_string(), "application/json".to_string()),
|
||||
("user-agent".to_string(), CHATGPT_WEB_USER_AGENT.to_string()),
|
||||
("origin".to_string(), base_url.clone()),
|
||||
("referer".to_string(), format!("{base_url}/")),
|
||||
(
|
||||
"accept-language".to_string(),
|
||||
"zh-CN,zh;q=0.9,en;q=0.8,en-US;q=0.7".to_string(),
|
||||
),
|
||||
("cache-control".to_string(), "no-cache".to_string()),
|
||||
("pragma".to_string(), "no-cache".to_string()),
|
||||
("priority".to_string(), "u=1, i".to_string()),
|
||||
("sec-ch-ua".to_string(), CHATGPT_WEB_SEC_CH_UA.to_string()),
|
||||
("sec-ch-ua-arch".to_string(), r#""x86""#.to_string()),
|
||||
("sec-ch-ua-bitness".to_string(), r#""64""#.to_string()),
|
||||
("sec-ch-ua-mobile".to_string(), "?0".to_string()),
|
||||
("sec-ch-ua-model".to_string(), r#""""#.to_string()),
|
||||
("sec-ch-ua-platform".to_string(), r#""Windows""#.to_string()),
|
||||
(
|
||||
"sec-ch-ua-platform-version".to_string(),
|
||||
r#""19.0.0""#.to_string(),
|
||||
),
|
||||
("sec-fetch-dest".to_string(), "empty".to_string()),
|
||||
("sec-fetch-mode".to_string(), "cors".to_string()),
|
||||
("sec-fetch-site".to_string(), "same-origin".to_string()),
|
||||
("oai-device-id".to_string(), device_id),
|
||||
("oai-session-id".to_string(), session_id),
|
||||
("oai-language".to_string(), "zh-CN".to_string()),
|
||||
(
|
||||
"oai-client-version".to_string(),
|
||||
CHATGPT_WEB_CLIENT_VERSION.to_string(),
|
||||
),
|
||||
(
|
||||
"oai-client-build-number".to_string(),
|
||||
CHATGPT_WEB_BUILD_NUMBER.to_string(),
|
||||
),
|
||||
(
|
||||
"x-openai-target-path".to_string(),
|
||||
CHATGPT_WEB_CONVERSATION_INIT_PATH.to_string(),
|
||||
),
|
||||
(
|
||||
"x-openai-target-route".to_string(),
|
||||
CHATGPT_WEB_CONVERSATION_INIT_PATH.to_string(),
|
||||
),
|
||||
]);
|
||||
headers.insert(authorization.0.to_ascii_lowercase(), authorization.1);
|
||||
|
||||
ProviderPoolQuotaRequestSpec {
|
||||
request_id: format!("chatgpt-web-quota:{key_id}"),
|
||||
provider_name: "chatgpt_web".to_string(),
|
||||
quota_kind: "chatgpt_web".to_string(),
|
||||
method: "POST".to_string(),
|
||||
url: format!("{base_url}{CHATGPT_WEB_CONVERSATION_INIT_PATH}"),
|
||||
headers,
|
||||
content_type: Some("application/json".to_string()),
|
||||
json_body: Some(json!({
|
||||
"gizmo_id": Value::Null,
|
||||
"requested_default_model": Value::Null,
|
||||
"conversation_id": Value::Null,
|
||||
"timezone_offset_min": -480,
|
||||
"system_hints": ["picture_v2"],
|
||||
})),
|
||||
client_api_format: "openai:image".to_string(),
|
||||
provider_api_format: "chatgpt_web:conversation_init".to_string(),
|
||||
model_name: Some("chatgpt-web-conversation-init".to_string()),
|
||||
accept_invalid_certs: true,
|
||||
}
|
||||
}
|
||||
|
||||
fn chatgpt_web_base_url(endpoint_base_url: &str) -> String {
|
||||
let base_url = endpoint_base_url.trim().trim_end_matches('/');
|
||||
if base_url.is_empty() {
|
||||
CHATGPT_WEB_DEFAULT_BASE_URL.to_string()
|
||||
} else {
|
||||
base_url.to_string()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn enrich_chatgpt_web_quota_metadata(metadata: &mut Value, auth_config: Option<&Value>) {
|
||||
let Some(object) = metadata.as_object_mut() else {
|
||||
return;
|
||||
};
|
||||
for (target, fields) in [
|
||||
("plan_type", &["plan_type", "tier", "plan"][..]),
|
||||
("email", &["email"][..]),
|
||||
("account_id", &["account_id", "accountId"][..]),
|
||||
("account_user_id", &["account_user_id", "accountUserId"][..]),
|
||||
("user_id", &["user_id", "userId"][..]),
|
||||
] {
|
||||
if object.contains_key(target) {
|
||||
continue;
|
||||
}
|
||||
if let Some(value) = chatgpt_web_auth_config_string(auth_config, fields) {
|
||||
object.insert(target.to_string(), json!(value));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn normalize_chatgpt_web_image_quota_limit(
|
||||
metadata: &mut Value,
|
||||
upstream_metadata: Option<&Value>,
|
||||
) {
|
||||
let existing_limit = existing_chatgpt_web_image_quota_limit(upstream_metadata);
|
||||
let Some(object) = metadata.as_object_mut() else {
|
||||
return;
|
||||
};
|
||||
|
||||
let remaining = provider_pool_json_f64(object.get("image_quota_remaining"));
|
||||
let plan_type = chatgpt_web_image_quota_plan_type(object)
|
||||
.map(ToOwned::to_owned)
|
||||
.or_else(|| {
|
||||
existing_limit
|
||||
.as_ref()
|
||||
.and_then(|existing| existing.plan_type.clone())
|
||||
});
|
||||
let raw_explicit_limit =
|
||||
provider_pool_json_f64(object.get("image_quota_total")).filter(|value| *value > 0.0);
|
||||
let explicit_limit_is_free_default = raw_explicit_limit.is_some_and(|limit| {
|
||||
is_legacy_chatgpt_web_free_default_limit_value(limit, None, plan_type.as_deref(), remaining)
|
||||
});
|
||||
if explicit_limit_is_free_default {
|
||||
object.remove("image_quota_total");
|
||||
object.remove("image_quota_limit_source");
|
||||
}
|
||||
let explicit_limit = raw_explicit_limit.filter(|_| !explicit_limit_is_free_default);
|
||||
let limit = explicit_limit
|
||||
.map(|limit| ChatGptWebImageQuotaLimit {
|
||||
value: limit,
|
||||
source: Some("upstream_total".to_string()),
|
||||
plan_type: plan_type.clone(),
|
||||
})
|
||||
.or_else(|| {
|
||||
infer_chatgpt_web_image_quota_limit(remaining, existing_limit, plan_type.as_deref())
|
||||
});
|
||||
|
||||
if let Some(limit) = limit {
|
||||
object.insert("image_quota_total".to_string(), json!(limit.value));
|
||||
if let Some(source) = limit.source.as_deref().filter(|value| !value.is_empty()) {
|
||||
object.insert("image_quota_limit_source".to_string(), json!(source));
|
||||
}
|
||||
|
||||
if !object.contains_key("image_quota_used") {
|
||||
if let Some(remaining) = remaining {
|
||||
object.insert(
|
||||
"image_quota_used".to_string(),
|
||||
json!((limit.value - remaining).max(0.0)),
|
||||
);
|
||||
} else if object.get("image_quota_blocked").and_then(Value::as_bool) == Some(true) {
|
||||
object.insert("image_quota_used".to_string(), json!(limit.value));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn chatgpt_web_auth_config_string(auth_config: Option<&Value>, fields: &[&str]) -> Option<String> {
|
||||
let object = auth_config.and_then(Value::as_object)?;
|
||||
fields.iter().find_map(|field| {
|
||||
object
|
||||
.get(*field)
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
})
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct ChatGptWebImageQuotaLimit {
|
||||
value: f64,
|
||||
source: Option<String>,
|
||||
plan_type: Option<String>,
|
||||
}
|
||||
|
||||
fn chatgpt_web_image_quota_plan_type(object: &Map<String, Value>) -> Option<&str> {
|
||||
object
|
||||
.get("plan_type")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
fn existing_chatgpt_web_image_quota_limit(
|
||||
upstream_metadata: Option<&Value>,
|
||||
) -> Option<ChatGptWebImageQuotaLimit> {
|
||||
let bucket = upstream_metadata
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|metadata| metadata.get("chatgpt_web"))
|
||||
.and_then(Value::as_object)?;
|
||||
let value =
|
||||
provider_pool_json_f64(bucket.get("image_quota_total")).filter(|value| *value > 0.0)?;
|
||||
let source = bucket
|
||||
.get("image_quota_limit_source")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned);
|
||||
let plan_type = chatgpt_web_image_quota_plan_type(bucket).map(ToOwned::to_owned);
|
||||
Some(ChatGptWebImageQuotaLimit {
|
||||
value,
|
||||
source,
|
||||
plan_type,
|
||||
})
|
||||
}
|
||||
|
||||
fn infer_chatgpt_web_image_quota_limit(
|
||||
remaining: Option<f64>,
|
||||
existing_limit: Option<ChatGptWebImageQuotaLimit>,
|
||||
plan_type: Option<&str>,
|
||||
) -> Option<ChatGptWebImageQuotaLimit> {
|
||||
if let Some(existing_limit) = existing_limit {
|
||||
if !is_legacy_chatgpt_web_free_default_limit(&existing_limit, plan_type, remaining) {
|
||||
return Some(existing_limit);
|
||||
}
|
||||
}
|
||||
|
||||
remaining
|
||||
.filter(|value| *value > 0.0)
|
||||
.map(|value| ChatGptWebImageQuotaLimit {
|
||||
value,
|
||||
source: Some("first_remaining".to_string()),
|
||||
plan_type: plan_type.map(ToOwned::to_owned),
|
||||
})
|
||||
}
|
||||
|
||||
fn is_legacy_chatgpt_web_free_default_limit(
|
||||
existing_limit: &ChatGptWebImageQuotaLimit,
|
||||
plan_type: Option<&str>,
|
||||
remaining: Option<f64>,
|
||||
) -> bool {
|
||||
is_legacy_chatgpt_web_free_default_limit_value(
|
||||
existing_limit.value,
|
||||
existing_limit.source.as_deref(),
|
||||
plan_type,
|
||||
remaining,
|
||||
)
|
||||
}
|
||||
|
||||
fn is_legacy_chatgpt_web_free_default_limit_value(
|
||||
value: f64,
|
||||
source: Option<&str>,
|
||||
plan_type: Option<&str>,
|
||||
remaining: Option<f64>,
|
||||
) -> bool {
|
||||
let plan_type_is_free = plan_type
|
||||
.map(str::trim)
|
||||
.is_some_and(|value| value.eq_ignore_ascii_case("free"));
|
||||
if !plan_type_is_free || source.is_some() {
|
||||
return false;
|
||||
}
|
||||
if (value - 25.0).abs() > f64::EPSILON {
|
||||
return false;
|
||||
}
|
||||
remaining.is_none_or(|remaining| remaining < value)
|
||||
}
|
||||
|
||||
pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
|
||||
if provider_pool_current_unix_secs().is_some_and(|now| {
|
||||
let mut image_quota = Map::new();
|
||||
if let Some(value) = bucket.get("image_quota_reset_at") {
|
||||
image_quota.insert("reset_at".to_string(), value.clone());
|
||||
}
|
||||
provider_pool_reset_deadline_elapsed(
|
||||
&image_quota,
|
||||
provider_pool_timestamp_unix_secs(bucket.get("updated_at")),
|
||||
now,
|
||||
)
|
||||
}) {
|
||||
return false;
|
||||
}
|
||||
if provider_pool_json_bool(bucket.get("image_quota_blocked")) == Some(true) {
|
||||
return true;
|
||||
}
|
||||
if provider_pool_json_f64(bucket.get("image_quota_remaining")).is_some_and(|value| value <= 0.0)
|
||||
{
|
||||
return true;
|
||||
}
|
||||
match (
|
||||
provider_pool_json_f64(bucket.get("image_quota_total")),
|
||||
provider_pool_json_f64(bucket.get("image_quota_used")),
|
||||
) {
|
||||
(Some(limit), Some(used)) if limit > 0.0 => used >= limit,
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,253 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogEndpoint;
|
||||
use aether_pool_core::PoolSchedulingPreset;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::capability::ProviderPoolCapabilities;
|
||||
use crate::provider::{
|
||||
provider_pool_endpoint_format_matches, provider_pool_matching_endpoint, ProviderPoolAdapter,
|
||||
ProviderPoolMemberInput,
|
||||
};
|
||||
use crate::quota::{
|
||||
provider_pool_current_unix_secs, provider_pool_json_bool, provider_pool_json_f64,
|
||||
provider_pool_metadata_bucket, provider_pool_quota_snapshot_exhausted_decision,
|
||||
provider_pool_reset_deadline_elapsed, provider_pool_timestamp_unix_secs,
|
||||
};
|
||||
use crate::quota_refresh::ProviderPoolQuotaRequestSpec;
|
||||
|
||||
pub const CODEX_WHAM_USAGE_URL: &str = "https://chatgpt.com/backend-api/wham/usage";
|
||||
pub const CODEX_WHAM_RESET_CREDITS_URL: &str =
|
||||
"https://chatgpt.com/backend-api/wham/rate-limit-reset-credits";
|
||||
pub const CODEX_WHAM_RESET_CREDITS_CONSUME_URL: &str =
|
||||
"https://chatgpt.com/backend-api/wham/rate-limit-reset-credits/consume";
|
||||
const PLACEHOLDER_API_KEY: &str = "__placeholder__";
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct CodexProviderPoolAdapter;
|
||||
|
||||
impl ProviderPoolAdapter for CodexProviderPoolAdapter {
|
||||
fn provider_type(&self) -> &'static str {
|
||||
"codex"
|
||||
}
|
||||
|
||||
fn capabilities(&self) -> ProviderPoolCapabilities {
|
||||
ProviderPoolCapabilities {
|
||||
plan_tier: true,
|
||||
quota_reset: true,
|
||||
quota_refresh: true,
|
||||
}
|
||||
}
|
||||
|
||||
fn default_scheduling_presets(&self) -> Vec<PoolSchedulingPreset> {
|
||||
vec![PoolSchedulingPreset {
|
||||
preset: "recent_refresh".to_string(),
|
||||
enabled: true,
|
||||
mode: None,
|
||||
}]
|
||||
}
|
||||
|
||||
fn quota_exhausted(&self, input: &ProviderPoolMemberInput<'_>) -> bool {
|
||||
if let Some(exhausted) =
|
||||
provider_pool_quota_snapshot_exhausted_decision(input.key, input.provider_type)
|
||||
{
|
||||
return exhausted;
|
||||
}
|
||||
provider_pool_metadata_bucket(input.key.upstream_metadata.as_ref(), input.provider_type)
|
||||
.is_some_and(quota_exhausted_from_bucket)
|
||||
}
|
||||
|
||||
fn quota_refresh_endpoint(
|
||||
&self,
|
||||
endpoints: &[StoredProviderCatalogEndpoint],
|
||||
include_inactive: bool,
|
||||
) -> Option<StoredProviderCatalogEndpoint> {
|
||||
provider_pool_matching_endpoint(endpoints, include_inactive, |endpoint| {
|
||||
provider_pool_endpoint_format_matches(endpoint, "openai:responses")
|
||||
})
|
||||
}
|
||||
|
||||
fn quota_refresh_missing_endpoint_message(&self) -> String {
|
||||
"找不到有效的 openai:responses 端点".to_string()
|
||||
}
|
||||
}
|
||||
|
||||
fn build_codex_wham_headers(
|
||||
resolved_oauth_auth: Option<(String, String)>,
|
||||
decrypted_api_key: Option<&str>,
|
||||
auth_config: Option<&Value>,
|
||||
) -> Result<BTreeMap<String, String>, String> {
|
||||
let mut headers = BTreeMap::new();
|
||||
headers.insert("accept".to_string(), "application/json".to_string());
|
||||
let auth_config_headers = auth_config
|
||||
.and_then(|value| value.get("headers"))
|
||||
.and_then(Value::as_object);
|
||||
let auth_config_authorization = auth_config_headers
|
||||
.and_then(|headers| {
|
||||
headers
|
||||
.get("authorization")
|
||||
.or_else(|| headers.get("Authorization"))
|
||||
})
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty());
|
||||
|
||||
if let Some(authorization) = auth_config_authorization {
|
||||
headers.insert("authorization".to_string(), authorization.to_string());
|
||||
} else if let Some((name, value)) = resolved_oauth_auth {
|
||||
headers.insert(name.to_ascii_lowercase(), value);
|
||||
} else {
|
||||
let decrypted_key = decrypted_api_key.unwrap_or_default().trim();
|
||||
if decrypted_key.is_empty() || decrypted_key == PLACEHOLDER_API_KEY {
|
||||
return Err("缺少 OAuth 认证信息,请先授权/刷新 Token".to_string());
|
||||
}
|
||||
headers.insert(
|
||||
"authorization".to_string(),
|
||||
format!("Bearer {decrypted_key}"),
|
||||
);
|
||||
}
|
||||
|
||||
let oauth_plan_type = auth_config
|
||||
.and_then(|value| value.get("plan_type"))
|
||||
.and_then(Value::as_str)
|
||||
.and_then(|value| crate::plan::normalize_provider_plan_tier(value, "codex"));
|
||||
let oauth_account_id = auth_config
|
||||
.and_then(|value| value.get("account_id"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty());
|
||||
if oauth_account_id.is_some() && oauth_plan_type.as_deref() != Some("free") {
|
||||
headers.insert(
|
||||
"chatgpt-account-id".to_string(),
|
||||
oauth_account_id.unwrap_or_default().to_string(),
|
||||
);
|
||||
}
|
||||
|
||||
Ok(headers)
|
||||
}
|
||||
|
||||
pub fn build_codex_pool_quota_request(
|
||||
key_id: &str,
|
||||
resolved_oauth_auth: Option<(String, String)>,
|
||||
decrypted_api_key: Option<&str>,
|
||||
auth_config: Option<&Value>,
|
||||
) -> Result<ProviderPoolQuotaRequestSpec, String> {
|
||||
let headers = build_codex_wham_headers(resolved_oauth_auth, decrypted_api_key, auth_config)?;
|
||||
|
||||
Ok(ProviderPoolQuotaRequestSpec {
|
||||
request_id: format!("codex-quota:{key_id}"),
|
||||
provider_name: "codex".to_string(),
|
||||
quota_kind: "codex".to_string(),
|
||||
method: "GET".to_string(),
|
||||
url: CODEX_WHAM_USAGE_URL.to_string(),
|
||||
headers,
|
||||
content_type: None,
|
||||
json_body: None,
|
||||
client_api_format: "openai:responses".to_string(),
|
||||
provider_api_format: "openai:responses".to_string(),
|
||||
model_name: Some("codex-wham-usage".to_string()),
|
||||
accept_invalid_certs: false,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn build_codex_pool_reset_credits_request(
|
||||
key_id: &str,
|
||||
resolved_oauth_auth: Option<(String, String)>,
|
||||
decrypted_api_key: Option<&str>,
|
||||
auth_config: Option<&Value>,
|
||||
) -> Result<ProviderPoolQuotaRequestSpec, String> {
|
||||
let headers = build_codex_wham_headers(resolved_oauth_auth, decrypted_api_key, auth_config)?;
|
||||
|
||||
Ok(ProviderPoolQuotaRequestSpec {
|
||||
request_id: format!("codex-reset-credits:{key_id}"),
|
||||
provider_name: "codex".to_string(),
|
||||
quota_kind: "codex_reset_credits".to_string(),
|
||||
method: "GET".to_string(),
|
||||
url: CODEX_WHAM_RESET_CREDITS_URL.to_string(),
|
||||
headers,
|
||||
content_type: None,
|
||||
json_body: None,
|
||||
client_api_format: "openai:responses".to_string(),
|
||||
provider_api_format: "openai:responses".to_string(),
|
||||
model_name: Some("codex-wham-reset-credits".to_string()),
|
||||
accept_invalid_certs: false,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn build_codex_pool_reset_credit_consume_request(
|
||||
key_id: &str,
|
||||
resolved_oauth_auth: Option<(String, String)>,
|
||||
decrypted_api_key: Option<&str>,
|
||||
auth_config: Option<&Value>,
|
||||
redeem_request_id: &str,
|
||||
) -> Result<ProviderPoolQuotaRequestSpec, String> {
|
||||
let redeem_request_id = redeem_request_id.trim();
|
||||
if redeem_request_id.is_empty() {
|
||||
return Err("缺少 Codex reset credit 幂等请求 ID".to_string());
|
||||
}
|
||||
|
||||
let mut headers =
|
||||
build_codex_wham_headers(resolved_oauth_auth, decrypted_api_key, auth_config)?;
|
||||
headers.insert("content-type".to_string(), "application/json".to_string());
|
||||
|
||||
Ok(ProviderPoolQuotaRequestSpec {
|
||||
request_id: format!("codex-reset-credit-consume:{key_id}:{redeem_request_id}"),
|
||||
provider_name: "codex".to_string(),
|
||||
quota_kind: "codex_reset_credit_consume".to_string(),
|
||||
method: "POST".to_string(),
|
||||
url: CODEX_WHAM_RESET_CREDITS_CONSUME_URL.to_string(),
|
||||
headers,
|
||||
content_type: Some("application/json".to_string()),
|
||||
json_body: Some(serde_json::json!({
|
||||
"redeem_request_id": redeem_request_id,
|
||||
})),
|
||||
client_api_format: "openai:responses".to_string(),
|
||||
provider_api_format: "openai:responses".to_string(),
|
||||
model_name: Some("codex-wham-reset-credit-consume".to_string()),
|
||||
accept_invalid_certs: false,
|
||||
})
|
||||
}
|
||||
|
||||
fn codex_window_reset_elapsed(bucket: &Map<String, Value>, prefix: &str) -> bool {
|
||||
let Some(now_unix_secs) = provider_pool_current_unix_secs() else {
|
||||
return false;
|
||||
};
|
||||
let mut window = Map::new();
|
||||
for (target, source) in [
|
||||
("reset_at", format!("{prefix}_reset_at")),
|
||||
("next_reset_at", format!("{prefix}_next_reset_at")),
|
||||
("reset_seconds", format!("{prefix}_reset_seconds")),
|
||||
(
|
||||
"reset_after_seconds",
|
||||
format!("{prefix}_reset_after_seconds"),
|
||||
),
|
||||
] {
|
||||
if let Some(value) = bucket.get(source.as_str()) {
|
||||
window.insert(target.to_string(), value.clone());
|
||||
}
|
||||
}
|
||||
provider_pool_reset_deadline_elapsed(
|
||||
&window,
|
||||
provider_pool_timestamp_unix_secs(bucket.get("updated_at")),
|
||||
now_unix_secs,
|
||||
)
|
||||
}
|
||||
|
||||
fn codex_window_used_percent_exhausted(bucket: &Map<String, Value>, prefix: &str) -> bool {
|
||||
let used_percent_key = format!("{prefix}_used_percent");
|
||||
provider_pool_json_f64(bucket.get(used_percent_key.as_str()))
|
||||
.is_some_and(|value| value >= 100.0 && !codex_window_reset_elapsed(bucket, prefix))
|
||||
}
|
||||
|
||||
pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
|
||||
if provider_pool_json_bool(bucket.get("credits_unlimited")) == Some(true) {
|
||||
return false;
|
||||
}
|
||||
let has_window_data = provider_pool_json_f64(bucket.get("primary_used_percent")).is_some()
|
||||
|| provider_pool_json_f64(bucket.get("secondary_used_percent")).is_some();
|
||||
if !has_window_data && provider_pool_json_bool(bucket.get("has_credits")) == Some(false) {
|
||||
return true;
|
||||
}
|
||||
codex_window_used_percent_exhausted(bucket, "primary")
|
||||
|| codex_window_used_percent_exhausted(bucket, "secondary")
|
||||
}
|
||||
@@ -0,0 +1,10 @@
|
||||
use crate::provider::ProviderPoolAdapter;
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct DefaultProviderPoolAdapter;
|
||||
|
||||
impl ProviderPoolAdapter for DefaultProviderPoolAdapter {
|
||||
fn provider_type(&self) -> &'static str {
|
||||
"default"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogEndpoint;
|
||||
use serde_json::json;
|
||||
|
||||
use crate::capability::ProviderPoolCapabilities;
|
||||
use crate::provider::{
|
||||
provider_pool_endpoint_format_matches, provider_pool_matching_endpoint, ProviderPoolAdapter,
|
||||
};
|
||||
use crate::quota_refresh::ProviderPoolQuotaRequestSpec;
|
||||
|
||||
pub const GEMINI_CLI_RETRIEVE_USER_QUOTA_PATH: &str = "/v1internal:retrieveUserQuota";
|
||||
pub const GEMINI_CLI_USER_AGENT: &str = "GeminiCLI/0.1.5 (Windows; AMD64)";
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct GeminiCliProviderPoolAdapter;
|
||||
|
||||
impl ProviderPoolAdapter for GeminiCliProviderPoolAdapter {
|
||||
fn provider_type(&self) -> &'static str {
|
||||
"gemini_cli"
|
||||
}
|
||||
|
||||
fn capabilities(&self) -> ProviderPoolCapabilities {
|
||||
ProviderPoolCapabilities {
|
||||
quota_refresh: true,
|
||||
..ProviderPoolCapabilities::default()
|
||||
}
|
||||
}
|
||||
|
||||
fn quota_refresh_endpoint(
|
||||
&self,
|
||||
endpoints: &[StoredProviderCatalogEndpoint],
|
||||
include_inactive: bool,
|
||||
) -> Option<StoredProviderCatalogEndpoint> {
|
||||
provider_pool_matching_endpoint(endpoints, include_inactive, |endpoint| {
|
||||
provider_pool_endpoint_format_matches(endpoint, "gemini:generate_content")
|
||||
})
|
||||
}
|
||||
|
||||
fn quota_refresh_missing_endpoint_message(&self) -> String {
|
||||
"找不到有效的 gemini:generate_content 端点".to_string()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn build_gemini_cli_pool_quota_request(
|
||||
key_id: &str,
|
||||
endpoint_base_url: &str,
|
||||
authorization: (String, String),
|
||||
project_id: &str,
|
||||
) -> ProviderPoolQuotaRequestSpec {
|
||||
let headers = BTreeMap::from([
|
||||
("authorization".to_string(), authorization.1),
|
||||
("content-type".to_string(), "application/json".to_string()),
|
||||
("accept".to_string(), "application/json".to_string()),
|
||||
("user-agent".to_string(), GEMINI_CLI_USER_AGENT.to_string()),
|
||||
]);
|
||||
|
||||
ProviderPoolQuotaRequestSpec {
|
||||
request_id: format!("gemini-cli-quota:{key_id}"),
|
||||
provider_name: "gemini_cli".to_string(),
|
||||
quota_kind: "gemini_cli".to_string(),
|
||||
method: "POST".to_string(),
|
||||
url: format!(
|
||||
"{}{}",
|
||||
endpoint_base_url.trim_end_matches('/'),
|
||||
GEMINI_CLI_RETRIEVE_USER_QUOTA_PATH
|
||||
),
|
||||
headers,
|
||||
content_type: Some("application/json".to_string()),
|
||||
json_body: Some(json!({
|
||||
"project": project_id,
|
||||
})),
|
||||
client_api_format: "gemini:generate_content".to_string(),
|
||||
provider_api_format: "gemini_cli:retrieve_user_quota".to_string(),
|
||||
model_name: Some("retrieveUserQuota".to_string()),
|
||||
accept_invalid_certs: false,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,225 @@
|
||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogEndpoint;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::capability::ProviderPoolCapabilities;
|
||||
use crate::provider::{
|
||||
provider_pool_endpoint_format_matches, provider_pool_matching_endpoint, ProviderPoolAdapter,
|
||||
ProviderPoolMemberInput,
|
||||
};
|
||||
use crate::quota::{
|
||||
provider_pool_current_unix_secs, provider_pool_json_bool, provider_pool_json_f64,
|
||||
provider_pool_metadata_bucket, provider_pool_quota_snapshot_exhausted_decision,
|
||||
provider_pool_reset_deadline_elapsed, provider_pool_timestamp_unix_secs,
|
||||
};
|
||||
|
||||
pub const GROK_QUOTA_WINDOWS_BASIC: &[(&str, &str)] = &[("quota_fast", "fast")];
|
||||
pub const GROK_QUOTA_WINDOWS_SUPER: &[(&str, &str)] = &[
|
||||
("quota_auto", "auto"),
|
||||
("quota_fast", "fast"),
|
||||
("quota_expert", "expert"),
|
||||
("quota_grok_4_3", "grok-420-computer-use-sa"),
|
||||
];
|
||||
pub const GROK_QUOTA_WINDOWS_HEAVY: &[(&str, &str)] = &[
|
||||
("quota_auto", "auto"),
|
||||
("quota_fast", "fast"),
|
||||
("quota_expert", "expert"),
|
||||
("quota_heavy", "heavy"),
|
||||
("quota_grok_4_3", "grok-420-computer-use-sa"),
|
||||
];
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct GrokProviderPoolAdapter;
|
||||
|
||||
impl ProviderPoolAdapter for GrokProviderPoolAdapter {
|
||||
fn provider_type(&self) -> &'static str {
|
||||
"grok"
|
||||
}
|
||||
|
||||
fn capabilities(&self) -> ProviderPoolCapabilities {
|
||||
ProviderPoolCapabilities {
|
||||
plan_tier: true,
|
||||
quota_reset: true,
|
||||
quota_refresh: true,
|
||||
}
|
||||
}
|
||||
|
||||
fn quota_exhausted(&self, input: &ProviderPoolMemberInput<'_>) -> bool {
|
||||
if let Some(exhausted) =
|
||||
provider_pool_quota_snapshot_exhausted_decision(input.key, input.provider_type)
|
||||
{
|
||||
return exhausted;
|
||||
}
|
||||
provider_pool_metadata_bucket(input.key.upstream_metadata.as_ref(), input.provider_type)
|
||||
.is_some_and(quota_exhausted_from_bucket)
|
||||
}
|
||||
|
||||
fn quota_refresh_endpoint(
|
||||
&self,
|
||||
endpoints: &[StoredProviderCatalogEndpoint],
|
||||
include_inactive: bool,
|
||||
) -> Option<StoredProviderCatalogEndpoint> {
|
||||
provider_pool_matching_endpoint(endpoints, include_inactive, |endpoint| {
|
||||
provider_pool_endpoint_format_matches(endpoint, "openai:chat")
|
||||
})
|
||||
.or_else(|| provider_pool_matching_endpoint(endpoints, include_inactive, |_| true))
|
||||
}
|
||||
|
||||
fn quota_refresh_missing_endpoint_message(&self) -> String {
|
||||
"找不到有效的 Grok 端点".to_string()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn grok_supported_quota_windows_for_tier(
|
||||
tier: Option<&str>,
|
||||
) -> &'static [(&'static str, &'static str)] {
|
||||
match grok_normalize_pool_tier(tier) {
|
||||
Some("basic") => GROK_QUOTA_WINDOWS_BASIC,
|
||||
Some("super") => GROK_QUOTA_WINDOWS_SUPER,
|
||||
Some("heavy") => GROK_QUOTA_WINDOWS_HEAVY,
|
||||
_ => GROK_QUOTA_WINDOWS_HEAVY,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn grok_pool_tier_from_quota_bucket(bucket: &Map<String, Value>) -> Option<&'static str> {
|
||||
if let Some(tier) = grok_normalize_pool_tier(
|
||||
grok_bucket_string(bucket, &["pool_tier", "tier", "plan_type", "plan"]).as_deref(),
|
||||
) {
|
||||
return Some(tier);
|
||||
}
|
||||
|
||||
if let Some(auto_total) = grok_quota_total(bucket, "quota_auto") {
|
||||
if (auto_total - 50.0).abs() < f64::EPSILON {
|
||||
return Some("super");
|
||||
}
|
||||
if (auto_total - 150.0).abs() < f64::EPSILON {
|
||||
return Some("heavy");
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(fast_total) = grok_quota_total(bucket, "quota_fast") {
|
||||
if (fast_total - 30.0).abs() < f64::EPSILON {
|
||||
return Some("basic");
|
||||
}
|
||||
if (fast_total - 140.0).abs() < f64::EPSILON {
|
||||
return Some("super");
|
||||
}
|
||||
if (fast_total - 400.0).abs() < f64::EPSILON {
|
||||
return Some("heavy");
|
||||
}
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
pub fn grok_quota_window_key_for_model(model: Option<&str>) -> Option<&'static str> {
|
||||
Some(match grok_mode_id_for_model(model) {
|
||||
"fast" => "quota_fast",
|
||||
"auto" => "quota_auto",
|
||||
"expert" => "quota_expert",
|
||||
"heavy" => "quota_heavy",
|
||||
"grok-420-computer-use-sa" => "quota_grok_4_3",
|
||||
_ => return None,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn grok_mode_id_for_model(model: Option<&str>) -> &'static str {
|
||||
let model = model.unwrap_or_default().to_ascii_lowercase();
|
||||
if model.contains("4.3") || model.contains("computer") {
|
||||
"grok-420-computer-use-sa"
|
||||
} else if model.contains("multi-agent") {
|
||||
"heavy"
|
||||
} else if model.contains("non-reasoning") || model.contains("fast") || model.contains("lite") {
|
||||
"fast"
|
||||
} else if model.contains("expert") || model.contains("reasoning") {
|
||||
"expert"
|
||||
} else if model.contains("0309-heavy") {
|
||||
"auto"
|
||||
} else if model.contains("heavy") {
|
||||
"heavy"
|
||||
} else {
|
||||
"auto"
|
||||
}
|
||||
}
|
||||
|
||||
fn grok_normalize_pool_tier(value: Option<&str>) -> Option<&'static str> {
|
||||
match value?.trim().to_ascii_lowercase().as_str() {
|
||||
"basic" => Some("basic"),
|
||||
"super" => Some("super"),
|
||||
"heavy" => Some("heavy"),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn grok_bucket_string(bucket: &Map<String, Value>, fields: &[&str]) -> Option<String> {
|
||||
fields.iter().find_map(|field| {
|
||||
bucket
|
||||
.get(*field)
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
})
|
||||
}
|
||||
|
||||
fn grok_quota_total(quota_by_model: &Map<String, Value>, key: &str) -> Option<f64> {
|
||||
let models = quota_by_model
|
||||
.get("quota_by_model")
|
||||
.or_else(|| quota_by_model.get("models"))
|
||||
.and_then(Value::as_object)
|
||||
.unwrap_or(quota_by_model);
|
||||
models
|
||||
.get(key)
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|quota| quota.get("total"))
|
||||
.and_then(Value::as_f64)
|
||||
.filter(|value| value.is_finite() && *value > 0.0)
|
||||
}
|
||||
|
||||
pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
|
||||
let models = bucket
|
||||
.get("quota_by_model")
|
||||
.or_else(|| bucket.get("models"))
|
||||
.and_then(Value::as_object);
|
||||
let Some(models) = models else {
|
||||
return false;
|
||||
};
|
||||
|
||||
let supported_mode_keys =
|
||||
grok_supported_quota_windows_for_tier(grok_pool_tier_from_quota_bucket(bucket))
|
||||
.iter()
|
||||
.map(|(quota_key, _)| *quota_key)
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let mut model_count = 0usize;
|
||||
let mut exhausted_count = 0usize;
|
||||
let now_unix_secs = provider_pool_current_unix_secs();
|
||||
let bucket_updated_at = provider_pool_timestamp_unix_secs(bucket.get("updated_at"));
|
||||
for (model_key, item) in models.iter() {
|
||||
if !supported_mode_keys.is_empty() && !supported_mode_keys.contains(&model_key.as_str()) {
|
||||
continue;
|
||||
}
|
||||
|
||||
let Some(item) = item.as_object() else {
|
||||
continue;
|
||||
};
|
||||
let has_quota_data = provider_pool_json_bool(item.get("is_exhausted")).is_some()
|
||||
|| provider_pool_json_f64(item.get("used_percent")).is_some()
|
||||
|| provider_pool_json_f64(item.get("remaining")).is_some()
|
||||
|| provider_pool_json_f64(item.get("remaining_fraction")).is_some();
|
||||
if !has_quota_data {
|
||||
continue;
|
||||
}
|
||||
model_count += 1;
|
||||
let quota_exhausted = provider_pool_json_bool(item.get("is_exhausted")) == Some(true)
|
||||
|| provider_pool_json_f64(item.get("used_percent")).is_some_and(|value| value >= 100.0)
|
||||
|| provider_pool_json_f64(item.get("remaining")).is_some_and(|value| value <= 0.0)
|
||||
|| provider_pool_json_f64(item.get("remaining_fraction"))
|
||||
.is_some_and(|value| value <= 0.0);
|
||||
let reset_elapsed = now_unix_secs
|
||||
.is_some_and(|now| provider_pool_reset_deadline_elapsed(item, bucket_updated_at, now));
|
||||
if quota_exhausted && !reset_elapsed {
|
||||
exhausted_count += 1;
|
||||
}
|
||||
}
|
||||
model_count > 0 && model_count == exhausted_count
|
||||
}
|
||||
@@ -0,0 +1,177 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogEndpoint;
|
||||
use serde_json::{Map, Value};
|
||||
use url::form_urlencoded;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::capability::ProviderPoolCapabilities;
|
||||
use crate::provider::{
|
||||
provider_pool_endpoint_format_matches, provider_pool_matching_endpoint, ProviderPoolAdapter,
|
||||
ProviderPoolMemberInput,
|
||||
};
|
||||
use crate::quota::{
|
||||
provider_pool_current_unix_secs, provider_pool_json_f64, provider_pool_metadata_bucket,
|
||||
provider_pool_quota_snapshot_exhausted_decision, provider_pool_reset_deadline_elapsed,
|
||||
provider_pool_timestamp_unix_secs,
|
||||
};
|
||||
use crate::quota_refresh::ProviderPoolQuotaRequestSpec;
|
||||
|
||||
pub const KIRO_USAGE_LIMITS_PATH: &str = "/getUsageLimits";
|
||||
pub const KIRO_USAGE_SDK_VERSION: &str = "1.0.0";
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct KiroPoolQuotaAuthInput {
|
||||
pub authorization_value: String,
|
||||
pub api_region: String,
|
||||
pub kiro_version: String,
|
||||
pub machine_id: String,
|
||||
pub profile_arn: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct KiroProviderPoolAdapter;
|
||||
|
||||
impl ProviderPoolAdapter for KiroProviderPoolAdapter {
|
||||
fn provider_type(&self) -> &'static str {
|
||||
"kiro"
|
||||
}
|
||||
|
||||
fn capabilities(&self) -> ProviderPoolCapabilities {
|
||||
ProviderPoolCapabilities {
|
||||
plan_tier: true,
|
||||
quota_reset: true,
|
||||
quota_refresh: true,
|
||||
}
|
||||
}
|
||||
|
||||
fn quota_exhausted(&self, input: &ProviderPoolMemberInput<'_>) -> bool {
|
||||
if let Some(exhausted) =
|
||||
provider_pool_quota_snapshot_exhausted_decision(input.key, input.provider_type)
|
||||
{
|
||||
return exhausted;
|
||||
}
|
||||
provider_pool_metadata_bucket(input.key.upstream_metadata.as_ref(), input.provider_type)
|
||||
.is_some_and(quota_exhausted_from_bucket)
|
||||
}
|
||||
|
||||
fn quota_refresh_endpoint(
|
||||
&self,
|
||||
endpoints: &[StoredProviderCatalogEndpoint],
|
||||
include_inactive: bool,
|
||||
) -> Option<StoredProviderCatalogEndpoint> {
|
||||
provider_pool_matching_endpoint(endpoints, include_inactive, |endpoint| {
|
||||
provider_pool_endpoint_format_matches(endpoint, "claude:messages")
|
||||
})
|
||||
.or_else(|| provider_pool_matching_endpoint(endpoints, include_inactive, |_| true))
|
||||
}
|
||||
|
||||
fn quota_refresh_missing_endpoint_message(&self) -> String {
|
||||
"找不到有效的 Kiro 端点".to_string()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn build_kiro_pool_quota_request(
|
||||
key_id: &str,
|
||||
auth: &KiroPoolQuotaAuthInput,
|
||||
) -> ProviderPoolQuotaRequestSpec {
|
||||
let host = format!("q.{}.amazonaws.com", normalize_region(&auth.api_region));
|
||||
let machine_id = auth.machine_id.trim();
|
||||
let ide_tag = if machine_id.is_empty() {
|
||||
format!("KiroIDE-{}", normalize_kiro_version(&auth.kiro_version))
|
||||
} else {
|
||||
format!(
|
||||
"KiroIDE-{}-{machine_id}",
|
||||
normalize_kiro_version(&auth.kiro_version)
|
||||
)
|
||||
};
|
||||
let mut serializer = form_urlencoded::Serializer::new(String::new());
|
||||
serializer.append_pair("origin", "AI_EDITOR");
|
||||
serializer.append_pair("resourceType", "AGENTIC_REQUEST");
|
||||
serializer.append_pair("isEmailRequired", "true");
|
||||
if let Some(profile_arn) = auth
|
||||
.profile_arn
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
serializer.append_pair("profileArn", profile_arn);
|
||||
}
|
||||
|
||||
ProviderPoolQuotaRequestSpec {
|
||||
request_id: format!("kiro-quota:{key_id}"),
|
||||
provider_name: "kiro".to_string(),
|
||||
quota_kind: "kiro".to_string(),
|
||||
method: "GET".to_string(),
|
||||
url: format!("https://{host}{KIRO_USAGE_LIMITS_PATH}?{}", serializer.finish()),
|
||||
headers: BTreeMap::from([
|
||||
(
|
||||
"x-amz-user-agent".to_string(),
|
||||
format!("aws-sdk-js/{KIRO_USAGE_SDK_VERSION} {ide_tag}"),
|
||||
),
|
||||
(
|
||||
"user-agent".to_string(),
|
||||
format!(
|
||||
"aws-sdk-js/{KIRO_USAGE_SDK_VERSION} ua/2.1 os/other#unknown lang/js md/nodejs#22.21.1 api/codewhispererruntime#1.0.0 m/N,E {ide_tag}"
|
||||
),
|
||||
),
|
||||
("host".to_string(), host),
|
||||
("amz-sdk-invocation-id".to_string(), Uuid::new_v4().to_string()),
|
||||
("amz-sdk-request".to_string(), "attempt=1; max=1".to_string()),
|
||||
(
|
||||
"authorization".to_string(),
|
||||
auth.authorization_value.clone(),
|
||||
),
|
||||
("connection".to_string(), "close".to_string()),
|
||||
]),
|
||||
content_type: None,
|
||||
json_body: None,
|
||||
client_api_format: "claude:messages".to_string(),
|
||||
provider_api_format: "kiro:usage".to_string(),
|
||||
model_name: Some("kiro-usage-limits".to_string()),
|
||||
accept_invalid_certs: false,
|
||||
}
|
||||
}
|
||||
|
||||
fn normalize_region(value: &str) -> &str {
|
||||
let value = value.trim();
|
||||
if value.is_empty() {
|
||||
"us-east-1"
|
||||
} else {
|
||||
value
|
||||
}
|
||||
}
|
||||
|
||||
fn normalize_kiro_version(value: &str) -> &str {
|
||||
let value = value.trim();
|
||||
if value.is_empty() {
|
||||
"0.3.210"
|
||||
} else {
|
||||
value
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
|
||||
if provider_pool_current_unix_secs().is_some_and(|now| {
|
||||
provider_pool_reset_deadline_elapsed(
|
||||
bucket,
|
||||
provider_pool_timestamp_unix_secs(bucket.get("updated_at")),
|
||||
now,
|
||||
)
|
||||
}) {
|
||||
return false;
|
||||
}
|
||||
if provider_pool_json_f64(bucket.get("remaining")).is_some_and(|value| value <= 0.0) {
|
||||
return true;
|
||||
}
|
||||
if provider_pool_json_f64(bucket.get("usage_percentage")).is_some_and(|value| value >= 100.0) {
|
||||
return true;
|
||||
}
|
||||
match (
|
||||
provider_pool_json_f64(bucket.get("usage_limit")),
|
||||
provider_pool_json_f64(bucket.get("current_usage")),
|
||||
) {
|
||||
(Some(limit), Some(current)) if limit > 0.0 => current >= limit,
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,52 @@
|
||||
pub mod antigravity;
|
||||
pub mod chatgpt_web;
|
||||
pub mod codex;
|
||||
pub mod default;
|
||||
pub mod gemini_cli;
|
||||
pub mod grok;
|
||||
pub mod kiro;
|
||||
pub mod unsupported;
|
||||
pub mod windsurf;
|
||||
|
||||
pub use antigravity::AntigravityProviderPoolAdapter;
|
||||
pub use antigravity::{
|
||||
build_antigravity_pool_quota_request, ANTIGRAVITY_FETCH_AVAILABLE_MODELS_PATH,
|
||||
};
|
||||
pub use chatgpt_web::ChatGptWebProviderPoolAdapter;
|
||||
pub use chatgpt_web::{
|
||||
build_chatgpt_web_pool_quota_request, enrich_chatgpt_web_quota_metadata,
|
||||
normalize_chatgpt_web_image_quota_limit, CHATGPT_WEB_CONVERSATION_INIT_PATH,
|
||||
CHATGPT_WEB_DEFAULT_BASE_URL,
|
||||
};
|
||||
pub use codex::CodexProviderPoolAdapter;
|
||||
pub use codex::{
|
||||
build_codex_pool_quota_request, build_codex_pool_reset_credit_consume_request,
|
||||
build_codex_pool_reset_credits_request, CODEX_WHAM_RESET_CREDITS_CONSUME_URL,
|
||||
CODEX_WHAM_RESET_CREDITS_URL, CODEX_WHAM_USAGE_URL,
|
||||
};
|
||||
pub use default::DefaultProviderPoolAdapter;
|
||||
pub use gemini_cli::GeminiCliProviderPoolAdapter;
|
||||
pub use gemini_cli::{
|
||||
build_gemini_cli_pool_quota_request, GEMINI_CLI_RETRIEVE_USER_QUOTA_PATH, GEMINI_CLI_USER_AGENT,
|
||||
};
|
||||
pub use grok::{
|
||||
grok_mode_id_for_model, grok_pool_tier_from_quota_bucket, grok_quota_window_key_for_model,
|
||||
grok_supported_quota_windows_for_tier, GrokProviderPoolAdapter,
|
||||
};
|
||||
pub use kiro::KiroProviderPoolAdapter;
|
||||
pub use kiro::{
|
||||
build_kiro_pool_quota_request, KiroPoolQuotaAuthInput, KIRO_USAGE_LIMITS_PATH,
|
||||
KIRO_USAGE_SDK_VERSION,
|
||||
};
|
||||
pub use unsupported::{
|
||||
UnsupportedQuotaProviderPoolAdapter, CLAUDE_CODE_PROVIDER_POOL_ADAPTER,
|
||||
VERTEX_AI_PROVIDER_POOL_ADAPTER,
|
||||
};
|
||||
pub use windsurf::{
|
||||
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, WindsurfProviderPoolAdapter,
|
||||
WINDSURF_DEFAULT_BASE_URL, WINDSURF_MODEL_CONFIGS_PATH, WINDSURF_RATE_LIMIT_PATH,
|
||||
WINDSURF_USER_STATUS_PATH,
|
||||
};
|
||||
@@ -0,0 +1,41 @@
|
||||
use crate::provider::ProviderPoolAdapter;
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct UnsupportedQuotaProviderPoolAdapter {
|
||||
provider_type: &'static str,
|
||||
quota_refresh_unsupported_message: &'static str,
|
||||
}
|
||||
|
||||
impl UnsupportedQuotaProviderPoolAdapter {
|
||||
pub const fn new(
|
||||
provider_type: &'static str,
|
||||
quota_refresh_unsupported_message: &'static str,
|
||||
) -> Self {
|
||||
Self {
|
||||
provider_type,
|
||||
quota_refresh_unsupported_message,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ProviderPoolAdapter for UnsupportedQuotaProviderPoolAdapter {
|
||||
fn provider_type(&self) -> &'static str {
|
||||
self.provider_type
|
||||
}
|
||||
|
||||
fn quota_refresh_unsupported_message(&self) -> String {
|
||||
self.quota_refresh_unsupported_message.to_string()
|
||||
}
|
||||
}
|
||||
|
||||
pub const CLAUDE_CODE_PROVIDER_POOL_ADAPTER: UnsupportedQuotaProviderPoolAdapter =
|
||||
UnsupportedQuotaProviderPoolAdapter::new(
|
||||
"claude_code",
|
||||
"Claude Code 暂不支持自动刷新额度:上游没有稳定可用的账号额度查询接口",
|
||||
);
|
||||
|
||||
pub const VERTEX_AI_PROVIDER_POOL_ADAPTER: UnsupportedQuotaProviderPoolAdapter =
|
||||
UnsupportedQuotaProviderPoolAdapter::new(
|
||||
"vertex_ai",
|
||||
"Vertex AI 暂不支持自动刷新额度:额度属于 Google Cloud 项目/区域配额",
|
||||
);
|
||||
@@ -0,0 +1,286 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
};
|
||||
use aether_pool_core::PoolSchedulingPreset;
|
||||
use serde_json::{json, Map, Value};
|
||||
|
||||
use crate::capability::ProviderPoolCapabilities;
|
||||
use crate::provider::{
|
||||
provider_pool_endpoint_format_matches, provider_pool_matching_endpoint, ProviderPoolAdapter,
|
||||
ProviderPoolMemberInput,
|
||||
};
|
||||
use crate::quota::{
|
||||
provider_pool_json_bool, provider_pool_json_f64, provider_pool_member_quota_snapshot,
|
||||
provider_pool_metadata_bucket, provider_pool_quota_snapshot_exhausted_decision,
|
||||
};
|
||||
use crate::quota_refresh::ProviderPoolQuotaRequestSpec;
|
||||
|
||||
pub const WINDSURF_DEFAULT_BASE_URL: &str = "https://server.codeium.com";
|
||||
pub const WINDSURF_USER_STATUS_PATH: &str =
|
||||
"/exa.seat_management_pb.SeatManagementService/GetUserStatus";
|
||||
pub const WINDSURF_MODEL_CONFIGS_PATH: &str =
|
||||
"/exa.api_server_pb.ApiServerService/GetCascadeModelConfigs";
|
||||
pub const WINDSURF_RATE_LIMIT_PATH: &str =
|
||||
"/exa.api_server_pb.ApiServerService/CheckUserMessageRateLimit";
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct WindsurfProviderPoolAdapter;
|
||||
|
||||
impl ProviderPoolAdapter for WindsurfProviderPoolAdapter {
|
||||
fn provider_type(&self) -> &'static str {
|
||||
"windsurf"
|
||||
}
|
||||
|
||||
fn capabilities(&self) -> ProviderPoolCapabilities {
|
||||
ProviderPoolCapabilities {
|
||||
plan_tier: true,
|
||||
quota_reset: true,
|
||||
quota_refresh: true,
|
||||
}
|
||||
}
|
||||
|
||||
fn default_scheduling_presets(&self) -> Vec<PoolSchedulingPreset> {
|
||||
vec![PoolSchedulingPreset {
|
||||
preset: "recent_refresh".to_string(),
|
||||
enabled: true,
|
||||
mode: None,
|
||||
}]
|
||||
}
|
||||
|
||||
fn quota_exhausted(&self, input: &ProviderPoolMemberInput<'_>) -> bool {
|
||||
if windsurf_quota_snapshot_hard_exhausted(input.key, input.provider_type) {
|
||||
return true;
|
||||
}
|
||||
if let Some(exhausted) =
|
||||
provider_pool_quota_snapshot_exhausted_decision(input.key, input.provider_type)
|
||||
{
|
||||
return exhausted;
|
||||
}
|
||||
provider_pool_metadata_bucket(input.key.upstream_metadata.as_ref(), input.provider_type)
|
||||
.is_some_and(windsurf_quota_exhausted_from_bucket)
|
||||
}
|
||||
|
||||
fn quota_refresh_endpoint(
|
||||
&self,
|
||||
endpoints: &[StoredProviderCatalogEndpoint],
|
||||
include_inactive: bool,
|
||||
) -> Option<StoredProviderCatalogEndpoint> {
|
||||
provider_pool_matching_endpoint(endpoints, include_inactive, |endpoint| {
|
||||
provider_pool_endpoint_format_matches(endpoint, "openai:chat")
|
||||
})
|
||||
}
|
||||
|
||||
fn quota_refresh_missing_endpoint_message(&self) -> String {
|
||||
"找不到有效的 openai:chat 端点".to_string()
|
||||
}
|
||||
}
|
||||
|
||||
fn windsurf_quota_snapshot_hard_exhausted(
|
||||
key: &StoredProviderCatalogKey,
|
||||
provider_type: &str,
|
||||
) -> bool {
|
||||
provider_pool_member_quota_snapshot(key, provider_type)
|
||||
.and_then(|quota| quota.get("code"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.map(str::to_ascii_lowercase)
|
||||
.is_some_and(|code| matches!(code.as_str(), "banned" | "forbidden" | "quarantined"))
|
||||
}
|
||||
|
||||
pub fn build_windsurf_pool_quota_request(
|
||||
key_id: &str,
|
||||
api_key: &str,
|
||||
) -> ProviderPoolQuotaRequestSpec {
|
||||
build_windsurf_pool_quota_request_with_base_url(key_id, WINDSURF_DEFAULT_BASE_URL, api_key)
|
||||
}
|
||||
|
||||
pub fn build_windsurf_pool_quota_request_with_base_url(
|
||||
key_id: &str,
|
||||
base_url: &str,
|
||||
api_key: &str,
|
||||
) -> ProviderPoolQuotaRequestSpec {
|
||||
build_windsurf_connect_rpc_request(
|
||||
format!("windsurf-quota:{key_id}"),
|
||||
"windsurf:user_status",
|
||||
"windsurf-user-status",
|
||||
base_url,
|
||||
WINDSURF_USER_STATUS_PATH,
|
||||
api_key,
|
||||
)
|
||||
}
|
||||
|
||||
pub fn build_windsurf_pool_model_configs_request(
|
||||
key_id: &str,
|
||||
api_key: &str,
|
||||
) -> ProviderPoolQuotaRequestSpec {
|
||||
build_windsurf_pool_model_configs_request_with_base_url(
|
||||
key_id,
|
||||
WINDSURF_DEFAULT_BASE_URL,
|
||||
api_key,
|
||||
)
|
||||
}
|
||||
|
||||
pub fn build_windsurf_pool_model_configs_request_with_base_url(
|
||||
key_id: &str,
|
||||
base_url: &str,
|
||||
api_key: &str,
|
||||
) -> ProviderPoolQuotaRequestSpec {
|
||||
build_windsurf_connect_rpc_request(
|
||||
format!("windsurf-models:{key_id}"),
|
||||
"windsurf:model_configs",
|
||||
"windsurf-model-configs",
|
||||
base_url,
|
||||
WINDSURF_MODEL_CONFIGS_PATH,
|
||||
api_key,
|
||||
)
|
||||
}
|
||||
|
||||
pub fn build_windsurf_pool_rate_limit_request(
|
||||
key_id: &str,
|
||||
api_key: &str,
|
||||
) -> ProviderPoolQuotaRequestSpec {
|
||||
build_windsurf_pool_rate_limit_request_with_base_url(key_id, WINDSURF_DEFAULT_BASE_URL, api_key)
|
||||
}
|
||||
|
||||
pub fn build_windsurf_pool_rate_limit_request_with_base_url(
|
||||
key_id: &str,
|
||||
base_url: &str,
|
||||
api_key: &str,
|
||||
) -> ProviderPoolQuotaRequestSpec {
|
||||
build_windsurf_connect_rpc_request(
|
||||
format!("windsurf-rate-limit:{key_id}"),
|
||||
"windsurf:rate_limit",
|
||||
"windsurf-rate-limit",
|
||||
base_url,
|
||||
WINDSURF_RATE_LIMIT_PATH,
|
||||
api_key,
|
||||
)
|
||||
}
|
||||
|
||||
fn build_windsurf_connect_rpc_request(
|
||||
request_id: String,
|
||||
provider_api_format: &str,
|
||||
model_name: &str,
|
||||
base_url: &str,
|
||||
path: &str,
|
||||
api_key: &str,
|
||||
) -> ProviderPoolQuotaRequestSpec {
|
||||
let mut headers = BTreeMap::new();
|
||||
headers.insert("content-type".to_string(), "application/json".to_string());
|
||||
headers.insert("accept".to_string(), "application/json".to_string());
|
||||
headers.insert("connect-protocol-version".to_string(), "1".to_string());
|
||||
headers.insert("user-agent".to_string(), "windsurf/1.9600.41".to_string());
|
||||
|
||||
ProviderPoolQuotaRequestSpec {
|
||||
request_id,
|
||||
provider_name: "windsurf".to_string(),
|
||||
quota_kind: "windsurf".to_string(),
|
||||
method: "POST".to_string(),
|
||||
url: format!("{}{}", base_url.trim_end_matches('/'), path),
|
||||
headers,
|
||||
content_type: Some("application/json".to_string()),
|
||||
json_body: Some(json!({
|
||||
"metadata": windsurf_metadata(api_key),
|
||||
})),
|
||||
client_api_format: "openai:chat".to_string(),
|
||||
provider_api_format: provider_api_format.to_string(),
|
||||
model_name: Some(model_name.to_string()),
|
||||
accept_invalid_certs: false,
|
||||
}
|
||||
}
|
||||
|
||||
fn windsurf_metadata(api_key: &str) -> Value {
|
||||
json!({
|
||||
"apiKey": api_key,
|
||||
"ideName": "windsurf",
|
||||
"ideVersion": "1.9600.41",
|
||||
"extensionName": "windsurf",
|
||||
"extensionVersion": "1.9600.41",
|
||||
"locale": "en",
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn windsurf_quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
|
||||
if provider_pool_json_bool(bucket.get("banned"))
|
||||
.or_else(|| provider_pool_json_bool(bucket.get("quarantined")))
|
||||
.unwrap_or(false)
|
||||
{
|
||||
return true;
|
||||
}
|
||||
let daily_remaining = provider_pool_json_f64(bucket.get("daily_remaining_percent"));
|
||||
let weekly_remaining = provider_pool_json_f64(bucket.get("weekly_remaining_percent"));
|
||||
daily_remaining.is_some_and(|value| value <= 0.0)
|
||||
|| weekly_remaining.is_some_and(|value| value <= 0.0)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{windsurf_quota_exhausted_from_bucket, windsurf_quota_snapshot_hard_exhausted};
|
||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
|
||||
use serde_json::json;
|
||||
|
||||
fn sample_key_with_quota(code: &str, exhausted: bool) -> StoredProviderCatalogKey {
|
||||
let mut key = StoredProviderCatalogKey::new(
|
||||
"key-windsurf".to_string(),
|
||||
"provider-windsurf".to_string(),
|
||||
"[email protected]".to_string(),
|
||||
"oauth".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("sample key should build");
|
||||
key.status_snapshot = Some(json!({
|
||||
"quota": {
|
||||
"provider_type": "windsurf",
|
||||
"code": code,
|
||||
"exhausted": exhausted,
|
||||
"windows": [{
|
||||
"code": "daily",
|
||||
"used_ratio": 0.0,
|
||||
"remaining_ratio": 1.0
|
||||
}]
|
||||
}
|
||||
}));
|
||||
key
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn windsurf_rate_limit_bucket_does_not_mark_quota_exhausted() {
|
||||
let bucket = json!({
|
||||
"rate_limit": {
|
||||
"limited": true,
|
||||
"retry_after_ms": 60_000u64
|
||||
},
|
||||
"daily_remaining_percent": 50.0,
|
||||
"weekly_remaining_percent": 50.0
|
||||
});
|
||||
let bucket = bucket.as_object().expect("bucket should be object");
|
||||
|
||||
assert!(!windsurf_quota_exhausted_from_bucket(bucket));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn windsurf_banned_and_quarantined_snapshot_codes_are_hard_exhausted() {
|
||||
for code in ["banned", "forbidden", "quarantined"] {
|
||||
let key = sample_key_with_quota(code, false);
|
||||
|
||||
assert!(
|
||||
windsurf_quota_snapshot_hard_exhausted(&key, "windsurf"),
|
||||
"{code} should be hard exhausted"
|
||||
);
|
||||
}
|
||||
|
||||
let cooldown_key = sample_key_with_quota("cooldown", false);
|
||||
assert!(!windsurf_quota_snapshot_hard_exhausted(
|
||||
&cooldown_key,
|
||||
"windsurf"
|
||||
));
|
||||
let rate_limited_key = sample_key_with_quota("rate_limited", false);
|
||||
assert!(!windsurf_quota_snapshot_hard_exhausted(
|
||||
&rate_limited_key,
|
||||
"windsurf"
|
||||
));
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user