feat(grok): add provider pool and transport support

This commit is contained in:
mayrain
2026-05-16 19:22:54 +08:00
parent e5f1f52759
commit cbfe1d378f
20 changed files with 2616 additions and 21 deletions

View File

@@ -10,9 +10,9 @@ pub use frame::{StreamFrame, StreamFramePayload, StreamFrameType};
pub use plan::{
ExecutionPlan, ExecutionTimeouts, ProxySnapshot, RequestBody, ResolvedTransportProfile,
EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER, EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER,
EXECUTION_REQUEST_HTTP1_ONLY_HEADER, TRANSPORT_BACKEND_HYPER_RUSTLS,
TRANSPORT_BACKEND_REQWEST_RUSTLS, TRANSPORT_HTTP_MODE_AUTO, TRANSPORT_HTTP_MODE_HTTP1_ONLY,
TRANSPORT_POOL_SCOPE_KEY,
EXECUTION_REQUEST_HTTP1_ONLY_HEADER, TRANSPORT_BACKEND_BROWSER_WREQ,
TRANSPORT_BACKEND_HYPER_RUSTLS, TRANSPORT_BACKEND_REQWEST_RUSTLS, TRANSPORT_HTTP_MODE_AUTO,
TRANSPORT_HTTP_MODE_HTTP1_ONLY, TRANSPORT_POOL_SCOPE_KEY,
};
pub use result::{ExecutionResult, ExecutionTelemetry, ResponseBody};
pub use usage::{ExecutionStreamTerminalSummary, StandardizedUsage};

View File

