shaw
b3f796972e
Anthropic OAuth/setup-token accounts now expose a Fable-only 7d window in 429 responses (anthropic-ratelimit-unified-7d_oi-*). Previously, a 429 triggered solely by this window fell through to the legacy path and marked the whole account rate limited (picking the sooner 5h/7d reset) even though 5h/7d were still "allowed". Backend: - HandleUpstreamError: when only 7d_oi is rejected/exceeded, set a model-level rate limit (scope "claude-fable-5") and skip account-level SetRateLimited / session-window rewrite / temp-unsched entirely. Real 5h/7d exhaustion keeps the exact legacy account-level behavior. - modelRateLimitKeysForRequest: append the Fable family scope for any mapped Anthropic model containing "fable", so claude-fable-5[1m] and other variants are excluded from scheduling together. - Passive sampling: collect 7d_oi utilization/reset into Extra (passive_usage_7d_oi_*), including from the 429 response itself so the usage bar doesn't freeze at a stale value during the limit window; cleared on 5h window reset like the existing 7d keys. - Usage API: new seven_day_fable window in UsageInfo. GetPassiveUsage builds it from passive samples; active queries parse the upstream seven_day_overage_included field (claim name from the representative-claim header) and fall back to passive data when the upstream usage API doesn't return it, so the active-query button keeps the bar populated. seven_day_sonnet handling is preserved unchanged. Frontend: - /admin/accounts usage cell renders a "7d F" (amber) progress bar next to the existing 5h / 7d / 7d S bars for Anthropic OAuth/setup-token accounts.
546 lines
14 KiB
Go
546 lines
14 KiB
Go
package service
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestIsModelRateLimited(t *testing.T) {
|
|
now := time.Now()
|
|
future := now.Add(10 * time.Minute).Format(time.RFC3339)
|
|
past := now.Add(-10 * time.Minute).Format(time.RFC3339)
|
|
|
|
tests := []struct {
|
|
name string
|
|
account *Account
|
|
requestedModel string
|
|
expected bool
|
|
}{
|
|
{
|
|
name: "official model ID hit - claude-sonnet-4-5",
|
|
account: &Account{
|
|
Extra: map[string]any{
|
|
modelRateLimitsKey: map[string]any{
|
|
"claude-sonnet-4-5": map[string]any{
|
|
"rate_limit_reset_at": future,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
requestedModel: "claude-sonnet-4-5",
|
|
expected: true,
|
|
},
|
|
{
|
|
name: "official model ID hit via mapping - request claude-3-5-sonnet, mapped to claude-sonnet-4-5",
|
|
account: &Account{
|
|
Credentials: map[string]any{
|
|
"model_mapping": map[string]any{
|
|
"claude-3-5-sonnet": "claude-sonnet-4-5",
|
|
},
|
|
},
|
|
Extra: map[string]any{
|
|
modelRateLimitsKey: map[string]any{
|
|
"claude-sonnet-4-5": map[string]any{
|
|
"rate_limit_reset_at": future,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
requestedModel: "claude-3-5-sonnet",
|
|
expected: true,
|
|
},
|
|
{
|
|
name: "no rate limit - expired",
|
|
account: &Account{
|
|
Extra: map[string]any{
|
|
modelRateLimitsKey: map[string]any{
|
|
"claude-sonnet-4-5": map[string]any{
|
|
"rate_limit_reset_at": past,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
requestedModel: "claude-sonnet-4-5",
|
|
expected: false,
|
|
},
|
|
{
|
|
name: "no rate limit - no matching key",
|
|
account: &Account{
|
|
Extra: map[string]any{
|
|
modelRateLimitsKey: map[string]any{
|
|
"gemini-3-flash": map[string]any{
|
|
"rate_limit_reset_at": future,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
requestedModel: "claude-sonnet-4-5",
|
|
expected: false,
|
|
},
|
|
{
|
|
name: "no rate limit - unsupported model",
|
|
account: &Account{},
|
|
requestedModel: "gpt-4",
|
|
expected: false,
|
|
},
|
|
{
|
|
name: "no rate limit - empty model",
|
|
account: &Account{},
|
|
requestedModel: "",
|
|
expected: false,
|
|
},
|
|
{
|
|
name: "gemini model hit",
|
|
account: &Account{
|
|
Extra: map[string]any{
|
|
modelRateLimitsKey: map[string]any{
|
|
"gemini-3-pro-high": map[string]any{
|
|
"rate_limit_reset_at": future,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
requestedModel: "gemini-3-pro-high",
|
|
expected: true,
|
|
},
|
|
{
|
|
name: "antigravity platform - gemini-3-pro-preview mapped to gemini-3-pro-high",
|
|
account: &Account{
|
|
Platform: PlatformAntigravity,
|
|
Extra: map[string]any{
|
|
modelRateLimitsKey: map[string]any{
|
|
"gemini-3-pro-high": map[string]any{
|
|
"rate_limit_reset_at": future,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
requestedModel: "gemini-3-pro-preview",
|
|
expected: true,
|
|
},
|
|
{
|
|
name: "antigravity platform - gemini family rate limit blocks mapped preview",
|
|
account: &Account{
|
|
Platform: PlatformAntigravity,
|
|
Extra: map[string]any{
|
|
modelRateLimitsKey: map[string]any{
|
|
antigravityGeminiModelRateLimitKey: map[string]any{
|
|
"rate_limit_reset_at": future,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
requestedModel: "gemini-3-pro-preview",
|
|
expected: true,
|
|
},
|
|
{
|
|
name: "antigravity platform - gemini family rate limit does not block claude",
|
|
account: &Account{
|
|
Platform: PlatformAntigravity,
|
|
Extra: map[string]any{
|
|
modelRateLimitsKey: map[string]any{
|
|
antigravityGeminiModelRateLimitKey: map[string]any{
|
|
"rate_limit_reset_at": future,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
requestedModel: "claude-sonnet-4-5",
|
|
expected: false,
|
|
},
|
|
{
|
|
name: "non-antigravity platform - gemini-3-pro-preview NOT mapped",
|
|
account: &Account{
|
|
Platform: PlatformGemini,
|
|
Extra: map[string]any{
|
|
modelRateLimitsKey: map[string]any{
|
|
"gemini-3-pro-high": map[string]any{
|
|
"rate_limit_reset_at": future,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
requestedModel: "gemini-3-pro-preview",
|
|
expected: false, // gemini 平台不走 antigravity 映射
|
|
},
|
|
{
|
|
name: "antigravity platform - claude-opus-4-5-thinking mapped to opus-4-6-thinking",
|
|
account: &Account{
|
|
Platform: PlatformAntigravity,
|
|
Extra: map[string]any{
|
|
modelRateLimitsKey: map[string]any{
|
|
"claude-opus-4-6-thinking": map[string]any{
|
|
"rate_limit_reset_at": future,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
requestedModel: "claude-opus-4-5-thinking",
|
|
expected: true,
|
|
},
|
|
{
|
|
name: "no scope fallback - claude_sonnet should not match",
|
|
account: &Account{
|
|
Extra: map[string]any{
|
|
modelRateLimitsKey: map[string]any{
|
|
"claude_sonnet": map[string]any{
|
|
"rate_limit_reset_at": future,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
requestedModel: "claude-3-5-sonnet-20241022",
|
|
expected: false,
|
|
},
|
|
{
|
|
name: "openai image generation family key blocks image model",
|
|
account: &Account{
|
|
Platform: PlatformOpenAI,
|
|
Extra: map[string]any{
|
|
modelRateLimitsKey: map[string]any{
|
|
openAIImageGenerationRateLimitKey: map[string]any{
|
|
"rate_limit_reset_at": future,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
requestedModel: "gpt-image-2",
|
|
expected: true,
|
|
},
|
|
{
|
|
name: "openai image generation family key does not block text model",
|
|
account: &Account{
|
|
Platform: PlatformOpenAI,
|
|
Extra: map[string]any{
|
|
modelRateLimitsKey: map[string]any{
|
|
openAIImageGenerationRateLimitKey: map[string]any{
|
|
"rate_limit_reset_at": future,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
requestedModel: "gpt-5.4",
|
|
expected: false,
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
result := tt.account.isModelRateLimitedWithContext(context.Background(), tt.requestedModel)
|
|
if result != tt.expected {
|
|
t.Errorf("isModelRateLimited(%q) = %v, want %v", tt.requestedModel, result, tt.expected)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestIsModelRateLimited_OpenAIImageGenerationIntentBlocksTextModelImageTool(t *testing.T) {
|
|
future := time.Now().Add(10 * time.Minute).Format(time.RFC3339)
|
|
account := &Account{
|
|
Platform: PlatformOpenAI,
|
|
Extra: map[string]any{
|
|
modelRateLimitsKey: map[string]any{
|
|
openAIImageGenerationRateLimitKey: map[string]any{
|
|
"rate_limit_reset_at": future,
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
require.False(t, account.isModelRateLimitedWithContext(context.Background(), "gpt-5.4"))
|
|
require.True(t, account.isModelRateLimitedWithContext(WithOpenAIImageGenerationIntent(context.Background()), "gpt-5.4"))
|
|
}
|
|
|
|
func TestIsModelRateLimited_Antigravity_ThinkingAffectsModelKey(t *testing.T) {
|
|
now := time.Now()
|
|
future := now.Add(10 * time.Minute).Format(time.RFC3339)
|
|
|
|
account := &Account{
|
|
Platform: PlatformAntigravity,
|
|
Extra: map[string]any{
|
|
modelRateLimitsKey: map[string]any{
|
|
"claude-sonnet-4-5-thinking": map[string]any{
|
|
"rate_limit_reset_at": future,
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
ctx := context.WithValue(context.Background(), ctxkey.ThinkingEnabled, true)
|
|
if !account.isModelRateLimitedWithContext(ctx, "claude-sonnet-4-5") {
|
|
t.Errorf("expected model to be rate limited")
|
|
}
|
|
}
|
|
|
|
func TestGetModelRateLimitRemainingTime(t *testing.T) {
|
|
now := time.Now()
|
|
future10m := now.Add(10 * time.Minute).Format(time.RFC3339)
|
|
future5m := now.Add(5 * time.Minute).Format(time.RFC3339)
|
|
past := now.Add(-10 * time.Minute).Format(time.RFC3339)
|
|
|
|
tests := []struct {
|
|
name string
|
|
account *Account
|
|
requestedModel string
|
|
minExpected time.Duration
|
|
maxExpected time.Duration
|
|
}{
|
|
{
|
|
name: "nil account",
|
|
account: nil,
|
|
requestedModel: "claude-sonnet-4-5",
|
|
minExpected: 0,
|
|
maxExpected: 0,
|
|
},
|
|
{
|
|
name: "model rate limited - direct hit",
|
|
account: &Account{
|
|
Extra: map[string]any{
|
|
modelRateLimitsKey: map[string]any{
|
|
"claude-sonnet-4-5": map[string]any{
|
|
"rate_limit_reset_at": future10m,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
requestedModel: "claude-sonnet-4-5",
|
|
minExpected: 9 * time.Minute,
|
|
maxExpected: 11 * time.Minute,
|
|
},
|
|
{
|
|
name: "model rate limited - via mapping",
|
|
account: &Account{
|
|
Credentials: map[string]any{
|
|
"model_mapping": map[string]any{
|
|
"claude-3-5-sonnet": "claude-sonnet-4-5",
|
|
},
|
|
},
|
|
Extra: map[string]any{
|
|
modelRateLimitsKey: map[string]any{
|
|
"claude-sonnet-4-5": map[string]any{
|
|
"rate_limit_reset_at": future5m,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
requestedModel: "claude-3-5-sonnet",
|
|
minExpected: 4 * time.Minute,
|
|
maxExpected: 6 * time.Minute,
|
|
},
|
|
{
|
|
name: "expired rate limit",
|
|
account: &Account{
|
|
Extra: map[string]any{
|
|
modelRateLimitsKey: map[string]any{
|
|
"claude-sonnet-4-5": map[string]any{
|
|
"rate_limit_reset_at": past,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
requestedModel: "claude-sonnet-4-5",
|
|
minExpected: 0,
|
|
maxExpected: 0,
|
|
},
|
|
{
|
|
name: "no rate limit data",
|
|
account: &Account{},
|
|
requestedModel: "claude-sonnet-4-5",
|
|
minExpected: 0,
|
|
maxExpected: 0,
|
|
},
|
|
{
|
|
name: "no scope fallback",
|
|
account: &Account{
|
|
Extra: map[string]any{
|
|
modelRateLimitsKey: map[string]any{
|
|
"claude_sonnet": map[string]any{
|
|
"rate_limit_reset_at": future5m,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
requestedModel: "claude-3-5-sonnet-20241022",
|
|
minExpected: 0,
|
|
maxExpected: 0,
|
|
},
|
|
{
|
|
name: "antigravity platform - claude-opus-4-5-thinking mapped to opus-4-6-thinking",
|
|
account: &Account{
|
|
Platform: PlatformAntigravity,
|
|
Extra: map[string]any{
|
|
modelRateLimitsKey: map[string]any{
|
|
"claude-opus-4-6-thinking": map[string]any{
|
|
"rate_limit_reset_at": future5m,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
requestedModel: "claude-opus-4-5-thinking",
|
|
minExpected: 4 * time.Minute,
|
|
maxExpected: 6 * time.Minute,
|
|
},
|
|
{
|
|
name: "antigravity platform - gemini family rate limit remaining",
|
|
account: &Account{
|
|
Platform: PlatformAntigravity,
|
|
Extra: map[string]any{
|
|
modelRateLimitsKey: map[string]any{
|
|
antigravityGeminiModelRateLimitKey: map[string]any{
|
|
"rate_limit_reset_at": future10m,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
requestedModel: "gemini-3-pro-preview",
|
|
minExpected: 9 * time.Minute,
|
|
maxExpected: 11 * time.Minute,
|
|
},
|
|
{
|
|
name: "antigravity platform - gemini family remaining ignored for claude",
|
|
account: &Account{
|
|
Platform: PlatformAntigravity,
|
|
Extra: map[string]any{
|
|
modelRateLimitsKey: map[string]any{
|
|
antigravityGeminiModelRateLimitKey: map[string]any{
|
|
"rate_limit_reset_at": future10m,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
requestedModel: "claude-sonnet-4-5",
|
|
minExpected: 0,
|
|
maxExpected: 0,
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
result := tt.account.GetModelRateLimitRemainingTimeWithContext(context.Background(), tt.requestedModel)
|
|
if result < tt.minExpected || result > tt.maxExpected {
|
|
t.Errorf("GetModelRateLimitRemainingTime() = %v, want between %v and %v", result, tt.minExpected, tt.maxExpected)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestGetRateLimitRemainingTime(t *testing.T) {
|
|
now := time.Now()
|
|
future15m := now.Add(15 * time.Minute).Format(time.RFC3339)
|
|
future5m := now.Add(5 * time.Minute).Format(time.RFC3339)
|
|
|
|
tests := []struct {
|
|
name string
|
|
account *Account
|
|
requestedModel string
|
|
minExpected time.Duration
|
|
maxExpected time.Duration
|
|
}{
|
|
{
|
|
name: "nil account",
|
|
account: nil,
|
|
requestedModel: "claude-sonnet-4-5",
|
|
minExpected: 0,
|
|
maxExpected: 0,
|
|
},
|
|
{
|
|
name: "model rate limited - 15 minutes",
|
|
account: &Account{
|
|
Platform: PlatformAntigravity,
|
|
Extra: map[string]any{
|
|
modelRateLimitsKey: map[string]any{
|
|
"claude-sonnet-4-5": map[string]any{
|
|
"rate_limit_reset_at": future15m,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
requestedModel: "claude-sonnet-4-5",
|
|
minExpected: 14 * time.Minute,
|
|
maxExpected: 16 * time.Minute,
|
|
},
|
|
{
|
|
name: "only model rate limited",
|
|
account: &Account{
|
|
Platform: PlatformAntigravity,
|
|
Extra: map[string]any{
|
|
modelRateLimitsKey: map[string]any{
|
|
"claude-sonnet-4-5": map[string]any{
|
|
"rate_limit_reset_at": future5m,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
requestedModel: "claude-sonnet-4-5",
|
|
minExpected: 4 * time.Minute,
|
|
maxExpected: 6 * time.Minute,
|
|
},
|
|
{
|
|
name: "neither rate limited",
|
|
account: &Account{
|
|
Platform: PlatformAntigravity,
|
|
},
|
|
requestedModel: "claude-sonnet-4-5",
|
|
minExpected: 0,
|
|
maxExpected: 0,
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
result := tt.account.GetRateLimitRemainingTimeWithContext(context.Background(), tt.requestedModel)
|
|
if result < tt.minExpected || result > tt.maxExpected {
|
|
t.Errorf("GetRateLimitRemainingTime() = %v, want between %v and %v", result, tt.minExpected, tt.maxExpected)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestIsModelRateLimited_AnthropicFableFamilyKey(t *testing.T) {
|
|
now := time.Now()
|
|
future := now.Add(48 * time.Hour).Format(time.RFC3339)
|
|
|
|
account := &Account{
|
|
Platform: PlatformAnthropic,
|
|
Extra: map[string]any{
|
|
modelRateLimitsKey: map[string]any{
|
|
anthropicFableRateLimitKey: map[string]any{
|
|
"rate_limit_reset_at": future,
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
tests := []struct {
|
|
requestedModel string
|
|
expected bool
|
|
}{
|
|
{"claude-fable-5", true},
|
|
{"claude-fable-5[1m]", true}, // 家族 key 覆盖变体
|
|
{"Claude-Fable-5-20260601", true}, // 大小写不敏感
|
|
{"claude-sonnet-4-6", false}, // 其他模型不受影响
|
|
{"claude-opus-4-8", false},
|
|
}
|
|
|
|
for _, tc := range tests {
|
|
t.Run(tc.requestedModel, func(t *testing.T) {
|
|
got := account.isModelRateLimitedWithContext(context.Background(), tc.requestedModel)
|
|
require.Equal(t, tc.expected, got)
|
|
remaining := account.GetModelRateLimitRemainingTimeWithContext(context.Background(), tc.requestedModel)
|
|
require.Equal(t, tc.expected, remaining > 0)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestIsAnthropicFableModel(t *testing.T) {
|
|
require.True(t, isAnthropicFableModel("claude-fable-5"))
|
|
require.True(t, isAnthropicFableModel("claude-fable-5[1m]"))
|
|
require.True(t, isAnthropicFableModel("Claude-Fable-5"))
|
|
require.False(t, isAnthropicFableModel("claude-sonnet-4-6"))
|
|
require.False(t, isAnthropicFableModel(""))
|
|
}
|