feat(llm-access-kiro): spread new sessions + band latency for fair account scheduling

[acking-you]

Problem

Two fairness defects in the Kiro account selector (selection_ordered_kiro_routes):

1. No-sample accounts get starved. Selection compared exact smoothed latency scores as the second key. An account without latency samples falls back to the global average; any account measured even marginally faster than the average permanently out-sorts it. The last_started_at round-robin tier almost never fired, so the single fastest account won every request. "No sample" is not "high latency" — but the old ordering treated it that way.

2. New sessions all stack on one account. A new session (has a session id but no prior affinity) followed pure latency and was pinned to the same fastest account. Per-session affinity itself was correct (same session → same account for 6h), but additional new sessions never spread — they all landed on that one account too.

Both stem from the same root: latency was too dominant, producing greedy winner-take-all.

Changes

1. Latency banding (kiro_latency.rs)

• Replace route_score_ms with route_score_band: quantize the smoothed score into bands of 15% of the global average (BAND_FRACTION). Accounts in the same band tie and fall through to the round-robin / balance tiers. A small latency edge no longer pins traffic.
• Scoring logic extracted into PreparedKiroLatencySnapshot::route_score (shared); the three guards (latency_routing_enabled, snapshot staleness, finite global avg) are preserved so legacy-order fallbacks are unchanged.

2. New-session spread (kiro_session_affinity.rs, route_selection.rs, kiro_dispatch.rs)

• KiroSessionAffinity::account_session_counts() tallies non-expired sessions per account by iterating the LRU (iter does not perturb recency), skipping TTL-expired entries.
• selection_ordered_kiro_routes / select_kiro_route_with_account_permit take an optional session_counts. When present (new session only), fewest bound sessions is the primary balance key, with latency band as tiebreaker.
• Wired into both the main and websearch dispatch paths. Computed once per request (stable across failover retries, since affinity only mutates on success).

Resulting sort order

• New session (session id, no affinity): proxy health → fewest bound sessions → latency band → last_started → credits → name
• Stateless / sticky-hit: proxy health → latency band → last_started → credits → name (unchanged shape, exact-score → band is the only diff)

Existing affinity stays sticky; the preferred-account fast path is untouched.

Notes

• Self-balancing: as sessions bind, counts diverge, so subsequent new sessions flow to the laggards. Cold start (all counts 0) ties through to band + last_started.
• Snapshot semantics: session_counts is a point-in-time snapshot; it does not react to sessions bound by concurrent in-flight requests between retries — acceptable for fairness.
• State remains per-process (no cross-instance sharing), same as before.

Tests

• account_session_counts_tallies_live_entries_per_account + ..._skips_expired_entries
• kiro_selection_spreads_new_session_to_least_bound_account (slower account with fewer sessions wins for a new session)
• kiro_selection_treats_same_band_accounts_as_equal_for_round_robin (no-sample account not starved by marginal edge)
• All 5 existing ordering tests still pass (band preserves the three guards).

cargo test -p llm-access → 295 passed. cargo clippy -p llm-access -- -D warnings → clean. rustfmt on changed files only.

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request introduces latency-score bands and session-balancing to prevent starving accounts with fewer samples and to distribute new sessions more evenly. Specifically, it quantizes latency scores into bands using a new BAND_FRACTION constant and adds session-affinity tracking to count active sessions per account, prioritizing accounts with fewer bound sessions during route selection. Feedback on these changes highlights a performance and lock contention hazard in account_session_counts_at, where string allocations are performed on every iteration while holding a mutex lock; using a borrowed lookup can optimize this path.

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.

[gemini-code-assist]

Performance & Lock Contention Hazard

Calling entry.account_name.to_string() on every single non-expired entry in the LRU cache (which can hold up to 65,536 entries) allocates a new String on every iteration, even if the account is already present in the counts map.

This is highly inefficient and holds the global entries mutex lock for much longer than necessary, causing severe lock contention on the hot path of every new session request.

By using counts.get_mut(entry.account_name.as_ref()), we can query the map using a borrowed &str without allocating. We only allocate a String when inserting a new account name for the first time, which reduces the number of allocations from O(N) to O(A) (where A is the number of unique accounts, typically very small).

        let mut counts: HashMap<String, usize> = HashMap::new();
        for (_, entry) in entries.iter() {
            if now.saturating_duration_since(entry.updated_at) > self.ttl {
                continue;
            }
            if let Some(count) = counts.get_mut(entry.account_name.as_ref()) {
                *count += 1;
            } else {
                counts.insert(entry.account_name.to_string(), 1);
            }
        }

References

1. Avoid allocating temporary collections (such as Vec<String>) and cloning elements solely for hashing. Instead, hash borrowed segments directly by chaining iterators or updating the hasher incrementally, especially on performance-critical hot paths.