@@ -63,6 +63,7 @@ pub struct ProxySnapshot {
pub const TRANSPORT_BACKEND_REQWEST_RUSTLS: &str = "reqwest_rustls";
pub const TRANSPORT_BACKEND_HYPER_RUSTLS: &str = "hyper_rustls";
pub const TRANSPORT_BACKEND_BROWSER_WREQ: &str = "browser_wreq";
pub const TRANSPORT_HTTP_MODE_AUTO: &str = "auto";
pub const TRANSPORT_HTTP_MODE_HTTP1_ONLY: &str = "http1_only";
pub const TRANSPORT_POOL_SCOPE_KEY: &str = "key";

View File

@@ -316,6 +316,13 @@ fn key_auth_channel_matches(row: &StoredMinimalCandidateSelectionRow, api_format
"gemini_cli" | "antigravity" => {
auth_type == "oauth" && api_format == "gemini:generate_content"
}
"grok" => {
auth_type == "oauth"
&& matches!(
api_format.as_str(),
"openai:chat" | "openai:responses" | "claude:messages" | "openai:image"
)
}
"vertex_ai" => {
(auth_type == "api_key" && api_format == "gemini:generate_content")
|| (matches!(auth_type.as_str(), "service_account" | "vertex_ai")
@@ -419,6 +426,35 @@ mod tests {
assert_eq!(rows[0].provider_id, "provider-1");
}
#[tokio::test]
async fn includes_grok_oauth_rows_for_chat_models() {
let mut row = sample_row(
"provider-grok",
"openai:chat",
"grok-4.20-0309-non-reasoning",
10,
);
row.provider_type = "grok".to_string();
row.provider_name = "grok".to_string();
row.key_auth_type = "oauth".to_string();
row.key_api_formats = Some(vec![
"openai:chat".to_string(),
"openai:responses".to_string(),
"claude:messages".to_string(),
"openai:image".to_string(),
]);
let repository = InMemoryMinimalCandidateSelectionReadRepository::seed(vec![row]);
let rows = repository
.list_for_exact_api_format("openai:chat")
.await
.expect("list should succeed");
assert_eq!(rows.len(), 1);
assert_eq!(rows[0].provider_type, "grok");
assert_eq!(rows[0].global_model_name, "grok-4.20-0309-non-reasoning");
}
#[tokio::test]
async fn requested_model_filter_respects_endpoint_scoped_default_mapping() {
let mut selected = sample_row("provider-1", "openai:chat", "deepseek-v4-pro", 10);

View File

@@ -445,6 +445,13 @@ fn key_auth_channel_matches(row: &CandidateSelectionRow, api_format: &str) -> bo
"gemini_cli" | "antigravity" => {
auth_type == "oauth" && api_format == "gemini:generate_content"
}
"grok" => {
auth_type == "oauth"
&& matches!(
api_format.as_str(),
"openai:chat" | "openai:responses" | "claude:messages" | "openai:image"
)
}
"vertex_ai" => {
(auth_type == "api_key" && api_format == "gemini:generate_content")
|| (matches!(auth_type.as_str(), "service_account" | "vertex_ai")

View File

@@ -103,6 +103,11 @@ WHERE p.is_active = TRUE
)
)
)
OR (
LOWER(BTRIM(p.provider_type)) = 'grok'
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
AND LOWER($3) IN ('openai:chat', 'openai:responses', 'claude:messages', 'openai:image')
)
OR (
LOWER(BTRIM(p.provider_type)) IN ('gemini_cli', 'antigravity')
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
@@ -127,6 +132,7 @@ WHERE p.is_active = TRUE
'claude_code',
'codex',
'gemini_cli',
'grok',
'vertex_ai',
'antigravity',
'kiro'
@@ -286,6 +292,11 @@ WHERE p.is_active = TRUE
)
)
)
OR (
LOWER(BTRIM(p.provider_type)) = 'grok'
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
AND LOWER($4) IN ('openai:chat', 'openai:responses', 'claude:messages', 'openai:image')
)
OR (
LOWER(BTRIM(p.provider_type)) IN ('gemini_cli', 'antigravity')
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
@@ -310,6 +321,7 @@ WHERE p.is_active = TRUE
'claude_code',
'codex',
'gemini_cli',
'grok',
'vertex_ai',
'antigravity',
'kiro'
@@ -468,6 +480,11 @@ WHERE p.is_active = TRUE
)
)
)
OR (
LOWER(BTRIM(p.provider_type)) = 'grok'
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
AND LOWER($6) IN ('openai:chat', 'openai:responses', 'claude:messages', 'openai:image')
)
OR (
LOWER(BTRIM(p.provider_type)) IN ('gemini_cli', 'antigravity')
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
@@ -492,6 +509,7 @@ WHERE p.is_active = TRUE
'claude_code',
'codex',
'gemini_cli',
'grok',
'vertex_ai',
'antigravity',
'kiro'
@@ -1287,6 +1305,23 @@ mod tests {
}
}
#[test]
fn candidate_selection_sql_allows_grok_oauth_chat_auth() {
let requested_model_sql = requested_model_selection_sql();
for sql in [
LIST_FOR_EXACT_API_FORMAT_SQL,
LIST_FOR_EXACT_API_FORMAT_AND_GLOBAL_MODEL_SQL,
LIST_POOL_KEYS_FOR_GROUP_SQL,
requested_model_sql.as_str(),
] {
assert!(sql.contains("LOWER(BTRIM(p.provider_type)) = 'grok'"));
assert!(sql.contains("LOWER(BTRIM(pak.auth_type)) = 'oauth'"));
assert!(sql
.contains("'openai:chat', 'openai:responses', 'claude:messages', 'openai:image'"));
assert!(sql.contains("'grok',"));
}
}
#[test]
fn requested_model_selection_page_sql_adds_limit_and_offset() {
let sql = requested_model_selection_page_sql();

View File

@@ -445,6 +445,13 @@ fn key_auth_channel_matches(row: &CandidateSelectionRow, api_format: &str) -> bo
"gemini_cli" | "antigravity" => {
auth_type == "oauth" && api_format == "gemini:generate_content"
}
"grok" => {
auth_type == "oauth"
&& matches!(
api_format.as_str(),
"openai:chat" | "openai:responses" | "claude:messages" | "openai:image"
)
}
"vertex_ai" => {
(auth_type == "api_key" && api_format == "gemini:generate_content")
|| (matches!(auth_type.as_str(), "service_account" | "vertex_ai")

View File

@@ -229,7 +229,7 @@ pub fn endpoint_supports_rust_models_fetch(api_format: &str) -> bool {
pub fn provider_type_uses_preset_models(provider_type: &str) -> bool {
matches!(
provider_type.trim().to_ascii_lowercase().as_str(),
"claude_code" | "gemini_cli"
"claude_code" | "gemini_cli" | "grok"
)
}
@@ -272,6 +272,27 @@ pub fn preset_models_for_provider(provider_type: &str) -> Option<Vec<Value>> {
preset_model("gpt-5.3-codex", "openai", "GPT-5.3 Codex", "openai:responses"),
preset_model("gpt-5.3-codex-spark", "openai", "GPT-5.3 Codex Spark", "openai:responses"),
],
"grok" => vec![
preset_model("grok-4.20-0309-non-reasoning", "xai", "Grok 4.20 0309 Non-Reasoning", "openai:chat"),
preset_model("grok-4.20-0309", "xai", "Grok 4.20 0309", "openai:chat"),
preset_model("grok-4.20-0309-reasoning", "xai", "Grok 4.20 0309 Reasoning", "openai:chat"),
preset_model("grok-4.20-0309-non-reasoning-super", "xai", "Grok 4.20 0309 Non-Reasoning Super", "openai:chat"),
preset_model("grok-4.20-0309-super", "xai", "Grok 4.20 0309 Super", "openai:chat"),
preset_model("grok-4.20-0309-reasoning-super", "xai", "Grok 4.20 0309 Reasoning Super", "openai:chat"),
preset_model("grok-4.20-0309-non-reasoning-heavy", "xai", "Grok 4.20 0309 Non-Reasoning Heavy", "openai:chat"),
preset_model("grok-4.20-0309-heavy", "xai", "Grok 4.20 0309 Heavy", "openai:chat"),
preset_model("grok-4.20-0309-reasoning-heavy", "xai", "Grok 4.20 0309 Reasoning Heavy", "openai:chat"),
preset_model("grok-4.20-multi-agent-0309", "xai", "Grok 4.20 Multi-Agent 0309", "openai:chat"),
preset_model("grok-4.20-auto", "xai", "Grok 4.20 Auto", "openai:chat"),
preset_model("grok-4.20-fast", "xai", "Grok 4.20 Fast", "openai:chat"),
preset_model("grok-4.20-expert", "xai", "Grok 4.20 Expert", "openai:chat"),
preset_model("grok-4.20-heavy", "xai", "Grok 4.20 Heavy", "openai:chat"),
preset_model("grok-4.3-beta", "xai", "Grok 4.3 Beta", "openai:chat"),
preset_model("grok-imagine-image-lite", "xai", "Grok Imagine Image Lite", "openai:image"),
preset_model("grok-imagine-image", "xai", "Grok Imagine Image", "openai:image"),
preset_model("grok-imagine-image-pro", "xai", "Grok Imagine Image Pro", "openai:image"),
preset_model("grok-imagine-image-edit", "xai", "Grok Imagine Image Edit", "openai:image"),
],
_ => return None,
};
Some(models)
@@ -934,4 +955,42 @@ mod tests {
.iter()
.all(|model| model["api_formats"] == json!(["claude:messages"])));
}
#[test]
fn preset_models_cover_grok_non_video_catalog() {
let models = preset_models_for_provider("grok").expect("preset models should exist");
let model_ids = models
.iter()
.map(|model| model["id"].as_str().expect("model id"))
.collect::<Vec<_>>();
assert_eq!(
model_ids,
vec![
"grok-4.20-0309-non-reasoning",
"grok-4.20-0309",
"grok-4.20-0309-reasoning",
"grok-4.20-0309-non-reasoning-super",
"grok-4.20-0309-super",
"grok-4.20-0309-reasoning-super",
"grok-4.20-0309-non-reasoning-heavy",
"grok-4.20-0309-heavy",
"grok-4.20-0309-reasoning-heavy",
"grok-4.20-multi-agent-0309",
"grok-4.20-auto",
"grok-4.20-fast",
"grok-4.20-expert",
"grok-4.20-heavy",
"grok-4.3-beta",
"grok-imagine-image-lite",
"grok-imagine-image",
"grok-imagine-image-pro",
"grok-imagine-image-edit",
]
);
assert!(!model_ids.contains(&"grok-imagine-video"));
assert_eq!(models[0]["api_formats"], json!(["openai:chat"]));
assert_eq!(models[10]["api_formats"], json!(["openai:chat"]));
assert_eq!(models[15]["api_formats"], json!(["openai:image"]));
assert_eq!(models[18]["api_formats"], json!(["openai:image"]));
}
}

View File

@@ -17,9 +17,11 @@ pub use provider::{ProviderPoolAdapter, ProviderPoolMemberInput};
pub use providers::{
build_antigravity_pool_quota_request, build_chatgpt_web_pool_quota_request,
build_codex_pool_quota_request, build_kiro_pool_quota_request,
enrich_chatgpt_web_quota_metadata, normalize_chatgpt_web_image_quota_limit,
AntigravityProviderPoolAdapter, ChatGptWebProviderPoolAdapter, CodexProviderPoolAdapter,
DefaultProviderPoolAdapter, KiroPoolQuotaAuthInput, KiroProviderPoolAdapter,
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,
GrokProviderPoolAdapter, KiroPoolQuotaAuthInput, KiroProviderPoolAdapter,
UnsupportedQuotaProviderPoolAdapter, ANTIGRAVITY_FETCH_AVAILABLE_MODELS_PATH,
CHATGPT_WEB_CONVERSATION_INIT_PATH, CHATGPT_WEB_DEFAULT_BASE_URL, CODEX_WHAM_USAGE_URL,
KIRO_USAGE_LIMITS_PATH, KIRO_USAGE_SDK_VERSION,
@@ -65,6 +67,7 @@ mod tests {
"claude_code",
"codex",
"gemini_cli",
"grok",
"kiro",
"vertex_ai"
]
@@ -82,10 +85,11 @@ mod tests {
assert_eq!(
service.provider_types_for_capability(ProviderPoolCapability::QuotaRefresh),
["antigravity", "chatgpt_web", "codex", "kiro"]
["antigravity", "chatgpt_web", "codex", "grok", "kiro"]
);
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_eq!(
service.quota_refresh_unsupported_message("claude_code"),
@@ -247,8 +251,11 @@ mod tests {
.find(|item| item["name"] == "recent_refresh")
.expect("recent_refresh should exist");
assert_eq!(free_first["providers"], json!(["codex", "kiro"]));
assert_eq!(recent_refresh["providers"], json!(["codex", "kiro"]));
assert_eq!(free_first["providers"], json!(["codex", "grok", "kiro"]));
assert_eq!(
recent_refresh["providers"],
json!(["codex", "grok", "kiro"])
);
}
#[test]
@@ -331,6 +338,37 @@ mod tests {
}))),
"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": {
@@ -342,6 +380,68 @@ mod tests {
));
}
#[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!({

