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

feat: scope model mappings by endpoint

Browse Source
This commit is contained in:
fawney19
2026-05-07 00:48:15 +08:00
parent 4e9f063385
commit 012f8bcdf7
58 changed files with 731 additions and 101 deletions
Download Patch File Download Diff File Expand all files Collapse all files

97
crates/aether-ai-formats/src/provider_compat/proxy/rules.rs
View File

@@ -34,6 +34,16 @@ pub fn header_rules_are_locally_supported(rules: Option<&Value>) -> bool {
rules.is_array()
}
pub fn header_rules_have_enabled_rules(rules: Option<&Value>) -> bool {
let Some(rules) = rules.and_then(Value::as_array) else {
return false;
};
rules
.iter()
.any(|rule| rule.as_object().is_some_and(header_rule_is_enabled))
}
pub fn apply_local_header_rules(
headers: &mut BTreeMap<String, String>,
rules: Option<&Value>,
@@ -85,6 +95,9 @@ fn apply_local_header_rules_inner(
let Some(rule) = rule.as_object() else {
continue;
};
if !header_rule_is_enabled(rule) {
continue;
}
if let Some(condition) = rule.get("condition").filter(|value| !value.is_null()) {
if !condition_is_locally_supported(condition) {
continue;
@@ -159,6 +172,10 @@ fn header_rule_value_to_string(value: &Value) -> String {
.unwrap_or_else(|| value.to_string())
}
fn header_rule_is_enabled(rule: &Map<String, Value>) -> bool {
rule.get("enabled").and_then(Value::as_bool) != Some(false)
}
pub fn body_rules_are_locally_supported(rules: Option<&Value>) -> bool {
let Some(rules) = rules else {
return true;
@@ -166,6 +183,16 @@ pub fn body_rules_are_locally_supported(rules: Option<&Value>) -> bool {
rules.is_array()
}
pub fn body_rules_have_enabled_rules(rules: Option<&Value>) -> bool {
let Some(rules) = rules.and_then(Value::as_array) else {
return false;
};
rules
.iter()
.any(|rule| rule.as_object().is_some_and(body_rule_is_enabled))
}
pub fn body_rules_handle_path(rules: Option<&Value>, path: &str) -> bool {
let Some(target_path) = parse_body_path(path) else {
return false;
@@ -178,6 +205,9 @@ pub fn body_rules_handle_path(rules: Option<&Value>, path: &str) -> bool {
let Some(rule) = rule.as_object() else {
return false;
};
if !body_rule_is_enabled(rule) {
return false;
}
match rule
.get("action")
.and_then(Value::as_str)
@@ -249,6 +279,9 @@ fn apply_local_body_rules_inner(
let Some(rule) = rule.as_object() else {
continue;
};
if !body_rule_is_enabled(rule) {
continue;
}
let condition = rule.get("condition").filter(|value| !value.is_null());
let item_condition = condition.is_some_and(condition_has_item_ref);
@@ -451,6 +484,10 @@ fn apply_local_body_rules_inner(
true
}
fn body_rule_is_enabled(rule: &Map<String, Value>) -> bool {
rule.get("enabled").and_then(Value::as_bool) != Some(false)
}
fn condition_is_locally_supported(condition: &Value) -> bool {
let Some(condition) = condition.as_object() else {
return false;
@@ -1295,8 +1332,8 @@ mod tests {
use super::{
apply_local_body_rules, apply_local_body_rules_with_request_headers,
apply_local_header_rules, apply_local_header_rules_with_request_headers,
body_rules_are_locally_supported, body_rules_handle_path,
header_rules_are_locally_supported,
body_rules_are_locally_supported, body_rules_handle_path, body_rules_have_enabled_rules,
header_rules_are_locally_supported, header_rules_have_enabled_rules,
};
#[test]
@@ -1609,6 +1646,33 @@ mod tests {
assert_eq!(body["ok"], true);
}
#[test]
fn body_rules_skip_disabled_entries() {
let rules = serde_json::json!([
{"enabled":false,"action":"set","path":"metadata.disabled","value":true},
{"enabled":false,"action":"drop","path":"keep"},
{"action":"set","path":"metadata.enabled","value":true}
]);
let mut body = serde_json::json!({
"keep": true,
"metadata": {}
});
assert!(body_rules_have_enabled_rules(Some(&rules)));
assert!(apply_local_body_rules(&mut body, Some(&rules), None));
assert_eq!(body["keep"], true);
assert!(body["metadata"].get("disabled").is_none());
assert_eq!(body["metadata"]["enabled"], true);
assert!(!body_rules_handle_path(Some(&rules), "metadata.disabled"));
assert!(body_rules_handle_path(Some(&rules), "metadata.enabled"));
let disabled_only = serde_json::json!([
{"enabled":false,"action":"set","path":"metadata.disabled","value":true}
]);
assert!(!body_rules_have_enabled_rules(Some(&disabled_only)));
}
#[test]
fn header_rules_skip_invalid_entries_without_rejecting_whole_headers() {
let rules = serde_json::json!([
@@ -1645,6 +1709,35 @@ mod tests {
assert!(!headers.contains_key("x-legacy"));
}
#[test]
fn header_rules_skip_disabled_entries() {
let rules = serde_json::json!([
{"enabled":false,"action":"set","key":"x-disabled","value":"bad"},
{"enabled":false,"action":"drop","key":"x-keep"},
{"action":"set","key":"x-enabled","value":"ok"}
]);
assert!(header_rules_have_enabled_rules(Some(&rules)));
let mut headers =
std::collections::BTreeMap::from([("x-keep".to_string(), "yes".to_string())]);
assert!(apply_local_header_rules(
&mut headers,
Some(&rules),
&[],
&serde_json::json!({}),
None,
));
assert_eq!(headers.get("x-keep").map(String::as_str), Some("yes"));
assert!(!headers.contains_key("x-disabled"));
assert_eq!(headers.get("x-enabled").map(String::as_str), Some("ok"));
let disabled_only = serde_json::json!([
{"enabled":false,"action":"set","key":"x-disabled","value":"bad"}
]);
assert!(!header_rules_have_enabled_rules(Some(&disabled_only)));
}
#[test]
fn body_rules_handle_exact_and_wildcard_paths() {
let rules = serde_json::json!([

2
crates/aether-data-contracts/src/repository/candidate_selection/types.rs
View File

@@ -5,6 +5,8 @@ pub struct StoredProviderModelMapping {
pub name: String,
pub priority: i32,
pub api_formats: Option<Vec<String>>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub endpoint_ids: Option<Vec<String>>,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]

5
crates/aether-data/src/repository/candidate_selection/memory.rs
View File

@@ -167,6 +167,10 @@ fn row_matches_requested_model(
formats
.iter()
.any(|value| api_format_matches(value, api_format))
}) && mapping.endpoint_ids.as_ref().is_none_or(|endpoint_ids| {
endpoint_ids
.iter()
.any(|endpoint_id| endpoint_id == &row.endpoint_id)
}) && mapping.name == requested_model_name
})
})
@@ -281,6 +285,7 @@ mod tests {
name: "alias-gpt-4.1".to_string(),
priority: 0,
api_formats: Some(vec!["openai:chat".to_string()]),
endpoint_ids: None,
},
]);
let repository = InMemoryMinimalCandidateSelectionReadRepository::seed(vec![

11
crates/aether-data/src/repository/candidate_selection/mysql.rs
View File

@@ -272,6 +272,10 @@ fn row_matches_requested_model(
formats
.iter()
.any(|value| api_format_matches(value, api_format))
}) && mapping.endpoint_ids.as_ref().is_none_or(|endpoint_ids| {
endpoint_ids
.iter()
.any(|endpoint_id| endpoint_id == &row.endpoint_id)
}) && mapping.name == requested_model_name
})
})
@@ -514,6 +518,7 @@ fn parse_embedded_provider_model_mappings(
name: raw.to_string(),
priority: 1,
api_formats: None,
endpoint_ids: None,
}]))
}
@@ -533,6 +538,7 @@ fn parse_provider_model_mappings_array(
name: raw.trim().to_string(),
priority: 1,
api_formats: None,
endpoint_ids: None,
});
}
_ => {}
@@ -573,6 +579,10 @@ fn parse_provider_model_mapping_object_lenient(
.map(|value| normalize_api_format(&value))
.collect()
});
let endpoint_ids = parse_string_list(
object.get("endpoint_ids").cloned(),
"models.provider_model_mappings.endpoint_ids",
)?;
Ok(Some(StoredProviderModelMapping {
name: name.to_string(),
@@ -582,6 +592,7 @@ fn parse_provider_model_mapping_object_lenient(
))
})?,
api_formats,
endpoint_ids,
}))
}