[acking-you]

Good catch — adopted in ac98e5e. account_session_counts_at now borrows the name for the lookup (get_mut(entry.account_name.as_ref())) and only allocates a String when first inserting an account, so allocations drop from O(entries) to O(distinct accounts) and the affinity mutex is held for much less time on the new-session hot path.

I also tightened the same hot path further: the per-account scan is now skipped entirely when a key has a single candidate account (route_strategy=fixed) via a routes.len() > 1 guard at both dispatch sites, since there is nothing to spread there. Behavior is unchanged.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 5138d99f9b

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Aggregate session counts by routing identity

When a route list contains multiple account names for the same upstream Kiro user, this new-session key counts only the individual account_name, while the scheduler, cooldowns, and quota exhaustion operate on route.routing_identity (the next line already reads last-started by identity, and quota failover groups aliases by identity). In a set like alpha/beta sharing one routing identity plus gamma on another, existing sessions on alpha do not raise beta's count, so new sessions can keep being steered to the same upstream account through its alias instead of the actually less-bound identity. Aggregate the affinity counts over the candidate routes' routing identities before applying the fewest-sessions sort key.

Useful? React with 👍 / 👎.

[acking-you]

Agreed, real bug — fixed in da51954. selection_ordered_kiro_routes now folds the per-account affinity counts onto the candidate routes' routing_identity before applying the fewest-sessions key, so aliases of one upstream account share a single load figure — matching how the scheduler, cooldowns, last-started, and quota failover already group. O(routes), keys borrow from routes (zero-alloc). Added kiro_selection_spreads_new_session_by_routing_identity_not_account_name: with alpha/beta on one identity (alpha=3 sessions) and gamma on another, a new session now goes to gamma instead of beta (which would win on per-account counts as the high-credit alias of the loaded identity).

[chatgpt-codex-connector]

Anchor latency bands around the neutral score

Because the band index is floor(score / width) with buckets anchored at zero, the neutral no-sample score (global_avg) often sits close to a bucket boundary, so accounts that are still within the configured 15% tolerance can land in a lower band and continue to win before round-robin. For example, with a 500 ms global average the width is 75 ms: a no-sample route is band 6, while a sampled 449 ms route (only ~10% faster) is band 5 and will still permanently sort ahead until its measured latency crosses an arbitrary boundary. If the intent is to treat near-global/no-sample accounts as equal, compute the band relative to the global baseline or otherwise center the neutral score inside the same tolerance bucket.

Useful? React with 👍 / 👎.

[acking-you]

Good catch — fixed in da51954. Bands are now centered on the neutral baseline: round((score - global_avg) / (2 * tolerance)) with tolerance = global_avg * BAND_FRACTION, so band 0 is global_avg ± tolerance and the no-sample score sits dead-center instead of on a bucket boundary. Your 500ms/449ms example now shares band 0. Reworked the same-band test to a 430ms route score (~14% under the 500ms average): the old floor-from-zero scheme bucketed it at band 5 vs neutral's band 6 (permanent win), the centered scheme keeps both in band 0 so round-robin takes over. BAND_FRACTION is now documented as the tolerance radius.

[chatgpt-codex-connector]

Avoid scanning affinity cache on every new session

For every brand-new session with multiple routes, this walks the entire session-affinity LRU while holding the synchronous mutex; with the configured limit raised toward LLM_ACCESS_KIRO_SESSION_AFFINITY_MAX_ENTRIES (up to 65,536 entries), a burst of new Kiro conversations can spend request-handler time serially scanning old sessions before route selection. This is avoidable by maintaining per-account counters or a smaller per-key/route snapshot instead of rebuilding counts from all cached entries on each request.

Useful? React with 👍 / 👎.

[acking-you]

Considered, but keeping the point-in-time scan for now — I think the cost is bounded and cold-path:

• It runs only on the first request of a brand-new session: subsequent turns in that conversation hit affinity and take the sticky fast path, which never scans. And it's already skipped entirely for single-account keys (routes.len() > 1 guard).
• At the default cap (4096) that's a tens-of-µs walk under the lock, once per new conversation, with allocation already cut to O(distinct accounts).

A maintained per-account/identity counter would have to stay consistent across LRU's implicit capacity eviction and lazy TTL expiry (entries are only reclaimed on lookup), which is materially more error-prone than rebuilding a fold each time. The scan only becomes a hot-path concern if someone raises LLM_ACCESS_KIRO_SESSION_AFFINITY_MAX_ENTRIES toward 65,536 — at which point I'd revisit with a maintained counter backed by profiling. Flagging this as a conscious trade-off rather than an oversight.