feat(kiro): add model group routing preferences (#61)

* feat(kiro): add model group routing preferences

* fix(kiro): repair model group routing CI

* ci: avoid repeated postgres test migrations

Author: acking-you

crates/frontend/src/api.rs

@@ -6132,6 +6132,8 @@ pub struct AdminLlmGatewayKeyView {
     #[serde(default = "default_anthropic_upstream_pool_mode")]
     pub kiro_anthropic_upstream_pool_mode: String,
     pub model_name_map: Option<BTreeMap<String, String>>,
+    #[serde(default)]
+    pub kiro_model_group_preferences: BTreeMap<String, String>,
     pub request_max_concurrency: Option<u64>,
     pub request_min_start_interval_ms: Option<u64>,
     #[serde(default = "default_true")]
@@ -9030,6 +9032,7 @@ pub async fn create_admin_llm_gateway_key(
             preferred_pool_strategy: default_kiro_pool_strategy(),
             kiro_anthropic_upstream_pool_mode: default_anthropic_upstream_pool_mode(),
             model_name_map: None,
+            kiro_model_group_preferences: BTreeMap::new(),
             request_max_concurrency,
             request_min_start_interval_ms,
             kiro_request_validation_enabled: true,
@@ -9096,6 +9099,7 @@ pub struct PatchAdminLlmGatewayKeyRequest<'a> {
     pub preferred_pool_strategy: Option<&'a str>,
     pub kiro_anthropic_upstream_pool_mode: Option<&'a str>,
     pub model_name_map: Option<&'a BTreeMap<String, String>>,
+    pub kiro_model_group_preferences: Option<&'a BTreeMap<String, String>>,
     pub request_max_concurrency: Option<u64>,
     pub request_min_start_interval_ms: Option<u64>,
     pub codex_fast_enabled: Option<bool>,
@@ -9136,6 +9140,7 @@ pub async fn patch_admin_llm_gateway_key(
             request.preferred_pool_strategy,
             request.kiro_anthropic_upstream_pool_mode,
             request.model_name_map,
+            request.kiro_model_group_preferences,
             request.request_max_concurrency,
             request.request_min_start_interval_ms,
             request.codex_fast_enabled,
@@ -9236,6 +9241,11 @@ pub async fn patch_admin_llm_gateway_key(
                 .map_err(|e| format!("Serialize error: {:?}", e))?;
             body.insert("model_name_map".to_string(), value);
         }
+        if let Some(preferences) = request.kiro_model_group_preferences {
+            let value = serde_json::to_value(preferences)
+                .map_err(|e| format!("Serialize error: {:?}", e))?;
+            body.insert("kiro_model_group_preferences".to_string(), value);
+        }
         if let Some(request_max_concurrency) = request.request_max_concurrency {
             body.insert(
                 "request_max_concurrency".to_string(),
@@ -11445,6 +11455,7 @@ pub async fn create_admin_kiro_key(
             preferred_pool_strategy: default_kiro_pool_strategy(),
             kiro_anthropic_upstream_pool_mode: default_anthropic_upstream_pool_mode(),
             model_name_map: None,
+            kiro_model_group_preferences: BTreeMap::new(),
             request_max_concurrency: None,
             request_min_start_interval_ms: None,
             kiro_request_validation_enabled: true,
@@ -11513,6 +11524,7 @@ pub async fn patch_admin_kiro_key(
             request.preferred_pool_strategy,
             request.kiro_anthropic_upstream_pool_mode,
             request.model_name_map,
+            request.kiro_model_group_preferences,
             request.request_max_concurrency,
             request.request_min_start_interval_ms,
             request.kiro_request_validation_enabled,
@@ -11608,6 +11620,11 @@ pub async fn patch_admin_kiro_key(
                 .map_err(|e| format!("Serialize error: {:?}", e))?;
             body.insert("model_name_map".to_string(), value);
         }
+        if let Some(preferences) = request.kiro_model_group_preferences {
+            let value = serde_json::to_value(preferences)
+                .map_err(|e| format!("Serialize error: {:?}", e))?;
+            body.insert("kiro_model_group_preferences".to_string(), value);
+        }
         if let Some(kiro_request_validation_enabled) = request.kiro_request_validation_enabled {
             body.insert(
                 "kiro_request_validation_enabled".to_string(),

crates/frontend/src/pages/admin_kiro_gateway.rs

@@ -1489,6 +1489,7 @@ fn key_editor_card(props: &KeyEditorCardProps) -> Html {
                     preferred_pool_strategy: None,
                     kiro_anthropic_upstream_pool_mode: None,
                     model_name_map: None,
+                    kiro_model_group_preferences: None,
                     request_max_concurrency: request_max_concurrency_value,
                     request_min_start_interval_ms: request_min_start_interval_ms_value,
                     codex_fast_enabled: Some(codex_fast_enabled_value),

crates/frontend/src/pages/admin_llm_gateway.rs

@@ -431,6 +431,7 @@ impl AdminKeyStore for EmptyAdminKeyStore {
             auto_account_names: None,
             preferred_pool_strategy: default_kiro_pool_strategy(),
             model_name_map: None,
+            kiro_model_group_preferences: std::collections::BTreeMap::new(),
             request_max_concurrency: key.request_max_concurrency,
             request_min_start_interval_ms: key.request_min_start_interval_ms,
             codex_fast_enabled: true,

crates/llm-access-core/src/store/empty.rs

@@ -70,6 +70,9 @@ pub struct AdminKey {
     pub preferred_pool_strategy: String,
     /// Model name mapping.
     pub model_name_map: Option<BTreeMap<String, String>>,
+    /// Exact Kiro request-model to preferred account-group map.
+    #[serde(default)]
+    pub kiro_model_group_preferences: BTreeMap<String, String>,
     /// Per-key request concurrency cap.
     pub request_max_concurrency: Option<u64>,
     /// Per-key request pacing interval.
@@ -382,6 +385,8 @@ pub struct AdminKeyPatch {
     pub preferred_pool_strategy: Option<String>,
     /// New model name map.
     pub model_name_map: Option<Option<BTreeMap<String, String>>>,
+    /// New exact Kiro request-model to preferred account-group map.
+    pub kiro_model_group_preferences: Option<BTreeMap<String, String>>,
     /// New per-key request concurrency cap.
     pub request_max_concurrency: Option<Option<u64>>,
     /// New per-key request pacing interval.
@@ -455,6 +460,7 @@ mod tests {
             auto_account_names: None,
             preferred_pool_strategy: "balanced".to_string(),
             model_name_map: None,
+            kiro_model_group_preferences: BTreeMap::new(),
             request_max_concurrency: None,
             request_min_start_interval_ms: None,
             codex_fast_enabled: true,

crates/llm-access-core/src/store/keys.rs

@@ -0,0 +1,68 @@
+//! Kiro model -> account-group preference helpers.
+//!
+//! ```text
+//! raw request model
+//!        |
+//!        +-- exact preference hit --> preferred account set
+//!        |
+//!        +-- no hit ---------------> existing affinity/order rules
+//! ```
+
+use std::collections::BTreeMap;
+
+/// Per-key exact model-name to Kiro account-group preference map.
+pub type KiroModelGroupPreferences = BTreeMap<String, String>;
+
+/// Normalize admin-supplied model preference rules.
+pub fn normalize_kiro_model_group_preferences(
+    input: KiroModelGroupPreferences,
+) -> KiroModelGroupPreferences {
+    input
+        .into_iter()
+        .filter_map(|(model, group_id)| {
+            let model = model.trim();
+            let group_id = group_id.trim();
+            (!model.is_empty() && !group_id.is_empty())
+                .then(|| (model.to_string(), group_id.to_string()))
+        })
+        .collect()
+}
+
+/// Resolve one exact model-name preference.
+pub fn kiro_model_group_preference<'a>(
+    preferences: &'a KiroModelGroupPreferences,
+    model: &str,
+) -> Option<&'a str> {
+    preferences.get(model.trim()).map(String::as_str)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn normalizes_model_group_preferences_by_trimming_and_dropping_empty_entries() {
+        let normalized = normalize_kiro_model_group_preferences(BTreeMap::from([
+            (" claude-sonnet-4 ".to_string(), " group-sonnet ".to_string()),
+            ("".to_string(), "group-empty-model".to_string()),
+            ("claude-opus-4".to_string(), " ".to_string()),
+        ]));
+
+        assert_eq!(
+            normalized,
+            BTreeMap::from([("claude-sonnet-4".to_string(), "group-sonnet".to_string())])
+        );
+    }
+
+    #[test]
+    fn resolves_model_group_preferences_by_exact_model_name() {
+        let preferences =
+            BTreeMap::from([("claude-sonnet-4".to_string(), "group-sonnet".to_string())]);
+
+        assert_eq!(
+            kiro_model_group_preference(&preferences, "claude-sonnet-4"),
+            Some("group-sonnet")
+        );
+        assert_eq!(kiro_model_group_preference(&preferences, "claude-sonnet"), None);
+    }
+}

crates/llm-access-core/src/store/kiro_model_routing.rs

@@ -14,6 +14,7 @@ mod empty;
 mod groups;
 mod keys;
 mod kiro_account;
+mod kiro_model_routing;
 mod proxy;
 mod public;
 mod routes;
@@ -71,6 +72,9 @@ pub use kiro_account::{
     ADMIN_KIRO_ACCOUNT_ISSUE_ABNORMAL, ADMIN_KIRO_ACCOUNT_ISSUE_AUTH_401,
     ADMIN_KIRO_ACCOUNT_ISSUE_DISABLED, ADMIN_KIRO_ACCOUNT_ISSUE_ERROR,
 };
+pub use kiro_model_routing::{
+    kiro_model_group_preference, normalize_kiro_model_group_preferences, KiroModelGroupPreferences,
+};
 pub use proxy::{
     default_proxy_bindings, AdminProxyBinding, AdminProxyConfig, AdminProxyConfigPatch,
     AdminProxyEndpointCheck, AdminProxyEndpointCheckUpdate, AdminProxyTrafficSnapshot,

crates/llm-access-core/src/store/mod.rs

@@ -1,6 +1,8 @@
 //! Provider routing contracts: authenticated key, provider proxy config,
 //! Codex/Kiro route + auth-update views, and JWT/auth-error helpers.
 
+use std::collections::BTreeMap;
+
 use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
 use serde_json::Value;
 
@@ -234,6 +236,8 @@ pub struct ProviderKiroRoute {
     pub cctest_proxy_api_key: Option<String>,
     /// JSON object mapping public model names to upstream Kiro model names.
     pub model_name_map_json: String,
+    /// Exact request-model to preferred account names resolved from key config.
+    pub model_group_preferred_account_names: BTreeMap<String, Vec<String>>,
     /// Effective Kiro cache k-model JSON for this key.
     pub cache_kmodels_json: String,
     /// Effective Kiro cache policy JSON for this key.

crates/llm-access-core/src/store/routes.rs

@@ -26,6 +26,9 @@ CREATE TABLE IF NOT EXISTS llm_key_route_config (
         preferred_pool_strategy IN ('balanced', 'credit_first')
     ),
     model_name_map_json JSONB,
+    kiro_model_group_preferences_json JSONB NOT NULL DEFAULT '{}'::jsonb CHECK (
+        jsonb_typeof(kiro_model_group_preferences_json) = 'object'
+    ),
     request_max_concurrency BIGINT,
     request_min_start_interval_ms BIGINT,
     codex_fast_enabled BOOLEAN NOT NULL DEFAULT TRUE,

crates/llm-access-migrations/migrations/postgres/0001_init.sql

@@ -0,0 +1,21 @@
+ALTER TABLE IF EXISTS llm_key_route_config
+    ADD COLUMN IF NOT EXISTS kiro_model_group_preferences_json JSONB NOT NULL DEFAULT '{}'::jsonb;
+
+UPDATE llm_key_route_config
+SET kiro_model_group_preferences_json = '{}'::jsonb
+WHERE kiro_model_group_preferences_json IS NULL
+   OR jsonb_typeof(kiro_model_group_preferences_json) <> 'object';
+
+DO $$
+BEGIN
+    IF NOT EXISTS (
+        SELECT 1
+        FROM pg_constraint
+        WHERE conname = 'ck_llm_key_route_config_kiro_model_group_preferences_object'
+          AND conrelid = 'llm_key_route_config'::regclass
+    ) THEN
+        ALTER TABLE llm_key_route_config
+            ADD CONSTRAINT ck_llm_key_route_config_kiro_model_group_preferences_object
+            CHECK (jsonb_typeof(kiro_model_group_preferences_json) = 'object');
+    END IF;
+END $$;