feat: add weighted round-robin credential selection#3935
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a new weighted-round-robin routing strategy, allowing credentials to be selected based on a configured selection-weight within their priority bucket. It updates configuration structures, API handlers, the TUI, documentation, and plugin APIs to support this new strategy and its associated weights. Feedback on the changes suggests improving consistency and robustness in selection weight parsing: specifically, adding int32 support to authFileStrictIntValue and json.Number support to parseSelectionWeightValue to align with other parsing functions in the codebase.
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.
| func authFileStrictIntValue(value any) (int, bool) { | ||
| switch typed := value.(type) { | ||
| case int: | ||
| return typed, true | ||
| case int64: | ||
| return int(typed), true | ||
| case float64: | ||
| if math.Trunc(typed) != typed { | ||
| return 0, false | ||
| } | ||
| return int(typed), true | ||
| case json.Number: | ||
| if i, err := typed.Int64(); err == nil { | ||
| return int(i), true | ||
| } | ||
| case string: | ||
| if i, err := strconv.Atoi(strings.TrimSpace(typed)); err == nil { | ||
| return i, true | ||
| } | ||
| } | ||
| return 0, false | ||
| } |
There was a problem hiding this comment.
The authFileStrictIntValue function is missing support for int32 types, which are handled in other selection weight parsing functions like parseSelectionWeightValue and metadataSelectionWeight. Adding int32 support ensures consistency and robustness across different integer representations.
func authFileStrictIntValue(value any) (int, bool) {
switch typed := value.(type) {
case int:
return typed, true
case int32:
return int(typed), true
case int64:
return int(typed), true
case float64:
if math.Trunc(typed) != typed {
return 0, false
}
return int(typed), true
case json.Number:
if i, err := typed.Int64(); err == nil {
return int(i), true
}
case string:
if i, err := strconv.Atoi(strings.TrimSpace(typed)); err == nil {
return i, true
}
}
return 0, false
}There was a problem hiding this comment.
Addressed in caf8993b: authFileStrictIntValue now accepts int32, and TestAuthFileSelectionWeightValueProgrammaticIntegerTypes covers the programmatic integer path.
| func parseSelectionWeightValue(raw any) (int, bool) { | ||
| var parsed int | ||
| var ok bool | ||
| switch v := raw.(type) { | ||
| case int: | ||
| parsed, ok = v, true | ||
| case int32: | ||
| parsed, ok = int(v), true | ||
| case int64: | ||
| parsed, ok = int(v), true | ||
| case float64: | ||
| if math.Trunc(v) != v { | ||
| return 0, false | ||
| } | ||
| parsed, ok = int(v), true | ||
| case string: | ||
| value, err := strconv.Atoi(strings.TrimSpace(v)) | ||
| if err == nil { | ||
| parsed, ok = value, true | ||
| } | ||
| } | ||
| if !ok || parsed < 0 { | ||
| return 0, false | ||
| } | ||
| return parsed, true | ||
| } |
There was a problem hiding this comment.
The parseSelectionWeightValue function is missing support for json.Number types, which are handled in other selection weight parsing functions like authFileStrictIntValue and metadataSelectionWeight. If the JSON decoder is configured to use json.Number (e.g., via UseNumber()), parsing will fail. Adding json.Number support ensures consistent and robust parsing.
func parseSelectionWeightValue(raw any) (int, bool) {
var parsed int
var ok bool
switch v := raw.(type) {
case int:
parsed, ok = v, true
case int32:
parsed, ok = int(v), true
case int64:
parsed, ok = int(v), true
case float64:
if math.Trunc(v) != v {
return 0, false
}
parsed, ok = int(v), true
case json.Number:
if i, err := v.Int64(); err == nil {
parsed, ok = int(i), true
}
case string:
value, err := strconv.Atoi(strings.TrimSpace(v))
if err == nil {
parsed, ok = value, true
}
}
if !ok || parsed < 0 {
return 0, false
}
return parsed, true
}There was a problem hiding this comment.
Addressed in caf8993b: parseSelectionWeightValue now accepts json.Number, with TestParseSelectionWeightValueProgrammaticNumberTypes covering both integral and fractional values.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 81af308afc
ℹ️ 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".
| if predicate != nil && !predicate(entry) { | ||
| continue | ||
| } | ||
| if scheduledAuthSelectionWeight(entry) > 0 { |
There was a problem hiding this comment.
When weighted routing is enabled, this check makes the fast-path scheduler report no selectable ready auth if the only matching entry has selection_weight: 0, even when the request predicate has already narrowed selection to a pinned_auth_id. That breaks long-lived Responses/WebSocket requests that pin the previously selected credential: setting weight 0 to drain new traffic will make subsequent pinned messages fail with auth_unavailable instead of continuing on the pinned credential.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Addressed in caf8993b: the scheduler now uses a non-weighted pick strategy when pinned_auth_id is present, so selection_weight: 0 drains new weighted picks without making an otherwise available pinned auth unselectable. Covered by TestManagerWeightedSchedulerPinnedZeroWeightAuth.
| available := availableByPriority[priority] | ||
| hasPositiveWeight := false | ||
| for _, candidate := range available { | ||
| if authSelectionWeight(candidate) > 0 { |
There was a problem hiding this comment.
In the legacy/session-affinity/route-aware path, pickNextLegacy has already filtered candidates down to pinned_auth_id before calling the selector, but this positive-weight gate then drops that drained credential from consideration. A pinned request or websocket continuation that should keep using an otherwise available auth will fail as soon as the operator sets its selection weight to 0, even though weight 0 is meant to drain new selections rather than disable the credential.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Addressed in caf8993b: WeightedRoundRobinSelector now returns an available pinned auth before running weighted selection, preserving continuations for drained credentials. Covered by TestWeightedRoundRobinSelectorPick_PinnedZeroWeightAuth.
bagaking
force-pushed
the
feat/weighted-round-robin-routing
branch
from
81af308 to
caf8993
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: caf8993b78
ℹ️ 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".
| if err != nil { | ||
| return nil, err | ||
| } | ||
| available = preferCodexWebsocketAuths(ctx, provider, available) |
There was a problem hiding this comment.
When the legacy weighted selector is used for a Codex websocket request (for example after route-aware model aliases or through a session-affinity fallback), this filters the already-selected priority bucket down to websocket-enabled auths without checking their selection_weight. If the only websocket-enabled auth in that bucket has weight 0 to drain it, while another same-priority auth has positive weight, pickWeightedAuth sees total weight 0 and returns auth_unavailable instead of using the positive-weight credential; the scheduler path avoids this by checking for a selectable websocket entry before switching views.
Useful? React with 👍 / 👎.
1 participant
Summary
Adds a
weighted-round-robinrouting strategy that distributes requests by credential selection weight while preserving existing priority semantics. Higherprioritycredentials are still considered first; within the selected priority bucket,selection-weightcontrols the proportional share. A weight of0drains the credential from new weighted selections and allows fallback to lower-priority ready credentials when needed.This wires
selection-weightthrough config parsing, auth files, management APIs, TUI, plugin APIs, scheduler/conductor behavior, config diff, watcher synthesis, and examples so the behavior can be configured and inspected consistently.Details
selection-weightfor Gemini, Claude, Codex, Vertex, OpenAI-compatible providers, and OpenAI-compatible per-key entries.selection_weight/selection-weightin auth files, including strict non-negative integer validation and null clearing through the management fields API.1; missing or invalid priority resolves to0.Related
Validation
go test ./internal/watcher/synthesizer -run 'TestFileSynthesizer_Synthesize_SelectionWeightParsing|TestMetadataSelectionWeight_ProgrammaticIntegerTypes' -count=1go test ./internal/api/handlers/management -run 'TestPatchAuthFileFields_SelectionWeight|TestBuildAuthFromFileDataSelectionWeight|Test(PatchOpenAICompatRejectsNegativeAPIKeyEntrySelectionWeight|PutConfigListsRejectNegativeSelectionWeight)' -count=1go test ./internal/watcher/diff ./internal/watcher/synthesizer ./internal/pluginhost ./sdk/cliproxy/auth -count=1go test -count=1inexamples/plugin/scheduler/gogo build -o test-output ./cmd/server && rm test-outputgit diff --check