794 lines · 30.4 KB · Go
RawDownload
package 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})

Lines 1–500 of 794.