Something went wrong. Try again.
sloproxy
Something went wrong. Try again.
31 kB · 794 lines
Go
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795package executor
import ( "context" "errors" "io" "net/http" "net/http/httptest" "strconv" "strings" "sync/atomic" "testing" "time"
"github.com/google/uuid" "github.com/router-for-me/CLIProxyAPI/v7/internal/config" "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator")
type retryAfterProvider interface { RetryAfter() *time.Duration}
func TestClaudeExecutor_HonorsAnthropicRateLimitHeaders_Execute(t *testing.T) { now := time.Now() sevenDayReset := now.Add(7 * 24 * time.Hour).Unix() fiveHourReset := now.Add(5 * time.Hour).Unix()
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") w.Header().Set("Anthropic-Ratelimit-Unified-Status", "rejected") w.Header().Set("Anthropic-Ratelimit-Unified-5h-Status", "allowed") w.Header().Set("Anthropic-Ratelimit-Unified-5h-Reset", strconv.FormatInt(fiveHourReset, 10)) w.Header().Set("Anthropic-Ratelimit-Unified-7d-Status", "rejected") w.Header().Set("Anthropic-Ratelimit-Unified-7d-Reset", strconv.FormatInt(sevenDayReset, 10)) w.Header().Set("Anthropic-Ratelimit-Unified-Representative-Claim", "seven_day") w.Header().Set("Anthropic-Ratelimit-Unified-Reset", strconv.FormatInt(sevenDayReset, 10)) w.Header().Set("Retry-After", strconv.FormatInt(7*24*3600, 10)) w.WriteHeader(http.StatusTooManyRequests) _, _ = w.Write([]byte(`{"type":"error","error":{"type":"rate_limit_error","message":"Number of requests has exceeded your 7-day rate limit."}}`)) })) defer server.Close()
executor := NewClaudeExecutor(&config.Config{}) auth := &cliproxyauth.Auth{ ID: "claude-auth-1", Provider: "claude", Attributes: map[string]string{ "api_key": "test-key", "base_url": server.URL, }, }
payload := []byte(`{"messages":[{"role":"user","content":[{"type":"text","text":"hi"}]}]}`) _, err := executor.Execute(context.Background(), auth, cliproxyexecutor.Request{ Model: "claude-3-5-sonnet-20241022", Payload: payload, }, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude}) if err == nil { t.Fatal("expected error from Execute, got nil") }
var rap retryAfterProvider if !errors.As(err, &rap) || rap == nil { t.Fatalf("expected error %T to implement RetryAfter() *time.Duration", err) }
retryAfter := rap.RetryAfter() if retryAfter == nil { t.Fatalf("expected non-nil RetryAfter, got nil") }
// Should be at least 7 days (reported reset) and at most 7 days + 35s (fuzz upper bound). minExpected := 7*24*time.Hour - 5*time.Second maxExpected := 7*24*time.Hour + 35*time.Second if *retryAfter < minExpected || *retryAfter > maxExpected { t.Fatalf("RetryAfter = %v, want between %v and %v", *retryAfter, minExpected, maxExpected) }
// Verify one-time fuzz stability: repeat calls return exact same value if second := rap.RetryAfter(); second == nil || *second != *retryAfter { t.Fatalf("RetryAfter changed across calls: %v vs %v", *second, *retryAfter) }}
func TestClaudeExecutor_HonorsAnthropicRateLimitHeaders_ExecuteStream(t *testing.T) { now := time.Now() fiveHourReset := now.Add(5 * time.Hour).Unix()
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") w.Header().Set("Anthropic-Ratelimit-Unified-Status", "rejected") w.Header().Set("Anthropic-Ratelimit-Unified-5h-Status", "rejected") w.Header().Set("Anthropic-Ratelimit-Unified-5h-Reset", strconv.FormatInt(fiveHourReset, 10)) w.Header().Set("Anthropic-Ratelimit-Unified-7d-Status", "allowed") w.Header().Set("Anthropic-Ratelimit-Unified-7d-Reset", strconv.FormatInt(now.Add(7*24*time.Hour).Unix(), 10)) w.WriteHeader(http.StatusTooManyRequests) _, _ = w.Write([]byte(`{"type":"error","error":{"type":"rate_limit_error","message":"5-hour limit exceeded."}}`)) })) defer server.Close()
executor := NewClaudeExecutor(&config.Config{}) auth := &cliproxyauth.Auth{ ID: "claude-auth-1", Provider: "claude", Attributes: map[string]string{ "api_key": "test-key", "base_url": server.URL, }, }
payload := []byte(`{"messages":[{"role":"user","content":[{"type":"text","text":"hi"}]}]}`) _, err := executor.ExecuteStream(context.Background(), auth, cliproxyexecutor.Request{ Model: "claude-3-5-sonnet-20241022", Payload: payload, }, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude}) if err == nil { t.Fatal("expected error from ExecuteStream, got nil") }
var rap retryAfterProvider if !errors.As(err, &rap) || rap == nil { t.Fatalf("expected error %T to implement RetryAfter() *time.Duration", err) }
retryAfter := rap.RetryAfter() if retryAfter == nil { t.Fatalf("expected non-nil RetryAfter, got nil") }
minExpected := 5*time.Hour - 5*time.Second maxExpected := 5*time.Hour + 35*time.Second if *retryAfter < minExpected || *retryAfter > maxExpected { t.Fatalf("RetryAfter = %v, want between %v and %v (5h window)", *retryAfter, minExpected, maxExpected) }}
func TestClaudeExecutor_RateLimit_BothRejectedUsesLongest(t *testing.T) { now := time.Now() fiveHourReset := now.Add(5 * time.Hour).Unix() sevenDayReset := now.Add(7 * 24 * time.Hour).Unix()
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") w.Header().Set("Anthropic-Ratelimit-Unified-5h-Status", "rejected") w.Header().Set("Anthropic-Ratelimit-Unified-5h-Reset", strconv.FormatInt(fiveHourReset, 10)) w.Header().Set("Anthropic-Ratelimit-Unified-7d-Status", "rejected") w.Header().Set("Anthropic-Ratelimit-Unified-7d-Reset", strconv.FormatInt(sevenDayReset, 10)) w.WriteHeader(http.StatusTooManyRequests) _, _ = w.Write([]byte(`{"type":"error","error":{"type":"rate_limit_error","message":"Both limits exceeded."}}`)) })) defer server.Close()
executor := NewClaudeExecutor(&config.Config{}) auth := &cliproxyauth.Auth{ ID: "claude-auth-1", Provider: "claude", Attributes: map[string]string{ "api_key": "test-key", "base_url": server.URL, }, }
payload := []byte(`{"messages":[{"role":"user","content":[{"type":"text","text":"hi"}]}]}`) _, err := executor.Execute(context.Background(), auth, cliproxyexecutor.Request{ Model: "claude-3-5-sonnet-20241022", Payload: payload, }, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude}) if err == nil { t.Fatal("expected error, got nil") }
var rap retryAfterProvider if !errors.As(err, &rap) || rap == nil { t.Fatalf("expected error %T to implement RetryAfter() *time.Duration", err) }
retryAfter := rap.RetryAfter() if retryAfter == nil { t.Fatalf("expected non-nil RetryAfter, got nil") }
minExpected := 7*24*time.Hour - 5*time.Second maxExpected := 7*24*time.Hour + 35*time.Second if *retryAfter < minExpected || *retryAfter > maxExpected { t.Fatalf("RetryAfter = %v, want between %v and %v", *retryAfter, minExpected, maxExpected) }}
func TestClaudeExecutor_RateLimit_CountTokensHonorsRateLimitReset(t *testing.T) { now := time.Now() sevenDayReset := now.Add(7 * 24 * time.Hour).Unix()
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") w.Header().Set("Anthropic-Ratelimit-Unified-Status", "rejected") w.Header().Set("Anthropic-Ratelimit-Unified-7d-Status", "rejected") w.Header().Set("Anthropic-Ratelimit-Unified-7d-Reset", strconv.FormatInt(sevenDayReset, 10)) w.Header().Set("Retry-After", "604800") w.WriteHeader(http.StatusTooManyRequests) _, _ = w.Write([]byte(`{"type":"error","error":{"type":"rate_limit_error","message":"7-day rate limit exceeded."}}`)) })) defer server.Close()
executor := NewClaudeExecutor(&config.Config{}) auth := &cliproxyauth.Auth{ ID: "claude-auth-1", Provider: "claude", Attributes: map[string]string{ "api_key": "test-key", "base_url": server.URL, }, }
payload := []byte(`{"messages":[{"role":"user","content":[{"type":"text","text":"hi"}]}]}`) ctx := cliproxyexecutor.WithUpstreamAttemptTracker(context.Background()) _, err := executor.countTokensUpstream(ctx, auth, cliproxyexecutor.Request{ Model: "claude-3-5-sonnet-20241022", Payload: payload, }, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude}) if err == nil { t.Fatal("expected error from CountTokens, got nil") } if !cliproxyexecutor.UpstreamAttempted(ctx) { t.Fatal("CountTokens() did not mark the HTTP request as an upstream attempt") }
var rap retryAfterProvider if !errors.As(err, &rap) || rap == nil || rap.RetryAfter() == nil { t.Fatalf("expected CountTokens rate limit error to implement RetryAfter, got %v", err) }
type credentialScopedProvider interface { IsCredentialScoped() bool } var csp credentialScopedProvider if !errors.As(err, &csp) || csp == nil || !csp.IsCredentialScoped() { t.Fatalf("expected CountTokens rate limit error to be credential-scoped, got %v", err) }
minExpected := 7*24*time.Hour - 5*time.Second maxExpected := 7*24*time.Hour + 35*time.Second if *rap.RetryAfter() < minExpected || *rap.RetryAfter() > maxExpected { t.Fatalf("RetryAfter = %v, want between %v and %v", *rap.RetryAfter(), minExpected, maxExpected) }}
func TestClaudeExecutor_RateLimit_CaseInsensitiveRawHeaderMap(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header()["anthropic-ratelimit-unified-status"] = []string{"rejected"} w.Header()["anthropic-ratelimit-unified-7d-status"] = []string{"rejected"} w.Header()["anthropic-ratelimit-unified-7d-reset"] = []string{strconv.FormatInt(time.Now().Add(2*time.Hour).Unix(), 10)} w.Header()["retry-after"] = []string{"7200"} w.WriteHeader(http.StatusTooManyRequests) _, _ = w.Write([]byte(`{"type":"error","error":{"type":"rate_limit_error","message":"Too many requests."}}`)) })) defer server.Close()
executor := NewClaudeExecutor(&config.Config{}) auth := &cliproxyauth.Auth{ ID: "claude-auth-1", Provider: "claude", Attributes: map[string]string{ "api_key": "test-key", "base_url": server.URL, }, }
payload := []byte(`{"messages":[{"role":"user","content":[{"type":"text","text":"hi"}]}]}`) _, err := executor.Execute(context.Background(), auth, cliproxyexecutor.Request{ Model: "claude-3-5-sonnet-20241022", Payload: payload, }, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude}) if err == nil { t.Fatal("expected error, got nil") }
var rap retryAfterProvider if !errors.As(err, &rap) || rap == nil || rap.RetryAfter() == nil { t.Fatalf("expected RetryAfter for non-canonical header map, got %v", err) }
minExpected := 2*time.Hour - 5*time.Second maxExpected := 2*time.Hour + 35*time.Second if *rap.RetryAfter() < minExpected || *rap.RetryAfter() > maxExpected { t.Fatalf("RetryAfter = %v, want between %v and %v", *rap.RetryAfter(), minExpected, maxExpected) }}
func TestClaudeExecutor_RateLimit_FastModeAuthoritativeRejectionHeadersOverrideBody(t *testing.T) { var attemptsCred1 atomic.Int32 var attemptsCred2 atomic.Int32
server1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { attemptsCred1.Add(1) w.Header().Set("Content-Type", "application/json") w.Header().Set("Anthropic-Ratelimit-Unified-Status", "rejected") w.Header().Set("Anthropic-Ratelimit-Unified-7d-Status", "rejected") w.Header().Set("Anthropic-Ratelimit-Unified-7d-Reset", strconv.FormatInt(time.Now().Add(7*24*time.Hour).Unix(), 10)) w.Header().Set("Retry-After", "604800") w.WriteHeader(http.StatusTooManyRequests) // Body text mentioning fast request rejected, but headers explicitly reject unified quota _, _ = w.Write([]byte(`{"type":"error","error":{"type":"rate_limit_error","message":"Fast request rejected"}}`)) })) defer server1.Close()
server2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { attemptsCred2.Add(1) w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte(`{"id":"msg-fast-ok","type":"message","role":"assistant","content":[{"type":"text","text":"hello from cred2"}]}`)) })) defer server2.Close()
cfg := &config.Config{DisableCooling: false} manager := cliproxyauth.NewManager(nil, nil, nil) manager.SetRetryConfig(0, 0, 2)
executor := NewClaudeExecutor(cfg) manager.RegisterExecutor(executor)
baseID := uuid.NewString() auth1 := &cliproxyauth.Auth{ID: baseID + "-fast-override-1", Provider: "claude", Attributes: map[string]string{"api_key": "k1", "base_url": server1.URL}} auth2 := &cliproxyauth.Auth{ID: baseID + "-fast-override-2", Provider: "claude", Attributes: map[string]string{"api_key": "k2", "base_url": server2.URL}}
reg := registry.GetGlobalRegistry() reg.RegisterClient(auth1.ID, "claude", []*registry.ModelInfo{{ID: "claude-3-5-sonnet-20241022"}}) reg.RegisterClient(auth2.ID, "claude", []*registry.ModelInfo{{ID: "claude-3-5-sonnet-20241022"}}) t.Cleanup(func() { reg.UnregisterClient(auth1.ID) reg.UnregisterClient(auth2.ID) })
if _, err := manager.Register(context.Background(), auth1); err != nil { t.Fatalf("register auth1: %v", err) } if _, err := manager.Register(context.Background(), auth2); err != nil { t.Fatalf("register auth2: %v", err) }
payload := []byte(`{"speed":"fast","messages":[{"role":"user","content":[{"type":"text","text":"hi"}]}]}`) resp, err := manager.Execute(context.Background(), []string{"claude"}, cliproxyexecutor.Request{ Model: "claude-3-5-sonnet-20241022", Payload: payload, }, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude}) if err != nil { t.Fatalf("expected failover to cred2 when authoritative rate limit headers present, got: %v", err) } if len(resp.Payload) == 0 { t.Fatal("expected response from cred2") }
if attemptsCred1.Load() != 1 { t.Fatalf("attempts on cred1 = %d, want 1", attemptsCred1.Load()) } if attemptsCred2.Load() != 1 { t.Fatalf("attempts on cred2 = %d, want 1", attemptsCred2.Load()) }
// Verify cred1 was cooled down at credential level registeredAuth, ok := manager.GetByID(auth1.ID) if !ok || registeredAuth == nil { t.Fatal("auth1 not found") } if !registeredAuth.Unavailable || !registeredAuth.Quota.Exceeded { t.Fatalf("cred1 was not cooled down: unavailable=%v quota=%+v", registeredAuth.Unavailable, registeredAuth.Quota) }}
func TestClaudeExecutor_RateLimit_FastEntitlementWithRetryAfterRemainsRequestScoped(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") w.Header().Set("Retry-After", "120") w.WriteHeader(http.StatusTooManyRequests) _, _ = w.Write([]byte(`{"type":"error","error":{"type":"rate_limit_error","message":"Usage credits are required for fast mode."}}`)) })) defer server.Close()
cfg := &config.Config{DisableCooling: false} manager := cliproxyauth.NewManager(nil, nil, nil) manager.SetRetryConfig(0, 0, 2)
executor := NewClaudeExecutor(cfg) manager.RegisterExecutor(executor)
baseID := uuid.NewString() auth := &cliproxyauth.Auth{ID: baseID + "-fast-entitlement", Provider: "claude", Attributes: map[string]string{"api_key": "k1", "base_url": server.URL}}
reg := registry.GetGlobalRegistry() reg.RegisterClient(auth.ID, "claude", []*registry.ModelInfo{{ID: "claude-3-5-sonnet-20241022"}}) t.Cleanup(func() { reg.UnregisterClient(auth.ID) })
if _, err := manager.Register(context.Background(), auth); err != nil { t.Fatalf("register auth: %v", err) }
payload := []byte(`{"speed":"fast","messages":[{"role":"user","content":[{"type":"text","text":"hi"}]}]}`) _, err := manager.Execute(context.Background(), []string{"claude"}, cliproxyexecutor.Request{ Model: "claude-3-5-sonnet-20241022", Payload: payload, }, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude}) if err == nil { t.Fatal("expected error, got nil") }
registeredAuth, ok := manager.GetByID(auth.ID) if !ok || registeredAuth == nil { t.Fatal("auth not found") } if registeredAuth.Unavailable || registeredAuth.Quota.Exceeded { t.Fatalf("fast entitlement refusal incorrectly cooled down the credential: unavailable=%v quota=%+v", registeredAuth.Unavailable, registeredAuth.Quota) }}
func TestClaudeExecutor_AuthManager_CredentialScopeBlocksAllModelsAndAliases(t *testing.T) { var upstreamAttempts atomic.Int32 now := time.Now() sevenDayReset := now.Add(7 * 24 * time.Hour).Unix()
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { upstreamAttempts.Add(1) w.Header().Set("Content-Type", "application/json") w.Header().Set("Anthropic-Ratelimit-Unified-Status", "rejected") w.Header().Set("Anthropic-Ratelimit-Unified-7d-Status", "rejected") w.Header().Set("Anthropic-Ratelimit-Unified-7d-Reset", strconv.FormatInt(sevenDayReset, 10)) w.Header().Set("Retry-After", strconv.FormatInt(7*24*3600, 10)) w.WriteHeader(http.StatusTooManyRequests) _, _ = w.Write([]byte(`{"type":"error","error":{"type":"rate_limit_error","message":"7d limit rejected."}}`)) })) defer server.Close()
cfg := &config.Config{ DisableCooling: false, } manager := cliproxyauth.NewManager(nil, nil, nil) manager.SetRetryConfig(0, 0, 0)
executor := NewClaudeExecutor(cfg) manager.RegisterExecutor(executor)
baseID := uuid.NewString() auth := &cliproxyauth.Auth{ ID: baseID + "-claude-cred", Provider: "claude", Attributes: map[string]string{ "api_key": "test-key", "base_url": server.URL, }, }
reg := registry.GetGlobalRegistry() reg.RegisterClient(auth.ID, "claude", []*registry.ModelInfo{ {ID: "claude-3-5-sonnet-20241022"}, {ID: "claude-3-opus-20240229"}, {ID: "claude-3-7-sonnet-20250219"}, }) t.Cleanup(func() { reg.UnregisterClient(auth.ID) })
if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { t.Fatalf("failed to register auth: %v", errRegister) }
// 1. Initial request on sonnet triggers 429 and records 7d cooldown payload := []byte(`{"messages":[{"role":"user","content":[{"type":"text","text":"hi"}]}]}`) _, err := manager.Execute(context.Background(), []string{"claude"}, cliproxyexecutor.Request{ Model: "claude-3-5-sonnet-20241022", Payload: payload, }, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude}) if err == nil { t.Fatal("expected error on first execute, got nil") }
if attempts := upstreamAttempts.Load(); attempts != 1 { t.Fatalf("upstream attempts = %d, want 1", attempts) }
// 2. Try requesting a completely different model (opus) on the same credential -> must be blocked locally _, errOpus := manager.Execute(context.Background(), []string{"claude"}, cliproxyexecutor.Request{ Model: "claude-3-opus-20240229", Payload: payload, }, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude}) if errOpus == nil { t.Fatal("expected error for opus, got nil") } if attempts := upstreamAttempts.Load(); attempts != 1 { t.Fatalf("upstream attempts after opus = %d, want 1 (must be blocked locally)", attempts) }
// 3. Try requesting a thinking suffix alias on the same credential -> must also be blocked locally _, errThinking := manager.Execute(context.Background(), []string{"claude"}, cliproxyexecutor.Request{ Model: "claude-3-7-sonnet-20250219-thinking-16k", Payload: payload, }, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude}) if errThinking == nil { t.Fatal("expected error for thinking suffix, got nil") } if attempts := upstreamAttempts.Load(); attempts != 1 { t.Fatalf("upstream attempts after thinking suffix = %d, want 1 (must be blocked locally)", attempts) }
// 4. Try streaming execution for opus on the same cooling credential -> must also be blocked locally _, errStream := manager.ExecuteStream(context.Background(), []string{"claude"}, cliproxyexecutor.Request{ Model: "claude-3-opus-20240229", Payload: payload, }, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude}) if errStream == nil { t.Fatal("expected error for streaming opus, got nil") } if attempts := upstreamAttempts.Load(); attempts != 1 { t.Fatalf("upstream attempts after streaming opus = %d, want 1 (must be blocked locally)", attempts) }}
func TestClaudeExecutor_AuthManager_OrdinaryModel429DoesNotBlockSiblingModels(t *testing.T) { var attemptsSonnet atomic.Int32 var attemptsOpus atomic.Int32
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { body, _ := io.ReadAll(r.Body) if strings.Contains(string(body), "claude-3-5-sonnet") { attemptsSonnet.Add(1) w.Header().Set("Content-Type", "application/json") w.Header().Set("Anthropic-Ratelimit-Unified-5h-Status", "allowed") w.Header().Set("Anthropic-Ratelimit-Unified-7d-Status", "allowed") w.WriteHeader(http.StatusTooManyRequests) _, _ = w.Write([]byte(`{"type":"error","error":{"type":"rate_limit_error","message":"Model rate limit exceeded."}}`)) return } attemptsOpus.Add(1) w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte(`{"id":"msg-opus","type":"message","role":"assistant","content":[{"type":"text","text":"hello from opus"}]}`)) })) defer server.Close()
cfg := &config.Config{DisableCooling: false} manager := cliproxyauth.NewManager(nil, nil, nil) manager.SetRetryConfig(0, 0, 0)
executor := NewClaudeExecutor(cfg) manager.RegisterExecutor(executor)
baseID := uuid.NewString() auth := &cliproxyauth.Auth{ ID: baseID + "-ordinary-429", Provider: "claude", Attributes: map[string]string{ "api_key": "test-key", "base_url": server.URL, }, }
reg := registry.GetGlobalRegistry() reg.RegisterClient(auth.ID, "claude", []*registry.ModelInfo{ {ID: "claude-3-5-sonnet-20241022"}, {ID: "claude-3-opus-20240229"}, }) t.Cleanup(func() { reg.UnregisterClient(auth.ID) })
if _, err := manager.Register(context.Background(), auth); err != nil { t.Fatalf("register auth: %v", err) }
// 1. Initial request on sonnet triggers ordinary model 429 payloadSonnet := []byte(`{"messages":[{"role":"user","content":[{"type":"text","text":"hi"}]}]}`) _, errSonnet := manager.Execute(context.Background(), []string{"claude"}, cliproxyexecutor.Request{ Model: "claude-3-5-sonnet-20241022", Payload: payloadSonnet, }, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude}) if errSonnet == nil { t.Fatal("expected error on sonnet execute, got nil") } if attemptsSonnet.Load() != 1 { t.Fatalf("sonnet attempts = %d, want 1", attemptsSonnet.Load()) }
// 2. Request on opus MUST succeed on the same credential (not blocked by ordinary model-level 429) payloadOpus := []byte(`{"messages":[{"role":"user","content":[{"type":"text","text":"hi opus"}]}]}`) respOpus, errOpus := manager.Execute(context.Background(), []string{"claude"}, cliproxyexecutor.Request{ Model: "claude-3-opus-20240229", Payload: payloadOpus, }, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude}) if errOpus != nil { t.Fatalf("expected opus to succeed on same credential, got error: %v", errOpus) } if len(respOpus.Payload) == 0 { t.Fatal("expected non-empty response for opus") } if attemptsOpus.Load() != 1 { t.Fatalf("opus attempts = %d, want 1", attemptsOpus.Load()) }}
func TestClaudeExecutor_AuthManager_AlternativeCredentialCanBeSelected(t *testing.T) { var attemptsCred1 atomic.Int32 var attemptsCred2 atomic.Int32 now := time.Now() sevenDayReset := now.Add(7 * 24 * time.Hour).Unix()
server1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { attemptsCred1.Add(1) w.Header().Set("Content-Type", "application/json") w.Header().Set("Anthropic-Ratelimit-Unified-Status", "rejected") w.Header().Set("Anthropic-Ratelimit-Unified-7d-Status", "rejected") w.Header().Set("Anthropic-Ratelimit-Unified-7d-Reset", strconv.FormatInt(sevenDayReset, 10)) w.WriteHeader(http.StatusTooManyRequests) _, _ = w.Write([]byte(`{"type":"error","error":{"type":"rate_limit_error","message":"7d limit rejected."}}`)) })) defer server1.Close()
server2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { attemptsCred2.Add(1) w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte(`{"id":"msg-123","type":"message","role":"assistant","content":[{"type":"text","text":"hello from cred2"}]}`)) })) defer server2.Close()
cfg := &config.Config{ DisableCooling: false, } manager := cliproxyauth.NewManager(nil, nil, nil) manager.SetRetryConfig(0, 0, 2)
executor := NewClaudeExecutor(cfg) manager.RegisterExecutor(executor)
baseID := uuid.NewString() auth1 := &cliproxyauth.Auth{ ID: baseID + "-claude-cred-1", Provider: "claude", Attributes: map[string]string{ "api_key": "test-key-1", "base_url": server1.URL, }, } auth2 := &cliproxyauth.Auth{ ID: baseID + "-claude-cred-2", Provider: "claude", Attributes: map[string]string{ "api_key": "test-key-2", "base_url": server2.URL, }, }
reg := registry.GetGlobalRegistry() reg.RegisterClient(auth1.ID, "claude", []*registry.ModelInfo{{ID: "claude-3-5-sonnet-20241022"}}) reg.RegisterClient(auth2.ID, "claude", []*registry.ModelInfo{{ID: "claude-3-5-sonnet-20241022"}}) t.Cleanup(func() { reg.UnregisterClient(auth1.ID) reg.UnregisterClient(auth2.ID) })
if _, err := manager.Register(context.Background(), auth1); err != nil { t.Fatalf("register auth1: %v", err) } if _, err := manager.Register(context.Background(), auth2); err != nil { t.Fatalf("register auth2: %v", err) }
payload := []byte(`{"messages":[{"role":"user","content":[{"type":"text","text":"hi"}]}]}`) resp, err := manager.Execute(context.Background(), []string{"claude"}, cliproxyexecutor.Request{ Model: "claude-3-5-sonnet-20241022", Payload: payload, }, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude}) if err != nil { t.Fatalf("expected successful failover to cred2, got error: %v", err) }
if attemptsCred1.Load() != 1 { t.Fatalf("attempts on cred1 = %d, want 1", attemptsCred1.Load()) } if attemptsCred2.Load() != 1 { t.Fatalf("attempts on cred2 = %d, want 1", attemptsCred2.Load()) } if len(resp.Payload) == 0 { t.Fatal("expected non-empty response payload from cred2") }
// Next request should directly use cred2 without attempting cred1 (which is cooling down) resp2, err2 := manager.Execute(context.Background(), []string{"claude"}, cliproxyexecutor.Request{ Model: "claude-3-5-sonnet-20241022", Payload: payload, }, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude}) if err2 != nil { t.Fatalf("expected successful request on cred2, got error: %v", err2) } if attemptsCred1.Load() != 1 { t.Fatalf("attempts on cred1 after 2nd request = %d, want 1 (must stay 1)", attemptsCred1.Load()) } if attemptsCred2.Load() != 2 { t.Fatalf("attempts on cred2 after 2nd request = %d, want 2", attemptsCred2.Load()) } if len(resp2.Payload) == 0 { t.Fatal("expected non-empty response payload from 2nd request") }}
func TestClaudeExecutor_AuthManager_MultiModelPoolStreamStopsProbingOn429(t *testing.T) { var attemptsCred1 atomic.Int32 var attemptsCred2 atomic.Int32
server1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { attemptsCred1.Add(1) w.Header().Set("Content-Type", "application/json") w.Header().Set("Anthropic-Ratelimit-Unified-Status", "rejected") w.Header().Set("Anthropic-Ratelimit-Unified-7d-Status", "rejected") w.Header().Set("Anthropic-Ratelimit-Unified-7d-Reset", strconv.FormatInt(time.Now().Add(7*24*time.Hour).Unix(), 10)) w.Header().Set("Retry-After", "604800") w.WriteHeader(http.StatusTooManyRequests) _, _ = w.Write([]byte(`{"type":"error","error":{"type":"rate_limit_error","message":"rate limited"}}`)) })) defer server1.Close()
server2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { attemptsCred2.Add(1) w.Header().Set("Content-Type", "text/event-stream") w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg-1\",\"model\":\"claude-3-5-sonnet-20241022\"}}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n")) })) defer server2.Close()
cfg := &config.Config{DisableCooling: false} manager := cliproxyauth.NewManager(nil, nil, nil) manager.SetRetryConfig(0, 0, 2)
manager.SetOAuthModelAlias(map[string][]config.OAuthModelAlias{ "claude": { {Name: "claude-3-5-sonnet-20241022", Alias: "claude-pool-alias"}, {Name: "claude-3-opus-20240229", Alias: "claude-pool-alias"}, }, })
executor := NewClaudeExecutor(cfg) manager.RegisterExecutor(executor)
baseID := uuid.NewString() auth1 := &cliproxyauth.Auth{ID: baseID + "-pool-1", Provider: "claude", Attributes: map[string]string{"api_key": "k1", "base_url": server1.URL}} auth2 := &cliproxyauth.Auth{ID: baseID + "-pool-2", Provider: "claude", Attributes: map[string]string{"api_key": "k2", "base_url": server2.URL}}
reg := registry.GetGlobalRegistry() reg.RegisterClient(auth1.ID, "claude", []*registry.ModelInfo{ {ID: "claude-pool-alias"}, {ID: "claude-3-5-sonnet-20241022"}, {ID: "claude-3-opus-20240229"}, }) reg.RegisterClient(auth2.ID, "claude", []*registry.ModelInfo{ {ID: "claude-pool-alias"}, {ID: "claude-3-5-sonnet-20241022"}, {ID: "claude-3-opus-20240229"}, }) t.Cleanup(func() { reg.UnregisterClient(auth1.ID) reg.UnregisterClient(auth2.ID) })
if _, err := manager.Register(context.Background(), auth1); err != nil { t.Fatalf("register auth1: %v", err) } if _, err := manager.Register(context.Background(), auth2); err != nil { t.Fatalf("register auth2: %v", err) }
payload := []byte(`{"messages":[{"role":"user","content":[{"type":"text","text":"hi"}]}]}`) res, err := manager.ExecuteStream(context.Background(), []string{"claude"}, cliproxyexecutor.Request{ Model: "claude-pool-alias", Payload: payload, }, cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude}) if err != nil { t.Fatalf("ExecuteStream failed: %v", err) } for chunk := range res.Chunks { if chunk.Err != nil { t.Fatalf("unexpected chunk error: %v", chunk.Err) } }
// Must have tried cred1 exactly once (did NOT probe the 2nd model on cred1 after 429) and failed over to cred2 if got := attemptsCred1.Load(); got != 1 { t.Fatalf("attempts on cred1 = %d, want 1 (must not probe other models on cooled cred)", got) } if got := attemptsCred2.Load(); got != 1 { t.Fatalf("attempts on cred2 = %d, want 1", got) }}