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 2654 additions and 181 deletions
@@ -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");
@@ -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");
}
@@ -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")