View File

@@ -0,0 +1,220 @@
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_json_bool, provider_pool_json_f64, provider_pool_metadata_bucket,
provider_pool_quota_snapshot_exhausted_decision,
};
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;
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;
if 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)
{
exhausted_count += 1;
}
}
model_count > 0 && model_count == exhausted_count
}

View File

@@ -2,6 +2,7 @@ pub mod antigravity;
pub mod chatgpt_web;
pub mod codex;
pub mod default;
pub mod grok;
pub mod kiro;
pub mod unsupported;
@@ -18,6 +19,10 @@ pub use chatgpt_web::{
pub use codex::CodexProviderPoolAdapter;
pub use codex::{build_codex_pool_quota_request, CODEX_WHAM_USAGE_URL};
pub use default::DefaultProviderPoolAdapter;
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,

View File

@@ -12,8 +12,9 @@ use crate::presets::normalize_provider_scheduling_presets;
use crate::provider::{ProviderPoolAdapter, ProviderPoolMemberInput};
use crate::providers::{
AntigravityProviderPoolAdapter, ChatGptWebProviderPoolAdapter, CodexProviderPoolAdapter,
DefaultProviderPoolAdapter, KiroProviderPoolAdapter, CLAUDE_CODE_PROVIDER_POOL_ADAPTER,
GEMINI_CLI_PROVIDER_POOL_ADAPTER, VERTEX_AI_PROVIDER_POOL_ADAPTER,
DefaultProviderPoolAdapter, GrokProviderPoolAdapter, KiroProviderPoolAdapter,
CLAUDE_CODE_PROVIDER_POOL_ADAPTER, GEMINI_CLI_PROVIDER_POOL_ADAPTER,
VERTEX_AI_PROVIDER_POOL_ADAPTER,
};
#[derive(Clone)]
@@ -50,6 +51,7 @@ impl ProviderPoolService {
.with_adapter(Arc::new(CLAUDE_CODE_PROVIDER_POOL_ADAPTER))
.with_adapter(Arc::new(CodexProviderPoolAdapter))
.with_adapter(Arc::new(GEMINI_CLI_PROVIDER_POOL_ADAPTER))
.with_adapter(Arc::new(GrokProviderPoolAdapter))
.with_adapter(Arc::new(KiroProviderPoolAdapter))
.with_adapter(Arc::new(ChatGptWebProviderPoolAdapter))
.with_adapter(Arc::new(VERTEX_AI_PROVIDER_POOL_ADAPTER))

View File

@@ -8,6 +8,7 @@ use crate::conversion::{
request_conversion_enabled_for_transport, request_conversion_transport_unsupported_reason,
request_pair_allowed_for_transport,
};
use crate::grok::grok_browser_resolved_transport_profile_from_auth_config;
use crate::network::{
resolve_transport_profile, resolve_transport_profile_id, transport_proxy_is_locally_supported,
};
@@ -89,6 +90,31 @@ pub fn build_transport_diagnostics(
.and_then(|value| value.get("transport_profile"))
.cloned()
.unwrap_or(Value::Null);
let configured_legacy_grok_transport_profile = if transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("grok")
{
transport
.key
.decrypted_auth_config
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.and_then(|value| serde_json::from_str::<Value>(value).ok())
.and_then(|value| value.as_object().cloned())
.and_then(|auth_config| {
grok_browser_resolved_transport_profile_from_auth_config(
&auth_config,
"grok_auth_config",
)
.and_then(|profile| serde_json::to_value(profile).ok())
})
.unwrap_or(Value::Null)
} else {
Value::Null
};
let has_oauth_config = transport.key.decrypted_auth_config.is_some();
let oauth_resolution_supported =
!has_oauth_config || crate::supports_local_oauth_request_auth_resolution(transport);
@@ -126,6 +152,7 @@ pub fn build_transport_diagnostics(
"fingerprint": transport.key.fingerprint,
"configured_key_transport_profile": configured_key_transport_profile,
"configured_provider_transport_profile": configured_provider_transport_profile,
"configured_legacy_grok_transport_profile": configured_legacy_grok_transport_profile,
"resolved_transport_profile_id": resolved_transport_profile_id,
"resolved_transport_profile": resolved_transport_profile,
"request_pair": {
@@ -423,6 +450,43 @@ mod tests {
);
}
fn sample_grok_transport_with_legacy_user_agent() -> GatewayProviderTransportSnapshot {
let mut transport = sample_transport();
transport.provider.provider_type = "grok".to_string();
transport.key.fingerprint = None;
transport.provider.config = None;
transport.key.decrypted_auth_config = Some(
json!({
"sso_token": "sso-token",
"user_agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/137.0.0.0 Safari/537.36"
})
.to_string(),
);
transport
}
#[test]
fn transport_diagnostics_include_legacy_grok_transport_profile() {
let diagnostics = build_transport_diagnostics(
&sample_grok_transport_with_legacy_user_agent(),
"openai:chat",
"openai:chat",
);
assert_eq!(
diagnostics["configured_legacy_grok_transport_profile"]["profile_id"],
"chrome137"
);
assert_eq!(
diagnostics["resolved_transport_profile"]["profile_id"],
"chrome137"
);
assert_eq!(
diagnostics["resolved_transport_profile"]["backend"],
"browser_wreq"
);
}
#[test]
fn request_trace_proxy_value_sanitizes_url_and_marks_config_source() {
let transport = sample_transport();

File diff suppressed because it is too large Load Diff

View File

@@ -7,6 +7,7 @@ pub mod conversion;
mod diagnostics;
mod gemini_files;
mod generic_oauth;
pub mod grok;
mod headers;
pub mod kiro;
mod network;
@@ -44,6 +45,17 @@ pub use gemini_files::{
pub use generic_oauth::{
supports_local_generic_oauth_request_auth_resolution, GenericOAuthRefreshAdapter,
};
pub use grok::{
build_grok_app_chat_body, build_grok_browser_headers, build_grok_upstream_url, grok_base_url,
grok_browser_profile_id_from_user_agent,
grok_browser_profile_metadata_from_resolved_transport_profile,
grok_browser_resolved_transport_profile,
grok_browser_resolved_transport_profile_from_auth_config,
grok_browser_transport_fingerprint_from_auth_config, is_grok_provider_transport,
resolve_grok_session_auth, GrokBrowserProfileMetadata, GrokHeaderInput, GROK_CHAT_PATH,
GROK_DEFAULT_BASE_URL, GROK_DEFAULT_BROWSER_PROFILE, GROK_DEFAULT_USER_AGENT,
GROK_INTERNAL_HEADER, GROK_RATE_LIMITS_PATH,
};
pub use headers::{should_skip_request_header, should_skip_upstream_passthrough_header};
pub use network::{
resolve_transport_execution_timeouts, resolve_transport_profile, resolve_transport_profile_id,

View File

@@ -6,6 +6,8 @@ use async_trait::async_trait;
use serde_json::{json, Map, Value};
use tracing::warn;
use crate::grok::grok_browser_resolved_transport_profile_from_auth_config;
use super::snapshot::GatewayProviderTransportSnapshot;
const TUNNEL_BASE_URL_EXTRA_KEY: &str = "tunnel_base_url";
@@ -152,9 +154,39 @@ pub fn resolve_transport_profile(
) -> Option<ResolvedTransportProfile> {
resolve_transport_profile_from_fingerprint(transport.key.fingerprint.as_ref()).or_else(|| {
resolve_transport_profile_from_provider_config(transport.provider.config.as_ref())
.or_else(|| resolve_grok_browser_transport_profile(transport))
})
}
fn resolve_grok_browser_transport_profile(
transport: &GatewayProviderTransportSnapshot,
) -> Option<ResolvedTransportProfile> {
if !transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("grok")
{
return None;
}
let auth_config = transport
.key
.decrypted_auth_config
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.and_then(|value| serde_json::from_str::<Value>(value).ok())?;
let object = auth_config.as_object()?;
let has_session = json_string_field(object, "sso_token")
.or_else(|| json_string_field(object, "access_token"))
.or_else(|| json_string_field(object, "token"))
.is_some();
if !has_session {
return None;
}
grok_browser_resolved_transport_profile_from_auth_config(object, "grok_auth_config")
}
fn resolve_transport_profile_from_provider_config(
config: Option<&Value>,
) -> Option<ResolvedTransportProfile> {
@@ -529,4 +561,218 @@ mod tests {
assert!(resolve_transport_profile(&transport).is_none());
assert!(!transport_profile_is_configured(&transport));
}
#[test]
fn resolves_grok_browser_transport_profile_from_session_auth_config() {
let mut transport = sample_transport();
transport.provider.provider_type = "grok".to_string();
transport.key.fingerprint = None;
transport.provider.config = None;
transport.key.decrypted_auth_config = Some(
json!({
"sso_token": "sso-token",
"browser_profile": "chrome136",
"cf_clearance": "clearance"
})
.to_string(),
);
let profile = resolve_transport_profile(&transport).expect("profile");
assert_eq!(profile.profile_id, "chrome136");
assert_eq!(profile.backend, "browser_wreq");
assert_eq!(profile.http_mode, "auto");
assert_eq!(profile.pool_scope, "key");
assert_eq!(
profile
.extra
.as_ref()
.and_then(|value| value.get("browser_profile"))
.and_then(Value::as_str),
Some("chrome136")
);
}
#[test]
fn resolves_grok_browser_transport_profile_default_from_session_auth_config() {
let mut transport = sample_transport();
transport.provider.provider_type = "grok".to_string();
transport.key.fingerprint = None;
transport.provider.config = None;
transport.key.decrypted_auth_config = Some(
json!({
"sso_token": "sso-token"
})
.to_string(),
);
let profile = resolve_transport_profile(&transport).expect("profile");
assert_eq!(profile.profile_id, "chrome136");
assert_eq!(profile.backend, "browser_wreq");
assert_eq!(
profile
.extra
.as_ref()
.and_then(|value| value.get("source"))
.and_then(Value::as_str),
Some("grok_auth_config")
);
}
#[test]
fn resolves_grok_browser_transport_profile_normalizes_auth_config_alias() {
let mut transport = sample_transport();
transport.provider.provider_type = "grok".to_string();
transport.key.fingerprint = None;
transport.provider.config = None;
transport.key.decrypted_auth_config = Some(
json!({
"sso_token": "sso-token",
"browser_profile": "Chrome-137"
})
.to_string(),
);
let profile = resolve_transport_profile(&transport).expect("profile");
assert_eq!(profile.profile_id, "chrome137");
assert_eq!(
profile
.extra
.as_ref()
.and_then(|value| value.get("browser_profile"))
.and_then(Value::as_str),
Some("chrome137")
);
}
#[test]
fn resolves_grok_browser_transport_profile_from_legacy_user_agent() {
let mut transport = sample_transport();
transport.provider.provider_type = "grok".to_string();
transport.key.fingerprint = None;
transport.provider.config = None;
transport.key.decrypted_auth_config = Some(
json!({
"sso_token": "sso-token",
"user_agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/137.0.0.0 Safari/537.36"
})
.to_string(),
);
let profile = resolve_transport_profile(&transport).expect("profile");
assert_eq!(profile.profile_id, "chrome137");
assert_eq!(
profile
.extra
.as_ref()
.and_then(|value| value.get("browser_profile"))
.and_then(Value::as_str),
Some("chrome137")
);
}
#[test]
fn key_fingerprint_wins_over_grok_auth_config_fallback() {
let mut transport = sample_transport();
transport.provider.provider_type = "grok".to_string();
transport.provider.config = None;
transport.key.fingerprint = Some(json!({
"transport_profile": {
"profile_id": "chrome136",
"backend": "browser_wreq",
"extra": {"browser_profile": "chrome136", "source": "key"}
}
}));
transport.key.decrypted_auth_config = Some(
json!({
"sso_token": "sso-token",
"browser_profile": "chrome137"
})
.to_string(),
);
let profile = resolve_transport_profile(&transport).expect("profile");
assert_eq!(profile.profile_id, "chrome136");
assert_eq!(
profile
.extra
.as_ref()
.and_then(|value| value.get("source"))
.and_then(Value::as_str),
Some("key")
);
}
#[test]
fn provider_fingerprint_wins_over_grok_auth_config_fallback() {
let mut transport = sample_transport();
transport.provider.provider_type = "grok".to_string();
transport.key.fingerprint = None;
transport.provider.config = Some(json!({
"fingerprint": {
"transport_profile": {
"profile_id": "chrome136",
"backend": "browser_wreq",
"extra": {"browser_profile": "chrome136", "source": "provider"}
}
}
}));
transport.key.decrypted_auth_config = Some(
json!({
"sso_token": "sso-token",
"browser_profile": "chrome137"
})
.to_string(),
);
let profile = resolve_transport_profile(&transport).expect("profile");
assert_eq!(profile.profile_id, "chrome136");
assert_eq!(
profile
.extra
.as_ref()
.and_then(|value| value.get("source"))
.and_then(Value::as_str),
Some("provider")
);
}
#[test]
fn rejects_unsupported_grok_auth_config_browser_profile() {
let mut transport = sample_transport();
transport.provider.provider_type = "grok".to_string();
transport.key.fingerprint = None;
transport.provider.config = None;
transport.key.decrypted_auth_config = Some(
json!({
"sso_token": "sso-token",
"browser_profile": "safari999"
})
.to_string(),
);
assert!(resolve_transport_profile(&transport).is_none());
}
#[test]
fn rejects_unsupported_grok_auth_config_user_agent_profile() {
let mut transport = sample_transport();
transport.provider.provider_type = "grok".to_string();
transport.key.fingerprint = None;
transport.provider.config = None;
transport.key.decrypted_auth_config = Some(
json!({
"sso_token": "sso-token",
"user_agent": "Mozilla/5.0 Version/18.0 Safari/605.1.15"
})
.to_string(),
);
assert!(resolve_transport_profile(&transport).is_none());
}
}

View File

@@ -3,6 +3,7 @@ use std::collections::BTreeMap;
use serde_json::Value;
use crate::auth::{build_passthrough_headers_with_auth, resolve_local_openai_bearer_auth};
use crate::grok::{is_grok_provider_transport, resolve_grok_session_auth};
use crate::policy::local_standard_transport_unsupported_reason_with_network;
use crate::rules::apply_local_header_rules_with_request_headers;
use crate::snapshot::GatewayProviderTransportSnapshot;
@@ -23,21 +24,33 @@ pub fn openai_image_transport_unsupported_reason(
api_format: &str,
) -> Option<&'static str> {
let reason = local_standard_transport_unsupported_reason_with_network(transport, api_format);
if reason == Some("transport_provider_type_unsupported")
&& transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("chatgpt_web")
if is_dedicated_openai_image_provider(transport)
&& matches!(
reason,
Some("transport_provider_type_unsupported")
| Some("transport_oauth_resolution_unsupported")
)
{
return None;
}
reason
}
fn is_dedicated_openai_image_provider(transport: &GatewayProviderTransportSnapshot) -> bool {
transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("chatgpt_web")
|| is_grok_provider_transport(transport)
}
pub fn resolve_openai_image_auth(
transport: &GatewayProviderTransportSnapshot,
) -> Option<(String, String)> {
if is_grok_provider_transport(transport) {
return resolve_grok_session_auth(transport);
}
resolve_local_openai_bearer_auth(transport)
}
@@ -159,6 +172,45 @@ mod tests {
);
}
#[test]
fn grok_is_supported_by_dedicated_openai_image_transport_policy() {
let mut transport = sample_transport();
transport.provider.provider_type = "grok".to_string();
assert_eq!(
openai_image_transport_unsupported_reason(&transport, "openai:image"),
None
);
}
#[test]
fn grok_oauth_session_is_supported_by_dedicated_openai_image_transport_policy() {
let mut transport = sample_transport();
transport.provider.provider_type = "grok".to_string();
transport.key.auth_type = "oauth".to_string();
transport.key.decrypted_api_key = String::new();
transport.key.decrypted_auth_config = Some(json!({"sso_token":"abc"}).to_string());
assert_eq!(
openai_image_transport_unsupported_reason(&transport, "openai:image"),
None
);
}
#[test]
fn chatgpt_web_oauth_is_supported_by_dedicated_openai_image_transport_policy() {
let mut transport = sample_transport();
transport.provider.provider_type = "chatgpt_web".to_string();
transport.key.auth_type = "oauth".to_string();
transport.key.decrypted_api_key = String::new();
transport.key.decrypted_auth_config = Some(json!({"access_token":"token"}).to_string());
assert_eq!(
openai_image_transport_unsupported_reason(&transport, "openai:image"),
None
);
}
#[test]
fn builds_json_eventstream_headers_and_applies_rules() {
let headers = build_openai_image_headers(ProviderOpenAiImageHeadersInput {

View File

@@ -238,6 +238,15 @@ const ANTIGRAVITY_RUNTIME_POLICY: ProviderRuntimePolicy = ProviderRuntimePolicy
supports_local_same_format_transport: false,
..STANDARD_RUNTIME_POLICY
};
const GROK_RUNTIME_POLICY: ProviderRuntimePolicy = ProviderRuntimePolicy {
fixed_provider: true,
api_format_inheritance: ProviderApiFormatInheritance::OAuth,
enable_format_conversion_by_default: true,
supports_model_fetch: false,
supports_local_openai_chat_transport: false,
supports_local_same_format_transport: false,
..STANDARD_RUNTIME_POLICY
};
const CLAUDE_CODE_FIXED_PROVIDER_TEMPLATE: FixedProviderTemplate = FixedProviderTemplate {
provider_type: "claude_code",
@@ -352,6 +361,39 @@ const ANTIGRAVITY_FIXED_PROVIDER_TEMPLATE: FixedProviderTemplate = FixedProvider
runtime_policy: ANTIGRAVITY_RUNTIME_POLICY,
};
const GROK_FIXED_PROVIDER_TEMPLATE: FixedProviderTemplate = FixedProviderTemplate {
provider_type: "grok",
version: 1,
base_url: "https://grok.com",
endpoints: &[
FixedProviderEndpointTemplate {
item_key: "openai:chat",
api_format: "openai:chat",
custom_path: None,
config_defaults: EMPTY_ENDPOINT_CONFIG_DEFAULTS,
},
FixedProviderEndpointTemplate {
item_key: "openai:responses",
api_format: "openai:responses",
custom_path: None,
config_defaults: EMPTY_ENDPOINT_CONFIG_DEFAULTS,
},
FixedProviderEndpointTemplate {
item_key: "claude:messages",
api_format: "claude:messages",
custom_path: None,
config_defaults: EMPTY_ENDPOINT_CONFIG_DEFAULTS,
},
FixedProviderEndpointTemplate {
item_key: "openai:image",
api_format: "openai:image",
custom_path: None,
config_defaults: EMPTY_ENDPOINT_CONFIG_DEFAULTS,
},
],
runtime_policy: GROK_RUNTIME_POLICY,
};
pub fn provider_type_is_fixed(provider_type: &str) -> bool {
provider_runtime_policy(provider_type).fixed_provider
}
@@ -398,6 +440,7 @@ pub fn fixed_provider_template(provider_type: &str) -> Option<&'static FixedProv
"codex" => Some(&CODEX_FIXED_PROVIDER_TEMPLATE),
"chatgpt_web" => Some(&CHATGPT_WEB_FIXED_PROVIDER_TEMPLATE),
"kiro" => Some(&KIRO_FIXED_PROVIDER_TEMPLATE),
"grok" => Some(&GROK_FIXED_PROVIDER_TEMPLATE),
"gemini_cli" => Some(&GEMINI_CLI_FIXED_PROVIDER_TEMPLATE),
"vertex_ai" => Some(&VERTEX_AI_FIXED_PROVIDER_TEMPLATE),
"antigravity" => Some(&ANTIGRAVITY_FIXED_PROVIDER_TEMPLATE),
@@ -598,6 +641,29 @@ mod tests {
);
}
#[test]
fn grok_fixed_provider_template_exposes_chat_responses_messages_and_image() {
let template = fixed_provider_template("grok").expect("grok template should exist");
assert_eq!(template.base_url, "https://grok.com");
assert_eq!(template.version, 1);
assert_eq!(
template
.endpoints
.iter()
.map(|item| item.api_format)
.collect::<Vec<_>>(),
vec![
"openai:chat",
"openai:responses",
"claude:messages",
"openai:image"
]
);
assert!(!template.runtime_policy.supports_model_fetch);
assert!(!template.runtime_policy.supports_local_openai_chat_transport);
assert!(!template.runtime_policy.supports_local_same_format_transport);
}
#[test]
fn fixed_provider_key_inheritance_keeps_oauth_and_kiro_configured_bearer_keys_open() {
assert!(fixed_provider_key_inherits_api_formats(

View File

@@ -9,6 +9,7 @@ use crate::auth::{
};
use crate::claude_code::build_claude_code_passthrough_headers;
use crate::claude_code::local_claude_code_transport_unsupported_reason_with_network;
use crate::grok::{is_grok_provider_transport, resolve_grok_session_auth};
use crate::kiro::{
build_kiro_provider_headers, build_kiro_provider_request_body, is_kiro_provider_transport,
local_kiro_request_transport_unsupported_reason_with_network, KiroAuthConfig,
@@ -351,6 +352,10 @@ pub fn same_format_provider_transport_unsupported_reason(
family: SameFormatProviderFamily,
api_format: &str,
) -> Option<&'static str> {
if is_grok_provider_transport(transport) && matches!(family, SameFormatProviderFamily::Standard)
{
return None;
}
if behavior.is_kiro {
local_kiro_request_transport_unsupported_reason_with_network(transport)
} else if behavior.is_antigravity {
@@ -434,6 +439,10 @@ pub fn resolve_same_format_provider_direct_auth(
transport: &GatewayProviderTransportSnapshot,
family: SameFormatProviderFamily,
) -> Option<(String, String)> {
if is_grok_provider_transport(transport) && matches!(family, SameFormatProviderFamily::Standard)
{
return resolve_grok_session_auth(transport);
}
if behavior.is_vertex {
None
} else {