diff --git a/sdk/cliproxy/auth/conductor_selection.go b/sdk/cliproxy/auth/conductor_selection.go index 43c76397..bcf597bf 100644 --- a/sdk/cliproxy/auth/conductor_selection.go +++ b/sdk/cliproxy/auth/conductor_selection.go @@ -368,6 +368,17 @@ func selectionArgForSelector(selector Selector, routeModel string) string { return routeModel } +func restoreModelCooldownErrorModel(err error, requestedModel string) error { + if err == nil || requestedModel == "" { + return err + } + var cooldownErr *modelCooldownError + if !errors.As(err, &cooldownErr) || cooldownErr == nil || cooldownErr.model != "" { + return err + } + return newModelCooldownError(requestedModel, cooldownErr.provider, cooldownErr.resetIn) +} + func schedulerAttributeSensitive(key string) bool { key = strings.ToLower(strings.TrimSpace(key)) normalized := strings.NewReplacer("-", "_", ".", "_", " ", "_").Replace(key) @@ -1046,6 +1057,9 @@ func (m *Manager) pickNextLegacy(ctx context.Context, provider, model string, op selectorCtx := withWeightedSelectorStateModel(ctx, selector, model) selected, errPick = selector.Pick(selectorCtx, provider, selectionArgForSelector(selector, model), opts, selectorAuths) if errPick != nil { + if isBuiltInSelector(selector) { + errPick = restoreModelCooldownErrorModel(errPick, model) + } return nil, nil, errPick } } @@ -1363,6 +1377,9 @@ func (m *Manager) pickNextMixedLegacy(ctx context.Context, providers []string, m selectorCtx := withWeightedSelectorStateModel(ctx, selector, model) selected, errPick = selector.Pick(selectorCtx, "mixed", selectionArgForSelector(selector, model), opts, selectorAuths) if errPick != nil { + if isBuiltInSelector(selector) { + errPick = restoreModelCooldownErrorModel(errPick, model) + } return nil, nil, "", errPick } } diff --git a/sdk/cliproxy/auth/conductor_selection_cooldown_test.go b/sdk/cliproxy/auth/conductor_selection_cooldown_test.go new file mode 100644 index 00000000..1403347b --- /dev/null +++ b/sdk/cliproxy/auth/conductor_selection_cooldown_test.go @@ -0,0 +1,60 @@ +package auth + +import ( + "context" + "errors" + "testing" + "time" + + cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" +) + +func TestBuiltInSelectorCooldownErrorPreservesRouteModel(t *testing.T) { + t.Parallel() + + const routeModel = "client-opus(high)" + next := time.Now().Add(time.Hour) + auth := &Auth{ + ID: "cooling-auth", + Unavailable: true, + NextRetryAfter: next, + Quota: QuotaState{ + Exceeded: true, + NextRecoverAt: next, + }, + ModelStates: map[string]*ModelState{ + "other-model": {Status: StatusActive}, + }, + } + + selectors := map[string]Selector{ + "round-robin": &RoundRobinSelector{}, + "weighted-round-robin": &WeightedRoundRobinSelector{}, + "fill-first": &FillFirstSelector{}, + } + for name, selector := range selectors { + t.Run(name, func(t *testing.T) { + t.Parallel() + + _, errPick := selector.Pick( + context.Background(), + "mixed", + selectionArgForSelector(selector, routeModel), + cliproxyexecutor.Options{}, + []*Auth{auth}, + ) + if errPick == nil { + t.Fatal("Pick() error = nil, want model cooldown") + } + + errPick = restoreModelCooldownErrorModel(errPick, routeModel) + var cooldownErr *modelCooldownError + if !errors.As(errPick, &cooldownErr) { + t.Fatalf("Pick() error = %T, want *modelCooldownError", errPick) + } + if cooldownErr.model != routeModel { + t.Fatalf("cooldown model = %q, want %q", cooldownErr.model, routeModel) + } + }) + } +}