mirror of
https://github.com/router-for-me/CLIProxyAPI.git
synced 2026-02-02 04:20:50 +08:00
feat: add fill-first routing strategy
This commit is contained in:
@@ -60,6 +60,9 @@ type Config struct {
|
||||
// QuotaExceeded defines the behavior when a quota is exceeded.
|
||||
QuotaExceeded QuotaExceeded `yaml:"quota-exceeded" json:"quota-exceeded"`
|
||||
|
||||
// Routing controls credential selection behavior.
|
||||
Routing RoutingConfig `yaml:"routing" json:"routing"`
|
||||
|
||||
// WebsocketAuth enables or disables authentication for the WebSocket API.
|
||||
WebsocketAuth bool `yaml:"ws-auth" json:"ws-auth"`
|
||||
|
||||
@@ -124,6 +127,13 @@ type QuotaExceeded struct {
|
||||
SwitchPreviewModel bool `yaml:"switch-preview-model" json:"switch-preview-model"`
|
||||
}
|
||||
|
||||
// RoutingConfig configures how credentials are selected for requests.
|
||||
type RoutingConfig struct {
|
||||
// Strategy selects the credential selection strategy.
|
||||
// Supported values: "fill-first" (default), "round-robin".
|
||||
Strategy string `yaml:"strategy,omitempty" json:"strategy,omitempty"`
|
||||
}
|
||||
|
||||
// AmpModelMapping defines a model name mapping for Amp CLI requests.
|
||||
// When Amp requests a model that isn't available locally, this mapping
|
||||
// allows routing to an alternative model that IS available.
|
||||
|
||||
@@ -120,7 +120,7 @@ type Manager struct {
|
||||
// NewManager constructs a manager with optional custom selector and hook.
|
||||
func NewManager(store Store, selector Selector, hook Hook) *Manager {
|
||||
if selector == nil {
|
||||
selector = &RoundRobinSelector{}
|
||||
selector = &FillFirstSelector{}
|
||||
}
|
||||
if hook == nil {
|
||||
hook = NoopHook{}
|
||||
|
||||
@@ -20,6 +20,11 @@ type RoundRobinSelector struct {
|
||||
cursors map[string]int
|
||||
}
|
||||
|
||||
// FillFirstSelector selects the first available credential (deterministic ordering).
|
||||
// This "burns" one account before moving to the next, which can help stagger
|
||||
// rolling-window subscription caps (e.g. chat message limits).
|
||||
type FillFirstSelector struct{}
|
||||
|
||||
type blockReason int
|
||||
|
||||
const (
|
||||
@@ -98,20 +103,8 @@ func (e *modelCooldownError) Headers() http.Header {
|
||||
return headers
|
||||
}
|
||||
|
||||
// Pick selects the next available auth for the provider in a round-robin manner.
|
||||
func (s *RoundRobinSelector) Pick(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, auths []*Auth) (*Auth, error) {
|
||||
_ = ctx
|
||||
_ = opts
|
||||
if len(auths) == 0 {
|
||||
return nil, &Error{Code: "auth_not_found", Message: "no auth candidates"}
|
||||
}
|
||||
if s.cursors == nil {
|
||||
s.cursors = make(map[string]int)
|
||||
}
|
||||
available := make([]*Auth, 0, len(auths))
|
||||
now := time.Now()
|
||||
cooldownCount := 0
|
||||
var earliest time.Time
|
||||
func collectAvailable(auths []*Auth, model string, now time.Time) (available []*Auth, cooldownCount int, earliest time.Time) {
|
||||
available = make([]*Auth, 0, len(auths))
|
||||
for i := 0; i < len(auths); i++ {
|
||||
candidate := auths[i]
|
||||
blocked, reason, next := isAuthBlockedForModel(candidate, model, now)
|
||||
@@ -126,6 +119,18 @@ func (s *RoundRobinSelector) Pick(ctx context.Context, provider, model string, o
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(available) > 1 {
|
||||
sort.Slice(available, func(i, j int) bool { return available[i].ID < available[j].ID })
|
||||
}
|
||||
return available, cooldownCount, earliest
|
||||
}
|
||||
|
||||
func getAvailableAuths(auths []*Auth, provider, model string, now time.Time) ([]*Auth, error) {
|
||||
if len(auths) == 0 {
|
||||
return nil, &Error{Code: "auth_not_found", Message: "no auth candidates"}
|
||||
}
|
||||
|
||||
available, cooldownCount, earliest := collectAvailable(auths, model, now)
|
||||
if len(available) == 0 {
|
||||
if cooldownCount == len(auths) && !earliest.IsZero() {
|
||||
resetIn := earliest.Sub(now)
|
||||
@@ -136,9 +141,21 @@ func (s *RoundRobinSelector) Pick(ctx context.Context, provider, model string, o
|
||||
}
|
||||
return nil, &Error{Code: "auth_unavailable", Message: "no auth available"}
|
||||
}
|
||||
// Make round-robin deterministic even if caller's candidate order is unstable.
|
||||
if len(available) > 1 {
|
||||
sort.Slice(available, func(i, j int) bool { return available[i].ID < available[j].ID })
|
||||
|
||||
return available, nil
|
||||
}
|
||||
|
||||
// Pick selects the next available auth for the provider in a round-robin manner.
|
||||
func (s *RoundRobinSelector) Pick(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, auths []*Auth) (*Auth, error) {
|
||||
_ = ctx
|
||||
_ = opts
|
||||
if s.cursors == nil {
|
||||
s.cursors = make(map[string]int)
|
||||
}
|
||||
now := time.Now()
|
||||
available, err := getAvailableAuths(auths, provider, model, now)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
key := provider + ":" + model
|
||||
s.mu.Lock()
|
||||
@@ -154,6 +171,18 @@ func (s *RoundRobinSelector) Pick(ctx context.Context, provider, model string, o
|
||||
return available[index%len(available)], nil
|
||||
}
|
||||
|
||||
// Pick selects the first available auth for the provider in a deterministic manner.
|
||||
func (s *FillFirstSelector) Pick(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, auths []*Auth) (*Auth, error) {
|
||||
_ = ctx
|
||||
_ = opts
|
||||
now := time.Now()
|
||||
available, err := getAvailableAuths(auths, provider, model, now)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return available[0], nil
|
||||
}
|
||||
|
||||
func isAuthBlockedForModel(auth *Auth, model string, now time.Time) (bool, blockReason, time.Time) {
|
||||
if auth == nil {
|
||||
return true, blockReasonOther, time.Time{}
|
||||
|
||||
@@ -5,6 +5,7 @@ package cliproxy
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/api"
|
||||
sdkaccess "github.com/router-for-me/CLIProxyAPI/v6/sdk/access"
|
||||
@@ -197,7 +198,20 @@ func (b *Builder) Build() (*Service, error) {
|
||||
if dirSetter, ok := tokenStore.(interface{ SetBaseDir(string) }); ok && b.cfg != nil {
|
||||
dirSetter.SetBaseDir(b.cfg.AuthDir)
|
||||
}
|
||||
coreManager = coreauth.NewManager(tokenStore, nil, nil)
|
||||
|
||||
strategy := ""
|
||||
if b.cfg != nil {
|
||||
strategy = strings.ToLower(strings.TrimSpace(b.cfg.Routing.Strategy))
|
||||
}
|
||||
var selector coreauth.Selector
|
||||
switch strategy {
|
||||
case "round-robin", "roundrobin", "rr":
|
||||
selector = &coreauth.RoundRobinSelector{}
|
||||
default:
|
||||
selector = &coreauth.FillFirstSelector{}
|
||||
}
|
||||
|
||||
coreManager = coreauth.NewManager(tokenStore, selector, nil)
|
||||
}
|
||||
// Attach a default RoundTripper provider so providers can opt-in per-auth transports.
|
||||
coreManager.SetRoundTripperProvider(newDefaultRoundTripperProvider())
|
||||
|
||||
Reference in New Issue
Block a user