feat: add fill-first routing strategy

This commit is contained in:
gwizz
2025-12-22 17:40:35 +11:00
parent 27b43ed63f
commit b078be4613
4 changed files with 72 additions and 19 deletions

View File

@@ -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.

View File

@@ -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{}

View File

@@ -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{}

View File

@@ -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())