feat(selector): add priority support for auth selection

This commit is contained in:
Luis Pater
2026-01-15 07:08:24 +08:00
parent b163f8ed9e
commit 6f8a8f8136
6 changed files with 381 additions and 24 deletions

View File

@@ -8,6 +8,7 @@ import (
"net/http"
"sort"
"strconv"
"strings"
"sync"
"time"
@@ -103,13 +104,29 @@ func (e *modelCooldownError) Headers() http.Header {
return headers
}
func collectAvailable(auths []*Auth, model string, now time.Time) (available []*Auth, cooldownCount int, earliest time.Time) {
available = make([]*Auth, 0, len(auths))
func authPriority(auth *Auth) int {
if auth == nil || auth.Attributes == nil {
return 0
}
raw := strings.TrimSpace(auth.Attributes["priority"])
if raw == "" {
return 0
}
parsed, err := strconv.Atoi(raw)
if err != nil {
return 0
}
return parsed
}
func collectAvailableByPriority(auths []*Auth, model string, now time.Time) (available map[int][]*Auth, cooldownCount int, earliest time.Time) {
available = make(map[int][]*Auth)
for i := 0; i < len(auths); i++ {
candidate := auths[i]
blocked, reason, next := isAuthBlockedForModel(candidate, model, now)
if !blocked {
available = append(available, candidate)
priority := authPriority(candidate)
available[priority] = append(available[priority], candidate)
continue
}
if reason == blockReasonCooldown {
@@ -119,9 +136,6 @@ func collectAvailable(auths []*Auth, model string, now time.Time) (available []*
}
}
}
if len(available) > 1 {
sort.Slice(available, func(i, j int) bool { return available[i].ID < available[j].ID })
}
return available, cooldownCount, earliest
}
@@ -130,18 +144,35 @@ func getAvailableAuths(auths []*Auth, provider, model string, now time.Time) ([]
return nil, &Error{Code: "auth_not_found", Message: "no auth candidates"}
}
available, cooldownCount, earliest := collectAvailable(auths, model, now)
if len(available) == 0 {
availableByPriority, cooldownCount, earliest := collectAvailableByPriority(auths, model, now)
if len(availableByPriority) == 0 {
if cooldownCount == len(auths) && !earliest.IsZero() {
providerForError := provider
if providerForError == "mixed" {
providerForError = ""
}
resetIn := earliest.Sub(now)
if resetIn < 0 {
resetIn = 0
}
return nil, newModelCooldownError(model, provider, resetIn)
return nil, newModelCooldownError(model, providerForError, resetIn)
}
return nil, &Error{Code: "auth_unavailable", Message: "no auth available"}
}
bestPriority := 0
found := false
for priority := range availableByPriority {
if !found || priority > bestPriority {
bestPriority = priority
found = true
}
}
available := availableByPriority[bestPriority]
if len(available) > 1 {
sort.Slice(available, func(i, j int) bool { return available[i].ID < available[j].ID })
}
return available, nil
}