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

Merge remote-tracking branch 'origin/main' into codex/pool-key-bulk-management-20260714

Browse Source 
# Conflicts:
#	apps/aether-gateway/src/handlers/admin/request/provider/tasks.rs
#	frontend/src/api/endpoints/pool.ts
This commit is contained in:
MMEXA
2026-07-16 23:43:04 +08:00
parent 9b8724453b 312583f055
commit ec1783a35c
1257 changed files with 80521 additions and 35495 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/handlers/admin/observability/monitoring/cache_affinity.rs
+15
View File
@@ -104,6 +104,7 @@ fn is_known_admin_monitoring_api_format(value: &str) -> bool {
"openai:chat"
| "openai:responses"
| "openai:responses:compact"
| "openai:search"
| "openai:image"
| "openai:video"
| "openai:embedding"
@@ -492,6 +493,20 @@ mod tests {
);
}
#[test]
fn parses_v2_scheduler_affinity_key_for_openai_search() {
let parsed = parse_admin_monitoring_scheduler_affinity_key(
"scheduler_affinity:v2:user-key-1:openai:search:gpt-5.6-sol:codex:sessionhash",
)
.expect("Search scheduler key should parse");
assert_eq!(parsed.affinity_key, "user-key-1");
assert_eq!(parsed.api_format, "openai:search");
assert_eq!(parsed.model_name, "gpt-5.6-sol");
assert_eq!(parsed.client_family.as_deref(), Some("codex"));
assert_eq!(parsed.session_hash.as_deref(), Some("sessionhash"));
}
#[test]
fn parses_v2_scheduler_affinity_key_with_three_segment_api_format() {
let parsed = parse_admin_monitoring_scheduler_affinity_key(
apps/aether-gateway/src/handlers/admin/observability/stats/analytics_routes.rs
+5
View File
@@ -259,6 +259,11 @@ pub(super) async fn maybe_build_local_admin_stats_analytics_response(
"has_format_conversion",
),
slow_threshold_ms,
include_timeline: query_param_optional_bool(
request_context.query_string(),
"include_timeline",
)
.unwrap_or(true),
})
.await?;
return Ok(Some(build_admin_stats_provider_performance_response(
apps/aether-gateway/src/handlers/admin/provider/crud/writes.rs
+8 -4
View File
@@ -6,7 +6,8 @@ use crate::handlers::admin::provider::shared::payloads::{
AdminProviderCreateRequest, AdminProviderUpdatePatch,
};
use crate::handlers::admin::provider::write::provider::{
reconcile_admin_fixed_provider_template_endpoints, reconcile_admin_fixed_provider_template_keys,
reconcile_admin_fixed_provider_template_endpoints,
reconcile_admin_fixed_provider_template_endpoints_after_update,
};
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::handlers::admin::shared::attach_admin_audit_response;
@@ -80,7 +81,6 @@ pub(crate) async fn maybe_build_local_admin_provider_writes_response(
.is_some()
{
reconcile_admin_fixed_provider_template_endpoints(state, &created_provider).await?;
reconcile_admin_fixed_provider_template_keys(state, &created_provider).await?;
}
return Ok(Some(attach_admin_audit_response(
Json(json!({
@@ -161,8 +161,12 @@ pub(crate) async fn maybe_build_local_admin_provider_writes_response(
.fixed_provider_template(&updated_record.provider_type)
.is_some()
{
reconcile_admin_fixed_provider_template_endpoints(state, &updated_record).await?;
reconcile_admin_fixed_provider_template_keys(state, &updated_record).await?;
reconcile_admin_fixed_provider_template_endpoints_after_update(
state,
&existing_provider,
&updated_record,
)
.await?;
}
return Ok(Some(
match state
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/codex/mod.rs
+64 -20
View File
@@ -51,14 +51,6 @@ fn merge_codex_quota_metadata(
serde_json::Value::Object(merged)
}
fn codex_reset_credits_available_count(metadata: &Map<String, Value>) -> Option<u64> {
metadata
.get("reset_credits")
.and_then(Value::as_object)
.and_then(|reset_credits| reset_credits.get("available_count"))
.and_then(aether_admin::provider::quota::coerce_json_u64)
}
fn truncate_codex_reset_credit_detail_error(message: impl Into<String>) -> String {
let message = message.into();
let mut sanitized = message.replace('\n', " ");
@@ -85,11 +77,7 @@ fn merge_codex_reset_credit_detail_metadata(
.cloned()
.unwrap_or_default();
let has_usage_available_count = reset_credits.contains_key("available_count");
for (key, value) in detail_reset_credits {
if key == "available_count" && has_usage_available_count {
continue;
}
reset_credits.insert(key.clone(), value.clone());
}
codex_metadata.insert("reset_credits".to_string(), Value::Object(reset_credits));
@@ -97,6 +85,7 @@ fn merge_codex_reset_credit_detail_metadata(
fn mark_codex_reset_credit_detail_failed(
codex_metadata: &mut Map<String, Value>,
updated_at_unix_secs: u64,
detail_error: impl Into<String>,
) {
let mut reset_credits = codex_metadata
@@ -104,6 +93,7 @@ fn mark_codex_reset_credit_detail_failed(
.and_then(Value::as_object)
.cloned()
.unwrap_or_default();
reset_credits.insert("updated_at".to_string(), json!(updated_at_unix_secs));
reset_credits.insert("detail_source".to_string(), json!("wham_readonly"));
reset_credits.insert("detail_status".to_string(), json!("failed"));
reset_credits.insert(
@@ -124,16 +114,11 @@ async fn enrich_codex_reset_credit_details(
codex_metadata: &mut Map<String, Value>,
now_unix_secs: u64,
) -> Result<(), GatewayError> {
let available_count = codex_reset_credits_available_count(codex_metadata).unwrap_or(0);
if available_count == 0 {
return Ok(());
}
let request_spec = match build_codex_reset_credits_request_spec(transport, resolved_oauth_auth)
{
Ok(request_spec) => request_spec,
Err(message) => {
mark_codex_reset_credit_detail_failed(codex_metadata, message);
mark_codex_reset_credit_detail_failed(codex_metadata, now_unix_secs, message);
return Ok(());
}
};
@@ -146,6 +131,7 @@ async fn enrich_codex_reset_credit_details(
ProviderQuotaExecutionOutcome::Failure(detail) => {
mark_codex_reset_credit_detail_failed(
codex_metadata,
now_unix_secs,
format!("reset credit detail 请求执行失败: {detail}"),
);
return Ok(());
@@ -157,6 +143,7 @@ async fn enrich_codex_reset_credit_details(
.unwrap_or_else(|| format!("HTTP {}", result.status_code));
mark_codex_reset_credit_detail_failed(
codex_metadata,
now_unix_secs,
format!(
"reset credit detail 返回状态码 {}: {detail}",
result.status_code
@@ -170,7 +157,11 @@ async fn enrich_codex_reset_credit_details(
.as_ref()
.and_then(|body| body.json_body.as_ref())
else {
mark_codex_reset_credit_detail_failed(codex_metadata, "无法解析 reset credit detail 响应");
mark_codex_reset_credit_detail_failed(
codex_metadata,
now_unix_secs,
"无法解析 reset credit detail 响应",
);
return Ok(());
};
if let Some(detail_metadata) =
@@ -178,7 +169,11 @@ async fn enrich_codex_reset_credit_details(
{
merge_codex_reset_credit_detail_metadata(codex_metadata, &detail_metadata);
} else {
mark_codex_reset_credit_detail_failed(codex_metadata, "reset credit detail 响应为空");
mark_codex_reset_credit_detail_failed(
codex_metadata,
now_unix_secs,
"reset credit detail 响应为空",
);
}
Ok(())
@@ -776,3 +771,52 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
"auto_removed_hard_banned": auto_removed_hard_banned_count,
})))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn codex_reset_credit_detail_count_overrides_usage_count() {
let mut metadata = json!({
"reset_credits": {
"available_count": 0,
"detail_source": "wham_usage"
}
})
.as_object()
.cloned()
.expect("metadata object");
let detail = json!({
"reset_credits": {
"available_count": 2,
"detail_source": "wham_readonly"
}
});
merge_codex_reset_credit_detail_metadata(&mut metadata, &detail);
assert_eq!(
metadata
.get("reset_credits")
.and_then(Value::as_object)
.and_then(|credits| credits.get("available_count")),
Some(&json!(2u64))
);
}
#[test]
fn codex_reset_credit_detail_failure_records_attempt_time() {
let mut metadata = Map::new();
mark_codex_reset_credit_detail_failed(&mut metadata, 1_777_000_000, "request failed");
assert_eq!(
metadata
.get("reset_credits")
.and_then(Value::as_object)
.and_then(|credits| credits.get("updated_at")),
Some(&json!(1_777_000_000u64))
);
}
}
apps/aether-gateway/src/handlers/admin/provider/pool/runtime/writes.rs
+1 -1
View File
@@ -600,7 +600,7 @@ mod tests {
};
use crate::AppState;
use aether_runtime_state::{RedisClientConfig, RuntimeState, RuntimeStateConfig};
use aether_testkit::ManagedRedisServer;
use aether_test_support::ManagedRedisServer;
use std::collections::BTreeMap;
async fn start_managed_redis_or_skip() -> Option<ManagedRedisServer> {
apps/aether-gateway/src/handlers/admin/provider/pool_admin/batch_routes/import.rs
-7
View File
@@ -55,13 +55,6 @@ pub(super) async fn build_admin_pool_batch_import_response(
}
};
if payload.keys.len() > 500 {
return Ok(build_admin_pool_error_response(
http::StatusCode::BAD_REQUEST,
"keys length must be less than or equal to 500",
));
}
state
.build_admin_pool_batch_import_response(&provider_id, payload)
.await
apps/aether-gateway/src/handlers/admin/provider/query/models/mod.rs
+64 -19
View File
@@ -64,7 +64,8 @@ use aether_data_contracts::repository::provider_catalog::{
};
use aether_model_fetch::{
aggregate_models_for_cache, fetch_models_from_transports, json_string_list,
merge_upstream_metadata, preset_models_for_provider, selected_models_fetch_endpoints,
model_catalog_upstream_metadata, preset_models_for_provider, selected_models_fetch_endpoints,
upstream_metadata_namespace_updates,
};
use axum::{
body::{to_bytes, Body},
@@ -301,6 +302,7 @@ fn provider_query_attach_model_test_capabilities(
fn provider_query_codex_preset_fallback(
provider: &StoredProviderCatalogProvider,
fetch_error: &str,
) -> Option<ProviderQueryKeyFetchResult> {
if !provider.provider_type.trim().eq_ignore_ascii_case("codex") {
return None;
@@ -309,12 +311,59 @@ fn provider_query_codex_preset_fallback(
Some(ProviderQueryKeyFetchResult {
models: aggregate_models_for_cache(&models),
error: None,
warning: None,
warning: Some(format!(
"Codex 动态模型目录不可用，已使用内置模型卡：{fetch_error}"
)),
from_cache: false,
has_success: true,
})
}
async fn provider_query_persist_preset_models(
state: &AdminAppState<'_>,
provider: &StoredProviderCatalogProvider,
key: &StoredProviderCatalogKey,
models: &[Value],
) -> Result<(), GatewayError> {
if models.is_empty() {
return Ok(());
}
<AppState as ModelFetchRuntimeState>::write_upstream_models_cache(
state.app(),
&provider.id,
&key.id,
models,
)
.await;
if let Some(catalog_metadata) = model_catalog_upstream_metadata(&provider.provider_type, models)
{
provider_query_persist_upstream_metadata(state, key, &catalog_metadata).await?;
}
Ok(())
}
async fn provider_query_persist_upstream_metadata(
state: &AdminAppState<'_>,
key: &StoredProviderCatalogKey,
upstream_metadata: &Value,
) -> Result<(), GatewayError> {
let updated_at = current_unix_secs();
for (namespace, value) in
upstream_metadata_namespace_updates(key.upstream_metadata.as_ref(), upstream_metadata)
{
state
.app()
.upsert_provider_catalog_key_upstream_metadata_namespace(
&key.id,
&namespace,
&value,
Some(updated_at),
)
.await?;
}
Ok(())
}
mod model_test;
pub(crate) use self::model_test::{
@@ -459,11 +508,9 @@ 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),
);
let models = aggregate_models_for_cache(&models);
provider_query_persist_preset_models(state, provider, key, &models).await?;
let models = provider_query_filter_models_for_key(provider, key, models);
return Ok(ProviderQueryKeyFetchResult {
models,
error: None,
@@ -512,7 +559,11 @@ async fn provider_query_fetch_models_for_key(
Ok(outcome) => outcome,
Err(err) => {
all_errors.push(err);
if let Some(fallback) = provider_query_codex_preset_fallback(provider) {
if let Some(fallback) =
provider_query_codex_preset_fallback(provider, &all_errors.join("; "))
{
provider_query_persist_preset_models(state, provider, key, &fallback.models)
.await?;
return Ok(fallback);
}
return Ok(ProviderQueryKeyFetchResult {
@@ -537,20 +588,14 @@ async fn provider_query_fetch_models_for_key(
.await;
}
if let Some(upstream_metadata) = outcome.upstream_metadata.as_ref() {
let merged_metadata =
merge_upstream_metadata(key.upstream_metadata.as_ref(), upstream_metadata);
state
.app()
.update_provider_catalog_key_upstream_metadata(
&key.id,
Some(&merged_metadata),
Some(current_unix_secs()),
)
.await?;
provider_query_persist_upstream_metadata(state, key, upstream_metadata).await?;
}
if unique_models.is_empty() && !all_errors.is_empty() {
if let Some(fallback) = provider_query_codex_preset_fallback(provider) {
if let Some(fallback) =
provider_query_codex_preset_fallback(provider, &all_errors.join("; "))
{
provider_query_persist_preset_models(state, provider, key, &fallback.models).await?;
return Ok(fallback);
}
}
apps/aether-gateway/src/handlers/admin/provider/query/models/model_test.rs
+253 -62
View File
@@ -29,10 +29,7 @@ use crate::handlers::shared::{
provider_key_status_snapshot_payload,
};
use crate::model_fetch::ModelFetchRuntimeState;
use crate::provider_key_auth::{
provider_key_auth_semantics, provider_key_configured_api_formats,
provider_key_inherits_provider_api_formats,
};
use crate::provider_key_auth::provider_key_auth_semantics;
use crate::provider_transport::antigravity::{
build_antigravity_safe_v1internal_request, build_antigravity_static_identity_headers,
classify_local_antigravity_request_support, AntigravityEnvelopeRequestType,
@@ -543,6 +540,22 @@ fn provider_query_build_test_request_body_for_api_format(
model: &str,
route_path: &str,
client_api_format: &str,
) -> Value {
provider_query_build_test_request_body_for_api_format_with_search_session(
payload,
model,
route_path,
client_api_format,
None,
)
}
fn provider_query_build_test_request_body_for_api_format_with_search_session(
payload: &Value,
model: &str,
route_path: &str,
client_api_format: &str,
search_session_id: Option<&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")
@@ -568,7 +581,7 @@ fn provider_query_build_test_request_body_for_api_format(
);
} else if matches!(
client_api_format.as_str(),
"openai:responses" | "openai:responses:compact"
"openai:responses" | "openai:responses:compact" | "openai:search"
) && !value_has_non_empty_text(object.get("input"))
{
if let Some(prompt) = object
@@ -580,7 +593,7 @@ fn provider_query_build_test_request_body_for_api_format(
}
if matches!(
client_api_format.as_str(),
"openai:responses" | "openai:responses:compact"
"openai:responses" | "openai:responses:compact" | "openai:search"
) && value_has_non_empty_text(object.get("input"))
{
object.remove("prompt");
@@ -590,6 +603,9 @@ fn provider_query_build_test_request_body_for_api_format(
{
object.remove("messages");
}
if client_api_format == "openai:search" {
provider_query_ensure_search_test_fields(object, payload, search_session_id);
}
}
return body;
}
@@ -620,6 +636,15 @@ fn provider_query_build_test_request_body_for_api_format(
"temperature": 0.7,
"stream": true,
}),
"openai:search" => json!({
"id": provider_query_search_test_session_id(search_session_id),
"model": model,
"input": message,
"commands": {
"search_query": [{"q": message}]
},
"max_output_tokens": 256,
}),
"claude:messages" => json!({
"model": model,
"messages": [{
@@ -680,7 +705,7 @@ fn provider_query_insert_default_test_conversation(
.entry("top_n".to_string())
.or_insert_with(|| Value::from(4_u64));
}
"openai:responses" | "openai:responses:compact" => {
"openai:responses" | "openai:responses:compact" | "openai:search" => {
object.insert("input".to_string(), Value::String(message));
}
"claude:messages" => {
@@ -698,6 +723,48 @@ fn provider_query_insert_default_test_conversation(
}
}
fn provider_query_search_test_session_id(search_session_id: Option<&str>) -> String {
search_session_id
.map(str::trim)
.filter(|value| !value.is_empty())
.map(|value| format!("aether-model-test-{value}"))
.unwrap_or_else(|| format!("aether-model-test-{}", Uuid::new_v4().simple()))
}
fn provider_query_ensure_search_test_fields(
object: &mut Map<String, Value>,
payload: &Value,
search_session_id: Option<&str>,
) {
let query = object
.get("input")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.or_else(|| provider_query_extract_message(payload))
.unwrap_or_else(|| DEFAULT_PROVIDER_QUERY_TEST_MESSAGE.to_string());
let has_session_id = object
.get("id")
.and_then(Value::as_str)
.is_some_and(|value| !value.trim().is_empty());
if !has_session_id {
object.insert(
"id".to_string(),
provider_query_search_test_session_id(search_session_id).into(),
);
}
object
.entry("input".to_string())
.or_insert_with(|| Value::String(query.clone()));
object
.entry("commands".to_string())
.or_insert_with(|| json!({"search_query": [{"q": query}]}));
object
.entry("max_output_tokens".to_string())
.or_insert_with(|| Value::from(256_u64));
}
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",
@@ -767,7 +834,7 @@ fn provider_query_request_body_has_conversation_for_api_format(
client_api_format: &str,
) -> bool {
match provider_query_normalize_api_format_alias(client_api_format).as_str() {
"openai:responses" | "openai:responses:compact" => {
"openai:responses" | "openai:responses:compact" | "openai:search" => {
value_has_non_empty_text(body.get("input"))
|| value_has_non_empty_text(body.get("prompt"))
}
@@ -828,14 +895,12 @@ fn provider_query_resolve_standard_test_upstream_is_stream(
provider_type: &str,
provider_api_format: &str,
) -> bool {
let hard_requires_streaming = crate::ai_serving::force_upstream_streaming_for_provider(
crate::ai_serving::resolve_upstream_is_stream_for_provider(
endpoint_config,
provider_type,
provider_api_format,
);
crate::ai_serving::resolve_upstream_is_stream_from_endpoint_config(
endpoint_config,
false,
hard_requires_streaming,
false,
)
}
@@ -876,15 +941,11 @@ fn provider_query_key_supports_endpoint(
provider_type: &str,
endpoint_api_format: &str,
) -> bool {
if provider_key_inherits_provider_api_formats(key, provider_type) {
return true;
}
let formats = provider_key_configured_api_formats(key);
let endpoint_api_format = provider_query_normalize_api_format_alias(endpoint_api_format);
formats.is_empty()
|| formats
.iter()
.any(|value| provider_query_normalize_api_format_alias(value) == endpoint_api_format)
crate::handlers::shared::provider_catalog_key_supports_format(
key,
provider_type,
endpoint_api_format,
)
}
async fn provider_query_select_preferred_non_kiro_endpoint(
@@ -1626,6 +1687,15 @@ fn provider_query_standard_execution_response_body(
{
return None;
}
if result.status_code < 400
&& provider_query_normalize_api_format_alias(provider_api_format) == "openai:search"
&& !body
.get("output")
.and_then(Value::as_str)
.is_some_and(|value| !value.trim().is_empty())
{
return None;
}
Some(body)
}
@@ -2104,8 +2174,18 @@ async fn provider_query_execute_openai_image_test_candidate(
route_path,
);
let incoming_request_headers = provider_query_extract_request_headers(payload);
let image_request_path = if request_body.get("image").is_some()
|| request_body
.get("images")
.and_then(Value::as_array)
.is_some_and(|images| !images.is_empty())
{
"/v1/images/edits"
} else {
"/v1/images/generations"
};
let mut synthetic_request = http::Request::builder()
.uri("/v1/images/generations")
.uri(image_request_path)
.body(())
.map_err(|err| GatewayError::Internal(err.to_string()))?;
*synthetic_request.headers_mut() = incoming_request_headers;
@@ -2116,11 +2196,18 @@ async fn provider_query_execute_openai_image_test_candidate(
&parts,
&request_body,
None,
provider_query_openai_image_normalize_options(provider_type),
provider_query_openai_image_normalize_options(
provider_type,
Some(candidate.effective_model.as_str()),
),
) else {
return Ok(provider_query_skipped_execution_outcome(
request_body.clone(),
provider_query_openai_image_normalize_failure_message(provider_type, &request_body),
provider_query_openai_image_normalize_failure_message(
provider_type,
Some(candidate.effective_model.as_str()),
&request_body,
),
));
};
@@ -2139,29 +2226,50 @@ async fn provider_query_execute_openai_image_test_candidate(
.provider_type
.trim()
.eq_ignore_ascii_case("codex");
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,
Err(err) => err.to_error_json(),
}
} else if is_codex || is_grok {
crate::ai_serving::build_openai_image_provider_request_body(&normalized_request)
let upstream_is_stream = crate::ai_serving::resolve_upstream_is_stream_for_provider(
transport.endpoint.config.as_ref(),
transport.provider.provider_type.as_str(),
"openai:image",
request_body
.get("stream")
.and_then(Value::as_bool)
.unwrap_or(false),
false,
);
let provider_request_body = if is_chatgpt_web {
Some(
match crate::ai_serving::build_chatgpt_web_image_request_body(
&parts,
&request_body,
None,
) {
Ok(body) => body,
Err(err) => err.to_error_json(),
},
)
} else if is_codex {
crate::ai_serving::build_codex_openai_image_api_provider_request_body(
&normalized_request,
Some(candidate.effective_model.as_str()),
upstream_is_stream,
)
} else if is_grok {
Some(crate::ai_serving::build_openai_image_provider_request_body(
&normalized_request,
))
} else {
crate::ai_serving::build_openai_image_api_provider_request_body(
&normalized_request,
Some(candidate.effective_model.as_str()),
upstream_is_stream,
)
};
if !is_chatgpt_web {
crate::ai_serving::apply_codex_openai_responses_special_body_edits(
&mut provider_request_body,
transport.provider.provider_type.as_str(),
"openai:image",
transport.endpoint.body_rules.as_ref(),
Some(candidate.key.id.as_str()),
);
}
let Some(provider_request_body) = provider_request_body else {
return Ok(provider_query_skipped_execution_outcome(
request_body,
"Provider request is outside the Codex Images contract",
));
};
let oauth_auth = state.resolve_local_oauth_header_auth(&transport).await?;
let Some((auth_header, auth_value)) =
crate::provider_transport::resolve_openai_image_auth(&transport).or(oauth_auth)
@@ -2193,10 +2301,12 @@ async fn provider_query_execute_openai_image_test_candidate(
headers: &parts.headers,
auth_header: &auth_header,
auth_value: &auth_value,
accept: if is_codex || is_chatgpt_web {
"text/event-stream"
accept: if is_codex {
None
} else if upstream_is_stream {
Some("text/event-stream")
} else {
"application/json"
Some("application/json")
},
header_rules: transport.endpoint.header_rules.as_ref(),
provider_request_body: &provider_request_body,
@@ -2221,7 +2331,7 @@ async fn provider_query_execute_openai_image_test_candidate(
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(
crate::ai_serving::apply_codex_openai_special_headers(
&mut request_headers,
&provider_request_body,
&parts.headers,
@@ -2260,13 +2370,9 @@ async fn provider_query_execute_openai_image_test_candidate(
};
let request_url = provider_query_openai_image_test_upstream_url(
&transport,
Some(parts.uri.path()),
Some(image_request_path),
parts.uri.query(),
);
let upstream_is_stream = provider_request_body
.get("stream")
.and_then(Value::as_bool)
.unwrap_or(true);
let plan = ExecutionPlan {
request_id: trace_id.to_string(),
@@ -2842,12 +2948,14 @@ async fn provider_query_execute_standard_test_candidate(
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,
);
let original_request_body =
provider_query_build_test_request_body_for_api_format_with_search_session(
payload,
&candidate.effective_model,
route_path,
client_api_format,
Some(trace_id),
);
if crate::provider_transport::is_windsurf_provider_transport(&transport)
&& provider_query_normalize_api_format_alias(candidate.endpoint.api_format.as_str())
== "openai:chat"
@@ -2992,6 +3100,45 @@ async fn provider_query_execute_standard_test_candidate(
);
provider_request_body
}
"openai:search" => {
let Some(mut provider_request_body) =
crate::provider_transport::build_same_format_provider_request_body(
crate::provider_transport::SameFormatProviderRequestBodyInput {
body_json: &request_body,
mapped_model: request_model,
client_api_format,
provider_api_format,
source_model: request_body.get("model").and_then(Value::as_str),
family: crate::provider_transport::SameFormatProviderFamily::Standard,
body_rules: transport.endpoint.body_rules.as_ref(),
request_headers: Some(&incoming_request_headers),
upstream_is_stream,
force_body_stream_field: require_body_stream_field,
kiro_auth_config: None,
is_claude_code: false,
enable_model_directives: false,
},
)
else {
return Ok(provider_query_skipped_execution_outcome(
request_body.clone(),
format!("Provider request body could not be built for {provider_api_format}"),
));
};
if let Err(err) = crate::provider_transport::apply_transport_request_body_semantics(
&mut provider_request_body,
&transport,
normalized_provider_api_format.as_str(),
) {
return Ok(provider_query_skipped_execution_outcome(
provider_request_body,
format!(
"Provider request body is not compatible with transport semantics: {err}"
),
));
}
provider_request_body
}
"openai:embedding"
| "gemini:embedding"
| "jina:embedding"
@@ -3046,6 +3193,37 @@ async fn provider_query_execute_standard_test_candidate(
upstream_is_stream,
require_body_stream_field,
);
let source_model = provider_query_request_body_model(&request_body, request_model);
let codex_model_capabilities = crate::ai_serving::codex_model_capabilities_for_transport(
&transport,
provider_api_format,
request_model,
source_model,
);
if matches!(
normalized_provider_api_format.as_str(),
"openai:chat" | "openai:responses" | "openai:responses:compact" | "openai:search"
) && crate::ai_serving::finalize_openai_provider_request_with_codex_model_capabilities(
&mut provider_request_body,
crate::ai_serving::OpenAiProviderRequestFinalization {
source_api_format: client_api_format,
provider_api_format,
provider_type: transport.provider.provider_type.as_str(),
provider_model: request_model,
source_model,
body_rules: transport.endpoint.body_rules.as_ref(),
upstream_is_stream,
require_body_stream_field,
},
codex_model_capabilities.as_ref(),
)
.is_err()
{
return Ok(provider_query_skipped_execution_outcome(
provider_request_body,
"Provider request body violates the OpenAI provider contract",
));
}
if crate::provider_transport::is_gemini_cli_provider_transport(&transport)
&& normalized_provider_api_format == "gemini:generate_content"
{
@@ -3122,6 +3300,7 @@ async fn provider_query_execute_standard_test_candidate(
"openai:chat"
| "openai:responses"
| "openai:responses:compact"
| "openai:search"
| "claude:messages"
| "gemini:generate_content"
| "openai:embedding"
@@ -3137,6 +3316,7 @@ async fn provider_query_execute_standard_test_candidate(
"openai:chat"
| "openai:responses"
| "openai:responses:compact"
| "openai:search"
| "openai:embedding"
| "jina:embedding"
| "doubao:embedding"
@@ -3202,7 +3382,7 @@ async fn provider_query_execute_standard_test_candidate(
&BTreeMap::new(),
Some("application/json"),
),
"openai:responses" | "openai:responses:compact" => {
"openai:responses" | "openai:responses:compact" | "openai:search" => {
crate::provider_transport::auth::build_complete_passthrough_headers_with_auth(
&parts.headers,
auth_header.as_deref().unwrap_or_default(),
@@ -3268,8 +3448,10 @@ async fn provider_query_execute_standard_test_candidate(
response_body: None,
});
}
if crate::ai_serving::is_openai_responses_format(provider_api_format) {
crate::ai_serving::apply_codex_openai_responses_special_headers(
if crate::ai_serving::is_openai_responses_family_format(provider_api_format)
|| crate::ai_serving::api_format_alias_matches(provider_api_format, "openai:search")
{
crate::ai_serving::apply_codex_openai_special_headers(
&mut request_headers,
&provider_request_body,
&parts.headers,
@@ -3278,9 +3460,18 @@ async fn provider_query_execute_standard_test_candidate(
Some(trace_id),
transport.key.decrypted_auth_config.as_deref(),
);
crate::provider_transport::apply_local_auth_config_header_overrides(
let final_provider_model = provider_request_body
.get("model")
.and_then(Value::as_str)
.unwrap_or(request_model);
crate::ai_serving::apply_codex_openai_responses_lite_header_for_request_body_with_capabilities(
&mut request_headers,
transport.key.decrypted_auth_config.as_deref(),
Some(&provider_request_body),
transport.provider.provider_type.as_str(),
provider_api_format,
final_provider_model,
source_model,
codex_model_capabilities.as_ref(),
);
}
if !uses_vertex_query_auth {
apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/adapter.rs
+5
View File
@@ -33,6 +33,8 @@ pub(super) fn provider_query_standard_test_client_api_format(
let normalized_api_format = crate::ai_serving::normalize_api_format_alias(provider_api_format);
if normalized_api_format == "openai:responses:compact" {
"openai:responses:compact"
} else if normalized_api_format == "openai:search" {
"openai:search"
} else if crate::ai_serving::is_embedding_api_format(&normalized_api_format) {
"openai:embedding"
} else if crate::ai_serving::is_rerank_api_format(&normalized_api_format) {
@@ -71,6 +73,7 @@ pub(super) fn provider_query_standard_test_unsupported_reason(
}
"openai:responses"
| "openai:responses:compact"
| "openai:search"
| "claude:messages"
| "openai:embedding"
| "jina:embedding"
@@ -261,6 +264,7 @@ pub(super) fn provider_query_test_adapter_for_provider_api_format(
"openai:chat"
| "openai:responses"
| "openai:responses:compact"
| "openai:search"
| "claude:messages"
| "gemini:generate_content"
| "gemini:interactions"
@@ -365,6 +369,7 @@ pub(super) fn provider_query_transport_supports_model_test_execution(
}
"openai:responses"
| "openai:responses:compact"
| "openai:search"
| "openai:embedding"
| "jina:embedding"
| "doubao:embedding"
apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/capabilities.rs
+6 -2
View File
@@ -9,16 +9,19 @@ pub(super) struct ProviderQueryOpenAiImageTestCapability(AdminProviderOpenAiImag
pub(super) fn provider_query_openai_image_test_capability(
provider_type: &str,
provider_model: Option<&str>,
) -> ProviderQueryOpenAiImageTestCapability {
ProviderQueryOpenAiImageTestCapability(admin_provider_openai_image_test_capability(
provider_type,
provider_model,
))
}
pub(super) fn provider_query_openai_image_normalize_options(
provider_type: &str,
provider_model: Option<&str>,
) -> crate::ai_serving::OpenAiImageNormalizeOptions {
admin_provider_openai_image_normalize_options(provider_type)
admin_provider_openai_image_normalize_options(provider_type, provider_model)
}
pub(super) fn provider_query_openai_image_requested_count(request_body: &Value) -> Option<u64> {
@@ -35,9 +38,10 @@ pub(super) fn provider_query_openai_image_requested_count(request_body: &Value)
pub(super) fn provider_query_openai_image_normalize_failure_message(
provider_type: &str,
provider_model: Option<&str>,
request_body: &Value,
) -> String {
let capability = provider_query_openai_image_test_capability(provider_type);
let capability = provider_query_openai_image_test_capability(provider_type, provider_model);
if provider_query_openai_image_requested_count(request_body)
.is_some_and(|value| !capability.0.supports_generation_count(value))
{
apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/model_mapping.rs
+57 -2
View File
@@ -101,8 +101,7 @@ fn provider_query_model_mapping_matches_endpoint(
) -> bool {
let api_format_matches = mapping.api_formats.as_ref().is_none_or(|api_formats| {
api_formats.iter().any(|value| {
aether_scheduler_core::normalize_api_format(value)
== aether_scheduler_core::normalize_api_format(&endpoint.api_format)
crate::ai_serving::api_format_permission_covers(value, &endpoint.api_format)
})
});
if !api_format_matches {
@@ -214,6 +213,7 @@ fn provider_query_parse_embedded_provider_model_mappings(
priority: 1,
api_formats: None,
endpoint_ids: None,
operations: None,
}]))
}
@@ -238,6 +238,7 @@ fn provider_query_parse_provider_model_mappings_array(
priority: 1,
api_formats: None,
endpoint_ids: None,
operations: None,
});
}
}
@@ -294,6 +295,11 @@ fn provider_query_parse_provider_model_mapping_object_lenient(
object.get("endpoint_ids"),
"models.provider_model_mappings.endpoint_ids",
)?;
let operations = provider_query_parse_mapping_string_list(
object.get("operations"),
"models.provider_model_mappings.operations",
)?
.and_then(provider_query_normalize_request_operations);
Ok(Some(StoredProviderModelMapping {
name: name.to_string(),
@@ -304,9 +310,19 @@ fn provider_query_parse_provider_model_mapping_object_lenient(
})?,
api_formats,
endpoint_ids,
operations,
}))
}
fn provider_query_normalize_request_operations(values: Vec<String>) -> Option<Vec<String>> {
let operations = values
.into_iter()
.map(|value| value.trim().to_ascii_lowercase())
.filter(|value| !value.is_empty())
.collect::<Vec<_>>();
(!operations.is_empty()).then_some(operations)
}
fn provider_query_parse_mapping_string_list(
value: Option<&Value>,
field_name: &str,
@@ -367,3 +383,42 @@ fn provider_query_parse_mapping_string_list_array(
}
Ok(parsed)
}
#[cfg(test)]
mod tests {
use super::*;
fn sample_endpoint(api_format: &str) -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
"endpoint-1".to_string(),
"provider-1".to_string(),
api_format.to_string(),
None,
None,
true,
)
.expect("endpoint should build")
}
fn sample_mapping(api_format: &str) -> StoredProviderModelMapping {
StoredProviderModelMapping {
name: "gpt-5.6-luna".to_string(),
priority: 1,
api_formats: Some(vec![api_format.to_string()]),
endpoint_ids: None,
operations: None,
}
}
#[test]
fn responses_mapping_scope_covers_search_in_one_direction() {
assert!(provider_query_model_mapping_matches_endpoint(
&sample_mapping("openai:responses"),
&sample_endpoint("openai:search"),
));
assert!(!provider_query_model_mapping_matches_endpoint(
&sample_mapping("openai:search"),
&sample_endpoint("openai:responses"),
));
}
}
apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/tests.rs
+164 -3
View File
@@ -337,6 +337,11 @@ fn provider_query_standard_test_resolves_codex_responses_upstream_streaming() {
"codex",
"openai:responses:compact",
));
assert!(!provider_query_resolve_standard_test_upstream_is_stream(
Some(&json!({"upstream_stream_policy": "force_stream"})),
"codex",
"openai:responses:compact",
));
assert!(!provider_query_resolve_standard_test_upstream_is_stream(
None,
"custom",
@@ -470,6 +475,118 @@ fn provider_query_compact_test_request_body_defaults_to_responses_input() {
assert!(body.get("messages").is_none());
}
#[test]
fn provider_query_search_test_request_body_defaults_to_typed_search_input() {
let payload = json!({"message": "find current documentation"});
let client_api_format = provider_query_standard_test_client_api_format("openai:search");
let body = provider_query_build_test_request_body_for_api_format(
&payload,
"gpt-5.6-sol",
"/api/admin/provider-query/test-model",
client_api_format,
);
assert_eq!(client_api_format, "openai:search");
assert!(body["id"]
.as_str()
.is_some_and(|value| value.starts_with("aether-model-test-")));
assert_eq!(body["model"], json!("gpt-5.6-sol"));
assert_eq!(body["input"], json!("find current documentation"));
assert_eq!(
body["commands"]["search_query"][0]["q"],
json!("find current documentation")
);
assert_eq!(body["max_output_tokens"], json!(256));
assert!(body.get("messages").is_none());
assert!(body.get("stream").is_none());
}
#[test]
fn provider_query_search_test_completes_missing_protocol_fields() {
let payload = json!({
"request_body": {
"model": "gpt-5.6-luna",
"input": "find release notes"
}
});
let body = provider_query_build_test_request_body_for_api_format_with_search_session(
&payload,
"gpt-5.6-luna",
"/api/admin/provider-query/test-model",
"openai:search",
Some("trace-model-test-1"),
);
assert_eq!(body["id"], json!("aether-model-test-trace-model-test-1"));
assert_eq!(
body["commands"]["search_query"][0]["q"],
json!("find release notes")
);
assert_eq!(body["max_output_tokens"], json!(256));
}
#[test]
fn provider_query_search_test_preserves_a_non_empty_client_session_id() {
let payload = json!({
"request_body": {
"id": "client-search-session",
"model": "gpt-5.6-luna",
"input": "find release notes"
}
});
let body = provider_query_build_test_request_body_for_api_format_with_search_session(
&payload,
"gpt-5.6-luna",
"/api/admin/provider-query/test-model",
"openai:search",
Some("trace-model-test-1"),
);
assert_eq!(body["id"], json!("client-search-session"));
}
#[test]
fn provider_query_search_success_requires_non_empty_output() {
fn result(body: Value) -> aether_contracts::ExecutionResult {
aether_contracts::ExecutionResult {
request_id: "provider-search-test".to_string(),
candidate_id: Some("candidate-0".to_string()),
status_code: 200,
headers: BTreeMap::new(),
body: Some(aether_contracts::ResponseBody {
json_body: Some(body),
body_bytes_b64: None,
}),
telemetry: None,
error: None,
}
}
assert!(provider_query_standard_execution_response_body(
"openai:search",
&result(json!({})),
None,
)
.is_none());
assert!(provider_query_standard_execution_response_body(
"openai:search",
&result(json!({"output": " "})),
None,
)
.is_none());
assert_eq!(
provider_query_standard_execution_response_body(
"openai:search",
&result(json!({"output": "search result", "encrypted_output": "ciphertext"})),
None,
),
Some(json!({"output": "search result", "encrypted_output": "ciphertext"}))
);
}
#[test]
fn provider_query_embedding_test_request_body_defaults_to_embedding_input() {
let payload = json!({"message": "hello from embedding"});
@@ -632,6 +749,10 @@ fn provider_query_test_adapter_routes_fixed_provider_endpoint_types() {
provider_query_test_adapter_for_provider_api_format("codex", "openai:responses:compact"),
Some(ProviderQueryTestAdapter::Standard)
);
assert_eq!(
provider_query_test_adapter_for_provider_api_format("codex", "openai:search"),
Some(ProviderQueryTestAdapter::Standard)
);
assert_eq!(
provider_query_test_adapter_for_provider_api_format("chatgpt_web", "openai:image"),
Some(ProviderQueryTestAdapter::OpenAiImage)
@@ -717,6 +838,10 @@ fn provider_query_endpoint_priority_prefers_text_before_cli_and_image() {
provider_query_model_test_endpoint_priority("codex", "openai:responses:compact"),
Some(1)
);
assert_eq!(
provider_query_model_test_endpoint_priority("codex", "openai:search"),
Some(1)
);
assert_eq!(
provider_query_model_test_endpoint_priority("chatgpt_web", "openai:image"),
Some(2)
@@ -997,7 +1122,7 @@ fn provider_query_grok_image_test_allows_multi_generation_count() {
&parts,
&body,
None,
provider_query_openai_image_normalize_options("grok"),
provider_query_openai_image_normalize_options("grok", Some("grok-imagine-image")),
)
.expect("grok image model tests should allow multi-image generation");
let provider_body = crate::ai_serving::build_openai_image_provider_request_body(&normalized);
@@ -1065,12 +1190,48 @@ fn provider_query_non_grok_image_test_keeps_single_generation_boundary() {
&parts,
&body,
None,
provider_query_openai_image_normalize_options("chatgpt_web"),
provider_query_openai_image_normalize_options("chatgpt_web", Some("gpt-image-2")),
)
.is_none()
);
assert_eq!(
provider_query_openai_image_normalize_failure_message("chatgpt_web", &body),
provider_query_openai_image_normalize_failure_message(
"chatgpt_web",
Some("gpt-image-2"),
&body,
),
"Provider request body could not be normalized for openai:image: selected provider supports n=1..1 for generation"
);
}
#[test]
fn provider_query_dall_e_3_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": "dall-e-3",
"prompt": "draw",
"n": 2
});
assert!(
crate::ai_serving::normalize_openai_image_request_with_options(
&parts,
&body,
None,
provider_query_openai_image_normalize_options("openai", Some("dall-e-3")),
)
.is_none()
);
assert_eq!(
provider_query_openai_image_normalize_failure_message(
"openai",
Some("dall-e-3"),
&body,
),
"Provider request body could not be normalized for openai:image: selected provider supports n=1..1 for generation"
);
}
apps/aether-gateway/src/handlers/admin/provider/shared/model_test_capabilities.rs
+17 -4
View File
@@ -1,5 +1,6 @@
use crate::image_capabilities::{
openai_image_normalize_options_for_provider, openai_image_provider_max_generation_count,
openai_image_normalize_options_for_provider,
openai_image_provider_max_generation_count_for_model,
};
use serde_json::{json, Value};
@@ -24,16 +25,21 @@ impl AdminProviderOpenAiImageTestCapability {
pub(crate) fn admin_provider_openai_image_test_capability(
provider_type: &str,
provider_model: Option<&str>,
) -> AdminProviderOpenAiImageTestCapability {
AdminProviderOpenAiImageTestCapability {
max_generation_count: openai_image_provider_max_generation_count(provider_type),
max_generation_count: openai_image_provider_max_generation_count_for_model(
provider_type,
provider_model,
),
}
}
pub(crate) fn admin_provider_openai_image_normalize_options(
provider_type: &str,
provider_model: Option<&str>,
) -> crate::ai_serving::OpenAiImageNormalizeOptions {
openai_image_normalize_options_for_provider(provider_type)
openai_image_normalize_options_for_provider(provider_type, provider_model)
}
pub(crate) fn admin_provider_model_test_capabilities_payload(
@@ -47,7 +53,7 @@ pub(crate) fn admin_provider_model_test_capabilities_payload(
provider_type.eq_ignore_ascii_case("grok") && model_id == GROK_IMAGE_EDIT_MODEL_ID;
let openai_image = if supports_image_generation {
Some(json!({
"max_generation_count": admin_provider_openai_image_test_capability(provider_type).max_generation_count,
"max_generation_count": admin_provider_openai_image_test_capability(provider_type, Some(model_id)).max_generation_count,
"supports_generation": !is_grok_image_edit,
"supports_edit": is_grok_image_edit,
}))
@@ -105,6 +111,13 @@ mod tests {
assert!(payload["openai:image"].is_null());
}
#[test]
fn dall_e_3_reports_its_model_specific_generation_limit() {
let payload = admin_provider_model_test_capabilities_payload("openai", "dall-e-3", true);
assert_eq!(payload["openai:image"]["max_generation_count"], 1);
}
#[test]
fn grok_image_support_uses_catalog_model_ids_not_global_fallback() {
assert!(admin_provider_model_supports_image_generation(
apps/aether-gateway/src/handlers/admin/provider/write/keys/create.rs
+2 -2
View File
@@ -2,7 +2,7 @@ use crate::handlers::admin::provider::shared::payloads::AdminProviderKeyCreateRe
use crate::handlers::admin::provider::write::normalize::{
normalize_allow_auth_channel_mismatch_formats, normalize_api_format_json_object_keys,
normalize_api_format_list, normalize_auth_type, normalize_auth_type_by_format,
normalize_max_probe_interval_minutes, validate_vertex_api_formats,
normalize_max_probe_interval_minutes, normalize_rate_multipliers, validate_vertex_api_formats,
};
use crate::handlers::admin::request::AdminAppState;
use crate::handlers::admin::shared::{
@@ -165,7 +165,7 @@ pub(crate) async fn build_admin_create_provider_key_record(
},
encrypted_api_key,
encrypted_auth_config,
normalize_api_format_json_object_keys(payload.rate_multipliers, "rate_multipliers")?,
normalize_rate_multipliers(payload.rate_multipliers)?,
None,
normalize_string_list(payload.allowed_models).map(|value| json!(value)),
None,
apps/aether-gateway/src/handlers/admin/provider/write/keys/update.rs
+2 -3
View File
@@ -2,7 +2,7 @@ use crate::handlers::admin::provider::shared::payloads::AdminProviderKeyUpdatePa
use crate::handlers::admin::provider::write::normalize::{
normalize_allow_auth_channel_mismatch_formats, normalize_api_format_json_object_keys,
normalize_api_format_list, normalize_auth_type, normalize_auth_type_by_format,
normalize_max_probe_interval_minutes, validate_vertex_api_formats,
normalize_max_probe_interval_minutes, normalize_rate_multipliers, validate_vertex_api_formats,
};
use crate::handlers::admin::request::AdminAppState;
use crate::handlers::admin::shared::{
@@ -276,8 +276,7 @@ pub(crate) fn build_admin_update_provider_key_record_with_existing_keys(
updated.name = trimmed.to_string();
}
if fields.contains("rate_multipliers") {
updated.rate_multipliers =
normalize_api_format_json_object_keys(payload.rate_multipliers, "rate_multipliers")?;
updated.rate_multipliers = normalize_rate_multipliers(payload.rate_multipliers)?;
}
if let Some(internal_priority) = payload.internal_priority {
updated.internal_priority = internal_priority;
apps/aether-gateway/src/handlers/admin/provider/write/normalize.rs
+41 -1
View File
@@ -42,6 +42,31 @@ pub(crate) fn normalize_api_format_json_object_keys(
Ok(Some(serde_json::Value::Object(normalized)))
}
pub(crate) fn normalize_rate_multipliers(
value: Option<serde_json::Value>,
) -> Result<Option<serde_json::Value>, String> {
let Some(value) = normalize_json_like_object(value, "rate_multipliers")? else {
return Ok(None);
};
let serde_json::Value::Object(map) = value else {
return Ok(Some(value));
};
let mut normalized = serde_json::Map::new();
for (key, value) in map {
let canonical = crate::ai_serving::normalize_api_format_alias(&key);
let multiplier = value
.as_f64()
.filter(|value| value.is_finite() && *value >= 0.0)
.ok_or_else(|| format!("rate_multipliers.{canonical} 必须是大于或等于 0 的有限数值"))?;
normalized.insert(canonical, serde_json::Value::from(multiplier));
}
if normalized.is_empty() {
Ok(None)
} else {
Ok(Some(serde_json::Value::Object(normalized)))
}
}
pub(crate) fn normalize_auth_type_by_format(
value: Option<serde_json::Value>,
field_name: &str,
@@ -212,7 +237,7 @@ mod tests {
normalize_allow_auth_channel_mismatch_formats, normalize_api_format_json_object_keys,
normalize_api_format_list, normalize_auth_type, normalize_auth_type_by_format,
normalize_chat_pii_redaction_config, normalize_pool_advanced_config,
normalize_provider_type_input, validate_vertex_api_formats,
normalize_provider_type_input, normalize_rate_multipliers, validate_vertex_api_formats,
};
use serde_json::json;
@@ -224,6 +249,21 @@ mod tests {
);
}
#[test]
fn rate_multipliers_require_non_negative_finite_numbers() {
assert_eq!(
normalize_rate_multipliers(Some(json!({" OPENAI:RESPONSES ": 1.25})))
.expect("valid multiplier should normalize"),
Some(json!({"openai:responses": 1.25}))
);
for value in [
json!({"openai:responses": -0.1}),
json!({"openai:responses": "1.0"}),
] {
assert!(normalize_rate_multipliers(Some(value)).is_err());
}
}
#[test]
fn normalize_pool_advanced_rejects_legacy_booleans() {
assert_eq!(
apps/aether-gateway/src/handlers/admin/provider/write/provider.rs
+69 -1
View File
@@ -3,11 +3,79 @@ mod endpoint;
mod template;
mod update;
fn normalize_provider_request_timeout(value: Option<f64>) -> Result<Option<f64>, String> {
let max = aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_SECS as f64;
match value {
Some(value) if (1.0..=max).contains(&value) => Ok(Some(value)),
Some(_) => Err(format!(
"request_timeout 必须是 1 到 {} 之间的数字",
aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_SECS
)),
None => Ok(None),
}
}
fn normalize_provider_stream_first_byte_timeout(value: Option<f64>) -> Result<Option<f64>, String> {
let max = aether_contracts::MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_SECS as f64;
match value {
Some(value) if (1.0..=max).contains(&value) => Ok(Some(value)),
Some(_) => Err(format!(
"stream_first_byte_timeout 必须是 1 到 {} 之间的数字",
aether_contracts::MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_SECS
)),
None => Ok(None),
}
}
pub(crate) use self::create::build_admin_create_provider_record;
pub(crate) use self::endpoint::build_admin_fixed_provider_endpoint_record;
pub(crate) use self::template::{
apply_admin_fixed_provider_endpoint_template_overrides,
reconcile_admin_fixed_provider_template_endpoints,
reconcile_admin_fixed_provider_template_keys,
reconcile_admin_fixed_provider_template_endpoints_after_update,
};
pub(crate) use self::update::build_admin_update_provider_record;
#[cfg(test)]
mod tests {
use super::{normalize_provider_request_timeout, normalize_provider_stream_first_byte_timeout};
#[test]
fn provider_request_timeout_accepts_the_execution_protocol_range() {
assert_eq!(normalize_provider_request_timeout(Some(1.0)), Ok(Some(1.0)));
assert_eq!(
normalize_provider_request_timeout(Some(
aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_SECS as f64,
)),
Ok(Some(
aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_SECS as f64,
))
);
assert!(normalize_provider_request_timeout(Some(0.0)).is_err());
assert!(normalize_provider_request_timeout(Some(
aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_SECS as f64 + 1.0,
))
.is_err());
assert!(normalize_provider_request_timeout(Some(f64::NAN)).is_err());
}
#[test]
fn provider_stream_first_byte_timeout_keeps_its_protocol_range() {
assert_eq!(
normalize_provider_stream_first_byte_timeout(Some(1.0)),
Ok(Some(1.0))
);
assert_eq!(
normalize_provider_stream_first_byte_timeout(Some(
aether_contracts::MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_SECS as f64,
)),
Ok(Some(
aether_contracts::MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_SECS as f64,
))
);
assert!(normalize_provider_stream_first_byte_timeout(Some(
aether_contracts::MAX_EXECUTION_STREAM_FIRST_BYTE_TIMEOUT_SECS as f64 + 1.0,
))
.is_err());
}
}
apps/aether-gateway/src/handlers/admin/provider/write/provider/create.rs
+3 -10
View File
@@ -107,16 +107,9 @@ pub(crate) async fn build_admin_create_provider_record(
None => Some(2),
};
let proxy = normalize_json_object(payload.proxy, "proxy")?;
let stream_first_byte_timeout_secs = match payload.stream_first_byte_timeout {
Some(value) if (1.0..=300.0).contains(&value) => Some(value),
Some(_) => return Err("stream_first_byte_timeout 必须是 1 到 300 之间的数字".to_string()),
None => None,
};
let request_timeout_secs = match payload.request_timeout {
Some(value) if (1.0..=600.0).contains(&value) => Some(value),
Some(_) => return Err("request_timeout 必须是 1 到 600 之间的数字".to_string()),
None => None,
};
let stream_first_byte_timeout_secs =
super::normalize_provider_stream_first_byte_timeout(payload.stream_first_byte_timeout)?;
let request_timeout_secs = super::normalize_provider_request_timeout(payload.request_timeout)?;
let mut config_map = normalize_json_object(payload.config, "config")?
.and_then(|value| value.as_object().cloned())
apps/aether-gateway/src/handlers/admin/provider/write/provider/template.rs
+204 -48
View File
@@ -1,11 +1,11 @@
use super::endpoint::{
build_admin_fixed_provider_endpoint_defaults, build_admin_fixed_provider_endpoint_record,
AdminFixedProviderEndpointDefaults,
};
use crate::handlers::admin::request::AdminAppState;
use crate::provider_key_auth::provider_key_is_oauth_managed;
use crate::GatewayError;
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
StoredProviderCatalogEndpoint, StoredProviderCatalogProvider,
};
use aether_provider_transport::provider_types::{
fixed_provider_template, FixedProviderEndpointTemplate, FixedProviderTemplate,
@@ -36,6 +36,30 @@ struct FixedProviderEndpointMetadata {
pub(crate) async fn reconcile_admin_fixed_provider_template_endpoints(
state: &AdminAppState<'_>,
provider: &StoredProviderCatalogProvider,
) -> Result<(), GatewayError> {
reconcile_admin_fixed_provider_template_endpoints_with_adoption_provider(
state, provider, provider,
)
.await
}
pub(crate) async fn reconcile_admin_fixed_provider_template_endpoints_after_update(
state: &AdminAppState<'_>,
previous_provider: &StoredProviderCatalogProvider,
updated_provider: &StoredProviderCatalogProvider,
) -> Result<(), GatewayError> {
reconcile_admin_fixed_provider_template_endpoints_with_adoption_provider(
state,
updated_provider,
previous_provider,
)
.await
}
async fn reconcile_admin_fixed_provider_template_endpoints_with_adoption_provider(
state: &AdminAppState<'_>,
provider: &StoredProviderCatalogProvider,
adoption_provider: &StoredProviderCatalogProvider,
) -> Result<(), GatewayError> {
let Some(template) = state.fixed_provider_template(&provider.provider_type) else {
return Ok(());
@@ -53,8 +77,9 @@ pub(crate) async fn reconcile_admin_fixed_provider_template_endpoints(
match existing_endpoint {
Some(existing_endpoint) => {
matched_endpoint_ids.insert(existing_endpoint.id.clone());
let updated = reconcile_fixed_provider_endpoint(
let updated = reconcile_fixed_provider_endpoint_with_adoption_provider(
provider,
adoption_provider,
existing_endpoint,
template,
endpoint_template,
@@ -115,31 +140,6 @@ pub(crate) async fn reconcile_admin_fixed_provider_template_endpoints(
Ok(())
}
pub(crate) async fn reconcile_admin_fixed_provider_template_keys(
state: &AdminAppState<'_>,
provider: &StoredProviderCatalogProvider,
) -> Result<(), GatewayError> {
let Some(_) = state.fixed_provider_template(&provider.provider_type) else {
return Ok(());
};
let existing_keys = state
.list_provider_catalog_keys_by_provider_ids(std::slice::from_ref(&provider.id))
.await?;
for existing_key in existing_keys {
let Some(updated_key) = reconcile_fixed_provider_key(provider, &existing_key) else {
continue;
};
let Some(_) = state.update_provider_catalog_key(&updated_key).await? else {
return Err(GatewayError::Internal(
"provider catalog key writer unavailable".to_string(),
));
};
}
Ok(())
}
pub(crate) fn apply_admin_fixed_provider_endpoint_template_overrides(
provider: &StoredProviderCatalogProvider,
existing_endpoint: &StoredProviderCatalogEndpoint,
@@ -156,8 +156,14 @@ pub(crate) fn apply_admin_fixed_provider_endpoint_template_overrides(
let defaults =
build_admin_fixed_provider_endpoint_defaults(provider, template, endpoint_template)?;
let mut metadata = fixed_provider_endpoint_metadata(existing_endpoint)
.unwrap_or_else(|| managed_fixed_provider_endpoint_metadata(template, endpoint_template));
let mut metadata = fixed_provider_endpoint_metadata(existing_endpoint).unwrap_or_else(|| {
adopt_fixed_provider_endpoint_metadata(
existing_endpoint,
&defaults,
template,
endpoint_template,
)
});
let mut overrides = metadata.overrides.clone();
sync_override_if_changed(
@@ -250,17 +256,48 @@ pub(crate) fn apply_admin_fixed_provider_endpoint_template_overrides(
Ok(())
}
#[cfg(test)]
fn reconcile_fixed_provider_endpoint(
provider: &StoredProviderCatalogProvider,
existing_endpoint: &StoredProviderCatalogEndpoint,
template: &FixedProviderTemplate,
endpoint_template: &FixedProviderEndpointTemplate,
) -> Result<StoredProviderCatalogEndpoint, String> {
reconcile_fixed_provider_endpoint_with_adoption_provider(
provider,
provider,
existing_endpoint,
template,
endpoint_template,
)
}
fn reconcile_fixed_provider_endpoint_with_adoption_provider(
provider: &StoredProviderCatalogProvider,
adoption_provider: &StoredProviderCatalogProvider,
existing_endpoint: &StoredProviderCatalogEndpoint,
template: &FixedProviderTemplate,
endpoint_template: &FixedProviderEndpointTemplate,
) -> Result<StoredProviderCatalogEndpoint, String> {
let defaults =
build_admin_fixed_provider_endpoint_defaults(provider, template, endpoint_template)?;
let mut updated = existing_endpoint.clone();
let metadata = fixed_provider_endpoint_metadata(existing_endpoint)
.unwrap_or_else(|| managed_fixed_provider_endpoint_metadata(template, endpoint_template));
let metadata = match fixed_provider_endpoint_metadata(existing_endpoint) {
Some(metadata) => metadata,
None => {
let adoption_defaults = build_admin_fixed_provider_endpoint_defaults(
adoption_provider,
template,
endpoint_template,
)?;
adopt_fixed_provider_endpoint_metadata(
existing_endpoint,
&adoption_defaults,
template,
endpoint_template,
)
}
};
updated.api_format = defaults.api_format.clone();
updated.api_family = Some(defaults.api_family.clone());
@@ -455,6 +492,77 @@ fn managed_fixed_provider_endpoint_metadata(
}
}
fn adopt_fixed_provider_endpoint_metadata(
existing_endpoint: &StoredProviderCatalogEndpoint,
defaults: &AdminFixedProviderEndpointDefaults,
template: &FixedProviderTemplate,
endpoint_template: &FixedProviderEndpointTemplate,
) -> FixedProviderEndpointMetadata {
let mut metadata = managed_fixed_provider_endpoint_metadata(template, endpoint_template);
sync_override(
&mut metadata.overrides,
OVERRIDE_BASE_URL,
&existing_endpoint.base_url,
&defaults.base_url,
);
sync_override(
&mut metadata.overrides,
OVERRIDE_CUSTOM_PATH,
&existing_endpoint.custom_path,
&defaults.custom_path,
);
sync_override(
&mut metadata.overrides,
OVERRIDE_HEADER_RULES,
&existing_endpoint.header_rules,
&defaults.header_rules,
);
sync_override(
&mut metadata.overrides,
OVERRIDE_BODY_RULES,
&existing_endpoint.body_rules,
&defaults.body_rules,
);
sync_override(
&mut metadata.overrides,
OVERRIDE_MAX_RETRIES,
&existing_endpoint.max_retries,
&defaults.max_retries,
);
sync_override(
&mut metadata.overrides,
OVERRIDE_IS_ACTIVE,
&existing_endpoint.is_active,
&defaults.is_active,
);
sync_override(
&mut metadata.overrides,
OVERRIDE_PROXY,
&existing_endpoint.proxy,
&defaults.proxy,
);
sync_override(
&mut metadata.overrides,
OVERRIDE_FORMAT_ACCEPTANCE_CONFIG,
&existing_endpoint.format_acceptance_config,
&defaults.format_acceptance_config,
);
let existing_config = endpoint_config_without_metadata(existing_endpoint.config.as_ref());
for (key, desired) in fixed_provider_endpoint_config_defaults(endpoint_template) {
let Some(actual) = existing_config.get(&key) else {
continue;
};
sync_override(
&mut metadata.overrides,
&config_override_key(&key),
actual,
&desired,
);
}
metadata
}
fn upsert_fixed_provider_endpoint_metadata(
endpoint: &mut StoredProviderCatalogEndpoint,
metadata: &FixedProviderEndpointMetadata,
@@ -513,22 +621,6 @@ fn current_unix_secs() -> u64 {
.unwrap_or(0)
}
fn reconcile_fixed_provider_key(
provider: &StoredProviderCatalogProvider,
existing_key: &StoredProviderCatalogKey,
) -> Option<StoredProviderCatalogKey> {
if !provider_key_is_oauth_managed(existing_key, &provider.provider_type)
|| existing_key.api_formats.is_none()
{
return None;
}
let mut updated = existing_key.clone();
updated.api_formats = None;
updated.updated_at_unix_secs = Some(current_unix_secs());
Some(updated)
}
fn sync_override<T>(overrides: &mut BTreeSet<String>, key: &str, actual: &T, desired: &T)
where
T: PartialEq,
@@ -624,4 +716,68 @@ mod tests {
.expect("endpoint should reconcile");
assert_eq!(reconciled.base_url, "http://127.0.0.1:18181/v1");
}
#[test]
fn fixed_provider_endpoint_reconcile_adopts_existing_customization() {
let provider = sample_codex_provider();
let template = fixed_provider_template("codex").expect("codex template should exist");
let endpoint_template = template
.endpoints
.iter()
.find(|endpoint| endpoint.api_format == "openai:responses")
.expect("responses endpoint template should exist");
let mut existing = sample_codex_endpoint("http://127.0.0.1:18181/backend-api/codex");
existing.is_active = false;
existing.max_retries = Some(9);
existing.proxy = Some(serde_json::json!({"url": "http://proxy.internal:8080"}));
existing.config = Some(serde_json::json!({
"upstream_stream_policy": "force_non_stream",
"custom_transport_option": true
}));
let reconciled =
reconcile_fixed_provider_endpoint(&provider, &existing, template, endpoint_template)
.expect("endpoint should reconcile");
assert_eq!(
reconciled.base_url,
"http://127.0.0.1:18181/backend-api/codex"
);
assert!(!reconciled.is_active);
assert_eq!(reconciled.max_retries, Some(9));
assert_eq!(
reconciled.proxy,
Some(serde_json::json!({"url": "http://proxy.internal:8080"}))
);
assert_eq!(
reconciled
.config
.as_ref()
.and_then(|value| value.get("upstream_stream_policy")),
Some(&serde_json::json!("force_non_stream"))
);
assert_eq!(
reconciled
.config
.as_ref()
.and_then(|value| value.get("custom_transport_option")),
Some(&serde_json::json!(true))
);
let metadata = fixed_provider_endpoint_metadata(&reconciled)
.expect("fixed provider metadata should exist");
for key in [
"base_url",
"is_active",
"max_retries",
"proxy",
"config.upstream_stream_policy",
] {
assert!(metadata.overrides.contains(key), "missing override {key}");
}
let second =
reconcile_fixed_provider_endpoint(&provider, &reconciled, template, endpoint_template)
.expect("endpoint should reconcile idempotently");
assert_eq!(second, reconciled);
}
}
apps/aether-gateway/src/handlers/admin/provider/write/provider/update.rs
+4 -12
View File
@@ -202,21 +202,13 @@ pub(crate) async fn build_admin_update_provider_record(
}
if fields.contains("stream_first_byte_timeout") {
updated.stream_first_byte_timeout_secs = match payload.stream_first_byte_timeout {
Some(value) if (1.0..=300.0).contains(&value) => Some(value),
Some(_) => {
return Err("stream_first_byte_timeout 必须是 1 到 300 之间的数字".to_string());
}
None => None,
};
updated.stream_first_byte_timeout_secs =
super::normalize_provider_stream_first_byte_timeout(payload.stream_first_byte_timeout)?;
}
if fields.contains("request_timeout") {
updated.request_timeout_secs = match payload.request_timeout {
Some(value) if (1.0..=600.0).contains(&value) => Some(value),
Some(_) => return Err("request_timeout 必须是 1 到 600 之间的数字".to_string()),
None => None,
};
updated.request_timeout_secs =
super::normalize_provider_request_timeout(payload.request_timeout)?;
}
if fields.contains("enable_format_conversion") {
apps/aether-gateway/src/handlers/admin/request/models.rs
+6 -3
View File
@@ -15,7 +15,7 @@ use serde_json::json;
use std::collections::{BTreeMap, BTreeSet};
use uuid::Uuid;
fn normalize_provider_model_mappings_api_formats(
fn normalize_provider_model_mapping_scopes(
value: Option<serde_json::Value>,
) -> Option<serde_json::Value> {
let Some(mut value) = value else {
@@ -36,6 +36,9 @@ fn normalize_provider_model_mappings_api_formats(
normalize_provider_model_mapping_string_array_field(object, "endpoint_ids", |value| {
value.trim().to_string()
});
normalize_provider_model_mapping_string_array_field(object, "operations", |value| {
value.trim().to_ascii_lowercase()
});
}
Some(value)
}
@@ -155,7 +158,7 @@ impl<'a> AdminAppState<'a> {
"price_per_request",
)?;
let tiered_pricing = normalize_json_object(payload.tiered_pricing, "tiered_pricing")?;
let provider_model_mappings = normalize_provider_model_mappings_api_formats(
let provider_model_mappings = normalize_provider_model_mapping_scopes(
normalize_json_array(payload.provider_model_mappings, "provider_model_mappings")?,
);
let config = normalize_json_object(payload.config, "config")?;
@@ -241,7 +244,7 @@ impl<'a> AdminAppState<'a> {
existing.tiered_pricing.clone()
};
let provider_model_mappings = if fields.contains("provider_model_mappings") {
normalize_provider_model_mappings_api_formats(normalize_json_array(
normalize_provider_model_mapping_scopes(normalize_json_array(
payload.provider_model_mappings,
"provider_model_mappings",
)?)
apps/aether-gateway/src/handlers/admin/request/provider/builders.rs
+66
View File
@@ -1,5 +1,26 @@
use super::*;
fn validate_admin_endpoint_stream_policy(
api_format: &str,
config: Option<&serde_json::Value>,
) -> Result<(), String> {
if !crate::ai_serving::api_format_alias_matches(api_format, "openai:search") {
return Ok(());
}
let requested = config
.and_then(serde_json::Value::as_object)
.and_then(|config| {
config
.get("upstream_stream_policy")
.or_else(|| config.get("upstreamStreamPolicy"))
.or_else(|| config.get("upstream_stream"))
});
if requested.is_some_and(crate::handlers::public::admin_requested_force_stream) {
return Err("OpenAI Search 端点仅支持非流式上游请求".to_string());
}
Ok(())
}
impl<'a> AdminAppState<'a> {
pub(crate) async fn build_admin_keys_grouped_by_format_payload(
&self,
@@ -232,6 +253,7 @@ impl<'a> AdminAppState<'a> {
let (normalized_api_format, api_family, endpoint_kind) =
admin_endpoint_signature_parts(&payload.api_format)
.ok_or_else(|| format!("无效的 api_format: {}", payload.api_format))?;
validate_admin_endpoint_stream_policy(normalized_api_format, payload.config.as_ref())?;
let base_url = normalize_admin_base_url(&payload.base_url)?;
let existing_endpoints = self
@@ -337,6 +359,13 @@ impl<'a> AdminAppState<'a> {
&update_fields,
)?;
if fields.contains("config") {
validate_admin_endpoint_stream_policy(
existing_endpoint.api_format.as_str(),
updated.config.as_ref(),
)?;
}
if provider_type == "codex"
&& crate::ai_serving::is_openai_responses_format(&existing_endpoint.api_format)
{
@@ -389,3 +418,40 @@ impl<'a> AdminAppState<'a> {
Ok(updated)
}
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::validate_admin_endpoint_stream_policy;
#[test]
fn search_endpoint_rejects_explicit_streaming_policy_for_all_config_keys() {
for (api_format, key, value) in [
(
"openai:search",
"upstream_stream_policy",
json!("force_stream"),
),
("openai:search", "upstreamStreamPolicy", json!(true)),
("/v1/alpha/search", "upstream_stream", json!("sse")),
] {
let config = json!({(key): value});
assert!(validate_admin_endpoint_stream_policy(api_format, Some(&config),).is_err());
}
}
#[test]
fn search_endpoint_accepts_non_streaming_and_unrelated_config() {
assert!(validate_admin_endpoint_stream_policy(
"openai:search",
Some(&json!({"upstream_stream_policy": "force_non_stream"})),
)
.is_ok());
assert!(validate_admin_endpoint_stream_policy(
"openai:responses",
Some(&json!({"upstream_stream_policy": "force_stream"})),
)
.is_ok());
}
}
apps/aether-gateway/src/handlers/admin/request/provider/tasks.rs
+156 -21
View File
@@ -12,7 +12,7 @@ use axum::{
response::{IntoResponse, Response},
Json,
};
use serde_json::json;
use serde_json::{json, Map, Value};
use std::collections::{BTreeMap, BTreeSet};
use std::time::{SystemTime, UNIX_EPOCH};
@@ -107,9 +107,9 @@ impl<'a> AdminAppState<'a> {
let existing_keys = self
.list_provider_catalog_keys_by_provider_ids(std::slice::from_ref(&provider.id))
.await?;
let api_formats =
let available_api_formats =
admin_provider_pool_pure::admin_pool_resolved_api_formats(&endpoints, &existing_keys);
if api_formats.is_empty() {
if available_api_formats.is_empty() {
return Ok((
http::StatusCode::BAD_REQUEST,
Json(json!({ "detail": "Provider 没有可用 endpoint 或现有 key，无法推断 api_formats" })),
@@ -117,8 +117,74 @@ impl<'a> AdminAppState<'a> {
.into_response());
}
let proxy =
admin_provider_pool_pure::admin_pool_key_proxy_value(payload.proxy_node_id.as_deref());
let requested_api_formats = payload
.api_formats
.iter()
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
.collect::<BTreeSet<_>>()
.into_iter()
.collect::<Vec<_>>();
let available_api_format_set = available_api_formats
.iter()
.cloned()
.collect::<BTreeSet<_>>();
let api_formats = if requested_api_formats.is_empty() {
available_api_formats.clone()
} else {
if let Some(unsupported) = requested_api_formats
.iter()
.find(|value| !available_api_format_set.contains(*value))
{
return Ok((
http::StatusCode::BAD_REQUEST,
Json(json!({ "detail": format!("Provider 不支持 api_format: {unsupported}") })),
)
.into_response());
}
requested_api_formats
};
let mut settings_map = match payload.settings {
Some(Value::Object(map)) => map,
Some(_) => {
return Ok((
http::StatusCode::BAD_REQUEST,
Json(json!({ "detail": "settings payload must be an object" })),
)
.into_response());
}
None => Map::new(),
};
if let Some(proxy_node_id) = payload.proxy_node_id {
settings_map
.entry("proxy_node_id".to_string())
.or_insert(Value::String(proxy_node_id));
}
let shared_settings = (!settings_map.is_empty()).then_some(Value::Object(settings_map));
if let Some(settings) = shared_settings.as_ref() {
if let Err(detail) =
admin_provider_pool_pure::validate_admin_pool_key_settings_payload(settings)
{
return Ok((
http::StatusCode::BAD_REQUEST,
Json(json!({ "detail": detail })),
)
.into_response());
}
}
let mut known_names = existing_keys
.iter()
.map(|key| key.name.trim().to_string())
.filter(|name| !name.is_empty())
.collect::<BTreeSet<_>>();
let mut known_api_keys = existing_keys
.iter()
.filter_map(|key| key.encrypted_api_key.as_deref())
.filter_map(|ciphertext| self.decrypt_catalog_secret_with_fallbacks(ciphertext))
.filter(|value| value != "__placeholder__")
.collect::<BTreeSet<_>>();
let mut imported = 0usize;
let skipped = 0usize;
let mut errors = Vec::new();
@@ -138,6 +204,79 @@ impl<'a> AdminAppState<'a> {
continue;
}
let name = item.name.trim();
if name.is_empty() {
errors.push(json!({
"index": index,
"reason": "name is empty",
}));
continue;
}
if known_names.contains(name) {
errors.push(json!({
"index": index,
"reason": "该名称已存在于当前 Provider 或本次导入中",
}));
continue;
}
let auth_type = item.auth_type.trim().to_ascii_lowercase();
let auth_type = if auth_type.is_empty() {
"api_key".to_string()
} else {
auth_type
};
if !matches!(auth_type.as_str(), "api_key" | "bearer") {
errors.push(json!({
"index": index,
"reason": "auth_type must be api_key or bearer",
}));
continue;
}
let requested_item_api_formats = item
.api_formats
.iter()
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
.collect::<BTreeSet<_>>()
.into_iter()
.collect::<Vec<_>>();
let item_api_formats = if requested_item_api_formats.is_empty() {
api_formats.clone()
} else {
if let Some(unsupported) = requested_item_api_formats
.iter()
.find(|value| !available_api_format_set.contains(*value))
{
errors.push(json!({
"index": index,
"reason": format!("Provider 不支持 api_format: {unsupported}"),
}));
continue;
}
requested_item_api_formats
};
let item_settings = match admin_provider_pool_pure::resolve_admin_pool_key_settings(
shared_settings.as_ref(),
item.settings.as_ref(),
) {
Ok(value) => value,
Err(detail) => {
errors.push(json!({
"index": index,
"reason": detail,
}));
continue;
}
};
if known_api_keys.contains(api_key) {
errors.push(json!({
"index": index,
"reason": "该 API Key 已存在于当前 Provider 或本次导入中",
}));
continue;
}
let Some(encrypted_api_key) = self.encrypt_catalog_secret_with_fallbacks(api_key)
else {
errors.push(json!({
@@ -146,26 +285,15 @@ impl<'a> AdminAppState<'a> {
}));
continue;
};
let auth_type = item.auth_type.trim().to_ascii_lowercase();
let auth_type = if auth_type.is_empty() {
"api_key".to_string()
} else {
auth_type
};
let name = item.name.trim();
let record = match admin_provider_pool_pure::build_admin_pool_batch_import_key_record(
uuid::Uuid::new_v4().to_string(),
provider.id.clone(),
if name.is_empty() {
format!("imported-{index}")
} else {
name.to_string()
},
name.to_string(),
auth_type,
api_formats.clone(),
item_api_formats,
encrypted_api_key,
proxy.clone(),
None,
item_settings.as_ref(),
now_unix_secs,
) {
Ok(value) => value,
@@ -177,7 +305,6 @@ impl<'a> AdminAppState<'a> {
continue;
}
};
let Some(_) = self.create_provider_catalog_key(&record).await? else {
return Ok((
http::StatusCode::SERVICE_UNAVAILABLE,
@@ -187,6 +314,8 @@ impl<'a> AdminAppState<'a> {
)
.into_response());
};
known_names.insert(name.to_string());
known_api_keys.insert(api_key.to_string());
imported += 1;
}
@@ -475,6 +604,12 @@ impl<'a> AdminAppState<'a> {
AdminPoolBatchActionKind::Disable => key.is_active = false,
AdminPoolBatchActionKind::ClearProxy => key.proxy = None,
AdminPoolBatchActionKind::SetProxy => key.proxy = plan.proxy_payload.clone(),
AdminPoolBatchActionKind::UpdateSettings => {
if let Some(settings) = plan.settings_payload.as_ref() {
admin_provider_pool_pure::apply_admin_pool_key_settings(&mut key, settings)
.map_err(GatewayError::Internal)?;
}
}
AdminPoolBatchActionKind::RegenerateFingerprint => {
key.fingerprint =
Some(aether_provider_transport::claude_code::generate_random_fingerprint())
apps/aether-gateway/src/handlers/proxy/body_buffer.rs
+91 -124
View File
@@ -2,11 +2,11 @@ use crate::api::response::build_local_http_error_response;
use crate::control::GatewayPublicRequestContext;
use crate::headers::RequestBodyNormalizationError;
use crate::{AppState, GatewayError};
use axum::body::{to_bytes, Body, Bytes};
use aether_gateway_frontdoor::{BodyBufferError, BodyBufferPolicy as FrontdoorBodyBufferPolicy};
use axum::body::{Body, Bytes};
use axum::http::{self, Response};
use std::error::Error as StdError;
use std::sync::Arc;
use std::time::{Duration, Instant};
use std::time::Duration;
use tokio::sync::Semaphore;
use tracing::{info, warn};
@@ -16,23 +16,22 @@ const REQUEST_BODY_READ_FAILED_DETAIL: &str = "Failed to read request body";
#[derive(Debug, Clone)]
pub(super) struct RequestBodyBufferPolicy {
max_bytes: u64,
read_timeout: Duration,
queue_timeout: Duration,
budget_bytes: usize,
budget: Arc<Semaphore>,
inner: FrontdoorBodyBufferPolicy,
}
impl RequestBodyBufferPolicy {
pub(super) fn from_state(state: &AppState) -> Self {
Self {
max_bytes: crate::headers::max_request_body_bytes(),
read_timeout: state.frontdoor_runtime_guards.request_body_read_timeout,
queue_timeout: state.frontdoor_runtime_guards.internal_gate_queue_budget,
budget_bytes: state
.frontdoor_runtime_guards
.request_body_buffer_budget_bytes,
budget: Arc::clone(&state.request_body_buffer_budget),
inner: FrontdoorBodyBufferPolicy::with_permit_bytes(
crate::headers::max_request_body_bytes(),
state.frontdoor_runtime_guards.request_body_read_timeout,
state.frontdoor_runtime_guards.internal_gate_queue_budget,
state
.frontdoor_runtime_guards
.request_body_buffer_budget_bytes,
crate::state::REQUEST_BODY_BUFFER_PERMIT_BYTES,
Arc::clone(&state.request_body_buffer_budget),
),
}
}
@@ -42,14 +41,17 @@ impl RequestBodyBufferPolicy {
.unwrap_or(usize::MAX)
.max(crate::state::REQUEST_BODY_BUFFER_PERMIT_BYTES);
Self {
max_bytes,
read_timeout,
queue_timeout: read_timeout,
budget_bytes,
budget: Arc::new(Semaphore::new(
budget_bytes.saturating_add(crate::state::REQUEST_BODY_BUFFER_PERMIT_BYTES - 1)
/ crate::state::REQUEST_BODY_BUFFER_PERMIT_BYTES,
)),
inner: FrontdoorBodyBufferPolicy::with_permit_bytes(
max_bytes,
read_timeout,
read_timeout,
budget_bytes,
crate::state::REQUEST_BODY_BUFFER_PERMIT_BYTES,
Arc::new(Semaphore::new(
budget_bytes.saturating_add(crate::state::REQUEST_BODY_BUFFER_PERMIT_BYTES - 1)
/ crate::state::REQUEST_BODY_BUFFER_PERMIT_BYTES,
)),
),
}
}
@@ -62,13 +64,35 @@ impl RequestBodyBufferPolicy {
budget: Arc<Semaphore>,
) -> Self {
Self {
max_bytes,
read_timeout,
queue_timeout,
budget_bytes,
budget,
inner: FrontdoorBodyBufferPolicy::with_permit_bytes(
max_bytes,
read_timeout,
queue_timeout,
budget_bytes,
crate::state::REQUEST_BODY_BUFFER_PERMIT_BYTES,
budget,
),
}
}
fn max_bytes(&self) -> u64 {
self.inner.max_bytes()
}
fn budget_bytes(&self) -> usize {
self.inner.budget_bytes()
}
fn read_timeout(&self) -> Duration {
self.inner.read_timeout()
}
async fn reserve(
&self,
headers: &http::HeaderMap,
) -> Result<aether_gateway_frontdoor::BodyBufferReservation, BodyBufferError> {
self.inner.reserve(headers).await
}
}
#[derive(Debug)]
@@ -135,30 +159,23 @@ impl RequestBodyBufferError {
}
}
fn request_body_buffer_reservation_bytes(headers: &http::HeaderMap, max_bytes: u64) -> usize {
let max_bytes = usize::try_from(max_bytes).unwrap_or(usize::MAX);
let encoded = headers
.get(http::header::CONTENT_ENCODING)
.and_then(|value| value.to_str().ok())
.map(str::trim)
.is_some_and(|value| !value.is_empty() && !value.eq_ignore_ascii_case("identity"));
if encoded {
return max_bytes;
impl From<BodyBufferError> for RequestBodyBufferError {
fn from(error: BodyBufferError) -> Self {
match error {
BodyBufferError::TooLarge { limit_bytes } => Self::TooLarge { limit_bytes },
BodyBufferError::Overloaded {
requested_bytes,
budget_bytes,
timeout_ms,
} => Self::Overloaded {
requested_bytes,
budget_bytes,
timeout_ms,
},
BodyBufferError::Timeout { timeout_ms } => Self::Timeout { timeout_ms },
BodyBufferError::ReadFailed { message } => Self::ReadFailed { message },
}
}
headers
.get(http::header::CONTENT_LENGTH)
.and_then(|value| value.to_str().ok())
.and_then(|value| value.trim().parse::<usize>().ok())
.map(|value| value.min(max_bytes))
.unwrap_or(max_bytes)
}
fn request_body_buffer_reservation_permits(reservation_bytes: usize) -> u32 {
let permits = reservation_bytes
.max(1)
.saturating_add(crate::state::REQUEST_BODY_BUFFER_PERMIT_BYTES - 1)
/ crate::state::REQUEST_BODY_BUFFER_PERMIT_BYTES;
u32::try_from(permits).unwrap_or(u32::MAX).max(1)
}
pub(super) async fn buffer_and_normalize_request_body(
@@ -171,33 +188,12 @@ pub(super) async fn buffer_and_normalize_request_body(
phase: &'static str,
policy: RequestBodyBufferPolicy,
) -> Result<Bytes, RequestBodyBufferError> {
if let Err(err) =
crate::headers::check_request_content_length_with_limit(headers, policy.max_bytes)
{
return Err(RequestBodyBufferError::Normalization(err));
}
let reservation = policy
.reserve(headers)
.await
.map_err(RequestBodyBufferError::from)?;
let reservation_bytes = reservation.requested_bytes();
let reservation_bytes = request_body_buffer_reservation_bytes(headers, policy.max_bytes);
let reservation_permits = request_body_buffer_reservation_permits(reservation_bytes);
let queue_timeout_ms = policy.queue_timeout.as_millis() as u64;
let _budget_permit = match tokio::time::timeout(
policy.queue_timeout,
Arc::clone(&policy.budget).acquire_many_owned(reservation_permits),
)
.await
{
Ok(Ok(permit)) => permit,
Ok(Err(_)) | Err(_) => {
return Err(RequestBodyBufferError::Overloaded {
requested_bytes: reservation_bytes,
budget_bytes: policy.budget_bytes,
timeout_ms: queue_timeout_ms,
});
}
};
let read_started_at = Instant::now();
let timeout_ms = policy.read_timeout.as_millis() as u64;
info!(
event_name = "frontdoor_request_body_buffer_started",
log_type = "event",
@@ -205,45 +201,27 @@ pub(super) async fn buffer_and_normalize_request_body(
method = %method,
path = %path_and_query,
phase,
max_body_bytes = policy.max_bytes,
max_body_bytes = policy.max_bytes(),
reserved_body_bytes = reservation_bytes,
body_buffer_budget_bytes = policy.budget_bytes,
timeout_ms,
body_buffer_budget_bytes = policy.budget_bytes(),
timeout_ms = policy.read_timeout().as_millis() as u64,
"gateway started buffering request body"
);
let body_limit = usize::try_from(policy.max_bytes).unwrap_or(usize::MAX);
let body = match tokio::time::timeout(
policy.read_timeout,
to_bytes(
request_body.take().expect(body_owner_expectation),
body_limit,
),
)
.await
{
Ok(Ok(body)) => body,
Ok(Err(err)) if request_body_collection_exceeded_limit(&err) => {
return Err(RequestBodyBufferError::TooLarge {
limit_bytes: policy.max_bytes,
});
}
Ok(Err(err)) => {
return Err(RequestBodyBufferError::ReadFailed {
message: err.to_string(),
});
}
Err(_) => {
return Err(RequestBodyBufferError::Timeout { timeout_ms });
}
};
let normalized = crate::headers::normalize_request_body_headers_and_bytes_with_limit(
headers,
body,
policy.max_bytes,
)
.map_err(RequestBodyBufferError::Normalization)?;
let buffered = reservation
.collect(request_body.take().expect(body_owner_expectation))
.await
.map_err(RequestBodyBufferError::from)?;
let elapsed_ms = buffered.elapsed().as_millis() as u64;
let normalized = buffered
.try_map(|body| {
crate::headers::normalize_request_body_headers_and_bytes_with_limit(
headers,
body,
policy.max_bytes(),
)
})
.map_err(RequestBodyBufferError::Normalization)?;
info!(
event_name = "frontdoor_request_body_buffer_completed",
log_type = "event",
@@ -252,23 +230,12 @@ pub(super) async fn buffer_and_normalize_request_body(
path = %path_and_query,
phase,
body_bytes = normalized.len(),
elapsed_ms = read_started_at.elapsed().as_millis() as u64,
elapsed_ms,
"gateway completed buffering request body"
);
Ok(normalized)
}
fn request_body_collection_exceeded_limit(error: &(dyn StdError + 'static)) -> bool {
let mut current = Some(error);
while let Some(error) = current {
if error.to_string().contains("length limit exceeded") {
return true;
}
current = error.source();
}
false
}
pub(super) fn build_request_body_buffer_error_response(
trace_id: &str,
request_context: &GatewayPublicRequestContext,
apps/aether-gateway/src/handlers/proxy/mod.rs
+177 -14
View File
@@ -84,6 +84,8 @@ const OPENAI_RESPONSES_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL: &str =
"当前 OpenAI Responses 请求无法在本地执行：没有匹配到可用的执行路径";
const OPENAI_RESPONSES_COMPACT_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL: &str =
"当前 OpenAI Responses Compact 请求无法在本地执行：没有匹配到可用的执行路径";
const OPENAI_SEARCH_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL: &str =
"当前 OpenAI Search 请求无法在本地执行：没有匹配到可用的执行路径";
const OPENAI_VIDEO_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL: &str =
"当前 OpenAI Video 请求无法在本地执行：没有匹配到可用的执行路径";
const CLAUDE_MESSAGES_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL: &str =
@@ -381,8 +383,9 @@ async fn maybe_forward_public_request_to_tunnel_owner(
serde_json::from_slice::<serde_json::Value>(body.as_ref()).ok()
});
let client_session_affinity =
crate::client_session_affinity::client_session_affinity_from_parts(
parts,
crate::client_session_affinity::client_session_affinity_from_api_request(
api_format,
&parts.headers,
body_json.as_ref(),
);
let Some(target) = crate::scheduler::affinity::read_cached_scheduler_affinity_target(
@@ -452,7 +455,27 @@ async fn maybe_forward_public_request_to_tunnel_owner(
owner.relay_base_url.trim_end_matches('/'),
request_context.request_path_and_query()
);
let mut upstream_request = state.client.request(parts.method.clone(), owner_url);
let is_stream =
owner_forward_request_is_stream(parts, decision, buffered_body.unwrap_or(&empty_body));
let transport_timeouts =
crate::provider_transport::resolve_transport_execution_timeouts(&transport);
let non_stream_timeout =
crate::execution_runtime::transport::resolve_non_stream_total_timeout_for_request(
is_stream,
&transport.endpoint.api_format,
transport_timeouts.as_ref(),
);
let stream_first_byte_timeout =
crate::execution_runtime::transport::resolve_stream_first_byte_timeout_for_request(
is_stream,
transport_timeouts.as_ref(),
);
let mut upstream_request = state
.owner_forward_client
.request(parts.method.clone(), owner_url);
if let Some(timeout) = non_stream_timeout {
upstream_request = upstream_request.timeout(timeout);
}
for (name, value) in &parts.headers {
if should_skip_request_header(name.as_str()) || name == http::header::HOST {
continue;
@@ -498,14 +521,15 @@ async fn maybe_forward_public_request_to_tunnel_owner(
upstream_request.header(TRUSTED_AUTH_BALANCE_HEADER, balance_remaining.to_string());
}
let upstream_response = upstream_request
.body(buffered_body.cloned().unwrap_or_default())
.send()
.await
.map_err(|err| GatewayError::UpstreamUnavailable {
trace_id: request_context.trace_id.clone(),
message: format!("owner gateway affinity forward failed: {err}"),
})?;
let upstream_response = crate::tunnel::send_owner_forward_request(
upstream_request.body(buffered_body.cloned().unwrap_or_default()),
stream_first_byte_timeout,
)
.await
.map_err(|message| GatewayError::UpstreamUnavailable {
trace_id: request_context.trace_id.clone(),
message: format!("owner gateway affinity forward failed: {message}"),
})?;
let mut response = build_sync_aware_affinity_forward_response(
request_context,
@@ -523,6 +547,29 @@ async fn maybe_forward_public_request_to_tunnel_owner(
Ok(Some(response))
}
fn owner_forward_request_is_stream(
parts: &http::request::Parts,
decision: &GatewayControlDecision,
body_bytes: &Bytes,
) -> bool {
let Some(plan_kind) =
crate::ai_serving::api::resolve_execution_runtime_stream_plan_kind(parts, decision)
else {
return false;
};
let Some((body_json, body_base64)) =
crate::ai_serving::api::parse_direct_request_body(parts, body_bytes)
else {
return false;
};
crate::ai_serving::api::is_matching_stream_request(
plan_kind,
parts,
&body_json,
body_base64.as_deref(),
)
}
fn upstream_response_is_sse(headers: &reqwest::header::HeaderMap) -> bool {
headers
.get(http::header::CONTENT_TYPE)
@@ -2196,6 +2243,7 @@ fn local_execution_runtime_miss_route_label(
"/v1/chat/completions" => "OpenAI Chat Completions",
"/v1/responses" => "OpenAI Responses",
"/v1/responses/compact" => "OpenAI Responses Compact",
"/v1/alpha/search" => "OpenAI Search",
"/v1/messages" => "Claude Messages",
path if path.starts_with("/v1/videos") => "OpenAI Video",
path if path.starts_with("/upload/v1beta/files") || path.starts_with("/v1beta/files") => {
@@ -2238,6 +2286,7 @@ fn local_execution_runtime_miss_route_detail(
"/v1/responses/compact" => {
Some(OPENAI_RESPONSES_COMPACT_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL)
}
"/v1/alpha/search" => Some(OPENAI_SEARCH_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL),
"/v1/messages" => Some(CLAUDE_MESSAGES_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL),
path if path.starts_with("/v1/videos") => {
Some(OPENAI_VIDEO_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL)
@@ -2262,15 +2311,129 @@ mod tests {
use super::{
api_key_remote_ip_allowed, buffer_and_normalize_request_body,
diagnostic_is_auth_api_key_concurrency_limited, local_execution_runtime_miss_detail,
restore_redacted_stream_execution_response, restore_redacted_sync_execution_response,
GatewayControlDecision, LocalExecutionRuntimeMissDiagnostic, RequestBodyBufferError,
RequestBodyBufferPolicy,
owner_forward_request_is_stream, restore_redacted_stream_execution_response,
restore_redacted_sync_execution_response, GatewayControlDecision,
LocalExecutionRuntimeMissDiagnostic, RequestBodyBufferError, RequestBodyBufferPolicy,
};
use axum::body::{to_bytes, Body, Bytes};
use axum::http::{header, HeaderMap, HeaderValue, Method, Response};
use serde_json::json;
use tokio::sync::Semaphore;
#[test]
fn owner_forward_uses_search_protocol_timeout_semantics() {
let request = http::Request::builder()
.method(Method::POST)
.uri("/v1/alpha/search")
.header(header::CONTENT_TYPE, "application/json")
.body(())
.expect("request should build");
let (parts, _) = request.into_parts();
let decision = GatewayControlDecision::synthetic(
"/v1/alpha/search",
Some("ai_public".to_string()),
Some("openai".to_string()),
Some("search".to_string()),
Some("openai:search".to_string()),
);
let body =
Bytes::from_static(br#"{"model":"gpt-5.6-sol","input":"find docs","stream":true}"#);
let is_stream = owner_forward_request_is_stream(&parts, &decision, &body);
let timeouts = aether_contracts::ExecutionTimeouts {
total_ms: Some(aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_MS),
first_byte_ms: Some(10),
..aether_contracts::ExecutionTimeouts::default()
};
assert!(!is_stream);
assert_eq!(
crate::execution_runtime::transport::resolve_non_stream_total_timeout_for_request(
is_stream,
"openai:search",
Some(&timeouts),
),
Some(Duration::from_millis(
aether_contracts::MAX_EXECUTION_REQUEST_TIMEOUT_MS
))
);
assert_eq!(
crate::execution_runtime::transport::resolve_stream_first_byte_timeout_for_request(
is_stream,
Some(&timeouts),
),
None
);
}
#[test]
fn owner_forward_keeps_streaming_for_stream_capable_protocols() {
let chat_request = http::Request::builder()
.method(Method::POST)
.uri("/v1/chat/completions")
.header(header::CONTENT_TYPE, "application/json")
.body(())
.expect("request should build");
let (chat_parts, _) = chat_request.into_parts();
let chat_decision = GatewayControlDecision::synthetic(
"/v1/chat/completions",
Some("ai_public".to_string()),
Some("openai".to_string()),
Some("chat".to_string()),
Some("openai:chat".to_string()),
);
assert!(owner_forward_request_is_stream(
&chat_parts,
&chat_decision,
&Bytes::from_static(br#"{"model":"gpt-5.6-sol","stream":true}"#),
));
assert!(!owner_forward_request_is_stream(
&chat_parts,
&chat_decision,
&Bytes::from_static(br#"{"model":"gpt-5.6-sol","stream":false}"#),
));
let image_request = http::Request::builder()
.method(Method::POST)
.uri("/v1/images/generations")
.header(header::CONTENT_TYPE, "application/json")
.body(())
.expect("request should build");
let (image_parts, _) = image_request.into_parts();
let image_decision = GatewayControlDecision::synthetic(
"/v1/images/generations",
Some("ai_public".to_string()),
Some("openai".to_string()),
Some("image".to_string()),
Some("openai:image".to_string()),
);
assert!(owner_forward_request_is_stream(
&image_parts,
&image_decision,
&Bytes::from_static(br#"{"model":"gpt-image-1","stream":true}"#),
));
let compact_request = http::Request::builder()
.method(Method::POST)
.uri("/v1/responses/compact")
.header(header::CONTENT_TYPE, "application/json")
.body(())
.expect("request should build");
let (compact_parts, _) = compact_request.into_parts();
let compact_decision = GatewayControlDecision::synthetic(
"/v1/responses/compact",
Some("ai_public".to_string()),
Some("openai".to_string()),
Some("responses:compact".to_string()),
Some("openai:responses:compact".to_string()),
);
assert!(!owner_forward_request_is_stream(
&compact_parts,
&compact_decision,
&Bytes::from_static(br#"{"model":"gpt-5.6-sol","stream":true}"#),
));
}
#[test]
fn api_key_remote_ip_allows_unrestricted_keys() {
let remote_ip = "203.0.113.10".parse().expect("valid ip");
apps/aether-gateway/src/handlers/public/ai_public.rs
+30 -12
View File
@@ -1,9 +1,8 @@
use crate::ai_serving::normalize_openai_image_quality;
use crate::async_task::CancelVideoTaskError;
use crate::control::GatewayControlDecision;
use crate::control::GatewayPublicRequestContext;
use crate::image_capabilities::{
openai_image_gateway_max_generation_count, openai_image_gateway_max_generation_count_for_model,
};
use crate::image_capabilities::openai_image_gateway_max_generation_count;
use crate::{AppState, GatewayError};
use aether_data_contracts::repository::video_tasks::{
StoredVideoTask, VideoTaskQueryFilter, VideoTaskStatus,
@@ -26,7 +25,7 @@ const OPENAI_IMAGE_PARTIAL_IMAGES_DETAIL: &str =
const OPENAI_IMAGE_STYLE_DETAIL: &str = "当前 Codex 图片反代暂不支持 style 参数";
const OPENAI_IMAGE_RESPONSE_FORMAT_DETAIL: &str = "response_format 仅支持 url 或 b64_json";
const OPENAI_IMAGE_OUTPUT_FORMAT_DETAIL: &str = "output_format 仅支持 png、jpeg 或 webp";
const OPENAI_IMAGE_QUALITY_DETAIL: &str = "quality 仅支持 low、medium、high、standard 或 hd";
const OPENAI_IMAGE_QUALITY_DETAIL: &str = "quality 仅支持 auto、low、medium、high、standard 或 hd";
const OPENAI_IMAGE_BACKGROUND_DETAIL: &str = "background 仅支持 auto、opaque 或 transparent";
const OPENAI_IMAGE_MODERATION_DETAIL: &str = "moderation 仅支持 auto 或 low";
const OPENAI_IMAGE_INPUT_FIDELITY_DETAIL: &str = "input_fidelity 仅支持 low 或 high";
@@ -303,7 +302,7 @@ fn maybe_build_local_openai_request_validation_response(
if validation
.quality
.as_deref()
.is_some_and(|value| !matches!(value, "low" | "medium" | "high" | "standard" | "hd"))
.is_some_and(|value| normalize_openai_image_quality(value).is_none())
{
return Some(build_ai_public_error_response(
http::StatusCode::BAD_REQUEST,
@@ -366,8 +365,7 @@ fn openai_image_n_detail(max_generation_count: u64) -> String {
}
fn validate_openai_image_n(validation: &OpenAiImageValidationInput) -> Option<String> {
let max_generation_count =
openai_image_gateway_max_generation_count_for_model(validation.model.as_deref());
let max_generation_count = openai_image_gateway_max_generation_count();
validation
.n
.is_some_and(|value| value == 0 || value > max_generation_count)
@@ -1761,7 +1759,7 @@ mod tests {
}
#[test]
fn image_validation_restricts_multi_image_count_to_grok_models() {
fn image_validation_applies_the_global_count_limit_before_model_mapping() {
let openai_body = Bytes::from_static(br#"{"model":"gpt-image-2","prompt":"draw","n":2}"#);
let openai_validation = parse_openai_image_validation_input(
OpenAiImageOperation::Generate,
@@ -1770,10 +1768,7 @@ mod tests {
)
.expect("valid image payload should parse");
assert_eq!(
validate_openai_image_n(&openai_validation).as_deref(),
Some("当前图片模型仅支持 n=1..1")
);
assert!(validate_openai_image_n(&openai_validation).is_none());
let grok_body =
Bytes::from_static(br#"{"model":"grok-imagine-image-lite","prompt":"draw","n":4}"#);
@@ -1785,5 +1780,28 @@ mod tests {
.expect("valid grok image payload should parse");
assert!(validate_openai_image_n(&grok_validation).is_none());
let alias_body =
Bytes::from_static(br#"{"model":"production-image-alias","prompt":"draw","n":10}"#);
let alias_validation = parse_openai_image_validation_input(
OpenAiImageOperation::Generate,
Some("application/json"),
&alias_body,
)
.expect("valid image alias payload should parse");
assert!(validate_openai_image_n(&alias_validation).is_none());
let excessive_body =
Bytes::from_static(br#"{"model":"production-image-alias","prompt":"draw","n":11}"#);
let excessive_validation = parse_openai_image_validation_input(
OpenAiImageOperation::Generate,
Some("application/json"),
&excessive_body,
)
.expect("image payload should parse before count validation");
assert_eq!(
validate_openai_image_n(&excessive_validation).as_deref(),
Some("当前图片反代仅支持 n=1..10")
);
}
}
apps/aether-gateway/src/handlers/public/support/dashboard_filters.rs
+97 -72
View File
@@ -5,8 +5,8 @@ use super::{
};
use aether_data_contracts::repository::usage::{
StoredUsageCostSavingsSummary, StoredUsageDashboardDailyBreakdownRow,
StoredUsageDashboardSummary, UsageAuditAggregationGroupBy, UsageAuditAggregationQuery,
UsageCostSavingsSummaryQuery, UsageDashboardDailyBreakdownQuery,
StoredUsageDashboardStatsSummary, StoredUsageDashboardSummary, UsageAuditAggregationGroupBy,
UsageAuditAggregationQuery, UsageDashboardDailyBreakdownQuery,
UsageDashboardProviderCountsQuery, UsageDashboardSummaryQuery,
};
use axum::{
@@ -550,6 +550,39 @@ async fn dashboard_summary_for_range(
dashboard_summary_for_range_raw(state, range, user_id, error_context).await
}
async fn dashboard_stats_for_range(
state: &AppState,
range: DashboardDateRange,
user_id: Option<&str>,
error_context: &str,
) -> Result<StoredUsageDashboardStatsSummary, Response<Body>> {
let Some((created_from_unix_secs, created_until_unix_secs)) =
dashboard_range_bounds_unix(range)
else {
return Err(build_auth_error_response(
http::StatusCode::INTERNAL_SERVER_ERROR,
format!("{error_context}: invalid time range"),
false,
));
};
match state
.summarize_dashboard_stats(&UsageDashboardSummaryQuery {
created_from_unix_secs,
created_until_unix_secs,
user_id: user_id.map(ToOwned::to_owned),
})
.await
{
Ok(value) => Ok(value),
Err(err) => Err(build_auth_error_response(
http::StatusCode::INTERNAL_SERVER_ERROR,
format!("{error_context}: {err:?}"),
false,
)),
}
}
async fn dashboard_daily_breakdown_for_range(
state: &AppState,
range: DashboardDateRange,
@@ -855,28 +888,6 @@ fn dashboard_cache_savings_usd(summary: &StoredUsageCostSavingsSummary) -> f64 {
)
}
async fn dashboard_load_cache_savings(
state: &AppState,
range: DashboardDateRange,
user_id: Option<&str>,
) -> Result<f64, GatewayError> {
let Some((created_from_unix_secs, created_until_unix_secs)) =
dashboard_range_bounds_unix(range)
else {
return Ok(0.0);
};
let summary = state
.summarize_usage_cost_savings(&UsageCostSavingsSummaryQuery {
created_from_unix_secs,
created_until_unix_secs,
user_id: user_id.map(ToOwned::to_owned),
provider_name: None,
model: None,
})
.await?;
Ok(dashboard_cache_savings_usd(&summary))
}
pub(super) async fn handle_dashboard_stats_get(
state: &AppState,
request_context: &GatewayPublicRequestContext,
@@ -902,7 +913,7 @@ pub(super) async fn handle_dashboard_stats_get(
.as_deref()
.unwrap_or("");
let cache_key = format!("stats:{cache_identity}:{query_string}");
let cache_ttl = std::time::Duration::from_secs(15);
let cache_ttl = std::time::Duration::from_secs(30);
if let Some(cached) = state.dashboard_response_cache.get(&cache_key, cache_ttl) {
return Response::builder()
@@ -924,31 +935,63 @@ pub(super) async fn handle_dashboard_stats_get(
tz_offset_minutes: summary_range.tz_offset_minutes,
};
let user_filter = (!is_admin).then_some(auth.user.id.as_str());
let period_summary = match dashboard_summary_for_range(
state,
summary_range,
user_filter,
"dashboard stats lookup failed",
)
.await
{
Ok(value) => value,
Err(response) => return response,
let (period_totals, today_totals, admin_cost_savings) = if is_admin {
let (period_result, today_result) = tokio::join!(
dashboard_stats_for_range(
state,
summary_range,
user_filter,
"dashboard stats lookup failed",
),
dashboard_stats_for_range(
state,
today_range,
user_filter,
"dashboard today stats lookup failed",
),
);
let period = match period_result {
Ok(value) => value,
Err(response) => return response,
};
let today = match today_result {
Ok(value) => value,
Err(response) => return response,
};
(
dashboard_usage_totals_from_summary(&period.usage),
dashboard_usage_totals_from_summary(&today.usage),
Some((period.cost_savings, today.cost_savings)),
)
} else {
let period_summary = match dashboard_summary_for_range(
state,
summary_range,
user_filter,
"dashboard stats lookup failed",
)
.await
{
Ok(value) => value,
Err(response) => return response,
};
let today_summary = match dashboard_summary_for_range(
state,
today_range,
user_filter,
"dashboard today stats lookup failed",
)
.await
{
Ok(value) => value,
Err(response) => return response,
};
(
dashboard_usage_totals_from_summary(&period_summary),
dashboard_usage_totals_from_summary(&today_summary),
None,
)
};
let today_summary = match dashboard_summary_for_range(
state,
today_range,
user_filter,
"dashboard today stats lookup failed",
)
.await
{
Ok(value) => value,
Err(response) => return response,
};
let period_totals = dashboard_usage_totals_from_summary(&period_summary);
let today_totals = dashboard_usage_totals_from_summary(&today_summary);
let api_key_counts = match dashboard_load_api_key_counts(state, is_admin, &auth.user.id).await {
Ok(value) => value,
@@ -1030,28 +1073,10 @@ pub(super) async fn handle_dashboard_stats_get(
/ today_totals.requests as f64
* 100.0
};
let today_cost_savings =
match dashboard_load_cache_savings(state, today_range, user_filter).await {
Ok(value) => value,
Err(err) => {
return build_auth_error_response(
http::StatusCode::INTERNAL_SERVER_ERROR,
format!("dashboard today cache savings lookup failed: {err:?}"),
false,
);
}
};
let period_cost_savings =
match dashboard_load_cache_savings(state, summary_range, user_filter).await {
Ok(value) => value,
Err(err) => {
return build_auth_error_response(
http::StatusCode::INTERNAL_SERVER_ERROR,
format!("dashboard cache savings lookup failed: {err:?}"),
false,
);
}
};
let (period_cost_savings_summary, today_cost_savings_summary) =
admin_cost_savings.unwrap_or_default();
let today_cost_savings = dashboard_cache_savings_usd(&today_cost_savings_summary);
let period_cost_savings = dashboard_cache_savings_usd(&period_cost_savings_summary);
let stats = json!([
{
"name": "今日请求 / 费用",
@@ -1248,7 +1273,7 @@ pub(super) async fn handle_dashboard_daily_stats_get(
.as_deref()
.unwrap_or("");
let cache_key = format!("daily:{cache_identity}:{query_string}");
let cache_ttl = std::time::Duration::from_secs(30);
let cache_ttl = std::time::Duration::from_secs(60);
if let Some(cached) = state.dashboard_response_cache.get(&cache_key, cache_ttl) {
return Response::builder()
apps/aether-gateway/src/handlers/public/support/models/responses.rs
+5
View File
@@ -89,6 +89,7 @@ pub(super) fn build_models_not_found_response(model_id: &str, api_format: &str)
pub(super) fn build_empty_models_list_response(api_format: &str) -> Response<Body> {
match api_format {
"openai:responses" => Json(json!({ "models": [] })).into_response(),
"claude:messages" => Json(json!({
"data": [],
"has_more": false,
@@ -101,6 +102,10 @@ pub(super) fn build_empty_models_list_response(api_format: &str) -> Response<Bod
}
}
pub(super) fn build_codex_models_list_response(models: Vec<serde_json::Value>) -> Response<Body> {
Json(json!({ "models": models })).into_response()
}
pub(super) fn build_openai_models_list_response(
rows: &[StoredMinimalCandidateSelectionRow],
) -> Response<Body> {
apps/aether-gateway/src/handlers/public/support/models/route.rs
+159 -10
View File
@@ -1,21 +1,24 @@
use std::collections::{BTreeMap, BTreeSet};
use std::fmt::Debug;
use std::future::Future;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use aether_data_contracts::repository::candidate_selection::StoredMinimalCandidateSelectionRow;
use axum::{body::Body, response::Response};
use serde_json::Value;
use tokio::time::timeout;
use tracing::warn;
use super::models_responses::{
build_claude_model_detail_response, build_claude_models_list_response,
build_empty_models_list_response, build_gemini_model_detail_response,
build_gemini_models_list_response, build_models_auth_error_response,
build_models_not_found_response, build_openai_model_detail_response,
build_openai_models_list_response,
build_codex_models_list_response, build_empty_models_list_response,
build_gemini_model_detail_response, build_gemini_models_list_response,
build_models_auth_error_response, build_models_not_found_response,
build_openai_model_detail_response, build_openai_models_list_response,
};
use super::models_shared::{
filter_rows_for_models, models_api_format, models_detail_id, models_query_api_formats,
filter_eligible_model_rows, filter_rows_for_models, models_api_format, models_detail_id,
models_query_api_formats,
};
use super::{query_param_value, AppState, GatewayPublicRequestContext};
@@ -23,6 +26,7 @@ use super::{query_param_value, AppState, GatewayPublicRequestContext};
const MODELS_ROUTE_READ_TIMEOUT: Duration = Duration::from_secs(5);
#[cfg(test)]
const MODELS_ROUTE_READ_TIMEOUT: Duration = Duration::from_millis(50);
const CODEX_MODELS_QUERY_API_FORMATS: &[&str] = &["openai:responses"];
async fn await_models_route_read<T, E, Fut>(operation: &'static str, future: Fut) -> Option<T>
where
@@ -72,7 +76,7 @@ fn build_models_read_fallback_response(
}
}
fn sort_and_dedup_model_rows(
fn sort_model_rows(
mut rows: Vec<StoredMinimalCandidateSelectionRow>,
) -> Vec<StoredMinimalCandidateSelectionRow> {
rows.sort_by(|left, right| {
@@ -85,9 +89,15 @@ fn sort_and_dedup_model_rows(
.then(left.key_id.cmp(&right.key_id))
.then(left.model_id.cmp(&right.model_id))
});
rows
}
fn sort_and_dedup_model_rows(
rows: Vec<StoredMinimalCandidateSelectionRow>,
) -> Vec<StoredMinimalCandidateSelectionRow> {
let mut deduped = Vec::with_capacity(rows.len());
let mut last_model_name: Option<String> = None;
for row in rows {
for row in sort_model_rows(rows) {
if last_model_name.as_deref() == Some(row.global_model_name.as_str()) {
continue;
}
@@ -97,22 +107,157 @@ fn sort_and_dedup_model_rows(
deduped
}
fn is_codex_models_api_format(api_format: &str) -> bool {
crate::ai_serving::normalize_api_format_alias(api_format) == "openai:responses"
}
fn is_codex_provider_row(row: &StoredMinimalCandidateSelectionRow) -> bool {
row.provider_type.trim().eq_ignore_ascii_case("codex")
}
fn codex_model_card_is_complete(card: &serde_json::Map<String, Value>) -> bool {
card.get("slug").and_then(Value::as_str).is_some()
&& card.get("display_name").and_then(Value::as_str).is_some()
&& card
.get("supported_reasoning_levels")
.and_then(Value::as_array)
.is_some()
&& card.get("shell_type").and_then(Value::as_str).is_some()
&& card.get("visibility").and_then(Value::as_str).is_some()
&& card
.get("supported_in_api")
.and_then(Value::as_bool)
.is_some()
&& card.get("priority").and_then(Value::as_i64).is_some()
&& card
.get("base_instructions")
.and_then(Value::as_str)
.is_some()
&& card
.get("supports_reasoning_summary_parameter")
.is_none_or(Value::is_boolean)
&& card
.get("support_verbosity")
.and_then(Value::as_bool)
.is_some()
&& card
.get("truncation_policy")
.and_then(Value::as_object)
.is_some()
&& card
.get("supports_parallel_tool_calls")
.and_then(Value::as_bool)
.is_some()
&& card
.get("experimental_supported_tools")
.and_then(Value::as_array)
.is_some()
}
fn project_codex_model_card(
cached_models: &[Value],
source_model: &str,
global_model: &str,
) -> Option<Value> {
let mut card = cached_models
.iter()
.find(|model| {
model.get("id").and_then(Value::as_str) == Some(source_model)
|| model.get("slug").and_then(Value::as_str) == Some(source_model)
})?
.as_object()?
.clone();
if !codex_model_card_is_complete(&card) {
return None;
}
card.remove("id");
card.remove("api_formats");
card.insert("slug".to_string(), Value::String(global_model.to_string()));
Some(Value::Object(card))
}
async fn load_codex_model_cards(
state: &AppState,
rows: &[StoredMinimalCandidateSelectionRow],
) -> Vec<Value> {
let cache_keys = rows
.iter()
.filter(|row| is_codex_provider_row(row))
.map(|row| format!("upstream_models:{}:{}", row.provider_id, row.key_id))
.collect::<BTreeSet<_>>()
.into_iter()
.collect::<Vec<_>>();
let cached_values = await_models_route_read(
"codex_models_cache",
state.runtime_state.kv_get_many(&cache_keys),
)
.await
.unwrap_or_default();
let cached_models_by_key = cache_keys
.into_iter()
.zip(cached_values)
.filter_map(|(key, raw)| {
let models = serde_json::from_str::<Vec<Value>>(raw.as_deref()?).ok()?;
Some((key, models))
})
.collect::<BTreeMap<_, _>>();
let mut seen_global_models = BTreeSet::new();
let mut cards = Vec::new();
for row in rows.iter().filter(|row| is_codex_provider_row(row)) {
if seen_global_models.contains(&row.global_model_name) {
continue;
}
let cache_key = format!("upstream_models:{}:{}", row.provider_id, row.key_id);
let Some(cached_models) = cached_models_by_key.get(&cache_key) else {
continue;
};
let source_model =
aether_scheduler_core::select_provider_model_name(row, "openai:responses");
let Some(card) = project_codex_model_card(
cached_models,
source_model.as_str(),
row.global_model_name.as_str(),
) else {
continue;
};
seen_global_models.insert(row.global_model_name.clone());
cards.push(card);
}
cards
}
async fn list_model_rows_for_client_format(
state: &AppState,
api_format: &str,
auth_snapshot: Option<&crate::data::auth::GatewayAuthApiKeySnapshot>,
) -> Option<Vec<StoredMinimalCandidateSelectionRow>> {
let mut collected = Vec::new();
for query_format in models_query_api_formats(api_format) {
let query_api_formats = if is_codex_models_api_format(api_format) {
CODEX_MODELS_QUERY_API_FORMATS
} else {
models_query_api_formats(api_format)
};
for query_format in query_api_formats {
let rows = await_models_route_read(
"candidate_selection_by_api_format",
state.list_minimal_candidate_selection_rows_for_api_format(query_format),
)
.await?;
let mut filtered = filter_rows_for_models(rows, auth_snapshot, query_format);
let mut filtered = if is_codex_models_api_format(api_format) {
filter_eligible_model_rows(rows, auth_snapshot, query_format)
} else {
filter_rows_for_models(rows, auth_snapshot, query_format)
};
collected.append(&mut filtered);
}
Some(sort_and_dedup_model_rows(collected))
if is_codex_models_api_format(api_format) {
collected.retain(is_codex_provider_row);
Some(sort_model_rows(collected))
} else {
Some(sort_and_dedup_model_rows(collected))
}
}
async fn list_model_rows_for_client_format_and_global_model(
@@ -190,6 +335,10 @@ pub(super) async fn maybe_build_local_models_route_response(
if rows.is_empty() {
return Some(build_empty_models_list_response(api_format));
}
if is_codex_models_api_format(api_format) {
let models = load_codex_model_cards(state, &rows).await;
return Some(build_codex_models_list_response(models));
}
let response = match api_format {
"claude:messages" => {
let before_id = query_param_value(
apps/aether-gateway/src/handlers/public/support/models/shared.rs
+11 -4
View File
@@ -194,13 +194,12 @@ fn row_exposes_global_model_for_models(
false
}
pub(super) fn filter_rows_for_models(
pub(super) fn filter_eligible_model_rows(
rows: Vec<StoredMinimalCandidateSelectionRow>,
auth_snapshot: Option<&crate::data::auth::GatewayAuthApiKeySnapshot>,
api_format: &str,
) -> Vec<StoredMinimalCandidateSelectionRow> {
let mut filtered = rows
.into_iter()
rows.into_iter()
.filter(|row| {
auth_snapshot_allows_provider_for_models(
auth_snapshot,
@@ -211,7 +210,15 @@ pub(super) fn filter_rows_for_models(
})
.filter(|row| auth_snapshot_allows_model_for_models(auth_snapshot, &row.global_model_name))
.filter(|row| row_exposes_global_model_for_models(row, api_format))
.collect::<Vec<_>>();
.collect()
}
pub(super) fn filter_rows_for_models(
rows: Vec<StoredMinimalCandidateSelectionRow>,
auth_snapshot: Option<&crate::data::auth::GatewayAuthApiKeySnapshot>,
api_format: &str,
) -> Vec<StoredMinimalCandidateSelectionRow> {
let mut filtered = filter_eligible_model_rows(rows, auth_snapshot, api_format);
filtered.sort_by(|left, right| left.global_model_name.cmp(&right.global_model_name));
let mut deduped = Vec::new();
let mut last_model_name: Option<String> = None;
apps/aether-gateway/src/handlers/public/support/user_me_usage.rs
+30 -8
View File
@@ -220,14 +220,7 @@ fn users_me_usage_api_format_defaults_to_non_stream(item: &StoredRequestUsageAud
let Some(value) = api_format else {
return false;
};
matches!(
crate::ai_serving::normalize_api_format_alias(value).as_str(),
"openai:chat"
| "openai:responses"
| "openai:responses:compact"
| "openai:image"
| "claude:messages"
)
crate::ai_serving::api_format_defaults_to_non_stream(value)
}
fn users_me_usage_request_body_implies_default_non_stream(item: &StoredRequestUsageAudit) -> bool {
@@ -495,6 +488,7 @@ fn build_users_me_usage_record_payload(
"cache_read_input_tokens": item.cache_read_input_tokens,
"status_code": item.status_code,
"error_message": item.error_message,
"request_type": item.request_type,
"input_price_per_1m": input_price_per_1m,
"output_price_per_1m": output_price_per_1m,
"cache_creation_price_per_1m": cache_creation_price_per_1m,
@@ -515,6 +509,9 @@ fn build_users_me_usage_record_payload(
if let Some(service_tier) = item.provider_service_tier() {
payload["service_tier"] = json!(service_tier);
}
if let Some(actual_service_tier) = item.provider_actual_service_tier() {
payload["actual_service_tier"] = json!(actual_service_tier);
}
if include_actual_cost {
payload["actual_cost"] = json!(round_to(item.actual_total_cost_usd, 6));
payload["rate_multiplier"] = json!(rate_multiplier);
@@ -529,6 +526,7 @@ fn build_users_me_usage_active_payload(item: &StoredRequestUsageAudit) -> serde_
let mut payload = json!({
"id": item.id,
"status": item.status,
"request_type": item.request_type,
"input_tokens": item.input_tokens,
"effective_input_tokens": users_me_usage_effective_input_tokens(item),
"output_tokens": item.output_tokens,
@@ -582,6 +580,9 @@ fn build_users_me_usage_active_payload(item: &StoredRequestUsageAudit) -> serde_
if let Some(service_tier) = item.provider_service_tier() {
payload["service_tier"] = json!(service_tier);
}
if let Some(actual_service_tier) = item.provider_actual_service_tier() {
payload["actual_service_tier"] = json!(actual_service_tier);
}
payload
}
@@ -1840,6 +1841,27 @@ mod tests {
assert_eq!(active_payload["client_is_stream"], false);
}
#[test]
fn user_usage_stream_defaults_to_non_stream_for_openai_search() {
let item = StoredRequestUsageAudit {
is_stream: false,
api_format: Some("openai:search".to_string()),
request_body: Some(json!({
"id": "session-search-1",
"model": "gpt-5.6-sol",
"input": "current documentation"
})),
..sample_usage("completed")
};
assert!(!users_me_usage_client_is_stream(&item));
let record_payload =
build_users_me_usage_record_payload(&item, false, &BTreeMap::new(), false);
assert_eq!(record_payload["client_requested_stream"], false);
assert_eq!(record_payload["client_is_stream"], false);
}
#[test]
fn user_usage_upstream_stream_prefers_request_metadata_flag() {
let item = StoredRequestUsageAudit {
apps/aether-gateway/src/handlers/shared/catalog.rs
+103 -2
View File
@@ -38,7 +38,7 @@ pub(crate) fn provider_catalog_key_supports_format(
}
formats
.iter()
.any(|candidate| crate::ai_serving::api_format_alias_matches(candidate, api_format))
.any(|candidate| crate::ai_serving::api_format_permission_covers(candidate, api_format))
}
pub(crate) fn decrypt_catalog_secret_with_fallbacks(
@@ -2054,6 +2054,29 @@ fn quota_snapshot_has_materialized_data(
})
}
fn codex_upstream_metadata_is_at_least_as_fresh(
quota_snapshot: Option<&Map<String, Value>>,
upstream_metadata: Option<&Value>,
) -> bool {
let Some(metadata) = provider_quota_metadata_bucket(upstream_metadata, "codex") else {
return false;
};
let Some(metadata_updated_at) = metadata
.get("updated_at")
.and_then(admin_provider_quota_pure::coerce_json_u64)
else {
return false;
};
let snapshot_updated_at = quota_snapshot.and_then(|quota| {
quota
.get("updated_at")
.or_else(|| quota.get("observed_at"))
.and_then(admin_provider_quota_pure::coerce_json_u64)
});
snapshot_updated_at.is_none_or(|updated_at| metadata_updated_at >= updated_at)
}
fn windsurf_quota_snapshot_has_stale_cooldown(quota_snapshot: &Map<String, Value>) -> bool {
let code = quota_snapshot
.get("code")
@@ -2121,8 +2144,15 @@ pub(crate) fn provider_key_status_snapshot_payload(
.and_then(Value::as_object)
.and_then(|snapshot| snapshot.get("quota"))
.and_then(Value::as_object);
let refresh_codex_snapshot = provider_type.trim().eq_ignore_ascii_case("codex")
&& codex_upstream_metadata_is_at_least_as_fresh(
quota_snapshot,
key.upstream_metadata.as_ref(),
);
let payload = if quota_snapshot_has_materialized_data(quota_snapshot, provider_type) {
let payload = if quota_snapshot_has_materialized_data(quota_snapshot, provider_type)
&& !refresh_codex_snapshot
{
status_snapshot
.cloned()
.unwrap_or_else(default_provider_key_status_snapshot)
@@ -2732,6 +2762,25 @@ mod tests {
.expect("key transport should build")
}
#[test]
fn responses_key_scope_covers_search_in_one_direction() {
let mut responses_key = sample_catalog_key();
responses_key.api_formats = Some(json!(["openai:responses"]));
assert!(provider_catalog_key_supports_format(
&responses_key,
"codex",
"openai:search",
));
let mut search_key = sample_catalog_key();
search_key.api_formats = Some(json!(["openai:search"]));
assert!(!provider_catalog_key_supports_format(
&search_key,
"codex",
"openai:responses",
));
}
#[test]
fn masked_catalog_api_key_handles_unicode_plaintext_without_panicking() {
let state = AppState::new().expect("gateway should build");
@@ -3493,6 +3542,58 @@ mod tests {
);
}
#[test]
fn provider_key_status_snapshot_payload_restores_complete_codex_cache() {
let mut key = sample_catalog_key();
key.upstream_metadata = Some(json!({
"codex": {
"updated_at": 200u64,
"plan_type": "plus",
"primary_used_percent": 89.0,
"primary_reset_at": 1_900_000_000u64,
"spark_primary_used_percent": 40.0,
"spark_primary_reset_at": 1_900_100_000u64,
"reset_credits": {
"available_count": 3,
"updated_at": 200u64,
"detail_status": "available",
"credits": []
}
}
}));
key.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "codex",
"updated_at": 200u64,
"windows": [{
"code": "weekly",
"used_ratio": 1.0,
"reset_at": 1_900_000_000u64
}]
}
}));
let payload = provider_key_status_snapshot_payload(&key, "codex");
let windows = payload["quota"]["windows"]
.as_array()
.expect("quota windows should exist");
assert_eq!(payload.pointer("/quota/plan_type"), Some(&json!("plus")));
assert_eq!(
payload.pointer("/quota/reset_credits/available_count"),
Some(&json!(3u64))
);
assert!(windows.iter().any(|window| window["code"] == "spark_5h"));
assert_eq!(
windows
.iter()
.find(|window| window["code"] == "weekly")
.and_then(|window| window.get("used_ratio")),
Some(&json!(0.89))
);
}
#[test]
fn sync_provider_key_quota_status_snapshot_preserves_codex_usage_state() {
let current_status_snapshot = json!({
apps/aether-gateway/src/handlers/shared/request_utils.rs
+5
View File
@@ -221,6 +221,11 @@ pub(crate) fn admin_proxy_local_requires_buffered_body(
| (Some("endpoints_manage"), http::Method::POST, Some("create_endpoint"))
| (Some("endpoints_manage"), http::Method::POST, Some("batch_delete_keys"))
| (Some("endpoints_manage"), http::Method::POST, Some("refresh_quota"))
| (
Some("endpoints_manage"),
http::Method::POST,
Some("codex_reset_credit_consume"),
)
| (Some("endpoints_manage"), http::Method::PUT, Some("update_key"))
| (Some("endpoints_manage"), http::Method::PUT, Some("update_endpoint"))
| (Some("modules_manage"), http::Method::PUT, Some("set_enabled"))