20
crates/aether-data/src/repository/candidate_selection/postgres.rs
View File

@@ -694,6 +694,15 @@ fn requested_model_selection_sql() -> String {
WHERE LOWER(BTRIM(fmt.value)) = ANY($3::text[])
)
)
AND (
mapping.value -> 'endpoint_ids' IS NULL
OR jsonb_typeof(mapping.value -> 'endpoint_ids') <> 'array'
OR EXISTS (
SELECT 1
FROM jsonb_array_elements_text(mapping.value -> 'endpoint_ids') AS endpoint(value)
WHERE endpoint.value = pe.id
)
)
)
)
)"#,
@@ -933,6 +942,7 @@ fn parse_embedded_provider_model_mappings(
name: raw.to_string(),
priority: 1,
api_formats: None,
endpoint_ids: None,
}]))
}
@@ -954,6 +964,7 @@ fn parse_provider_model_mappings_array(
name: raw.to_string(),
priority: 1,
api_formats: None,
endpoint_ids: None,
});
}
}
@@ -1006,6 +1017,10 @@ fn parse_provider_model_mapping_object_lenient(
.map(|value| aether_ai_formats::normalize_api_format_alias(&value))
.collect()
});
let endpoint_ids = parse_string_list(
object.get("endpoint_ids").cloned(),
"models.provider_model_mappings.endpoint_ids",
)?;
Ok(Some(StoredProviderModelMapping {
name: name.to_string(),
@@ -1015,6 +1030,7 @@ fn parse_provider_model_mapping_object_lenient(
))
})?,
api_formats,
endpoint_ids,
}))
}
@@ -1119,6 +1135,7 @@ mod tests {
name: "gpt-5.2".to_string(),
priority: 2,
api_formats: Some(vec!["openai:chat".to_string()]),
endpoint_ids: None,
}])
);
}
@@ -1134,6 +1151,7 @@ mod tests {
name: "gpt-5.2".to_string(),
priority: 1,
api_formats: None,
endpoint_ids: None,
}])
);
}
@@ -1156,11 +1174,13 @@ mod tests {
name: "gpt-5.2".to_string(),
priority: 1,
api_formats: None,
endpoint_ids: None,
},
StoredProviderModelMapping {
name: "gpt-5.2-mini".to_string(),
priority: 1,
api_formats: None,
endpoint_ids: None,
}
])
);

