fix(codex): prioritize websocket-enabled credentials across priority tiers in scheduler logic
This commit is contained in:
@@ -219,6 +219,19 @@ func (s *authScheduler) pickMixed(ctx context.Context, providers []string, model
|
|||||||
if len(normalized) == 0 {
|
if len(normalized) == 0 {
|
||||||
return nil, "", &Error{Code: "provider_not_found", Message: "no provider supplied"}
|
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)
|
pinnedAuthID := pinnedAuthIDFromMetadata(opts.Metadata)
|
||||||
modelKey := canonicalModelKey(model)
|
modelKey := canonicalModelKey(model)
|
||||||
|
|
||||||
@@ -696,16 +709,25 @@ func (m *modelScheduler) highestReadyPriorityLocked(preferWebsocket bool, predic
|
|||||||
if m == nil {
|
if m == nil {
|
||||||
return 0, false
|
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 {
|
for _, priority := range m.priorityOrder {
|
||||||
bucket := m.readyByPriority[priority]
|
bucket := m.readyByPriority[priority]
|
||||||
if bucket == nil {
|
if bucket == nil {
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
view := &bucket.all
|
if bucket.all.pickFirst(predicate) != nil {
|
||||||
if preferWebsocket && len(bucket.ws.flat) > 0 {
|
|
||||||
view = &bucket.ws
|
|
||||||
}
|
|
||||||
if view.pickFirst(predicate) != nil {
|
|
||||||
return priority, true
|
return priority, true
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -723,7 +745,7 @@ func (m *modelScheduler) pickReadyAtPriorityLocked(preferWebsocket bool, priorit
|
|||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
view := &bucket.all
|
view := &bucket.all
|
||||||
if preferWebsocket && len(bucket.ws.flat) > 0 {
|
if preferWebsocket && bucket.ws.pickFirst(predicate) != nil {
|
||||||
view = &bucket.ws
|
view = &bucket.ws
|
||||||
}
|
}
|
||||||
var picked *scheduledAuth
|
var picked *scheduledAuth
|
||||||
|
|||||||
@@ -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) {
|
func TestSchedulerPick_MixedProvidersUsesWeightedProviderRotationOverReadyCandidates(t *testing.T) {
|
||||||
t.Parallel()
|
t.Parallel()
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user