mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-15 23:22:59 +08:00
song
c0d5119375
Use shared model routing for image forwarding and account tests, retain Luna orchestration for legacy models, and add Image 2.5 pricing and usage accounting.
359 lines
15 KiB
Go
359 lines
15 KiB
Go
//go:build unit
|
|
|
|
package service
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/config"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestIsOpenAIImageRateLimitError(t *testing.T) {
|
|
imageBody := []byte(`{"error":{"message":"Rate limit reached for gpt-image-2-codex (for limit gpt-image) in organization org on input-images per min: Limit 4000, Used 4000. Please try again in 467ms."}}`)
|
|
textBody := []byte(`{"error":{"message":"Rate limit reached for gpt-5.4 in organization org on tokens per min: Limit 30000, Used 30000. Please try again in 1s."}}`)
|
|
|
|
require.True(t, isOpenAIImageRateLimitError(http.StatusTooManyRequests, imageBody))
|
|
require.False(t, isOpenAIImageRateLimitError(http.StatusTooManyRequests, textBody))
|
|
require.False(t, isOpenAIImageRateLimitError(http.StatusBadRequest, imageBody))
|
|
}
|
|
|
|
func TestRateLimitService_HandleOpenAIImageRateLimit_ParsesTryAgainCooldown(t *testing.T) {
|
|
repo := &modelNotFoundAccountRepoStub{}
|
|
svc := &RateLimitService{accountRepo: repo}
|
|
account := &Account{ID: 201, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
|
|
body := []byte(`{"error":{"type":"rate_limit_exceeded","message":"Rate limit reached for gpt-image-2-codex (for limit gpt-image) on input-images per min. Please try again in 2s."}}`)
|
|
|
|
before := time.Now()
|
|
handled := svc.HandleOpenAIImageRateLimit(context.Background(), account, http.StatusTooManyRequests, http.Header{}, body)
|
|
|
|
require.True(t, handled)
|
|
require.Len(t, repo.modelRateLimitCalls, 1)
|
|
call := repo.modelRateLimitCalls[0]
|
|
require.Equal(t, account.ID, call.accountID)
|
|
require.Equal(t, openAIImageGenerationRateLimitKey, call.scope)
|
|
require.Equal(t, openAIImageRateLimitReason, call.reason)
|
|
require.WithinDuration(t, before.Add(2*time.Second), call.resetAt, time.Second)
|
|
}
|
|
|
|
func TestRateLimitService_HandleOpenAIImageRateLimit_DefaultsToOneMinute(t *testing.T) {
|
|
repo := &modelNotFoundAccountRepoStub{}
|
|
svc := &RateLimitService{accountRepo: repo}
|
|
account := &Account{ID: 202, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
|
|
body := []byte(`{"error":{"type":"rate_limit_exceeded","message":"Rate limit reached for gpt-image-2-codex (for limit gpt-image) on input-images per min."}}`)
|
|
|
|
before := time.Now()
|
|
handled := svc.HandleOpenAIImageRateLimit(context.Background(), account, http.StatusTooManyRequests, http.Header{}, body)
|
|
|
|
require.True(t, handled)
|
|
require.Len(t, repo.modelRateLimitCalls, 1)
|
|
call := repo.modelRateLimitCalls[0]
|
|
require.Equal(t, openAIImageGenerationRateLimitKey, call.scope)
|
|
require.Equal(t, openAIImageRateLimitReason, call.reason)
|
|
require.WithinDuration(t, before.Add(time.Minute), call.resetAt, time.Second)
|
|
}
|
|
|
|
func TestOpenAIGatewayService_HandleOpenAIAccountUpstreamError_ImageRateLimitDoesNotBlockWholeAccount(t *testing.T) {
|
|
repo := &modelNotFoundAccountRepoStub{}
|
|
svc := &OpenAIGatewayService{rateLimitService: &RateLimitService{accountRepo: repo}}
|
|
account := &Account{ID: 203, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
|
|
body := []byte(`{"error":{"type":"rate_limit_exceeded","message":"Rate limit reached for gpt-image-2-codex (for limit gpt-image) on input-images per min. Please try again in 1s."}}`)
|
|
|
|
disabled := svc.handleOpenAIAccountUpstreamError(context.Background(), account, http.StatusTooManyRequests, http.Header{}, body, "gpt-image-2")
|
|
|
|
require.False(t, disabled)
|
|
require.Len(t, repo.modelRateLimitCalls, 1)
|
|
require.Equal(t, openAIImageGenerationRateLimitKey, repo.modelRateLimitCalls[0].scope)
|
|
_, wholeAccountBlocked := svc.openaiAccountRuntimeBlockUntil.Load(account.ID)
|
|
require.False(t, wholeAccountBlocked)
|
|
}
|
|
|
|
func TestOpenAIGatewayServiceForwardImages_ImageRateLimitReturnsFailoverAndCoolsCapability(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
repo := &modelNotFoundAccountRepoStub{}
|
|
body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat"}`)
|
|
errorBody := `{"error":{"type":"rate_limit_exceeded","message":"Rate limit reached for gpt-image-2-codex (for limit gpt-image) in organization org on input-images per min: Limit 4000, Used 4000. Please try again in 1s."}}`
|
|
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body))
|
|
req.Header.Set("Content-Type", "application/json")
|
|
rec := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(rec)
|
|
c.Request = req
|
|
|
|
svc := &OpenAIGatewayService{
|
|
rateLimitService: &RateLimitService{accountRepo: repo},
|
|
httpUpstream: &httpUpstreamRecorder{
|
|
resp: &http.Response{
|
|
StatusCode: http.StatusTooManyRequests,
|
|
Header: http.Header{"X-Request-Id": []string{"req_img_rate_limited"}},
|
|
Body: io.NopCloser(strings.NewReader(errorBody)),
|
|
},
|
|
},
|
|
}
|
|
parsed, err := svc.ParseOpenAIImagesRequest(c, body)
|
|
require.NoError(t, err)
|
|
account := &Account{
|
|
ID: 204,
|
|
Name: "openai-oauth",
|
|
Platform: PlatformOpenAI,
|
|
Type: AccountTypeOAuth,
|
|
Credentials: map[string]any{
|
|
"access_token": "token-123",
|
|
},
|
|
}
|
|
|
|
result, err := svc.ForwardImages(context.Background(), c, account, body, parsed, "")
|
|
|
|
require.Nil(t, result)
|
|
var failoverErr *UpstreamFailoverError
|
|
require.ErrorAs(t, err, &failoverErr)
|
|
require.Equal(t, http.StatusTooManyRequests, failoverErr.StatusCode)
|
|
require.Contains(t, string(failoverErr.ResponseBody), "input-images per min")
|
|
require.Len(t, repo.modelRateLimitCalls, 1)
|
|
require.Equal(t, openAIImageGenerationRateLimitKey, repo.modelRateLimitCalls[0].scope)
|
|
}
|
|
|
|
// issue #6171：上游"回文字没回图"是**这一轮**的结果（模型选择了说话），不是账号能力
|
|
// 失效。它同时被判为可重试（502）并驱动 failover，若还写 30 分钟账号级冷却，一次闲聊
|
|
// 回复就会沿号池把每个被重试到的账号依次冷却掉。冷却仍保留给结构化上游证据，见
|
|
// TestOpenAIGatewayServiceForwardImages_StructuredUnavailableCoolsImageCapability。
|
|
func TestOpenAIGatewayServiceForwardImages_TextFallbackDoesNotCoolImageCapability(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
repo := &modelNotFoundAccountRepoStub{}
|
|
body := []byte(`{"model":"gpt-image-1","prompt":"draw a cat"}`)
|
|
upstreamSSE := "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"r\",\"status\":\"completed\",\"model\":\"gpt-5.4-mini\",\"output\":[{\"type\":\"message\",\"content\":[{\"type\":\"output_text\",\"text\":\"Here's a polished image prompt for your request.\"}]}]}}\n\n"
|
|
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body))
|
|
req.Header.Set("Content-Type", "application/json")
|
|
rec := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(rec)
|
|
c.Request = req
|
|
|
|
svc := &OpenAIGatewayService{
|
|
accountRepo: repo,
|
|
httpUpstream: &httpUpstreamRecorder{
|
|
resp: &http.Response{
|
|
StatusCode: http.StatusOK,
|
|
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
|
Body: io.NopCloser(strings.NewReader(upstreamSSE)),
|
|
},
|
|
},
|
|
}
|
|
parsed, err := svc.ParseOpenAIImagesRequest(c, body)
|
|
require.NoError(t, err)
|
|
account := &Account{
|
|
ID: 205,
|
|
Name: "openai-oauth",
|
|
Platform: PlatformOpenAI,
|
|
Type: AccountTypeOAuth,
|
|
Credentials: map[string]any{
|
|
"access_token": "token-123",
|
|
},
|
|
}
|
|
|
|
result, err := svc.ForwardImages(context.Background(), c, account, body, parsed, "")
|
|
|
|
require.Nil(t, result)
|
|
require.Error(t, err)
|
|
var failoverErr *UpstreamFailoverError
|
|
require.ErrorAs(t, err, &failoverErr)
|
|
require.False(t, failoverErr.RetryableOnSameAccount)
|
|
// 换号行为不变：该判据仍足以放弃本账号重试这一次请求……
|
|
require.Equal(t, http.StatusBadGateway, failoverErr.StatusCode)
|
|
// ……但不再写任何账号级状态，否则重试会把冷却一路刷到整个号池。
|
|
require.Empty(t, repo.modelRateLimitCalls,
|
|
"模型回文字只说明这一轮没出图，不构成账号 30 分钟不可用的证据")
|
|
}
|
|
|
|
// 对照不变式：上游 error 帧点名 image_generation_unavailable 时仍写冷却，
|
|
// 保证 #6171 的修复没有把这项能力保护整个废掉。
|
|
func TestOpenAIGatewayServiceForwardImages_StructuredUnavailableCoolsImageCapability(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
repo := &modelNotFoundAccountRepoStub{}
|
|
body := []byte(`{"model":"gpt-image-1","prompt":"draw a cat"}`)
|
|
upstreamSSE := "data: {\"type\":\"response.failed\",\"response\":{\"id\":\"r\",\"error\":" +
|
|
"{\"type\":\"upstream_error\",\"code\":\"image_generation_unavailable\"," +
|
|
"\"message\":\"image generation tool is not available for this account\"}}}\n\n"
|
|
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body))
|
|
req.Header.Set("Content-Type", "application/json")
|
|
rec := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(rec)
|
|
c.Request = req
|
|
|
|
svc := &OpenAIGatewayService{
|
|
accountRepo: repo,
|
|
httpUpstream: &httpUpstreamRecorder{
|
|
resp: &http.Response{
|
|
StatusCode: http.StatusOK,
|
|
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
|
Body: io.NopCloser(strings.NewReader(upstreamSSE)),
|
|
},
|
|
},
|
|
}
|
|
parsed, err := svc.ParseOpenAIImagesRequest(c, body)
|
|
require.NoError(t, err)
|
|
account := &Account{
|
|
ID: 206,
|
|
Name: "openai-oauth",
|
|
Platform: PlatformOpenAI,
|
|
Type: AccountTypeOAuth,
|
|
Credentials: map[string]any{
|
|
"access_token": "token-123",
|
|
},
|
|
}
|
|
|
|
before := time.Now()
|
|
result, err := svc.ForwardImages(context.Background(), c, account, body, parsed, "")
|
|
|
|
require.Nil(t, result)
|
|
require.Error(t, err)
|
|
require.Len(t, repo.modelRateLimitCalls, 1)
|
|
call := repo.modelRateLimitCalls[0]
|
|
require.Equal(t, account.ID, call.accountID)
|
|
require.Equal(t, openAIImageGenerationRateLimitKey, call.scope)
|
|
require.Equal(t, openAIImagesOAuthUnavailableReason, call.reason)
|
|
require.WithinDuration(t, before.Add(openAIImagesOAuthUnavailableDefaultCooldown), call.resetAt, time.Second)
|
|
}
|
|
|
|
func TestOpenAIGatewayService_CoolOpenAIImagesOAuthToolUsesConfiguredCooldown(t *testing.T) {
|
|
accountRepo := &modelNotFoundAccountRepoStub{}
|
|
settingRepo := newMockSettingRepo()
|
|
settingRepo.data[SettingKeyOpenAIImagesOAuthUnavailableCooldownSettings] = `{"cooldown_minutes":7}`
|
|
svc := &OpenAIGatewayService{
|
|
accountRepo: accountRepo,
|
|
settingService: NewSettingService(settingRepo, &config.Config{}),
|
|
}
|
|
|
|
before := time.Now()
|
|
svc.coolOpenAIImagesOAuthTool(context.Background(), &Account{ID: 206, Platform: PlatformOpenAI, Type: AccountTypeOAuth})
|
|
|
|
require.Len(t, accountRepo.modelRateLimitCalls, 1)
|
|
require.WithinDuration(t, before.Add(7*time.Minute), accountRepo.modelRateLimitCalls[0].resetAt, time.Second)
|
|
}
|
|
|
|
func TestOpenAIGatewayServiceForwardImages_CapabilityLossCoolsImageScope(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
repo := &modelNotFoundAccountRepoStub{}
|
|
body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat"}`)
|
|
errorBody := `{"error":{"message":"Tool choice 'image_generation' not found in 'tools' parameter.","param":"tool_choice","type":"invalid_request_error"}}`
|
|
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body))
|
|
req.Header.Set("Content-Type", "application/json")
|
|
rec := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(rec)
|
|
c.Request = req
|
|
|
|
svc := &OpenAIGatewayService{
|
|
rateLimitService: &RateLimitService{accountRepo: repo},
|
|
httpUpstream: &httpUpstreamRecorder{
|
|
resp: &http.Response{
|
|
StatusCode: http.StatusBadRequest,
|
|
Header: http.Header{"X-Request-Id": []string{"req_img_capability_lost"}},
|
|
Body: io.NopCloser(strings.NewReader(errorBody)),
|
|
},
|
|
},
|
|
}
|
|
parsed, err := svc.ParseOpenAIImagesRequest(c, body)
|
|
require.NoError(t, err)
|
|
account := &Account{
|
|
ID: 205,
|
|
Name: "openai-oauth",
|
|
Platform: PlatformOpenAI,
|
|
Type: AccountTypeOAuth,
|
|
Credentials: map[string]any{
|
|
"access_token": "token-123",
|
|
},
|
|
}
|
|
|
|
before := time.Now()
|
|
result, err := svc.ForwardImages(context.Background(), c, account, body, parsed, "")
|
|
|
|
require.Nil(t, result)
|
|
require.Error(t, err)
|
|
require.Len(t, repo.modelRateLimitCalls, 1)
|
|
call := repo.modelRateLimitCalls[0]
|
|
require.Equal(t, account.ID, call.accountID)
|
|
require.Equal(t, openAIImageGenerationRateLimitKey, call.scope)
|
|
require.Equal(t, openAIImageCapabilityLossReason, call.reason)
|
|
require.WithinDuration(t, before.Add(openAIImageCapabilityLossCooldown), call.resetAt, time.Second)
|
|
}
|
|
|
|
func TestOpenAIGatewayServiceHandleUpstreamError_PassthroughCapabilityLossDoesNotCool(t *testing.T) {
|
|
repo := &modelNotFoundAccountRepoStub{}
|
|
svc := &OpenAIGatewayService{rateLimitService: &RateLimitService{accountRepo: repo}}
|
|
account := &Account{ID: 206, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
|
|
body := []byte(`{"error":{"message":"Tool choice 'image_generation' not found in 'tools' parameter.","param":"tool_choice","type":"invalid_request_error"}}`)
|
|
|
|
disabled := svc.handleOpenAIAccountUpstreamError(context.Background(), account, http.StatusBadRequest, http.Header{}, body, "gpt-5.5")
|
|
|
|
require.False(t, disabled)
|
|
require.Empty(t, repo.modelRateLimitCalls)
|
|
_, wholeAccountBlocked := svc.openaiAccountRuntimeBlockUntil.Load(account.ID)
|
|
require.False(t, wholeAccountBlocked)
|
|
}
|
|
|
|
func TestRateLimitServiceHandleOpenAIImageCapabilityLoss_IgnoresGenericBadRequest(t *testing.T) {
|
|
repo := &modelNotFoundAccountRepoStub{}
|
|
svc := &RateLimitService{accountRepo: repo}
|
|
account := &Account{ID: 207, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
|
|
body := []byte(`{"error":{"message":"Invalid type for input[0].arguments"}}`)
|
|
|
|
handled := svc.HandleOpenAIImageCapabilityLoss(context.Background(), account, http.StatusBadRequest, body)
|
|
|
|
require.False(t, handled)
|
|
require.Empty(t, repo.modelRateLimitCalls)
|
|
}
|
|
|
|
func TestRateLimitServiceHandleOpenAIImageCapabilityLoss_RespectsPlatformAndErrorCodePolicy(t *testing.T) {
|
|
body := []byte(`{"error":{"message":"Tool choice 'image_generation' not found in 'tools' parameter.","param":"tool_choice","type":"invalid_request_error"}}`)
|
|
|
|
t.Run("non_openai_platform", func(t *testing.T) {
|
|
repo := &modelNotFoundAccountRepoStub{}
|
|
svc := &RateLimitService{accountRepo: repo}
|
|
account := &Account{ID: 208, Platform: PlatformAnthropic, Type: AccountTypeOAuth}
|
|
|
|
handled := svc.HandleOpenAIImageCapabilityLoss(context.Background(), account, http.StatusBadRequest, body)
|
|
|
|
require.False(t, handled)
|
|
require.Empty(t, repo.modelRateLimitCalls)
|
|
})
|
|
|
|
t.Run("custom_error_code_policy_excludes_400", func(t *testing.T) {
|
|
repo := &modelNotFoundAccountRepoStub{}
|
|
svc := &RateLimitService{accountRepo: repo}
|
|
account := &Account{
|
|
ID: 209,
|
|
Platform: PlatformOpenAI,
|
|
Type: AccountTypeAPIKey,
|
|
Credentials: map[string]any{
|
|
"custom_error_codes_enabled": true,
|
|
"custom_error_codes": []any{float64(http.StatusTooManyRequests)},
|
|
},
|
|
}
|
|
|
|
require.False(t, account.ShouldHandleErrorCode(http.StatusBadRequest))
|
|
handled := svc.HandleOpenAIImageCapabilityLoss(context.Background(), account, http.StatusBadRequest, body)
|
|
|
|
require.False(t, handled)
|
|
require.Empty(t, repo.modelRateLimitCalls)
|
|
})
|
|
}
|
|
|
|
func TestIsOpenAIImageCapabilityLossError(t *testing.T) {
|
|
capabilityLossBody := []byte(`{"error":{"message":"Tool choice 'image_generation' not found in 'tools' parameter.","param":"tool_choice","type":"invalid_request_error"}}`)
|
|
genericBadRequestBody := []byte(`{"error":{"message":"Invalid type for input[0].arguments"}}`)
|
|
|
|
require.True(t, isOpenAIImageCapabilityLossError(http.StatusBadRequest, capabilityLossBody))
|
|
require.False(t, isOpenAIImageCapabilityLossError(http.StatusBadRequest, genericBadRequestBody))
|
|
require.False(t, isOpenAIImageCapabilityLossError(http.StatusTooManyRequests, capabilityLossBody))
|
|
}
|