11
crates/aether-data/src/repository/candidate_selection/sqlite.rs
View File

@@ -272,6 +272,10 @@ fn row_matches_requested_model(
formats
.iter()
.any(|value| api_format_matches(value, api_format))
}) && mapping.endpoint_ids.as_ref().is_none_or(|endpoint_ids| {
endpoint_ids
.iter()
.any(|endpoint_id| endpoint_id == &row.endpoint_id)
}) && mapping.name == requested_model_name
})
})
@@ -514,6 +518,7 @@ fn parse_embedded_provider_model_mappings(
name: raw.to_string(),
priority: 1,
api_formats: None,
endpoint_ids: None,
}]))
}
@@ -533,6 +538,7 @@ fn parse_provider_model_mappings_array(
name: raw.trim().to_string(),
priority: 1,
api_formats: None,
endpoint_ids: None,
});
}
_ => {}
@@ -573,6 +579,10 @@ fn parse_provider_model_mapping_object_lenient(
.map(|value| normalize_api_format(&value))
.collect()
});
let endpoint_ids = parse_string_list(
object.get("endpoint_ids").cloned(),
"models.provider_model_mappings.endpoint_ids",
)?;
Ok(Some(StoredProviderModelMapping {
name: name.to_string(),
@@ -582,6 +592,7 @@ fn parse_provider_model_mapping_object_lenient(
))
})?,
api_formats,
endpoint_ids,
}))
}

5
crates/aether-provider-transport/src/antigravity/policy.rs
View File

