mirror of
https://github.com/router-for-me/CLIProxyAPIPlus.git
synced 2026-03-11 08:15:14 +00:00
Compare commits
22 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
a862984dca | ||
|
|
f0365f0465 | ||
|
|
6d1e20e940 | ||
|
|
e52b542e22 | ||
|
|
8f6abb8a86 | ||
|
|
ed8eaae964 | ||
|
|
e8de87ee90 | ||
|
|
4e572ec8b9 | ||
|
|
6c7f18c448 | ||
|
|
24bc9cba67 | ||
|
|
97356b1a04 | ||
|
|
1084b53fba | ||
|
|
b1aecc2bf1 | ||
|
|
83b90e106f | ||
|
|
f52114dab2 | ||
|
|
5106caf641 | ||
|
|
12370ee84e | ||
|
|
b84ccc6e7a | ||
|
|
e19ddb53e7 | ||
|
|
2a0100b2d6 | ||
|
|
c020fa60d0 | ||
|
|
b078be4613 |
BIN
assets/packycode.png
Normal file
BIN
assets/packycode.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 8.1 KiB |
@@ -71,6 +71,10 @@ quota-exceeded:
|
||||
switch-project: true # Whether to automatically switch to another project when a quota is exceeded
|
||||
switch-preview-model: true # Whether to automatically switch to a preview model when a quota is exceeded
|
||||
|
||||
# Routing strategy for selecting credentials when multiple match.
|
||||
routing:
|
||||
strategy: "round-robin" # round-robin (default), fill-first
|
||||
|
||||
# When true, enable authentication for the WebSocket API (/v1/ws).
|
||||
ws-auth: false
|
||||
|
||||
|
||||
@@ -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"`
|
||||
|
||||
@@ -136,6 +139,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: "round-robin" (default), "fill-first".
|
||||
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.
|
||||
|
||||
@@ -33,18 +33,18 @@ import (
|
||||
)
|
||||
|
||||
const (
|
||||
antigravityBaseURLDaily = "https://daily-cloudcode-pa.sandbox.googleapis.com"
|
||||
// antigravityBaseURLAutopush = "https://autopush-cloudcode-pa.sandbox.googleapis.com"
|
||||
antigravityBaseURLProd = "https://cloudcode-pa.googleapis.com"
|
||||
antigravityCountTokensPath = "/v1internal:countTokens"
|
||||
antigravityStreamPath = "/v1internal:streamGenerateContent"
|
||||
antigravityGeneratePath = "/v1internal:generateContent"
|
||||
antigravityModelsPath = "/v1internal:fetchAvailableModels"
|
||||
antigravityClientID = "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com"
|
||||
antigravityClientSecret = "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"
|
||||
defaultAntigravityAgent = "antigravity/1.11.5 windows/amd64"
|
||||
antigravityAuthType = "antigravity"
|
||||
refreshSkew = 3000 * time.Second
|
||||
antigravityBaseURLDaily = "https://daily-cloudcode-pa.googleapis.com"
|
||||
antigravitySandboxBaseURLDaily = "https://daily-cloudcode-pa.sandbox.googleapis.com"
|
||||
antigravityBaseURLProd = "https://cloudcode-pa.googleapis.com"
|
||||
antigravityCountTokensPath = "/v1internal:countTokens"
|
||||
antigravityStreamPath = "/v1internal:streamGenerateContent"
|
||||
antigravityGeneratePath = "/v1internal:generateContent"
|
||||
antigravityModelsPath = "/v1internal:fetchAvailableModels"
|
||||
antigravityClientID = "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com"
|
||||
antigravityClientSecret = "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"
|
||||
defaultAntigravityAgent = "antigravity/1.11.5 windows/amd64"
|
||||
antigravityAuthType = "antigravity"
|
||||
refreshSkew = 3000 * time.Second
|
||||
)
|
||||
|
||||
var (
|
||||
@@ -1156,7 +1156,7 @@ func antigravityBaseURLFallbackOrder(auth *cliproxyauth.Auth) []string {
|
||||
}
|
||||
return []string{
|
||||
antigravityBaseURLDaily,
|
||||
// antigravityBaseURLAutopush,
|
||||
antigravitySandboxBaseURLDaily,
|
||||
antigravityBaseURLProd,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -662,7 +662,14 @@ func decodeResponseBody(body io.ReadCloser, contentEncoding string) (io.ReadClos
|
||||
}
|
||||
|
||||
func applyClaudeHeaders(r *http.Request, auth *cliproxyauth.Auth, apiKey string, stream bool, extraBetas []string) {
|
||||
r.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
useAPIKey := auth != nil && auth.Attributes != nil && strings.TrimSpace(auth.Attributes["api_key"]) != ""
|
||||
isAnthropicBase := r.URL != nil && strings.EqualFold(r.URL.Scheme, "https") && strings.EqualFold(r.URL.Host, "api.anthropic.com")
|
||||
if isAnthropicBase && useAPIKey {
|
||||
r.Header.Del("Authorization")
|
||||
r.Header.Set("x-api-key", apiKey)
|
||||
} else {
|
||||
r.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
}
|
||||
r.Header.Set("Content-Type", "application/json")
|
||||
|
||||
var ginHeaders http.Header
|
||||
|
||||
@@ -86,6 +86,10 @@ func ConvertClaudeRequestToAntigravity(modelName string, inputRawJSON []byte, _
|
||||
hasSystemInstruction = true
|
||||
}
|
||||
}
|
||||
} else if systemResult.Type == gjson.String {
|
||||
systemInstructionJSON = `{"role":"user","parts":[{"text":""}]}`
|
||||
systemInstructionJSON, _ = sjson.Set(systemInstructionJSON, "parts.0.text", systemResult.String())
|
||||
hasSystemInstruction = true
|
||||
}
|
||||
|
||||
// contents
|
||||
@@ -121,32 +125,31 @@ func ConvertClaudeRequestToAntigravity(modelName string, inputRawJSON []byte, _
|
||||
// Use GetThinkingText to handle wrapped thinking objects
|
||||
thinkingText := util.GetThinkingText(contentResult)
|
||||
signatureResult := contentResult.Get("signature")
|
||||
clientSignature := ""
|
||||
if signatureResult.Exists() && signatureResult.String() != "" {
|
||||
clientSignature = signatureResult.String()
|
||||
}
|
||||
|
||||
// Always try cached signature first (more reliable than client-provided)
|
||||
// Client may send stale or invalid signatures from different sessions
|
||||
signature := ""
|
||||
if sessionID != "" && thinkingText != "" {
|
||||
if cachedSig := cache.GetCachedSignature(sessionID, thinkingText); cachedSig != "" {
|
||||
signature = cachedSig
|
||||
log.Debugf("Using cached signature for thinking block")
|
||||
clientSignature := ""
|
||||
if signatureResult.Exists() && signatureResult.String() != "" {
|
||||
clientSignature = signatureResult.String()
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback to client signature only if cache miss and client signature is valid
|
||||
if signature == "" && cache.HasValidSignature(clientSignature) {
|
||||
signature = clientSignature
|
||||
log.Debugf("Using client-provided signature for thinking block")
|
||||
}
|
||||
// Always try cached signature first (more reliable than client-provided)
|
||||
// Client may send stale or invalid signatures from different sessions
|
||||
signature := ""
|
||||
if sessionID != "" && thinkingText != "" {
|
||||
if cachedSig := cache.GetCachedSignature(sessionID, thinkingText); cachedSig != "" {
|
||||
signature = cachedSig
|
||||
log.Debugf("Using cached signature for thinking block")
|
||||
}
|
||||
}
|
||||
|
||||
// Store for subsequent tool_use in the same message
|
||||
if cache.HasValidSignature(signature) {
|
||||
currentMessageThinkingSignature = signature
|
||||
}
|
||||
// Fallback to client signature only if cache miss and client signature is valid
|
||||
if signature == "" && cache.HasValidSignature(clientSignature) {
|
||||
signature = clientSignature
|
||||
log.Debugf("Using client-provided signature for thinking block")
|
||||
}
|
||||
|
||||
// Store for subsequent tool_use in the same message
|
||||
if cache.HasValidSignature(signature) {
|
||||
currentMessageThinkingSignature = signature
|
||||
}
|
||||
|
||||
// Skip trailing unsigned thinking blocks on last assistant message
|
||||
isUnsigned := !cache.HasValidSignature(signature)
|
||||
@@ -321,6 +324,7 @@ func ConvertClaudeRequestToAntigravity(modelName string, inputRawJSON []byte, _
|
||||
// tools
|
||||
toolsJSON := ""
|
||||
toolDeclCount := 0
|
||||
allowedToolKeys := []string{"name", "description", "behavior", "parameters", "parametersJsonSchema", "response", "responseJsonSchema"}
|
||||
toolsResult := gjson.GetBytes(rawJSON, "tools")
|
||||
if toolsResult.IsArray() {
|
||||
toolsJSON = `[{"functionDeclarations":[]}]`
|
||||
@@ -333,10 +337,12 @@ func ConvertClaudeRequestToAntigravity(modelName string, inputRawJSON []byte, _
|
||||
inputSchema := util.CleanJSONSchemaForAntigravity(inputSchemaResult.Raw)
|
||||
tool, _ := sjson.Delete(toolResult.Raw, "input_schema")
|
||||
tool, _ = sjson.SetRaw(tool, "parametersJsonSchema", inputSchema)
|
||||
tool, _ = sjson.Delete(tool, "strict")
|
||||
tool, _ = sjson.Delete(tool, "input_examples")
|
||||
tool, _ = sjson.Delete(tool, "type")
|
||||
tool, _ = sjson.Delete(tool, "cache_control")
|
||||
for toolKey := range gjson.Parse(tool).Map() {
|
||||
if util.InArray(allowedToolKeys, toolKey) {
|
||||
continue
|
||||
}
|
||||
tool, _ = sjson.Delete(tool, toolKey)
|
||||
}
|
||||
toolsJSON, _ = sjson.SetRaw(toolsJSON, "0.functionDeclarations.-1", tool)
|
||||
toolDeclCount++
|
||||
}
|
||||
|
||||
@@ -192,6 +192,14 @@ func ConvertOpenAIRequestToAntigravity(modelName string, inputRawJSON []byte, _
|
||||
} else if content.IsObject() && content.Get("type").String() == "text" {
|
||||
out, _ = sjson.SetBytes(out, "request.systemInstruction.role", "user")
|
||||
out, _ = sjson.SetBytes(out, "request.systemInstruction.parts.0.text", content.Get("text").String())
|
||||
} else if content.IsArray() {
|
||||
contents := content.Array()
|
||||
if len(contents) > 0 {
|
||||
out, _ = sjson.SetBytes(out, "request.systemInstruction.role", "user")
|
||||
for j := 0; j < len(contents); j++ {
|
||||
out, _ = sjson.SetBytes(out, fmt.Sprintf("request.systemInstruction.parts.%d.text", j), contents[j].Get("text").String())
|
||||
}
|
||||
}
|
||||
}
|
||||
} else if role == "user" || (role == "system" && len(arr) == 1) {
|
||||
// Build single user content node to avoid splitting into multiple contents
|
||||
@@ -258,7 +266,11 @@ func ConvertOpenAIRequestToAntigravity(modelName string, inputRawJSON []byte, _
|
||||
fargs := tc.Get("function.arguments").String()
|
||||
node, _ = sjson.SetBytes(node, "parts."+itoa(p)+".functionCall.id", fid)
|
||||
node, _ = sjson.SetBytes(node, "parts."+itoa(p)+".functionCall.name", fname)
|
||||
node, _ = sjson.SetRawBytes(node, "parts."+itoa(p)+".functionCall.args", []byte(fargs))
|
||||
if gjson.Valid(fargs) {
|
||||
node, _ = sjson.SetRawBytes(node, "parts."+itoa(p)+".functionCall.args", []byte(fargs))
|
||||
} else {
|
||||
node, _ = sjson.SetBytes(node, "parts."+itoa(p)+".functionCall.args.params", []byte(fargs))
|
||||
}
|
||||
node, _ = sjson.SetBytes(node, "parts."+itoa(p)+".thoughtSignature", geminiCLIFunctionThoughtSignature)
|
||||
p++
|
||||
if fid != "" {
|
||||
|
||||
@@ -62,6 +62,8 @@ func ConvertClaudeRequestToCLI(modelName string, inputRawJSON []byte, _ bool) []
|
||||
if hasSystemParts {
|
||||
out, _ = sjson.SetRaw(out, "request.systemInstruction", systemInstruction)
|
||||
}
|
||||
} else if systemResult.Type == gjson.String {
|
||||
out, _ = sjson.Set(out, "request.systemInstruction.parts.-1.text", systemResult.String())
|
||||
}
|
||||
|
||||
// contents
|
||||
|
||||
@@ -160,6 +160,14 @@ func ConvertOpenAIRequestToGeminiCLI(modelName string, inputRawJSON []byte, _ bo
|
||||
} else if content.IsObject() && content.Get("type").String() == "text" {
|
||||
out, _ = sjson.SetBytes(out, "request.systemInstruction.role", "user")
|
||||
out, _ = sjson.SetBytes(out, "request.systemInstruction.parts.0.text", content.Get("text").String())
|
||||
} else if content.IsArray() {
|
||||
contents := content.Array()
|
||||
if len(contents) > 0 {
|
||||
out, _ = sjson.SetBytes(out, "request.systemInstruction.role", "user")
|
||||
for j := 0; j < len(contents); j++ {
|
||||
out, _ = sjson.SetBytes(out, fmt.Sprintf("request.systemInstruction.parts.%d.text", j), contents[j].Get("text").String())
|
||||
}
|
||||
}
|
||||
}
|
||||
} else if role == "user" || (role == "system" && len(arr) == 1) {
|
||||
// Build single user content node to avoid splitting into multiple contents
|
||||
|
||||
@@ -55,6 +55,8 @@ func ConvertClaudeRequestToGemini(modelName string, inputRawJSON []byte, _ bool)
|
||||
if hasSystemParts {
|
||||
out, _ = sjson.SetRaw(out, "system_instruction", systemInstruction)
|
||||
}
|
||||
} else if systemResult.Type == gjson.String {
|
||||
out, _ = sjson.Set(out, "request.system_instruction.parts.-1.text", systemResult.String())
|
||||
}
|
||||
|
||||
// contents
|
||||
|
||||
@@ -178,6 +178,14 @@ func ConvertOpenAIRequestToGemini(modelName string, inputRawJSON []byte, _ bool)
|
||||
} else if content.IsObject() && content.Get("type").String() == "text" {
|
||||
out, _ = sjson.SetBytes(out, "system_instruction.role", "user")
|
||||
out, _ = sjson.SetBytes(out, "system_instruction.parts.0.text", content.Get("text").String())
|
||||
} else if content.IsArray() {
|
||||
contents := content.Array()
|
||||
if len(contents) > 0 {
|
||||
out, _ = sjson.SetBytes(out, "request.systemInstruction.role", "user")
|
||||
for j := 0; j < len(contents); j++ {
|
||||
out, _ = sjson.SetBytes(out, fmt.Sprintf("request.systemInstruction.parts.%d.text", j), contents[j].Get("text").String())
|
||||
}
|
||||
}
|
||||
}
|
||||
} else if role == "user" || (role == "system" && len(arr) == 1) {
|
||||
// Build single user content node to avoid splitting into multiple contents
|
||||
|
||||
@@ -135,6 +135,18 @@ func NewManager(store Store, selector Selector, hook Hook) *Manager {
|
||||
}
|
||||
}
|
||||
|
||||
func (m *Manager) SetSelector(selector Selector) {
|
||||
if m == nil {
|
||||
return
|
||||
}
|
||||
if selector == nil {
|
||||
selector = &RoundRobinSelector{}
|
||||
}
|
||||
m.mu.Lock()
|
||||
m.selector = selector
|
||||
m.mu.Unlock()
|
||||
}
|
||||
|
||||
// SetStore swaps the underlying persistence store.
|
||||
func (m *Manager) SetStore(store Store) {
|
||||
m.mu.Lock()
|
||||
|
||||
@@ -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,12 +141,24 @@ 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
|
||||
now := time.Now()
|
||||
available, err := getAvailableAuths(auths, provider, model, now)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
key := provider + ":" + model
|
||||
s.mu.Lock()
|
||||
if s.cursors == nil {
|
||||
s.cursors = make(map[string]int)
|
||||
}
|
||||
index := s.cursors[key]
|
||||
|
||||
if index >= 2_147_483_640 {
|
||||
@@ -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{}
|
||||
|
||||
113
sdk/cliproxy/auth/selector_test.go
Normal file
113
sdk/cliproxy/auth/selector_test.go
Normal file
@@ -0,0 +1,113 @@
|
||||
package auth
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor"
|
||||
)
|
||||
|
||||
func TestFillFirstSelectorPick_Deterministic(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
selector := &FillFirstSelector{}
|
||||
auths := []*Auth{
|
||||
{ID: "b"},
|
||||
{ID: "a"},
|
||||
{ID: "c"},
|
||||
}
|
||||
|
||||
got, err := selector.Pick(context.Background(), "gemini", "", cliproxyexecutor.Options{}, auths)
|
||||
if err != nil {
|
||||
t.Fatalf("Pick() error = %v", err)
|
||||
}
|
||||
if got == nil {
|
||||
t.Fatalf("Pick() auth = nil")
|
||||
}
|
||||
if got.ID != "a" {
|
||||
t.Fatalf("Pick() auth.ID = %q, want %q", got.ID, "a")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRoundRobinSelectorPick_CyclesDeterministic(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
selector := &RoundRobinSelector{}
|
||||
auths := []*Auth{
|
||||
{ID: "b"},
|
||||
{ID: "a"},
|
||||
{ID: "c"},
|
||||
}
|
||||
|
||||
want := []string{"a", "b", "c", "a", "b"}
|
||||
for i, id := range want {
|
||||
got, err := selector.Pick(context.Background(), "gemini", "", cliproxyexecutor.Options{}, auths)
|
||||
if err != nil {
|
||||
t.Fatalf("Pick() #%d error = %v", i, err)
|
||||
}
|
||||
if got == nil {
|
||||
t.Fatalf("Pick() #%d auth = nil", i)
|
||||
}
|
||||
if got.ID != id {
|
||||
t.Fatalf("Pick() #%d auth.ID = %q, want %q", i, got.ID, id)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRoundRobinSelectorPick_Concurrent(t *testing.T) {
|
||||
selector := &RoundRobinSelector{}
|
||||
auths := []*Auth{
|
||||
{ID: "b"},
|
||||
{ID: "a"},
|
||||
{ID: "c"},
|
||||
}
|
||||
|
||||
start := make(chan struct{})
|
||||
var wg sync.WaitGroup
|
||||
errCh := make(chan error, 1)
|
||||
|
||||
goroutines := 32
|
||||
iterations := 100
|
||||
for i := 0; i < goroutines; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
<-start
|
||||
for j := 0; j < iterations; j++ {
|
||||
got, err := selector.Pick(context.Background(), "gemini", "", cliproxyexecutor.Options{}, auths)
|
||||
if err != nil {
|
||||
select {
|
||||
case errCh <- err:
|
||||
default:
|
||||
}
|
||||
return
|
||||
}
|
||||
if got == nil {
|
||||
select {
|
||||
case errCh <- errors.New("Pick() returned nil auth"):
|
||||
default:
|
||||
}
|
||||
return
|
||||
}
|
||||
if got.ID == "" {
|
||||
select {
|
||||
case errCh <- errors.New("Pick() returned auth with empty ID"):
|
||||
default:
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
close(start)
|
||||
wg.Wait()
|
||||
|
||||
select {
|
||||
case err := <-errCh:
|
||||
t.Fatalf("concurrent Pick() error = %v", err)
|
||||
default:
|
||||
}
|
||||
}
|
||||
@@ -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 "fill-first", "fillfirst", "ff":
|
||||
selector = &coreauth.FillFirstSelector{}
|
||||
default:
|
||||
selector = &coreauth.RoundRobinSelector{}
|
||||
}
|
||||
|
||||
coreManager = coreauth.NewManager(tokenStore, selector, nil)
|
||||
}
|
||||
// Attach a default RoundTripper provider so providers can opt-in per-auth transports.
|
||||
coreManager.SetRoundTripperProvider(newDefaultRoundTripperProvider())
|
||||
|
||||
@@ -510,6 +510,13 @@ func (s *Service) Run(ctx context.Context) error {
|
||||
|
||||
var watcherWrapper *WatcherWrapper
|
||||
reloadCallback := func(newCfg *config.Config) {
|
||||
previousStrategy := ""
|
||||
s.cfgMu.RLock()
|
||||
if s.cfg != nil {
|
||||
previousStrategy = strings.ToLower(strings.TrimSpace(s.cfg.Routing.Strategy))
|
||||
}
|
||||
s.cfgMu.RUnlock()
|
||||
|
||||
if newCfg == nil {
|
||||
s.cfgMu.RLock()
|
||||
newCfg = s.cfg
|
||||
@@ -518,6 +525,30 @@ func (s *Service) Run(ctx context.Context) error {
|
||||
if newCfg == nil {
|
||||
return
|
||||
}
|
||||
|
||||
nextStrategy := strings.ToLower(strings.TrimSpace(newCfg.Routing.Strategy))
|
||||
normalizeStrategy := func(strategy string) string {
|
||||
switch strategy {
|
||||
case "fill-first", "fillfirst", "ff":
|
||||
return "fill-first"
|
||||
default:
|
||||
return "round-robin"
|
||||
}
|
||||
}
|
||||
previousStrategy = normalizeStrategy(previousStrategy)
|
||||
nextStrategy = normalizeStrategy(nextStrategy)
|
||||
if s.coreManager != nil && previousStrategy != nextStrategy {
|
||||
var selector coreauth.Selector
|
||||
switch nextStrategy {
|
||||
case "fill-first":
|
||||
selector = &coreauth.FillFirstSelector{}
|
||||
default:
|
||||
selector = &coreauth.RoundRobinSelector{}
|
||||
}
|
||||
s.coreManager.SetSelector(selector)
|
||||
log.Infof("routing strategy updated to %s", nextStrategy)
|
||||
}
|
||||
|
||||
s.applyRetryConfig(newCfg)
|
||||
if s.server != nil {
|
||||
s.server.UpdateClients(newCfg)
|
||||
|
||||
Reference in New Issue
Block a user