Files
sub2api/backend/internal/service/ratelimit_service_model_not_found_test.go
T
李建琦 6d655c9903
Release / update-version (push) Has been cancelled
Release / build-frontend (push) Has been cancelled
Release / release (push) Has been cancelled
Release / sync-version-file (push) Has been cancelled
CI / shell (push) Canceled after 0s
CI / test (push) Canceled after 0s
CI / frontend (push) Canceled after 0s
CI / golangci-lint (push) Canceled after 0s
Security Scan / backend-security (push) Canceled after 0s
Security Scan / frontend-security (push) Canceled after 0s
Sub2API v1.0 - AI API 网关(二开初始版本,基于上游 Wei-Shaw/sub2api)
2026-08-21 18:30:13 +08:00

528 lines
17 KiB
Go
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
//go:build unit
package service
import (
"context"
"encoding/json"
"errors"
"net/http"
"testing"
"time"
"github.com/stretchr/testify/require"
)
type modelNotFoundRateLimitCall struct {
accountID int64
scope string
resetAt time.Time
reason string
}
type modelNotFoundAccountRepoStub struct {
mockAccountRepoForGemini
tempCalls int
modelRateLimitCalls []modelNotFoundRateLimitCall
modelRateLimitErr error
}
func (r *modelNotFoundAccountRepoStub) SetTempUnschedulable(ctx context.Context, id int64, until time.Time, reason string) error {
r.tempCalls++
return nil
}
func (r *modelNotFoundAccountRepoStub) SetModelRateLimit(ctx context.Context, id int64, scope string, resetAt time.Time, reason ...string) error {
call := modelNotFoundRateLimitCall{
accountID: id,
scope: scope,
resetAt: resetAt,
}
if len(reason) > 0 {
call.reason = reason[0]
}
r.modelRateLimitCalls = append(r.modelRateLimitCalls, call)
return r.modelRateLimitErr
}
func TestRateLimitService_HandleUpstreamError_ModelNotFoundUsesModelRateLimit(t *testing.T) {
repo := &modelNotFoundAccountRepoStub{}
svc := &RateLimitService{accountRepo: repo}
account := openAIModelNotFoundTempAccount()
handled := svc.HandleUpstreamError(
context.Background(),
account,
http.StatusNotFound,
http.Header{},
[]byte(`{"error":{"code":"model_not_found","message":"model not found"}}`),
"gpt-5.4",
)
require.True(t, handled)
require.Zero(t, repo.tempCalls)
require.Len(t, repo.modelRateLimitCalls, 1)
call := repo.modelRateLimitCalls[0]
require.Equal(t, account.ID, call.accountID)
require.Equal(t, "gpt-5.4", call.scope)
require.Equal(t, upstreamModelNotFoundReason, call.reason)
require.WithinDuration(t, time.Now().Add(upstreamModelNotFoundCooldown), call.resetAt, 5*time.Second)
}
func TestRateLimitService_HandleUpstreamError_ModelNotFoundWriteFailureDoesNotTempUnschedule(t *testing.T) {
repo := &modelNotFoundAccountRepoStub{modelRateLimitErr: errors.New("write failed")}
svc := &RateLimitService{accountRepo: repo}
account := openAIModelNotFoundTempAccount()
handled := svc.HandleUpstreamError(
context.Background(),
account,
http.StatusNotFound,
http.Header{},
[]byte(`{"error":{"code":"model_not_found","message":"model not found"}}`),
"gpt-5.4",
)
require.True(t, handled)
require.Zero(t, repo.tempCalls)
require.Len(t, repo.modelRateLimitCalls, 1)
}
func TestRateLimitService_HandleUpstreamError_Bare404UsesModelScopedTempUnschedulableWhenModelKnown(t *testing.T) {
repo := &modelNotFoundAccountRepoStub{}
svc := &RateLimitService{accountRepo: repo}
account := openAIModelNotFoundTempAccount()
handled := svc.HandleUpstreamError(
context.Background(),
account,
http.StatusNotFound,
http.Header{},
[]byte(`{"error":{"message":"endpoint not found"}}`),
"gpt-5.4",
)
require.True(t, handled)
require.Zero(t, repo.tempCalls)
require.Len(t, repo.modelRateLimitCalls, 1)
call := repo.modelRateLimitCalls[0]
require.Equal(t, account.ID, call.accountID)
require.Equal(t, "gpt-5.4", call.scope)
require.WithinDuration(t, time.Now().Add(10*time.Minute), call.resetAt, 5*time.Second)
var state TempUnschedState
require.NoError(t, json.Unmarshal([]byte(call.reason), &state))
require.Equal(t, http.StatusNotFound, state.StatusCode)
require.Equal(t, "not found", state.MatchedKeyword)
}
func TestRateLimitService_HandleUpstreamError_Bare404WithoutModelKeepsAccountTempUnschedulable(t *testing.T) {
repo := &modelNotFoundAccountRepoStub{}
svc := &RateLimitService{accountRepo: repo}
account := openAIModelNotFoundTempAccount()
handled := svc.HandleUpstreamError(
context.Background(),
account,
http.StatusNotFound,
http.Header{},
[]byte(`{"error":{"message":"endpoint not found"}}`),
)
require.True(t, handled)
require.Equal(t, 1, repo.tempCalls)
require.Empty(t, repo.modelRateLimitCalls)
}
func TestRateLimitService_HandleUpstreamError_ModelTempWriteFailureNeverWidensToAccount(t *testing.T) {
repo := &modelNotFoundAccountRepoStub{modelRateLimitErr: errors.New("write failed")}
svc := &RateLimitService{accountRepo: repo}
account := openAIModelNotFoundTempAccount()
handled := svc.HandleUpstreamError(
context.Background(),
account,
http.StatusNotFound,
http.Header{},
[]byte(`{"error":{"message":"endpoint not found"}}`),
"gpt-5.4",
)
require.True(t, handled)
require.Zero(t, repo.tempCalls)
require.Len(t, repo.modelRateLimitCalls, 1)
}
func TestRateLimitService_HandleTempUnschedulable_PoolModeWithoutCustomPolicySkipsState(t *testing.T) {
repo := &modelNotFoundAccountRepoStub{}
svc := &RateLimitService{accountRepo: repo}
account := openAIModelNotFoundTempAccount()
account.Credentials["pool_mode"] = true
handled := svc.HandleTempUnschedulable(
context.Background(),
account,
http.StatusNotFound,
[]byte(`{"error":{"message":"endpoint not found"}}`),
"gpt-5.4",
)
require.False(t, handled)
require.Zero(t, repo.tempCalls)
require.Empty(t, repo.modelRateLimitCalls)
}
func TestRateLimitService_HandleTempUnschedulable_PoolModeCustomPolicyUsesModelScope(t *testing.T) {
repo := &modelNotFoundAccountRepoStub{}
svc := &RateLimitService{accountRepo: repo}
account := openAIModelNotFoundTempAccount()
account.Credentials["pool_mode"] = true
account.Credentials["custom_error_codes_enabled"] = true
account.Credentials["custom_error_codes"] = []any{float64(http.StatusNotFound)}
handled := svc.HandleTempUnschedulable(
context.Background(),
account,
http.StatusNotFound,
[]byte(`{"error":{"message":"endpoint not found"}}`),
"gpt-5.4",
)
require.True(t, handled)
require.Zero(t, repo.tempCalls)
require.Len(t, repo.modelRateLimitCalls, 1)
require.Equal(t, "gpt-5.4", repo.modelRateLimitCalls[0].scope)
}
func TestRateLimitService_TempUnschedulableContextPreservesModelForPoolDependency(t *testing.T) {
repo := &modelNotFoundAccountRepoStub{}
svc := &RateLimitService{accountRepo: repo}
account := openAIModelNotFoundTempAccount()
account.Credentials["pool_mode"] = true
ctx := withTempUnschedulableModel(context.Background(), []string{"gpt-5.4"})
// #4496 calls tryTempUnschedulable from the pool-mode branch without an
// explicit model argument. The request context must preserve the canonical
// model so that combined behavior remains model-scoped after it lands.
handled := svc.tryTempUnschedulable(
ctx,
account,
http.StatusNotFound,
[]byte(`{"error":{"message":"endpoint not found"}}`),
)
require.True(t, handled)
require.Zero(t, repo.tempCalls)
require.Len(t, repo.modelRateLimitCalls, 1)
require.Equal(t, "gpt-5.4", repo.modelRateLimitCalls[0].scope)
}
func TestRateLimitService_HandleUpstreamError_CustomPolicyExclusionSkipsAllState(t *testing.T) {
repo := &modelNotFoundAccountRepoStub{}
svc := &RateLimitService{accountRepo: repo}
account := openAIModelNotFoundTempAccount()
account.Credentials["custom_error_codes_enabled"] = true
account.Credentials["custom_error_codes"] = []any{float64(http.StatusServiceUnavailable)}
handled := svc.HandleUpstreamError(
context.Background(),
account,
http.StatusNotFound,
http.Header{},
[]byte(`{"error":{"message":"endpoint not found"}}`),
"gpt-5.4",
)
require.False(t, handled)
require.Zero(t, repo.tempCalls)
require.Empty(t, repo.modelRateLimitCalls)
}
func TestRateLimitService_HandleTempUnschedulable_AuthenticationFailureStaysAccountScoped(t *testing.T) {
repo := &modelNotFoundAccountRepoStub{}
svc := &RateLimitService{accountRepo: repo}
account := openAIModelNotFoundTempAccount()
account.TempUnschedulableReason = "legacy non-JSON reason"
account.Credentials["temp_unschedulable_rules"] = []any{
map[string]any{
"error_code": float64(http.StatusUnauthorized),
"keywords": []any{"unauthorized"},
"duration_minutes": float64(10),
},
}
handled := svc.HandleTempUnschedulable(
context.Background(),
account,
http.StatusUnauthorized,
[]byte(`{"error":{"message":"unauthorized"}}`),
"gpt-5.4",
)
require.True(t, handled)
require.Equal(t, 1, repo.tempCalls)
require.Empty(t, repo.modelRateLimitCalls)
}
func TestRateLimitService_ModelTempUnschedulableIsolatesSchedulerByModel(t *testing.T) {
repo := &modelNotFoundAccountRepoStub{}
svc := &RateLimitService{accountRepo: repo}
account := openAIModelNotFoundTempAccount()
account.Credentials["model_mapping"] = map[string]any{
"public-a": "upstream-a",
"upstream-a": "upstream-b",
}
handled := svc.HandleUpstreamError(
context.Background(),
account,
http.StatusNotFound,
http.Header{},
[]byte(`{"error":{"message":"endpoint not found"}}`),
"upstream-a",
)
require.True(t, handled)
require.Len(t, repo.modelRateLimitCalls, 1)
call := repo.modelRateLimitCalls[0]
require.Equal(t, "upstream-a", call.scope, "canonical upstream model must not be mapped a second time")
account.Extra = map[string]any{
modelRateLimitsKey: map[string]any{
call.scope: map[string]any{
"rate_limit_reset_at": call.resetAt.UTC().Format(time.RFC3339),
},
},
}
require.False(t, account.IsSchedulableForModelWithContext(context.Background(), "public-a"))
require.True(t, account.IsSchedulableForModelWithContext(context.Background(), "gpt-5.6-sol"))
}
func openAIModelNotFoundTempAccount() *Account {
return &Account{
ID: 101,
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Status: StatusActive,
Schedulable: true,
Credentials: map[string]any{
"temp_unschedulable_enabled": true,
"temp_unschedulable_rules": []any{
map[string]any{
"error_code": float64(http.StatusNotFound),
"keywords": []any{"not found"},
"duration_minutes": float64(10),
},
},
},
}
}
func TestRateLimitService_HandleUpstreamError_CodexPlanGatedModelUsesModelRateLimit(t *testing.T) {
repo := &modelNotFoundAccountRepoStub{}
svc := &RateLimitService{accountRepo: repo}
account := openAICodexPlanGatedOAuthAccount()
handled := svc.HandleUpstreamError(
context.Background(),
account,
http.StatusBadRequest,
http.Header{},
[]byte(`{"detail":"The 'gpt-5.6-sol' model is not supported when using Codex with a ChatGPT account."}`),
"gpt-5.6-sol",
)
require.True(t, handled)
require.Zero(t, repo.tempCalls)
require.Len(t, repo.modelRateLimitCalls, 1)
call := repo.modelRateLimitCalls[0]
require.Equal(t, account.ID, call.accountID)
require.Equal(t, "gpt-5.6-sol", call.scope)
require.Equal(t, upstreamCodexPlanGatedModelReason, call.reason)
require.WithinDuration(t, time.Now().Add(upstreamCodexPlanGatedModelCooldown), call.resetAt, 5*time.Second)
}
func TestRateLimitService_HandleUpstreamError_CodexPlanGatedModelRespectsModelMapping(t *testing.T) {
repo := &modelNotFoundAccountRepoStub{}
svc := &RateLimitService{accountRepo: repo}
account := openAICodexPlanGatedOAuthAccount()
account.Credentials["model_mapping"] = map[string]any{"gpt-5.6-sol": "gpt-5.6-sol-upstream"}
handled := svc.HandleUpstreamError(
context.Background(),
account,
http.StatusBadRequest,
http.Header{},
[]byte(`{"detail":"The 'gpt-5.6-sol-upstream' model is not supported when using Codex with a ChatGPT account."}`),
"gpt-5.6-sol",
)
require.True(t, handled)
require.Len(t, repo.modelRateLimitCalls, 1)
require.Equal(t, "gpt-5.6-sol-upstream", repo.modelRateLimitCalls[0].scope)
}
func TestRateLimitService_HandleUpstreamError_CodexPlanGatedModelIgnoresAPIKeyAccount(t *testing.T) {
repo := &modelNotFoundAccountRepoStub{}
svc := &RateLimitService{accountRepo: repo}
account := openAICodexPlanGatedOAuthAccount()
account.Type = AccountTypeAPIKey
handled := svc.HandleUpstreamError(
context.Background(),
account,
http.StatusBadRequest,
http.Header{},
[]byte(`{"detail":"The 'gpt-5.6-sol' model is not supported when using Codex with a ChatGPT account."}`),
"gpt-5.6-sol",
)
require.False(t, handled)
require.Empty(t, repo.modelRateLimitCalls)
}
func TestRateLimitService_HandleUpstreamError_CodexPlanGatedImageModelSkipsCooldown(t *testing.T) {
for _, model := range []string{"gpt-image-1", "gpt-image-1.5", "gpt-image-2"} {
t.Run(model, func(t *testing.T) {
repo := &modelNotFoundAccountRepoStub{}
svc := &RateLimitService{accountRepo: repo}
account := openAICodexPlanGatedOAuthAccount()
handled := svc.HandleUpstreamError(
context.Background(),
account,
http.StatusBadRequest,
http.Header{},
[]byte(`{"detail":"The '`+model+`' model is not supported when using Codex with a ChatGPT account."}`),
model,
)
require.True(t, handled, "attempt should still fail over")
require.Empty(t, repo.modelRateLimitCalls,
"image models must not be cooled down: the account still serves them over /v1/images/*")
require.Zero(t, repo.tempCalls)
})
}
}
func TestRateLimitService_HandleUpstreamError_CodexPlanGatedTextModelStillCoolsDown(t *testing.T) {
repo := &modelNotFoundAccountRepoStub{}
svc := &RateLimitService{accountRepo: repo}
account := openAICodexPlanGatedOAuthAccount()
handled := svc.HandleUpstreamError(
context.Background(),
account,
http.StatusBadRequest,
http.Header{},
[]byte(`{"detail":"The 'gpt-5.6-sol' model is not supported when using Codex with a ChatGPT account."}`),
"gpt-5.6-sol",
)
require.True(t, handled)
require.Len(t, repo.modelRateLimitCalls, 1, "non-image plan-gated models keep the existing cooldown")
require.Equal(t, upstreamCodexPlanGatedModelReason, repo.modelRateLimitCalls[0].reason)
}
func openAICodexPlanGatedOAuthAccount() *Account {
return &Account{
ID: 202,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Status: StatusActive,
Schedulable: true,
Credentials: map[string]any{},
}
}
// 请求本身就走 /v1/images/* 时必须保留冷却。
//
// OAuth 账号的 /v1/images/* 上游同样是 Codex Responsesopenai_images_responses.go
// → handleOpenAIImagesErrorResponse → handleOpenAIAccountUpstreamError →
// HandleUpstreamModelNotFound),所以这条路径也会命中 plan-gated 分支。账号确实
// 不具备生图能力时,冷却是唯一的刹车:调度层靠 model_rate_limits 跳过该账号后
// 快速 503;一旦跳过冷却,每个请求都会完整走一遍号池,对上游形成无上界的 400 放大。
func TestRateLimitService_HandleUpstreamError_CodexPlanGatedImageModelKeepsCooldownOnImagesEndpoint(t *testing.T) {
repo := &modelNotFoundAccountRepoStub{}
svc := &RateLimitService{accountRepo: repo}
account := openAICodexPlanGatedOAuthAccount()
handled := svc.HandleUpstreamError(
WithOpenAIImagesEndpoint(context.Background()),
account,
http.StatusBadRequest,
http.Header{},
[]byte(`{"detail":"The 'gpt-image-2' model is not supported when using Codex with a ChatGPT account."}`),
"gpt-image-2",
)
require.True(t, handled)
require.Len(t, repo.modelRateLimitCalls, 1,
"/v1/images/* 上的 plan-gated 拒绝是真实的能力缺失,必须保留冷却刹车")
require.Equal(t, "gpt-image-2", repo.modelRateLimitCalls[0].scope)
require.Equal(t, upstreamCodexPlanGatedModelReason, repo.modelRateLimitCalls[0].reason)
}
// 仅 WithOpenAIImageGenerationIntent/v1/responses 因模型名自动置位)不算专用生图
// 端点,仍按"用错端点"处理。
func TestRateLimitService_HandleUpstreamError_CodexPlanGatedImageModelSkipsCooldownOnIntentOnly(t *testing.T) {
repo := &modelNotFoundAccountRepoStub{}
svc := &RateLimitService{accountRepo: repo}
account := openAICodexPlanGatedOAuthAccount()
handled := svc.HandleUpstreamError(
WithOpenAIImageGenerationIntent(context.Background()),
account,
http.StatusBadRequest,
http.Header{},
[]byte(`{"detail":"The 'gpt-image-2' model is not supported when using Codex with a ChatGPT account."}`),
"gpt-image-2",
)
require.True(t, handled)
require.Empty(t, repo.modelRateLimitCalls)
}
// 守卫口径必须与冷却键一致:冷却键走 account.GetMappedModel,账号可以把文本别名
// 映射到 gpt-image-*,只判请求模型会漏掉这种形态,原 bug 原样复现。
func TestRateLimitService_HandleUpstreamError_CodexPlanGatedImageModelSkipsCooldownViaModelMapping(t *testing.T) {
repo := &modelNotFoundAccountRepoStub{}
svc := &RateLimitService{accountRepo: repo}
account := openAICodexPlanGatedOAuthAccount()
account.Credentials["model_mapping"] = map[string]any{"my-draw-alias": "gpt-image-2"}
handled := svc.HandleUpstreamError(
context.Background(),
account,
http.StatusBadRequest,
http.Header{},
[]byte(`{"detail":"The 'gpt-image-2' model is not supported when using Codex with a ChatGPT account."}`),
"my-draw-alias",
)
require.True(t, handled)
require.Empty(t, repo.modelRateLimitCalls,
"映射后的上游模型是图片模型,冷却键会写到 gpt-image-2 上,守卫必须一并识别")
}
// 404 model-not-found 分支不受守卫影响:即使是图片模型也照常冷却。
func TestRateLimitService_HandleUpstreamError_ModelNotFoundImageModelStillCoolsDown(t *testing.T) {
repo := &modelNotFoundAccountRepoStub{}
svc := &RateLimitService{accountRepo: repo}
account := openAICodexPlanGatedOAuthAccount()
handled := svc.HandleUpstreamError(
context.Background(),
account,
http.StatusNotFound,
http.Header{},
[]byte(`{"error":{"message":"The model 'gpt-image-2' does not exist","code":"model_not_found"}}`),
"gpt-image-2",
)
require.True(t, handled)
require.Len(t, repo.modelRateLimitCalls, 1, "守卫只作用于 codex plan-gated 分支")
require.Equal(t, upstreamModelNotFoundReason, repo.modelRateLimitCalls[0].reason)
}