Skip to content

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

Merged
acking-you merged 3 commits into
masterfrom
feat/kiro-model-group-routing
Jul 2, 2026
Merged

feat(kiro): add model group routing preferences#61
acking-you merged 3 commits into
masterfrom
feat/kiro-model-group-routing

Conversation

@acking-you

Copy link
Copy Markdown
Owner

Summary

Add per-key Kiro model-to-account-group routing preferences so selected models can prefer a specific Kiro account group without expanding the main key editor card.

Root cause

Kiro routing only considered session affinity, route scope, pool strategy, latency, and local limits. There was no first-class way to bias a specific requested model toward a dedicated account group, and adding that inline to the existing admin card would make the Kiro key editor too long and heavier to render.

What changed

  • Added kiro_model_group_preferences to key config, Postgres schema/migrations, request cache snapshots, and route resolution.
  • Made model-group preference outrank Kiro session affinity; when a model preference is present, sticky affinity is ignored for that request.
  • Kept proxy cooldown and account throttling as availability gates, so preferred groups soft-fallback when their accounts cannot accept the request.
  • Added admin-side normalization and Kiro account group reference validation before persisting config.
  • Added a compact Kiro key-card summary plus a second-level drawer for configuring model-to-group routing rules.

Verification

  • cargo test -p llm-access-core -p llm-access-store -p llm-access --jobs 4
  • cargo test -p llm-access kiro_selection_model_group_preference --jobs 4 -- --nocapture
  • cargo check -p static-flow-frontend --target wasm32-unknown-unknown --jobs 4
  • cargo clippy -p llm-access-core -p llm-access-store -p llm-access --jobs 4 -- -D warnings
  • cargo clippy -p static-flow-frontend --target wasm32-unknown-unknown --jobs 4 -- -D warnings

Deployment note

Cloud llm-access needs the new Postgres migration before rolling out this API path.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces "Model Group Routing" for Kiro, allowing administrators to map specific models to preferred account groups. The changes span database migrations, Postgres repository updates, API endpoints, frontend UI configuration, and route selection logic to prioritize preferred accounts. The review feedback identifies a critical bug where inserting a column in the middle of a SELECT query shifts indices and breaks hardcoded database decoding. Additionally, suggestions are provided to improve robustness by gracefully falling back when configured account groups are missing or invalid, and to optimize route selection by only bypassing sticky affinity when the preferred account set is non-empty.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

Comment on lines 349 to 355
k.created_at_ms, k.updated_at_ms,
r.route_strategy, r.fixed_account_name, r.auto_account_names_json,
r.account_group_id, r.model_name_map_json,
r.kiro_model_group_preferences_json,
r.request_max_concurrency, r.request_min_start_interval_ms,
r.codex_fast_enabled, r.codex_strict_session_rejection_enabled,
r.codex_image_generation_enabled,

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

critical

Critical Bug: Column Index Shift in SELECT Query

In decode_key_bundle (defined in crates/llm-access-store/src/postgres/decode.rs), several columns are decoded using hardcoded indices (e.g., model_name_map_json at index 15, request_max_concurrency at index 16, request_min_start_interval_ms at index 17, and codex_fast_enabled at index 18).

By inserting r.kiro_model_group_preferences_json in the middle of the SELECT list (at line 352), you have shifted the indices of all subsequent columns by 1. This will cause decode_key_bundle to decode incorrect fields or fail with a database decoding error at runtime.

To fix this safely without breaking the hardcoded indices, move r.kiro_model_group_preferences_json to the end of the SELECT list (after r.codex_image_generation_enabled), matching the pattern used in the other queries in this file.

Suggested change
k.created_at_ms, k.updated_at_ms,
r.route_strategy, r.fixed_account_name, r.auto_account_names_json,
r.account_group_id, r.model_name_map_json,
r.kiro_model_group_preferences_json,
r.request_max_concurrency, r.request_min_start_interval_ms,
r.codex_fast_enabled, r.codex_strict_session_rejection_enabled,
r.codex_image_generation_enabled,
k.created_at_ms, k.updated_at_ms,
r.route_strategy, r.fixed_account_name, r.auto_account_names_json,
r.account_group_id, r.model_name_map_json,
r.request_max_concurrency, r.request_min_start_interval_ms,
r.codex_fast_enabled, r.codex_strict_session_rejection_enabled,
r.codex_image_generation_enabled,
r.kiro_model_group_preferences_json

Comment on lines +466 to +476
let group = self
.get_admin_account_group_row(&group_id)
.await?
.with_context(|| {
format!("configured kiro model group preference `{group_id}` does not exist")
})?;
if group.provider_type != core_store::PROVIDER_KIRO {
anyhow::bail!(
"configured kiro model group preference belongs to a different provider"
);
}

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Robustness Issue: Graceful Fallback for Missing/Invalid Account Groups

Currently, if an account group configured in kiro_model_group_preferences is deleted or belongs to a different provider, resolve_kiro_model_group_preferred_account_names will return an Err (via anyhow::bail! or with_context). This error bubbles up and causes the entire key's route resolution or cache snapshot generation to fail, making the key completely unusable.

To make the routing more robust, log a warning and skip the invalid preference instead of failing the entire operation. This ensures that other models or fallback routing still work even if one preference is misconfigured or pointing to a deleted group.

            let group = match self.get_admin_account_group_row(&group_id).await {
                Ok(Some(g)) => g,
                Ok(None) => {
                    tracing::warn!(
                        group_id = %group_id,
                        model = %model,
                        "Configured Kiro model group preference does not exist, skipping"
                    );
                    continue;
                }
                Err(err) => {
                    tracing::error!(
                        group_id = %group_id,
                        model = %model,
                        ?err,
                        "Failed to load Kiro model group preference, skipping"
                    );
                    continue;
                }
            };
            if group.provider_type != core_store::PROVIDER_KIRO {
                tracing::warn!(
                    group_id = %group_id,
                    model = %model,
                    provider_type = %group.provider_type,
                    "Configured Kiro model group preference belongs to a different provider, skipping"
                );
                continue;
            }

"kiro route is not configured",
));
}
let model_preference_active = model_preferred_account_names.is_some();

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Optimization: Only Ignore Sticky Affinity if Preferred Accounts are Available

Currently, model_preference_active is set to true as long as model_preferred_account_names is Some, even if the set of preferred accounts is empty (e.g., because the preferred group has no accounts or none of them overlap with the key's authorized accounts). This causes sticky affinity to be ignored, degrading session consistency and cache hit rates for no benefit since no preferred accounts can be selected anyway.

We should only treat the model preference as active if the preferred account set is non-empty.

Suggested change
let model_preference_active = model_preferred_account_names.is_some();
let model_preference_active = model_preferred_account_names.is_some_and(|set| !set.is_empty());

@acking-you acking-you merged commit 42ce06a into master Jul 2, 2026
3 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant