feat: add model directive management

This commit is contained in:
fawney19
2026-05-03 14:48:25 +08:00
parent fe27fb17fb
commit c4ea042eb4
53 changed files with 2655 additions and 182 deletions

View File

@@ -63,23 +63,44 @@ pub fn auth_constraints_allow_model(
constraints: Option<&SchedulerAuthConstraints>,
requested_model_name: &str,
resolved_global_model_name: &str,
) -> bool {
auth_constraints_allow_model_with_model_directives(
constraints,
requested_model_name,
resolved_global_model_name,
false,
)
}
pub fn auth_constraints_allow_model_with_model_directives(
constraints: Option<&SchedulerAuthConstraints>,
requested_model_name: &str,
resolved_global_model_name: &str,
enable_model_directives: bool,
) -> bool {
let Some(allowed) = constraints.and_then(|constraints| constraints.allowed_models.as_deref())
else {
return true;
};
allowed
.iter()
.any(|value| value == requested_model_name || value == resolved_global_model_name)
let base_model = enable_model_directives
.then(|| aether_ai_formats::model_directive_base_model(requested_model_name))
.flatten();
allowed.iter().any(|value| {
value == requested_model_name
|| value == resolved_global_model_name
|| base_model
.as_ref()
.is_some_and(|base_model| value == base_model)
})
}
#[cfg(test)]
mod tests {
use super::{
api_format_matches_allowed_value, auth_constraints_allow_api_format,
auth_constraints_allow_model, auth_constraints_allow_provider,
provider_matches_allowed_value, SchedulerAuthConstraints,
auth_constraints_allow_model, auth_constraints_allow_model_with_model_directives,
auth_constraints_allow_provider, provider_matches_allowed_value, SchedulerAuthConstraints,
};
fn sample_constraints() -> SchedulerAuthConstraints {
@@ -200,6 +221,23 @@ mod tests {
));
}
#[test]
fn model_directive_base_model_requires_explicit_enablement() {
let constraints = sample_constraints();
assert!(!auth_constraints_allow_model(
Some(&constraints),
"gpt-5-high",
"gpt-5-high"
));
assert!(auth_constraints_allow_model_with_model_directives(
Some(&constraints),
"gpt-5-high",
"gpt-5-high",
true
));
}
#[test]
fn api_format_allowed_value_matches_current_signatures_only() {
assert!(api_format_matches_allowed_value(