@@ -9,6 +9,7 @@ use super::request::{
classify_antigravity_safe_request_body, AntigravityEnvelopeRequestType,
AntigravityRequestEnvelopeUnsupportedReason,
};
use crate::rules::{body_rules_have_enabled_rules, header_rules_have_enabled_rules};
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AntigravityRequestSideSpec {
@@ -76,12 +77,12 @@ pub fn classify_local_antigravity_request_support(
AntigravityRequestSideUnsupportedReason::UnsupportedCustomPath,
);
}
if transport.endpoint.header_rules.is_some() {
if header_rules_have_enabled_rules(transport.endpoint.header_rules.as_ref()) {
return AntigravityRequestSideSupport::Unsupported(
AntigravityRequestSideUnsupportedReason::UnsupportedHeaderRules,
);
}
if transport.endpoint.body_rules.is_some() {
if body_rules_have_enabled_rules(transport.endpoint.body_rules.as_ref()) {
return AntigravityRequestSideSupport::Unsupported(
AntigravityRequestSideUnsupportedReason::UnsupportedBodyRules,
);

3
crates/aether-provider-transport/src/gemini_files.rs
View File

@@ -6,6 +6,7 @@ use crate::auth::{build_passthrough_headers_with_auth, resolve_local_gemini_auth
use crate::policy::local_gemini_transport_unsupported_reason_with_network;
use crate::rules::{
apply_local_body_rules_with_request_headers, apply_local_header_rules_with_request_headers,
body_rules_have_enabled_rules,
};
use crate::snapshot::GatewayProviderTransportSnapshot;
use crate::url::build_gemini_files_passthrough_url;
@@ -87,7 +88,7 @@ pub fn build_gemini_files_request_body(
} else {
None
};
if provider_request_body_base64.is_some() && body_rules.is_some() {
if provider_request_body_base64.is_some() && body_rules_have_enabled_rules(body_rules) {
return Err(GeminiFilesRequestBodyError::BodyRulesUnsupportedForBinaryUpload);
}
if let Some(body) = provider_request_body.as_mut() {

3
crates/aether-provider-transport/src/lib.rs
View File

@@ -77,7 +77,8 @@ pub use request_url::{
pub use rules::{
apply_local_body_rules, apply_local_body_rules_with_request_headers, apply_local_header_rules,
apply_local_header_rules_with_request_headers, body_rules_are_locally_supported,
body_rules_handle_path, header_rules_are_locally_supported,
body_rules_handle_path, body_rules_have_enabled_rules, header_rules_are_locally_supported,
header_rules_have_enabled_rules,
};
pub use same_format_provider::{
build_same_format_provider_headers, build_same_format_provider_request_body,

1
crates/aether-scheduler-core/src/candidate/mod.rs
View File

@@ -70,6 +70,7 @@ mod tests {
name: format!("gpt-5-canary-{id}"),
priority: 1,
api_formats: Some(vec!["openai:chat".to_string()]),
endpoint_ids: None,
}]),
model_supports_streaming: None,
model_is_active: true,

35
crates/aether-scheduler-core/src/model.rs
View File

@@ -41,7 +41,7 @@ pub fn resolve_requested_global_model_name_with_model_directives(
.as_ref()
.is_some_and(|mappings| {
mappings.iter().any(|mapping| {
mapping_scope_matches(mapping, api_format)
mapping_scope_matches(mapping, row, api_format)
&& mapping.name == requested_model_name
})
})
@@ -93,7 +93,7 @@ fn row_supports_requested_model_exact(
.as_ref()
.is_some_and(|mappings| {
mappings.iter().any(|mapping| {
mapping_scope_matches(mapping, api_format)
mapping_scope_matches(mapping, row, api_format)
&& mapping.name == requested_model_name
})
})
@@ -201,7 +201,7 @@ pub fn select_provider_model_name(
mappings
.iter()
.filter(|mapping| mapping_scope_matches(mapping, api_format))
.filter(|mapping| mapping_scope_matches(mapping, row, api_format))
.min_by(|left, right| {
left.priority
.cmp(&right.priority)
@@ -218,7 +218,7 @@ pub fn candidate_model_names(
let mut names = BTreeSet::from([row.model_provider_model_name.clone()]);
if let Some(mappings) = row.model_provider_model_mappings.as_ref() {
for mapping in mappings {
if mapping_scope_matches(mapping, api_format) {
if mapping_scope_matches(mapping, row, api_format) {
names.insert(mapping.name.clone());
}
}
@@ -226,14 +226,25 @@ pub fn candidate_model_names(
names
}
fn mapping_scope_matches(mapping: &StoredProviderModelMapping, api_format: &str) -> bool {
let Some(api_formats) = mapping.api_formats.as_ref() else {
return true;
};
fn mapping_scope_matches(
mapping: &StoredProviderModelMapping,
row: &StoredMinimalCandidateSelectionRow,
api_format: &str,
) -> bool {
let api_format_matches_scope = mapping.api_formats.as_ref().is_none_or(|api_formats| {
api_formats
.iter()
.any(|value| api_format_matches(value, api_format))
});
if !api_format_matches_scope {
return false;
}
api_formats
.iter()
.any(|value| api_format_matches(value, api_format))
mapping.endpoint_ids.as_ref().is_none_or(|endpoint_ids| {
endpoint_ids
.iter()
.any(|endpoint_id| endpoint_id == &row.endpoint_id)
})
}
pub fn row_supports_required_capability(
@@ -353,7 +364,7 @@ fn row_has_candidate_model_name(
.as_ref()
.is_some_and(|mappings| {
mappings.iter().any(|mapping| {
mapping_scope_matches(mapping, api_format) && mapping.name == model_name
mapping_scope_matches(mapping, row, api_format) && mapping.name == model_name
})
})
}