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

feat: add model directive management

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

2
apps/aether-gateway/src/scheduler/candidate/enumeration.rs
View File

@@ -14,6 +14,7 @@ pub(super) async fn enumerate_scheduler_candidates(
require_streaming: bool,
required_capabilities: Option<&serde_json::Value>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
enable_model_directives: bool,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
enumerate_minimal_candidate_selection_with_required_capabilities(
selection_row_source,
@@ -22,6 +23,7 @@ pub(super) async fn enumerate_scheduler_candidates(
require_streaming,
auth_snapshot,
required_capabilities,
enable_model_directives,
)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))

5
apps/aether-gateway/src/scheduler/candidate/mod.rs
View File

@@ -57,6 +57,7 @@ pub(crate) async fn list_selectable_candidates(
required_capabilities: Option<&serde_json::Value>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
now_unix_secs: u64,
enable_model_directives: bool,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
collect_selectable_candidates(
selection_row_source,
@@ -67,6 +68,7 @@ pub(crate) async fn list_selectable_candidates(
required_capabilities,
auth_snapshot,
now_unix_secs,
enable_model_directives,
)
.await
}
@@ -87,6 +89,7 @@ pub(crate) async fn list_selectable_candidates_with_skip_reasons(
required_capabilities: Option<&serde_json::Value>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
now_unix_secs: u64,
enable_model_directives: bool,
) -> Result<
(
Vec<SchedulerMinimalCandidateSelectionCandidate>,
@@ -103,6 +106,7 @@ pub(crate) async fn list_selectable_candidates_with_skip_reasons(
required_capabilities,
auth_snapshot,
now_unix_secs,
enable_model_directives,
)
.await
}
@@ -181,6 +185,7 @@ pub(crate) async fn list_selectable_candidates_for_required_capability_without_r
required_capabilities.as_ref(),
auth_snapshot,
now_unix_secs,
false,
)
.await?;
all_attempts_blocked_by_auth_limit &=

6
apps/aether-gateway/src/scheduler/candidate/selection.rs
View File

@@ -42,6 +42,7 @@ pub(super) async fn select_minimal_candidate(
required_capabilities: Option<&serde_json::Value>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
now_unix_secs: u64,
enable_model_directives: bool,
) -> Result<Option<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
let affinity_cache_key =
build_scheduler_affinity_cache_key(auth_snapshot, api_format, global_model_name);
@@ -54,6 +55,7 @@ pub(super) async fn select_minimal_candidate(
required_capabilities,
auth_snapshot,
now_unix_secs,
enable_model_directives,
)
.await?
.into_iter()
@@ -73,6 +75,7 @@ pub(super) async fn collect_selectable_candidates(
required_capabilities: Option<&serde_json::Value>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
now_unix_secs: u64,
enable_model_directives: bool,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
Ok(collect_selectable_candidates_with_skip_reasons(
selection_row_source,
@@ -83,6 +86,7 @@ pub(super) async fn collect_selectable_candidates(
required_capabilities,
auth_snapshot,
now_unix_secs,
enable_model_directives,
)
.await?
.0)
@@ -97,6 +101,7 @@ pub(super) async fn collect_selectable_candidates_with_skip_reasons(
required_capabilities: Option<&serde_json::Value>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
now_unix_secs: u64,
enable_model_directives: bool,
) -> Result<
(
Vec<SchedulerMinimalCandidateSelectionCandidate>,
@@ -114,6 +119,7 @@ pub(super) async fn collect_selectable_candidates_with_skip_reasons(
require_streaming,
required_capabilities,
auth_snapshot,
enable_model_directives,
)
.await?;
let runtime_snapshot =

3
apps/aether-gateway/src/scheduler/candidate/tests/affinity.rs
View File

@@ -46,6 +46,7 @@ async fn select_candidate(
None,
auth_snapshot,
now_unix_secs,
false,
)
.await
}
@@ -92,7 +93,7 @@ async fn same_priority_candidates_are_distributed_by_affinity_key() {
let auth_snapshot = sample_auth_snapshot("affinity-key-1");
let (_resolved_global_model_name, rows) =
read_requested_model_rows(&state, "openai:chat", "gpt-4.1")
read_requested_model_rows(&state, "openai:chat", "gpt-4.1", false)
.await
.expect("selection rows should read")
.expect("selection rows should match requested model");

43
apps/aether-gateway/src/scheduler/candidate/tests/model.rs
View File

@@ -142,6 +142,7 @@ async fn enumerate_minimal_candidate_selection_resolves_provider_model_alias() {
false,
None,
None,
false,
)
.await
.expect("selection should succeed");
@@ -191,6 +192,7 @@ async fn enumerate_minimal_candidate_selection_keeps_only_resolved_global_model_
false,
None,
None,
false,
)
.await
.expect("selection should succeed");
@@ -231,6 +233,7 @@ async fn enumerate_minimal_candidate_selection_allows_resolved_global_model_in_a
false,
Some(&auth_snapshot),
None,
false,
)
.await
.expect("selection should succeed");
@@ -238,3 +241,43 @@ async fn enumerate_minimal_candidate_selection_allows_resolved_global_model_in_a
assert_eq!(selection.len(), 1);
assert_eq!(selection[0].global_model_name, "gpt-5");
}
#[tokio::test]
async fn enumerate_minimal_candidate_selection_gates_model_directive_fallback() {
let mut row = sample_row();
row.global_model_name = "gpt-5.4".to_string();
row.model_provider_model_name = "gpt-5.4-upstream".to_string();
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
row,
]));
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
let state = GatewayDataState::with_candidate_selection_and_quota_for_tests(candidates, quotas);
let disabled = enumerate_minimal_candidate_selection_with_required_capabilities(
&state,
"openai:chat",
"gpt-5.4-high",
false,
None,
None,
false,
)
.await
.expect("selection should succeed");
assert!(disabled.is_empty());
let enabled = enumerate_minimal_candidate_selection_with_required_capabilities(
&state,
"openai:chat",
"gpt-5.4-high",
false,
None,
None,
true,
)
.await
.expect("selection should succeed");
assert_eq!(enabled.len(), 1);
assert_eq!(enabled[0].global_model_name, "gpt-5.4");
}

4
apps/aether-gateway/src/scheduler/candidate/tests/selection.rs
View File

@@ -48,6 +48,7 @@ async fn select_candidate(
None,
auth_snapshot,
now_unix_secs,
false,
)
.await
}
@@ -70,6 +71,7 @@ async fn collect_selectable_candidates(
None,
auth_snapshot,
now_unix_secs,
false,
)
.await
}
@@ -98,6 +100,7 @@ async fn collect_selectable_candidates_with_skip_reasons(
None,
auth_snapshot,
now_unix_secs,
false,
)
.await
}
@@ -404,6 +407,7 @@ async fn scheduler_selection_prefers_required_capability_matches_before_priority
Some(&required_capabilities),
None,
100,
false,
)
.await
.expect("selection should succeed")