mirror of
https://gitee.com/wanwujie/sub2api
synced 2026-04-25 17:14:45 +08:00
Merge branch 'main' into test-sora
This commit is contained in:
@@ -425,6 +425,22 @@ func (a *Account) GetBaseURL() string {
|
||||
if baseURL == "" {
|
||||
return "https://api.anthropic.com"
|
||||
}
|
||||
if a.Platform == PlatformAntigravity {
|
||||
return strings.TrimRight(baseURL, "/") + "/antigravity"
|
||||
}
|
||||
return baseURL
|
||||
}
|
||||
|
||||
// GetGeminiBaseURL 返回 Gemini 兼容端点的 base URL。
|
||||
// Antigravity 平台的 APIKey 账号自动拼接 /antigravity。
|
||||
func (a *Account) GetGeminiBaseURL(defaultBaseURL string) string {
|
||||
baseURL := strings.TrimSpace(a.GetCredential("base_url"))
|
||||
if baseURL == "" {
|
||||
return defaultBaseURL
|
||||
}
|
||||
if a.Platform == PlatformAntigravity && a.Type == AccountTypeAPIKey {
|
||||
return strings.TrimRight(baseURL, "/") + "/antigravity"
|
||||
}
|
||||
return baseURL
|
||||
}
|
||||
|
||||
|
||||
160
backend/internal/service/account_base_url_test.go
Normal file
160
backend/internal/service/account_base_url_test.go
Normal file
@@ -0,0 +1,160 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestGetBaseURL(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
account Account
|
||||
expected string
|
||||
}{
|
||||
{
|
||||
name: "non-apikey type returns empty",
|
||||
account: Account{
|
||||
Type: AccountTypeOAuth,
|
||||
Platform: PlatformAnthropic,
|
||||
},
|
||||
expected: "",
|
||||
},
|
||||
{
|
||||
name: "apikey without base_url returns default anthropic",
|
||||
account: Account{
|
||||
Type: AccountTypeAPIKey,
|
||||
Platform: PlatformAnthropic,
|
||||
Credentials: map[string]any{},
|
||||
},
|
||||
expected: "https://api.anthropic.com",
|
||||
},
|
||||
{
|
||||
name: "apikey with custom base_url",
|
||||
account: Account{
|
||||
Type: AccountTypeAPIKey,
|
||||
Platform: PlatformAnthropic,
|
||||
Credentials: map[string]any{"base_url": "https://custom.example.com"},
|
||||
},
|
||||
expected: "https://custom.example.com",
|
||||
},
|
||||
{
|
||||
name: "antigravity apikey auto-appends /antigravity",
|
||||
account: Account{
|
||||
Type: AccountTypeAPIKey,
|
||||
Platform: PlatformAntigravity,
|
||||
Credentials: map[string]any{"base_url": "https://upstream.example.com"},
|
||||
},
|
||||
expected: "https://upstream.example.com/antigravity",
|
||||
},
|
||||
{
|
||||
name: "antigravity apikey trims trailing slash before appending",
|
||||
account: Account{
|
||||
Type: AccountTypeAPIKey,
|
||||
Platform: PlatformAntigravity,
|
||||
Credentials: map[string]any{"base_url": "https://upstream.example.com/"},
|
||||
},
|
||||
expected: "https://upstream.example.com/antigravity",
|
||||
},
|
||||
{
|
||||
name: "antigravity non-apikey returns empty",
|
||||
account: Account{
|
||||
Type: AccountTypeOAuth,
|
||||
Platform: PlatformAntigravity,
|
||||
Credentials: map[string]any{"base_url": "https://upstream.example.com"},
|
||||
},
|
||||
expected: "",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := tt.account.GetBaseURL()
|
||||
if result != tt.expected {
|
||||
t.Errorf("GetBaseURL() = %q, want %q", result, tt.expected)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetGeminiBaseURL(t *testing.T) {
|
||||
const defaultGeminiURL = "https://generativelanguage.googleapis.com"
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
account Account
|
||||
expected string
|
||||
}{
|
||||
{
|
||||
name: "apikey without base_url returns default",
|
||||
account: Account{
|
||||
Type: AccountTypeAPIKey,
|
||||
Platform: PlatformGemini,
|
||||
Credentials: map[string]any{},
|
||||
},
|
||||
expected: defaultGeminiURL,
|
||||
},
|
||||
{
|
||||
name: "apikey with custom base_url",
|
||||
account: Account{
|
||||
Type: AccountTypeAPIKey,
|
||||
Platform: PlatformGemini,
|
||||
Credentials: map[string]any{"base_url": "https://custom-gemini.example.com"},
|
||||
},
|
||||
expected: "https://custom-gemini.example.com",
|
||||
},
|
||||
{
|
||||
name: "antigravity apikey auto-appends /antigravity",
|
||||
account: Account{
|
||||
Type: AccountTypeAPIKey,
|
||||
Platform: PlatformAntigravity,
|
||||
Credentials: map[string]any{"base_url": "https://upstream.example.com"},
|
||||
},
|
||||
expected: "https://upstream.example.com/antigravity",
|
||||
},
|
||||
{
|
||||
name: "antigravity apikey trims trailing slash",
|
||||
account: Account{
|
||||
Type: AccountTypeAPIKey,
|
||||
Platform: PlatformAntigravity,
|
||||
Credentials: map[string]any{"base_url": "https://upstream.example.com/"},
|
||||
},
|
||||
expected: "https://upstream.example.com/antigravity",
|
||||
},
|
||||
{
|
||||
name: "antigravity oauth does NOT append /antigravity",
|
||||
account: Account{
|
||||
Type: AccountTypeOAuth,
|
||||
Platform: PlatformAntigravity,
|
||||
Credentials: map[string]any{"base_url": "https://upstream.example.com"},
|
||||
},
|
||||
expected: "https://upstream.example.com",
|
||||
},
|
||||
{
|
||||
name: "oauth without base_url returns default",
|
||||
account: Account{
|
||||
Type: AccountTypeOAuth,
|
||||
Platform: PlatformAntigravity,
|
||||
Credentials: map[string]any{},
|
||||
},
|
||||
expected: defaultGeminiURL,
|
||||
},
|
||||
{
|
||||
name: "nil credentials returns default",
|
||||
account: Account{
|
||||
Type: AccountTypeAPIKey,
|
||||
Platform: PlatformGemini,
|
||||
},
|
||||
expected: defaultGeminiURL,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := tt.account.GetGeminiBaseURL(defaultGeminiURL)
|
||||
if result != tt.expected {
|
||||
t.Errorf("GetGeminiBaseURL() = %q, want %q", result, tt.expected)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -28,6 +28,9 @@ type AccountRepository interface {
|
||||
// FindByExtraField 根据 extra 字段中的键值对查找账号(限定 platform='sora')
|
||||
// 用于查找通过 linked_openai_account_id 关联的 Sora 账号
|
||||
FindByExtraField(ctx context.Context, key string, value any) ([]Account, error)
|
||||
// ListCRSAccountIDs returns a map of crs_account_id -> local account ID
|
||||
// for all accounts that have been synced from CRS.
|
||||
ListCRSAccountIDs(ctx context.Context) (map[string]int64, error)
|
||||
Update(ctx context.Context, account *Account) error
|
||||
Delete(ctx context.Context, id int64) error
|
||||
|
||||
@@ -53,7 +56,6 @@ type AccountRepository interface {
|
||||
ListSchedulableByGroupIDAndPlatforms(ctx context.Context, groupID int64, platforms []string) ([]Account, error)
|
||||
|
||||
SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error
|
||||
SetAntigravityQuotaScopeLimit(ctx context.Context, id int64, scope AntigravityQuotaScope, resetAt time.Time) error
|
||||
SetModelRateLimit(ctx context.Context, id int64, scope string, resetAt time.Time) error
|
||||
SetOverloaded(ctx context.Context, id int64, until time.Time) error
|
||||
SetTempUnschedulable(ctx context.Context, id int64, until time.Time, reason string) error
|
||||
|
||||
@@ -54,10 +54,14 @@ func (s *accountRepoStub) GetByCRSAccountID(ctx context.Context, crsAccountID st
|
||||
panic("unexpected GetByCRSAccountID call")
|
||||
}
|
||||
|
||||
func (s *accountRepoStub) FindByExtraField(ctx context.Context, key string, value interface{}) ([]Account, error) {
|
||||
func (s *accountRepoStub) FindByExtraField(ctx context.Context, key string, value any) ([]Account, error) {
|
||||
panic("unexpected FindByExtraField call")
|
||||
}
|
||||
|
||||
func (s *accountRepoStub) ListCRSAccountIDs(ctx context.Context) (map[string]int64, error) {
|
||||
panic("unexpected ListCRSAccountIDs call")
|
||||
}
|
||||
|
||||
func (s *accountRepoStub) Update(ctx context.Context, account *Account) error {
|
||||
panic("unexpected Update call")
|
||||
}
|
||||
@@ -147,10 +151,6 @@ func (s *accountRepoStub) SetRateLimited(ctx context.Context, id int64, resetAt
|
||||
panic("unexpected SetRateLimited call")
|
||||
}
|
||||
|
||||
func (s *accountRepoStub) SetAntigravityQuotaScopeLimit(ctx context.Context, id int64, scope AntigravityQuotaScope, resetAt time.Time) error {
|
||||
panic("unexpected SetAntigravityQuotaScopeLimit call")
|
||||
}
|
||||
|
||||
func (s *accountRepoStub) SetModelRateLimit(ctx context.Context, id int64, scope string, resetAt time.Time) error {
|
||||
panic("unexpected SetModelRateLimit call")
|
||||
}
|
||||
|
||||
@@ -250,7 +250,6 @@ func (s *AccountTestService) testClaudeAccountConnection(c *gin.Context, account
|
||||
// Set common headers
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("anthropic-version", "2023-06-01")
|
||||
req.Header.Set("anthropic-beta", claude.DefaultBetaHeader)
|
||||
|
||||
// Apply Claude Code client headers
|
||||
for key, value := range claude.DefaultHeaders {
|
||||
@@ -259,8 +258,10 @@ func (s *AccountTestService) testClaudeAccountConnection(c *gin.Context, account
|
||||
|
||||
// Set authentication header
|
||||
if useBearer {
|
||||
req.Header.Set("anthropic-beta", claude.DefaultBetaHeader)
|
||||
req.Header.Set("Authorization", "Bearer "+authToken)
|
||||
} else {
|
||||
req.Header.Set("anthropic-beta", claude.APIKeyBetaHeader)
|
||||
req.Header.Set("x-api-key", authToken)
|
||||
}
|
||||
|
||||
|
||||
@@ -36,6 +36,7 @@ type AdminService interface {
|
||||
UpdateGroup(ctx context.Context, id int64, input *UpdateGroupInput) (*Group, error)
|
||||
DeleteGroup(ctx context.Context, id int64) error
|
||||
GetGroupAPIKeys(ctx context.Context, groupID int64, page, pageSize int) ([]APIKey, int64, error)
|
||||
UpdateGroupSortOrders(ctx context.Context, updates []GroupSortOrderUpdate) error
|
||||
|
||||
// Account management
|
||||
ListAccounts(ctx context.Context, page, pageSize int, platform, accountType, status, search string) ([]Account, int64, error)
|
||||
@@ -1048,6 +1049,10 @@ func (s *adminServiceImpl) GetGroupAPIKeys(ctx context.Context, groupID int64, p
|
||||
return keys, result.Total, nil
|
||||
}
|
||||
|
||||
func (s *adminServiceImpl) UpdateGroupSortOrders(ctx context.Context, updates []GroupSortOrderUpdate) error {
|
||||
return s.groupRepo.UpdateSortOrders(ctx, updates)
|
||||
}
|
||||
|
||||
// Account management implementations
|
||||
func (s *adminServiceImpl) ListAccounts(ctx context.Context, page, pageSize int, platform, accountType, status, search string) ([]Account, int64, error) {
|
||||
params := pagination.PaginationParams{Page: page, PageSize: pageSize}
|
||||
|
||||
@@ -172,6 +172,10 @@ func (s *groupRepoStub) GetAccountIDsByGroupIDs(ctx context.Context, groupIDs []
|
||||
panic("unexpected GetAccountIDsByGroupIDs call")
|
||||
}
|
||||
|
||||
func (s *groupRepoStub) UpdateSortOrders(ctx context.Context, updates []GroupSortOrderUpdate) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
type proxyRepoStub struct {
|
||||
deleteErr error
|
||||
countErr error
|
||||
|
||||
@@ -116,6 +116,10 @@ func (s *groupRepoStubForAdmin) GetAccountIDsByGroupIDs(_ context.Context, _ []i
|
||||
panic("unexpected GetAccountIDsByGroupIDs call")
|
||||
}
|
||||
|
||||
func (s *groupRepoStubForAdmin) UpdateSortOrders(_ context.Context, _ []GroupSortOrderUpdate) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// TestAdminService_CreateGroup_WithImagePricing 测试创建分组时 ImagePrice 字段正确传递
|
||||
func TestAdminService_CreateGroup_WithImagePricing(t *testing.T) {
|
||||
repo := &groupRepoStubForAdmin{}
|
||||
@@ -395,6 +399,10 @@ func (s *groupRepoStubForFallbackCycle) GetAccountIDsByGroupIDs(_ context.Contex
|
||||
panic("unexpected GetAccountIDsByGroupIDs call")
|
||||
}
|
||||
|
||||
func (s *groupRepoStubForFallbackCycle) UpdateSortOrders(_ context.Context, _ []GroupSortOrderUpdate) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
type groupRepoStubForInvalidRequestFallback struct {
|
||||
groups map[int64]*Group
|
||||
created *Group
|
||||
@@ -466,6 +474,10 @@ func (s *groupRepoStubForInvalidRequestFallback) BindAccountsToGroup(_ context.C
|
||||
panic("unexpected BindAccountsToGroup call")
|
||||
}
|
||||
|
||||
func (s *groupRepoStubForInvalidRequestFallback) UpdateSortOrders(_ context.Context, _ []GroupSortOrderUpdate) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestAdminService_CreateGroup_InvalidRequestFallbackRejectsUnsupportedPlatform(t *testing.T) {
|
||||
fallbackID := int64(10)
|
||||
repo := &groupRepoStubForInvalidRequestFallback{
|
||||
|
||||
79
backend/internal/service/anthropic_session.go
Normal file
79
backend/internal/service/anthropic_session.go
Normal file
@@ -0,0 +1,79 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Anthropic 会话 Fallback 相关常量
|
||||
const (
|
||||
// anthropicSessionTTLSeconds Anthropic 会话缓存 TTL(5 分钟)
|
||||
anthropicSessionTTLSeconds = 300
|
||||
|
||||
// anthropicDigestSessionKeyPrefix Anthropic 摘要 fallback 会话 key 前缀
|
||||
anthropicDigestSessionKeyPrefix = "anthropic:digest:"
|
||||
)
|
||||
|
||||
// AnthropicSessionTTL 返回 Anthropic 会话缓存 TTL
|
||||
func AnthropicSessionTTL() time.Duration {
|
||||
return anthropicSessionTTLSeconds * time.Second
|
||||
}
|
||||
|
||||
// BuildAnthropicDigestChain 根据 Anthropic 请求生成摘要链
|
||||
// 格式: s:<hash>-u:<hash>-a:<hash>-u:<hash>-...
|
||||
// s = system, u = user, a = assistant
|
||||
func BuildAnthropicDigestChain(parsed *ParsedRequest) string {
|
||||
if parsed == nil {
|
||||
return ""
|
||||
}
|
||||
|
||||
var parts []string
|
||||
|
||||
// 1. system prompt
|
||||
if parsed.System != nil {
|
||||
systemData, _ := json.Marshal(parsed.System)
|
||||
if len(systemData) > 0 && string(systemData) != "null" {
|
||||
parts = append(parts, "s:"+shortHash(systemData))
|
||||
}
|
||||
}
|
||||
|
||||
// 2. messages
|
||||
for _, msg := range parsed.Messages {
|
||||
msgMap, ok := msg.(map[string]any)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
role, _ := msgMap["role"].(string)
|
||||
prefix := rolePrefix(role)
|
||||
content := msgMap["content"]
|
||||
contentData, _ := json.Marshal(content)
|
||||
parts = append(parts, prefix+":"+shortHash(contentData))
|
||||
}
|
||||
|
||||
return strings.Join(parts, "-")
|
||||
}
|
||||
|
||||
// rolePrefix 将 Anthropic 的 role 映射为单字符前缀
|
||||
func rolePrefix(role string) string {
|
||||
switch role {
|
||||
case "assistant":
|
||||
return "a"
|
||||
default:
|
||||
return "u"
|
||||
}
|
||||
}
|
||||
|
||||
// GenerateAnthropicDigestSessionKey 生成 Anthropic 摘要 fallback 的 sessionKey
|
||||
// 组合 prefixHash 前 8 位 + uuid 前 8 位,确保不同会话产生不同的 sessionKey
|
||||
func GenerateAnthropicDigestSessionKey(prefixHash, uuid string) string {
|
||||
prefix := prefixHash
|
||||
if len(prefixHash) >= 8 {
|
||||
prefix = prefixHash[:8]
|
||||
}
|
||||
uuidPart := uuid
|
||||
if len(uuid) >= 8 {
|
||||
uuidPart = uuid[:8]
|
||||
}
|
||||
return anthropicDigestSessionKeyPrefix + prefix + ":" + uuidPart
|
||||
}
|
||||
320
backend/internal/service/anthropic_session_test.go
Normal file
320
backend/internal/service/anthropic_session_test.go
Normal file
@@ -0,0 +1,320 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestBuildAnthropicDigestChain_NilRequest(t *testing.T) {
|
||||
result := BuildAnthropicDigestChain(nil)
|
||||
if result != "" {
|
||||
t.Errorf("expected empty string for nil request, got: %s", result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildAnthropicDigestChain_EmptyMessages(t *testing.T) {
|
||||
parsed := &ParsedRequest{
|
||||
Messages: []any{},
|
||||
}
|
||||
result := BuildAnthropicDigestChain(parsed)
|
||||
if result != "" {
|
||||
t.Errorf("expected empty string for empty messages, got: %s", result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildAnthropicDigestChain_SingleUserMessage(t *testing.T) {
|
||||
parsed := &ParsedRequest{
|
||||
Messages: []any{
|
||||
map[string]any{"role": "user", "content": "hello"},
|
||||
},
|
||||
}
|
||||
result := BuildAnthropicDigestChain(parsed)
|
||||
parts := splitChain(result)
|
||||
if len(parts) != 1 {
|
||||
t.Fatalf("expected 1 part, got %d: %s", len(parts), result)
|
||||
}
|
||||
if !strings.HasPrefix(parts[0], "u:") {
|
||||
t.Errorf("expected prefix 'u:', got: %s", parts[0])
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildAnthropicDigestChain_UserAndAssistant(t *testing.T) {
|
||||
parsed := &ParsedRequest{
|
||||
Messages: []any{
|
||||
map[string]any{"role": "user", "content": "hello"},
|
||||
map[string]any{"role": "assistant", "content": "hi there"},
|
||||
},
|
||||
}
|
||||
result := BuildAnthropicDigestChain(parsed)
|
||||
parts := splitChain(result)
|
||||
if len(parts) != 2 {
|
||||
t.Fatalf("expected 2 parts, got %d: %s", len(parts), result)
|
||||
}
|
||||
if !strings.HasPrefix(parts[0], "u:") {
|
||||
t.Errorf("part[0] expected prefix 'u:', got: %s", parts[0])
|
||||
}
|
||||
if !strings.HasPrefix(parts[1], "a:") {
|
||||
t.Errorf("part[1] expected prefix 'a:', got: %s", parts[1])
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildAnthropicDigestChain_WithSystemString(t *testing.T) {
|
||||
parsed := &ParsedRequest{
|
||||
System: "You are a helpful assistant",
|
||||
Messages: []any{
|
||||
map[string]any{"role": "user", "content": "hello"},
|
||||
},
|
||||
}
|
||||
result := BuildAnthropicDigestChain(parsed)
|
||||
parts := splitChain(result)
|
||||
if len(parts) != 2 {
|
||||
t.Fatalf("expected 2 parts (s + u), got %d: %s", len(parts), result)
|
||||
}
|
||||
if !strings.HasPrefix(parts[0], "s:") {
|
||||
t.Errorf("part[0] expected prefix 's:', got: %s", parts[0])
|
||||
}
|
||||
if !strings.HasPrefix(parts[1], "u:") {
|
||||
t.Errorf("part[1] expected prefix 'u:', got: %s", parts[1])
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildAnthropicDigestChain_WithSystemContentBlocks(t *testing.T) {
|
||||
parsed := &ParsedRequest{
|
||||
System: []any{
|
||||
map[string]any{"type": "text", "text": "You are a helpful assistant"},
|
||||
},
|
||||
Messages: []any{
|
||||
map[string]any{"role": "user", "content": "hello"},
|
||||
},
|
||||
}
|
||||
result := BuildAnthropicDigestChain(parsed)
|
||||
parts := splitChain(result)
|
||||
if len(parts) != 2 {
|
||||
t.Fatalf("expected 2 parts (s + u), got %d: %s", len(parts), result)
|
||||
}
|
||||
if !strings.HasPrefix(parts[0], "s:") {
|
||||
t.Errorf("part[0] expected prefix 's:', got: %s", parts[0])
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildAnthropicDigestChain_ConversationPrefixRelationship(t *testing.T) {
|
||||
// 核心测试:验证对话增长时链的前缀关系
|
||||
// 上一轮的完整链一定是下一轮链的前缀
|
||||
system := "You are a helpful assistant"
|
||||
|
||||
// 第 1 轮: system + user
|
||||
round1 := &ParsedRequest{
|
||||
System: system,
|
||||
Messages: []any{
|
||||
map[string]any{"role": "user", "content": "hello"},
|
||||
},
|
||||
}
|
||||
chain1 := BuildAnthropicDigestChain(round1)
|
||||
|
||||
// 第 2 轮: system + user + assistant + user
|
||||
round2 := &ParsedRequest{
|
||||
System: system,
|
||||
Messages: []any{
|
||||
map[string]any{"role": "user", "content": "hello"},
|
||||
map[string]any{"role": "assistant", "content": "hi there"},
|
||||
map[string]any{"role": "user", "content": "how are you?"},
|
||||
},
|
||||
}
|
||||
chain2 := BuildAnthropicDigestChain(round2)
|
||||
|
||||
// 第 3 轮: system + user + assistant + user + assistant + user
|
||||
round3 := &ParsedRequest{
|
||||
System: system,
|
||||
Messages: []any{
|
||||
map[string]any{"role": "user", "content": "hello"},
|
||||
map[string]any{"role": "assistant", "content": "hi there"},
|
||||
map[string]any{"role": "user", "content": "how are you?"},
|
||||
map[string]any{"role": "assistant", "content": "I'm doing well"},
|
||||
map[string]any{"role": "user", "content": "great"},
|
||||
},
|
||||
}
|
||||
chain3 := BuildAnthropicDigestChain(round3)
|
||||
|
||||
t.Logf("Chain1: %s", chain1)
|
||||
t.Logf("Chain2: %s", chain2)
|
||||
t.Logf("Chain3: %s", chain3)
|
||||
|
||||
// chain1 是 chain2 的前缀
|
||||
if !strings.HasPrefix(chain2, chain1) {
|
||||
t.Errorf("chain1 should be prefix of chain2:\n chain1: %s\n chain2: %s", chain1, chain2)
|
||||
}
|
||||
|
||||
// chain2 是 chain3 的前缀
|
||||
if !strings.HasPrefix(chain3, chain2) {
|
||||
t.Errorf("chain2 should be prefix of chain3:\n chain2: %s\n chain3: %s", chain2, chain3)
|
||||
}
|
||||
|
||||
// chain1 也是 chain3 的前缀(传递性)
|
||||
if !strings.HasPrefix(chain3, chain1) {
|
||||
t.Errorf("chain1 should be prefix of chain3:\n chain1: %s\n chain3: %s", chain1, chain3)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildAnthropicDigestChain_DifferentSystemProducesDifferentChain(t *testing.T) {
|
||||
parsed1 := &ParsedRequest{
|
||||
System: "System A",
|
||||
Messages: []any{
|
||||
map[string]any{"role": "user", "content": "hello"},
|
||||
},
|
||||
}
|
||||
parsed2 := &ParsedRequest{
|
||||
System: "System B",
|
||||
Messages: []any{
|
||||
map[string]any{"role": "user", "content": "hello"},
|
||||
},
|
||||
}
|
||||
|
||||
chain1 := BuildAnthropicDigestChain(parsed1)
|
||||
chain2 := BuildAnthropicDigestChain(parsed2)
|
||||
|
||||
if chain1 == chain2 {
|
||||
t.Error("Different system prompts should produce different chains")
|
||||
}
|
||||
|
||||
// 但 user 部分的 hash 应该相同
|
||||
parts1 := splitChain(chain1)
|
||||
parts2 := splitChain(chain2)
|
||||
if parts1[1] != parts2[1] {
|
||||
t.Error("Same user message should produce same hash regardless of system")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildAnthropicDigestChain_DifferentContentProducesDifferentChain(t *testing.T) {
|
||||
parsed1 := &ParsedRequest{
|
||||
Messages: []any{
|
||||
map[string]any{"role": "user", "content": "hello"},
|
||||
map[string]any{"role": "assistant", "content": "ORIGINAL reply"},
|
||||
map[string]any{"role": "user", "content": "next"},
|
||||
},
|
||||
}
|
||||
parsed2 := &ParsedRequest{
|
||||
Messages: []any{
|
||||
map[string]any{"role": "user", "content": "hello"},
|
||||
map[string]any{"role": "assistant", "content": "TAMPERED reply"},
|
||||
map[string]any{"role": "user", "content": "next"},
|
||||
},
|
||||
}
|
||||
|
||||
chain1 := BuildAnthropicDigestChain(parsed1)
|
||||
chain2 := BuildAnthropicDigestChain(parsed2)
|
||||
|
||||
if chain1 == chain2 {
|
||||
t.Error("Different content should produce different chains")
|
||||
}
|
||||
|
||||
parts1 := splitChain(chain1)
|
||||
parts2 := splitChain(chain2)
|
||||
// 第一个 user message hash 应该相同
|
||||
if parts1[0] != parts2[0] {
|
||||
t.Error("First user message hash should be the same")
|
||||
}
|
||||
// assistant reply hash 应该不同
|
||||
if parts1[1] == parts2[1] {
|
||||
t.Error("Assistant reply hash should differ")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildAnthropicDigestChain_Deterministic(t *testing.T) {
|
||||
parsed := &ParsedRequest{
|
||||
System: "test system",
|
||||
Messages: []any{
|
||||
map[string]any{"role": "user", "content": "hello"},
|
||||
map[string]any{"role": "assistant", "content": "hi"},
|
||||
},
|
||||
}
|
||||
|
||||
chain1 := BuildAnthropicDigestChain(parsed)
|
||||
chain2 := BuildAnthropicDigestChain(parsed)
|
||||
|
||||
if chain1 != chain2 {
|
||||
t.Errorf("BuildAnthropicDigestChain not deterministic: %s vs %s", chain1, chain2)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGenerateAnthropicDigestSessionKey(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
prefixHash string
|
||||
uuid string
|
||||
want string
|
||||
}{
|
||||
{
|
||||
name: "normal 16 char hash with uuid",
|
||||
prefixHash: "abcdefgh12345678",
|
||||
uuid: "550e8400-e29b-41d4-a716-446655440000",
|
||||
want: "anthropic:digest:abcdefgh:550e8400",
|
||||
},
|
||||
{
|
||||
name: "exactly 8 chars",
|
||||
prefixHash: "12345678",
|
||||
uuid: "abcdefgh",
|
||||
want: "anthropic:digest:12345678:abcdefgh",
|
||||
},
|
||||
{
|
||||
name: "short values",
|
||||
prefixHash: "abc",
|
||||
uuid: "xyz",
|
||||
want: "anthropic:digest:abc:xyz",
|
||||
},
|
||||
{
|
||||
name: "empty values",
|
||||
prefixHash: "",
|
||||
uuid: "",
|
||||
want: "anthropic:digest::",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := GenerateAnthropicDigestSessionKey(tt.prefixHash, tt.uuid)
|
||||
if got != tt.want {
|
||||
t.Errorf("GenerateAnthropicDigestSessionKey(%q, %q) = %q, want %q", tt.prefixHash, tt.uuid, got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// 验证不同 uuid 产生不同 sessionKey
|
||||
t.Run("different uuid different key", func(t *testing.T) {
|
||||
hash := "sameprefix123456"
|
||||
result1 := GenerateAnthropicDigestSessionKey(hash, "uuid0001-session-a")
|
||||
result2 := GenerateAnthropicDigestSessionKey(hash, "uuid0002-session-b")
|
||||
if result1 == result2 {
|
||||
t.Errorf("Different UUIDs should produce different session keys: %s vs %s", result1, result2)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestAnthropicSessionTTL(t *testing.T) {
|
||||
ttl := AnthropicSessionTTL()
|
||||
if ttl.Seconds() != 300 {
|
||||
t.Errorf("expected 300 seconds, got: %v", ttl.Seconds())
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildAnthropicDigestChain_ContentBlocks(t *testing.T) {
|
||||
// 测试 content 为 content blocks 数组的情况
|
||||
parsed := &ParsedRequest{
|
||||
Messages: []any{
|
||||
map[string]any{
|
||||
"role": "user",
|
||||
"content": []any{
|
||||
map[string]any{"type": "text", "text": "describe this image"},
|
||||
map[string]any{"type": "image", "source": map[string]any{"type": "base64"}},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
result := BuildAnthropicDigestChain(parsed)
|
||||
parts := splitChain(result)
|
||||
if len(parts) != 1 {
|
||||
t.Fatalf("expected 1 part, got %d: %s", len(parts), result)
|
||||
}
|
||||
if !strings.HasPrefix(parts[0], "u:") {
|
||||
t.Errorf("expected prefix 'u:', got: %s", parts[0])
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -4,18 +4,42 @@ import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// antigravityFailingWriter 模拟客户端断开连接的 gin.ResponseWriter
|
||||
type antigravityFailingWriter struct {
|
||||
gin.ResponseWriter
|
||||
failAfter int // 允许成功写入的次数,之后所有写入返回错误
|
||||
writes int
|
||||
}
|
||||
|
||||
func (w *antigravityFailingWriter) Write(p []byte) (int, error) {
|
||||
if w.writes >= w.failAfter {
|
||||
return 0, errors.New("write failed: client disconnected")
|
||||
}
|
||||
w.writes++
|
||||
return w.ResponseWriter.Write(p)
|
||||
}
|
||||
|
||||
// newAntigravityTestService 创建用于流式测试的 AntigravityGatewayService
|
||||
func newAntigravityTestService(cfg *config.Config) *AntigravityGatewayService {
|
||||
return &AntigravityGatewayService{
|
||||
settingService: &SettingService{cfg: cfg},
|
||||
}
|
||||
}
|
||||
|
||||
func TestStripSignatureSensitiveBlocksFromClaudeRequest(t *testing.T) {
|
||||
req := &antigravity.ClaudeRequest{
|
||||
Model: "claude-sonnet-4-5",
|
||||
@@ -338,8 +362,8 @@ func TestAntigravityGatewayService_Forward_StickySessionForceCacheBilling(t *tes
|
||||
require.True(t, failoverErr.ForceCacheBilling, "ForceCacheBilling should be true for sticky session switch")
|
||||
}
|
||||
|
||||
// TestAntigravityGatewayService_ForwardGemini_StickySessionForceCacheBilling
|
||||
// 验证:ForwardGemini 粘性会话切换时,UpstreamFailoverError.ForceCacheBilling 应为 true
|
||||
// TestAntigravityGatewayService_ForwardGemini_StickySessionForceCacheBilling verifies
|
||||
// that ForwardGemini sets ForceCacheBilling=true for sticky session switch.
|
||||
func TestAntigravityGatewayService_ForwardGemini_StickySessionForceCacheBilling(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
writer := httptest.NewRecorder()
|
||||
@@ -393,10 +417,16 @@ func TestAntigravityGatewayService_ForwardGemini_StickySessionForceCacheBilling(
|
||||
require.True(t, failoverErr.ForceCacheBilling, "ForceCacheBilling should be true for sticky session switch")
|
||||
}
|
||||
|
||||
func TestAntigravityStreamUpstreamResponse_UsageAndFirstToken(t *testing.T) {
|
||||
// TestStreamUpstreamResponse_UsageAndFirstToken
|
||||
// 验证:usage 字段可被累积/覆盖更新,并且能记录首 token 时间
|
||||
func TestStreamUpstreamResponse_UsageAndFirstToken(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
writer := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(writer)
|
||||
svc := newAntigravityTestService(&config.Config{
|
||||
Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize},
|
||||
})
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
|
||||
|
||||
pr, pw := io.Pipe()
|
||||
@@ -404,25 +434,458 @@ func TestAntigravityStreamUpstreamResponse_UsageAndFirstToken(t *testing.T) {
|
||||
|
||||
go func() {
|
||||
defer func() { _ = pw.Close() }()
|
||||
_, _ = pw.Write([]byte("data: {\"usage\":{\"input_tokens\":1,\"output_tokens\":2,\"cache_read_input_tokens\":3,\"cache_creation_input_tokens\":4}}\n"))
|
||||
_, _ = pw.Write([]byte("data: {\"usage\":{\"output_tokens\":5}}\n"))
|
||||
fmt.Fprintln(pw, `data: {"usage":{"input_tokens":1,"output_tokens":2,"cache_read_input_tokens":3,"cache_creation_input_tokens":4}}`)
|
||||
fmt.Fprintln(pw, `data: {"usage":{"output_tokens":5}}`)
|
||||
}()
|
||||
|
||||
svc := &AntigravityGatewayService{}
|
||||
start := time.Now().Add(-10 * time.Millisecond)
|
||||
usage, firstTokenMs := svc.streamUpstreamResponse(c, resp, start)
|
||||
result := svc.streamUpstreamResponse(c, resp, start)
|
||||
_ = pr.Close()
|
||||
|
||||
require.NotNil(t, usage)
|
||||
require.Equal(t, 1, usage.InputTokens)
|
||||
require.NotNil(t, result)
|
||||
require.NotNil(t, result.usage)
|
||||
require.Equal(t, 1, result.usage.InputTokens)
|
||||
// 第二次事件覆盖 output_tokens
|
||||
require.Equal(t, 5, usage.OutputTokens)
|
||||
require.Equal(t, 3, usage.CacheReadInputTokens)
|
||||
require.Equal(t, 4, usage.CacheCreationInputTokens)
|
||||
require.Equal(t, 5, result.usage.OutputTokens)
|
||||
require.Equal(t, 3, result.usage.CacheReadInputTokens)
|
||||
require.Equal(t, 4, result.usage.CacheCreationInputTokens)
|
||||
require.NotNil(t, result.firstTokenMs)
|
||||
|
||||
if firstTokenMs == nil {
|
||||
t.Fatalf("expected firstTokenMs to be set")
|
||||
}
|
||||
// 确保有透传输出
|
||||
require.True(t, strings.Contains(writer.Body.String(), "data:"))
|
||||
require.Contains(t, rec.Body.String(), "data:")
|
||||
}
|
||||
|
||||
// --- 流式 happy path 测试 ---
|
||||
|
||||
// TestStreamUpstreamResponse_NormalComplete
|
||||
// 验证:正常流式转发完成时,数据正确透传、usage 正确收集、clientDisconnect=false
|
||||
func TestStreamUpstreamResponse_NormalComplete(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc := newAntigravityTestService(&config.Config{
|
||||
Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize},
|
||||
})
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||||
|
||||
pr, pw := io.Pipe()
|
||||
resp := &http.Response{StatusCode: http.StatusOK, Body: pr, Header: http.Header{}}
|
||||
|
||||
go func() {
|
||||
defer func() { _ = pw.Close() }()
|
||||
fmt.Fprintln(pw, `event: message_start`)
|
||||
fmt.Fprintln(pw, `data: {"type":"message_start","message":{"usage":{"input_tokens":10}}}`)
|
||||
fmt.Fprintln(pw, "")
|
||||
fmt.Fprintln(pw, `event: content_block_delta`)
|
||||
fmt.Fprintln(pw, `data: {"type":"content_block_delta","delta":{"text":"hello"}}`)
|
||||
fmt.Fprintln(pw, "")
|
||||
fmt.Fprintln(pw, `event: message_delta`)
|
||||
fmt.Fprintln(pw, `data: {"type":"message_delta","usage":{"output_tokens":5}}`)
|
||||
fmt.Fprintln(pw, "")
|
||||
}()
|
||||
|
||||
result := svc.streamUpstreamResponse(c, resp, time.Now())
|
||||
_ = pr.Close()
|
||||
|
||||
require.NotNil(t, result)
|
||||
require.False(t, result.clientDisconnect, "normal completion should not set clientDisconnect")
|
||||
require.NotNil(t, result.usage)
|
||||
require.Equal(t, 5, result.usage.OutputTokens, "should collect output_tokens from message_delta")
|
||||
require.NotNil(t, result.firstTokenMs, "should record first token time")
|
||||
|
||||
// 验证数据被透传到客户端
|
||||
body := rec.Body.String()
|
||||
require.Contains(t, body, "event: message_start")
|
||||
require.Contains(t, body, "content_block_delta")
|
||||
require.Contains(t, body, "message_delta")
|
||||
}
|
||||
|
||||
// TestHandleGeminiStreamingResponse_NormalComplete
|
||||
// 验证:正常 Gemini 流式转发,数据正确透传、usage 正确收集
|
||||
func TestHandleGeminiStreamingResponse_NormalComplete(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc := newAntigravityTestService(&config.Config{
|
||||
Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize},
|
||||
})
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||||
|
||||
pr, pw := io.Pipe()
|
||||
resp := &http.Response{StatusCode: http.StatusOK, Body: pr, Header: http.Header{}}
|
||||
|
||||
go func() {
|
||||
defer func() { _ = pw.Close() }()
|
||||
// 第一个 chunk(部分内容)
|
||||
fmt.Fprintln(pw, `data: {"candidates":[{"content":{"parts":[{"text":"Hello"}]}}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":3}}`)
|
||||
fmt.Fprintln(pw, "")
|
||||
// 第二个 chunk(最终内容+完整 usage)
|
||||
fmt.Fprintln(pw, `data: {"candidates":[{"content":{"parts":[{"text":" world"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":8,"cachedContentTokenCount":2}}`)
|
||||
fmt.Fprintln(pw, "")
|
||||
}()
|
||||
|
||||
result, err := svc.handleGeminiStreamingResponse(c, resp, time.Now())
|
||||
_ = pr.Close()
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.False(t, result.clientDisconnect, "normal completion should not set clientDisconnect")
|
||||
require.NotNil(t, result.usage)
|
||||
// Gemini usage: promptTokenCount=10, candidatesTokenCount=8, cachedContentTokenCount=2
|
||||
// → InputTokens=10-2=8, OutputTokens=8, CacheReadInputTokens=2
|
||||
require.Equal(t, 8, result.usage.InputTokens)
|
||||
require.Equal(t, 8, result.usage.OutputTokens)
|
||||
require.Equal(t, 2, result.usage.CacheReadInputTokens)
|
||||
require.NotNil(t, result.firstTokenMs, "should record first token time")
|
||||
|
||||
// 验证数据被透传到客户端
|
||||
body := rec.Body.String()
|
||||
require.Contains(t, body, "Hello")
|
||||
require.Contains(t, body, "world")
|
||||
// 不应包含错误事件
|
||||
require.NotContains(t, body, "event: error")
|
||||
}
|
||||
|
||||
// TestHandleClaudeStreamingResponse_NormalComplete
|
||||
// 验证:正常 Claude 流式转发(Gemini→Claude 转换),数据正确转换并输出
|
||||
func TestHandleClaudeStreamingResponse_NormalComplete(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc := newAntigravityTestService(&config.Config{
|
||||
Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize},
|
||||
})
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||||
|
||||
pr, pw := io.Pipe()
|
||||
resp := &http.Response{StatusCode: http.StatusOK, Body: pr, Header: http.Header{}}
|
||||
|
||||
go func() {
|
||||
defer func() { _ = pw.Close() }()
|
||||
// v1internal 包装格式:Gemini 数据嵌套在 "response" 字段下
|
||||
// ProcessLine 先尝试反序列化为 V1InternalResponse,裸格式会导致 Response.UsageMetadata 为空
|
||||
fmt.Fprintln(pw, `data: {"response":{"candidates":[{"content":{"parts":[{"text":"Hi there"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":5,"candidatesTokenCount":3}}}`)
|
||||
fmt.Fprintln(pw, "")
|
||||
}()
|
||||
|
||||
result, err := svc.handleClaudeStreamingResponse(c, resp, time.Now(), "claude-sonnet-4-5")
|
||||
_ = pr.Close()
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.False(t, result.clientDisconnect, "normal completion should not set clientDisconnect")
|
||||
require.NotNil(t, result.usage)
|
||||
// Gemini→Claude 转换的 usage:promptTokenCount=5→InputTokens=5, candidatesTokenCount=3→OutputTokens=3
|
||||
require.Equal(t, 5, result.usage.InputTokens)
|
||||
require.Equal(t, 3, result.usage.OutputTokens)
|
||||
require.NotNil(t, result.firstTokenMs, "should record first token time")
|
||||
|
||||
// 验证输出是 Claude SSE 格式(processor 会转换)
|
||||
body := rec.Body.String()
|
||||
require.Contains(t, body, "event: message_start", "should contain Claude message_start event")
|
||||
require.Contains(t, body, "event: message_stop", "should contain Claude message_stop event")
|
||||
// 不应包含错误事件
|
||||
require.NotContains(t, body, "event: error")
|
||||
}
|
||||
|
||||
// --- 流式客户端断开检测测试 ---
|
||||
|
||||
// TestStreamUpstreamResponse_ClientDisconnectDrainsUsage
|
||||
// 验证:客户端写入失败后,streamUpstreamResponse 继续读取上游以收集 usage
|
||||
func TestStreamUpstreamResponse_ClientDisconnectDrainsUsage(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc := newAntigravityTestService(&config.Config{
|
||||
Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize},
|
||||
})
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||||
c.Writer = &antigravityFailingWriter{ResponseWriter: c.Writer, failAfter: 0}
|
||||
|
||||
pr, pw := io.Pipe()
|
||||
resp := &http.Response{StatusCode: http.StatusOK, Body: pr, Header: http.Header{}}
|
||||
|
||||
go func() {
|
||||
defer func() { _ = pw.Close() }()
|
||||
fmt.Fprintln(pw, `event: message_start`)
|
||||
fmt.Fprintln(pw, `data: {"type":"message_start","message":{"usage":{"input_tokens":10}}}`)
|
||||
fmt.Fprintln(pw, "")
|
||||
fmt.Fprintln(pw, `event: message_delta`)
|
||||
fmt.Fprintln(pw, `data: {"type":"message_delta","usage":{"output_tokens":20}}`)
|
||||
fmt.Fprintln(pw, "")
|
||||
}()
|
||||
|
||||
result := svc.streamUpstreamResponse(c, resp, time.Now())
|
||||
_ = pr.Close()
|
||||
|
||||
require.NotNil(t, result)
|
||||
require.True(t, result.clientDisconnect)
|
||||
require.NotNil(t, result.usage)
|
||||
require.Equal(t, 20, result.usage.OutputTokens)
|
||||
}
|
||||
|
||||
// TestStreamUpstreamResponse_ContextCanceled
|
||||
// 验证:context 取消时返回 usage 且标记 clientDisconnect
|
||||
func TestStreamUpstreamResponse_ContextCanceled(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc := newAntigravityTestService(&config.Config{
|
||||
Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize},
|
||||
})
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil).WithContext(ctx)
|
||||
|
||||
resp := &http.Response{StatusCode: http.StatusOK, Body: cancelReadCloser{}, Header: http.Header{}}
|
||||
|
||||
result := svc.streamUpstreamResponse(c, resp, time.Now())
|
||||
|
||||
require.NotNil(t, result)
|
||||
require.True(t, result.clientDisconnect)
|
||||
require.NotContains(t, rec.Body.String(), "event: error")
|
||||
}
|
||||
|
||||
// TestStreamUpstreamResponse_Timeout
|
||||
// 验证:上游超时时返回已收集的 usage
|
||||
func TestStreamUpstreamResponse_Timeout(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc := newAntigravityTestService(&config.Config{
|
||||
Gateway: config.GatewayConfig{StreamDataIntervalTimeout: 1, MaxLineSize: defaultMaxLineSize},
|
||||
})
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||||
|
||||
pr, pw := io.Pipe()
|
||||
resp := &http.Response{StatusCode: http.StatusOK, Body: pr, Header: http.Header{}}
|
||||
|
||||
result := svc.streamUpstreamResponse(c, resp, time.Now())
|
||||
_ = pw.Close()
|
||||
_ = pr.Close()
|
||||
|
||||
require.NotNil(t, result)
|
||||
require.False(t, result.clientDisconnect)
|
||||
}
|
||||
|
||||
// TestStreamUpstreamResponse_TimeoutAfterClientDisconnect
|
||||
// 验证:客户端断开后上游超时,返回 usage 并标记 clientDisconnect
|
||||
func TestStreamUpstreamResponse_TimeoutAfterClientDisconnect(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc := newAntigravityTestService(&config.Config{
|
||||
Gateway: config.GatewayConfig{StreamDataIntervalTimeout: 1, MaxLineSize: defaultMaxLineSize},
|
||||
})
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||||
c.Writer = &antigravityFailingWriter{ResponseWriter: c.Writer, failAfter: 0}
|
||||
|
||||
pr, pw := io.Pipe()
|
||||
resp := &http.Response{StatusCode: http.StatusOK, Body: pr, Header: http.Header{}}
|
||||
|
||||
go func() {
|
||||
fmt.Fprintln(pw, `data: {"type":"message_start","message":{"usage":{"input_tokens":5}}}`)
|
||||
fmt.Fprintln(pw, "")
|
||||
// 不关闭 pw → 等待超时
|
||||
}()
|
||||
|
||||
result := svc.streamUpstreamResponse(c, resp, time.Now())
|
||||
_ = pw.Close()
|
||||
_ = pr.Close()
|
||||
|
||||
require.NotNil(t, result)
|
||||
require.True(t, result.clientDisconnect)
|
||||
}
|
||||
|
||||
// TestHandleGeminiStreamingResponse_ClientDisconnect
|
||||
// 验证:Gemini 流式转发中客户端断开后继续 drain 上游
|
||||
func TestHandleGeminiStreamingResponse_ClientDisconnect(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc := newAntigravityTestService(&config.Config{
|
||||
Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize},
|
||||
})
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||||
c.Writer = &antigravityFailingWriter{ResponseWriter: c.Writer, failAfter: 0}
|
||||
|
||||
pr, pw := io.Pipe()
|
||||
resp := &http.Response{StatusCode: http.StatusOK, Body: pr, Header: http.Header{}}
|
||||
|
||||
go func() {
|
||||
defer func() { _ = pw.Close() }()
|
||||
fmt.Fprintln(pw, `data: {"candidates":[{"content":{"parts":[{"text":"hi"}]}}],"usageMetadata":{"promptTokenCount":5,"candidatesTokenCount":10}}`)
|
||||
fmt.Fprintln(pw, "")
|
||||
}()
|
||||
|
||||
result, err := svc.handleGeminiStreamingResponse(c, resp, time.Now())
|
||||
_ = pr.Close()
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.True(t, result.clientDisconnect)
|
||||
require.NotContains(t, rec.Body.String(), "write_failed")
|
||||
}
|
||||
|
||||
// TestHandleGeminiStreamingResponse_ContextCanceled
|
||||
// 验证:context 取消时不注入错误事件
|
||||
func TestHandleGeminiStreamingResponse_ContextCanceled(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc := newAntigravityTestService(&config.Config{
|
||||
Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize},
|
||||
})
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil).WithContext(ctx)
|
||||
|
||||
resp := &http.Response{StatusCode: http.StatusOK, Body: cancelReadCloser{}, Header: http.Header{}}
|
||||
|
||||
result, err := svc.handleGeminiStreamingResponse(c, resp, time.Now())
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.True(t, result.clientDisconnect)
|
||||
require.NotContains(t, rec.Body.String(), "event: error")
|
||||
}
|
||||
|
||||
// TestHandleClaudeStreamingResponse_ClientDisconnect
|
||||
// 验证:Claude 流式转发中客户端断开后继续 drain 上游
|
||||
func TestHandleClaudeStreamingResponse_ClientDisconnect(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc := newAntigravityTestService(&config.Config{
|
||||
Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize},
|
||||
})
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||||
c.Writer = &antigravityFailingWriter{ResponseWriter: c.Writer, failAfter: 0}
|
||||
|
||||
pr, pw := io.Pipe()
|
||||
resp := &http.Response{StatusCode: http.StatusOK, Body: pr, Header: http.Header{}}
|
||||
|
||||
go func() {
|
||||
defer func() { _ = pw.Close() }()
|
||||
// v1internal 包装格式
|
||||
fmt.Fprintln(pw, `data: {"response":{"candidates":[{"content":{"parts":[{"text":"hello"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":8,"candidatesTokenCount":15}}}`)
|
||||
fmt.Fprintln(pw, "")
|
||||
}()
|
||||
|
||||
result, err := svc.handleClaudeStreamingResponse(c, resp, time.Now(), "claude-sonnet-4-5")
|
||||
_ = pr.Close()
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.True(t, result.clientDisconnect)
|
||||
}
|
||||
|
||||
// TestHandleClaudeStreamingResponse_ContextCanceled
|
||||
// 验证:context 取消时不注入错误事件
|
||||
func TestHandleClaudeStreamingResponse_ContextCanceled(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc := newAntigravityTestService(&config.Config{
|
||||
Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize},
|
||||
})
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil).WithContext(ctx)
|
||||
|
||||
resp := &http.Response{StatusCode: http.StatusOK, Body: cancelReadCloser{}, Header: http.Header{}}
|
||||
|
||||
result, err := svc.handleClaudeStreamingResponse(c, resp, time.Now(), "claude-sonnet-4-5")
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.True(t, result.clientDisconnect)
|
||||
require.NotContains(t, rec.Body.String(), "event: error")
|
||||
}
|
||||
|
||||
// TestExtractSSEUsage 验证 extractSSEUsage 从 SSE data 行正确提取 usage
|
||||
func TestExtractSSEUsage(t *testing.T) {
|
||||
svc := &AntigravityGatewayService{}
|
||||
tests := []struct {
|
||||
name string
|
||||
line string
|
||||
expected ClaudeUsage
|
||||
}{
|
||||
{
|
||||
name: "message_delta with output_tokens",
|
||||
line: `data: {"type":"message_delta","usage":{"output_tokens":42}}`,
|
||||
expected: ClaudeUsage{OutputTokens: 42},
|
||||
},
|
||||
{
|
||||
name: "non-data line ignored",
|
||||
line: `event: message_start`,
|
||||
expected: ClaudeUsage{},
|
||||
},
|
||||
{
|
||||
name: "top-level usage with all fields",
|
||||
line: `data: {"usage":{"input_tokens":10,"output_tokens":20,"cache_read_input_tokens":5,"cache_creation_input_tokens":3}}`,
|
||||
expected: ClaudeUsage{InputTokens: 10, OutputTokens: 20, CacheReadInputTokens: 5, CacheCreationInputTokens: 3},
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
usage := &ClaudeUsage{}
|
||||
svc.extractSSEUsage(tt.line, usage)
|
||||
require.Equal(t, tt.expected, *usage)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestAntigravityClientWriter 验证 antigravityClientWriter 的断开检测
|
||||
func TestAntigravityClientWriter(t *testing.T) {
|
||||
t.Run("normal write succeeds", func(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
flusher, _ := c.Writer.(http.Flusher)
|
||||
cw := newAntigravityClientWriter(c.Writer, flusher, "test")
|
||||
|
||||
ok := cw.Write([]byte("hello"))
|
||||
require.True(t, ok)
|
||||
require.False(t, cw.Disconnected())
|
||||
require.Contains(t, rec.Body.String(), "hello")
|
||||
})
|
||||
|
||||
t.Run("write failure marks disconnected", func(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
fw := &antigravityFailingWriter{ResponseWriter: c.Writer, failAfter: 0}
|
||||
flusher, _ := c.Writer.(http.Flusher)
|
||||
cw := newAntigravityClientWriter(fw, flusher, "test")
|
||||
|
||||
ok := cw.Write([]byte("hello"))
|
||||
require.False(t, ok)
|
||||
require.True(t, cw.Disconnected())
|
||||
})
|
||||
|
||||
t.Run("subsequent writes are no-op", func(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
fw := &antigravityFailingWriter{ResponseWriter: c.Writer, failAfter: 0}
|
||||
flusher, _ := c.Writer.(http.Flusher)
|
||||
cw := newAntigravityClientWriter(fw, flusher, "test")
|
||||
|
||||
cw.Write([]byte("first"))
|
||||
ok := cw.Fprintf("second %d", 2)
|
||||
require.False(t, ok)
|
||||
require.True(t, cw.Disconnected())
|
||||
})
|
||||
}
|
||||
|
||||
@@ -2,63 +2,23 @@ package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"slices"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
const antigravityQuotaScopesKey = "antigravity_quota_scopes"
|
||||
|
||||
// AntigravityQuotaScope 表示 Antigravity 的配额域
|
||||
type AntigravityQuotaScope string
|
||||
|
||||
const (
|
||||
AntigravityQuotaScopeClaude AntigravityQuotaScope = "claude"
|
||||
AntigravityQuotaScopeGeminiText AntigravityQuotaScope = "gemini_text"
|
||||
AntigravityQuotaScopeGeminiImage AntigravityQuotaScope = "gemini_image"
|
||||
)
|
||||
|
||||
// IsScopeSupported 检查给定的 scope 是否在分组支持的 scope 列表中
|
||||
func IsScopeSupported(supportedScopes []string, scope AntigravityQuotaScope) bool {
|
||||
if len(supportedScopes) == 0 {
|
||||
// 未配置时默认全部支持
|
||||
return true
|
||||
}
|
||||
supported := slices.Contains(supportedScopes, string(scope))
|
||||
return supported
|
||||
}
|
||||
|
||||
// ResolveAntigravityQuotaScope 根据模型名称解析配额域(导出版本)
|
||||
func ResolveAntigravityQuotaScope(requestedModel string) (AntigravityQuotaScope, bool) {
|
||||
return resolveAntigravityQuotaScope(requestedModel)
|
||||
}
|
||||
|
||||
// resolveAntigravityQuotaScope 根据模型名称解析配额域
|
||||
func resolveAntigravityQuotaScope(requestedModel string) (AntigravityQuotaScope, bool) {
|
||||
model := normalizeAntigravityModelName(requestedModel)
|
||||
if model == "" {
|
||||
return "", false
|
||||
}
|
||||
switch {
|
||||
case strings.HasPrefix(model, "claude-"):
|
||||
return AntigravityQuotaScopeClaude, true
|
||||
case strings.HasPrefix(model, "gemini-"):
|
||||
if isImageGenerationModel(model) {
|
||||
return AntigravityQuotaScopeGeminiImage, true
|
||||
}
|
||||
return AntigravityQuotaScopeGeminiText, true
|
||||
default:
|
||||
return "", false
|
||||
}
|
||||
}
|
||||
|
||||
func normalizeAntigravityModelName(model string) string {
|
||||
normalized := strings.ToLower(strings.TrimSpace(model))
|
||||
normalized = strings.TrimPrefix(normalized, "models/")
|
||||
return normalized
|
||||
}
|
||||
|
||||
// IsSchedulableForModel 结合 Antigravity 配额域限流判断是否可调度。
|
||||
// resolveAntigravityModelKey 根据请求的模型名解析限流 key
|
||||
// 返回空字符串表示无法解析
|
||||
func resolveAntigravityModelKey(requestedModel string) string {
|
||||
return normalizeAntigravityModelName(requestedModel)
|
||||
}
|
||||
|
||||
// IsSchedulableForModel 结合模型级限流判断是否可调度。
|
||||
// 保持旧签名以兼容既有调用方;默认使用 context.Background()。
|
||||
func (a *Account) IsSchedulableForModel(requestedModel string) bool {
|
||||
return a.IsSchedulableForModelWithContext(context.Background(), requestedModel)
|
||||
@@ -74,107 +34,20 @@ func (a *Account) IsSchedulableForModelWithContext(ctx context.Context, requeste
|
||||
if a.isModelRateLimitedWithContext(ctx, requestedModel) {
|
||||
return false
|
||||
}
|
||||
if a.Platform != PlatformAntigravity {
|
||||
return true
|
||||
}
|
||||
scope, ok := resolveAntigravityQuotaScope(requestedModel)
|
||||
if !ok {
|
||||
return true
|
||||
}
|
||||
resetAt := a.antigravityQuotaScopeResetAt(scope)
|
||||
if resetAt == nil {
|
||||
return true
|
||||
}
|
||||
now := time.Now()
|
||||
return !now.Before(*resetAt)
|
||||
return true
|
||||
}
|
||||
|
||||
func (a *Account) antigravityQuotaScopeResetAt(scope AntigravityQuotaScope) *time.Time {
|
||||
if a == nil || a.Extra == nil || scope == "" {
|
||||
return nil
|
||||
}
|
||||
rawScopes, ok := a.Extra[antigravityQuotaScopesKey].(map[string]any)
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
rawScope, ok := rawScopes[string(scope)].(map[string]any)
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
resetAtRaw, ok := rawScope["rate_limit_reset_at"].(string)
|
||||
if !ok || strings.TrimSpace(resetAtRaw) == "" {
|
||||
return nil
|
||||
}
|
||||
resetAt, err := time.Parse(time.RFC3339, resetAtRaw)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
return &resetAt
|
||||
}
|
||||
|
||||
var antigravityAllScopes = []AntigravityQuotaScope{
|
||||
AntigravityQuotaScopeClaude,
|
||||
AntigravityQuotaScopeGeminiText,
|
||||
AntigravityQuotaScopeGeminiImage,
|
||||
}
|
||||
|
||||
func (a *Account) GetAntigravityScopeRateLimits() map[string]int64 {
|
||||
if a == nil || a.Platform != PlatformAntigravity {
|
||||
return nil
|
||||
}
|
||||
now := time.Now()
|
||||
result := make(map[string]int64)
|
||||
for _, scope := range antigravityAllScopes {
|
||||
resetAt := a.antigravityQuotaScopeResetAt(scope)
|
||||
if resetAt != nil && now.Before(*resetAt) {
|
||||
remainingSec := int64(time.Until(*resetAt).Seconds())
|
||||
if remainingSec > 0 {
|
||||
result[string(scope)] = remainingSec
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(result) == 0 {
|
||||
return nil
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// GetQuotaScopeRateLimitRemainingTime 获取模型域限流剩余时间
|
||||
// 返回 0 表示未限流或已过期
|
||||
func (a *Account) GetQuotaScopeRateLimitRemainingTime(requestedModel string) time.Duration {
|
||||
if a == nil || a.Platform != PlatformAntigravity {
|
||||
return 0
|
||||
}
|
||||
scope, ok := resolveAntigravityQuotaScope(requestedModel)
|
||||
if !ok {
|
||||
return 0
|
||||
}
|
||||
resetAt := a.antigravityQuotaScopeResetAt(scope)
|
||||
if resetAt == nil {
|
||||
return 0
|
||||
}
|
||||
if remaining := time.Until(*resetAt); remaining > 0 {
|
||||
return remaining
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// GetRateLimitRemainingTime 获取限流剩余时间(模型限流和模型域限流取最大值)
|
||||
// GetRateLimitRemainingTime 获取限流剩余时间(模型级限流)
|
||||
// 返回 0 表示未限流或已过期
|
||||
func (a *Account) GetRateLimitRemainingTime(requestedModel string) time.Duration {
|
||||
return a.GetRateLimitRemainingTimeWithContext(context.Background(), requestedModel)
|
||||
}
|
||||
|
||||
// GetRateLimitRemainingTimeWithContext 获取限流剩余时间(模型限流和模型域限流取最大值)
|
||||
// GetRateLimitRemainingTimeWithContext 获取限流剩余时间(模型级限流)
|
||||
// 返回 0 表示未限流或已过期
|
||||
func (a *Account) GetRateLimitRemainingTimeWithContext(ctx context.Context, requestedModel string) time.Duration {
|
||||
if a == nil {
|
||||
return 0
|
||||
}
|
||||
modelRemaining := a.GetModelRateLimitRemainingTimeWithContext(ctx, requestedModel)
|
||||
scopeRemaining := a.GetQuotaScopeRateLimitRemainingTime(requestedModel)
|
||||
if modelRemaining > scopeRemaining {
|
||||
return modelRemaining
|
||||
}
|
||||
return scopeRemaining
|
||||
return a.GetModelRateLimitRemainingTimeWithContext(ctx, requestedModel)
|
||||
}
|
||||
|
||||
@@ -65,12 +65,6 @@ func (s *stubAntigravityUpstream) DoWithTLS(req *http.Request, proxyURL string,
|
||||
return s.Do(req, proxyURL, accountID, accountConcurrency)
|
||||
}
|
||||
|
||||
type scopeLimitCall struct {
|
||||
accountID int64
|
||||
scope AntigravityQuotaScope
|
||||
resetAt time.Time
|
||||
}
|
||||
|
||||
type rateLimitCall struct {
|
||||
accountID int64
|
||||
resetAt time.Time
|
||||
@@ -84,16 +78,10 @@ type modelRateLimitCall struct {
|
||||
|
||||
type stubAntigravityAccountRepo struct {
|
||||
AccountRepository
|
||||
scopeCalls []scopeLimitCall
|
||||
rateCalls []rateLimitCall
|
||||
modelRateLimitCalls []modelRateLimitCall
|
||||
}
|
||||
|
||||
func (s *stubAntigravityAccountRepo) SetAntigravityQuotaScopeLimit(ctx context.Context, id int64, scope AntigravityQuotaScope, resetAt time.Time) error {
|
||||
s.scopeCalls = append(s.scopeCalls, scopeLimitCall{accountID: id, scope: scope, resetAt: resetAt})
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *stubAntigravityAccountRepo) SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error {
|
||||
s.rateCalls = append(s.rateCalls, rateLimitCall{accountID: id, resetAt: resetAt})
|
||||
return nil
|
||||
@@ -137,10 +125,9 @@ func TestAntigravityRetryLoop_URLFallback_UsesLatestSuccess(t *testing.T) {
|
||||
accessToken: "token",
|
||||
action: "generateContent",
|
||||
body: []byte(`{"input":"test"}`),
|
||||
quotaScope: AntigravityQuotaScopeClaude,
|
||||
httpUpstream: upstream,
|
||||
requestedModel: "claude-sonnet-4-5",
|
||||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, quotaScope AntigravityQuotaScope, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||||
handleErrorCalled = true
|
||||
return nil
|
||||
},
|
||||
@@ -161,23 +148,6 @@ func TestAntigravityRetryLoop_URLFallback_UsesLatestSuccess(t *testing.T) {
|
||||
require.Equal(t, base2, available[0])
|
||||
}
|
||||
|
||||
func TestAntigravityHandleUpstreamError_UsesScopeLimit(t *testing.T) {
|
||||
// 分区限流始终开启,不再支持通过环境变量关闭
|
||||
repo := &stubAntigravityAccountRepo{}
|
||||
svc := &AntigravityGatewayService{accountRepo: repo}
|
||||
account := &Account{ID: 9, Name: "acc-9", Platform: PlatformAntigravity}
|
||||
|
||||
body := buildGeminiRateLimitBody("3s")
|
||||
svc.handleUpstreamError(context.Background(), "[test]", account, http.StatusTooManyRequests, http.Header{}, body, AntigravityQuotaScopeClaude, 0, "", false)
|
||||
|
||||
require.Len(t, repo.scopeCalls, 1)
|
||||
require.Empty(t, repo.rateCalls)
|
||||
call := repo.scopeCalls[0]
|
||||
require.Equal(t, account.ID, call.accountID)
|
||||
require.Equal(t, AntigravityQuotaScopeClaude, call.scope)
|
||||
require.WithinDuration(t, time.Now().Add(3*time.Second), call.resetAt, 2*time.Second)
|
||||
}
|
||||
|
||||
// TestHandleUpstreamError_429_ModelRateLimit 测试 429 模型限流场景
|
||||
func TestHandleUpstreamError_429_ModelRateLimit(t *testing.T) {
|
||||
repo := &stubAntigravityAccountRepo{}
|
||||
@@ -195,7 +165,7 @@ func TestHandleUpstreamError_429_ModelRateLimit(t *testing.T) {
|
||||
}
|
||||
}`)
|
||||
|
||||
result := svc.handleUpstreamError(context.Background(), "[test]", account, http.StatusTooManyRequests, http.Header{}, body, AntigravityQuotaScopeClaude, 0, "", false)
|
||||
result := svc.handleUpstreamError(context.Background(), "[test]", account, http.StatusTooManyRequests, http.Header{}, body, "claude-sonnet-4-5", 0, "", false)
|
||||
|
||||
// 应该触发模型限流
|
||||
require.NotNil(t, result)
|
||||
@@ -206,22 +176,22 @@ func TestHandleUpstreamError_429_ModelRateLimit(t *testing.T) {
|
||||
require.Equal(t, "claude-sonnet-4-5", repo.modelRateLimitCalls[0].modelKey)
|
||||
}
|
||||
|
||||
// TestHandleUpstreamError_429_NonModelRateLimit 测试 429 非模型限流场景(走 scope 限流)
|
||||
// TestHandleUpstreamError_429_NonModelRateLimit 测试 429 非模型限流场景(走模型级限流兜底)
|
||||
func TestHandleUpstreamError_429_NonModelRateLimit(t *testing.T) {
|
||||
repo := &stubAntigravityAccountRepo{}
|
||||
svc := &AntigravityGatewayService{accountRepo: repo}
|
||||
account := &Account{ID: 2, Name: "acc-2", Platform: PlatformAntigravity}
|
||||
|
||||
// 429 + 普通限流响应(无 RATE_LIMIT_EXCEEDED reason)→ scope 限流
|
||||
// 429 + 普通限流响应(无 RATE_LIMIT_EXCEEDED reason)→ 走模型级限流兜底
|
||||
body := buildGeminiRateLimitBody("5s")
|
||||
|
||||
result := svc.handleUpstreamError(context.Background(), "[test]", account, http.StatusTooManyRequests, http.Header{}, body, AntigravityQuotaScopeClaude, 0, "", false)
|
||||
result := svc.handleUpstreamError(context.Background(), "[test]", account, http.StatusTooManyRequests, http.Header{}, body, "claude-sonnet-4-5", 0, "", false)
|
||||
|
||||
// 不应该触发模型限流,应该走 scope 限流
|
||||
// handleModelRateLimit 不会处理(因为没有 RATE_LIMIT_EXCEEDED),
|
||||
// 但 429 兜底逻辑会使用 requestedModel 设置模型级限流
|
||||
require.Nil(t, result)
|
||||
require.Empty(t, repo.modelRateLimitCalls)
|
||||
require.Len(t, repo.scopeCalls, 1)
|
||||
require.Equal(t, AntigravityQuotaScopeClaude, repo.scopeCalls[0].scope)
|
||||
require.Len(t, repo.modelRateLimitCalls, 1)
|
||||
require.Equal(t, "claude-sonnet-4-5", repo.modelRateLimitCalls[0].modelKey)
|
||||
}
|
||||
|
||||
// TestHandleUpstreamError_503_ModelRateLimit 测试 503 模型限流场景
|
||||
@@ -241,7 +211,7 @@ func TestHandleUpstreamError_503_ModelRateLimit(t *testing.T) {
|
||||
}
|
||||
}`)
|
||||
|
||||
result := svc.handleUpstreamError(context.Background(), "[test]", account, http.StatusServiceUnavailable, http.Header{}, body, AntigravityQuotaScopeGeminiText, 0, "", false)
|
||||
result := svc.handleUpstreamError(context.Background(), "[test]", account, http.StatusServiceUnavailable, http.Header{}, body, "gemini-3-pro-high", 0, "", false)
|
||||
|
||||
// 应该触发模型限流
|
||||
require.NotNil(t, result)
|
||||
@@ -269,12 +239,11 @@ func TestHandleUpstreamError_503_NonModelRateLimit(t *testing.T) {
|
||||
}
|
||||
}`)
|
||||
|
||||
result := svc.handleUpstreamError(context.Background(), "[test]", account, http.StatusServiceUnavailable, http.Header{}, body, AntigravityQuotaScopeGeminiText, 0, "", false)
|
||||
result := svc.handleUpstreamError(context.Background(), "[test]", account, http.StatusServiceUnavailable, http.Header{}, body, "gemini-3-pro-high", 0, "", false)
|
||||
|
||||
// 503 非模型限流不应该做任何处理
|
||||
require.Nil(t, result)
|
||||
require.Empty(t, repo.modelRateLimitCalls, "503 non-model rate limit should not trigger model rate limit")
|
||||
require.Empty(t, repo.scopeCalls, "503 non-model rate limit should not trigger scope rate limit")
|
||||
require.Empty(t, repo.rateCalls, "503 non-model rate limit should not trigger account rate limit")
|
||||
}
|
||||
|
||||
@@ -287,12 +256,11 @@ func TestHandleUpstreamError_503_EmptyBody(t *testing.T) {
|
||||
// 503 + 空响应体 → 不做任何处理
|
||||
body := []byte(`{}`)
|
||||
|
||||
result := svc.handleUpstreamError(context.Background(), "[test]", account, http.StatusServiceUnavailable, http.Header{}, body, AntigravityQuotaScopeGeminiText, 0, "", false)
|
||||
result := svc.handleUpstreamError(context.Background(), "[test]", account, http.StatusServiceUnavailable, http.Header{}, body, "gemini-3-pro-high", 0, "", false)
|
||||
|
||||
// 503 空响应不应该做任何处理
|
||||
require.Nil(t, result)
|
||||
require.Empty(t, repo.modelRateLimitCalls)
|
||||
require.Empty(t, repo.scopeCalls)
|
||||
require.Empty(t, repo.rateCalls)
|
||||
}
|
||||
|
||||
@@ -313,15 +281,7 @@ func TestAccountIsSchedulableForModel_AntigravityRateLimits(t *testing.T) {
|
||||
require.False(t, account.IsSchedulableForModel("gemini-3-flash"))
|
||||
|
||||
account.RateLimitResetAt = nil
|
||||
account.Extra = map[string]any{
|
||||
antigravityQuotaScopesKey: map[string]any{
|
||||
"claude": map[string]any{
|
||||
"rate_limit_reset_at": future.Format(time.RFC3339),
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
require.False(t, account.IsSchedulableForModel("claude-sonnet-4-5"))
|
||||
require.True(t, account.IsSchedulableForModel("claude-sonnet-4-5"))
|
||||
require.True(t, account.IsSchedulableForModel("gemini-3-flash"))
|
||||
}
|
||||
|
||||
@@ -641,6 +601,7 @@ func TestShouldTriggerAntigravitySmartRetry(t *testing.T) {
|
||||
}`,
|
||||
expectedShouldRetry: false,
|
||||
expectedShouldRateLimit: true,
|
||||
minWait: 7 * time.Second,
|
||||
modelName: "gemini-pro",
|
||||
},
|
||||
{
|
||||
@@ -658,6 +619,7 @@ func TestShouldTriggerAntigravitySmartRetry(t *testing.T) {
|
||||
}`,
|
||||
expectedShouldRetry: false,
|
||||
expectedShouldRateLimit: true,
|
||||
minWait: 39 * time.Second,
|
||||
modelName: "gemini-3-pro-high",
|
||||
},
|
||||
{
|
||||
@@ -675,6 +637,7 @@ func TestShouldTriggerAntigravitySmartRetry(t *testing.T) {
|
||||
}`,
|
||||
expectedShouldRetry: false,
|
||||
expectedShouldRateLimit: true,
|
||||
minWait: 30 * time.Second,
|
||||
modelName: "gemini-2.5-flash",
|
||||
},
|
||||
{
|
||||
@@ -692,6 +655,7 @@ func TestShouldTriggerAntigravitySmartRetry(t *testing.T) {
|
||||
}`,
|
||||
expectedShouldRetry: false,
|
||||
expectedShouldRateLimit: true,
|
||||
minWait: 30 * time.Second,
|
||||
modelName: "claude-sonnet-4-5",
|
||||
},
|
||||
}
|
||||
@@ -710,6 +674,11 @@ func TestShouldTriggerAntigravitySmartRetry(t *testing.T) {
|
||||
t.Errorf("wait = %v, want >= %v", wait, tt.minWait)
|
||||
}
|
||||
}
|
||||
if shouldRateLimit && tt.minWait > 0 {
|
||||
if wait < tt.minWait {
|
||||
t.Errorf("rate limit wait = %v, want >= %v", wait, tt.minWait)
|
||||
}
|
||||
}
|
||||
if (shouldRetry || shouldRateLimit) && model != tt.modelName {
|
||||
t.Errorf("modelName = %q, want %q", model, tt.modelName)
|
||||
}
|
||||
@@ -809,7 +778,7 @@ func TestSetModelRateLimitByModelName_NotConvertToScope(t *testing.T) {
|
||||
require.NotEqual(t, "claude_sonnet", call.modelKey, "should NOT be scope")
|
||||
}
|
||||
|
||||
func TestAntigravityRetryLoop_PreCheck_WaitsWhenRemainingBelowThreshold(t *testing.T) {
|
||||
func TestAntigravityRetryLoop_PreCheck_SwitchesWhenRateLimited(t *testing.T) {
|
||||
upstream := &recordingOKUpstream{}
|
||||
account := &Account{
|
||||
ID: 1,
|
||||
@@ -821,19 +790,15 @@ func TestAntigravityRetryLoop_PreCheck_WaitsWhenRemainingBelowThreshold(t *testi
|
||||
Extra: map[string]any{
|
||||
modelRateLimitsKey: map[string]any{
|
||||
"claude-sonnet-4-5": map[string]any{
|
||||
// RFC3339 here is second-precision; keep it safely in the future.
|
||||
"rate_limit_reset_at": time.Now().Add(2 * time.Second).Format(time.RFC3339),
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond)
|
||||
defer cancel()
|
||||
|
||||
svc := &AntigravityGatewayService{}
|
||||
result, err := svc.antigravityRetryLoop(antigravityRetryLoopParams{
|
||||
ctx: ctx,
|
||||
ctx: context.Background(),
|
||||
prefix: "[test]",
|
||||
account: account,
|
||||
accessToken: "token",
|
||||
@@ -842,17 +807,21 @@ func TestAntigravityRetryLoop_PreCheck_WaitsWhenRemainingBelowThreshold(t *testi
|
||||
requestedModel: "claude-sonnet-4-5",
|
||||
httpUpstream: upstream,
|
||||
isStickySession: true,
|
||||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, quotaScope AntigravityQuotaScope, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||||
return nil
|
||||
},
|
||||
})
|
||||
|
||||
require.ErrorIs(t, err, context.DeadlineExceeded)
|
||||
require.Nil(t, result)
|
||||
require.Equal(t, 0, upstream.calls, "should not call upstream while waiting on pre-check")
|
||||
var switchErr *AntigravityAccountSwitchError
|
||||
require.ErrorAs(t, err, &switchErr)
|
||||
require.Equal(t, account.ID, switchErr.OriginalAccountID)
|
||||
require.Equal(t, "claude-sonnet-4-5", switchErr.RateLimitedModel)
|
||||
require.True(t, switchErr.IsStickySession)
|
||||
require.Equal(t, 0, upstream.calls, "should not call upstream when switching on pre-check")
|
||||
}
|
||||
|
||||
func TestAntigravityRetryLoop_PreCheck_SwitchesWhenRemainingAtOrAboveThreshold(t *testing.T) {
|
||||
func TestAntigravityRetryLoop_PreCheck_SwitchesWhenRemainingLong(t *testing.T) {
|
||||
upstream := &recordingOKUpstream{}
|
||||
account := &Account{
|
||||
ID: 2,
|
||||
@@ -881,7 +850,7 @@ func TestAntigravityRetryLoop_PreCheck_SwitchesWhenRemainingAtOrAboveThreshold(t
|
||||
requestedModel: "claude-sonnet-4-5",
|
||||
httpUpstream: upstream,
|
||||
isStickySession: true,
|
||||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, quotaScope AntigravityQuotaScope, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||||
return nil
|
||||
},
|
||||
})
|
||||
|
||||
@@ -13,6 +13,23 @@ import (
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// stubSmartRetryCache 用于 handleSmartRetry 测试的 GatewayCache mock
|
||||
// 仅关注 DeleteSessionAccountID 的调用记录
|
||||
type stubSmartRetryCache struct {
|
||||
GatewayCache // 嵌入接口,未实现的方法 panic(确保只调用预期方法)
|
||||
deleteCalls []deleteSessionCall
|
||||
}
|
||||
|
||||
type deleteSessionCall struct {
|
||||
groupID int64
|
||||
sessionHash string
|
||||
}
|
||||
|
||||
func (c *stubSmartRetryCache) DeleteSessionAccountID(_ context.Context, groupID int64, sessionHash string) error {
|
||||
c.deleteCalls = append(c.deleteCalls, deleteSessionCall{groupID: groupID, sessionHash: sessionHash})
|
||||
return nil
|
||||
}
|
||||
|
||||
// mockSmartRetryUpstream 用于 handleSmartRetry 测试的 mock upstream
|
||||
type mockSmartRetryUpstream struct {
|
||||
responses []*http.Response
|
||||
@@ -58,7 +75,7 @@ func TestHandleSmartRetry_URLLevelRateLimit(t *testing.T) {
|
||||
accessToken: "token",
|
||||
action: "generateContent",
|
||||
body: []byte(`{"input":"test"}`),
|
||||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, quotaScope AntigravityQuotaScope, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||||
return nil
|
||||
},
|
||||
}
|
||||
@@ -110,7 +127,7 @@ func TestHandleSmartRetry_LongDelay_ReturnsSwitchError(t *testing.T) {
|
||||
body: []byte(`{"input":"test"}`),
|
||||
accountRepo: repo,
|
||||
isStickySession: true,
|
||||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, quotaScope AntigravityQuotaScope, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||||
return nil
|
||||
},
|
||||
}
|
||||
@@ -177,7 +194,7 @@ func TestHandleSmartRetry_ShortDelay_SmartRetrySuccess(t *testing.T) {
|
||||
action: "generateContent",
|
||||
body: []byte(`{"input":"test"}`),
|
||||
httpUpstream: upstream,
|
||||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, quotaScope AntigravityQuotaScope, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||||
return nil
|
||||
},
|
||||
}
|
||||
@@ -198,7 +215,7 @@ func TestHandleSmartRetry_ShortDelay_SmartRetrySuccess(t *testing.T) {
|
||||
|
||||
// TestHandleSmartRetry_ShortDelay_SmartRetryFailed_ReturnsSwitchError 测试智能重试失败后返回 switchError
|
||||
func TestHandleSmartRetry_ShortDelay_SmartRetryFailed_ReturnsSwitchError(t *testing.T) {
|
||||
// 智能重试后仍然返回 429(需要提供 3 个响应,因为智能重试最多 3 次)
|
||||
// 智能重试后仍然返回 429(需要提供 1 个响应,因为智能重试最多 1 次)
|
||||
failRespBody := `{
|
||||
"error": {
|
||||
"status": "RESOURCE_EXHAUSTED",
|
||||
@@ -213,19 +230,9 @@ func TestHandleSmartRetry_ShortDelay_SmartRetryFailed_ReturnsSwitchError(t *test
|
||||
Header: http.Header{},
|
||||
Body: io.NopCloser(strings.NewReader(failRespBody)),
|
||||
}
|
||||
failResp2 := &http.Response{
|
||||
StatusCode: http.StatusTooManyRequests,
|
||||
Header: http.Header{},
|
||||
Body: io.NopCloser(strings.NewReader(failRespBody)),
|
||||
}
|
||||
failResp3 := &http.Response{
|
||||
StatusCode: http.StatusTooManyRequests,
|
||||
Header: http.Header{},
|
||||
Body: io.NopCloser(strings.NewReader(failRespBody)),
|
||||
}
|
||||
upstream := &mockSmartRetryUpstream{
|
||||
responses: []*http.Response{failResp1, failResp2, failResp3},
|
||||
errors: []error{nil, nil, nil},
|
||||
responses: []*http.Response{failResp1},
|
||||
errors: []error{nil},
|
||||
}
|
||||
|
||||
repo := &stubAntigravityAccountRepo{}
|
||||
@@ -236,7 +243,7 @@ func TestHandleSmartRetry_ShortDelay_SmartRetryFailed_ReturnsSwitchError(t *test
|
||||
Platform: PlatformAntigravity,
|
||||
}
|
||||
|
||||
// 3s < 7s 阈值,应该触发智能重试(最多 3 次)
|
||||
// 3s < 7s 阈值,应该触发智能重试(最多 1 次)
|
||||
respBody := []byte(`{
|
||||
"error": {
|
||||
"status": "RESOURCE_EXHAUSTED",
|
||||
@@ -262,7 +269,7 @@ func TestHandleSmartRetry_ShortDelay_SmartRetryFailed_ReturnsSwitchError(t *test
|
||||
httpUpstream: upstream,
|
||||
accountRepo: repo,
|
||||
isStickySession: false,
|
||||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, quotaScope AntigravityQuotaScope, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||||
return nil
|
||||
},
|
||||
}
|
||||
@@ -284,7 +291,7 @@ func TestHandleSmartRetry_ShortDelay_SmartRetryFailed_ReturnsSwitchError(t *test
|
||||
// 验证模型限流已设置
|
||||
require.Len(t, repo.modelRateLimitCalls, 1)
|
||||
require.Equal(t, "gemini-3-flash", repo.modelRateLimitCalls[0].modelKey)
|
||||
require.Len(t, upstream.calls, 3, "should have made three retry calls (max attempts)")
|
||||
require.Len(t, upstream.calls, 1, "should have made one retry call (max attempts)")
|
||||
}
|
||||
|
||||
// TestHandleSmartRetry_503_ModelCapacityExhausted_ReturnsSwitchError 测试 503 MODEL_CAPACITY_EXHAUSTED 返回 switchError
|
||||
@@ -324,7 +331,7 @@ func TestHandleSmartRetry_503_ModelCapacityExhausted_ReturnsSwitchError(t *testi
|
||||
body: []byte(`{"input":"test"}`),
|
||||
accountRepo: repo,
|
||||
isStickySession: true,
|
||||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, quotaScope AntigravityQuotaScope, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||||
return nil
|
||||
},
|
||||
}
|
||||
@@ -380,7 +387,7 @@ func TestHandleSmartRetry_NonAntigravityAccount_ContinuesDefaultLogic(t *testing
|
||||
accessToken: "token",
|
||||
action: "generateContent",
|
||||
body: []byte(`{"input":"test"}`),
|
||||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, quotaScope AntigravityQuotaScope, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||||
return nil
|
||||
},
|
||||
}
|
||||
@@ -429,7 +436,7 @@ func TestHandleSmartRetry_NonModelRateLimit_ContinuesDefaultLogic(t *testing.T)
|
||||
accessToken: "token",
|
||||
action: "generateContent",
|
||||
body: []byte(`{"input":"test"}`),
|
||||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, quotaScope AntigravityQuotaScope, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||||
return nil
|
||||
},
|
||||
}
|
||||
@@ -480,7 +487,7 @@ func TestHandleSmartRetry_ExactlyAtThreshold_ReturnsSwitchError(t *testing.T) {
|
||||
action: "generateContent",
|
||||
body: []byte(`{"input":"test"}`),
|
||||
accountRepo: repo,
|
||||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, quotaScope AntigravityQuotaScope, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||||
return nil
|
||||
},
|
||||
}
|
||||
@@ -541,7 +548,7 @@ func TestAntigravityRetryLoop_HandleSmartRetry_SwitchError_Propagates(t *testing
|
||||
httpUpstream: upstream,
|
||||
accountRepo: repo,
|
||||
isStickySession: true,
|
||||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, quotaScope AntigravityQuotaScope, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||||
return nil
|
||||
},
|
||||
})
|
||||
@@ -556,19 +563,15 @@ func TestAntigravityRetryLoop_HandleSmartRetry_SwitchError_Propagates(t *testing
|
||||
require.True(t, switchErr.IsStickySession)
|
||||
}
|
||||
|
||||
// TestHandleSmartRetry_NetworkError_ContinuesRetry 测试网络错误时继续重试
|
||||
func TestHandleSmartRetry_NetworkError_ContinuesRetry(t *testing.T) {
|
||||
// 第一次网络错误,第二次成功
|
||||
successResp := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{},
|
||||
Body: io.NopCloser(strings.NewReader(`{"result":"ok"}`)),
|
||||
}
|
||||
// TestHandleSmartRetry_NetworkError_ExhaustsRetry 测试网络错误时(maxAttempts=1)直接耗尽重试并切换账号
|
||||
func TestHandleSmartRetry_NetworkError_ExhaustsRetry(t *testing.T) {
|
||||
// 唯一一次重试遇到网络错误(nil response)
|
||||
upstream := &mockSmartRetryUpstream{
|
||||
responses: []*http.Response{nil, successResp}, // 第一次返回 nil(模拟网络错误)
|
||||
errors: []error{nil, nil}, // mock 不返回 error,靠 nil response 触发
|
||||
responses: []*http.Response{nil}, // 返回 nil(模拟网络错误)
|
||||
errors: []error{nil}, // mock 不返回 error,靠 nil response 触发
|
||||
}
|
||||
|
||||
repo := &stubAntigravityAccountRepo{}
|
||||
account := &Account{
|
||||
ID: 8,
|
||||
Name: "acc-8",
|
||||
@@ -600,7 +603,8 @@ func TestHandleSmartRetry_NetworkError_ContinuesRetry(t *testing.T) {
|
||||
action: "generateContent",
|
||||
body: []byte(`{"input":"test"}`),
|
||||
httpUpstream: upstream,
|
||||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, quotaScope AntigravityQuotaScope, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||||
accountRepo: repo,
|
||||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||||
return nil
|
||||
},
|
||||
}
|
||||
@@ -612,10 +616,15 @@ func TestHandleSmartRetry_NetworkError_ContinuesRetry(t *testing.T) {
|
||||
|
||||
require.NotNil(t, result)
|
||||
require.Equal(t, smartRetryActionBreakWithResp, result.action)
|
||||
require.NotNil(t, result.resp, "should return successful response after network error recovery")
|
||||
require.Equal(t, http.StatusOK, result.resp.StatusCode)
|
||||
require.Nil(t, result.switchError, "should not return switchError on success")
|
||||
require.Len(t, upstream.calls, 2, "should have made two retry calls")
|
||||
require.Nil(t, result.resp, "should not return resp when switchError is set")
|
||||
require.NotNil(t, result.switchError, "should return switchError after network error exhausted retry")
|
||||
require.Equal(t, account.ID, result.switchError.OriginalAccountID)
|
||||
require.Equal(t, "claude-sonnet-4-5", result.switchError.RateLimitedModel)
|
||||
require.Len(t, upstream.calls, 1, "should have made one retry call")
|
||||
|
||||
// 验证模型限流已设置
|
||||
require.Len(t, repo.modelRateLimitCalls, 1)
|
||||
require.Equal(t, "claude-sonnet-4-5", repo.modelRateLimitCalls[0].modelKey)
|
||||
}
|
||||
|
||||
// TestHandleSmartRetry_NoRetryDelay_UsesDefaultRateLimit 测试无 retryDelay 时使用默认 1 分钟限流
|
||||
@@ -653,7 +662,7 @@ func TestHandleSmartRetry_NoRetryDelay_UsesDefaultRateLimit(t *testing.T) {
|
||||
body: []byte(`{"input":"test"}`),
|
||||
accountRepo: repo,
|
||||
isStickySession: true,
|
||||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, quotaScope AntigravityQuotaScope, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||||
return nil
|
||||
},
|
||||
}
|
||||
@@ -674,3 +683,617 @@ func TestHandleSmartRetry_NoRetryDelay_UsesDefaultRateLimit(t *testing.T) {
|
||||
require.Len(t, repo.modelRateLimitCalls, 1)
|
||||
require.Equal(t, "claude-sonnet-4-5", repo.modelRateLimitCalls[0].modelKey)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// 以下测试覆盖本次改动:
|
||||
// 1. antigravitySmartRetryMaxAttempts = 1(仅重试 1 次)
|
||||
// 2. 智能重试失败后清除粘性会话绑定(DeleteSessionAccountID)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// TestSmartRetryMaxAttempts_VerifyConstant 验证常量值为 1
|
||||
func TestSmartRetryMaxAttempts_VerifyConstant(t *testing.T) {
|
||||
require.Equal(t, 1, antigravitySmartRetryMaxAttempts,
|
||||
"antigravitySmartRetryMaxAttempts should be 1 to prevent repeated rate limiting")
|
||||
}
|
||||
|
||||
// TestHandleSmartRetry_ShortDelay_StickySession_FailedRetry_ClearsSession
|
||||
// 核心场景:粘性会话 + 短延迟重试失败 → 必须清除粘性绑定
|
||||
func TestHandleSmartRetry_ShortDelay_StickySession_FailedRetry_ClearsSession(t *testing.T) {
|
||||
failRespBody := `{
|
||||
"error": {
|
||||
"status": "RESOURCE_EXHAUSTED",
|
||||
"details": [
|
||||
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "claude-sonnet-4-5"}, "reason": "RATE_LIMIT_EXCEEDED"},
|
||||
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
|
||||
]
|
||||
}
|
||||
}`
|
||||
failResp := &http.Response{
|
||||
StatusCode: http.StatusTooManyRequests,
|
||||
Header: http.Header{},
|
||||
Body: io.NopCloser(strings.NewReader(failRespBody)),
|
||||
}
|
||||
upstream := &mockSmartRetryUpstream{
|
||||
responses: []*http.Response{failResp},
|
||||
errors: []error{nil},
|
||||
}
|
||||
|
||||
repo := &stubAntigravityAccountRepo{}
|
||||
cache := &stubSmartRetryCache{}
|
||||
account := &Account{
|
||||
ID: 10,
|
||||
Name: "acc-10",
|
||||
Type: AccountTypeOAuth,
|
||||
Platform: PlatformAntigravity,
|
||||
}
|
||||
|
||||
respBody := []byte(`{
|
||||
"error": {
|
||||
"status": "RESOURCE_EXHAUSTED",
|
||||
"details": [
|
||||
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "claude-sonnet-4-5"}, "reason": "RATE_LIMIT_EXCEEDED"},
|
||||
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
|
||||
]
|
||||
}
|
||||
}`)
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusTooManyRequests,
|
||||
Header: http.Header{},
|
||||
Body: io.NopCloser(bytes.NewReader(respBody)),
|
||||
}
|
||||
|
||||
params := antigravityRetryLoopParams{
|
||||
ctx: context.Background(),
|
||||
prefix: "[test]",
|
||||
account: account,
|
||||
accessToken: "token",
|
||||
action: "generateContent",
|
||||
body: []byte(`{"input":"test"}`),
|
||||
httpUpstream: upstream,
|
||||
accountRepo: repo,
|
||||
isStickySession: true,
|
||||
groupID: 42,
|
||||
sessionHash: "sticky-hash-abc",
|
||||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||||
return nil
|
||||
},
|
||||
}
|
||||
|
||||
availableURLs := []string{"https://ag-1.test"}
|
||||
|
||||
svc := &AntigravityGatewayService{cache: cache}
|
||||
result := svc.handleSmartRetry(params, resp, respBody, "https://ag-1.test", 0, availableURLs)
|
||||
|
||||
// 验证返回 switchError
|
||||
require.NotNil(t, result)
|
||||
require.Equal(t, smartRetryActionBreakWithResp, result.action)
|
||||
require.NotNil(t, result.switchError)
|
||||
require.True(t, result.switchError.IsStickySession, "switchError should carry IsStickySession=true")
|
||||
require.Equal(t, account.ID, result.switchError.OriginalAccountID)
|
||||
|
||||
// 核心断言:DeleteSessionAccountID 被调用,且参数正确
|
||||
require.Len(t, cache.deleteCalls, 1, "should call DeleteSessionAccountID exactly once")
|
||||
require.Equal(t, int64(42), cache.deleteCalls[0].groupID)
|
||||
require.Equal(t, "sticky-hash-abc", cache.deleteCalls[0].sessionHash)
|
||||
|
||||
// 验证仅重试 1 次
|
||||
require.Len(t, upstream.calls, 1, "should make exactly 1 retry call (maxAttempts=1)")
|
||||
|
||||
// 验证模型限流已设置
|
||||
require.Len(t, repo.modelRateLimitCalls, 1)
|
||||
require.Equal(t, "claude-sonnet-4-5", repo.modelRateLimitCalls[0].modelKey)
|
||||
}
|
||||
|
||||
// TestHandleSmartRetry_ShortDelay_NonStickySession_FailedRetry_NoDeleteSession
|
||||
// 非粘性会话 + 短延迟重试失败 → 不应调用 DeleteSessionAccountID(sessionHash 为空)
|
||||
func TestHandleSmartRetry_ShortDelay_NonStickySession_FailedRetry_NoDeleteSession(t *testing.T) {
|
||||
failRespBody := `{
|
||||
"error": {
|
||||
"status": "RESOURCE_EXHAUSTED",
|
||||
"details": [
|
||||
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-flash"}, "reason": "RATE_LIMIT_EXCEEDED"},
|
||||
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
|
||||
]
|
||||
}
|
||||
}`
|
||||
failResp := &http.Response{
|
||||
StatusCode: http.StatusTooManyRequests,
|
||||
Header: http.Header{},
|
||||
Body: io.NopCloser(strings.NewReader(failRespBody)),
|
||||
}
|
||||
upstream := &mockSmartRetryUpstream{
|
||||
responses: []*http.Response{failResp},
|
||||
errors: []error{nil},
|
||||
}
|
||||
|
||||
repo := &stubAntigravityAccountRepo{}
|
||||
cache := &stubSmartRetryCache{}
|
||||
account := &Account{
|
||||
ID: 11,
|
||||
Name: "acc-11",
|
||||
Type: AccountTypeOAuth,
|
||||
Platform: PlatformAntigravity,
|
||||
}
|
||||
|
||||
respBody := []byte(`{
|
||||
"error": {
|
||||
"status": "RESOURCE_EXHAUSTED",
|
||||
"details": [
|
||||
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-flash"}, "reason": "RATE_LIMIT_EXCEEDED"},
|
||||
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
|
||||
]
|
||||
}
|
||||
}`)
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusTooManyRequests,
|
||||
Header: http.Header{},
|
||||
Body: io.NopCloser(bytes.NewReader(respBody)),
|
||||
}
|
||||
|
||||
params := antigravityRetryLoopParams{
|
||||
ctx: context.Background(),
|
||||
prefix: "[test]",
|
||||
account: account,
|
||||
accessToken: "token",
|
||||
action: "generateContent",
|
||||
body: []byte(`{"input":"test"}`),
|
||||
httpUpstream: upstream,
|
||||
accountRepo: repo,
|
||||
isStickySession: false,
|
||||
groupID: 42,
|
||||
sessionHash: "", // 非粘性会话,sessionHash 为空
|
||||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||||
return nil
|
||||
},
|
||||
}
|
||||
|
||||
availableURLs := []string{"https://ag-1.test"}
|
||||
|
||||
svc := &AntigravityGatewayService{cache: cache}
|
||||
result := svc.handleSmartRetry(params, resp, respBody, "https://ag-1.test", 0, availableURLs)
|
||||
|
||||
require.NotNil(t, result)
|
||||
require.Equal(t, smartRetryActionBreakWithResp, result.action)
|
||||
require.NotNil(t, result.switchError)
|
||||
require.False(t, result.switchError.IsStickySession)
|
||||
|
||||
// 核心断言:sessionHash 为空时不应调用 DeleteSessionAccountID
|
||||
require.Len(t, cache.deleteCalls, 0, "should NOT call DeleteSessionAccountID when sessionHash is empty")
|
||||
}
|
||||
|
||||
// TestHandleSmartRetry_ShortDelay_StickySession_FailedRetry_NilCache_NoPanic
|
||||
// 边界:cache 为 nil 时不应 panic
|
||||
func TestHandleSmartRetry_ShortDelay_StickySession_FailedRetry_NilCache_NoPanic(t *testing.T) {
|
||||
failRespBody := `{
|
||||
"error": {
|
||||
"status": "RESOURCE_EXHAUSTED",
|
||||
"details": [
|
||||
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "claude-sonnet-4-5"}, "reason": "RATE_LIMIT_EXCEEDED"},
|
||||
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
|
||||
]
|
||||
}
|
||||
}`
|
||||
failResp := &http.Response{
|
||||
StatusCode: http.StatusTooManyRequests,
|
||||
Header: http.Header{},
|
||||
Body: io.NopCloser(strings.NewReader(failRespBody)),
|
||||
}
|
||||
upstream := &mockSmartRetryUpstream{
|
||||
responses: []*http.Response{failResp},
|
||||
errors: []error{nil},
|
||||
}
|
||||
|
||||
repo := &stubAntigravityAccountRepo{}
|
||||
account := &Account{
|
||||
ID: 12,
|
||||
Name: "acc-12",
|
||||
Type: AccountTypeOAuth,
|
||||
Platform: PlatformAntigravity,
|
||||
}
|
||||
|
||||
respBody := []byte(`{
|
||||
"error": {
|
||||
"status": "RESOURCE_EXHAUSTED",
|
||||
"details": [
|
||||
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "claude-sonnet-4-5"}, "reason": "RATE_LIMIT_EXCEEDED"},
|
||||
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
|
||||
]
|
||||
}
|
||||
}`)
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusTooManyRequests,
|
||||
Header: http.Header{},
|
||||
Body: io.NopCloser(bytes.NewReader(respBody)),
|
||||
}
|
||||
|
||||
params := antigravityRetryLoopParams{
|
||||
ctx: context.Background(),
|
||||
prefix: "[test]",
|
||||
account: account,
|
||||
accessToken: "token",
|
||||
action: "generateContent",
|
||||
body: []byte(`{"input":"test"}`),
|
||||
httpUpstream: upstream,
|
||||
accountRepo: repo,
|
||||
isStickySession: true,
|
||||
groupID: 42,
|
||||
sessionHash: "sticky-hash-nil-cache",
|
||||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||||
return nil
|
||||
},
|
||||
}
|
||||
|
||||
availableURLs := []string{"https://ag-1.test"}
|
||||
|
||||
// cache 为 nil,不应 panic
|
||||
svc := &AntigravityGatewayService{cache: nil}
|
||||
require.NotPanics(t, func() {
|
||||
result := svc.handleSmartRetry(params, resp, respBody, "https://ag-1.test", 0, availableURLs)
|
||||
require.NotNil(t, result)
|
||||
require.Equal(t, smartRetryActionBreakWithResp, result.action)
|
||||
require.NotNil(t, result.switchError)
|
||||
require.True(t, result.switchError.IsStickySession)
|
||||
})
|
||||
}
|
||||
|
||||
// TestHandleSmartRetry_ShortDelay_StickySession_SuccessRetry_NoDeleteSession
|
||||
// 重试成功时不应清除粘性会话(只有失败才清除)
|
||||
func TestHandleSmartRetry_ShortDelay_StickySession_SuccessRetry_NoDeleteSession(t *testing.T) {
|
||||
successResp := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{},
|
||||
Body: io.NopCloser(strings.NewReader(`{"result":"ok"}`)),
|
||||
}
|
||||
upstream := &mockSmartRetryUpstream{
|
||||
responses: []*http.Response{successResp},
|
||||
errors: []error{nil},
|
||||
}
|
||||
|
||||
cache := &stubSmartRetryCache{}
|
||||
account := &Account{
|
||||
ID: 13,
|
||||
Name: "acc-13",
|
||||
Type: AccountTypeOAuth,
|
||||
Platform: PlatformAntigravity,
|
||||
}
|
||||
|
||||
respBody := []byte(`{
|
||||
"error": {
|
||||
"status": "RESOURCE_EXHAUSTED",
|
||||
"details": [
|
||||
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "claude-opus-4"}, "reason": "RATE_LIMIT_EXCEEDED"},
|
||||
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.5s"}
|
||||
]
|
||||
}
|
||||
}`)
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusTooManyRequests,
|
||||
Header: http.Header{},
|
||||
Body: io.NopCloser(bytes.NewReader(respBody)),
|
||||
}
|
||||
|
||||
params := antigravityRetryLoopParams{
|
||||
ctx: context.Background(),
|
||||
prefix: "[test]",
|
||||
account: account,
|
||||
accessToken: "token",
|
||||
action: "generateContent",
|
||||
body: []byte(`{"input":"test"}`),
|
||||
httpUpstream: upstream,
|
||||
isStickySession: true,
|
||||
groupID: 42,
|
||||
sessionHash: "sticky-hash-success",
|
||||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||||
return nil
|
||||
},
|
||||
}
|
||||
|
||||
availableURLs := []string{"https://ag-1.test"}
|
||||
|
||||
svc := &AntigravityGatewayService{cache: cache}
|
||||
result := svc.handleSmartRetry(params, resp, respBody, "https://ag-1.test", 0, availableURLs)
|
||||
|
||||
require.NotNil(t, result)
|
||||
require.Equal(t, smartRetryActionBreakWithResp, result.action)
|
||||
require.NotNil(t, result.resp, "should return successful response")
|
||||
require.Equal(t, http.StatusOK, result.resp.StatusCode)
|
||||
require.Nil(t, result.switchError, "should not return switchError on success")
|
||||
|
||||
// 核心断言:重试成功时不应清除粘性会话
|
||||
require.Len(t, cache.deleteCalls, 0, "should NOT call DeleteSessionAccountID on successful retry")
|
||||
}
|
||||
|
||||
// TestHandleSmartRetry_LongDelay_StickySession_NoDeleteInHandleSmartRetry
|
||||
// 长延迟路径(情况1)在 handleSmartRetry 中不直接调用 DeleteSessionAccountID
|
||||
// (清除由 handler 层的 shouldClearStickySession 在下次请求时处理)
|
||||
func TestHandleSmartRetry_LongDelay_StickySession_NoDeleteInHandleSmartRetry(t *testing.T) {
|
||||
repo := &stubAntigravityAccountRepo{}
|
||||
cache := &stubSmartRetryCache{}
|
||||
account := &Account{
|
||||
ID: 14,
|
||||
Name: "acc-14",
|
||||
Type: AccountTypeOAuth,
|
||||
Platform: PlatformAntigravity,
|
||||
}
|
||||
|
||||
// 15s >= 7s 阈值 → 走长延迟路径
|
||||
respBody := []byte(`{
|
||||
"error": {
|
||||
"status": "RESOURCE_EXHAUSTED",
|
||||
"details": [
|
||||
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "claude-sonnet-4-5"}, "reason": "RATE_LIMIT_EXCEEDED"},
|
||||
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "15s"}
|
||||
]
|
||||
}
|
||||
}`)
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusTooManyRequests,
|
||||
Header: http.Header{},
|
||||
Body: io.NopCloser(bytes.NewReader(respBody)),
|
||||
}
|
||||
|
||||
params := antigravityRetryLoopParams{
|
||||
ctx: context.Background(),
|
||||
prefix: "[test]",
|
||||
account: account,
|
||||
accessToken: "token",
|
||||
action: "generateContent",
|
||||
body: []byte(`{"input":"test"}`),
|
||||
accountRepo: repo,
|
||||
isStickySession: true,
|
||||
groupID: 42,
|
||||
sessionHash: "sticky-hash-long-delay",
|
||||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||||
return nil
|
||||
},
|
||||
}
|
||||
|
||||
availableURLs := []string{"https://ag-1.test"}
|
||||
|
||||
svc := &AntigravityGatewayService{cache: cache}
|
||||
result := svc.handleSmartRetry(params, resp, respBody, "https://ag-1.test", 0, availableURLs)
|
||||
|
||||
require.NotNil(t, result)
|
||||
require.Equal(t, smartRetryActionBreakWithResp, result.action)
|
||||
require.NotNil(t, result.switchError)
|
||||
require.True(t, result.switchError.IsStickySession)
|
||||
|
||||
// 长延迟路径不在 handleSmartRetry 中调用 DeleteSessionAccountID
|
||||
// (由上游 handler 的 shouldClearStickySession 处理)
|
||||
require.Len(t, cache.deleteCalls, 0,
|
||||
"long delay path should NOT call DeleteSessionAccountID in handleSmartRetry (handled by handler layer)")
|
||||
}
|
||||
|
||||
// TestHandleSmartRetry_ShortDelay_NetworkError_StickySession_ClearsSession
|
||||
// 网络错误耗尽重试 + 粘性会话 → 也应清除粘性绑定
|
||||
func TestHandleSmartRetry_ShortDelay_NetworkError_StickySession_ClearsSession(t *testing.T) {
|
||||
upstream := &mockSmartRetryUpstream{
|
||||
responses: []*http.Response{nil}, // 网络错误
|
||||
errors: []error{nil},
|
||||
}
|
||||
|
||||
repo := &stubAntigravityAccountRepo{}
|
||||
cache := &stubSmartRetryCache{}
|
||||
account := &Account{
|
||||
ID: 15,
|
||||
Name: "acc-15",
|
||||
Type: AccountTypeOAuth,
|
||||
Platform: PlatformAntigravity,
|
||||
}
|
||||
|
||||
respBody := []byte(`{
|
||||
"error": {
|
||||
"status": "RESOURCE_EXHAUSTED",
|
||||
"details": [
|
||||
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-flash"}, "reason": "RATE_LIMIT_EXCEEDED"},
|
||||
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
|
||||
]
|
||||
}
|
||||
}`)
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusTooManyRequests,
|
||||
Header: http.Header{},
|
||||
Body: io.NopCloser(bytes.NewReader(respBody)),
|
||||
}
|
||||
|
||||
params := antigravityRetryLoopParams{
|
||||
ctx: context.Background(),
|
||||
prefix: "[test]",
|
||||
account: account,
|
||||
accessToken: "token",
|
||||
action: "generateContent",
|
||||
body: []byte(`{"input":"test"}`),
|
||||
httpUpstream: upstream,
|
||||
accountRepo: repo,
|
||||
isStickySession: true,
|
||||
groupID: 99,
|
||||
sessionHash: "sticky-net-error",
|
||||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||||
return nil
|
||||
},
|
||||
}
|
||||
|
||||
availableURLs := []string{"https://ag-1.test"}
|
||||
|
||||
svc := &AntigravityGatewayService{cache: cache}
|
||||
result := svc.handleSmartRetry(params, resp, respBody, "https://ag-1.test", 0, availableURLs)
|
||||
|
||||
require.NotNil(t, result)
|
||||
require.NotNil(t, result.switchError)
|
||||
require.True(t, result.switchError.IsStickySession)
|
||||
|
||||
// 核心断言:网络错误耗尽重试后也应清除粘性绑定
|
||||
require.Len(t, cache.deleteCalls, 1, "should call DeleteSessionAccountID after network error exhausts retry")
|
||||
require.Equal(t, int64(99), cache.deleteCalls[0].groupID)
|
||||
require.Equal(t, "sticky-net-error", cache.deleteCalls[0].sessionHash)
|
||||
}
|
||||
|
||||
// TestHandleSmartRetry_ShortDelay_503_StickySession_FailedRetry_ClearsSession
|
||||
// 503 + 短延迟 + 粘性会话 + 重试失败 → 清除粘性绑定
|
||||
func TestHandleSmartRetry_ShortDelay_503_StickySession_FailedRetry_ClearsSession(t *testing.T) {
|
||||
failRespBody := `{
|
||||
"error": {
|
||||
"code": 503,
|
||||
"status": "UNAVAILABLE",
|
||||
"details": [
|
||||
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-pro"}, "reason": "MODEL_CAPACITY_EXHAUSTED"},
|
||||
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.5s"}
|
||||
]
|
||||
}
|
||||
}`
|
||||
failResp := &http.Response{
|
||||
StatusCode: http.StatusServiceUnavailable,
|
||||
Header: http.Header{},
|
||||
Body: io.NopCloser(strings.NewReader(failRespBody)),
|
||||
}
|
||||
upstream := &mockSmartRetryUpstream{
|
||||
responses: []*http.Response{failResp},
|
||||
errors: []error{nil},
|
||||
}
|
||||
|
||||
repo := &stubAntigravityAccountRepo{}
|
||||
cache := &stubSmartRetryCache{}
|
||||
account := &Account{
|
||||
ID: 16,
|
||||
Name: "acc-16",
|
||||
Type: AccountTypeOAuth,
|
||||
Platform: PlatformAntigravity,
|
||||
}
|
||||
|
||||
respBody := []byte(`{
|
||||
"error": {
|
||||
"code": 503,
|
||||
"status": "UNAVAILABLE",
|
||||
"details": [
|
||||
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-pro"}, "reason": "MODEL_CAPACITY_EXHAUSTED"},
|
||||
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.5s"}
|
||||
]
|
||||
}
|
||||
}`)
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusServiceUnavailable,
|
||||
Header: http.Header{},
|
||||
Body: io.NopCloser(bytes.NewReader(respBody)),
|
||||
}
|
||||
|
||||
params := antigravityRetryLoopParams{
|
||||
ctx: context.Background(),
|
||||
prefix: "[test]",
|
||||
account: account,
|
||||
accessToken: "token",
|
||||
action: "generateContent",
|
||||
body: []byte(`{"input":"test"}`),
|
||||
httpUpstream: upstream,
|
||||
accountRepo: repo,
|
||||
isStickySession: true,
|
||||
groupID: 77,
|
||||
sessionHash: "sticky-503-short",
|
||||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||||
return nil
|
||||
},
|
||||
}
|
||||
|
||||
availableURLs := []string{"https://ag-1.test"}
|
||||
|
||||
svc := &AntigravityGatewayService{cache: cache}
|
||||
result := svc.handleSmartRetry(params, resp, respBody, "https://ag-1.test", 0, availableURLs)
|
||||
|
||||
require.NotNil(t, result)
|
||||
require.NotNil(t, result.switchError)
|
||||
require.True(t, result.switchError.IsStickySession)
|
||||
|
||||
// 验证粘性绑定被清除
|
||||
require.Len(t, cache.deleteCalls, 1)
|
||||
require.Equal(t, int64(77), cache.deleteCalls[0].groupID)
|
||||
require.Equal(t, "sticky-503-short", cache.deleteCalls[0].sessionHash)
|
||||
|
||||
// 验证模型限流已设置
|
||||
require.Len(t, repo.modelRateLimitCalls, 1)
|
||||
require.Equal(t, "gemini-3-pro", repo.modelRateLimitCalls[0].modelKey)
|
||||
}
|
||||
|
||||
// TestAntigravityRetryLoop_SmartRetryFailed_StickySession_SwitchErrorPropagates
|
||||
// 集成测试:antigravityRetryLoop → handleSmartRetry → switchError 传播
|
||||
// 验证 IsStickySession 正确传递到上层,且粘性绑定被清除
|
||||
func TestAntigravityRetryLoop_SmartRetryFailed_StickySession_SwitchErrorPropagates(t *testing.T) {
|
||||
// 初始 429 响应
|
||||
initialRespBody := []byte(`{
|
||||
"error": {
|
||||
"status": "RESOURCE_EXHAUSTED",
|
||||
"details": [
|
||||
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "claude-opus-4-6"}, "reason": "RATE_LIMIT_EXCEEDED"},
|
||||
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
|
||||
]
|
||||
}
|
||||
}`)
|
||||
initialResp := &http.Response{
|
||||
StatusCode: http.StatusTooManyRequests,
|
||||
Header: http.Header{},
|
||||
Body: io.NopCloser(bytes.NewReader(initialRespBody)),
|
||||
}
|
||||
|
||||
// 智能重试也返回 429
|
||||
retryRespBody := `{
|
||||
"error": {
|
||||
"status": "RESOURCE_EXHAUSTED",
|
||||
"details": [
|
||||
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "claude-opus-4-6"}, "reason": "RATE_LIMIT_EXCEEDED"},
|
||||
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
|
||||
]
|
||||
}
|
||||
}`
|
||||
retryResp := &http.Response{
|
||||
StatusCode: http.StatusTooManyRequests,
|
||||
Header: http.Header{},
|
||||
Body: io.NopCloser(strings.NewReader(retryRespBody)),
|
||||
}
|
||||
|
||||
upstream := &mockSmartRetryUpstream{
|
||||
responses: []*http.Response{initialResp, retryResp},
|
||||
errors: []error{nil, nil},
|
||||
}
|
||||
|
||||
repo := &stubAntigravityAccountRepo{}
|
||||
cache := &stubSmartRetryCache{}
|
||||
account := &Account{
|
||||
ID: 17,
|
||||
Name: "acc-17",
|
||||
Type: AccountTypeOAuth,
|
||||
Platform: PlatformAntigravity,
|
||||
Schedulable: true,
|
||||
Status: StatusActive,
|
||||
Concurrency: 1,
|
||||
}
|
||||
|
||||
svc := &AntigravityGatewayService{cache: cache}
|
||||
result, err := svc.antigravityRetryLoop(antigravityRetryLoopParams{
|
||||
ctx: context.Background(),
|
||||
prefix: "[test]",
|
||||
account: account,
|
||||
accessToken: "token",
|
||||
action: "generateContent",
|
||||
body: []byte(`{"input":"test"}`),
|
||||
httpUpstream: upstream,
|
||||
accountRepo: repo,
|
||||
isStickySession: true,
|
||||
groupID: 55,
|
||||
sessionHash: "sticky-loop-test",
|
||||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||||
return nil
|
||||
},
|
||||
})
|
||||
|
||||
require.Nil(t, result, "should not return result when switchError")
|
||||
require.NotNil(t, err, "should return error")
|
||||
|
||||
var switchErr *AntigravityAccountSwitchError
|
||||
require.ErrorAs(t, err, &switchErr, "error should be AntigravityAccountSwitchError")
|
||||
require.Equal(t, account.ID, switchErr.OriginalAccountID)
|
||||
require.Equal(t, "claude-opus-4-6", switchErr.RateLimitedModel)
|
||||
require.True(t, switchErr.IsStickySession, "IsStickySession must propagate through retryLoop")
|
||||
|
||||
// 验证粘性绑定被清除
|
||||
require.Len(t, cache.deleteCalls, 1, "should clear sticky session in handleSmartRetry")
|
||||
require.Equal(t, int64(55), cache.deleteCalls[0].groupID)
|
||||
require.Equal(t, "sticky-loop-test", cache.deleteCalls[0].sessionHash)
|
||||
}
|
||||
|
||||
112
backend/internal/service/crs_sync_helpers_test.go
Normal file
112
backend/internal/service/crs_sync_helpers_test.go
Normal file
@@ -0,0 +1,112 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestBuildSelectedSet(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
ids []string
|
||||
wantNil bool
|
||||
wantSize int
|
||||
}{
|
||||
{
|
||||
name: "nil input returns nil (backward compatible: create all)",
|
||||
ids: nil,
|
||||
wantNil: true,
|
||||
},
|
||||
{
|
||||
name: "empty slice returns empty map (create none)",
|
||||
ids: []string{},
|
||||
wantNil: false,
|
||||
wantSize: 0,
|
||||
},
|
||||
{
|
||||
name: "single ID",
|
||||
ids: []string{"abc-123"},
|
||||
wantNil: false,
|
||||
wantSize: 1,
|
||||
},
|
||||
{
|
||||
name: "multiple IDs",
|
||||
ids: []string{"a", "b", "c"},
|
||||
wantNil: false,
|
||||
wantSize: 3,
|
||||
},
|
||||
{
|
||||
name: "duplicate IDs are deduplicated",
|
||||
ids: []string{"a", "a", "b"},
|
||||
wantNil: false,
|
||||
wantSize: 2,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := buildSelectedSet(tt.ids)
|
||||
if tt.wantNil {
|
||||
if got != nil {
|
||||
t.Errorf("buildSelectedSet(%v) = %v, want nil", tt.ids, got)
|
||||
}
|
||||
return
|
||||
}
|
||||
if got == nil {
|
||||
t.Fatalf("buildSelectedSet(%v) = nil, want non-nil map", tt.ids)
|
||||
}
|
||||
if len(got) != tt.wantSize {
|
||||
t.Errorf("buildSelectedSet(%v) has %d entries, want %d", tt.ids, len(got), tt.wantSize)
|
||||
}
|
||||
// Verify all unique IDs are present
|
||||
for _, id := range tt.ids {
|
||||
if _, ok := got[id]; !ok {
|
||||
t.Errorf("buildSelectedSet(%v) missing key %q", tt.ids, id)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestShouldCreateAccount(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
crsID string
|
||||
selectedSet map[string]struct{}
|
||||
want bool
|
||||
}{
|
||||
{
|
||||
name: "nil set allows all (backward compatible)",
|
||||
crsID: "any-id",
|
||||
selectedSet: nil,
|
||||
want: true,
|
||||
},
|
||||
{
|
||||
name: "empty set blocks all",
|
||||
crsID: "any-id",
|
||||
selectedSet: map[string]struct{}{},
|
||||
want: false,
|
||||
},
|
||||
{
|
||||
name: "ID in set is allowed",
|
||||
crsID: "abc-123",
|
||||
selectedSet: map[string]struct{}{"abc-123": {}, "def-456": {}},
|
||||
want: true,
|
||||
},
|
||||
{
|
||||
name: "ID not in set is blocked",
|
||||
crsID: "xyz-789",
|
||||
selectedSet: map[string]struct{}{"abc-123": {}, "def-456": {}},
|
||||
want: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := shouldCreateAccount(tt.crsID, tt.selectedSet)
|
||||
if got != tt.want {
|
||||
t.Errorf("shouldCreateAccount(%q, %v) = %v, want %v",
|
||||
tt.crsID, tt.selectedSet, got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -45,10 +45,11 @@ func NewCRSSyncService(
|
||||
}
|
||||
|
||||
type SyncFromCRSInput struct {
|
||||
BaseURL string
|
||||
Username string
|
||||
Password string
|
||||
SyncProxies bool
|
||||
BaseURL string
|
||||
Username string
|
||||
Password string
|
||||
SyncProxies bool
|
||||
SelectedAccountIDs []string // if non-empty, only create new accounts with these CRS IDs
|
||||
}
|
||||
|
||||
type SyncFromCRSItemResult struct {
|
||||
@@ -190,25 +191,27 @@ type crsGeminiAPIKeyAccount struct {
|
||||
Extra map[string]any `json:"extra"`
|
||||
}
|
||||
|
||||
func (s *CRSSyncService) SyncFromCRS(ctx context.Context, input SyncFromCRSInput) (*SyncFromCRSResult, error) {
|
||||
// fetchCRSExport validates the connection parameters, authenticates with CRS,
|
||||
// and returns the exported accounts. Shared by SyncFromCRS and PreviewFromCRS.
|
||||
func (s *CRSSyncService) fetchCRSExport(ctx context.Context, baseURL, username, password string) (*crsExportResponse, error) {
|
||||
if s.cfg == nil {
|
||||
return nil, errors.New("config is not available")
|
||||
}
|
||||
baseURL := strings.TrimSpace(input.BaseURL)
|
||||
normalizedURL := strings.TrimSpace(baseURL)
|
||||
if s.cfg.Security.URLAllowlist.Enabled {
|
||||
normalized, err := normalizeBaseURL(baseURL, s.cfg.Security.URLAllowlist.CRSHosts, s.cfg.Security.URLAllowlist.AllowPrivateHosts)
|
||||
normalized, err := normalizeBaseURL(normalizedURL, s.cfg.Security.URLAllowlist.CRSHosts, s.cfg.Security.URLAllowlist.AllowPrivateHosts)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
baseURL = normalized
|
||||
normalizedURL = normalized
|
||||
} else {
|
||||
normalized, err := urlvalidator.ValidateURLFormat(baseURL, s.cfg.Security.URLAllowlist.AllowInsecureHTTP)
|
||||
normalized, err := urlvalidator.ValidateURLFormat(normalizedURL, s.cfg.Security.URLAllowlist.AllowInsecureHTTP)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("invalid base_url: %w", err)
|
||||
}
|
||||
baseURL = normalized
|
||||
normalizedURL = normalized
|
||||
}
|
||||
if strings.TrimSpace(input.Username) == "" || strings.TrimSpace(input.Password) == "" {
|
||||
if strings.TrimSpace(username) == "" || strings.TrimSpace(password) == "" {
|
||||
return nil, errors.New("username and password are required")
|
||||
}
|
||||
|
||||
@@ -221,12 +224,16 @@ func (s *CRSSyncService) SyncFromCRS(ctx context.Context, input SyncFromCRSInput
|
||||
client = &http.Client{Timeout: 20 * time.Second}
|
||||
}
|
||||
|
||||
adminToken, err := crsLogin(ctx, client, baseURL, input.Username, input.Password)
|
||||
adminToken, err := crsLogin(ctx, client, normalizedURL, username, password)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
exported, err := crsExportAccounts(ctx, client, baseURL, adminToken)
|
||||
return crsExportAccounts(ctx, client, normalizedURL, adminToken)
|
||||
}
|
||||
|
||||
func (s *CRSSyncService) SyncFromCRS(ctx context.Context, input SyncFromCRSInput) (*SyncFromCRSResult, error) {
|
||||
exported, err := s.fetchCRSExport(ctx, input.BaseURL, input.Username, input.Password)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -241,6 +248,8 @@ func (s *CRSSyncService) SyncFromCRS(ctx context.Context, input SyncFromCRSInput
|
||||
),
|
||||
}
|
||||
|
||||
selectedSet := buildSelectedSet(input.SelectedAccountIDs)
|
||||
|
||||
var proxies []Proxy
|
||||
if input.SyncProxies {
|
||||
proxies, _ = s.proxyRepo.ListActive(ctx)
|
||||
@@ -329,6 +338,13 @@ func (s *CRSSyncService) SyncFromCRS(ctx context.Context, input SyncFromCRSInput
|
||||
}
|
||||
|
||||
if existing == nil {
|
||||
if !shouldCreateAccount(src.ID, selectedSet) {
|
||||
item.Action = "skipped"
|
||||
item.Error = "not selected"
|
||||
result.Skipped++
|
||||
result.Items = append(result.Items, item)
|
||||
continue
|
||||
}
|
||||
account := &Account{
|
||||
Name: defaultName(src.Name, src.ID),
|
||||
Platform: PlatformAnthropic,
|
||||
@@ -446,6 +462,13 @@ func (s *CRSSyncService) SyncFromCRS(ctx context.Context, input SyncFromCRSInput
|
||||
}
|
||||
|
||||
if existing == nil {
|
||||
if !shouldCreateAccount(src.ID, selectedSet) {
|
||||
item.Action = "skipped"
|
||||
item.Error = "not selected"
|
||||
result.Skipped++
|
||||
result.Items = append(result.Items, item)
|
||||
continue
|
||||
}
|
||||
account := &Account{
|
||||
Name: defaultName(src.Name, src.ID),
|
||||
Platform: PlatformAnthropic,
|
||||
@@ -569,6 +592,13 @@ func (s *CRSSyncService) SyncFromCRS(ctx context.Context, input SyncFromCRSInput
|
||||
}
|
||||
|
||||
if existing == nil {
|
||||
if !shouldCreateAccount(src.ID, selectedSet) {
|
||||
item.Action = "skipped"
|
||||
item.Error = "not selected"
|
||||
result.Skipped++
|
||||
result.Items = append(result.Items, item)
|
||||
continue
|
||||
}
|
||||
account := &Account{
|
||||
Name: defaultName(src.Name, src.ID),
|
||||
Platform: PlatformOpenAI,
|
||||
@@ -690,6 +720,13 @@ func (s *CRSSyncService) SyncFromCRS(ctx context.Context, input SyncFromCRSInput
|
||||
}
|
||||
|
||||
if existing == nil {
|
||||
if !shouldCreateAccount(src.ID, selectedSet) {
|
||||
item.Action = "skipped"
|
||||
item.Error = "not selected"
|
||||
result.Skipped++
|
||||
result.Items = append(result.Items, item)
|
||||
continue
|
||||
}
|
||||
account := &Account{
|
||||
Name: defaultName(src.Name, src.ID),
|
||||
Platform: PlatformOpenAI,
|
||||
@@ -798,6 +835,13 @@ func (s *CRSSyncService) SyncFromCRS(ctx context.Context, input SyncFromCRSInput
|
||||
}
|
||||
|
||||
if existing == nil {
|
||||
if !shouldCreateAccount(src.ID, selectedSet) {
|
||||
item.Action = "skipped"
|
||||
item.Error = "not selected"
|
||||
result.Skipped++
|
||||
result.Items = append(result.Items, item)
|
||||
continue
|
||||
}
|
||||
account := &Account{
|
||||
Name: defaultName(src.Name, src.ID),
|
||||
Platform: PlatformGemini,
|
||||
@@ -909,6 +953,13 @@ func (s *CRSSyncService) SyncFromCRS(ctx context.Context, input SyncFromCRSInput
|
||||
}
|
||||
|
||||
if existing == nil {
|
||||
if !shouldCreateAccount(src.ID, selectedSet) {
|
||||
item.Action = "skipped"
|
||||
item.Error = "not selected"
|
||||
result.Skipped++
|
||||
result.Items = append(result.Items, item)
|
||||
continue
|
||||
}
|
||||
account := &Account{
|
||||
Name: defaultName(src.Name, src.ID),
|
||||
Platform: PlatformGemini,
|
||||
@@ -1253,3 +1304,102 @@ func (s *CRSSyncService) refreshOAuthToken(ctx context.Context, account *Account
|
||||
|
||||
return newCredentials
|
||||
}
|
||||
|
||||
// buildSelectedSet converts a slice of selected CRS account IDs to a set for O(1) lookup.
|
||||
// Returns nil if ids is nil (field not sent → backward compatible: create all).
|
||||
// Returns an empty map if ids is non-nil but empty (user selected none → create none).
|
||||
func buildSelectedSet(ids []string) map[string]struct{} {
|
||||
if ids == nil {
|
||||
return nil
|
||||
}
|
||||
set := make(map[string]struct{}, len(ids))
|
||||
for _, id := range ids {
|
||||
set[id] = struct{}{}
|
||||
}
|
||||
return set
|
||||
}
|
||||
|
||||
// shouldCreateAccount checks if a new CRS account should be created based on user selection.
|
||||
// Returns true if selectedSet is nil (backward compatible: create all) or if crsID is in the set.
|
||||
func shouldCreateAccount(crsID string, selectedSet map[string]struct{}) bool {
|
||||
if selectedSet == nil {
|
||||
return true
|
||||
}
|
||||
_, ok := selectedSet[crsID]
|
||||
return ok
|
||||
}
|
||||
|
||||
// PreviewFromCRSResult contains the preview of accounts from CRS before sync.
|
||||
type PreviewFromCRSResult struct {
|
||||
NewAccounts []CRSPreviewAccount `json:"new_accounts"`
|
||||
ExistingAccounts []CRSPreviewAccount `json:"existing_accounts"`
|
||||
}
|
||||
|
||||
// CRSPreviewAccount represents a single account in the preview result.
|
||||
type CRSPreviewAccount struct {
|
||||
CRSAccountID string `json:"crs_account_id"`
|
||||
Kind string `json:"kind"`
|
||||
Name string `json:"name"`
|
||||
Platform string `json:"platform"`
|
||||
Type string `json:"type"`
|
||||
}
|
||||
|
||||
// PreviewFromCRS connects to CRS, fetches all accounts, and classifies them
|
||||
// as new or existing by batch-querying local crs_account_id mappings.
|
||||
func (s *CRSSyncService) PreviewFromCRS(ctx context.Context, input SyncFromCRSInput) (*PreviewFromCRSResult, error) {
|
||||
exported, err := s.fetchCRSExport(ctx, input.BaseURL, input.Username, input.Password)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Batch query all existing CRS account IDs
|
||||
existingCRSIDs, err := s.accountRepo.ListCRSAccountIDs(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to list existing CRS accounts: %w", err)
|
||||
}
|
||||
|
||||
result := &PreviewFromCRSResult{
|
||||
NewAccounts: make([]CRSPreviewAccount, 0),
|
||||
ExistingAccounts: make([]CRSPreviewAccount, 0),
|
||||
}
|
||||
|
||||
classify := func(crsID, kind, name, platform, accountType string) {
|
||||
preview := CRSPreviewAccount{
|
||||
CRSAccountID: crsID,
|
||||
Kind: kind,
|
||||
Name: defaultName(name, crsID),
|
||||
Platform: platform,
|
||||
Type: accountType,
|
||||
}
|
||||
if _, exists := existingCRSIDs[crsID]; exists {
|
||||
result.ExistingAccounts = append(result.ExistingAccounts, preview)
|
||||
} else {
|
||||
result.NewAccounts = append(result.NewAccounts, preview)
|
||||
}
|
||||
}
|
||||
|
||||
for _, src := range exported.Data.ClaudeAccounts {
|
||||
authType := strings.TrimSpace(src.AuthType)
|
||||
if authType == "" {
|
||||
authType = AccountTypeOAuth
|
||||
}
|
||||
classify(src.ID, src.Kind, src.Name, PlatformAnthropic, authType)
|
||||
}
|
||||
for _, src := range exported.Data.ClaudeConsoleAccounts {
|
||||
classify(src.ID, src.Kind, src.Name, PlatformAnthropic, AccountTypeAPIKey)
|
||||
}
|
||||
for _, src := range exported.Data.OpenAIOAuthAccounts {
|
||||
classify(src.ID, src.Kind, src.Name, PlatformOpenAI, AccountTypeOAuth)
|
||||
}
|
||||
for _, src := range exported.Data.OpenAIResponsesAccounts {
|
||||
classify(src.ID, src.Kind, src.Name, PlatformOpenAI, AccountTypeAPIKey)
|
||||
}
|
||||
for _, src := range exported.Data.GeminiOAuthAccounts {
|
||||
classify(src.ID, src.Kind, src.Name, PlatformGemini, AccountTypeOAuth)
|
||||
}
|
||||
for _, src := range exported.Data.GeminiAPIKeyAccounts {
|
||||
classify(src.ID, src.Kind, src.Name, PlatformGemini, AccountTypeAPIKey)
|
||||
}
|
||||
|
||||
return result, nil
|
||||
}
|
||||
|
||||
69
backend/internal/service/digest_session_store.go
Normal file
69
backend/internal/service/digest_session_store.go
Normal file
@@ -0,0 +1,69 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
gocache "github.com/patrickmn/go-cache"
|
||||
)
|
||||
|
||||
// digestSessionTTL 摘要会话默认 TTL
|
||||
const digestSessionTTL = 5 * time.Minute
|
||||
|
||||
// sessionEntry flat cache 条目
|
||||
type sessionEntry struct {
|
||||
uuid string
|
||||
accountID int64
|
||||
}
|
||||
|
||||
// DigestSessionStore 内存摘要会话存储(flat cache 实现)
|
||||
// key: "{groupID}:{prefixHash}|{digestChain}" → *sessionEntry
|
||||
type DigestSessionStore struct {
|
||||
cache *gocache.Cache
|
||||
}
|
||||
|
||||
// NewDigestSessionStore 创建内存摘要会话存储
|
||||
func NewDigestSessionStore() *DigestSessionStore {
|
||||
return &DigestSessionStore{
|
||||
cache: gocache.New(digestSessionTTL, time.Minute),
|
||||
}
|
||||
}
|
||||
|
||||
// Save 保存摘要会话。oldDigestChain 为 Find 返回的 matchedChain,用于删旧 key。
|
||||
func (s *DigestSessionStore) Save(groupID int64, prefixHash, digestChain, uuid string, accountID int64, oldDigestChain string) {
|
||||
if digestChain == "" {
|
||||
return
|
||||
}
|
||||
ns := buildNS(groupID, prefixHash)
|
||||
s.cache.Set(ns+digestChain, &sessionEntry{uuid: uuid, accountID: accountID}, gocache.DefaultExpiration)
|
||||
if oldDigestChain != "" && oldDigestChain != digestChain {
|
||||
s.cache.Delete(ns + oldDigestChain)
|
||||
}
|
||||
}
|
||||
|
||||
// Find 查找摘要会话,从完整 chain 逐段截断,返回最长匹配及对应 matchedChain。
|
||||
func (s *DigestSessionStore) Find(groupID int64, prefixHash, digestChain string) (uuid string, accountID int64, matchedChain string, found bool) {
|
||||
if digestChain == "" {
|
||||
return "", 0, "", false
|
||||
}
|
||||
ns := buildNS(groupID, prefixHash)
|
||||
chain := digestChain
|
||||
for {
|
||||
if val, ok := s.cache.Get(ns + chain); ok {
|
||||
if e, ok := val.(*sessionEntry); ok {
|
||||
return e.uuid, e.accountID, chain, true
|
||||
}
|
||||
}
|
||||
i := strings.LastIndex(chain, "-")
|
||||
if i < 0 {
|
||||
return "", 0, "", false
|
||||
}
|
||||
chain = chain[:i]
|
||||
}
|
||||
}
|
||||
|
||||
// buildNS 构建 namespace 前缀
|
||||
func buildNS(groupID int64, prefixHash string) string {
|
||||
return strconv.FormatInt(groupID, 10) + ":" + prefixHash + "|"
|
||||
}
|
||||
312
backend/internal/service/digest_session_store_test.go
Normal file
312
backend/internal/service/digest_session_store_test.go
Normal file
@@ -0,0 +1,312 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
gocache "github.com/patrickmn/go-cache"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestDigestSessionStore_SaveAndFind(t *testing.T) {
|
||||
store := NewDigestSessionStore()
|
||||
|
||||
store.Save(1, "prefix", "s:a1-u:b2-m:c3", "uuid-1", 100, "")
|
||||
|
||||
uuid, accountID, _, found := store.Find(1, "prefix", "s:a1-u:b2-m:c3")
|
||||
require.True(t, found)
|
||||
assert.Equal(t, "uuid-1", uuid)
|
||||
assert.Equal(t, int64(100), accountID)
|
||||
}
|
||||
|
||||
func TestDigestSessionStore_PrefixMatch(t *testing.T) {
|
||||
store := NewDigestSessionStore()
|
||||
|
||||
// 保存短链
|
||||
store.Save(1, "prefix", "u:a-m:b", "uuid-short", 10, "")
|
||||
|
||||
// 用长链查找,应前缀匹配到短链
|
||||
uuid, accountID, matchedChain, found := store.Find(1, "prefix", "u:a-m:b-u:c-m:d")
|
||||
require.True(t, found)
|
||||
assert.Equal(t, "uuid-short", uuid)
|
||||
assert.Equal(t, int64(10), accountID)
|
||||
assert.Equal(t, "u:a-m:b", matchedChain)
|
||||
}
|
||||
|
||||
func TestDigestSessionStore_LongestPrefixMatch(t *testing.T) {
|
||||
store := NewDigestSessionStore()
|
||||
|
||||
store.Save(1, "prefix", "u:a", "uuid-1", 1, "")
|
||||
store.Save(1, "prefix", "u:a-m:b", "uuid-2", 2, "")
|
||||
store.Save(1, "prefix", "u:a-m:b-u:c", "uuid-3", 3, "")
|
||||
|
||||
// 应匹配最深的 "u:a-m:b-u:c"(从完整 chain 逐段截断,先命中最长的)
|
||||
uuid, accountID, _, found := store.Find(1, "prefix", "u:a-m:b-u:c-m:d-u:e")
|
||||
require.True(t, found)
|
||||
assert.Equal(t, "uuid-3", uuid)
|
||||
assert.Equal(t, int64(3), accountID)
|
||||
|
||||
// 查找中等长度,应匹配到 "u:a-m:b"
|
||||
uuid, accountID, _, found = store.Find(1, "prefix", "u:a-m:b-u:x")
|
||||
require.True(t, found)
|
||||
assert.Equal(t, "uuid-2", uuid)
|
||||
assert.Equal(t, int64(2), accountID)
|
||||
}
|
||||
|
||||
func TestDigestSessionStore_SaveDeletesOldChain(t *testing.T) {
|
||||
store := NewDigestSessionStore()
|
||||
|
||||
// 第一轮:保存 "u:a-m:b"
|
||||
store.Save(1, "prefix", "u:a-m:b", "uuid-1", 100, "")
|
||||
|
||||
// 第二轮:同一 uuid 保存更长的链,传入旧 chain
|
||||
store.Save(1, "prefix", "u:a-m:b-u:c-m:d", "uuid-1", 100, "u:a-m:b")
|
||||
|
||||
// 旧链 "u:a-m:b" 应已被删除
|
||||
_, _, _, found := store.Find(1, "prefix", "u:a-m:b")
|
||||
assert.False(t, found, "old chain should be deleted")
|
||||
|
||||
// 新链应能找到
|
||||
uuid, accountID, _, found := store.Find(1, "prefix", "u:a-m:b-u:c-m:d")
|
||||
require.True(t, found)
|
||||
assert.Equal(t, "uuid-1", uuid)
|
||||
assert.Equal(t, int64(100), accountID)
|
||||
}
|
||||
|
||||
func TestDigestSessionStore_DifferentSessionsNoInterference(t *testing.T) {
|
||||
store := NewDigestSessionStore()
|
||||
|
||||
// 相同系统提示词,不同用户提示词
|
||||
store.Save(1, "prefix", "s:sys-u:user1", "uuid-1", 100, "")
|
||||
store.Save(1, "prefix", "s:sys-u:user2", "uuid-2", 200, "")
|
||||
|
||||
uuid, accountID, _, found := store.Find(1, "prefix", "s:sys-u:user1-m:reply1")
|
||||
require.True(t, found)
|
||||
assert.Equal(t, "uuid-1", uuid)
|
||||
assert.Equal(t, int64(100), accountID)
|
||||
|
||||
uuid, accountID, _, found = store.Find(1, "prefix", "s:sys-u:user2-m:reply2")
|
||||
require.True(t, found)
|
||||
assert.Equal(t, "uuid-2", uuid)
|
||||
assert.Equal(t, int64(200), accountID)
|
||||
}
|
||||
|
||||
func TestDigestSessionStore_NoMatch(t *testing.T) {
|
||||
store := NewDigestSessionStore()
|
||||
|
||||
store.Save(1, "prefix", "u:a-m:b", "uuid-1", 100, "")
|
||||
|
||||
// 完全不同的 chain
|
||||
_, _, _, found := store.Find(1, "prefix", "u:x-m:y")
|
||||
assert.False(t, found)
|
||||
}
|
||||
|
||||
func TestDigestSessionStore_DifferentPrefixHash(t *testing.T) {
|
||||
store := NewDigestSessionStore()
|
||||
|
||||
store.Save(1, "prefix1", "u:a-m:b", "uuid-1", 100, "")
|
||||
|
||||
// 不同 prefixHash 应隔离
|
||||
_, _, _, found := store.Find(1, "prefix2", "u:a-m:b")
|
||||
assert.False(t, found)
|
||||
}
|
||||
|
||||
func TestDigestSessionStore_DifferentGroupID(t *testing.T) {
|
||||
store := NewDigestSessionStore()
|
||||
|
||||
store.Save(1, "prefix", "u:a-m:b", "uuid-1", 100, "")
|
||||
|
||||
// 不同 groupID 应隔离
|
||||
_, _, _, found := store.Find(2, "prefix", "u:a-m:b")
|
||||
assert.False(t, found)
|
||||
}
|
||||
|
||||
func TestDigestSessionStore_EmptyDigestChain(t *testing.T) {
|
||||
store := NewDigestSessionStore()
|
||||
|
||||
// 空链不应保存
|
||||
store.Save(1, "prefix", "", "uuid-1", 100, "")
|
||||
_, _, _, found := store.Find(1, "prefix", "")
|
||||
assert.False(t, found)
|
||||
}
|
||||
|
||||
func TestDigestSessionStore_TTLExpiration(t *testing.T) {
|
||||
store := &DigestSessionStore{
|
||||
cache: gocache.New(100*time.Millisecond, 50*time.Millisecond),
|
||||
}
|
||||
|
||||
store.Save(1, "prefix", "u:a-m:b", "uuid-1", 100, "")
|
||||
|
||||
// 立即应该能找到
|
||||
_, _, _, found := store.Find(1, "prefix", "u:a-m:b")
|
||||
require.True(t, found)
|
||||
|
||||
// 等待过期 + 清理周期
|
||||
time.Sleep(300 * time.Millisecond)
|
||||
|
||||
// 过期后应找不到
|
||||
_, _, _, found = store.Find(1, "prefix", "u:a-m:b")
|
||||
assert.False(t, found)
|
||||
}
|
||||
|
||||
func TestDigestSessionStore_ConcurrentSafety(t *testing.T) {
|
||||
store := NewDigestSessionStore()
|
||||
|
||||
var wg sync.WaitGroup
|
||||
const goroutines = 50
|
||||
const operations = 100
|
||||
|
||||
wg.Add(goroutines)
|
||||
for g := 0; g < goroutines; g++ {
|
||||
go func(id int) {
|
||||
defer wg.Done()
|
||||
prefix := fmt.Sprintf("prefix-%d", id%5)
|
||||
for i := 0; i < operations; i++ {
|
||||
chain := fmt.Sprintf("u:%d-m:%d", id, i)
|
||||
uuid := fmt.Sprintf("uuid-%d-%d", id, i)
|
||||
store.Save(1, prefix, chain, uuid, int64(id), "")
|
||||
store.Find(1, prefix, chain)
|
||||
}
|
||||
}(g)
|
||||
}
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
func TestDigestSessionStore_MultipleSessions(t *testing.T) {
|
||||
store := NewDigestSessionStore()
|
||||
|
||||
sessions := []struct {
|
||||
chain string
|
||||
uuid string
|
||||
accountID int64
|
||||
}{
|
||||
{"u:session1", "uuid-1", 1},
|
||||
{"u:session2-m:reply2", "uuid-2", 2},
|
||||
{"u:session3-m:reply3-u:msg3", "uuid-3", 3},
|
||||
}
|
||||
|
||||
for _, sess := range sessions {
|
||||
store.Save(1, "prefix", sess.chain, sess.uuid, sess.accountID, "")
|
||||
}
|
||||
|
||||
// 验证每个会话都能正确查找
|
||||
for _, sess := range sessions {
|
||||
uuid, accountID, _, found := store.Find(1, "prefix", sess.chain)
|
||||
require.True(t, found, "should find session: %s", sess.chain)
|
||||
assert.Equal(t, sess.uuid, uuid)
|
||||
assert.Equal(t, sess.accountID, accountID)
|
||||
}
|
||||
|
||||
// 验证继续对话的场景
|
||||
uuid, accountID, _, found := store.Find(1, "prefix", "u:session2-m:reply2-u:newmsg")
|
||||
require.True(t, found)
|
||||
assert.Equal(t, "uuid-2", uuid)
|
||||
assert.Equal(t, int64(2), accountID)
|
||||
}
|
||||
|
||||
func TestDigestSessionStore_Performance1000Sessions(t *testing.T) {
|
||||
store := NewDigestSessionStore()
|
||||
|
||||
// 插入 1000 个会话
|
||||
for i := 0; i < 1000; i++ {
|
||||
chain := fmt.Sprintf("s:sys-u:user%d-m:reply%d", i, i)
|
||||
store.Save(1, "prefix", chain, fmt.Sprintf("uuid-%d", i), int64(i), "")
|
||||
}
|
||||
|
||||
// 查找性能测试
|
||||
start := time.Now()
|
||||
const lookups = 10000
|
||||
for i := 0; i < lookups; i++ {
|
||||
idx := i % 1000
|
||||
chain := fmt.Sprintf("s:sys-u:user%d-m:reply%d-u:newmsg", idx, idx)
|
||||
_, _, _, found := store.Find(1, "prefix", chain)
|
||||
assert.True(t, found)
|
||||
}
|
||||
elapsed := time.Since(start)
|
||||
t.Logf("%d lookups in %v (%.0f ns/op)", lookups, elapsed, float64(elapsed.Nanoseconds())/lookups)
|
||||
}
|
||||
|
||||
func TestDigestSessionStore_FindReturnsMatchedChain(t *testing.T) {
|
||||
store := NewDigestSessionStore()
|
||||
|
||||
store.Save(1, "prefix", "u:a-m:b-u:c", "uuid-1", 100, "")
|
||||
|
||||
// 精确匹配
|
||||
_, _, matchedChain, found := store.Find(1, "prefix", "u:a-m:b-u:c")
|
||||
require.True(t, found)
|
||||
assert.Equal(t, "u:a-m:b-u:c", matchedChain)
|
||||
|
||||
// 前缀匹配(截断后命中)
|
||||
_, _, matchedChain, found = store.Find(1, "prefix", "u:a-m:b-u:c-m:d-u:e")
|
||||
require.True(t, found)
|
||||
assert.Equal(t, "u:a-m:b-u:c", matchedChain)
|
||||
}
|
||||
|
||||
func TestDigestSessionStore_CacheItemCountStable(t *testing.T) {
|
||||
store := NewDigestSessionStore()
|
||||
|
||||
// 模拟 100 个独立会话,每个进行 10 轮对话
|
||||
// 正确传递 oldDigestChain 时,每个会话始终只保留 1 个 key
|
||||
for conv := 0; conv < 100; conv++ {
|
||||
var prevMatchedChain string
|
||||
for round := 0; round < 10; round++ {
|
||||
chain := fmt.Sprintf("s:sys-u:user%d", conv)
|
||||
for r := 0; r < round; r++ {
|
||||
chain += fmt.Sprintf("-m:a%d-u:q%d", r, r+1)
|
||||
}
|
||||
uuid := fmt.Sprintf("uuid-conv%d", conv)
|
||||
|
||||
_, _, matched, _ := store.Find(1, "prefix", chain)
|
||||
store.Save(1, "prefix", chain, uuid, int64(conv), matched)
|
||||
prevMatchedChain = matched
|
||||
_ = prevMatchedChain
|
||||
}
|
||||
}
|
||||
|
||||
// 100 个会话 × 1 key/会话 = 应该 ≤ 100 个 key
|
||||
// 允许少量并发残留,但绝不能接近 100×10=1000
|
||||
itemCount := store.cache.ItemCount()
|
||||
assert.LessOrEqual(t, itemCount, 100, "cache should have at most 100 items (1 per conversation), got %d", itemCount)
|
||||
t.Logf("Cache item count after 100 conversations × 10 rounds: %d", itemCount)
|
||||
}
|
||||
|
||||
func TestDigestSessionStore_TTLPreventsUnboundedGrowth(t *testing.T) {
|
||||
// 使用极短 TTL 验证大量写入后 cache 能被清理
|
||||
store := &DigestSessionStore{
|
||||
cache: gocache.New(100*time.Millisecond, 50*time.Millisecond),
|
||||
}
|
||||
|
||||
// 插入 500 个不同的 key(无 oldDigestChain,模拟最坏场景:全是新会话首轮)
|
||||
for i := 0; i < 500; i++ {
|
||||
chain := fmt.Sprintf("u:user%d", i)
|
||||
store.Save(1, "prefix", chain, fmt.Sprintf("uuid-%d", i), int64(i), "")
|
||||
}
|
||||
|
||||
assert.Equal(t, 500, store.cache.ItemCount())
|
||||
|
||||
// 等待 TTL + 清理周期
|
||||
time.Sleep(300 * time.Millisecond)
|
||||
|
||||
assert.Equal(t, 0, store.cache.ItemCount(), "all items should be expired and cleaned up")
|
||||
}
|
||||
|
||||
func TestDigestSessionStore_SaveSameChainNoDelete(t *testing.T) {
|
||||
store := NewDigestSessionStore()
|
||||
|
||||
// 保存 chain
|
||||
store.Save(1, "prefix", "u:a-m:b", "uuid-1", 100, "")
|
||||
|
||||
// 用户重发相同消息:oldDigestChain == digestChain,不应删掉刚设置的 key
|
||||
store.Save(1, "prefix", "u:a-m:b", "uuid-1", 100, "u:a-m:b")
|
||||
|
||||
// 仍然能找到
|
||||
uuid, accountID, _, found := store.Find(1, "prefix", "u:a-m:b")
|
||||
require.True(t, found)
|
||||
assert.Equal(t, "uuid-1", uuid)
|
||||
assert.Equal(t, int64(100), accountID)
|
||||
}
|
||||
366
backend/internal/service/error_policy_integration_test.go
Normal file
366
backend/internal/service/error_policy_integration_test.go
Normal file
@@ -0,0 +1,366 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Mocks (scoped to this file by naming convention)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// epFixedUpstream returns a fixed response for every request.
|
||||
type epFixedUpstream struct {
|
||||
statusCode int
|
||||
body string
|
||||
calls int
|
||||
}
|
||||
|
||||
func (u *epFixedUpstream) Do(req *http.Request, proxyURL string, accountID int64, accountConcurrency int) (*http.Response, error) {
|
||||
u.calls++
|
||||
return &http.Response{
|
||||
StatusCode: u.statusCode,
|
||||
Header: http.Header{},
|
||||
Body: io.NopCloser(strings.NewReader(u.body)),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (u *epFixedUpstream) DoWithTLS(req *http.Request, proxyURL string, accountID int64, accountConcurrency int, enableTLSFingerprint bool) (*http.Response, error) {
|
||||
return u.Do(req, proxyURL, accountID, accountConcurrency)
|
||||
}
|
||||
|
||||
// epAccountRepo records SetTempUnschedulable / SetError calls.
|
||||
type epAccountRepo struct {
|
||||
mockAccountRepoForGemini
|
||||
tempCalls int
|
||||
setErrCalls int
|
||||
}
|
||||
|
||||
func (r *epAccountRepo) SetTempUnschedulable(_ context.Context, _ int64, _ time.Time, _ string) error {
|
||||
r.tempCalls++
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *epAccountRepo) SetError(_ context.Context, _ int64, _ string) error {
|
||||
r.setErrCalls++
|
||||
return nil
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func saveAndSetBaseURLs(t *testing.T) {
|
||||
t.Helper()
|
||||
oldBaseURLs := append([]string(nil), antigravity.BaseURLs...)
|
||||
oldAvail := antigravity.DefaultURLAvailability
|
||||
antigravity.BaseURLs = []string{"https://ep-test.example"}
|
||||
antigravity.DefaultURLAvailability = antigravity.NewURLAvailability(time.Minute)
|
||||
t.Cleanup(func() {
|
||||
antigravity.BaseURLs = oldBaseURLs
|
||||
antigravity.DefaultURLAvailability = oldAvail
|
||||
})
|
||||
}
|
||||
|
||||
func newRetryParams(account *Account, upstream HTTPUpstream, handleError func(context.Context, string, *Account, int, http.Header, []byte, string, int64, string, bool) *handleModelRateLimitResult) antigravityRetryLoopParams {
|
||||
return antigravityRetryLoopParams{
|
||||
ctx: context.Background(),
|
||||
prefix: "[ep-test]",
|
||||
account: account,
|
||||
accessToken: "token",
|
||||
action: "generateContent",
|
||||
body: []byte(`{"input":"test"}`),
|
||||
httpUpstream: upstream,
|
||||
requestedModel: "claude-sonnet-4-5",
|
||||
handleError: handleError,
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// TestRetryLoop_ErrorPolicy_CustomErrorCodes
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestRetryLoop_ErrorPolicy_CustomErrorCodes(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
upstreamStatus int
|
||||
upstreamBody string
|
||||
customCodes []any
|
||||
expectHandleError int
|
||||
expectUpstream int
|
||||
expectStatusCode int
|
||||
}{
|
||||
{
|
||||
name: "429_in_custom_codes_matched",
|
||||
upstreamStatus: 429,
|
||||
upstreamBody: `{"error":"rate limited"}`,
|
||||
customCodes: []any{float64(429)},
|
||||
expectHandleError: 1,
|
||||
expectUpstream: 1,
|
||||
expectStatusCode: 429,
|
||||
},
|
||||
{
|
||||
name: "429_not_in_custom_codes_skipped",
|
||||
upstreamStatus: 429,
|
||||
upstreamBody: `{"error":"rate limited"}`,
|
||||
customCodes: []any{float64(500)},
|
||||
expectHandleError: 0,
|
||||
expectUpstream: 1,
|
||||
expectStatusCode: 429,
|
||||
},
|
||||
{
|
||||
name: "500_in_custom_codes_matched",
|
||||
upstreamStatus: 500,
|
||||
upstreamBody: `{"error":"internal"}`,
|
||||
customCodes: []any{float64(500)},
|
||||
expectHandleError: 1,
|
||||
expectUpstream: 1,
|
||||
expectStatusCode: 500,
|
||||
},
|
||||
{
|
||||
name: "500_not_in_custom_codes_skipped",
|
||||
upstreamStatus: 500,
|
||||
upstreamBody: `{"error":"internal"}`,
|
||||
customCodes: []any{float64(429)},
|
||||
expectHandleError: 0,
|
||||
expectUpstream: 1,
|
||||
expectStatusCode: 500,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
saveAndSetBaseURLs(t)
|
||||
|
||||
upstream := &epFixedUpstream{statusCode: tt.upstreamStatus, body: tt.upstreamBody}
|
||||
repo := &epAccountRepo{}
|
||||
rlSvc := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
|
||||
|
||||
account := &Account{
|
||||
ID: 100,
|
||||
Type: AccountTypeAPIKey,
|
||||
Platform: PlatformAntigravity,
|
||||
Schedulable: true,
|
||||
Status: StatusActive,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"custom_error_codes_enabled": true,
|
||||
"custom_error_codes": tt.customCodes,
|
||||
},
|
||||
}
|
||||
|
||||
svc := &AntigravityGatewayService{rateLimitService: rlSvc}
|
||||
|
||||
var handleErrorCount int
|
||||
p := newRetryParams(account, upstream, func(_ context.Context, _ string, _ *Account, _ int, _ http.Header, _ []byte, _ string, _ int64, _ string, _ bool) *handleModelRateLimitResult {
|
||||
handleErrorCount++
|
||||
return nil
|
||||
})
|
||||
|
||||
result, err := svc.antigravityRetryLoop(p)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.NotNil(t, result.resp)
|
||||
defer func() { _ = result.resp.Body.Close() }()
|
||||
|
||||
require.Equal(t, tt.expectStatusCode, result.resp.StatusCode)
|
||||
require.Equal(t, tt.expectHandleError, handleErrorCount, "handleError call count")
|
||||
require.Equal(t, tt.expectUpstream, upstream.calls, "upstream call count")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// TestRetryLoop_ErrorPolicy_TempUnschedulable
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestRetryLoop_ErrorPolicy_TempUnschedulable(t *testing.T) {
|
||||
tempRulesAccount := func(rules []any) *Account {
|
||||
return &Account{
|
||||
ID: 200,
|
||||
Type: AccountTypeOAuth,
|
||||
Platform: PlatformAntigravity,
|
||||
Schedulable: true,
|
||||
Status: StatusActive,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"temp_unschedulable_enabled": true,
|
||||
"temp_unschedulable_rules": rules,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
overloadedRule := map[string]any{
|
||||
"error_code": float64(503),
|
||||
"keywords": []any{"overloaded"},
|
||||
"duration_minutes": float64(10),
|
||||
}
|
||||
|
||||
rateLimitRule := map[string]any{
|
||||
"error_code": float64(429),
|
||||
"keywords": []any{"rate limited keyword"},
|
||||
"duration_minutes": float64(5),
|
||||
}
|
||||
|
||||
t.Run("503_overloaded_matches_rule", func(t *testing.T) {
|
||||
saveAndSetBaseURLs(t)
|
||||
|
||||
upstream := &epFixedUpstream{statusCode: 503, body: `overloaded`}
|
||||
repo := &epAccountRepo{}
|
||||
rlSvc := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
|
||||
svc := &AntigravityGatewayService{rateLimitService: rlSvc}
|
||||
|
||||
account := tempRulesAccount([]any{overloadedRule})
|
||||
p := newRetryParams(account, upstream, func(_ context.Context, _ string, _ *Account, _ int, _ http.Header, _ []byte, _ string, _ int64, _ string, _ bool) *handleModelRateLimitResult {
|
||||
t.Error("handleError should not be called for temp unschedulable")
|
||||
return nil
|
||||
})
|
||||
|
||||
result, err := svc.antigravityRetryLoop(p)
|
||||
|
||||
require.Nil(t, result)
|
||||
var switchErr *AntigravityAccountSwitchError
|
||||
require.ErrorAs(t, err, &switchErr)
|
||||
require.Equal(t, account.ID, switchErr.OriginalAccountID)
|
||||
require.Equal(t, 1, upstream.calls, "should not retry")
|
||||
})
|
||||
|
||||
t.Run("429_rate_limited_keyword_matches_rule", func(t *testing.T) {
|
||||
saveAndSetBaseURLs(t)
|
||||
|
||||
upstream := &epFixedUpstream{statusCode: 429, body: `rate limited keyword`}
|
||||
repo := &epAccountRepo{}
|
||||
rlSvc := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
|
||||
svc := &AntigravityGatewayService{rateLimitService: rlSvc}
|
||||
|
||||
account := tempRulesAccount([]any{rateLimitRule})
|
||||
p := newRetryParams(account, upstream, func(_ context.Context, _ string, _ *Account, _ int, _ http.Header, _ []byte, _ string, _ int64, _ string, _ bool) *handleModelRateLimitResult {
|
||||
t.Error("handleError should not be called for temp unschedulable")
|
||||
return nil
|
||||
})
|
||||
|
||||
result, err := svc.antigravityRetryLoop(p)
|
||||
|
||||
require.Nil(t, result)
|
||||
var switchErr *AntigravityAccountSwitchError
|
||||
require.ErrorAs(t, err, &switchErr)
|
||||
require.Equal(t, account.ID, switchErr.OriginalAccountID)
|
||||
require.Equal(t, 1, upstream.calls, "should not retry")
|
||||
})
|
||||
|
||||
t.Run("503_body_no_match_continues_default_retry", func(t *testing.T) {
|
||||
saveAndSetBaseURLs(t)
|
||||
|
||||
upstream := &epFixedUpstream{statusCode: 503, body: `random`}
|
||||
repo := &epAccountRepo{}
|
||||
rlSvc := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
|
||||
svc := &AntigravityGatewayService{rateLimitService: rlSvc}
|
||||
|
||||
account := tempRulesAccount([]any{overloadedRule})
|
||||
|
||||
// Use a short-lived context: the backoff sleep (~1s) will be
|
||||
// interrupted, proving the code entered the default retry path
|
||||
// instead of breaking early via error policy.
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
|
||||
defer cancel()
|
||||
|
||||
p := newRetryParams(account, upstream, func(_ context.Context, _ string, _ *Account, _ int, _ http.Header, _ []byte, _ string, _ int64, _ string, _ bool) *handleModelRateLimitResult {
|
||||
return nil
|
||||
})
|
||||
p.ctx = ctx
|
||||
|
||||
result, err := svc.antigravityRetryLoop(p)
|
||||
|
||||
// Context cancellation during backoff proves default retry was entered
|
||||
require.Nil(t, result)
|
||||
require.ErrorIs(t, err, context.DeadlineExceeded)
|
||||
require.GreaterOrEqual(t, upstream.calls, 1, "should have called upstream at least once")
|
||||
})
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// TestRetryLoop_ErrorPolicy_NilRateLimitService
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestRetryLoop_ErrorPolicy_NilRateLimitService(t *testing.T) {
|
||||
saveAndSetBaseURLs(t)
|
||||
|
||||
upstream := &epFixedUpstream{statusCode: 429, body: `{"error":"rate limited"}`}
|
||||
// rateLimitService is nil — must not panic
|
||||
svc := &AntigravityGatewayService{rateLimitService: nil}
|
||||
|
||||
account := &Account{
|
||||
ID: 300,
|
||||
Type: AccountTypeOAuth,
|
||||
Platform: PlatformAntigravity,
|
||||
Schedulable: true,
|
||||
Status: StatusActive,
|
||||
Concurrency: 1,
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
|
||||
defer cancel()
|
||||
|
||||
p := newRetryParams(account, upstream, func(_ context.Context, _ string, _ *Account, _ int, _ http.Header, _ []byte, _ string, _ int64, _ string, _ bool) *handleModelRateLimitResult {
|
||||
return nil
|
||||
})
|
||||
p.ctx = ctx
|
||||
|
||||
// Should not panic; enters the default retry path (eventually times out)
|
||||
result, err := svc.antigravityRetryLoop(p)
|
||||
|
||||
require.Nil(t, result)
|
||||
require.ErrorIs(t, err, context.DeadlineExceeded)
|
||||
require.GreaterOrEqual(t, upstream.calls, 1)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// TestRetryLoop_ErrorPolicy_NoPolicy_OriginalBehavior
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestRetryLoop_ErrorPolicy_NoPolicy_OriginalBehavior(t *testing.T) {
|
||||
saveAndSetBaseURLs(t)
|
||||
|
||||
upstream := &epFixedUpstream{statusCode: 429, body: `{"error":"rate limited"}`}
|
||||
repo := &epAccountRepo{}
|
||||
rlSvc := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
|
||||
svc := &AntigravityGatewayService{rateLimitService: rlSvc}
|
||||
|
||||
// Plain OAuth account with no error policy configured
|
||||
account := &Account{
|
||||
ID: 400,
|
||||
Type: AccountTypeOAuth,
|
||||
Platform: PlatformAntigravity,
|
||||
Schedulable: true,
|
||||
Status: StatusActive,
|
||||
Concurrency: 1,
|
||||
}
|
||||
|
||||
var handleErrorCount int
|
||||
p := newRetryParams(account, upstream, func(_ context.Context, _ string, _ *Account, _ int, _ http.Header, _ []byte, _ string, _ int64, _ string, _ bool) *handleModelRateLimitResult {
|
||||
handleErrorCount++
|
||||
return nil
|
||||
})
|
||||
|
||||
result, err := svc.antigravityRetryLoop(p)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.NotNil(t, result.resp)
|
||||
defer func() { _ = result.resp.Body.Close() }()
|
||||
|
||||
require.Equal(t, http.StatusTooManyRequests, result.resp.StatusCode)
|
||||
require.Equal(t, antigravityMaxRetries, upstream.calls, "should exhaust all retries")
|
||||
require.Equal(t, 1, handleErrorCount, "handleError should be called once after retries exhausted")
|
||||
}
|
||||
289
backend/internal/service/error_policy_test.go
Normal file
289
backend/internal/service/error_policy_test.go
Normal file
@@ -0,0 +1,289 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// TestCheckErrorPolicy — 6 table-driven cases for the pure logic function
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestCheckErrorPolicy(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
account *Account
|
||||
statusCode int
|
||||
body []byte
|
||||
expected ErrorPolicyResult
|
||||
}{
|
||||
{
|
||||
name: "no_policy_oauth_returns_none",
|
||||
account: &Account{
|
||||
ID: 1,
|
||||
Type: AccountTypeOAuth,
|
||||
Platform: PlatformAntigravity,
|
||||
// no custom error codes, no temp rules
|
||||
},
|
||||
statusCode: 500,
|
||||
body: []byte(`"error"`),
|
||||
expected: ErrorPolicyNone,
|
||||
},
|
||||
{
|
||||
name: "custom_error_codes_hit_returns_matched",
|
||||
account: &Account{
|
||||
ID: 2,
|
||||
Type: AccountTypeAPIKey,
|
||||
Platform: PlatformAntigravity,
|
||||
Credentials: map[string]any{
|
||||
"custom_error_codes_enabled": true,
|
||||
"custom_error_codes": []any{float64(429), float64(500)},
|
||||
},
|
||||
},
|
||||
statusCode: 500,
|
||||
body: []byte(`"error"`),
|
||||
expected: ErrorPolicyMatched,
|
||||
},
|
||||
{
|
||||
name: "custom_error_codes_miss_returns_skipped",
|
||||
account: &Account{
|
||||
ID: 3,
|
||||
Type: AccountTypeAPIKey,
|
||||
Platform: PlatformAntigravity,
|
||||
Credentials: map[string]any{
|
||||
"custom_error_codes_enabled": true,
|
||||
"custom_error_codes": []any{float64(429), float64(500)},
|
||||
},
|
||||
},
|
||||
statusCode: 503,
|
||||
body: []byte(`"error"`),
|
||||
expected: ErrorPolicySkipped,
|
||||
},
|
||||
{
|
||||
name: "temp_unschedulable_hit_returns_temp_unscheduled",
|
||||
account: &Account{
|
||||
ID: 4,
|
||||
Type: AccountTypeOAuth,
|
||||
Platform: PlatformAntigravity,
|
||||
Credentials: map[string]any{
|
||||
"temp_unschedulable_enabled": true,
|
||||
"temp_unschedulable_rules": []any{
|
||||
map[string]any{
|
||||
"error_code": float64(503),
|
||||
"keywords": []any{"overloaded"},
|
||||
"duration_minutes": float64(10),
|
||||
"description": "overloaded rule",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
statusCode: 503,
|
||||
body: []byte(`overloaded service`),
|
||||
expected: ErrorPolicyTempUnscheduled,
|
||||
},
|
||||
{
|
||||
name: "temp_unschedulable_body_miss_returns_none",
|
||||
account: &Account{
|
||||
ID: 5,
|
||||
Type: AccountTypeOAuth,
|
||||
Platform: PlatformAntigravity,
|
||||
Credentials: map[string]any{
|
||||
"temp_unschedulable_enabled": true,
|
||||
"temp_unschedulable_rules": []any{
|
||||
map[string]any{
|
||||
"error_code": float64(503),
|
||||
"keywords": []any{"overloaded"},
|
||||
"duration_minutes": float64(10),
|
||||
"description": "overloaded rule",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
statusCode: 503,
|
||||
body: []byte(`random msg`),
|
||||
expected: ErrorPolicyNone,
|
||||
},
|
||||
{
|
||||
name: "custom_error_codes_override_temp_unschedulable",
|
||||
account: &Account{
|
||||
ID: 6,
|
||||
Type: AccountTypeAPIKey,
|
||||
Platform: PlatformAntigravity,
|
||||
Credentials: map[string]any{
|
||||
"custom_error_codes_enabled": true,
|
||||
"custom_error_codes": []any{float64(503)},
|
||||
"temp_unschedulable_enabled": true,
|
||||
"temp_unschedulable_rules": []any{
|
||||
map[string]any{
|
||||
"error_code": float64(503),
|
||||
"keywords": []any{"overloaded"},
|
||||
"duration_minutes": float64(10),
|
||||
"description": "overloaded rule",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
statusCode: 503,
|
||||
body: []byte(`overloaded`),
|
||||
expected: ErrorPolicyMatched, // custom codes take precedence
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
repo := &errorPolicyRepoStub{}
|
||||
svc := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
|
||||
|
||||
result := svc.CheckErrorPolicy(context.Background(), tt.account, tt.statusCode, tt.body)
|
||||
require.Equal(t, tt.expected, result, "unexpected ErrorPolicyResult")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// TestApplyErrorPolicy — 4 table-driven cases for the wrapper method
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestApplyErrorPolicy(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
account *Account
|
||||
statusCode int
|
||||
body []byte
|
||||
expectedHandled bool
|
||||
expectedSwitchErr bool // expect *AntigravityAccountSwitchError
|
||||
handleErrorCalls int
|
||||
}{
|
||||
{
|
||||
name: "none_not_handled",
|
||||
account: &Account{
|
||||
ID: 10,
|
||||
Type: AccountTypeOAuth,
|
||||
Platform: PlatformAntigravity,
|
||||
},
|
||||
statusCode: 500,
|
||||
body: []byte(`"error"`),
|
||||
expectedHandled: false,
|
||||
handleErrorCalls: 0,
|
||||
},
|
||||
{
|
||||
name: "skipped_handled_no_handleError",
|
||||
account: &Account{
|
||||
ID: 11,
|
||||
Type: AccountTypeAPIKey,
|
||||
Platform: PlatformAntigravity,
|
||||
Credentials: map[string]any{
|
||||
"custom_error_codes_enabled": true,
|
||||
"custom_error_codes": []any{float64(429)},
|
||||
},
|
||||
},
|
||||
statusCode: 500, // not in custom codes
|
||||
body: []byte(`"error"`),
|
||||
expectedHandled: true,
|
||||
handleErrorCalls: 0,
|
||||
},
|
||||
{
|
||||
name: "matched_handled_calls_handleError",
|
||||
account: &Account{
|
||||
ID: 12,
|
||||
Type: AccountTypeAPIKey,
|
||||
Platform: PlatformAntigravity,
|
||||
Credentials: map[string]any{
|
||||
"custom_error_codes_enabled": true,
|
||||
"custom_error_codes": []any{float64(500)},
|
||||
},
|
||||
},
|
||||
statusCode: 500,
|
||||
body: []byte(`"error"`),
|
||||
expectedHandled: true,
|
||||
handleErrorCalls: 1,
|
||||
},
|
||||
{
|
||||
name: "temp_unscheduled_returns_switch_error",
|
||||
account: &Account{
|
||||
ID: 13,
|
||||
Type: AccountTypeOAuth,
|
||||
Platform: PlatformAntigravity,
|
||||
Credentials: map[string]any{
|
||||
"temp_unschedulable_enabled": true,
|
||||
"temp_unschedulable_rules": []any{
|
||||
map[string]any{
|
||||
"error_code": float64(503),
|
||||
"keywords": []any{"overloaded"},
|
||||
"duration_minutes": float64(10),
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
statusCode: 503,
|
||||
body: []byte(`overloaded`),
|
||||
expectedHandled: true,
|
||||
expectedSwitchErr: true,
|
||||
handleErrorCalls: 0,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
repo := &errorPolicyRepoStub{}
|
||||
rlSvc := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
|
||||
svc := &AntigravityGatewayService{
|
||||
rateLimitService: rlSvc,
|
||||
}
|
||||
|
||||
var handleErrorCount int
|
||||
p := antigravityRetryLoopParams{
|
||||
ctx: context.Background(),
|
||||
prefix: "[test]",
|
||||
account: tt.account,
|
||||
handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
|
||||
handleErrorCount++
|
||||
return nil
|
||||
},
|
||||
isStickySession: true,
|
||||
}
|
||||
|
||||
handled, retErr := svc.applyErrorPolicy(p, tt.statusCode, http.Header{}, tt.body)
|
||||
|
||||
require.Equal(t, tt.expectedHandled, handled, "handled mismatch")
|
||||
require.Equal(t, tt.handleErrorCalls, handleErrorCount, "handleError call count mismatch")
|
||||
|
||||
if tt.expectedSwitchErr {
|
||||
var switchErr *AntigravityAccountSwitchError
|
||||
require.ErrorAs(t, retErr, &switchErr)
|
||||
require.Equal(t, tt.account.ID, switchErr.OriginalAccountID)
|
||||
} else {
|
||||
require.NoError(t, retErr)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// errorPolicyRepoStub — minimal AccountRepository stub for error policy tests
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
type errorPolicyRepoStub struct {
|
||||
mockAccountRepoForGemini
|
||||
tempCalls int
|
||||
setErrCalls int
|
||||
lastErrorMsg string
|
||||
}
|
||||
|
||||
func (r *errorPolicyRepoStub) SetTempUnschedulable(ctx context.Context, id int64, until time.Time, reason string) error {
|
||||
r.tempCalls++
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *errorPolicyRepoStub) SetError(ctx context.Context, id int64, errorMsg string) error {
|
||||
r.setErrCalls++
|
||||
r.lastErrorMsg = errorMsg
|
||||
return nil
|
||||
}
|
||||
@@ -77,7 +77,12 @@ func (m *mockAccountRepoForPlatform) Create(ctx context.Context, account *Accoun
|
||||
func (m *mockAccountRepoForPlatform) GetByCRSAccountID(ctx context.Context, crsAccountID string) (*Account, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (m *mockAccountRepoForPlatform) FindByExtraField(ctx context.Context, key string, value interface{}) ([]Account, error) {
|
||||
|
||||
func (m *mockAccountRepoForPlatform) FindByExtraField(ctx context.Context, key string, value any) ([]Account, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (m *mockAccountRepoForPlatform) ListCRSAccountIDs(ctx context.Context) (map[string]int64, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (m *mockAccountRepoForPlatform) Update(ctx context.Context, account *Account) error {
|
||||
@@ -145,9 +150,6 @@ func (m *mockAccountRepoForPlatform) ListSchedulableByGroupIDAndPlatforms(ctx co
|
||||
func (m *mockAccountRepoForPlatform) SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error {
|
||||
return nil
|
||||
}
|
||||
func (m *mockAccountRepoForPlatform) SetAntigravityQuotaScopeLimit(ctx context.Context, id int64, scope AntigravityQuotaScope, resetAt time.Time) error {
|
||||
return nil
|
||||
}
|
||||
func (m *mockAccountRepoForPlatform) SetModelRateLimit(ctx context.Context, id int64, scope string, resetAt time.Time) error {
|
||||
return nil
|
||||
}
|
||||
@@ -219,22 +221,6 @@ func (m *mockGatewayCacheForPlatform) DeleteSessionAccountID(ctx context.Context
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *mockGatewayCacheForPlatform) IncrModelCallCount(ctx context.Context, accountID int64, model string) (int64, error) {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
func (m *mockGatewayCacheForPlatform) GetModelLoadBatch(ctx context.Context, accountIDs []int64, model string) (map[int64]*ModelLoadInfo, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (m *mockGatewayCacheForPlatform) FindGeminiSession(ctx context.Context, groupID int64, prefixHash, digestChain string) (uuid string, accountID int64, found bool) {
|
||||
return "", 0, false
|
||||
}
|
||||
|
||||
func (m *mockGatewayCacheForPlatform) SaveGeminiSession(ctx context.Context, groupID int64, prefixHash, digestChain, uuid string, accountID int64) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
type mockGroupRepoForGateway struct {
|
||||
groups map[int64]*Group
|
||||
getByIDCalls int
|
||||
@@ -293,6 +279,10 @@ func (m *mockGroupRepoForGateway) GetAccountIDsByGroupIDs(ctx context.Context, g
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (m *mockGroupRepoForGateway) UpdateSortOrders(ctx context.Context, updates []GroupSortOrderUpdate) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func ptr[T any](v T) *T {
|
||||
return &v
|
||||
}
|
||||
|
||||
@@ -6,9 +6,19 @@ import (
|
||||
"fmt"
|
||||
"math"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/domain"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
|
||||
)
|
||||
|
||||
// SessionContext 粘性会话上下文,用于区分不同来源的请求。
|
||||
// 仅在 GenerateSessionHash 第 3 级 fallback(消息内容 hash)时混入,
|
||||
// 避免不同用户发送相同消息产生相同 hash 导致账号集中。
|
||||
type SessionContext struct {
|
||||
ClientIP string
|
||||
UserAgent string
|
||||
APIKeyID int64
|
||||
}
|
||||
|
||||
// ParsedRequest 保存网关请求的预解析结果
|
||||
//
|
||||
// 性能优化说明:
|
||||
@@ -22,20 +32,22 @@ import (
|
||||
// 2. 将解析结果 ParsedRequest 传递给 Service 层
|
||||
// 3. 避免重复 json.Unmarshal,减少 CPU 和内存开销
|
||||
type ParsedRequest struct {
|
||||
Body []byte // 原始请求体(保留用于转发)
|
||||
Model string // 请求的模型名称
|
||||
Stream bool // 是否为流式请求
|
||||
MetadataUserID string // metadata.user_id(用于会话亲和)
|
||||
System any // system 字段内容
|
||||
Messages []any // messages 数组
|
||||
HasSystem bool // 是否包含 system 字段(包含 null 也视为显式传入)
|
||||
ThinkingEnabled bool // 是否开启 thinking(部分平台会影响最终模型名)
|
||||
MaxTokens int // max_tokens 值(用于探测请求拦截)
|
||||
Body []byte // 原始请求体(保留用于转发)
|
||||
Model string // 请求的模型名称
|
||||
Stream bool // 是否为流式请求
|
||||
MetadataUserID string // metadata.user_id(用于会话亲和)
|
||||
System any // system 字段内容
|
||||
Messages []any // messages 数组
|
||||
HasSystem bool // 是否包含 system 字段(包含 null 也视为显式传入)
|
||||
ThinkingEnabled bool // 是否开启 thinking(部分平台会影响最终模型名)
|
||||
MaxTokens int // max_tokens 值(用于探测请求拦截)
|
||||
SessionContext *SessionContext // 可选:请求上下文区分因子(nil 时行为不变)
|
||||
}
|
||||
|
||||
// ParseGatewayRequest 解析网关请求体并返回结构化结果
|
||||
// 性能优化:一次解析提取所有需要的字段,避免重复 Unmarshal
|
||||
func ParseGatewayRequest(body []byte) (*ParsedRequest, error) {
|
||||
// ParseGatewayRequest 解析网关请求体并返回结构化结果。
|
||||
// protocol 指定请求协议格式(domain.PlatformAnthropic / domain.PlatformGemini),
|
||||
// 不同协议使用不同的 system/messages 字段名。
|
||||
func ParseGatewayRequest(body []byte, protocol string) (*ParsedRequest, error) {
|
||||
var req map[string]any
|
||||
if err := json.Unmarshal(body, &req); err != nil {
|
||||
return nil, err
|
||||
@@ -64,14 +76,29 @@ func ParseGatewayRequest(body []byte) (*ParsedRequest, error) {
|
||||
parsed.MetadataUserID = userID
|
||||
}
|
||||
}
|
||||
// system 字段只要存在就视为显式提供(即使为 null),
|
||||
// 以避免客户端传 null 时被默认 system 误注入。
|
||||
if system, ok := req["system"]; ok {
|
||||
parsed.HasSystem = true
|
||||
parsed.System = system
|
||||
}
|
||||
if messages, ok := req["messages"].([]any); ok {
|
||||
parsed.Messages = messages
|
||||
|
||||
switch protocol {
|
||||
case domain.PlatformGemini:
|
||||
// Gemini 原生格式: systemInstruction.parts / contents
|
||||
if sysInst, ok := req["systemInstruction"].(map[string]any); ok {
|
||||
if parts, ok := sysInst["parts"].([]any); ok {
|
||||
parsed.System = parts
|
||||
}
|
||||
}
|
||||
if contents, ok := req["contents"].([]any); ok {
|
||||
parsed.Messages = contents
|
||||
}
|
||||
default:
|
||||
// Anthropic / OpenAI 格式: system / messages
|
||||
// system 字段只要存在就视为显式提供(即使为 null),
|
||||
// 以避免客户端传 null 时被默认 system 误注入。
|
||||
if system, ok := req["system"]; ok {
|
||||
parsed.HasSystem = true
|
||||
parsed.System = system
|
||||
}
|
||||
if messages, ok := req["messages"].([]any); ok {
|
||||
parsed.Messages = messages
|
||||
}
|
||||
}
|
||||
|
||||
// thinking: {type: "enabled"}
|
||||
|
||||
@@ -4,12 +4,13 @@ import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/domain"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestParseGatewayRequest(t *testing.T) {
|
||||
body := []byte(`{"model":"claude-3-7-sonnet","stream":true,"metadata":{"user_id":"session_123e4567-e89b-12d3-a456-426614174000"},"system":[{"type":"text","text":"hello","cache_control":{"type":"ephemeral"}}],"messages":[{"content":"hi"}]}`)
|
||||
parsed, err := ParseGatewayRequest(body)
|
||||
parsed, err := ParseGatewayRequest(body, "")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "claude-3-7-sonnet", parsed.Model)
|
||||
require.True(t, parsed.Stream)
|
||||
@@ -22,7 +23,7 @@ func TestParseGatewayRequest(t *testing.T) {
|
||||
|
||||
func TestParseGatewayRequest_ThinkingEnabled(t *testing.T) {
|
||||
body := []byte(`{"model":"claude-sonnet-4-5","thinking":{"type":"enabled"},"messages":[{"content":"hi"}]}`)
|
||||
parsed, err := ParseGatewayRequest(body)
|
||||
parsed, err := ParseGatewayRequest(body, "")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "claude-sonnet-4-5", parsed.Model)
|
||||
require.True(t, parsed.ThinkingEnabled)
|
||||
@@ -30,21 +31,21 @@ func TestParseGatewayRequest_ThinkingEnabled(t *testing.T) {
|
||||
|
||||
func TestParseGatewayRequest_MaxTokens(t *testing.T) {
|
||||
body := []byte(`{"model":"claude-haiku-4-5","max_tokens":1}`)
|
||||
parsed, err := ParseGatewayRequest(body)
|
||||
parsed, err := ParseGatewayRequest(body, "")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, parsed.MaxTokens)
|
||||
}
|
||||
|
||||
func TestParseGatewayRequest_MaxTokensNonIntegralIgnored(t *testing.T) {
|
||||
body := []byte(`{"model":"claude-haiku-4-5","max_tokens":1.5}`)
|
||||
parsed, err := ParseGatewayRequest(body)
|
||||
parsed, err := ParseGatewayRequest(body, "")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 0, parsed.MaxTokens)
|
||||
}
|
||||
|
||||
func TestParseGatewayRequest_SystemNull(t *testing.T) {
|
||||
body := []byte(`{"model":"claude-3","system":null}`)
|
||||
parsed, err := ParseGatewayRequest(body)
|
||||
parsed, err := ParseGatewayRequest(body, "")
|
||||
require.NoError(t, err)
|
||||
// 显式传入 system:null 也应视为“字段已存在”,避免默认 system 被注入。
|
||||
require.True(t, parsed.HasSystem)
|
||||
@@ -53,16 +54,112 @@ func TestParseGatewayRequest_SystemNull(t *testing.T) {
|
||||
|
||||
func TestParseGatewayRequest_InvalidModelType(t *testing.T) {
|
||||
body := []byte(`{"model":123}`)
|
||||
_, err := ParseGatewayRequest(body)
|
||||
_, err := ParseGatewayRequest(body, "")
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func TestParseGatewayRequest_InvalidStreamType(t *testing.T) {
|
||||
body := []byte(`{"stream":"true"}`)
|
||||
_, err := ParseGatewayRequest(body)
|
||||
_, err := ParseGatewayRequest(body, "")
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
// ============ Gemini 原生格式解析测试 ============
|
||||
|
||||
func TestParseGatewayRequest_GeminiContents(t *testing.T) {
|
||||
body := []byte(`{
|
||||
"contents": [
|
||||
{"role": "user", "parts": [{"text": "Hello"}]},
|
||||
{"role": "model", "parts": [{"text": "Hi there"}]},
|
||||
{"role": "user", "parts": [{"text": "How are you?"}]}
|
||||
]
|
||||
}`)
|
||||
parsed, err := ParseGatewayRequest(body, domain.PlatformGemini)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, parsed.Messages, 3, "should parse contents as Messages")
|
||||
require.False(t, parsed.HasSystem, "Gemini format should not set HasSystem")
|
||||
require.Nil(t, parsed.System, "no systemInstruction means nil System")
|
||||
}
|
||||
|
||||
func TestParseGatewayRequest_GeminiSystemInstruction(t *testing.T) {
|
||||
body := []byte(`{
|
||||
"systemInstruction": {
|
||||
"parts": [{"text": "You are a helpful assistant."}]
|
||||
},
|
||||
"contents": [
|
||||
{"role": "user", "parts": [{"text": "Hello"}]}
|
||||
]
|
||||
}`)
|
||||
parsed, err := ParseGatewayRequest(body, domain.PlatformGemini)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, parsed.System, "should parse systemInstruction.parts as System")
|
||||
parts, ok := parsed.System.([]any)
|
||||
require.True(t, ok)
|
||||
require.Len(t, parts, 1)
|
||||
partMap, ok := parts[0].(map[string]any)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, "You are a helpful assistant.", partMap["text"])
|
||||
require.Len(t, parsed.Messages, 1)
|
||||
}
|
||||
|
||||
func TestParseGatewayRequest_GeminiWithModel(t *testing.T) {
|
||||
body := []byte(`{
|
||||
"model": "gemini-2.5-pro",
|
||||
"contents": [{"role": "user", "parts": [{"text": "test"}]}]
|
||||
}`)
|
||||
parsed, err := ParseGatewayRequest(body, domain.PlatformGemini)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "gemini-2.5-pro", parsed.Model)
|
||||
require.Len(t, parsed.Messages, 1)
|
||||
}
|
||||
|
||||
func TestParseGatewayRequest_GeminiIgnoresAnthropicFields(t *testing.T) {
|
||||
// Gemini 格式下 system/messages 字段应被忽略
|
||||
body := []byte(`{
|
||||
"system": "should be ignored",
|
||||
"messages": [{"role": "user", "content": "ignored"}],
|
||||
"contents": [{"role": "user", "parts": [{"text": "real content"}]}]
|
||||
}`)
|
||||
parsed, err := ParseGatewayRequest(body, domain.PlatformGemini)
|
||||
require.NoError(t, err)
|
||||
require.False(t, parsed.HasSystem, "Gemini protocol should not parse Anthropic system field")
|
||||
require.Nil(t, parsed.System, "no systemInstruction = nil System")
|
||||
require.Len(t, parsed.Messages, 1, "should use contents, not messages")
|
||||
}
|
||||
|
||||
func TestParseGatewayRequest_GeminiEmptyContents(t *testing.T) {
|
||||
body := []byte(`{"contents": []}`)
|
||||
parsed, err := ParseGatewayRequest(body, domain.PlatformGemini)
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, parsed.Messages)
|
||||
}
|
||||
|
||||
func TestParseGatewayRequest_GeminiNoContents(t *testing.T) {
|
||||
body := []byte(`{"model": "gemini-2.5-flash"}`)
|
||||
parsed, err := ParseGatewayRequest(body, domain.PlatformGemini)
|
||||
require.NoError(t, err)
|
||||
require.Nil(t, parsed.Messages)
|
||||
require.Equal(t, "gemini-2.5-flash", parsed.Model)
|
||||
}
|
||||
|
||||
func TestParseGatewayRequest_AnthropicIgnoresGeminiFields(t *testing.T) {
|
||||
// Anthropic 格式下 contents/systemInstruction 字段应被忽略
|
||||
body := []byte(`{
|
||||
"system": "real system",
|
||||
"messages": [{"role": "user", "content": "real content"}],
|
||||
"contents": [{"role": "user", "parts": [{"text": "ignored"}]}],
|
||||
"systemInstruction": {"parts": [{"text": "ignored"}]}
|
||||
}`)
|
||||
parsed, err := ParseGatewayRequest(body, domain.PlatformAnthropic)
|
||||
require.NoError(t, err)
|
||||
require.True(t, parsed.HasSystem)
|
||||
require.Equal(t, "real system", parsed.System)
|
||||
require.Len(t, parsed.Messages, 1)
|
||||
msg, ok := parsed.Messages[0].(map[string]any)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, "real content", msg["content"])
|
||||
}
|
||||
|
||||
func TestFilterThinkingBlocks(t *testing.T) {
|
||||
containsThinkingBlock := func(body []byte) bool {
|
||||
var req map[string]any
|
||||
|
||||
@@ -5,7 +5,6 @@ import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
@@ -17,6 +16,7 @@ import (
|
||||
"os"
|
||||
"regexp"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
@@ -26,6 +26,7 @@ import (
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
||||
"github.com/Wei-Shaw/sub2api/internal/util/responseheaders"
|
||||
"github.com/Wei-Shaw/sub2api/internal/util/urlvalidator"
|
||||
"github.com/cespare/xxhash/v2"
|
||||
"github.com/google/uuid"
|
||||
"github.com/tidwall/gjson"
|
||||
"github.com/tidwall/sjson"
|
||||
@@ -245,9 +246,6 @@ var (
|
||||
// ErrClaudeCodeOnly 表示分组仅允许 Claude Code 客户端访问
|
||||
var ErrClaudeCodeOnly = errors.New("this group only allows Claude Code clients")
|
||||
|
||||
// ErrModelScopeNotSupported 表示请求的模型系列不在分组支持的范围内
|
||||
var ErrModelScopeNotSupported = errors.New("model scope not supported by this group")
|
||||
|
||||
// allowedHeaders 白名单headers(参考CRS项目)
|
||||
var allowedHeaders = map[string]bool{
|
||||
"accept": true,
|
||||
@@ -273,13 +271,6 @@ var allowedHeaders = map[string]bool{
|
||||
// GatewayCache 定义网关服务的缓存操作接口。
|
||||
// 提供粘性会话(Sticky Session)的存储、查询、刷新和删除功能。
|
||||
//
|
||||
// ModelLoadInfo 模型负载信息(用于 Antigravity 调度)
|
||||
// Model load info for Antigravity scheduling
|
||||
type ModelLoadInfo struct {
|
||||
CallCount int64 // 当前分钟调用次数 / Call count in current minute
|
||||
LastUsedAt time.Time // 最后调度时间(零值表示未调度过)/ Last scheduling time (zero means never scheduled)
|
||||
}
|
||||
|
||||
// GatewayCache defines cache operations for gateway service.
|
||||
// Provides sticky session storage, retrieval, refresh and deletion capabilities.
|
||||
type GatewayCache interface {
|
||||
@@ -295,24 +286,6 @@ type GatewayCache interface {
|
||||
// DeleteSessionAccountID 删除粘性会话绑定,用于账号不可用时主动清理
|
||||
// Delete sticky session binding, used to proactively clean up when account becomes unavailable
|
||||
DeleteSessionAccountID(ctx context.Context, groupID int64, sessionHash string) error
|
||||
|
||||
// IncrModelCallCount 增加模型调用次数并更新最后调度时间(Antigravity 专用)
|
||||
// Increment model call count and update last scheduling time (Antigravity only)
|
||||
// 返回更新后的调用次数
|
||||
IncrModelCallCount(ctx context.Context, accountID int64, model string) (int64, error)
|
||||
|
||||
// GetModelLoadBatch 批量获取账号的模型负载信息(Antigravity 专用)
|
||||
// Batch get model load info for accounts (Antigravity only)
|
||||
GetModelLoadBatch(ctx context.Context, accountIDs []int64, model string) (map[int64]*ModelLoadInfo, error)
|
||||
|
||||
// FindGeminiSession 查找 Gemini 会话(MGET 倒序匹配)
|
||||
// Find Gemini session using MGET reverse order matching
|
||||
// 返回最长匹配的会话信息(uuid, accountID)
|
||||
FindGeminiSession(ctx context.Context, groupID int64, prefixHash, digestChain string) (uuid string, accountID int64, found bool)
|
||||
|
||||
// SaveGeminiSession 保存 Gemini 会话
|
||||
// Save Gemini session binding
|
||||
SaveGeminiSession(ctx context.Context, groupID int64, prefixHash, digestChain, uuid string, accountID int64) error
|
||||
}
|
||||
|
||||
// derefGroupID safely dereferences *int64 to int64, returning 0 if nil
|
||||
@@ -323,21 +296,15 @@ func derefGroupID(groupID *int64) int64 {
|
||||
return *groupID
|
||||
}
|
||||
|
||||
// stickySessionRateLimitThreshold 定义清除粘性会话的限流时间阈值。
|
||||
// 当账号限流剩余时间超过此阈值时,清除粘性会话以便切换到其他账号。
|
||||
// 低于此阈值时保持粘性会话,等待短暂限流结束。
|
||||
const stickySessionRateLimitThreshold = 10 * time.Second
|
||||
|
||||
// shouldClearStickySession 检查账号是否处于不可调度状态,需要清理粘性会话绑定。
|
||||
// 当账号状态为错误、禁用、不可调度、处于临时不可调度期间,
|
||||
// 或模型限流剩余时间超过 stickySessionRateLimitThreshold 时,返回 true。
|
||||
// 或请求的模型处于限流状态时,返回 true。
|
||||
// 这确保后续请求不会继续使用不可用的账号。
|
||||
//
|
||||
// shouldClearStickySession checks if an account is in an unschedulable state
|
||||
// and the sticky session binding should be cleared.
|
||||
// Returns true when account status is error/disabled, schedulable is false,
|
||||
// within temporary unschedulable period, or model rate limit remaining time
|
||||
// exceeds stickySessionRateLimitThreshold.
|
||||
// within temporary unschedulable period, or the requested model is rate-limited.
|
||||
// This ensures subsequent requests won't continue using unavailable accounts.
|
||||
func shouldClearStickySession(account *Account, requestedModel string) bool {
|
||||
if account == nil {
|
||||
@@ -349,8 +316,8 @@ func shouldClearStickySession(account *Account, requestedModel string) bool {
|
||||
if account.TempUnschedulableUntil != nil && time.Now().Before(*account.TempUnschedulableUntil) {
|
||||
return true
|
||||
}
|
||||
// 检查模型限流和 scope 限流,只在超过阈值时清除粘性会话
|
||||
if remaining := account.GetRateLimitRemainingTimeWithContext(context.Background(), requestedModel); remaining > stickySessionRateLimitThreshold {
|
||||
// 检查模型限流和 scope 限流,有限流即清除粘性会话
|
||||
if remaining := account.GetRateLimitRemainingTimeWithContext(context.Background(), requestedModel); remaining > 0 {
|
||||
return true
|
||||
}
|
||||
return false
|
||||
@@ -417,6 +384,7 @@ type GatewayService struct {
|
||||
userSubRepo UserSubscriptionRepository
|
||||
userGroupRateRepo UserGroupRateRepository
|
||||
cache GatewayCache
|
||||
digestStore *DigestSessionStore
|
||||
cfg *config.Config
|
||||
schedulerSnapshot *SchedulerSnapshotService
|
||||
billingService *BillingService
|
||||
@@ -450,6 +418,7 @@ func NewGatewayService(
|
||||
deferredService *DeferredService,
|
||||
claudeTokenProvider *ClaudeTokenProvider,
|
||||
sessionLimitCache SessionLimitCache,
|
||||
digestStore *DigestSessionStore,
|
||||
) *GatewayService {
|
||||
return &GatewayService{
|
||||
accountRepo: accountRepo,
|
||||
@@ -459,6 +428,7 @@ func NewGatewayService(
|
||||
userSubRepo: userSubRepo,
|
||||
userGroupRateRepo: userGroupRateRepo,
|
||||
cache: cache,
|
||||
digestStore: digestStore,
|
||||
cfg: cfg,
|
||||
schedulerSnapshot: schedulerSnapshot,
|
||||
concurrencyService: concurrencyService,
|
||||
@@ -492,23 +462,45 @@ func (s *GatewayService) GenerateSessionHash(parsed *ParsedRequest) string {
|
||||
return s.hashContent(cacheableContent)
|
||||
}
|
||||
|
||||
// 3. Fallback: 使用 system 内容
|
||||
// 3. 最后 fallback: 使用 session上下文 + system + 所有消息的完整摘要串
|
||||
var combined strings.Builder
|
||||
// 混入请求上下文区分因子,避免不同用户相同消息产生相同 hash
|
||||
if parsed.SessionContext != nil {
|
||||
_, _ = combined.WriteString(parsed.SessionContext.ClientIP)
|
||||
_, _ = combined.WriteString(":")
|
||||
_, _ = combined.WriteString(parsed.SessionContext.UserAgent)
|
||||
_, _ = combined.WriteString(":")
|
||||
_, _ = combined.WriteString(strconv.FormatInt(parsed.SessionContext.APIKeyID, 10))
|
||||
_, _ = combined.WriteString("|")
|
||||
}
|
||||
if parsed.System != nil {
|
||||
systemText := s.extractTextFromSystem(parsed.System)
|
||||
if systemText != "" {
|
||||
return s.hashContent(systemText)
|
||||
_, _ = combined.WriteString(systemText)
|
||||
}
|
||||
}
|
||||
|
||||
// 4. 最后 fallback: 使用第一条消息
|
||||
if len(parsed.Messages) > 0 {
|
||||
if firstMsg, ok := parsed.Messages[0].(map[string]any); ok {
|
||||
msgText := s.extractTextFromContent(firstMsg["content"])
|
||||
if msgText != "" {
|
||||
return s.hashContent(msgText)
|
||||
for _, msg := range parsed.Messages {
|
||||
if m, ok := msg.(map[string]any); ok {
|
||||
if content, exists := m["content"]; exists {
|
||||
// Anthropic: messages[].content
|
||||
if msgText := s.extractTextFromContent(content); msgText != "" {
|
||||
_, _ = combined.WriteString(msgText)
|
||||
}
|
||||
} else if parts, ok := m["parts"].([]any); ok {
|
||||
// Gemini: contents[].parts[].text
|
||||
for _, part := range parts {
|
||||
if partMap, ok := part.(map[string]any); ok {
|
||||
if text, ok := partMap["text"].(string); ok {
|
||||
_, _ = combined.WriteString(text)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if combined.Len() > 0 {
|
||||
return s.hashContent(combined.String())
|
||||
}
|
||||
|
||||
return ""
|
||||
}
|
||||
@@ -536,19 +528,37 @@ func (s *GatewayService) GetCachedSessionAccountID(ctx context.Context, groupID
|
||||
|
||||
// FindGeminiSession 查找 Gemini 会话(基于内容摘要链的 Fallback 匹配)
|
||||
// 返回最长匹配的会话信息(uuid, accountID)
|
||||
func (s *GatewayService) FindGeminiSession(ctx context.Context, groupID int64, prefixHash, digestChain string) (uuid string, accountID int64, found bool) {
|
||||
if digestChain == "" || s.cache == nil {
|
||||
return "", 0, false
|
||||
func (s *GatewayService) FindGeminiSession(_ context.Context, groupID int64, prefixHash, digestChain string) (uuid string, accountID int64, matchedChain string, found bool) {
|
||||
if digestChain == "" || s.digestStore == nil {
|
||||
return "", 0, "", false
|
||||
}
|
||||
return s.cache.FindGeminiSession(ctx, groupID, prefixHash, digestChain)
|
||||
return s.digestStore.Find(groupID, prefixHash, digestChain)
|
||||
}
|
||||
|
||||
// SaveGeminiSession 保存 Gemini 会话
|
||||
func (s *GatewayService) SaveGeminiSession(ctx context.Context, groupID int64, prefixHash, digestChain, uuid string, accountID int64) error {
|
||||
if digestChain == "" || s.cache == nil {
|
||||
// SaveGeminiSession 保存 Gemini 会话。oldDigestChain 为 Find 返回的 matchedChain,用于删旧 key。
|
||||
func (s *GatewayService) SaveGeminiSession(_ context.Context, groupID int64, prefixHash, digestChain, uuid string, accountID int64, oldDigestChain string) error {
|
||||
if digestChain == "" || s.digestStore == nil {
|
||||
return nil
|
||||
}
|
||||
return s.cache.SaveGeminiSession(ctx, groupID, prefixHash, digestChain, uuid, accountID)
|
||||
s.digestStore.Save(groupID, prefixHash, digestChain, uuid, accountID, oldDigestChain)
|
||||
return nil
|
||||
}
|
||||
|
||||
// FindAnthropicSession 查找 Anthropic 会话(基于内容摘要链的 Fallback 匹配)
|
||||
func (s *GatewayService) FindAnthropicSession(_ context.Context, groupID int64, prefixHash, digestChain string) (uuid string, accountID int64, matchedChain string, found bool) {
|
||||
if digestChain == "" || s.digestStore == nil {
|
||||
return "", 0, "", false
|
||||
}
|
||||
return s.digestStore.Find(groupID, prefixHash, digestChain)
|
||||
}
|
||||
|
||||
// SaveAnthropicSession 保存 Anthropic 会话
|
||||
func (s *GatewayService) SaveAnthropicSession(_ context.Context, groupID int64, prefixHash, digestChain, uuid string, accountID int64, oldDigestChain string) error {
|
||||
if digestChain == "" || s.digestStore == nil {
|
||||
return nil
|
||||
}
|
||||
s.digestStore.Save(groupID, prefixHash, digestChain, uuid, accountID, oldDigestChain)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *GatewayService) extractCacheableContent(parsed *ParsedRequest) string {
|
||||
@@ -633,8 +643,8 @@ func (s *GatewayService) extractTextFromContent(content any) string {
|
||||
}
|
||||
|
||||
func (s *GatewayService) hashContent(content string) string {
|
||||
hash := sha256.Sum256([]byte(content))
|
||||
return hex.EncodeToString(hash[:16]) // 32字符
|
||||
h := xxhash.Sum64String(content)
|
||||
return strconv.FormatUint(h, 36)
|
||||
}
|
||||
|
||||
// replaceModelInBody 替换请求体中的model字段
|
||||
@@ -993,13 +1003,6 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
|
||||
log.Printf("[ModelRoutingDebug] load-aware enabled: group_id=%v model=%s session=%s platform=%s", derefGroupID(groupID), requestedModel, shortSessionHash(sessionHash), platform)
|
||||
}
|
||||
|
||||
// Antigravity 模型系列检查(在账号选择前检查,确保所有代码路径都经过此检查)
|
||||
if platform == PlatformAntigravity && groupID != nil && requestedModel != "" {
|
||||
if err := s.checkAntigravityModelScope(ctx, *groupID, requestedModel); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
accounts, useMixed, err := s.listSchedulableAccounts(ctx, groupID, platform, hasForcePlatform)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -1114,7 +1117,6 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
|
||||
result.ReleaseFunc() // 释放槽位
|
||||
// 继续到负载感知选择
|
||||
} else {
|
||||
_ = s.cache.RefreshSessionTTL(ctx, derefGroupID(groupID), sessionHash, stickySessionTTL)
|
||||
if s.debugModelRoutingEnabled() {
|
||||
log.Printf("[ModelRoutingDebug] routed sticky hit: group_id=%v model=%s session=%s account=%d", derefGroupID(groupID), requestedModel, shortSessionHash(sessionHash), stickyAccountID)
|
||||
}
|
||||
@@ -1194,6 +1196,7 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
|
||||
return a.account.LastUsedAt.Before(*b.account.LastUsedAt)
|
||||
}
|
||||
})
|
||||
shuffleWithinSortGroups(routingAvailable)
|
||||
|
||||
// 4. 尝试获取槽位
|
||||
for _, item := range routingAvailable {
|
||||
@@ -1268,7 +1271,6 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
|
||||
if !s.checkAndRegisterSession(ctx, account, sessionHash) {
|
||||
result.ReleaseFunc() // 释放槽位,继续到 Layer 2
|
||||
} else {
|
||||
_ = s.cache.RefreshSessionTTL(ctx, derefGroupID(groupID), sessionHash, stickySessionTTL)
|
||||
return &AccountSelectionResult{
|
||||
Account: account,
|
||||
Acquired: true,
|
||||
@@ -1348,10 +1350,6 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
|
||||
return result, nil
|
||||
}
|
||||
} else {
|
||||
// Antigravity 平台:获取模型负载信息
|
||||
var modelLoadMap map[int64]*ModelLoadInfo
|
||||
isAntigravity := platform == PlatformAntigravity
|
||||
|
||||
var available []accountWithLoad
|
||||
for _, acc := range candidates {
|
||||
loadInfo := loadMap[acc.ID]
|
||||
@@ -1366,109 +1364,44 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
|
||||
}
|
||||
}
|
||||
|
||||
// Antigravity 平台:按账号实际映射后的模型名获取模型负载(与 Forward 的统计保持一致)
|
||||
if isAntigravity && requestedModel != "" && s.cache != nil && len(available) > 0 {
|
||||
modelLoadMap = make(map[int64]*ModelLoadInfo, len(available))
|
||||
modelToAccountIDs := make(map[string][]int64)
|
||||
for _, item := range available {
|
||||
mappedModel := mapAntigravityModel(item.account, requestedModel)
|
||||
if mappedModel == "" {
|
||||
continue
|
||||
}
|
||||
modelToAccountIDs[mappedModel] = append(modelToAccountIDs[mappedModel], item.account.ID)
|
||||
// 分层过滤选择:优先级 → 负载率 → LRU
|
||||
for len(available) > 0 {
|
||||
// 1. 取优先级最小的集合
|
||||
candidates := filterByMinPriority(available)
|
||||
// 2. 取负载率最低的集合
|
||||
candidates = filterByMinLoadRate(candidates)
|
||||
// 3. LRU 选择最久未用的账号
|
||||
selected := selectByLRU(candidates, preferOAuth)
|
||||
if selected == nil {
|
||||
break
|
||||
}
|
||||
for model, ids := range modelToAccountIDs {
|
||||
batch, err := s.cache.GetModelLoadBatch(ctx, ids, model)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
for id, info := range batch {
|
||||
modelLoadMap[id] = info
|
||||
}
|
||||
}
|
||||
if len(modelLoadMap) == 0 {
|
||||
modelLoadMap = nil
|
||||
}
|
||||
}
|
||||
|
||||
// Antigravity 平台:优先级硬过滤 →(同优先级内)按调用次数选择(最少优先,新账号用平均值)
|
||||
// 其他平台:分层过滤选择:优先级 → 负载率 → LRU
|
||||
if isAntigravity {
|
||||
for len(available) > 0 {
|
||||
// 1. 取优先级最小的集合(硬过滤)
|
||||
candidates := filterByMinPriority(available)
|
||||
// 2. 同优先级内按调用次数选择(调用次数最少优先,新账号使用平均值)
|
||||
selected := selectByCallCount(candidates, modelLoadMap, preferOAuth)
|
||||
if selected == nil {
|
||||
break
|
||||
}
|
||||
|
||||
result, err := s.tryAcquireAccountSlot(ctx, selected.account.ID, selected.account.Concurrency)
|
||||
if err == nil && result.Acquired {
|
||||
// 会话数量限制检查
|
||||
if !s.checkAndRegisterSession(ctx, selected.account, sessionHash) {
|
||||
result.ReleaseFunc() // 释放槽位,继续尝试下一个账号
|
||||
} else {
|
||||
if sessionHash != "" && s.cache != nil {
|
||||
_ = s.cache.SetSessionAccountID(ctx, derefGroupID(groupID), sessionHash, selected.account.ID, stickySessionTTL)
|
||||
}
|
||||
return &AccountSelectionResult{
|
||||
Account: selected.account,
|
||||
Acquired: true,
|
||||
ReleaseFunc: result.ReleaseFunc,
|
||||
}, nil
|
||||
result, err := s.tryAcquireAccountSlot(ctx, selected.account.ID, selected.account.Concurrency)
|
||||
if err == nil && result.Acquired {
|
||||
// 会话数量限制检查
|
||||
if !s.checkAndRegisterSession(ctx, selected.account, sessionHash) {
|
||||
result.ReleaseFunc() // 释放槽位,继续尝试下一个账号
|
||||
} else {
|
||||
if sessionHash != "" && s.cache != nil {
|
||||
_ = s.cache.SetSessionAccountID(ctx, derefGroupID(groupID), sessionHash, selected.account.ID, stickySessionTTL)
|
||||
}
|
||||
return &AccountSelectionResult{
|
||||
Account: selected.account,
|
||||
Acquired: true,
|
||||
ReleaseFunc: result.ReleaseFunc,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// 移除已尝试的账号,重新选择
|
||||
selectedID := selected.account.ID
|
||||
newAvailable := make([]accountWithLoad, 0, len(available)-1)
|
||||
for _, acc := range available {
|
||||
if acc.account.ID != selectedID {
|
||||
newAvailable = append(newAvailable, acc)
|
||||
}
|
||||
}
|
||||
available = newAvailable
|
||||
}
|
||||
} else {
|
||||
for len(available) > 0 {
|
||||
// 1. 取优先级最小的集合
|
||||
candidates := filterByMinPriority(available)
|
||||
// 2. 取负载率最低的集合
|
||||
candidates = filterByMinLoadRate(candidates)
|
||||
// 3. LRU 选择最久未用的账号
|
||||
selected := selectByLRU(candidates, preferOAuth)
|
||||
if selected == nil {
|
||||
break
|
||||
}
|
||||
|
||||
result, err := s.tryAcquireAccountSlot(ctx, selected.account.ID, selected.account.Concurrency)
|
||||
if err == nil && result.Acquired {
|
||||
// 会话数量限制检查
|
||||
if !s.checkAndRegisterSession(ctx, selected.account, sessionHash) {
|
||||
result.ReleaseFunc() // 释放槽位,继续尝试下一个账号
|
||||
} else {
|
||||
if sessionHash != "" && s.cache != nil {
|
||||
_ = s.cache.SetSessionAccountID(ctx, derefGroupID(groupID), sessionHash, selected.account.ID, stickySessionTTL)
|
||||
}
|
||||
return &AccountSelectionResult{
|
||||
Account: selected.account,
|
||||
Acquired: true,
|
||||
ReleaseFunc: result.ReleaseFunc,
|
||||
}, nil
|
||||
}
|
||||
// 移除已尝试的账号,重新进行分层过滤
|
||||
selectedID := selected.account.ID
|
||||
newAvailable := make([]accountWithLoad, 0, len(available)-1)
|
||||
for _, acc := range available {
|
||||
if acc.account.ID != selectedID {
|
||||
newAvailable = append(newAvailable, acc)
|
||||
}
|
||||
|
||||
// 移除已尝试的账号,重新进行分层过滤
|
||||
selectedID := selected.account.ID
|
||||
newAvailable := make([]accountWithLoad, 0, len(available)-1)
|
||||
for _, acc := range available {
|
||||
if acc.account.ID != selectedID {
|
||||
newAvailable = append(newAvailable, acc)
|
||||
}
|
||||
}
|
||||
available = newAvailable
|
||||
}
|
||||
available = newAvailable
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2004,87 +1937,79 @@ func sortAccountsByPriorityAndLastUsed(accounts []*Account, preferOAuth bool) {
|
||||
return a.LastUsedAt.Before(*b.LastUsedAt)
|
||||
}
|
||||
})
|
||||
shuffleWithinPriorityAndLastUsed(accounts)
|
||||
}
|
||||
|
||||
// selectByCallCount 从候选账号中选择调用次数最少的账号(Antigravity 专用)
|
||||
// 新账号(CallCount=0)使用平均调用次数作为虚拟值,避免冷启动被猛调
|
||||
// 如果有多个账号具有相同的最小调用次数,则随机选择一个
|
||||
func selectByCallCount(accounts []accountWithLoad, modelLoadMap map[int64]*ModelLoadInfo, preferOAuth bool) *accountWithLoad {
|
||||
if len(accounts) == 0 {
|
||||
return nil
|
||||
// shuffleWithinSortGroups 对排序后的 accountWithLoad 切片,按 (Priority, LoadRate, LastUsedAt) 分组后组内随机打乱。
|
||||
// 防止并发请求读取同一快照时,确定性排序导致所有请求命中相同账号。
|
||||
func shuffleWithinSortGroups(accounts []accountWithLoad) {
|
||||
if len(accounts) <= 1 {
|
||||
return
|
||||
}
|
||||
if len(accounts) == 1 {
|
||||
return &accounts[0]
|
||||
}
|
||||
|
||||
// 如果没有负载信息,回退到 LRU
|
||||
if modelLoadMap == nil {
|
||||
return selectByLRU(accounts, preferOAuth)
|
||||
}
|
||||
|
||||
// 1. 计算平均调用次数(用于新账号冷启动)
|
||||
var totalCallCount int64
|
||||
var countWithCalls int
|
||||
for _, acc := range accounts {
|
||||
if info := modelLoadMap[acc.account.ID]; info != nil && info.CallCount > 0 {
|
||||
totalCallCount += info.CallCount
|
||||
countWithCalls++
|
||||
i := 0
|
||||
for i < len(accounts) {
|
||||
j := i + 1
|
||||
for j < len(accounts) && sameAccountWithLoadGroup(accounts[i], accounts[j]) {
|
||||
j++
|
||||
}
|
||||
}
|
||||
|
||||
var avgCallCount int64
|
||||
if countWithCalls > 0 {
|
||||
avgCallCount = totalCallCount / int64(countWithCalls)
|
||||
}
|
||||
|
||||
// 2. 获取每个账号的有效调用次数
|
||||
getEffectiveCallCount := func(acc accountWithLoad) int64 {
|
||||
if acc.account == nil {
|
||||
return 0
|
||||
if j-i > 1 {
|
||||
mathrand.Shuffle(j-i, func(a, b int) {
|
||||
accounts[i+a], accounts[i+b] = accounts[i+b], accounts[i+a]
|
||||
})
|
||||
}
|
||||
info := modelLoadMap[acc.account.ID]
|
||||
if info == nil || info.CallCount == 0 {
|
||||
return avgCallCount // 新账号使用平均值
|
||||
}
|
||||
return info.CallCount
|
||||
i = j
|
||||
}
|
||||
}
|
||||
|
||||
// 3. 找到最小调用次数
|
||||
minCount := getEffectiveCallCount(accounts[0])
|
||||
for _, acc := range accounts[1:] {
|
||||
if c := getEffectiveCallCount(acc); c < minCount {
|
||||
minCount = c
|
||||
}
|
||||
// sameAccountWithLoadGroup 判断两个 accountWithLoad 是否属于同一排序组
|
||||
func sameAccountWithLoadGroup(a, b accountWithLoad) bool {
|
||||
if a.account.Priority != b.account.Priority {
|
||||
return false
|
||||
}
|
||||
|
||||
// 4. 收集所有具有最小调用次数的账号
|
||||
var candidateIdxs []int
|
||||
for i, acc := range accounts {
|
||||
if getEffectiveCallCount(acc) == minCount {
|
||||
candidateIdxs = append(candidateIdxs, i)
|
||||
}
|
||||
if a.loadInfo.LoadRate != b.loadInfo.LoadRate {
|
||||
return false
|
||||
}
|
||||
return sameLastUsedAt(a.account.LastUsedAt, b.account.LastUsedAt)
|
||||
}
|
||||
|
||||
// 5. 如果只有一个候选,直接返回
|
||||
if len(candidateIdxs) == 1 {
|
||||
return &accounts[candidateIdxs[0]]
|
||||
// shuffleWithinPriorityAndLastUsed 对排序后的 []*Account 切片,按 (Priority, LastUsedAt) 分组后组内随机打乱。
|
||||
func shuffleWithinPriorityAndLastUsed(accounts []*Account) {
|
||||
if len(accounts) <= 1 {
|
||||
return
|
||||
}
|
||||
|
||||
// 6. preferOAuth 处理
|
||||
if preferOAuth {
|
||||
var oauthIdxs []int
|
||||
for _, idx := range candidateIdxs {
|
||||
if accounts[idx].account.Type == AccountTypeOAuth {
|
||||
oauthIdxs = append(oauthIdxs, idx)
|
||||
}
|
||||
i := 0
|
||||
for i < len(accounts) {
|
||||
j := i + 1
|
||||
for j < len(accounts) && sameAccountGroup(accounts[i], accounts[j]) {
|
||||
j++
|
||||
}
|
||||
if len(oauthIdxs) > 0 {
|
||||
candidateIdxs = oauthIdxs
|
||||
if j-i > 1 {
|
||||
mathrand.Shuffle(j-i, func(a, b int) {
|
||||
accounts[i+a], accounts[i+b] = accounts[i+b], accounts[i+a]
|
||||
})
|
||||
}
|
||||
i = j
|
||||
}
|
||||
}
|
||||
|
||||
// 7. 随机选择
|
||||
return &accounts[candidateIdxs[mathrand.Intn(len(candidateIdxs))]]
|
||||
// sameAccountGroup 判断两个 Account 是否属于同一排序组(Priority + LastUsedAt)
|
||||
func sameAccountGroup(a, b *Account) bool {
|
||||
if a.Priority != b.Priority {
|
||||
return false
|
||||
}
|
||||
return sameLastUsedAt(a.LastUsedAt, b.LastUsedAt)
|
||||
}
|
||||
|
||||
// sameLastUsedAt 判断两个 LastUsedAt 是否相同(精度到秒)
|
||||
func sameLastUsedAt(a, b *time.Time) bool {
|
||||
switch {
|
||||
case a == nil && b == nil:
|
||||
return true
|
||||
case a == nil || b == nil:
|
||||
return false
|
||||
default:
|
||||
return a.Unix() == b.Unix()
|
||||
}
|
||||
}
|
||||
|
||||
// sortCandidatesForFallback 根据配置选择排序策略
|
||||
@@ -2139,13 +2064,6 @@ func shuffleWithinPriority(accounts []*Account) {
|
||||
|
||||
// selectAccountForModelWithPlatform 选择单平台账户(完全隔离)
|
||||
func (s *GatewayService) selectAccountForModelWithPlatform(ctx context.Context, groupID *int64, sessionHash string, requestedModel string, excludedIDs map[int64]struct{}, platform string) (*Account, error) {
|
||||
// 对 Antigravity 平台,检查请求的模型系列是否在分组支持范围内
|
||||
if platform == PlatformAntigravity && groupID != nil && requestedModel != "" {
|
||||
if err := s.checkAntigravityModelScope(ctx, *groupID, requestedModel); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
preferOAuth := platform == PlatformGemini
|
||||
routingAccountIDs := s.routingAccountIDsForRequest(ctx, groupID, requestedModel, platform)
|
||||
|
||||
@@ -2173,9 +2091,6 @@ func (s *GatewayService) selectAccountForModelWithPlatform(ctx context.Context,
|
||||
_ = s.cache.DeleteSessionAccountID(ctx, derefGroupID(groupID), sessionHash)
|
||||
}
|
||||
if !clearSticky && s.isAccountInGroup(account, groupID) && account.Platform == platform && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && account.IsSchedulableForModelWithContext(ctx, requestedModel) {
|
||||
if err := s.cache.RefreshSessionTTL(ctx, derefGroupID(groupID), sessionHash, stickySessionTTL); err != nil {
|
||||
log.Printf("refresh session ttl failed: session=%s err=%v", sessionHash, err)
|
||||
}
|
||||
if s.debugModelRoutingEnabled() {
|
||||
log.Printf("[ModelRoutingDebug] legacy routed sticky hit: group_id=%v model=%s session=%s account=%d", derefGroupID(groupID), requestedModel, shortSessionHash(sessionHash), accountID)
|
||||
}
|
||||
@@ -2276,9 +2191,6 @@ func (s *GatewayService) selectAccountForModelWithPlatform(ctx context.Context,
|
||||
_ = s.cache.DeleteSessionAccountID(ctx, derefGroupID(groupID), sessionHash)
|
||||
}
|
||||
if !clearSticky && s.isAccountInGroup(account, groupID) && account.Platform == platform && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && account.IsSchedulableForModelWithContext(ctx, requestedModel) {
|
||||
if err := s.cache.RefreshSessionTTL(ctx, derefGroupID(groupID), sessionHash, stickySessionTTL); err != nil {
|
||||
log.Printf("refresh session ttl failed: session=%s err=%v", sessionHash, err)
|
||||
}
|
||||
return account, nil
|
||||
}
|
||||
}
|
||||
@@ -2387,9 +2299,6 @@ func (s *GatewayService) selectAccountWithMixedScheduling(ctx context.Context, g
|
||||
}
|
||||
if !clearSticky && s.isAccountInGroup(account, groupID) && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && account.IsSchedulableForModelWithContext(ctx, requestedModel) {
|
||||
if account.Platform == nativePlatform || (account.Platform == PlatformAntigravity && account.IsMixedSchedulingEnabled()) {
|
||||
if err := s.cache.RefreshSessionTTL(ctx, derefGroupID(groupID), sessionHash, stickySessionTTL); err != nil {
|
||||
log.Printf("refresh session ttl failed: session=%s err=%v", sessionHash, err)
|
||||
}
|
||||
if s.debugModelRoutingEnabled() {
|
||||
log.Printf("[ModelRoutingDebug] legacy mixed routed sticky hit: group_id=%v model=%s session=%s account=%d", derefGroupID(groupID), requestedModel, shortSessionHash(sessionHash), accountID)
|
||||
}
|
||||
@@ -2492,9 +2401,6 @@ func (s *GatewayService) selectAccountWithMixedScheduling(ctx context.Context, g
|
||||
}
|
||||
if !clearSticky && s.isAccountInGroup(account, groupID) && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && account.IsSchedulableForModelWithContext(ctx, requestedModel) {
|
||||
if account.Platform == nativePlatform || (account.Platform == PlatformAntigravity && account.IsMixedSchedulingEnabled()) {
|
||||
if err := s.cache.RefreshSessionTTL(ctx, derefGroupID(groupID), sessionHash, stickySessionTTL); err != nil {
|
||||
log.Printf("refresh session ttl failed: session=%s err=%v", sessionHash, err)
|
||||
}
|
||||
return account, nil
|
||||
}
|
||||
}
|
||||
@@ -5185,27 +5091,6 @@ func (s *GatewayService) validateUpstreamBaseURL(raw string) (string, error) {
|
||||
return normalized, nil
|
||||
}
|
||||
|
||||
// checkAntigravityModelScope 检查 Antigravity 平台的模型系列是否在分组支持范围内
|
||||
func (s *GatewayService) checkAntigravityModelScope(ctx context.Context, groupID int64, requestedModel string) error {
|
||||
scope, ok := ResolveAntigravityQuotaScope(requestedModel)
|
||||
if !ok {
|
||||
return nil // 无法解析 scope,跳过检查
|
||||
}
|
||||
|
||||
group, err := s.resolveGroupByID(ctx, groupID)
|
||||
if err != nil {
|
||||
return nil // 查询失败时放行
|
||||
}
|
||||
if group == nil {
|
||||
return nil // 分组不存在时放行
|
||||
}
|
||||
|
||||
if !IsScopeSupported(group.SupportedModelScopes, scope) {
|
||||
return ErrModelScopeNotSupported
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetAvailableModels returns the list of models available for a group
|
||||
// It aggregates model_mapping keys from all schedulable accounts in the group
|
||||
func (s *GatewayService) GetAvailableModels(ctx context.Context, groupID *int64, platform string) []string {
|
||||
|
||||
@@ -14,7 +14,7 @@ func BenchmarkGenerateSessionHash_Metadata(b *testing.B) {
|
||||
|
||||
b.ReportAllocs()
|
||||
for i := 0; i < b.N; i++ {
|
||||
parsed, err := ParseGatewayRequest(body)
|
||||
parsed, err := ParseGatewayRequest(body, "")
|
||||
if err != nil {
|
||||
b.Fatalf("解析请求失败: %v", err)
|
||||
}
|
||||
|
||||
384
backend/internal/service/gemini_error_policy_test.go
Normal file
384
backend/internal/service/gemini_error_policy_test.go
Normal file
@@ -0,0 +1,384 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// TestShouldFailoverGeminiUpstreamError — verifies the failover decision
|
||||
// for the ErrorPolicyNone path (original logic preserved).
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestShouldFailoverGeminiUpstreamError(t *testing.T) {
|
||||
svc := &GeminiMessagesCompatService{}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
statusCode int
|
||||
expected bool
|
||||
}{
|
||||
{"401_failover", 401, true},
|
||||
{"403_failover", 403, true},
|
||||
{"429_failover", 429, true},
|
||||
{"529_failover", 529, true},
|
||||
{"500_failover", 500, true},
|
||||
{"502_failover", 502, true},
|
||||
{"503_failover", 503, true},
|
||||
{"400_no_failover", 400, false},
|
||||
{"404_no_failover", 404, false},
|
||||
{"422_no_failover", 422, false},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := svc.shouldFailoverGeminiUpstreamError(tt.statusCode)
|
||||
require.Equal(t, tt.expected, got)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// TestCheckErrorPolicy_GeminiAccounts — verifies CheckErrorPolicy works
|
||||
// correctly for Gemini platform accounts (API Key type).
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestCheckErrorPolicy_GeminiAccounts(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
account *Account
|
||||
statusCode int
|
||||
body []byte
|
||||
expected ErrorPolicyResult
|
||||
}{
|
||||
{
|
||||
name: "gemini_apikey_custom_codes_hit",
|
||||
account: &Account{
|
||||
ID: 100,
|
||||
Type: AccountTypeAPIKey,
|
||||
Platform: PlatformGemini,
|
||||
Credentials: map[string]any{
|
||||
"custom_error_codes_enabled": true,
|
||||
"custom_error_codes": []any{float64(429), float64(500)},
|
||||
},
|
||||
},
|
||||
statusCode: 429,
|
||||
body: []byte(`{"error":"rate limited"}`),
|
||||
expected: ErrorPolicyMatched,
|
||||
},
|
||||
{
|
||||
name: "gemini_apikey_custom_codes_miss",
|
||||
account: &Account{
|
||||
ID: 101,
|
||||
Type: AccountTypeAPIKey,
|
||||
Platform: PlatformGemini,
|
||||
Credentials: map[string]any{
|
||||
"custom_error_codes_enabled": true,
|
||||
"custom_error_codes": []any{float64(429)},
|
||||
},
|
||||
},
|
||||
statusCode: 500,
|
||||
body: []byte(`{"error":"internal"}`),
|
||||
expected: ErrorPolicySkipped,
|
||||
},
|
||||
{
|
||||
name: "gemini_apikey_no_custom_codes_returns_none",
|
||||
account: &Account{
|
||||
ID: 102,
|
||||
Type: AccountTypeAPIKey,
|
||||
Platform: PlatformGemini,
|
||||
},
|
||||
statusCode: 500,
|
||||
body: []byte(`{"error":"internal"}`),
|
||||
expected: ErrorPolicyNone,
|
||||
},
|
||||
{
|
||||
name: "gemini_apikey_temp_unschedulable_hit",
|
||||
account: &Account{
|
||||
ID: 103,
|
||||
Type: AccountTypeAPIKey,
|
||||
Platform: PlatformGemini,
|
||||
Credentials: map[string]any{
|
||||
"temp_unschedulable_enabled": true,
|
||||
"temp_unschedulable_rules": []any{
|
||||
map[string]any{
|
||||
"error_code": float64(503),
|
||||
"keywords": []any{"overloaded"},
|
||||
"duration_minutes": float64(10),
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
statusCode: 503,
|
||||
body: []byte(`overloaded service`),
|
||||
expected: ErrorPolicyTempUnscheduled,
|
||||
},
|
||||
{
|
||||
name: "gemini_custom_codes_override_temp_unschedulable",
|
||||
account: &Account{
|
||||
ID: 104,
|
||||
Type: AccountTypeAPIKey,
|
||||
Platform: PlatformGemini,
|
||||
Credentials: map[string]any{
|
||||
"custom_error_codes_enabled": true,
|
||||
"custom_error_codes": []any{float64(503)},
|
||||
"temp_unschedulable_enabled": true,
|
||||
"temp_unschedulable_rules": []any{
|
||||
map[string]any{
|
||||
"error_code": float64(503),
|
||||
"keywords": []any{"overloaded"},
|
||||
"duration_minutes": float64(10),
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
statusCode: 503,
|
||||
body: []byte(`overloaded`),
|
||||
expected: ErrorPolicyMatched, // custom codes take precedence
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
repo := &errorPolicyRepoStub{}
|
||||
svc := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
|
||||
|
||||
result := svc.CheckErrorPolicy(context.Background(), tt.account, tt.statusCode, tt.body)
|
||||
require.Equal(t, tt.expected, result)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// TestGeminiErrorPolicyIntegration — verifies the Gemini error handling
|
||||
// paths produce the correct behavior for each ErrorPolicyResult.
|
||||
//
|
||||
// These tests simulate the inline error policy switch in handleClaudeCompat
|
||||
// and forwardNativeGemini by calling the same methods in the same order.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestGeminiErrorPolicyIntegration(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
account *Account
|
||||
statusCode int
|
||||
respBody []byte
|
||||
expectFailover bool // expect UpstreamFailoverError
|
||||
expectHandleError bool // expect handleGeminiUpstreamError to be called
|
||||
expectShouldFailover bool // for None path, whether shouldFailover triggers
|
||||
}{
|
||||
{
|
||||
name: "custom_codes_matched_429_failover",
|
||||
account: &Account{
|
||||
ID: 200,
|
||||
Type: AccountTypeAPIKey,
|
||||
Platform: PlatformGemini,
|
||||
Credentials: map[string]any{
|
||||
"custom_error_codes_enabled": true,
|
||||
"custom_error_codes": []any{float64(429)},
|
||||
},
|
||||
},
|
||||
statusCode: 429,
|
||||
respBody: []byte(`{"error":"rate limited"}`),
|
||||
expectFailover: true,
|
||||
expectHandleError: true,
|
||||
},
|
||||
{
|
||||
name: "custom_codes_skipped_500_no_failover",
|
||||
account: &Account{
|
||||
ID: 201,
|
||||
Type: AccountTypeAPIKey,
|
||||
Platform: PlatformGemini,
|
||||
Credentials: map[string]any{
|
||||
"custom_error_codes_enabled": true,
|
||||
"custom_error_codes": []any{float64(429)},
|
||||
},
|
||||
},
|
||||
statusCode: 500,
|
||||
respBody: []byte(`{"error":"internal"}`),
|
||||
expectFailover: false,
|
||||
expectHandleError: false,
|
||||
},
|
||||
{
|
||||
name: "temp_unschedulable_matched_failover",
|
||||
account: &Account{
|
||||
ID: 202,
|
||||
Type: AccountTypeAPIKey,
|
||||
Platform: PlatformGemini,
|
||||
Credentials: map[string]any{
|
||||
"temp_unschedulable_enabled": true,
|
||||
"temp_unschedulable_rules": []any{
|
||||
map[string]any{
|
||||
"error_code": float64(503),
|
||||
"keywords": []any{"overloaded"},
|
||||
"duration_minutes": float64(10),
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
statusCode: 503,
|
||||
respBody: []byte(`overloaded`),
|
||||
expectFailover: true,
|
||||
expectHandleError: true,
|
||||
},
|
||||
{
|
||||
name: "no_policy_429_failover_via_shouldFailover",
|
||||
account: &Account{
|
||||
ID: 203,
|
||||
Type: AccountTypeAPIKey,
|
||||
Platform: PlatformGemini,
|
||||
},
|
||||
statusCode: 429,
|
||||
respBody: []byte(`{"error":"rate limited"}`),
|
||||
expectFailover: true,
|
||||
expectHandleError: true,
|
||||
expectShouldFailover: true,
|
||||
},
|
||||
{
|
||||
name: "no_policy_400_no_failover",
|
||||
account: &Account{
|
||||
ID: 204,
|
||||
Type: AccountTypeAPIKey,
|
||||
Platform: PlatformGemini,
|
||||
},
|
||||
statusCode: 400,
|
||||
respBody: []byte(`{"error":"bad request"}`),
|
||||
expectFailover: false,
|
||||
expectHandleError: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
repo := &geminiErrorPolicyRepo{}
|
||||
rlSvc := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
|
||||
svc := &GeminiMessagesCompatService{
|
||||
accountRepo: repo,
|
||||
rateLimitService: rlSvc,
|
||||
}
|
||||
|
||||
writer := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(writer)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
|
||||
|
||||
// Simulate the Claude compat error handling path (same logic as native).
|
||||
// This mirrors the inline switch in handleClaudeCompat.
|
||||
var handleErrorCalled bool
|
||||
var gotFailover bool
|
||||
|
||||
ctx := context.Background()
|
||||
statusCode := tt.statusCode
|
||||
respBody := tt.respBody
|
||||
account := tt.account
|
||||
headers := http.Header{}
|
||||
|
||||
if svc.rateLimitService != nil {
|
||||
switch svc.rateLimitService.CheckErrorPolicy(ctx, account, statusCode, respBody) {
|
||||
case ErrorPolicySkipped:
|
||||
// Skipped → return error directly (no handleGeminiUpstreamError, no failover)
|
||||
gotFailover = false
|
||||
handleErrorCalled = false
|
||||
goto verify
|
||||
case ErrorPolicyMatched, ErrorPolicyTempUnscheduled:
|
||||
svc.handleGeminiUpstreamError(ctx, account, statusCode, headers, respBody)
|
||||
handleErrorCalled = true
|
||||
gotFailover = true
|
||||
goto verify
|
||||
}
|
||||
}
|
||||
|
||||
// ErrorPolicyNone → original logic
|
||||
svc.handleGeminiUpstreamError(ctx, account, statusCode, headers, respBody)
|
||||
handleErrorCalled = true
|
||||
if svc.shouldFailoverGeminiUpstreamError(statusCode) {
|
||||
gotFailover = true
|
||||
}
|
||||
|
||||
verify:
|
||||
require.Equal(t, tt.expectFailover, gotFailover, "failover mismatch")
|
||||
require.Equal(t, tt.expectHandleError, handleErrorCalled, "handleGeminiUpstreamError call mismatch")
|
||||
|
||||
if tt.expectShouldFailover {
|
||||
require.True(t, svc.shouldFailoverGeminiUpstreamError(statusCode),
|
||||
"shouldFailoverGeminiUpstreamError should return true for status %d", statusCode)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// TestGeminiErrorPolicy_NilRateLimitService — verifies nil safety
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestGeminiErrorPolicy_NilRateLimitService(t *testing.T) {
|
||||
svc := &GeminiMessagesCompatService{
|
||||
rateLimitService: nil,
|
||||
}
|
||||
|
||||
// When rateLimitService is nil, error policy is skipped → falls through to
|
||||
// shouldFailoverGeminiUpstreamError (original logic).
|
||||
// Verify this doesn't panic and follows expected behavior.
|
||||
|
||||
ctx := context.Background()
|
||||
account := &Account{
|
||||
ID: 300,
|
||||
Type: AccountTypeAPIKey,
|
||||
Platform: PlatformGemini,
|
||||
Credentials: map[string]any{
|
||||
"custom_error_codes_enabled": true,
|
||||
"custom_error_codes": []any{float64(429)},
|
||||
},
|
||||
}
|
||||
|
||||
// The nil check should prevent CheckErrorPolicy from being called
|
||||
if svc.rateLimitService != nil {
|
||||
t.Fatal("rateLimitService should be nil for this test")
|
||||
}
|
||||
|
||||
// shouldFailoverGeminiUpstreamError still works
|
||||
require.True(t, svc.shouldFailoverGeminiUpstreamError(429))
|
||||
require.False(t, svc.shouldFailoverGeminiUpstreamError(400))
|
||||
|
||||
// handleGeminiUpstreamError should not panic with nil rateLimitService
|
||||
require.NotPanics(t, func() {
|
||||
svc.handleGeminiUpstreamError(ctx, account, 500, http.Header{}, []byte(`error`))
|
||||
})
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// geminiErrorPolicyRepo — minimal AccountRepository stub for Gemini error
|
||||
// policy tests. Embeds mockAccountRepoForGemini and adds tracking.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
type geminiErrorPolicyRepo struct {
|
||||
mockAccountRepoForGemini
|
||||
setErrorCalls int
|
||||
setRateLimitedCalls int
|
||||
setTempCalls int
|
||||
}
|
||||
|
||||
func (r *geminiErrorPolicyRepo) SetError(_ context.Context, _ int64, _ string) error {
|
||||
r.setErrorCalls++
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *geminiErrorPolicyRepo) SetRateLimited(_ context.Context, _ int64, _ time.Time) error {
|
||||
r.setRateLimitedCalls++
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *geminiErrorPolicyRepo) SetTempUnschedulable(_ context.Context, _ int64, _ time.Time, _ string) error {
|
||||
r.setTempCalls++
|
||||
return nil
|
||||
}
|
||||
@@ -560,10 +560,7 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex
|
||||
return nil, "", errors.New("gemini api_key not configured")
|
||||
}
|
||||
|
||||
baseURL := strings.TrimSpace(account.GetCredential("base_url"))
|
||||
if baseURL == "" {
|
||||
baseURL = geminicli.AIStudioBaseURL
|
||||
}
|
||||
baseURL := account.GetGeminiBaseURL(geminicli.AIStudioBaseURL)
|
||||
normalizedBaseURL, err := s.validateUpstreamBaseURL(baseURL)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
@@ -640,10 +637,7 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex
|
||||
return upstreamReq, "x-request-id", nil
|
||||
} else {
|
||||
// Mode 2: AI Studio API with OAuth (like API key mode, but using Bearer token)
|
||||
baseURL := strings.TrimSpace(account.GetCredential("base_url"))
|
||||
if baseURL == "" {
|
||||
baseURL = geminicli.AIStudioBaseURL
|
||||
}
|
||||
baseURL := account.GetGeminiBaseURL(geminicli.AIStudioBaseURL)
|
||||
normalizedBaseURL, err := s.validateUpstreamBaseURL(baseURL)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
@@ -837,38 +831,47 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex
|
||||
|
||||
if resp.StatusCode >= 400 {
|
||||
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
|
||||
tempMatched := false
|
||||
// 统一错误策略:自定义错误码 + 临时不可调度
|
||||
if s.rateLimitService != nil {
|
||||
tempMatched = s.rateLimitService.HandleTempUnschedulable(ctx, account, resp.StatusCode, respBody)
|
||||
}
|
||||
s.handleGeminiUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
|
||||
if tempMatched {
|
||||
upstreamReqID := resp.Header.Get(requestIDHeader)
|
||||
if upstreamReqID == "" {
|
||||
upstreamReqID = resp.Header.Get("x-goog-request-id")
|
||||
}
|
||||
upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody))
|
||||
upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
|
||||
upstreamDetail := ""
|
||||
if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody {
|
||||
maxBytes := s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes
|
||||
if maxBytes <= 0 {
|
||||
maxBytes = 2048
|
||||
switch s.rateLimitService.CheckErrorPolicy(ctx, account, resp.StatusCode, respBody) {
|
||||
case ErrorPolicySkipped:
|
||||
upstreamReqID := resp.Header.Get(requestIDHeader)
|
||||
if upstreamReqID == "" {
|
||||
upstreamReqID = resp.Header.Get("x-goog-request-id")
|
||||
}
|
||||
upstreamDetail = truncateString(string(respBody), maxBytes)
|
||||
return nil, s.writeGeminiMappedError(c, account, resp.StatusCode, upstreamReqID, respBody)
|
||||
case ErrorPolicyMatched, ErrorPolicyTempUnscheduled:
|
||||
s.handleGeminiUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
|
||||
upstreamReqID := resp.Header.Get(requestIDHeader)
|
||||
if upstreamReqID == "" {
|
||||
upstreamReqID = resp.Header.Get("x-goog-request-id")
|
||||
}
|
||||
upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody))
|
||||
upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
|
||||
upstreamDetail := ""
|
||||
if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody {
|
||||
maxBytes := s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes
|
||||
if maxBytes <= 0 {
|
||||
maxBytes = 2048
|
||||
}
|
||||
upstreamDetail = truncateString(string(respBody), maxBytes)
|
||||
}
|
||||
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
|
||||
Platform: account.Platform,
|
||||
AccountID: account.ID,
|
||||
AccountName: account.Name,
|
||||
UpstreamStatusCode: resp.StatusCode,
|
||||
UpstreamRequestID: upstreamReqID,
|
||||
Kind: "failover",
|
||||
Message: upstreamMsg,
|
||||
Detail: upstreamDetail,
|
||||
})
|
||||
return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: respBody}
|
||||
}
|
||||
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
|
||||
Platform: account.Platform,
|
||||
AccountID: account.ID,
|
||||
AccountName: account.Name,
|
||||
UpstreamStatusCode: resp.StatusCode,
|
||||
UpstreamRequestID: upstreamReqID,
|
||||
Kind: "failover",
|
||||
Message: upstreamMsg,
|
||||
Detail: upstreamDetail,
|
||||
})
|
||||
return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: respBody}
|
||||
}
|
||||
|
||||
// ErrorPolicyNone → 原有逻辑
|
||||
s.handleGeminiUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
|
||||
if s.shouldFailoverGeminiUpstreamError(resp.StatusCode) {
|
||||
upstreamReqID := resp.Header.Get(requestIDHeader)
|
||||
if upstreamReqID == "" {
|
||||
@@ -1026,10 +1029,7 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
|
||||
return nil, "", errors.New("gemini api_key not configured")
|
||||
}
|
||||
|
||||
baseURL := strings.TrimSpace(account.GetCredential("base_url"))
|
||||
if baseURL == "" {
|
||||
baseURL = geminicli.AIStudioBaseURL
|
||||
}
|
||||
baseURL := account.GetGeminiBaseURL(geminicli.AIStudioBaseURL)
|
||||
normalizedBaseURL, err := s.validateUpstreamBaseURL(baseURL)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
@@ -1097,10 +1097,7 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
|
||||
return upstreamReq, "x-request-id", nil
|
||||
} else {
|
||||
// Mode 2: AI Studio API with OAuth (like API key mode, but using Bearer token)
|
||||
baseURL := strings.TrimSpace(account.GetCredential("base_url"))
|
||||
if baseURL == "" {
|
||||
baseURL = geminicli.AIStudioBaseURL
|
||||
}
|
||||
baseURL := account.GetGeminiBaseURL(geminicli.AIStudioBaseURL)
|
||||
normalizedBaseURL, err := s.validateUpstreamBaseURL(baseURL)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
@@ -1261,14 +1258,9 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
|
||||
|
||||
if resp.StatusCode >= 400 {
|
||||
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
|
||||
tempMatched := false
|
||||
if s.rateLimitService != nil {
|
||||
tempMatched = s.rateLimitService.HandleTempUnschedulable(ctx, account, resp.StatusCode, respBody)
|
||||
}
|
||||
s.handleGeminiUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
|
||||
|
||||
// Best-effort fallback for OAuth tokens missing AI Studio scopes when calling countTokens.
|
||||
// This avoids Gemini SDKs failing hard during preflight token counting.
|
||||
// Checked before error policy so it always works regardless of custom error codes.
|
||||
if action == "countTokens" && isOAuth && isGeminiInsufficientScope(resp.Header, respBody) {
|
||||
estimated := estimateGeminiCountTokens(body)
|
||||
c.JSON(http.StatusOK, map[string]any{"totalTokens": estimated})
|
||||
@@ -1282,30 +1274,46 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
|
||||
}, nil
|
||||
}
|
||||
|
||||
if tempMatched {
|
||||
evBody := unwrapIfNeeded(isOAuth, respBody)
|
||||
upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(evBody))
|
||||
upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
|
||||
upstreamDetail := ""
|
||||
if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody {
|
||||
maxBytes := s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes
|
||||
if maxBytes <= 0 {
|
||||
maxBytes = 2048
|
||||
// 统一错误策略:自定义错误码 + 临时不可调度
|
||||
if s.rateLimitService != nil {
|
||||
switch s.rateLimitService.CheckErrorPolicy(ctx, account, resp.StatusCode, respBody) {
|
||||
case ErrorPolicySkipped:
|
||||
respBody = unwrapIfNeeded(isOAuth, respBody)
|
||||
contentType := resp.Header.Get("Content-Type")
|
||||
if contentType == "" {
|
||||
contentType = "application/json"
|
||||
}
|
||||
upstreamDetail = truncateString(string(evBody), maxBytes)
|
||||
c.Data(resp.StatusCode, contentType, respBody)
|
||||
return nil, fmt.Errorf("gemini upstream error: %d (skipped by error policy)", resp.StatusCode)
|
||||
case ErrorPolicyMatched, ErrorPolicyTempUnscheduled:
|
||||
s.handleGeminiUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
|
||||
evBody := unwrapIfNeeded(isOAuth, respBody)
|
||||
upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(evBody))
|
||||
upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
|
||||
upstreamDetail := ""
|
||||
if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody {
|
||||
maxBytes := s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes
|
||||
if maxBytes <= 0 {
|
||||
maxBytes = 2048
|
||||
}
|
||||
upstreamDetail = truncateString(string(evBody), maxBytes)
|
||||
}
|
||||
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
|
||||
Platform: account.Platform,
|
||||
AccountID: account.ID,
|
||||
AccountName: account.Name,
|
||||
UpstreamStatusCode: resp.StatusCode,
|
||||
UpstreamRequestID: requestID,
|
||||
Kind: "failover",
|
||||
Message: upstreamMsg,
|
||||
Detail: upstreamDetail,
|
||||
})
|
||||
return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: respBody}
|
||||
}
|
||||
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
|
||||
Platform: account.Platform,
|
||||
AccountID: account.ID,
|
||||
AccountName: account.Name,
|
||||
UpstreamStatusCode: resp.StatusCode,
|
||||
UpstreamRequestID: requestID,
|
||||
Kind: "failover",
|
||||
Message: upstreamMsg,
|
||||
Detail: upstreamDetail,
|
||||
})
|
||||
return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: respBody}
|
||||
}
|
||||
|
||||
// ErrorPolicyNone → 原有逻辑
|
||||
s.handleGeminiUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
|
||||
if s.shouldFailoverGeminiUpstreamError(resp.StatusCode) {
|
||||
evBody := unwrapIfNeeded(isOAuth, respBody)
|
||||
upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(evBody))
|
||||
@@ -2420,10 +2428,7 @@ func (s *GeminiMessagesCompatService) ForwardAIStudioGET(ctx context.Context, ac
|
||||
return nil, errors.New("invalid path")
|
||||
}
|
||||
|
||||
baseURL := strings.TrimSpace(account.GetCredential("base_url"))
|
||||
if baseURL == "" {
|
||||
baseURL = geminicli.AIStudioBaseURL
|
||||
}
|
||||
baseURL := account.GetGeminiBaseURL(geminicli.AIStudioBaseURL)
|
||||
normalizedBaseURL, err := s.validateUpstreamBaseURL(baseURL)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
|
||||
@@ -66,7 +66,12 @@ func (m *mockAccountRepoForGemini) Create(ctx context.Context, account *Account)
|
||||
func (m *mockAccountRepoForGemini) GetByCRSAccountID(ctx context.Context, crsAccountID string) (*Account, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (m *mockAccountRepoForGemini) FindByExtraField(ctx context.Context, key string, value interface{}) ([]Account, error) {
|
||||
|
||||
func (m *mockAccountRepoForGemini) FindByExtraField(ctx context.Context, key string, value any) ([]Account, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (m *mockAccountRepoForGemini) ListCRSAccountIDs(ctx context.Context) (map[string]int64, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (m *mockAccountRepoForGemini) Update(ctx context.Context, account *Account) error { return nil }
|
||||
@@ -136,9 +141,6 @@ func (m *mockAccountRepoForGemini) ListSchedulableByGroupIDAndPlatforms(ctx cont
|
||||
func (m *mockAccountRepoForGemini) SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error {
|
||||
return nil
|
||||
}
|
||||
func (m *mockAccountRepoForGemini) SetAntigravityQuotaScopeLimit(ctx context.Context, id int64, scope AntigravityQuotaScope, resetAt time.Time) error {
|
||||
return nil
|
||||
}
|
||||
func (m *mockAccountRepoForGemini) SetModelRateLimit(ctx context.Context, id int64, scope string, resetAt time.Time) error {
|
||||
return nil
|
||||
}
|
||||
@@ -229,6 +231,10 @@ func (m *mockGroupRepoForGemini) GetAccountIDsByGroupIDs(ctx context.Context, gr
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (m *mockGroupRepoForGemini) UpdateSortOrders(ctx context.Context, updates []GroupSortOrderUpdate) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
var _ GroupRepository = (*mockGroupRepoForGemini)(nil)
|
||||
|
||||
// mockGatewayCacheForGemini Gemini 测试用的 cache mock
|
||||
@@ -268,22 +274,6 @@ func (m *mockGatewayCacheForGemini) DeleteSessionAccountID(ctx context.Context,
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *mockGatewayCacheForGemini) IncrModelCallCount(ctx context.Context, accountID int64, model string) (int64, error) {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
func (m *mockGatewayCacheForGemini) GetModelLoadBatch(ctx context.Context, accountIDs []int64, model string) (map[int64]*ModelLoadInfo, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (m *mockGatewayCacheForGemini) FindGeminiSession(ctx context.Context, groupID int64, prefixHash, digestChain string) (uuid string, accountID int64, found bool) {
|
||||
return "", 0, false
|
||||
}
|
||||
|
||||
func (m *mockGatewayCacheForGemini) SaveGeminiSession(ctx context.Context, groupID int64, prefixHash, digestChain, uuid string, accountID int64) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_GeminiPlatform 测试 Gemini 单平台选择
|
||||
func TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_GeminiPlatform(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
|
||||
@@ -6,26 +6,11 @@ import (
|
||||
"encoding/json"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
|
||||
"github.com/cespare/xxhash/v2"
|
||||
)
|
||||
|
||||
// Gemini 会话 ID Fallback 相关常量
|
||||
const (
|
||||
// geminiSessionTTLSeconds Gemini 会话缓存 TTL(5 分钟)
|
||||
geminiSessionTTLSeconds = 300
|
||||
|
||||
// geminiSessionKeyPrefix Gemini 会话 Redis key 前缀
|
||||
geminiSessionKeyPrefix = "gemini:sess:"
|
||||
)
|
||||
|
||||
// GeminiSessionTTL 返回 Gemini 会话缓存 TTL
|
||||
func GeminiSessionTTL() time.Duration {
|
||||
return geminiSessionTTLSeconds * time.Second
|
||||
}
|
||||
|
||||
// shortHash 使用 XXHash64 + Base36 生成短 hash(16 字符)
|
||||
// XXHash64 比 SHA256 快约 10 倍,Base36 比 Hex 短约 20%
|
||||
func shortHash(data []byte) string {
|
||||
@@ -79,35 +64,6 @@ func GenerateGeminiPrefixHash(userID, apiKeyID int64, ip, userAgent, platform, m
|
||||
return base64.RawURLEncoding.EncodeToString(hash[:12])
|
||||
}
|
||||
|
||||
// BuildGeminiSessionKey 构建 Gemini 会话 Redis key
|
||||
// 格式: gemini:sess:{groupID}:{prefixHash}:{digestChain}
|
||||
func BuildGeminiSessionKey(groupID int64, prefixHash, digestChain string) string {
|
||||
return geminiSessionKeyPrefix + strconv.FormatInt(groupID, 10) + ":" + prefixHash + ":" + digestChain
|
||||
}
|
||||
|
||||
// GenerateDigestChainPrefixes 生成摘要链的所有前缀(从长到短)
|
||||
// 用于 MGET 批量查询最长匹配
|
||||
func GenerateDigestChainPrefixes(chain string) []string {
|
||||
if chain == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
var prefixes []string
|
||||
c := chain
|
||||
|
||||
for c != "" {
|
||||
prefixes = append(prefixes, c)
|
||||
// 找到最后一个 "-" 的位置
|
||||
if i := strings.LastIndex(c, "-"); i > 0 {
|
||||
c = c[:i]
|
||||
} else {
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
return prefixes
|
||||
}
|
||||
|
||||
// ParseGeminiSessionValue 解析 Gemini 会话缓存值
|
||||
// 格式: {uuid}:{accountID}
|
||||
func ParseGeminiSessionValue(value string) (uuid string, accountID int64, ok bool) {
|
||||
@@ -139,15 +95,6 @@ func FormatGeminiSessionValue(uuid string, accountID int64) string {
|
||||
// geminiDigestSessionKeyPrefix Gemini 摘要 fallback 会话 key 前缀
|
||||
const geminiDigestSessionKeyPrefix = "gemini:digest:"
|
||||
|
||||
// geminiTrieKeyPrefix Gemini Trie 会话 key 前缀
|
||||
const geminiTrieKeyPrefix = "gemini:trie:"
|
||||
|
||||
// BuildGeminiTrieKey 构建 Gemini Trie Redis key
|
||||
// 格式: gemini:trie:{groupID}:{prefixHash}
|
||||
func BuildGeminiTrieKey(groupID int64, prefixHash string) string {
|
||||
return geminiTrieKeyPrefix + strconv.FormatInt(groupID, 10) + ":" + prefixHash
|
||||
}
|
||||
|
||||
// GenerateGeminiDigestSessionKey 生成 Gemini 摘要 fallback 的 sessionKey
|
||||
// 组合 prefixHash 前 8 位 + uuid 前 8 位,确保不同会话产生不同的 sessionKey
|
||||
// 用于在 SelectAccountWithLoadAwareness 中保持粘性会话
|
||||
|
||||
@@ -1,41 +1,14 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
|
||||
)
|
||||
|
||||
// mockGeminiSessionCache 模拟 Redis 缓存
|
||||
type mockGeminiSessionCache struct {
|
||||
sessions map[string]string // key -> value
|
||||
}
|
||||
|
||||
func newMockGeminiSessionCache() *mockGeminiSessionCache {
|
||||
return &mockGeminiSessionCache{sessions: make(map[string]string)}
|
||||
}
|
||||
|
||||
func (m *mockGeminiSessionCache) Save(groupID int64, prefixHash, digestChain, uuid string, accountID int64) {
|
||||
key := BuildGeminiSessionKey(groupID, prefixHash, digestChain)
|
||||
value := FormatGeminiSessionValue(uuid, accountID)
|
||||
m.sessions[key] = value
|
||||
}
|
||||
|
||||
func (m *mockGeminiSessionCache) Find(groupID int64, prefixHash, digestChain string) (uuid string, accountID int64, found bool) {
|
||||
prefixes := GenerateDigestChainPrefixes(digestChain)
|
||||
for _, p := range prefixes {
|
||||
key := BuildGeminiSessionKey(groupID, prefixHash, p)
|
||||
if val, ok := m.sessions[key]; ok {
|
||||
return ParseGeminiSessionValue(val)
|
||||
}
|
||||
}
|
||||
return "", 0, false
|
||||
}
|
||||
|
||||
// TestGeminiSessionContinuousConversation 测试连续会话的摘要链匹配
|
||||
func TestGeminiSessionContinuousConversation(t *testing.T) {
|
||||
cache := newMockGeminiSessionCache()
|
||||
store := NewDigestSessionStore()
|
||||
groupID := int64(1)
|
||||
prefixHash := "test_prefix_hash"
|
||||
sessionUUID := "session-uuid-12345"
|
||||
@@ -54,13 +27,13 @@ func TestGeminiSessionContinuousConversation(t *testing.T) {
|
||||
t.Logf("Round 1 chain: %s", chain1)
|
||||
|
||||
// 第一轮:没有找到会话,创建新会话
|
||||
_, _, found := cache.Find(groupID, prefixHash, chain1)
|
||||
_, _, _, found := store.Find(groupID, prefixHash, chain1)
|
||||
if found {
|
||||
t.Error("Round 1: should not find existing session")
|
||||
}
|
||||
|
||||
// 保存第一轮会话
|
||||
cache.Save(groupID, prefixHash, chain1, sessionUUID, accountID)
|
||||
// 保存第一轮会话(首轮无旧 chain)
|
||||
store.Save(groupID, prefixHash, chain1, sessionUUID, accountID, "")
|
||||
|
||||
// 模拟第二轮对话(用户继续对话)
|
||||
req2 := &antigravity.GeminiRequest{
|
||||
@@ -77,7 +50,7 @@ func TestGeminiSessionContinuousConversation(t *testing.T) {
|
||||
t.Logf("Round 2 chain: %s", chain2)
|
||||
|
||||
// 第二轮:应该能找到会话(通过前缀匹配)
|
||||
foundUUID, foundAccID, found := cache.Find(groupID, prefixHash, chain2)
|
||||
foundUUID, foundAccID, matchedChain, found := store.Find(groupID, prefixHash, chain2)
|
||||
if !found {
|
||||
t.Error("Round 2: should find session via prefix matching")
|
||||
}
|
||||
@@ -88,8 +61,8 @@ func TestGeminiSessionContinuousConversation(t *testing.T) {
|
||||
t.Errorf("Round 2: expected accountID %d, got %d", accountID, foundAccID)
|
||||
}
|
||||
|
||||
// 保存第二轮会话
|
||||
cache.Save(groupID, prefixHash, chain2, sessionUUID, accountID)
|
||||
// 保存第二轮会话,传入 Find 返回的 matchedChain 以删旧 key
|
||||
store.Save(groupID, prefixHash, chain2, sessionUUID, accountID, matchedChain)
|
||||
|
||||
// 模拟第三轮对话
|
||||
req3 := &antigravity.GeminiRequest{
|
||||
@@ -108,7 +81,7 @@ func TestGeminiSessionContinuousConversation(t *testing.T) {
|
||||
t.Logf("Round 3 chain: %s", chain3)
|
||||
|
||||
// 第三轮:应该能找到会话(通过第二轮的前缀匹配)
|
||||
foundUUID, foundAccID, found = cache.Find(groupID, prefixHash, chain3)
|
||||
foundUUID, foundAccID, _, found = store.Find(groupID, prefixHash, chain3)
|
||||
if !found {
|
||||
t.Error("Round 3: should find session via prefix matching")
|
||||
}
|
||||
@@ -118,13 +91,11 @@ func TestGeminiSessionContinuousConversation(t *testing.T) {
|
||||
if foundAccID != accountID {
|
||||
t.Errorf("Round 3: expected accountID %d, got %d", accountID, foundAccID)
|
||||
}
|
||||
|
||||
t.Log("✓ Continuous conversation session matching works correctly!")
|
||||
}
|
||||
|
||||
// TestGeminiSessionDifferentConversations 测试不同会话不会错误匹配
|
||||
func TestGeminiSessionDifferentConversations(t *testing.T) {
|
||||
cache := newMockGeminiSessionCache()
|
||||
store := NewDigestSessionStore()
|
||||
groupID := int64(1)
|
||||
prefixHash := "test_prefix_hash"
|
||||
|
||||
@@ -135,7 +106,7 @@ func TestGeminiSessionDifferentConversations(t *testing.T) {
|
||||
},
|
||||
}
|
||||
chain1 := BuildGeminiDigestChain(req1)
|
||||
cache.Save(groupID, prefixHash, chain1, "session-1", 100)
|
||||
store.Save(groupID, prefixHash, chain1, "session-1", 100, "")
|
||||
|
||||
// 第二个完全不同的会话
|
||||
req2 := &antigravity.GeminiRequest{
|
||||
@@ -146,61 +117,29 @@ func TestGeminiSessionDifferentConversations(t *testing.T) {
|
||||
chain2 := BuildGeminiDigestChain(req2)
|
||||
|
||||
// 不同会话不应该匹配
|
||||
_, _, found := cache.Find(groupID, prefixHash, chain2)
|
||||
_, _, _, found := store.Find(groupID, prefixHash, chain2)
|
||||
if found {
|
||||
t.Error("Different conversations should not match")
|
||||
}
|
||||
|
||||
t.Log("✓ Different conversations are correctly isolated!")
|
||||
}
|
||||
|
||||
// TestGeminiSessionPrefixMatchingOrder 测试前缀匹配的优先级(最长匹配优先)
|
||||
func TestGeminiSessionPrefixMatchingOrder(t *testing.T) {
|
||||
cache := newMockGeminiSessionCache()
|
||||
store := NewDigestSessionStore()
|
||||
groupID := int64(1)
|
||||
prefixHash := "test_prefix_hash"
|
||||
|
||||
// 创建一个三轮对话
|
||||
req := &antigravity.GeminiRequest{
|
||||
SystemInstruction: &antigravity.GeminiContent{
|
||||
Parts: []antigravity.GeminiPart{{Text: "System prompt"}},
|
||||
},
|
||||
Contents: []antigravity.GeminiContent{
|
||||
{Role: "user", Parts: []antigravity.GeminiPart{{Text: "Q1"}}},
|
||||
{Role: "model", Parts: []antigravity.GeminiPart{{Text: "A1"}}},
|
||||
{Role: "user", Parts: []antigravity.GeminiPart{{Text: "Q2"}}},
|
||||
},
|
||||
}
|
||||
fullChain := BuildGeminiDigestChain(req)
|
||||
prefixes := GenerateDigestChainPrefixes(fullChain)
|
||||
|
||||
t.Logf("Full chain: %s", fullChain)
|
||||
t.Logf("Prefixes (longest first): %v", prefixes)
|
||||
|
||||
// 验证前缀生成顺序(从长到短)
|
||||
if len(prefixes) != 4 {
|
||||
t.Errorf("Expected 4 prefixes, got %d", len(prefixes))
|
||||
}
|
||||
|
||||
// 保存不同轮次的会话到不同账号
|
||||
// 第一轮(最短前缀)-> 账号 1
|
||||
cache.Save(groupID, prefixHash, prefixes[3], "session-round1", 1)
|
||||
// 第二轮 -> 账号 2
|
||||
cache.Save(groupID, prefixHash, prefixes[2], "session-round2", 2)
|
||||
// 第三轮(最长前缀,完整链)-> 账号 3
|
||||
cache.Save(groupID, prefixHash, prefixes[0], "session-round3", 3)
|
||||
store.Save(groupID, prefixHash, "s:sys-u:q1", "session-round1", 1, "")
|
||||
store.Save(groupID, prefixHash, "s:sys-u:q1-m:a1", "session-round2", 2, "")
|
||||
store.Save(groupID, prefixHash, "s:sys-u:q1-m:a1-u:q2", "session-round3", 3, "")
|
||||
|
||||
// 查找应该返回最长匹配(账号 3)
|
||||
_, accID, found := cache.Find(groupID, prefixHash, fullChain)
|
||||
// 查找更长的链,应该返回最长匹配(账号 3)
|
||||
_, accID, _, found := store.Find(groupID, prefixHash, "s:sys-u:q1-m:a1-u:q2-m:a2")
|
||||
if !found {
|
||||
t.Error("Should find session")
|
||||
}
|
||||
if accID != 3 {
|
||||
t.Errorf("Should match longest prefix (account 3), got account %d", accID)
|
||||
}
|
||||
|
||||
t.Log("✓ Longest prefix matching works correctly!")
|
||||
}
|
||||
|
||||
// 确保 context 包被使用(避免未使用的导入警告)
|
||||
var _ = context.Background
|
||||
|
||||
@@ -152,61 +152,6 @@ func TestGenerateGeminiPrefixHash(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestGenerateDigestChainPrefixes(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
chain string
|
||||
want []string
|
||||
wantLen int
|
||||
}{
|
||||
{
|
||||
name: "empty",
|
||||
chain: "",
|
||||
wantLen: 0,
|
||||
},
|
||||
{
|
||||
name: "single part",
|
||||
chain: "u:abc123",
|
||||
want: []string{"u:abc123"},
|
||||
wantLen: 1,
|
||||
},
|
||||
{
|
||||
name: "two parts",
|
||||
chain: "s:xyz-u:abc",
|
||||
want: []string{"s:xyz-u:abc", "s:xyz"},
|
||||
wantLen: 2,
|
||||
},
|
||||
{
|
||||
name: "four parts",
|
||||
chain: "s:a-u:b-m:c-u:d",
|
||||
want: []string{"s:a-u:b-m:c-u:d", "s:a-u:b-m:c", "s:a-u:b", "s:a"},
|
||||
wantLen: 4,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := GenerateDigestChainPrefixes(tt.chain)
|
||||
|
||||
if len(result) != tt.wantLen {
|
||||
t.Errorf("expected %d prefixes, got %d: %v", tt.wantLen, len(result), result)
|
||||
}
|
||||
|
||||
if tt.want != nil {
|
||||
for i, want := range tt.want {
|
||||
if i >= len(result) {
|
||||
t.Errorf("missing prefix at index %d", i)
|
||||
continue
|
||||
}
|
||||
if result[i] != want {
|
||||
t.Errorf("prefix[%d]: expected %s, got %s", i, want, result[i])
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseGeminiSessionValue(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
@@ -442,40 +387,3 @@ func TestGenerateGeminiDigestSessionKey(t *testing.T) {
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestBuildGeminiTrieKey(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
groupID int64
|
||||
prefixHash string
|
||||
want string
|
||||
}{
|
||||
{
|
||||
name: "normal",
|
||||
groupID: 123,
|
||||
prefixHash: "abcdef12",
|
||||
want: "gemini:trie:123:abcdef12",
|
||||
},
|
||||
{
|
||||
name: "zero group",
|
||||
groupID: 0,
|
||||
prefixHash: "xyz",
|
||||
want: "gemini:trie:0:xyz",
|
||||
},
|
||||
{
|
||||
name: "empty prefix",
|
||||
groupID: 1,
|
||||
prefixHash: "",
|
||||
want: "gemini:trie:1:",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := BuildGeminiTrieKey(tt.groupID, tt.prefixHash)
|
||||
if got != tt.want {
|
||||
t.Errorf("BuildGeminiTrieKey(%d, %q) = %q, want %q", tt.groupID, tt.prefixHash, got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
1213
backend/internal/service/generate_session_hash_test.go
Normal file
1213
backend/internal/service/generate_session_hash_test.go
Normal file
File diff suppressed because it is too large
Load Diff
@@ -51,6 +51,9 @@ type Group struct {
|
||||
// 可选值: claude, gemini_text, gemini_image
|
||||
SupportedModelScopes []string
|
||||
|
||||
// 分组排序
|
||||
SortOrder int
|
||||
|
||||
CreatedAt time.Time
|
||||
UpdatedAt time.Time
|
||||
|
||||
|
||||
@@ -33,6 +33,14 @@ type GroupRepository interface {
|
||||
GetAccountIDsByGroupIDs(ctx context.Context, groupIDs []int64) ([]int64, error)
|
||||
// BindAccountsToGroup 将多个账号绑定到指定分组
|
||||
BindAccountsToGroup(ctx context.Context, groupID int64, accountIDs []int64) error
|
||||
// UpdateSortOrders 批量更新分组排序
|
||||
UpdateSortOrders(ctx context.Context, updates []GroupSortOrderUpdate) error
|
||||
}
|
||||
|
||||
// GroupSortOrderUpdate 分组排序更新
|
||||
type GroupSortOrderUpdate struct {
|
||||
ID int64 `json:"id"`
|
||||
SortOrder int `json:"sort_order"`
|
||||
}
|
||||
|
||||
// CreateGroupRequest 创建分组请求
|
||||
|
||||
@@ -318,110 +318,6 @@ func TestGetModelRateLimitRemainingTime(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetQuotaScopeRateLimitRemainingTime(t *testing.T) {
|
||||
now := time.Now()
|
||||
future10m := 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
|
||||
minExpected time.Duration
|
||||
maxExpected time.Duration
|
||||
}{
|
||||
{
|
||||
name: "nil account",
|
||||
account: nil,
|
||||
requestedModel: "claude-sonnet-4-5",
|
||||
minExpected: 0,
|
||||
maxExpected: 0,
|
||||
},
|
||||
{
|
||||
name: "non-antigravity platform",
|
||||
account: &Account{
|
||||
Platform: PlatformAnthropic,
|
||||
Extra: map[string]any{
|
||||
antigravityQuotaScopesKey: map[string]any{
|
||||
"claude": map[string]any{
|
||||
"rate_limit_reset_at": future10m,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
requestedModel: "claude-sonnet-4-5",
|
||||
minExpected: 0,
|
||||
maxExpected: 0,
|
||||
},
|
||||
{
|
||||
name: "claude scope rate limited",
|
||||
account: &Account{
|
||||
Platform: PlatformAntigravity,
|
||||
Extra: map[string]any{
|
||||
antigravityQuotaScopesKey: map[string]any{
|
||||
"claude": map[string]any{
|
||||
"rate_limit_reset_at": future10m,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
requestedModel: "claude-sonnet-4-5",
|
||||
minExpected: 9 * time.Minute,
|
||||
maxExpected: 11 * time.Minute,
|
||||
},
|
||||
{
|
||||
name: "gemini_text scope rate limited",
|
||||
account: &Account{
|
||||
Platform: PlatformAntigravity,
|
||||
Extra: map[string]any{
|
||||
antigravityQuotaScopesKey: map[string]any{
|
||||
"gemini_text": map[string]any{
|
||||
"rate_limit_reset_at": future10m,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
requestedModel: "gemini-3-flash",
|
||||
minExpected: 9 * time.Minute,
|
||||
maxExpected: 11 * time.Minute,
|
||||
},
|
||||
{
|
||||
name: "expired scope rate limit",
|
||||
account: &Account{
|
||||
Platform: PlatformAntigravity,
|
||||
Extra: map[string]any{
|
||||
antigravityQuotaScopesKey: map[string]any{
|
||||
"claude": map[string]any{
|
||||
"rate_limit_reset_at": past,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
requestedModel: "claude-sonnet-4-5",
|
||||
minExpected: 0,
|
||||
maxExpected: 0,
|
||||
},
|
||||
{
|
||||
name: "unsupported model",
|
||||
account: &Account{
|
||||
Platform: PlatformAntigravity,
|
||||
},
|
||||
requestedModel: "gpt-4",
|
||||
minExpected: 0,
|
||||
maxExpected: 0,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := tt.account.GetQuotaScopeRateLimitRemainingTime(tt.requestedModel)
|
||||
if result < tt.minExpected || result > tt.maxExpected {
|
||||
t.Errorf("GetQuotaScopeRateLimitRemainingTime() = %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)
|
||||
@@ -442,45 +338,19 @@ func TestGetRateLimitRemainingTime(t *testing.T) {
|
||||
maxExpected: 0,
|
||||
},
|
||||
{
|
||||
name: "model remaining > scope remaining - returns model",
|
||||
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, // 15 分钟
|
||||
},
|
||||
},
|
||||
antigravityQuotaScopesKey: map[string]any{
|
||||
"claude": map[string]any{
|
||||
"rate_limit_reset_at": future5m, // 5 分钟
|
||||
"rate_limit_reset_at": future15m,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
requestedModel: "claude-sonnet-4-5",
|
||||
minExpected: 14 * time.Minute, // 应返回较大的 15 分钟
|
||||
maxExpected: 16 * time.Minute,
|
||||
},
|
||||
{
|
||||
name: "scope remaining > model remaining - returns scope",
|
||||
account: &Account{
|
||||
Platform: PlatformAntigravity,
|
||||
Extra: map[string]any{
|
||||
modelRateLimitsKey: map[string]any{
|
||||
"claude-sonnet-4-5": map[string]any{
|
||||
"rate_limit_reset_at": future5m, // 5 分钟
|
||||
},
|
||||
},
|
||||
antigravityQuotaScopesKey: map[string]any{
|
||||
"claude": map[string]any{
|
||||
"rate_limit_reset_at": future15m, // 15 分钟
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
requestedModel: "claude-sonnet-4-5",
|
||||
minExpected: 14 * time.Minute, // 应返回较大的 15 分钟
|
||||
minExpected: 14 * time.Minute,
|
||||
maxExpected: 16 * time.Minute,
|
||||
},
|
||||
{
|
||||
@@ -499,22 +369,6 @@ func TestGetRateLimitRemainingTime(t *testing.T) {
|
||||
minExpected: 4 * time.Minute,
|
||||
maxExpected: 6 * time.Minute,
|
||||
},
|
||||
{
|
||||
name: "only scope rate limited",
|
||||
account: &Account{
|
||||
Platform: PlatformAntigravity,
|
||||
Extra: map[string]any{
|
||||
antigravityQuotaScopesKey: map[string]any{
|
||||
"claude": 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{
|
||||
|
||||
@@ -582,10 +582,6 @@ func (s *OpenAIGatewayService) SelectAccountWithLoadAwareness(ctx context.Contex
|
||||
}
|
||||
}
|
||||
} else {
|
||||
type accountWithLoad struct {
|
||||
account *Account
|
||||
loadInfo *AccountLoadInfo
|
||||
}
|
||||
var available []accountWithLoad
|
||||
for _, acc := range candidates {
|
||||
loadInfo := loadMap[acc.ID]
|
||||
@@ -620,6 +616,7 @@ func (s *OpenAIGatewayService) SelectAccountWithLoadAwareness(ctx context.Contex
|
||||
return a.account.LastUsedAt.Before(*b.account.LastUsedAt)
|
||||
}
|
||||
})
|
||||
shuffleWithinSortGroups(available)
|
||||
|
||||
for _, item := range available {
|
||||
result, err := s.tryAcquireAccountSlot(ctx, item.account.ID, item.account.Concurrency)
|
||||
|
||||
@@ -209,22 +209,6 @@ func (c *stubGatewayCache) DeleteSessionAccountID(ctx context.Context, groupID i
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *stubGatewayCache) IncrModelCallCount(ctx context.Context, accountID int64, model string) (int64, error) {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
func (c *stubGatewayCache) GetModelLoadBatch(ctx context.Context, accountIDs []int64, model string) (map[int64]*ModelLoadInfo, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (c *stubGatewayCache) FindGeminiSession(ctx context.Context, groupID int64, prefixHash, digestChain string) (uuid string, accountID int64, found bool) {
|
||||
return "", 0, false
|
||||
}
|
||||
|
||||
func (c *stubGatewayCache) SaveGeminiSession(ctx context.Context, groupID int64, prefixHash, digestChain, uuid string, accountID int64) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestOpenAISelectAccountWithLoadAwareness_FiltersUnschedulable(t *testing.T) {
|
||||
now := time.Now()
|
||||
resetAt := now.Add(10 * time.Minute)
|
||||
|
||||
@@ -66,7 +66,6 @@ func (s *OpsService) GetAccountAvailabilityStats(ctx context.Context, platformFi
|
||||
}
|
||||
|
||||
isAvailable := acc.Status == StatusActive && acc.Schedulable && !isRateLimited && !isOverloaded && !isTempUnsched
|
||||
scopeRateLimits := acc.GetAntigravityScopeRateLimits()
|
||||
|
||||
if acc.Platform != "" {
|
||||
if _, ok := platform[acc.Platform]; !ok {
|
||||
@@ -85,14 +84,6 @@ func (s *OpsService) GetAccountAvailabilityStats(ctx context.Context, platformFi
|
||||
if hasError {
|
||||
p.ErrorCount++
|
||||
}
|
||||
if len(scopeRateLimits) > 0 {
|
||||
if p.ScopeRateLimitCount == nil {
|
||||
p.ScopeRateLimitCount = make(map[string]int64)
|
||||
}
|
||||
for scope := range scopeRateLimits {
|
||||
p.ScopeRateLimitCount[scope]++
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for _, grp := range acc.Groups {
|
||||
@@ -117,14 +108,6 @@ func (s *OpsService) GetAccountAvailabilityStats(ctx context.Context, platformFi
|
||||
if hasError {
|
||||
g.ErrorCount++
|
||||
}
|
||||
if len(scopeRateLimits) > 0 {
|
||||
if g.ScopeRateLimitCount == nil {
|
||||
g.ScopeRateLimitCount = make(map[string]int64)
|
||||
}
|
||||
for scope := range scopeRateLimits {
|
||||
g.ScopeRateLimitCount[scope]++
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
displayGroupID := int64(0)
|
||||
@@ -157,9 +140,6 @@ func (s *OpsService) GetAccountAvailabilityStats(ctx context.Context, platformFi
|
||||
item.RateLimitRemainingSec = &remainingSec
|
||||
}
|
||||
}
|
||||
if len(scopeRateLimits) > 0 {
|
||||
item.ScopeRateLimits = scopeRateLimits
|
||||
}
|
||||
if isOverloaded && acc.OverloadUntil != nil {
|
||||
item.OverloadUntil = acc.OverloadUntil
|
||||
remainingSec := int64(time.Until(*acc.OverloadUntil).Seconds())
|
||||
|
||||
@@ -50,24 +50,22 @@ type UserConcurrencyInfo struct {
|
||||
|
||||
// PlatformAvailability aggregates account availability by platform.
|
||||
type PlatformAvailability struct {
|
||||
Platform string `json:"platform"`
|
||||
TotalAccounts int64 `json:"total_accounts"`
|
||||
AvailableCount int64 `json:"available_count"`
|
||||
RateLimitCount int64 `json:"rate_limit_count"`
|
||||
ScopeRateLimitCount map[string]int64 `json:"scope_rate_limit_count,omitempty"`
|
||||
ErrorCount int64 `json:"error_count"`
|
||||
Platform string `json:"platform"`
|
||||
TotalAccounts int64 `json:"total_accounts"`
|
||||
AvailableCount int64 `json:"available_count"`
|
||||
RateLimitCount int64 `json:"rate_limit_count"`
|
||||
ErrorCount int64 `json:"error_count"`
|
||||
}
|
||||
|
||||
// GroupAvailability aggregates account availability by group.
|
||||
type GroupAvailability struct {
|
||||
GroupID int64 `json:"group_id"`
|
||||
GroupName string `json:"group_name"`
|
||||
Platform string `json:"platform"`
|
||||
TotalAccounts int64 `json:"total_accounts"`
|
||||
AvailableCount int64 `json:"available_count"`
|
||||
RateLimitCount int64 `json:"rate_limit_count"`
|
||||
ScopeRateLimitCount map[string]int64 `json:"scope_rate_limit_count,omitempty"`
|
||||
ErrorCount int64 `json:"error_count"`
|
||||
GroupID int64 `json:"group_id"`
|
||||
GroupName string `json:"group_name"`
|
||||
Platform string `json:"platform"`
|
||||
TotalAccounts int64 `json:"total_accounts"`
|
||||
AvailableCount int64 `json:"available_count"`
|
||||
RateLimitCount int64 `json:"rate_limit_count"`
|
||||
ErrorCount int64 `json:"error_count"`
|
||||
}
|
||||
|
||||
// AccountAvailability represents current availability for a single account.
|
||||
@@ -85,11 +83,10 @@ type AccountAvailability struct {
|
||||
IsOverloaded bool `json:"is_overloaded"`
|
||||
HasError bool `json:"has_error"`
|
||||
|
||||
RateLimitResetAt *time.Time `json:"rate_limit_reset_at"`
|
||||
RateLimitRemainingSec *int64 `json:"rate_limit_remaining_sec"`
|
||||
ScopeRateLimits map[string]int64 `json:"scope_rate_limits,omitempty"`
|
||||
OverloadUntil *time.Time `json:"overload_until"`
|
||||
OverloadRemainingSec *int64 `json:"overload_remaining_sec"`
|
||||
ErrorMessage string `json:"error_message"`
|
||||
TempUnschedulableUntil *time.Time `json:"temp_unschedulable_until,omitempty"`
|
||||
RateLimitResetAt *time.Time `json:"rate_limit_reset_at"`
|
||||
RateLimitRemainingSec *int64 `json:"rate_limit_remaining_sec"`
|
||||
OverloadUntil *time.Time `json:"overload_until"`
|
||||
OverloadRemainingSec *int64 `json:"overload_remaining_sec"`
|
||||
ErrorMessage string `json:"error_message"`
|
||||
TempUnschedulableUntil *time.Time `json:"temp_unschedulable_until,omitempty"`
|
||||
}
|
||||
|
||||
@@ -12,6 +12,7 @@ import (
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/domain"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
"github.com/gin-gonic/gin"
|
||||
@@ -528,7 +529,7 @@ func (s *OpsService) selectAccountForRetry(ctx context.Context, reqType opsRetry
|
||||
func extractRetryModelAndStream(reqType opsRetryRequestType, errorLog *OpsErrorLogDetail, body []byte) (model string, stream bool, err error) {
|
||||
switch reqType {
|
||||
case opsRetryTypeMessages:
|
||||
parsed, parseErr := ParseGatewayRequest(body)
|
||||
parsed, parseErr := ParseGatewayRequest(body, domain.PlatformAnthropic)
|
||||
if parseErr != nil {
|
||||
return "", false, fmt.Errorf("failed to parse messages request body: %w", parseErr)
|
||||
}
|
||||
@@ -596,7 +597,7 @@ func (s *OpsService) executeWithAccount(ctx context.Context, reqType opsRetryReq
|
||||
if s.gatewayService == nil {
|
||||
return &opsRetryExecution{status: opsRetryStatusFailed, errorMessage: "gateway service not available"}
|
||||
}
|
||||
parsedReq, parseErr := ParseGatewayRequest(body)
|
||||
parsedReq, parseErr := ParseGatewayRequest(body, domain.PlatformAnthropic)
|
||||
if parseErr != nil {
|
||||
return &opsRetryExecution{status: opsRetryStatusFailed, errorMessage: "failed to parse request body"}
|
||||
}
|
||||
|
||||
@@ -62,6 +62,32 @@ func (s *RateLimitService) SetTokenCacheInvalidator(invalidator TokenCacheInvali
|
||||
s.tokenCacheInvalidator = invalidator
|
||||
}
|
||||
|
||||
// ErrorPolicyResult 表示错误策略检查的结果
|
||||
type ErrorPolicyResult int
|
||||
|
||||
const (
|
||||
ErrorPolicyNone ErrorPolicyResult = iota // 未命中任何策略,继续默认逻辑
|
||||
ErrorPolicySkipped // 自定义错误码开启但未命中,跳过处理
|
||||
ErrorPolicyMatched // 自定义错误码命中,应停止调度
|
||||
ErrorPolicyTempUnscheduled // 临时不可调度规则命中
|
||||
)
|
||||
|
||||
// CheckErrorPolicy 检查自定义错误码和临时不可调度规则。
|
||||
// 自定义错误码开启时覆盖后续所有逻辑(包括临时不可调度)。
|
||||
func (s *RateLimitService) CheckErrorPolicy(ctx context.Context, account *Account, statusCode int, responseBody []byte) ErrorPolicyResult {
|
||||
if account.IsCustomErrorCodesEnabled() {
|
||||
if account.ShouldHandleErrorCode(statusCode) {
|
||||
return ErrorPolicyMatched
|
||||
}
|
||||
slog.Info("account_error_code_skipped", "account_id", account.ID, "status_code", statusCode)
|
||||
return ErrorPolicySkipped
|
||||
}
|
||||
if s.tryTempUnschedulable(ctx, account, statusCode, responseBody) {
|
||||
return ErrorPolicyTempUnscheduled
|
||||
}
|
||||
return ErrorPolicyNone
|
||||
}
|
||||
|
||||
// HandleUpstreamError 处理上游错误响应,标记账号状态
|
||||
// 返回是否应该停止该账号的调度
|
||||
func (s *RateLimitService) HandleUpstreamError(ctx context.Context, account *Account, statusCode int, headers http.Header, responseBody []byte) (shouldDisable bool) {
|
||||
|
||||
318
backend/internal/service/scheduler_shuffle_test.go
Normal file
318
backend/internal/service/scheduler_shuffle_test.go
Normal file
@@ -0,0 +1,318 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// ============ shuffleWithinSortGroups 测试 ============
|
||||
|
||||
func TestShuffleWithinSortGroups_Empty(t *testing.T) {
|
||||
shuffleWithinSortGroups(nil)
|
||||
shuffleWithinSortGroups([]accountWithLoad{})
|
||||
}
|
||||
|
||||
func TestShuffleWithinSortGroups_SingleElement(t *testing.T) {
|
||||
accounts := []accountWithLoad{
|
||||
{account: &Account{ID: 1, Priority: 1}, loadInfo: &AccountLoadInfo{LoadRate: 10}},
|
||||
}
|
||||
shuffleWithinSortGroups(accounts)
|
||||
require.Equal(t, int64(1), accounts[0].account.ID)
|
||||
}
|
||||
|
||||
func TestShuffleWithinSortGroups_DifferentGroups_OrderPreserved(t *testing.T) {
|
||||
now := time.Now()
|
||||
earlier := now.Add(-1 * time.Hour)
|
||||
|
||||
accounts := []accountWithLoad{
|
||||
{account: &Account{ID: 1, Priority: 1, LastUsedAt: &earlier}, loadInfo: &AccountLoadInfo{LoadRate: 10}},
|
||||
{account: &Account{ID: 2, Priority: 1, LastUsedAt: &now}, loadInfo: &AccountLoadInfo{LoadRate: 20}},
|
||||
{account: &Account{ID: 3, Priority: 2, LastUsedAt: &earlier}, loadInfo: &AccountLoadInfo{LoadRate: 10}},
|
||||
}
|
||||
|
||||
// 每个元素都属于不同组(Priority 或 LoadRate 或 LastUsedAt 不同),顺序不变
|
||||
for i := 0; i < 20; i++ {
|
||||
cpy := make([]accountWithLoad, len(accounts))
|
||||
copy(cpy, accounts)
|
||||
shuffleWithinSortGroups(cpy)
|
||||
require.Equal(t, int64(1), cpy[0].account.ID)
|
||||
require.Equal(t, int64(2), cpy[1].account.ID)
|
||||
require.Equal(t, int64(3), cpy[2].account.ID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestShuffleWithinSortGroups_SameGroup_Shuffled(t *testing.T) {
|
||||
now := time.Now()
|
||||
// 同一秒的时间戳视为同一组
|
||||
sameSecond := time.Unix(now.Unix(), 0)
|
||||
sameSecond2 := time.Unix(now.Unix(), 500_000_000) // 同一秒但不同纳秒
|
||||
|
||||
accounts := []accountWithLoad{
|
||||
{account: &Account{ID: 1, Priority: 1, LastUsedAt: &sameSecond}, loadInfo: &AccountLoadInfo{LoadRate: 10}},
|
||||
{account: &Account{ID: 2, Priority: 1, LastUsedAt: &sameSecond2}, loadInfo: &AccountLoadInfo{LoadRate: 10}},
|
||||
{account: &Account{ID: 3, Priority: 1, LastUsedAt: &sameSecond}, loadInfo: &AccountLoadInfo{LoadRate: 10}},
|
||||
}
|
||||
|
||||
// 多次执行,验证所有 ID 都出现在第一个位置(说明确实被打乱了)
|
||||
seen := map[int64]bool{}
|
||||
for i := 0; i < 100; i++ {
|
||||
cpy := make([]accountWithLoad, len(accounts))
|
||||
copy(cpy, accounts)
|
||||
shuffleWithinSortGroups(cpy)
|
||||
seen[cpy[0].account.ID] = true
|
||||
// 无论怎么打乱,所有 ID 都应在候选中
|
||||
ids := map[int64]bool{}
|
||||
for _, a := range cpy {
|
||||
ids[a.account.ID] = true
|
||||
}
|
||||
require.True(t, ids[1] && ids[2] && ids[3])
|
||||
}
|
||||
// 至少 2 个不同的 ID 出现在首位(随机性验证)
|
||||
require.GreaterOrEqual(t, len(seen), 2, "shuffle should produce different orderings")
|
||||
}
|
||||
|
||||
func TestShuffleWithinSortGroups_NilLastUsedAt_SameGroup(t *testing.T) {
|
||||
accounts := []accountWithLoad{
|
||||
{account: &Account{ID: 1, Priority: 1, LastUsedAt: nil}, loadInfo: &AccountLoadInfo{LoadRate: 0}},
|
||||
{account: &Account{ID: 2, Priority: 1, LastUsedAt: nil}, loadInfo: &AccountLoadInfo{LoadRate: 0}},
|
||||
{account: &Account{ID: 3, Priority: 1, LastUsedAt: nil}, loadInfo: &AccountLoadInfo{LoadRate: 0}},
|
||||
}
|
||||
|
||||
seen := map[int64]bool{}
|
||||
for i := 0; i < 100; i++ {
|
||||
cpy := make([]accountWithLoad, len(accounts))
|
||||
copy(cpy, accounts)
|
||||
shuffleWithinSortGroups(cpy)
|
||||
seen[cpy[0].account.ID] = true
|
||||
}
|
||||
require.GreaterOrEqual(t, len(seen), 2, "nil LastUsedAt accounts should be shuffled")
|
||||
}
|
||||
|
||||
func TestShuffleWithinSortGroups_MixedGroups(t *testing.T) {
|
||||
now := time.Now()
|
||||
earlier := now.Add(-1 * time.Hour)
|
||||
sameAsNow := time.Unix(now.Unix(), 0)
|
||||
|
||||
// 组1: Priority=1, LoadRate=10, LastUsedAt=earlier (ID 1) — 单元素组
|
||||
// 组2: Priority=1, LoadRate=20, LastUsedAt=now (ID 2, 3) — 双元素组
|
||||
// 组3: Priority=2, LoadRate=10, LastUsedAt=earlier (ID 4) — 单元素组
|
||||
accounts := []accountWithLoad{
|
||||
{account: &Account{ID: 1, Priority: 1, LastUsedAt: &earlier}, loadInfo: &AccountLoadInfo{LoadRate: 10}},
|
||||
{account: &Account{ID: 2, Priority: 1, LastUsedAt: &now}, loadInfo: &AccountLoadInfo{LoadRate: 20}},
|
||||
{account: &Account{ID: 3, Priority: 1, LastUsedAt: &sameAsNow}, loadInfo: &AccountLoadInfo{LoadRate: 20}},
|
||||
{account: &Account{ID: 4, Priority: 2, LastUsedAt: &earlier}, loadInfo: &AccountLoadInfo{LoadRate: 10}},
|
||||
}
|
||||
|
||||
for i := 0; i < 20; i++ {
|
||||
cpy := make([]accountWithLoad, len(accounts))
|
||||
copy(cpy, accounts)
|
||||
shuffleWithinSortGroups(cpy)
|
||||
|
||||
// 组间顺序不变
|
||||
require.Equal(t, int64(1), cpy[0].account.ID, "group 1 position fixed")
|
||||
require.Equal(t, int64(4), cpy[3].account.ID, "group 3 position fixed")
|
||||
|
||||
// 组2 内部可以打乱,但仍在位置 1 和 2
|
||||
mid := map[int64]bool{cpy[1].account.ID: true, cpy[2].account.ID: true}
|
||||
require.True(t, mid[2] && mid[3], "group 2 elements should stay in positions 1-2")
|
||||
}
|
||||
}
|
||||
|
||||
// ============ shuffleWithinPriorityAndLastUsed 测试 ============
|
||||
|
||||
func TestShuffleWithinPriorityAndLastUsed_Empty(t *testing.T) {
|
||||
shuffleWithinPriorityAndLastUsed(nil)
|
||||
shuffleWithinPriorityAndLastUsed([]*Account{})
|
||||
}
|
||||
|
||||
func TestShuffleWithinPriorityAndLastUsed_SingleElement(t *testing.T) {
|
||||
accounts := []*Account{{ID: 1, Priority: 1}}
|
||||
shuffleWithinPriorityAndLastUsed(accounts)
|
||||
require.Equal(t, int64(1), accounts[0].ID)
|
||||
}
|
||||
|
||||
func TestShuffleWithinPriorityAndLastUsed_SameGroup_Shuffled(t *testing.T) {
|
||||
accounts := []*Account{
|
||||
{ID: 1, Priority: 1, LastUsedAt: nil},
|
||||
{ID: 2, Priority: 1, LastUsedAt: nil},
|
||||
{ID: 3, Priority: 1, LastUsedAt: nil},
|
||||
}
|
||||
|
||||
seen := map[int64]bool{}
|
||||
for i := 0; i < 100; i++ {
|
||||
cpy := make([]*Account, len(accounts))
|
||||
copy(cpy, accounts)
|
||||
shuffleWithinPriorityAndLastUsed(cpy)
|
||||
seen[cpy[0].ID] = true
|
||||
}
|
||||
require.GreaterOrEqual(t, len(seen), 2, "same group should be shuffled")
|
||||
}
|
||||
|
||||
func TestShuffleWithinPriorityAndLastUsed_DifferentPriority_OrderPreserved(t *testing.T) {
|
||||
accounts := []*Account{
|
||||
{ID: 1, Priority: 1, LastUsedAt: nil},
|
||||
{ID: 2, Priority: 2, LastUsedAt: nil},
|
||||
{ID: 3, Priority: 3, LastUsedAt: nil},
|
||||
}
|
||||
|
||||
for i := 0; i < 20; i++ {
|
||||
cpy := make([]*Account, len(accounts))
|
||||
copy(cpy, accounts)
|
||||
shuffleWithinPriorityAndLastUsed(cpy)
|
||||
require.Equal(t, int64(1), cpy[0].ID)
|
||||
require.Equal(t, int64(2), cpy[1].ID)
|
||||
require.Equal(t, int64(3), cpy[2].ID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestShuffleWithinPriorityAndLastUsed_DifferentLastUsedAt_OrderPreserved(t *testing.T) {
|
||||
now := time.Now()
|
||||
earlier := now.Add(-1 * time.Hour)
|
||||
|
||||
accounts := []*Account{
|
||||
{ID: 1, Priority: 1, LastUsedAt: nil},
|
||||
{ID: 2, Priority: 1, LastUsedAt: &earlier},
|
||||
{ID: 3, Priority: 1, LastUsedAt: &now},
|
||||
}
|
||||
|
||||
for i := 0; i < 20; i++ {
|
||||
cpy := make([]*Account, len(accounts))
|
||||
copy(cpy, accounts)
|
||||
shuffleWithinPriorityAndLastUsed(cpy)
|
||||
require.Equal(t, int64(1), cpy[0].ID)
|
||||
require.Equal(t, int64(2), cpy[1].ID)
|
||||
require.Equal(t, int64(3), cpy[2].ID)
|
||||
}
|
||||
}
|
||||
|
||||
// ============ sameLastUsedAt 测试 ============
|
||||
|
||||
func TestSameLastUsedAt(t *testing.T) {
|
||||
now := time.Now()
|
||||
sameSecond := time.Unix(now.Unix(), 0)
|
||||
sameSecondDiffNano := time.Unix(now.Unix(), 999_999_999)
|
||||
differentSecond := now.Add(1 * time.Second)
|
||||
|
||||
t.Run("both nil", func(t *testing.T) {
|
||||
require.True(t, sameLastUsedAt(nil, nil))
|
||||
})
|
||||
|
||||
t.Run("one nil one not", func(t *testing.T) {
|
||||
require.False(t, sameLastUsedAt(nil, &now))
|
||||
require.False(t, sameLastUsedAt(&now, nil))
|
||||
})
|
||||
|
||||
t.Run("same second different nanoseconds", func(t *testing.T) {
|
||||
require.True(t, sameLastUsedAt(&sameSecond, &sameSecondDiffNano))
|
||||
})
|
||||
|
||||
t.Run("different seconds", func(t *testing.T) {
|
||||
require.False(t, sameLastUsedAt(&now, &differentSecond))
|
||||
})
|
||||
|
||||
t.Run("exact same time", func(t *testing.T) {
|
||||
require.True(t, sameLastUsedAt(&now, &now))
|
||||
})
|
||||
}
|
||||
|
||||
// ============ sameAccountWithLoadGroup 测试 ============
|
||||
|
||||
func TestSameAccountWithLoadGroup(t *testing.T) {
|
||||
now := time.Now()
|
||||
sameSecond := time.Unix(now.Unix(), 0)
|
||||
|
||||
t.Run("same group", func(t *testing.T) {
|
||||
a := accountWithLoad{account: &Account{Priority: 1, LastUsedAt: &now}, loadInfo: &AccountLoadInfo{LoadRate: 10}}
|
||||
b := accountWithLoad{account: &Account{Priority: 1, LastUsedAt: &sameSecond}, loadInfo: &AccountLoadInfo{LoadRate: 10}}
|
||||
require.True(t, sameAccountWithLoadGroup(a, b))
|
||||
})
|
||||
|
||||
t.Run("different priority", func(t *testing.T) {
|
||||
a := accountWithLoad{account: &Account{Priority: 1, LastUsedAt: &now}, loadInfo: &AccountLoadInfo{LoadRate: 10}}
|
||||
b := accountWithLoad{account: &Account{Priority: 2, LastUsedAt: &now}, loadInfo: &AccountLoadInfo{LoadRate: 10}}
|
||||
require.False(t, sameAccountWithLoadGroup(a, b))
|
||||
})
|
||||
|
||||
t.Run("different load rate", func(t *testing.T) {
|
||||
a := accountWithLoad{account: &Account{Priority: 1, LastUsedAt: &now}, loadInfo: &AccountLoadInfo{LoadRate: 10}}
|
||||
b := accountWithLoad{account: &Account{Priority: 1, LastUsedAt: &now}, loadInfo: &AccountLoadInfo{LoadRate: 20}}
|
||||
require.False(t, sameAccountWithLoadGroup(a, b))
|
||||
})
|
||||
|
||||
t.Run("different last used at", func(t *testing.T) {
|
||||
later := now.Add(1 * time.Second)
|
||||
a := accountWithLoad{account: &Account{Priority: 1, LastUsedAt: &now}, loadInfo: &AccountLoadInfo{LoadRate: 10}}
|
||||
b := accountWithLoad{account: &Account{Priority: 1, LastUsedAt: &later}, loadInfo: &AccountLoadInfo{LoadRate: 10}}
|
||||
require.False(t, sameAccountWithLoadGroup(a, b))
|
||||
})
|
||||
|
||||
t.Run("both nil LastUsedAt", func(t *testing.T) {
|
||||
a := accountWithLoad{account: &Account{Priority: 1, LastUsedAt: nil}, loadInfo: &AccountLoadInfo{LoadRate: 0}}
|
||||
b := accountWithLoad{account: &Account{Priority: 1, LastUsedAt: nil}, loadInfo: &AccountLoadInfo{LoadRate: 0}}
|
||||
require.True(t, sameAccountWithLoadGroup(a, b))
|
||||
})
|
||||
}
|
||||
|
||||
// ============ sameAccountGroup 测试 ============
|
||||
|
||||
func TestSameAccountGroup(t *testing.T) {
|
||||
now := time.Now()
|
||||
|
||||
t.Run("same group", func(t *testing.T) {
|
||||
a := &Account{Priority: 1, LastUsedAt: nil}
|
||||
b := &Account{Priority: 1, LastUsedAt: nil}
|
||||
require.True(t, sameAccountGroup(a, b))
|
||||
})
|
||||
|
||||
t.Run("different priority", func(t *testing.T) {
|
||||
a := &Account{Priority: 1, LastUsedAt: nil}
|
||||
b := &Account{Priority: 2, LastUsedAt: nil}
|
||||
require.False(t, sameAccountGroup(a, b))
|
||||
})
|
||||
|
||||
t.Run("different LastUsedAt", func(t *testing.T) {
|
||||
later := now.Add(1 * time.Second)
|
||||
a := &Account{Priority: 1, LastUsedAt: &now}
|
||||
b := &Account{Priority: 1, LastUsedAt: &later}
|
||||
require.False(t, sameAccountGroup(a, b))
|
||||
})
|
||||
}
|
||||
|
||||
// ============ sortAccountsByPriorityAndLastUsed 集成随机化测试 ============
|
||||
|
||||
func TestSortAccountsByPriorityAndLastUsed_WithShuffle(t *testing.T) {
|
||||
t.Run("same priority and nil LastUsedAt are shuffled", func(t *testing.T) {
|
||||
accounts := []*Account{
|
||||
{ID: 1, Priority: 1, LastUsedAt: nil},
|
||||
{ID: 2, Priority: 1, LastUsedAt: nil},
|
||||
{ID: 3, Priority: 1, LastUsedAt: nil},
|
||||
}
|
||||
|
||||
seen := map[int64]bool{}
|
||||
for i := 0; i < 100; i++ {
|
||||
cpy := make([]*Account, len(accounts))
|
||||
copy(cpy, accounts)
|
||||
sortAccountsByPriorityAndLastUsed(cpy, false)
|
||||
seen[cpy[0].ID] = true
|
||||
}
|
||||
require.GreaterOrEqual(t, len(seen), 2, "identical sort keys should produce different orderings after shuffle")
|
||||
})
|
||||
|
||||
t.Run("different priorities still sorted correctly", func(t *testing.T) {
|
||||
now := time.Now()
|
||||
accounts := []*Account{
|
||||
{ID: 3, Priority: 3, LastUsedAt: &now},
|
||||
{ID: 1, Priority: 1, LastUsedAt: &now},
|
||||
{ID: 2, Priority: 2, LastUsedAt: &now},
|
||||
}
|
||||
|
||||
sortAccountsByPriorityAndLastUsed(accounts, false)
|
||||
require.Equal(t, int64(1), accounts[0].ID)
|
||||
require.Equal(t, int64(2), accounts[1].ID)
|
||||
require.Equal(t, int64(3), accounts[2].ID)
|
||||
})
|
||||
}
|
||||
@@ -23,8 +23,7 @@ import (
|
||||
// - 临时不可调度且未过期:清理
|
||||
// - 临时不可调度已过期:不清理
|
||||
// - 正常可调度状态:不清理
|
||||
// - 模型限流超过阈值:清理
|
||||
// - 模型限流未超过阈值:不清理
|
||||
// - 模型限流(任意时长):清理
|
||||
//
|
||||
// TestShouldClearStickySession tests the sticky session clearing logic.
|
||||
// Verifies correct behavior for various account states including:
|
||||
@@ -35,9 +34,9 @@ func TestShouldClearStickySession(t *testing.T) {
|
||||
future := now.Add(1 * time.Hour)
|
||||
past := now.Add(-1 * time.Hour)
|
||||
|
||||
// 短限流时间(低于阈值,不应清除粘性会话)
|
||||
// 短限流时间(有限流即清除粘性会话)
|
||||
shortRateLimitReset := now.Add(5 * time.Second).Format(time.RFC3339)
|
||||
// 长限流时间(超过阈值,应清除粘性会话)
|
||||
// 长限流时间(有限流即清除粘性会话)
|
||||
longRateLimitReset := now.Add(30 * time.Second).Format(time.RFC3339)
|
||||
|
||||
tests := []struct {
|
||||
@@ -53,7 +52,7 @@ func TestShouldClearStickySession(t *testing.T) {
|
||||
{name: "temp unschedulable", account: &Account{Status: StatusActive, Schedulable: true, TempUnschedulableUntil: &future}, requestedModel: "", want: true},
|
||||
{name: "temp unschedulable expired", account: &Account{Status: StatusActive, Schedulable: true, TempUnschedulableUntil: &past}, requestedModel: "", want: false},
|
||||
{name: "active schedulable", account: &Account{Status: StatusActive, Schedulable: true}, requestedModel: "", want: false},
|
||||
// 模型限流测试
|
||||
// 模型限流测试:有限流即清除
|
||||
{
|
||||
name: "model rate limited short duration",
|
||||
account: &Account{
|
||||
@@ -68,7 +67,7 @@ func TestShouldClearStickySession(t *testing.T) {
|
||||
},
|
||||
},
|
||||
requestedModel: "claude-sonnet-4",
|
||||
want: false, // 低于阈值,不清除
|
||||
want: true, // 有限流即清除
|
||||
},
|
||||
{
|
||||
name: "model rate limited long duration",
|
||||
@@ -84,7 +83,7 @@ func TestShouldClearStickySession(t *testing.T) {
|
||||
},
|
||||
},
|
||||
requestedModel: "claude-sonnet-4",
|
||||
want: true, // 超过阈值,清除
|
||||
want: true, // 有限流即清除
|
||||
},
|
||||
{
|
||||
name: "model rate limited different model",
|
||||
|
||||
@@ -295,4 +295,5 @@ var ProviderSet = wire.NewSet(
|
||||
NewUsageCache,
|
||||
NewTotpService,
|
||||
NewErrorPassthroughService,
|
||||
NewDigestSessionStore,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user