Add projects
This commit is contained in:
parent
2d3a9ad623
commit
8b607dd700
1802 changed files with 503346 additions and 2 deletions
315
backend/sdk/cliproxy/auth/claude_ratelimit_cooldown_test.go
Normal file
315
backend/sdk/cliproxy/auth/claude_ratelimit_cooldown_test.go
Normal file
|
|
@ -0,0 +1,315 @@
|
|||
package auth
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/registry"
|
||||
)
|
||||
|
||||
func TestAuthManager_ConcurrentSuccessDoesNotClearActiveCredentialCooldown(t *testing.T) {
|
||||
now := time.Now()
|
||||
sevenDayReset := now.Add(7 * 24 * time.Hour)
|
||||
|
||||
manager := NewManager(nil, nil, nil)
|
||||
|
||||
baseID := uuid.NewString()
|
||||
auth := &Auth{
|
||||
ID: baseID + "-claude-concurrent",
|
||||
Provider: "claude",
|
||||
Attributes: map[string]string{
|
||||
"api_key": "test-key",
|
||||
},
|
||||
}
|
||||
|
||||
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. Request A fails with 7d credential-scoped cooldown
|
||||
sevenDayDuration := 7 * 24 * time.Hour
|
||||
manager.MarkResult(context.Background(), Result{
|
||||
AuthID: auth.ID,
|
||||
Provider: "claude",
|
||||
Model: "claude-3-5-sonnet-20241022",
|
||||
Success: false,
|
||||
RetryAfter: &sevenDayDuration,
|
||||
CredentialScope: true,
|
||||
Error: &Error{HTTPStatus: http.StatusTooManyRequests, Message: "7d limit rejected"},
|
||||
})
|
||||
|
||||
// 2. An earlier in-flight request on opus returns 200 OK after the 429
|
||||
manager.MarkResult(context.Background(), Result{
|
||||
AuthID: auth.ID,
|
||||
Provider: "claude",
|
||||
Model: "claude-3-opus-20240229",
|
||||
Success: true,
|
||||
})
|
||||
|
||||
// 3. The credential MUST still be blocked for all models
|
||||
updatedAuth, ok := manager.GetByID(auth.ID)
|
||||
if !ok || updatedAuth == nil {
|
||||
t.Fatal("auth not found")
|
||||
}
|
||||
if !updatedAuth.Quota.Exceeded || !updatedAuth.Quota.NextRecoverAt.After(now.Add(6*24*time.Hour)) {
|
||||
t.Fatalf("auth quota was cleared or shortened by concurrent success: quota=%+v", updatedAuth.Quota)
|
||||
}
|
||||
|
||||
// Selecting any model on this credential must be blocked locally
|
||||
for _, m := range []string{"claude-3-5-sonnet-20241022", "claude-3-opus-20240229", "claude-3-7-sonnet-20250219"} {
|
||||
blocked, reason, next := isAuthBlockedForModel(updatedAuth, m, time.Now())
|
||||
if !blocked {
|
||||
t.Fatalf("model %q was unblocked despite active 7d credential cooldown", m)
|
||||
}
|
||||
if reason != blockReasonCooldown || next.Before(sevenDayReset.Add(-time.Minute)) {
|
||||
t.Fatalf("model %q block reason=%v next=%v, want cooldown ~7d", m, reason, next)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAuthManager_UpdatePreservesActiveCredentialCooldown(t *testing.T) {
|
||||
now := time.Now()
|
||||
manager := NewManager(nil, nil, nil)
|
||||
|
||||
baseID := uuid.NewString()
|
||||
auth := &Auth{
|
||||
ID: baseID + "-claude-update",
|
||||
Provider: "claude",
|
||||
Attributes: map[string]string{
|
||||
"api_key": "test-key",
|
||||
},
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
|
||||
sevenDayDuration := 7 * 24 * time.Hour
|
||||
manager.MarkResult(context.Background(), Result{
|
||||
AuthID: auth.ID,
|
||||
Provider: "claude",
|
||||
Model: "claude-3-5-sonnet-20241022",
|
||||
Success: false,
|
||||
RetryAfter: &sevenDayDuration,
|
||||
CredentialScope: true,
|
||||
Error: &Error{HTTPStatus: http.StatusTooManyRequests, Message: "7d limit rejected"},
|
||||
})
|
||||
|
||||
// Reload/update auth (e.g. config reload or token refresh)
|
||||
updatedAuth := &Auth{
|
||||
ID: auth.ID,
|
||||
Provider: "claude",
|
||||
Attributes: map[string]string{
|
||||
"api_key": "test-key-updated",
|
||||
},
|
||||
}
|
||||
if _, err := manager.Update(context.Background(), updatedAuth); err != nil {
|
||||
t.Fatalf("update auth: %v", err)
|
||||
}
|
||||
|
||||
persistedAuth, ok := manager.GetByID(auth.ID)
|
||||
if !ok || persistedAuth == nil {
|
||||
t.Fatal("auth not found after update")
|
||||
}
|
||||
if !persistedAuth.Quota.Exceeded || persistedAuth.Quota.Reason != "credential_quota" || !persistedAuth.Quota.NextRecoverAt.After(now.Add(6*24*time.Hour)) {
|
||||
t.Fatalf("credential cooldown was lost after Update: quota=%+v", persistedAuth.Quota)
|
||||
}
|
||||
|
||||
blocked, reason, _ := isAuthBlockedForModel(persistedAuth, "claude-3-5-sonnet-20241022", time.Now())
|
||||
if !blocked || reason != blockReasonCooldown {
|
||||
t.Fatalf("model unblocked after Update: blocked=%v reason=%v", blocked, reason)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAuthManager_DisableCoolingDoesNotPermanentlyBlock(t *testing.T) {
|
||||
SetQuotaCooldownDisabled(true)
|
||||
t.Cleanup(func() { SetQuotaCooldownDisabled(false) })
|
||||
|
||||
manager := NewManager(nil, nil, nil)
|
||||
|
||||
baseID := uuid.NewString()
|
||||
auth := &Auth{
|
||||
ID: baseID + "-claude-disable-cooling",
|
||||
Provider: "claude",
|
||||
Attributes: map[string]string{
|
||||
"api_key": "test-key",
|
||||
},
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
|
||||
// 429 arrives while cooling is disabled
|
||||
sevenDayDuration := 7 * 24 * time.Hour
|
||||
manager.MarkResult(context.Background(), Result{
|
||||
AuthID: auth.ID,
|
||||
Provider: "claude",
|
||||
Model: "claude-3-5-sonnet-20241022",
|
||||
Success: false,
|
||||
RetryAfter: &sevenDayDuration,
|
||||
CredentialScope: true,
|
||||
Error: &Error{HTTPStatus: http.StatusTooManyRequests, Message: "7d limit rejected"},
|
||||
})
|
||||
|
||||
// Must NOT be blocked when cooling is disabled
|
||||
for _, m := range []string{"claude-3-5-sonnet-20241022", "claude-3-opus-20240229"} {
|
||||
updatedAuth, _ := manager.GetByID(auth.ID)
|
||||
blocked, _, _ := isAuthBlockedForModel(updatedAuth, m, time.Now())
|
||||
if blocked {
|
||||
t.Fatalf("model %q was blocked even though cooling is disabled", m)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAuthManager_NonClaudeProvider_Model429DoesNotBlockSiblingModels(t *testing.T) {
|
||||
manager := NewManager(nil, nil, nil)
|
||||
|
||||
baseID := uuid.NewString()
|
||||
auth := &Auth{
|
||||
ID: baseID + "-openai-auth",
|
||||
Provider: "openai",
|
||||
Attributes: map[string]string{
|
||||
"api_key": "test-key",
|
||||
},
|
||||
}
|
||||
|
||||
reg := registry.GetGlobalRegistry()
|
||||
reg.RegisterClient(auth.ID, "openai", []*registry.ModelInfo{
|
||||
{ID: "gpt-4o"},
|
||||
{ID: "gpt-4o-mini"},
|
||||
})
|
||||
t.Cleanup(func() {
|
||||
reg.UnregisterClient(auth.ID)
|
||||
})
|
||||
|
||||
if _, err := manager.Register(context.Background(), auth); err != nil {
|
||||
t.Fatalf("register auth: %v", err)
|
||||
}
|
||||
|
||||
// Regular model 429 on gpt-4o (CredentialScope is false)
|
||||
manager.MarkResult(context.Background(), Result{
|
||||
AuthID: auth.ID,
|
||||
Provider: "openai",
|
||||
Model: "gpt-4o",
|
||||
Success: false,
|
||||
CredentialScope: false,
|
||||
Error: &Error{HTTPStatus: http.StatusTooManyRequests, Message: "rate limit"},
|
||||
})
|
||||
|
||||
// gpt-4o should be blocked
|
||||
updatedAuth, _ := manager.GetByID(auth.ID)
|
||||
blocked4o, _, _ := isAuthBlockedForModel(updatedAuth, "gpt-4o", time.Now())
|
||||
if !blocked4o {
|
||||
t.Fatal("gpt-4o should be blocked after 429")
|
||||
}
|
||||
|
||||
// gpt-4o-mini MUST remain selectable (unaffected by sibling model 429)
|
||||
blockedMini, _, _ := isAuthBlockedForModel(updatedAuth, "gpt-4o-mini", time.Now())
|
||||
if blockedMini {
|
||||
t.Fatal("gpt-4o-mini was incorrectly blocked by sibling model 429")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAuthManager_CooldownPersistenceAcrossRestore(t *testing.T) {
|
||||
manager := NewManager(nil, nil, nil)
|
||||
|
||||
baseID := uuid.NewString()
|
||||
auth := &Auth{
|
||||
ID: baseID + "-persistence-test",
|
||||
Provider: "claude",
|
||||
Attributes: map[string]string{
|
||||
"api_key": "k",
|
||||
},
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
|
||||
futureCooldown := 7 * 24 * time.Hour
|
||||
manager.MarkResult(context.Background(), Result{
|
||||
AuthID: auth.ID,
|
||||
Provider: "claude",
|
||||
Model: "claude-3-5-sonnet-20241022",
|
||||
Success: false,
|
||||
RetryAfter: &futureCooldown,
|
||||
CredentialScope: true,
|
||||
Error: &Error{HTTPStatus: http.StatusTooManyRequests, Message: "7d rejected"},
|
||||
})
|
||||
|
||||
records := manager.cooldownStateRecordsSnapshot()
|
||||
if len(records) == 0 {
|
||||
t.Fatal("expected cooldown state records to be captured")
|
||||
}
|
||||
|
||||
// Create a new manager instance and restore state
|
||||
newManager := NewManager(nil, nil, nil)
|
||||
newAuth := &Auth{
|
||||
ID: auth.ID,
|
||||
Provider: "claude",
|
||||
}
|
||||
if _, err := newManager.Register(context.Background(), newAuth); err != nil {
|
||||
t.Fatalf("register new auth: %v", err)
|
||||
}
|
||||
|
||||
newManager.SetCooldownStateStore(&mockCooldownStateStore{records: records})
|
||||
if err := newManager.RestoreCooldownStates(context.Background()); err != nil {
|
||||
t.Fatalf("RestoreCooldownStates error: %v", err)
|
||||
}
|
||||
|
||||
restoredAuth, ok := newManager.GetByID(auth.ID)
|
||||
if !ok || restoredAuth == nil {
|
||||
t.Fatal("restored auth not found")
|
||||
}
|
||||
if !restoredAuth.Quota.Exceeded || restoredAuth.Quota.NextRecoverAt.Before(time.Now().Add(6*24*time.Hour)) {
|
||||
t.Fatalf("restored auth quota was not preserved: quota=%+v", restoredAuth.Quota)
|
||||
}
|
||||
}
|
||||
|
||||
type mockCooldownStateStore struct {
|
||||
records []CooldownStateRecord
|
||||
}
|
||||
|
||||
func (s *mockCooldownStateStore) Load(context.Context) ([]CooldownStateRecord, error) {
|
||||
return s.records, nil
|
||||
}
|
||||
|
||||
func (s *mockCooldownStateStore) Save(context.Context, []CooldownStateRecord) error {
|
||||
return nil
|
||||
}
|
||||
Loading…
Reference in a new issue