i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-09 02:47:45 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(grok): add admin oauth and quota support

Browse Source
This commit is contained in:
mayrain
2026-05-16 21:15:38 +08:00
parent cbfe1d378f
commit 936e1ae37b
61 changed files with 4748 additions and 287 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/handlers/admin/provider/query/models/mod.rs
+341 -5
View File
@@ -17,6 +17,9 @@ use crate::ai_serving::{
};
use crate::clock::current_unix_ms;
use crate::execution_runtime;
use crate::handlers::admin::provider::shared::model_test_capabilities::{
admin_provider_model_supports_image_generation, admin_provider_model_test_capabilities_payload,
};
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
use crate::handlers::shared::provider_pool::{
admin_provider_pool_config_from_config_value, read_admin_provider_pool_runtime_state,
@@ -100,6 +103,181 @@ struct ProviderQueryKeyFetchResult {
has_success: bool,
}
fn provider_query_model_id(model: &Value) -> Option<&str> {
model
.get("id")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
}
fn provider_query_grok_required_tier_rank(model_id: &str) -> Option<u8> {
match model_id.trim() {
"grok-4.20-0309-non-reasoning" | "grok-4.20-fast" | "grok-imagine-image-lite" => Some(0),
"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-auto"
| "grok-4.20-expert"
| "grok-4.3-beta"
| "grok-imagine-image"
| "grok-imagine-image-pro"
| "grok-imagine-image-edit" => Some(1),
"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-heavy" => Some(2),
_ => None,
}
}
fn provider_query_normalize_grok_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 provider_query_grok_pool_tier_rank(value: Option<&str>) -> u8 {
match provider_query_normalize_grok_pool_tier(value).unwrap_or("basic") {
"heavy" => 2,
"super" => 1,
_ => 0,
}
}
fn provider_query_grok_quota_string(quota: &Map<String, Value>, fields: &[&str]) -> Option<String> {
fields.iter().find_map(|field| {
quota
.get(*field)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
})
}
fn provider_query_grok_window_limit(quota: &Map<String, Value>, model_name: &str) -> Option<f64> {
quota
.get("windows")
.and_then(Value::as_array)?
.iter()
.filter_map(Value::as_object)
.find(|window| {
window
.get("model")
.and_then(Value::as_str)
.is_some_and(|value| value.trim() == model_name)
})
.and_then(|window| window.get("limit_value"))
.and_then(Value::as_f64)
.filter(|value| value.is_finite() && *value > 0.0)
}
fn provider_query_grok_pool_tier_from_quota(quota: &Map<String, Value>) -> Option<&'static str> {
if let Some(tier) =
provider_query_grok_quota_string(quota, &["pool_tier", "tier", "plan_type", "plan"])
.and_then(|value| provider_query_normalize_grok_pool_tier(Some(&value)))
{
return Some(tier);
}
if let Some(auto_total) = provider_query_grok_window_limit(quota, "quota_auto") {
if (auto_total - 150.0).abs() < f64::EPSILON {
return Some("heavy");
}
if (auto_total - 50.0).abs() < f64::EPSILON {
return Some("super");
}
}
if let Some(fast_total) = provider_query_grok_window_limit(quota, "quota_fast") {
if (fast_total - 400.0).abs() < f64::EPSILON {
return Some("heavy");
}
if (fast_total - 140.0).abs() < f64::EPSILON {
return Some("super");
}
if (fast_total - 30.0).abs() < f64::EPSILON {
return Some("basic");
}
}
None
}
fn provider_query_grok_key_pool_tier(key: &StoredProviderCatalogKey) -> Option<&'static str> {
key.status_snapshot
.as_ref()
.and_then(Value::as_object)
.and_then(|snapshot| snapshot.get("quota"))
.and_then(Value::as_object)
.and_then(provider_query_grok_pool_tier_from_quota)
}
fn provider_query_filter_models_for_key(
provider: &StoredProviderCatalogProvider,
key: &StoredProviderCatalogKey,
models: Vec<Value>,
) -> Vec<Value> {
if !provider.provider_type.trim().eq_ignore_ascii_case("grok") {
return models;
}
let allowed_rank = provider_query_grok_pool_tier_rank(provider_query_grok_key_pool_tier(key));
models
.into_iter()
.filter(|model| {
provider_query_model_id(model)
.and_then(provider_query_grok_required_tier_rank)
.is_some_and(|required_rank| required_rank <= allowed_rank)
})
.collect()
}
fn provider_query_attach_model_test_capabilities(
provider: &StoredProviderCatalogProvider,
models: Vec<Value>,
) -> Vec<Value> {
models
.into_iter()
.map(|mut model| {
let Some(object) = model.as_object_mut() else {
return model;
};
let model_id = object
.get("id")
.and_then(Value::as_str)
.map(str::trim)
.unwrap_or_default()
.to_string();
let supports_image_generation = admin_provider_model_supports_image_generation(
&provider.provider_type,
&model_id,
object
.get("supports_image_generation")
.or_else(|| object.get("effective_supports_image_generation"))
.and_then(Value::as_bool)
.unwrap_or(false),
);
object.insert(
"model_test_capabilities".to_string(),
admin_provider_model_test_capabilities_payload(
&provider.provider_type,
&model_id,
supports_image_generation,
),
);
model
})
.collect()
}
fn provider_query_codex_preset_fallback(
provider: &StoredProviderCatalogProvider,
) -> Option<ProviderQueryKeyFetchResult> {
@@ -245,8 +423,9 @@ async fn provider_query_fetch_models_for_key(
if let Some(cached_models) =
provider_query_read_cached_models(state, &provider.id, &key.id).await
{
let models = provider_query_filter_models_for_key(provider, key, cached_models);
return Ok(ProviderQueryKeyFetchResult {
models: cached_models,
models,
error: None,
from_cache: true,
has_success: true,
@@ -257,8 +436,13 @@ async fn provider_query_fetch_models_for_key(
let selected_endpoints = selected_models_fetch_endpoints(endpoints, key);
if selected_endpoints.is_empty() {
if let Some(models) = preset_models_for_provider(&provider.provider_type) {
let models = provider_query_filter_models_for_key(
provider,
key,
aggregate_models_for_cache(&models),
);
return Ok(ProviderQueryKeyFetchResult {
models: aggregate_models_for_cache(&models),
models,
error: None,
from_cache: false,
has_success: true,
@@ -342,7 +526,7 @@ async fn provider_query_fetch_models_for_key(
}
Ok(ProviderQueryKeyFetchResult {
models: unique_models,
models: provider_query_filter_models_for_key(provider, key, unique_models),
error,
from_cache: false,
has_success: outcome.has_success,
@@ -397,11 +581,12 @@ pub(crate) async fn build_admin_provider_query_models_response(
force_refresh,
)
.await?;
let success = !result.models.is_empty();
let models = provider_query_attach_model_test_capabilities(&provider, result.models);
let success = !models.is_empty();
return Ok(Json(json!({
"success": success,
"data": {
"models": result.models,
"models": models,
"error": result.error,
"from_cache": result.from_cache,
},
@@ -429,6 +614,7 @@ pub(crate) async fn build_admin_provider_query_models_response(
{
if let Some(models) = provider_query_read_provider_cached_models(state, &provider.id).await
{
let models = provider_query_attach_model_test_capabilities(&provider, models);
return Ok(Json(json!({
"success": !models.is_empty(),
"data": {
@@ -504,6 +690,7 @@ pub(crate) async fn build_admin_provider_query_models_response(
if !success && error.is_none() {
error = Some(ADMIN_PROVIDER_QUERY_NO_MODELS_FROM_KEY_DETAIL.to_string());
}
let models = provider_query_attach_model_test_capabilities(&provider, models);
Ok(Json(json!({
"success": success,
@@ -519,3 +706,152 @@ pub(crate) async fn build_admin_provider_query_models_response(
}))
.into_response())
}
#[cfg(test)]
mod tests {
use super::*;
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
fn grok_provider() -> StoredProviderCatalogProvider {
let mut provider = StoredProviderCatalogProvider::new(
"provider-1".to_string(),
"Grok".to_string(),
None,
"grok".to_string(),
)
.expect("provider should build");
provider.provider_type = "grok".to_string();
provider
}
fn grok_key_with_quota(quota: Value) -> StoredProviderCatalogKey {
let mut key = StoredProviderCatalogKey::new(
"key-1".to_string(),
"provider-1".to_string(),
"key-1".to_string(),
"oauth".to_string(),
None,
true,
)
.expect("key should build");
key.status_snapshot = Some(json!({ "quota": quota }));
key
}
fn model(id: &str) -> Value {
json!({ "id": id })
}
fn filtered_ids(key: &StoredProviderCatalogKey) -> Vec<String> {
provider_query_filter_models_for_key(
&grok_provider(),
key,
vec![
model("grok-4.20-0309-non-reasoning"),
model("grok-4.20-auto"),
model("grok-4.20-heavy"),
model("grok-imagine-image-lite"),
model("grok-imagine-image"),
model("grok-imagine-image-edit"),
],
)
.into_iter()
.filter_map(|item| item.get("id").and_then(Value::as_str).map(str::to_string))
.collect()
}
#[test]
fn provider_query_grok_basic_tier_hides_super_and_heavy_models() {
let key = grok_key_with_quota(json!({ "pool_tier": "basic" }));
assert_eq!(
filtered_ids(&key),
["grok-4.20-0309-non-reasoning", "grok-imagine-image-lite"]
);
}
#[test]
fn provider_query_grok_super_tier_hides_heavy_models() {
let key = grok_key_with_quota(json!({ "plan_type": "super" }));
assert_eq!(
filtered_ids(&key),
[
"grok-4.20-0309-non-reasoning",
"grok-4.20-auto",
"grok-imagine-image-lite",
"grok-imagine-image",
"grok-imagine-image-edit"
]
);
}
#[test]
fn provider_query_grok_heavy_tier_keeps_full_non_video_catalog() {
let key = grok_key_with_quota(json!({ "pool_tier": "heavy" }));
assert_eq!(
filtered_ids(&key),
[
"grok-4.20-0309-non-reasoning",
"grok-4.20-auto",
"grok-4.20-heavy",
"grok-imagine-image-lite",
"grok-imagine-image",
"grok-imagine-image-edit"
]
);
}
#[test]
fn provider_query_grok_tier_falls_back_to_live_quota_windows() {
let key = grok_key_with_quota(json!({
"windows": [
{ "model": "quota_fast", "limit_value": 140.0 }
]
}));
assert_eq!(
filtered_ids(&key),
[
"grok-4.20-0309-non-reasoning",
"grok-4.20-auto",
"grok-imagine-image-lite",
"grok-imagine-image",
"grok-imagine-image-edit"
]
);
}
#[test]
fn provider_query_attaches_model_test_capabilities_to_models() {
let models = provider_query_attach_model_test_capabilities(
&grok_provider(),
vec![
model("grok-4.20-fast"),
model("grok-imagine-image"),
model("grok-imagine-image-edit"),
],
);
assert!(models[0]["model_test_capabilities"]["openai:image"].is_null());
assert_eq!(
models[1]["model_test_capabilities"]["openai:image"]["max_generation_count"],
json!(4)
);
assert_eq!(
models[1]["model_test_capabilities"]["openai:image"]["supports_generation"],
json!(true)
);
assert_eq!(
models[2]["model_test_capabilities"]["openai:image"]["supports_generation"],
json!(false)
);
assert_eq!(
models[2]["model_test_capabilities"]["openai:image"]["supports_edit"],
json!(true)
);
}
}
apps/aether-gateway/src/handlers/admin/provider/query/models/model_test.rs
+491 -44
View File
@@ -81,6 +81,7 @@ use tracing::{debug, warn};
use uuid::Uuid;
mod adapter;
mod capabilities;
mod model_mapping;
mod summary;
@@ -88,13 +89,17 @@ use self::adapter::{
provider_query_antigravity_test_unsupported_reason,
provider_query_antigravity_unsupported_reason,
provider_query_default_antigravity_endpoint_test_body,
provider_query_model_test_endpoint_priority, provider_query_normalize_api_format_alias,
provider_query_standard_test_client_api_format,
provider_query_grok_test_unsupported_reason, provider_query_model_test_endpoint_priority,
provider_query_normalize_api_format_alias, provider_query_standard_test_client_api_format,
provider_query_standard_test_unsupported_reason,
provider_query_test_adapter_for_provider_api_format,
provider_query_transport_supports_model_test_execution,
provider_query_unsupported_test_api_format_message, ProviderQueryTestAdapter,
};
use self::capabilities::{
provider_query_openai_image_normalize_failure_message,
provider_query_openai_image_normalize_options,
};
use self::model_mapping::{
provider_query_resolve_explicit_mapped_effective_model,
provider_query_resolve_global_effective_model,
@@ -530,12 +535,20 @@ fn provider_query_build_test_request_body_for_route(
provider_query_build_test_request_body_with_model_policy(payload, model, override_custom_model)
}
fn provider_query_build_test_request_body_with_model_policy(
fn provider_query_build_test_request_body_for_api_format(
payload: &Value,
model: &str,
override_custom_model: bool,
route_path: &str,
client_api_format: &str,
) -> Value {
let client_api_format = provider_query_normalize_api_format_alias(client_api_format);
let override_custom_model = route_path.ends_with("/test-model-failover")
|| provider_query_extract_mapped_model_name(payload).is_some();
if let Some(mut body) = provider_query_extract_request_body(payload) {
let has_conversation = provider_query_request_body_has_conversation_for_api_format(
&body,
client_api_format.as_str(),
);
if let Some(object) = body.as_object_mut() {
if override_custom_model {
object.insert("model".to_string(), Value::String(model.to_string()));
@@ -544,6 +557,123 @@ fn provider_query_build_test_request_body_with_model_policy(
.entry("model".to_string())
.or_insert_with(|| Value::String(model.to_string()));
}
if !has_conversation {
provider_query_insert_default_test_conversation(
object,
client_api_format.as_str(),
payload,
);
}
}
return body;
}
let message = provider_query_extract_message(payload)
.unwrap_or_else(|| DEFAULT_PROVIDER_QUERY_TEST_MESSAGE.to_string());
match client_api_format.as_str() {
"openai:responses" | "openai:responses:compact" => json!({
"model": model,
"input": message,
"max_output_tokens": 30,
"temperature": 0.7,
"stream": true,
}),
"claude:messages" => json!({
"model": model,
"messages": [{
"role": "user",
"content": message
}],
"max_tokens": 30,
"temperature": 0.7,
"stream": true,
}),
_ => json!({
"model": model,
"messages": [{
"role": "user",
"content": message
}],
"max_tokens": 30,
"temperature": 0.7,
"stream": true,
}),
}
}
fn provider_query_build_grok_test_request_body_for_api_format(
payload: &Value,
model: &str,
route_path: &str,
client_api_format: &str,
) -> Value {
provider_query_build_test_request_body_for_api_format(
payload,
model,
route_path,
client_api_format,
)
}
fn provider_query_insert_default_test_conversation(
object: &mut Map<String, Value>,
client_api_format: &str,
payload: &Value,
) {
let message = provider_query_extract_message(payload)
.unwrap_or_else(|| DEFAULT_PROVIDER_QUERY_TEST_MESSAGE.to_string());
match client_api_format {
"openai:responses" | "openai:responses:compact" => {
object.insert("input".to_string(), Value::String(message));
}
"claude:messages" => {
object.insert(
"messages".to_string(),
json!([{ "role": "user", "content": message }]),
);
}
_ => {
object.insert(
"messages".to_string(),
json!([{ "role": "user", "content": message }]),
);
}
}
}
fn provider_query_grok_test_client_api_format(provider_api_format: &str) -> &'static str {
match provider_query_normalize_api_format_alias(provider_api_format).as_str() {
"openai:responses" | "openai:responses:compact" => "openai:responses",
"claude:messages" => "claude:messages",
_ => "openai:chat",
}
}
fn provider_query_build_test_request_body_with_model_policy(
payload: &Value,
model: &str,
override_custom_model: bool,
) -> Value {
if let Some(mut body) = provider_query_extract_request_body(payload) {
let has_conversation = provider_query_request_body_has_conversation(&body);
if let Some(object) = body.as_object_mut() {
if override_custom_model {
object.insert("model".to_string(), Value::String(model.to_string()));
} else {
object
.entry("model".to_string())
.or_insert_with(|| Value::String(model.to_string()));
}
if !has_conversation {
object.insert(
"messages".to_string(),
json!([{
"role": "user",
"content": provider_query_extract_message(payload)
.unwrap_or_else(|| DEFAULT_PROVIDER_QUERY_TEST_MESSAGE.to_string())
}]),
);
}
}
return body;
}
@@ -561,6 +691,73 @@ fn provider_query_build_test_request_body_with_model_policy(
})
}
fn provider_query_request_body_has_conversation(body: &Value) -> bool {
body.get("messages")
.and_then(Value::as_array)
.map(|messages| {
messages
.iter()
.any(|message| value_has_non_empty_text(message.get("content")))
})
.unwrap_or(false)
|| value_has_non_empty_text(body.get("input"))
|| value_has_non_empty_text(body.get("prompt"))
|| value_has_non_empty_text(body.get("query"))
|| value_has_non_empty_text(body.get("system"))
}
fn provider_query_request_body_has_conversation_for_api_format(
body: &Value,
client_api_format: &str,
) -> bool {
match provider_query_normalize_api_format_alias(client_api_format).as_str() {
"openai:responses" | "openai:responses:compact" => {
value_has_non_empty_text(body.get("input"))
|| value_has_non_empty_text(body.get("prompt"))
}
"claude:messages" => {
body.get("messages")
.and_then(Value::as_array)
.map(|messages| {
messages
.iter()
.any(|message| value_has_non_empty_text(message.get("content")))
})
.unwrap_or(false)
|| value_has_non_empty_text(body.get("system"))
}
_ => provider_query_request_body_has_conversation(body),
}
}
fn provider_query_request_body_is_openai_responses_shape(body: &Value) -> bool {
let Some(object) = body.as_object() else {
return false;
};
[
"input",
"tools",
"tool_choice",
"instructions",
"previous_response_id",
]
.iter()
.any(|key| object.contains_key(*key))
}
fn value_has_non_empty_text(value: Option<&Value>) -> bool {
match value {
Some(Value::String(value)) => !value.trim().is_empty(),
Some(Value::Array(values)) => values
.iter()
.any(|value| value_has_non_empty_text(Some(value))),
Some(Value::Object(values)) => values
.values()
.any(|value| value_has_non_empty_text(Some(value))),
_ => false,
}
}
fn provider_query_request_body_model<'a>(request_body: &'a Value, fallback: &'a str) -> &'a str {
request_body
.get("model")
@@ -1562,6 +1759,32 @@ fn provider_query_chatgpt_web_image_internal_url(base_url: &str) -> String {
format!("{base_url}/__aether/chatgpt-web-image")
}
fn provider_query_openai_image_test_upstream_url(
transport: &AdminGatewayProviderTransportSnapshot,
request_query: Option<&str>,
) -> String {
if transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("chatgpt_web")
{
provider_query_chatgpt_web_image_internal_url(&transport.endpoint.base_url)
} else if transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("grok")
{
crate::provider_transport::build_grok_upstream_url(
transport,
crate::provider_transport::GROK_CHAT_PATH,
)
} else {
crate::provider_transport::build_openai_image_upstream_url(transport, request_query)
}
}
async fn provider_query_finalize_openai_image_result(
route_path: &str,
trace_id: &str,
@@ -1661,12 +1884,16 @@ async fn provider_query_execute_openai_image_test_candidate(
*synthetic_request.headers_mut() = incoming_request_headers;
let (parts, _) = synthetic_request.into_parts();
let Some(normalized_request) =
crate::ai_serving::normalize_openai_image_request(&parts, &request_body, None)
else {
let provider_type = transport.provider.provider_type.as_str();
let Some(normalized_request) = crate::ai_serving::normalize_openai_image_request_with_options(
&parts,
&request_body,
None,
provider_query_openai_image_normalize_options(provider_type),
) else {
return Ok(provider_query_skipped_execution_outcome(
request_body.clone(),
"Provider request body could not be normalized for openai:image",
provider_query_openai_image_normalize_failure_message(provider_type, &request_body),
));
};
@@ -1675,6 +1902,11 @@ async fn provider_query_execute_openai_image_test_candidate(
.provider_type
.trim()
.eq_ignore_ascii_case("chatgpt_web");
let is_grok = transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("grok");
let mut provider_request_body = if is_chatgpt_web {
match crate::ai_serving::build_chatgpt_web_image_request_body(&parts, &request_body, None) {
Ok(body) => body,
@@ -1702,17 +1934,33 @@ async fn provider_query_execute_openai_image_test_candidate(
"Provider auth is unavailable for openai:image",
));
};
let transport_profile = state.resolve_transport_profile(&transport);
let Some(mut request_headers) = crate::provider_transport::build_openai_image_headers(
crate::provider_transport::ProviderOpenAiImageHeadersInput {
headers: &parts.headers,
auth_header: &auth_header,
auth_value: &auth_value,
header_rules: transport.endpoint.header_rules.as_ref(),
provider_request_body: &provider_request_body,
original_request_body: &request_body,
},
) else {
let Some(mut request_headers) = (if is_grok {
crate::provider_transport::build_grok_browser_headers(
crate::provider_transport::GrokHeaderInput {
transport: &transport,
transport_profile: transport_profile.as_ref(),
request_headers: Some(&parts.headers),
content_type: "application/json",
accept: "*/*",
header_rules: transport.endpoint.header_rules.as_ref(),
provider_request_body: &provider_request_body,
original_request_body: &request_body,
},
)
} else {
crate::provider_transport::build_openai_image_headers(
crate::provider_transport::ProviderOpenAiImageHeadersInput {
headers: &parts.headers,
auth_header: &auth_header,
auth_value: &auth_value,
header_rules: transport.endpoint.header_rules.as_ref(),
provider_request_body: &provider_request_body,
original_request_body: &request_body,
},
)
}) else {
return Ok(ProviderQueryExecutionOutcome {
status: "failed",
skip_reason: None,
@@ -1728,6 +1976,7 @@ async fn provider_query_execute_openai_image_test_candidate(
};
if is_chatgpt_web {
request_headers.insert("x-aether-chatgpt-web-image".to_string(), "1".to_string());
} else if is_grok {
} else {
crate::ai_serving::apply_codex_openai_responses_special_headers(
&mut request_headers,
@@ -1761,16 +2010,12 @@ async fn provider_query_execute_openai_image_test_candidate(
.filter(|value| !value.is_empty())
.unwrap_or(request_model.as_str())
.to_string();
let image_request = if is_chatgpt_web {
let image_request = if is_chatgpt_web || is_grok {
provider_request_body.clone()
} else {
normalized_request.summary_json.clone()
};
let request_url = if is_chatgpt_web {
provider_query_chatgpt_web_image_internal_url(&transport.endpoint.base_url)
} else {
crate::provider_transport::build_openai_image_upstream_url(&transport, parts.uri.query())
};
let request_url = provider_query_openai_image_test_upstream_url(&transport, parts.uri.query());
let upstream_is_stream = provider_request_body
.get("stream")
.and_then(Value::as_bool)
@@ -1796,11 +2041,27 @@ async fn provider_query_execute_openai_image_test_candidate(
proxy: state
.resolve_transport_proxy_snapshot_with_tunnel_affinity(&transport)
.await,
transport_profile: state.resolve_transport_profile(&transport),
transport_profile: transport_profile.clone(),
timeouts: state.resolve_transport_execution_timeouts(&transport),
};
let result = if is_chatgpt_web {
let result = if is_grok {
let report_context = json!({
"client_api_format": "openai:image",
"provider_api_format": "openai:image",
"provider_type": "grok",
"model": request_model,
"mapped_model": mapped_model,
"image_request": image_request.clone(),
});
state
.execute_execution_runtime_sync_plan_with_report_context(
Some(trace_id),
&plan,
Some(&report_context),
)
.await?
} else if is_chatgpt_web {
let report_context = json!({
"client_api_format": "openai:image",
"provider_api_format": "openai:image",
@@ -2115,6 +2376,177 @@ async fn provider_query_execute_antigravity_test_candidate(
})
}
async fn provider_query_execute_grok_test_candidate(
state: &AdminAppState<'_>,
provider: &StoredProviderCatalogProvider,
candidate: &ProviderQueryTestCandidate,
payload: &Value,
route_path: &str,
trace_id: &str,
) -> Result<ProviderQueryExecutionOutcome, GatewayError> {
let Some(transport) = state
.read_provider_transport_snapshot(&provider.id, &candidate.endpoint.id, &candidate.key.id)
.await?
else {
return Ok(provider_query_skipped_execution_outcome(
Value::Null,
"Provider transport snapshot is unavailable",
));
};
let provider_api_format =
provider_query_normalize_api_format_alias(&candidate.endpoint.api_format);
let client_api_format = provider_query_grok_test_client_api_format(&provider_api_format);
let request_body = provider_query_build_grok_test_request_body_for_api_format(
payload,
&candidate.effective_model,
route_path,
client_api_format,
);
if let Some(reason) =
provider_query_grok_test_unsupported_reason(&transport, &provider_api_format)
{
return Ok(provider_query_skipped_execution_outcome(
request_body,
format!(
"{} ({reason})",
provider_query_unsupported_test_api_format_message(&candidate.endpoint.api_format)
),
));
}
let incoming_request_headers = provider_query_extract_request_headers(payload);
let mut synthetic_request = http::Request::builder()
.uri(route_path)
.body(())
.map_err(|err| GatewayError::Internal(err.to_string()))?;
*synthetic_request.headers_mut() = incoming_request_headers;
let (parts, _) = synthetic_request.into_parts();
let request_model =
provider_query_request_body_model(&request_body, &candidate.effective_model);
let request_url = crate::provider_transport::build_grok_upstream_url(
&transport,
crate::provider_transport::GROK_CHAT_PATH,
);
let provider_request_body = crate::provider_transport::build_grok_app_chat_body(
client_api_format,
Some(request_model),
&request_body,
);
let report_context = json!({
"provider_type": provider.provider_type,
"provider_api_format": provider_api_format,
"client_api_format": client_api_format,
"model": request_model,
"mapped_model": candidate.effective_model,
"request_path": route_path,
"request_body": request_body,
});
let transport_profile = state.resolve_transport_profile(&transport);
let Some(request_headers) = crate::provider_transport::build_grok_browser_headers(
crate::provider_transport::GrokHeaderInput {
transport: &transport,
transport_profile: transport_profile.as_ref(),
request_headers: Some(&parts.headers),
content_type: "application/json",
accept: "text/event-stream",
header_rules: transport.endpoint.header_rules.as_ref(),
provider_request_body: &provider_request_body,
original_request_body: &request_body,
},
) else {
return Ok(ProviderQueryExecutionOutcome {
status: "failed",
skip_reason: None,
error_message: Some("provider request headers build failed".to_string()),
status_code: None,
latency_ms: None,
request_url,
request_headers: BTreeMap::new(),
request_body: provider_request_body,
response_headers: BTreeMap::new(),
response_body: None,
});
};
let plan = ExecutionPlan {
request_id: trace_id.to_string(),
candidate_id: Some(format!("provider-query-{}", candidate.key.id)),
provider_name: Some(provider.name.clone()),
provider_id: provider.id.clone(),
endpoint_id: candidate.endpoint.id.clone(),
key_id: candidate.key.id.clone(),
method: "POST".to_string(),
url: request_url.clone(),
headers: request_headers.clone(),
content_type: Some("application/json".to_string()),
content_encoding: None,
body: RequestBody::from_json(request_body.clone()),
stream: true,
client_api_format: client_api_format.to_string(),
provider_api_format: provider_api_format.clone(),
model_name: Some(request_model.to_string()),
proxy: state
.resolve_transport_proxy_snapshot_with_tunnel_affinity(&transport)
.await,
transport_profile,
timeouts: state.resolve_transport_execution_timeouts(&transport),
};
let result = match state
.execute_execution_runtime_sync_plan_with_report_context(
Some(trace_id),
&plan,
Some(&report_context),
)
.await
{
Ok(result) => result,
Err(err) => {
return Ok(ProviderQueryExecutionOutcome {
status: "failed",
skip_reason: None,
error_message: Some(format!("model test execution failed: {err:?}")),
status_code: None,
latency_ms: None,
request_url,
request_headers,
request_body: provider_request_body,
response_headers: BTreeMap::new(),
response_body: None,
});
}
};
let response_body = result.body.as_ref().and_then(|body| body.json_body.clone());
let did_fail = result.status_code >= 400 || response_body.is_none();
let error_message = if did_fail {
provider_query_extract_error_message(&result).or_else(|| {
response_body.is_none().then(|| {
format!(
"Provider returned HTTP {} without a model-test response body",
result.status_code
)
})
})
} else {
None
};
Ok(ProviderQueryExecutionOutcome {
status: if did_fail { "failed" } else { "success" },
skip_reason: None,
error_message,
status_code: Some(result.status_code),
latency_ms: result.telemetry.as_ref().and_then(|value| value.elapsed_ms),
request_url,
request_headers,
request_body: provider_request_body,
response_headers: result.headers,
response_body,
})
}
async fn provider_query_execute_standard_test_candidate(
state: &AdminAppState<'_>,
provider: &StoredProviderCatalogProvider,
@@ -2132,22 +2564,25 @@ async fn provider_query_execute_standard_test_candidate(
"Provider transport snapshot is unavailable",
));
};
let original_request_body = provider_query_build_test_request_body_for_route(
let provider_api_format = candidate.endpoint.api_format.as_str();
let normalized_provider_api_format =
crate::ai_serving::normalize_api_format_alias(provider_api_format);
let client_api_format =
provider_query_standard_test_client_api_format(normalized_provider_api_format.as_str());
let original_request_body = provider_query_build_test_request_body_for_api_format(
payload,
&candidate.effective_model,
route_path,
client_api_format,
);
if !provider_query_transport_supports_model_test_execution(
state,
&transport,
candidate.endpoint.api_format.as_str(),
provider_api_format,
) {
return Ok(provider_query_skipped_execution_outcome(
original_request_body,
provider_query_standard_test_unsupported_reason(
&transport,
candidate.endpoint.api_format.as_str(),
),
provider_query_standard_test_unsupported_reason(&transport, provider_api_format),
));
}
@@ -2159,11 +2594,6 @@ async fn provider_query_execute_standard_test_candidate(
let request_model =
provider_query_request_body_model(&request_body, &candidate.effective_model);
let provider_api_format = candidate.endpoint.api_format.as_str();
let normalized_provider_api_format =
crate::ai_serving::normalize_api_format_alias(provider_api_format);
let client_api_format =
provider_query_standard_test_client_api_format(normalized_provider_api_format.as_str());
let upstream_is_stream = provider_query_resolve_standard_test_upstream_is_stream(
transport.endpoint.config.as_ref(),
transport.provider.provider_type.as_str(),
@@ -2229,12 +2659,20 @@ async fn provider_query_execute_standard_test_candidate(
}
"openai:responses" | "openai:responses:compact" => {
let Some(mut provider_request_body) =
crate::ai_serving::build_cross_format_openai_chat_request_body(
&request_body,
request_model,
normalized_provider_api_format.as_str(),
upstream_is_stream,
)
(if provider_query_request_body_is_openai_responses_shape(&request_body) {
crate::ai_serving::build_local_openai_responses_request_body(
&request_body,
request_model,
upstream_is_stream,
)
} else {
crate::ai_serving::build_cross_format_openai_chat_request_body(
&request_body,
request_model,
normalized_provider_api_format.as_str(),
upstream_is_stream,
)
})
else {
return Ok(provider_query_skipped_execution_outcome(
request_body.clone(),
@@ -2659,6 +3097,12 @@ async fn build_admin_provider_query_kiro_failover_response(
)
.await
}
Some(ProviderQueryTestAdapter::Grok) => {
provider_query_execute_grok_test_candidate(
state, &provider, candidate, payload, route_path, &trace_id,
)
.await
}
Some(ProviderQueryTestAdapter::Standard) => {
provider_query_execute_standard_test_candidate(
state, &provider, candidate, payload, route_path, &trace_id,
@@ -2718,7 +3162,10 @@ async fn build_admin_provider_query_kiro_failover_response(
));
if is_success {
success_body = response_body;
success_stream = matches!(adapter, Some(ProviderQueryTestAdapter::Kiro));
success_stream = matches!(
adapter,
Some(ProviderQueryTestAdapter::Kiro | ProviderQueryTestAdapter::Grok)
);
winning_candidate_index = Some(candidate_index);
break;
}
apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/adapter.rs
+82
View File
@@ -10,6 +10,7 @@ use serde_json::{json, Value};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) enum ProviderQueryTestAdapter {
Standard,
Grok,
Kiro,
OpenAiImage,
Antigravity,
@@ -134,6 +135,65 @@ pub(super) fn provider_query_antigravity_test_unsupported_reason(
}
}
}
pub(super) fn provider_query_grok_test_unsupported_reason(
transport: &AdminGatewayProviderTransportSnapshot,
api_format: &str,
) -> Option<&'static str> {
if !transport.provider.is_active {
return Some("provider_inactive");
}
if !transport.endpoint.is_active {
return Some("endpoint_inactive");
}
if !transport.key.is_active {
return Some("key_inactive");
}
if !transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("grok")
{
return Some("transport_provider_type_unsupported");
}
let normalized_api_format = provider_query_normalize_api_format_alias(api_format);
if !matches!(
normalized_api_format.as_str(),
"openai:chat" | "openai:responses" | "openai:responses:compact" | "claude:messages"
) {
return Some("transport_api_format_mismatch");
}
if provider_query_normalize_api_format_alias(&transport.endpoint.api_format)
!= normalized_api_format
{
return Some("transport_api_format_mismatch");
}
if crate::provider_transport::resolve_grok_session_auth(transport).is_none() {
return Some("transport_oauth_resolution_unsupported");
}
if !crate::provider_transport::header_rules_are_locally_supported(
transport.endpoint.header_rules.as_ref(),
) {
return Some("transport_header_rules_unsupported");
}
if !crate::provider_transport::body_rules_are_locally_supported(
transport.endpoint.body_rules.as_ref(),
) {
return Some("transport_body_rules_unsupported");
}
if !crate::provider_transport::transport_proxy_is_locally_supported(transport) {
return Some("transport_proxy_unsupported");
}
if crate::provider_transport::transport_profile_is_configured(transport)
&& crate::provider_transport::resolve_transport_profile(transport).is_none()
{
return Some("transport_profile_unsupported");
}
None
}
pub(super) fn provider_query_normalize_api_format_alias(value: &str) -> String {
crate::ai_serving::normalize_api_format_alias(value)
}
@@ -147,6 +207,15 @@ pub(super) fn provider_query_test_adapter_for_provider_api_format(
}
let normalized_api_format = provider_query_normalize_api_format_alias(api_format);
if provider_type.trim().eq_ignore_ascii_case("grok") {
return match normalized_api_format.as_str() {
"openai:chat" | "openai:responses" | "openai:responses:compact" | "claude:messages" => {
Some(ProviderQueryTestAdapter::Grok)
}
"openai:image" => Some(ProviderQueryTestAdapter::OpenAiImage),
_ => None,
};
}
if normalized_api_format == "openai:image" {
return Some(ProviderQueryTestAdapter::OpenAiImage);
}
@@ -182,6 +251,16 @@ pub(super) fn provider_query_model_test_endpoint_priority(
let normalized_api_format = provider_query_normalize_api_format_alias(api_format);
match provider_query_test_adapter_for_provider_api_format(provider_type, api_format)? {
ProviderQueryTestAdapter::Kiro => Some(0),
ProviderQueryTestAdapter::Grok => {
if matches!(
normalized_api_format.as_str(),
"openai:chat" | "openai:responses" | "openai:responses:compact" | "claude:messages"
) {
Some(0)
} else {
Some(2)
}
}
ProviderQueryTestAdapter::Antigravity => Some(1),
ProviderQueryTestAdapter::OpenAiImage => Some(2),
ProviderQueryTestAdapter::Standard => {
@@ -234,6 +313,9 @@ pub(super) fn provider_query_transport_supports_model_test_execution(
)
.is_none()
}
Some(ProviderQueryTestAdapter::Grok) => {
provider_query_grok_test_unsupported_reason(transport, api_format).is_none()
}
Some(ProviderQueryTestAdapter::Standard) => match crate::ai_serving::normalize_api_format_alias(api_format).as_str() {
"openai:chat" => {
crate::provider_transport::policy::supports_local_openai_chat_transport(transport)
apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/capabilities.rs
+50
View File
@@ -0,0 +1,50 @@
use crate::handlers::admin::provider::shared::model_test_capabilities::{
admin_provider_openai_image_normalize_options, admin_provider_openai_image_test_capability,
AdminProviderOpenAiImageTestCapability,
};
use serde_json::Value;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) struct ProviderQueryOpenAiImageTestCapability(AdminProviderOpenAiImageTestCapability);
pub(super) fn provider_query_openai_image_test_capability(
provider_type: &str,
) -> ProviderQueryOpenAiImageTestCapability {
ProviderQueryOpenAiImageTestCapability(admin_provider_openai_image_test_capability(
provider_type,
))
}
pub(super) fn provider_query_openai_image_normalize_options(
provider_type: &str,
) -> crate::ai_serving::OpenAiImageNormalizeOptions {
admin_provider_openai_image_normalize_options(provider_type)
}
pub(super) fn provider_query_openai_image_requested_count(request_body: &Value) -> Option<u64> {
request_body.get("n").and_then(|value| {
value.as_u64().or_else(|| {
value
.as_str()
.map(str::trim)
.filter(|value| !value.is_empty())
.and_then(|value| value.parse::<u64>().ok())
})
})
}
pub(super) fn provider_query_openai_image_normalize_failure_message(
provider_type: &str,
request_body: &Value,
) -> String {
let capability = provider_query_openai_image_test_capability(provider_type);
if provider_query_openai_image_requested_count(request_body)
.is_some_and(|value| !capability.0.supports_generation_count(value))
{
return format!(
"Provider request body could not be normalized for openai:image: selected provider supports n=1..{} for generation",
capability.0.max_generation_count
);
}
"Provider request body could not be normalized for openai:image".to_string()
}
apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/summary.rs
+68 -2
View File
@@ -1,3 +1,5 @@
use std::collections::BTreeMap;
use super::super::provider_query_key_display_name;
use super::{ProviderQueryExecutionOutcome, ProviderQueryTestCandidate};
use serde_json::{json, Value};
@@ -22,13 +24,43 @@ pub(super) fn provider_query_test_attempt_payload(
"status_code": execution.status_code,
"latency_ms": execution.latency_ms,
"request_url": execution.request_url,
"request_headers": execution.request_headers,
"request_headers": provider_query_redact_diagnostic_headers(&execution.request_headers),
"request_body": execution.request_body,
"response_headers": execution.response_headers,
"response_headers": provider_query_redact_diagnostic_headers(&execution.response_headers),
"response_body": execution.response_body,
})
}
fn provider_query_redact_diagnostic_headers(
headers: &BTreeMap<String, String>,
) -> BTreeMap<String, String> {
headers
.iter()
.map(|(name, value)| {
if provider_query_header_is_sensitive(name) {
(name.clone(), "<redacted>".to_string())
} else {
(name.clone(), value.clone())
}
})
.collect()
}
fn provider_query_header_is_sensitive(name: &str) -> bool {
matches!(
name.trim().to_ascii_lowercase().as_str(),
"authorization"
| "proxy-authorization"
| "cookie"
| "set-cookie"
| "x-api-key"
| "api-key"
| "x-goog-api-key"
| "anthropic-api-key"
| "openai-api-key"
)
}
pub(super) fn provider_query_candidate_summary_payload(
total_candidates: usize,
total_attempts: usize,
@@ -133,3 +165,37 @@ pub(super) fn provider_query_candidate_summary_payload(
.unwrap_or(Value::Null),
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn provider_query_diagnostic_headers_redact_credentials() {
let headers = BTreeMap::from([
("cookie".to_string(), "sso=secret".to_string()),
("authorization".to_string(), "Bearer secret".to_string()),
("x-goog-api-key".to_string(), "secret".to_string()),
("content-type".to_string(), "application/json".to_string()),
]);
let redacted = provider_query_redact_diagnostic_headers(&headers);
assert_eq!(
redacted.get("cookie").map(String::as_str),
Some("<redacted>")
);
assert_eq!(
redacted.get("authorization").map(String::as_str),
Some("<redacted>")
);
assert_eq!(
redacted.get("x-goog-api-key").map(String::as_str),
Some("<redacted>")
);
assert_eq!(
redacted.get("content-type").map(String::as_str),
Some("application/json")
);
}
}
apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/tests.rs
+376
View File
@@ -1,6 +1,68 @@
use super::*;
use crate::handlers::admin::request::AdminGatewayProviderTransportSnapshot;
use serde_json::json;
fn sample_openai_image_transport(provider_type: &str) -> AdminGatewayProviderTransportSnapshot {
AdminGatewayProviderTransportSnapshot {
provider: crate::provider_transport::snapshot::GatewayProviderTransportProvider {
id: "provider-1".to_string(),
name: "Provider".to_string(),
provider_type: provider_type.to_string(),
website: None,
is_active: true,
keep_priority_on_conversion: false,
enable_format_conversion: false,
concurrent_limit: None,
max_retries: None,
proxy: None,
request_timeout_secs: None,
stream_first_byte_timeout_secs: None,
config: None,
},
endpoint: crate::provider_transport::snapshot::GatewayProviderTransportEndpoint {
id: "endpoint-1".to_string(),
provider_id: "provider-1".to_string(),
api_format: "openai:image".to_string(),
api_family: None,
endpoint_kind: None,
is_active: true,
base_url: "https://grok.com/".to_string(),
header_rules: None,
body_rules: None,
max_retries: None,
custom_path: None,
config: None,
format_acceptance_config: None,
proxy: None,
},
key: crate::provider_transport::snapshot::GatewayProviderTransportKey {
id: "key-1".to_string(),
provider_id: "provider-1".to_string(),
name: "key".to_string(),
auth_type: "oauth".to_string(),
is_active: true,
api_formats: None,
auth_type_by_format: None,
allow_auth_channel_mismatch_formats: None,
allowed_models: None,
capabilities: None,
rate_multipliers: None,
global_priority_by_format: None,
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
decrypted_api_key: String::new(),
decrypted_auth_config: Some(
json!({
"sso_token": "abc",
"sso_rw_token": "rw"
})
.to_string(),
),
},
}
}
#[test]
fn provider_query_test_request_body_preserves_custom_model() {
let payload = json!({
@@ -28,6 +90,41 @@ fn provider_query_test_request_body_defaults_missing_model() {
assert_eq!(body["model"], json!("fallback-model"));
}
#[test]
fn provider_query_test_request_body_fills_empty_conversation() {
let payload = json!({
"request_body": {
"model": "custom-upstream-model",
"messages": []
}
});
let body = provider_query_build_test_request_body(&payload, "fallback-model");
assert_eq!(body["model"], json!("custom-upstream-model"));
assert_eq!(
body["messages"],
json!([{ "role": "user", "content": DEFAULT_PROVIDER_QUERY_TEST_MESSAGE }])
);
}
#[test]
fn provider_query_test_request_body_keeps_non_empty_conversation() {
let payload = json!({
"request_body": {
"model": "custom-upstream-model",
"messages": [{ "role": "user", "content": "custom prompt" }]
}
});
let body = provider_query_build_test_request_body(&payload, "fallback-model");
assert_eq!(
body["messages"],
json!([{ "role": "user", "content": "custom prompt" }])
);
}
#[test]
fn provider_query_failover_request_body_overrides_custom_model() {
let payload = json!({
@@ -168,6 +265,54 @@ fn provider_query_standard_test_aggregates_responses_stream_body() {
assert_eq!(body["output"][0]["content"][0]["text"], json!("Hello"));
}
#[test]
fn provider_query_standard_test_aggregates_responses_image_generation_call() {
let stream_body = concat!(
"event: response.created\n",
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_img_123\",\"object\":\"response\",\"model\":\"gpt-5.4-mini\",\"status\":\"in_progress\",\"output\":[]}}\n\n",
"event: response.output_item.done\n",
"data: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"id\":\"ig_123\",\"type\":\"image_generation_call\",\"status\":\"completed\",\"output_format\":\"png\",\"result\":\"aGVsbG8=\"}}\n\n",
"event: response.completed\n",
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_img_123\",\"object\":\"response\",\"model\":\"gpt-5.4-mini\",\"status\":\"completed\",\"output\":[]}}\n\n",
);
let result = aether_contracts::ExecutionResult {
request_id: "provider-test".to_string(),
candidate_id: Some("candidate-0".to_string()),
status_code: 200,
headers: BTreeMap::new(),
body: Some(aether_contracts::ResponseBody {
json_body: None,
body_bytes_b64: Some(
base64::engine::general_purpose::STANDARD.encode(stream_body.as_bytes()),
),
}),
telemetry: None,
error: None,
};
let body = provider_query_standard_execution_response_body("openai:responses", &result)
.expect("responses image stream body should aggregate");
assert_eq!(body["output"][0]["type"], json!("image_generation_call"));
assert_eq!(body["output"][0]["result"], json!("aGVsbG8="));
}
#[test]
fn provider_query_responses_test_request_body_defaults_to_responses_input() {
let payload = json!({"message": "hello from responses"});
let body = provider_query_build_test_request_body_for_api_format(
&payload,
"gpt-5.4-mini",
"/api/admin/provider-query/test-model",
"openai:responses",
);
assert_eq!(body["model"], json!("gpt-5.4-mini"));
assert_eq!(body["input"], json!("hello from responses"));
assert!(body.get("messages").is_none());
}
#[test]
fn provider_query_test_adapter_routes_fixed_provider_endpoint_types() {
assert_eq!(
@@ -204,6 +349,22 @@ fn provider_query_test_adapter_routes_fixed_provider_endpoint_types() {
),
Some(ProviderQueryTestAdapter::Antigravity)
);
assert_eq!(
provider_query_test_adapter_for_provider_api_format("grok", "openai:chat"),
Some(ProviderQueryTestAdapter::Grok)
);
assert_eq!(
provider_query_test_adapter_for_provider_api_format("grok", "openai:responses"),
Some(ProviderQueryTestAdapter::Grok)
);
assert_eq!(
provider_query_test_adapter_for_provider_api_format("grok", "claude:messages"),
Some(ProviderQueryTestAdapter::Grok)
);
assert_eq!(
provider_query_test_adapter_for_provider_api_format("grok", "openai:image"),
Some(ProviderQueryTestAdapter::OpenAiImage)
);
assert_eq!(
provider_query_test_adapter_for_provider_api_format("custom", "openai:embedding"),
Some(ProviderQueryTestAdapter::Standard)
@@ -244,12 +405,154 @@ fn provider_query_endpoint_priority_prefers_text_before_cli_and_image() {
provider_query_model_test_endpoint_priority("chatgpt_web", "openai:image"),
Some(2)
);
assert_eq!(
provider_query_model_test_endpoint_priority("grok", "openai:chat"),
Some(0)
);
assert_eq!(
provider_query_model_test_endpoint_priority("grok", "openai:responses"),
Some(0)
);
assert_eq!(
provider_query_model_test_endpoint_priority("antigravity", "gemini:generate_content"),
Some(1)
);
}
#[test]
fn provider_query_grok_model_test_body_maps_non_reasoning_model_to_fast_mode() {
let payload = json!({
"request_body": {
"model": "grok-4.20-0309-non-reasoning",
"messages": [
{"role": "system", "content": "be concise"},
{"role": "user", "content": "hello"}
]
}
});
let request_body = provider_query_build_test_request_body_for_route(
&payload,
"grok-4.20-0309-non-reasoning",
"/api/admin/provider-query/test-model",
);
let upstream_body = crate::provider_transport::build_grok_app_chat_body(
"openai:chat",
Some(provider_query_request_body_model(
&request_body,
"grok-4.20-0309-non-reasoning",
)),
&request_body,
);
assert_eq!(upstream_body["modeId"], json!("fast"));
assert_eq!(
upstream_body["message"],
json!("[system]: be concise\n\n[user]: hello")
);
}
#[test]
fn provider_query_grok_model_test_uses_responses_client_body_for_responses_endpoint() {
let payload = json!({
"request_body": {
"model": "grok-4.20-0309-non-reasoning",
"input": "hello from responses body"
}
});
let request_body = provider_query_build_grok_test_request_body_for_api_format(
&payload,
"grok-4.20-0309-non-reasoning",
"/api/admin/provider-query/test-model",
"openai:responses",
);
let upstream_body = crate::provider_transport::build_grok_app_chat_body(
provider_query_grok_test_client_api_format("openai:responses"),
Some(provider_query_request_body_model(
&request_body,
"grok-4.20-0309-non-reasoning",
)),
&request_body,
);
assert_eq!(upstream_body["modeId"], json!("fast"));
assert_eq!(upstream_body["message"], json!("hello from responses body"));
}
#[test]
fn provider_query_grok_model_test_uses_responses_input_when_existing_body_has_messages() {
let payload = json!({
"request_body": {
"model": "grok-4.20-0309-non-reasoning",
"messages": [{
"role": "user",
"content": "hello from stale chat body"
}]
}
});
let request_body = provider_query_build_grok_test_request_body_for_api_format(
&payload,
"grok-4.20-0309-non-reasoning",
"/api/admin/provider-query/test-model",
"openai:responses",
);
assert_eq!(request_body["model"], json!("grok-4.20-0309-non-reasoning"));
assert_eq!(
request_body["input"],
json!("Hello! This is a test message.")
);
assert!(request_body.get("messages").is_some());
let upstream_body = crate::provider_transport::build_grok_app_chat_body(
provider_query_grok_test_client_api_format("openai:responses"),
Some(provider_query_request_body_model(
&request_body,
"grok-4.20-0309-non-reasoning",
)),
&request_body,
);
assert_eq!(upstream_body["modeId"], json!("fast"));
assert_eq!(
upstream_body["message"],
json!("Hello! This is a test message.")
);
}
#[test]
fn provider_query_grok_model_test_defaults_claude_messages_body_for_claude_endpoint() {
let payload = json!({});
let request_body = provider_query_build_grok_test_request_body_for_api_format(
&payload,
"grok-4.20-0309-non-reasoning",
"/api/admin/provider-query/test-model",
"claude:messages",
);
assert_eq!(request_body["model"], json!("grok-4.20-0309-non-reasoning"));
assert_eq!(
request_body["messages"],
json!([{ "role": "user", "content": DEFAULT_PROVIDER_QUERY_TEST_MESSAGE }])
);
let upstream_body = crate::provider_transport::build_grok_app_chat_body(
provider_query_grok_test_client_api_format("claude:messages"),
Some(provider_query_request_body_model(
&request_body,
"grok-4.20-0309-non-reasoning",
)),
&request_body,
);
assert_eq!(upstream_body["modeId"], json!("fast"));
assert_eq!(
upstream_body["message"],
json!("[user]: Hello! This is a test message.")
);
}
#[test]
fn provider_query_candidate_summary_marks_unused_after_first_success() {
let attempts = vec![json!({
@@ -360,3 +663,76 @@ fn provider_query_failover_image_test_request_body_overrides_model() {
assert_eq!(body["model"], json!("new-image-model"));
}
#[test]
fn provider_query_grok_image_test_allows_multi_generation_count() {
let request = http::Request::builder()
.uri("/v1/images/generations")
.body(())
.expect("request should build");
let (parts, _) = request.into_parts();
let body = json!({
"model": "grok-imagine-image",
"prompt": "draw",
"n": 2
});
let normalized = crate::ai_serving::normalize_openai_image_request_with_options(
&parts,
&body,
None,
provider_query_openai_image_normalize_options("grok"),
)
.expect("grok image model tests should allow multi-image generation");
let provider_body = crate::ai_serving::build_openai_image_provider_request_body(&normalized);
assert_eq!(provider_body["n"], json!(2));
}
#[test]
fn provider_query_grok_image_test_uses_grok_app_chat_upstream_url() {
let transport = sample_openai_image_transport("grok");
assert_eq!(
provider_query_openai_image_test_upstream_url(&transport, Some("trace=1")),
"https://grok.com/rest/app-chat/conversations/new"
);
}
#[test]
fn provider_query_chatgpt_web_image_test_uses_internal_upstream_url() {
let transport = sample_openai_image_transport("chatgpt_web");
assert_eq!(
provider_query_openai_image_test_upstream_url(&transport, Some("trace=1")),
"https://grok.com/__aether/chatgpt-web-image"
);
}
#[test]
fn provider_query_non_grok_image_test_keeps_single_generation_boundary() {
let request = http::Request::builder()
.uri("/v1/images/generations")
.body(())
.expect("request should build");
let (parts, _) = request.into_parts();
let body = json!({
"model": "gpt-image-2",
"prompt": "draw",
"n": 2
});
assert!(
crate::ai_serving::normalize_openai_image_request_with_options(
&parts,
&body,
None,
provider_query_openai_image_normalize_options("chatgpt_web"),
)
.is_none()
);
assert_eq!(
provider_query_openai_image_normalize_failure_message("chatgpt_web", &body),
"Provider request body could not be normalized for openai:image: selected provider supports n=1..1 for generation"
);
}