fix(codex): prioritize websocket-enabled credentials across priority tiers in scheduler logic

This commit is contained in:
Luis Pater
2026-04-01 12:51:12 +08:00
parent ca11b236a7
commit 1734aa1664
2 changed files with 54 additions and 6 deletions

View File

@@ -219,6 +219,19 @@ func (s *authScheduler) pickMixed(ctx context.Context, providers []string, model
if len(normalized) == 0 {
return nil, "", &Error{Code: "provider_not_found", Message: "no provider supplied"}
}
if len(normalized) == 1 {
// When a single provider is eligible, reuse pickSingle so provider-specific preferences
// (for example Codex websocket transport) are applied consistently.
providerKey := normalized[0]
picked, errPick := s.pickSingle(ctx, providerKey, model, opts, tried)
if errPick != nil {
return nil, "", errPick
}
if picked == nil {
return nil, "", &Error{Code: "auth_not_found", Message: "no auth available"}
}
return picked, providerKey, nil
}
pinnedAuthID := pinnedAuthIDFromMetadata(opts.Metadata)
modelKey := canonicalModelKey(model)
@@ -696,16 +709,25 @@ func (m *modelScheduler) highestReadyPriorityLocked(preferWebsocket bool, predic
if m == nil {
return 0, false
}
if preferWebsocket {
// When downstream is websocket and Codex supports websocket transport, prefer websocket-enabled
// credentials even if they are in a lower priority tier than HTTP-only credentials.
for _, priority := range m.priorityOrder {
bucket := m.readyByPriority[priority]
if bucket == nil {
continue
}
if bucket.ws.pickFirst(predicate) != nil {
return priority, true
}
}
}
for _, priority := range m.priorityOrder {
bucket := m.readyByPriority[priority]
if bucket == nil {
continue
}
view := &bucket.all
if preferWebsocket && len(bucket.ws.flat) > 0 {
view = &bucket.ws
}
if view.pickFirst(predicate) != nil {
if bucket.all.pickFirst(predicate) != nil {
return priority, true
}
}
@@ -723,7 +745,7 @@ func (m *modelScheduler) pickReadyAtPriorityLocked(preferWebsocket bool, priorit
return nil
}
view := &bucket.all
if preferWebsocket && len(bucket.ws.flat) > 0 {
if preferWebsocket && bucket.ws.pickFirst(predicate) != nil {
view = &bucket.ws
}
var picked *scheduledAuth

View File

@@ -208,6 +208,32 @@ func TestSchedulerPick_CodexWebsocketPrefersWebsocketEnabledSubset(t *testing.T)
}
}
func TestSchedulerPick_CodexWebsocketPrefersWebsocketEnabledAcrossPriorities(t *testing.T) {
t.Parallel()
scheduler := newSchedulerForTest(
&RoundRobinSelector{},
&Auth{ID: "codex-http", Provider: "codex", Attributes: map[string]string{"priority": "10"}},
&Auth{ID: "codex-ws-a", Provider: "codex", Attributes: map[string]string{"priority": "0", "websockets": "true"}},
&Auth{ID: "codex-ws-b", Provider: "codex", Attributes: map[string]string{"priority": "0", "websockets": "true"}},
)
ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background())
want := []string{"codex-ws-a", "codex-ws-b", "codex-ws-a"}
for index, wantID := range want {
got, errPick := scheduler.pickSingle(ctx, "codex", "", cliproxyexecutor.Options{}, nil)
if errPick != nil {
t.Fatalf("pickSingle() #%d error = %v", index, errPick)
}
if got == nil {
t.Fatalf("pickSingle() #%d auth = nil", index)
}
if got.ID != wantID {
t.Fatalf("pickSingle() #%d auth.ID = %q, want %q", index, got.ID, wantID)
}
}
}
func TestSchedulerPick_MixedProvidersUsesWeightedProviderRotationOverReadyCandidates(t *testing.T) {
t.Parallel()