Merge branch 'Wei-Shaw:main' into main
This commit is contained in:
@@ -0,0 +1,30 @@
|
||||
package service
|
||||
|
||||
import "context"
|
||||
|
||||
type accountCredentialsUpdater interface {
|
||||
UpdateCredentials(ctx context.Context, id int64, credentials map[string]any) error
|
||||
}
|
||||
|
||||
func persistAccountCredentials(ctx context.Context, repo AccountRepository, account *Account, credentials map[string]any) error {
|
||||
if repo == nil || account == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
account.Credentials = cloneCredentials(credentials)
|
||||
if updater, ok := any(repo).(accountCredentialsUpdater); ok {
|
||||
return updater.UpdateCredentials(ctx, account.ID, account.Credentials)
|
||||
}
|
||||
return repo.Update(ctx, account)
|
||||
}
|
||||
|
||||
func cloneCredentials(in map[string]any) map[string]any {
|
||||
if in == nil {
|
||||
return map[string]any{}
|
||||
}
|
||||
out := make(map[string]any, len(in))
|
||||
for k, v := range in {
|
||||
out[k] = v
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -15,6 +15,7 @@ var (
|
||||
)
|
||||
|
||||
const AccountListGroupUngrouped int64 = -1
|
||||
const AccountPrivacyModeUnsetFilter = "__unset__"
|
||||
|
||||
type AccountRepository interface {
|
||||
Create(ctx context.Context, account *Account) error
|
||||
@@ -37,7 +38,7 @@ type AccountRepository interface {
|
||||
Delete(ctx context.Context, id int64) error
|
||||
|
||||
List(ctx context.Context, params pagination.PaginationParams) ([]Account, *pagination.PaginationResult, error)
|
||||
ListWithFilters(ctx context.Context, params pagination.PaginationParams, platform, accountType, status, search string, groupID int64) ([]Account, *pagination.PaginationResult, error)
|
||||
ListWithFilters(ctx context.Context, params pagination.PaginationParams, platform, accountType, status, search string, groupID int64, privacyMode string) ([]Account, *pagination.PaginationResult, error)
|
||||
ListByGroup(ctx context.Context, groupID int64) ([]Account, error)
|
||||
ListActive(ctx context.Context) ([]Account, error)
|
||||
ListByPlatform(ctx context.Context, platform string) ([]Account, error)
|
||||
|
||||
@@ -79,7 +79,7 @@ func (s *accountRepoStub) List(ctx context.Context, params pagination.Pagination
|
||||
panic("unexpected List call")
|
||||
}
|
||||
|
||||
func (s *accountRepoStub) ListWithFilters(ctx context.Context, params pagination.PaginationParams, platform, accountType, status, search string, groupID int64) ([]Account, *pagination.PaginationResult, error) {
|
||||
func (s *accountRepoStub) ListWithFilters(ctx context.Context, params pagination.PaginationParams, platform, accountType, status, search string, groupID int64, privacyMode string) ([]Account, *pagination.PaginationResult, error) {
|
||||
panic("unexpected ListWithFilters call")
|
||||
}
|
||||
|
||||
|
||||
@@ -54,7 +54,7 @@ type AdminService interface {
|
||||
ReplaceUserGroup(ctx context.Context, userID, oldGroupID, newGroupID int64) (*ReplaceUserGroupResult, error)
|
||||
|
||||
// Account management
|
||||
ListAccounts(ctx context.Context, page, pageSize int, platform, accountType, status, search string, groupID int64) ([]Account, int64, error)
|
||||
ListAccounts(ctx context.Context, page, pageSize int, platform, accountType, status, search string, groupID int64, privacyMode string) ([]Account, int64, error)
|
||||
GetAccount(ctx context.Context, id int64) (*Account, error)
|
||||
GetAccountsByIDs(ctx context.Context, ids []int64) ([]*Account, error)
|
||||
CreateAccount(ctx context.Context, input *CreateAccountInput) (*Account, error)
|
||||
@@ -1451,9 +1451,9 @@ func (s *adminServiceImpl) ReplaceUserGroup(ctx context.Context, userID, oldGrou
|
||||
}
|
||||
|
||||
// Account management implementations
|
||||
func (s *adminServiceImpl) ListAccounts(ctx context.Context, page, pageSize int, platform, accountType, status, search string, groupID int64) ([]Account, int64, error) {
|
||||
func (s *adminServiceImpl) ListAccounts(ctx context.Context, page, pageSize int, platform, accountType, status, search string, groupID int64, privacyMode string) ([]Account, int64, error) {
|
||||
params := pagination.PaginationParams{Page: page, PageSize: pageSize}
|
||||
accounts, result, err := s.accountRepo.ListWithFilters(ctx, params, platform, accountType, status, search, groupID)
|
||||
accounts, result, err := s.accountRepo.ListWithFilters(ctx, params, platform, accountType, status, search, groupID, privacyMode)
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
|
||||
@@ -19,18 +19,20 @@ type accountRepoStubForAdminList struct {
|
||||
listWithFiltersType string
|
||||
listWithFiltersStatus string
|
||||
listWithFiltersSearch string
|
||||
listWithFiltersPrivacy string
|
||||
listWithFiltersAccounts []Account
|
||||
listWithFiltersResult *pagination.PaginationResult
|
||||
listWithFiltersErr error
|
||||
}
|
||||
|
||||
func (s *accountRepoStubForAdminList) ListWithFilters(_ context.Context, params pagination.PaginationParams, platform, accountType, status, search string, groupID int64) ([]Account, *pagination.PaginationResult, error) {
|
||||
func (s *accountRepoStubForAdminList) ListWithFilters(_ context.Context, params pagination.PaginationParams, platform, accountType, status, search string, groupID int64, privacyMode string) ([]Account, *pagination.PaginationResult, error) {
|
||||
s.listWithFiltersCalls++
|
||||
s.listWithFiltersParams = params
|
||||
s.listWithFiltersPlatform = platform
|
||||
s.listWithFiltersType = accountType
|
||||
s.listWithFiltersStatus = status
|
||||
s.listWithFiltersSearch = search
|
||||
s.listWithFiltersPrivacy = privacyMode
|
||||
|
||||
if s.listWithFiltersErr != nil {
|
||||
return nil, nil, s.listWithFiltersErr
|
||||
@@ -168,7 +170,7 @@ func TestAdminService_ListAccounts_WithSearch(t *testing.T) {
|
||||
}
|
||||
svc := &adminServiceImpl{accountRepo: repo}
|
||||
|
||||
accounts, total, err := svc.ListAccounts(context.Background(), 1, 20, PlatformGemini, AccountTypeOAuth, StatusActive, "acc", 0)
|
||||
accounts, total, err := svc.ListAccounts(context.Background(), 1, 20, PlatformGemini, AccountTypeOAuth, StatusActive, "acc", 0, "")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(10), total)
|
||||
require.Equal(t, []Account{{ID: 1, Name: "acc"}}, accounts)
|
||||
@@ -182,6 +184,22 @@ func TestAdminService_ListAccounts_WithSearch(t *testing.T) {
|
||||
})
|
||||
}
|
||||
|
||||
func TestAdminService_ListAccounts_WithPrivacyMode(t *testing.T) {
|
||||
t.Run("privacy_mode 参数正常传递到 repository 层", func(t *testing.T) {
|
||||
repo := &accountRepoStubForAdminList{
|
||||
listWithFiltersAccounts: []Account{{ID: 2, Name: "acc2"}},
|
||||
listWithFiltersResult: &pagination.PaginationResult{Total: 1},
|
||||
}
|
||||
svc := &adminServiceImpl{accountRepo: repo}
|
||||
|
||||
accounts, total, err := svc.ListAccounts(context.Background(), 1, 20, PlatformOpenAI, AccountTypeOAuth, StatusActive, "acc2", 0, PrivacyModeCFBlocked)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(1), total)
|
||||
require.Equal(t, []Account{{ID: 2, Name: "acc2"}}, accounts)
|
||||
require.Equal(t, PrivacyModeCFBlocked, repo.listWithFiltersPrivacy)
|
||||
})
|
||||
}
|
||||
|
||||
func TestAdminService_ListProxies_WithSearch(t *testing.T) {
|
||||
t.Run("search 参数正常传递到 repository 层", func(t *testing.T) {
|
||||
repo := &proxyRepoStubForAdminList{
|
||||
|
||||
@@ -643,6 +643,7 @@ urlFallbackLoop:
|
||||
AccountID: p.account.ID,
|
||||
AccountName: p.account.Name,
|
||||
UpstreamStatusCode: 0,
|
||||
UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()),
|
||||
Kind: "request_error",
|
||||
Message: safeErr,
|
||||
})
|
||||
@@ -720,6 +721,7 @@ urlFallbackLoop:
|
||||
AccountName: p.account.Name,
|
||||
UpstreamStatusCode: resp.StatusCode,
|
||||
UpstreamRequestID: resp.Header.Get("x-request-id"),
|
||||
UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()),
|
||||
Kind: "retry",
|
||||
Message: upstreamMsg,
|
||||
Detail: getUpstreamDetail(respBody),
|
||||
@@ -754,6 +756,7 @@ urlFallbackLoop:
|
||||
AccountName: p.account.Name,
|
||||
UpstreamStatusCode: resp.StatusCode,
|
||||
UpstreamRequestID: resp.Header.Get("x-request-id"),
|
||||
UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()),
|
||||
Kind: "retry",
|
||||
Message: upstreamMsg,
|
||||
Detail: getUpstreamDetail(respBody),
|
||||
|
||||
@@ -138,7 +138,7 @@ func (p *AntigravityTokenProvider) GetAccessToken(ctx context.Context, account *
|
||||
p.markBackfillAttempted(account.ID)
|
||||
if projectID, err := p.antigravityOAuthService.FillProjectID(ctx, account, accessToken); err == nil && projectID != "" {
|
||||
account.Credentials["project_id"] = projectID
|
||||
if updateErr := p.accountRepo.Update(ctx, account); updateErr != nil {
|
||||
if updateErr := persistAccountCredentials(ctx, p.accountRepo, account, account.Credentials); updateErr != nil {
|
||||
slog.Warn("antigravity_project_id_backfill_persist_failed",
|
||||
"account_id", account.ID,
|
||||
"error", updateErr,
|
||||
|
||||
@@ -367,8 +367,7 @@ func (s *CRSSyncService) SyncFromCRS(ctx context.Context, input SyncFromCRSInput
|
||||
// 🔄 Refresh OAuth token after creation
|
||||
if targetType == AccountTypeOAuth {
|
||||
if refreshedCreds := s.refreshOAuthToken(ctx, account); refreshedCreds != nil {
|
||||
account.Credentials = refreshedCreds
|
||||
_ = s.accountRepo.Update(ctx, account)
|
||||
_ = persistAccountCredentials(ctx, s.accountRepo, account, refreshedCreds)
|
||||
}
|
||||
}
|
||||
item.Action = "created"
|
||||
@@ -402,8 +401,7 @@ func (s *CRSSyncService) SyncFromCRS(ctx context.Context, input SyncFromCRSInput
|
||||
// 🔄 Refresh OAuth token after update
|
||||
if targetType == AccountTypeOAuth {
|
||||
if refreshedCreds := s.refreshOAuthToken(ctx, existing); refreshedCreds != nil {
|
||||
existing.Credentials = refreshedCreds
|
||||
_ = s.accountRepo.Update(ctx, existing)
|
||||
_ = persistAccountCredentials(ctx, s.accountRepo, existing, refreshedCreds)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -620,8 +618,7 @@ func (s *CRSSyncService) SyncFromCRS(ctx context.Context, input SyncFromCRSInput
|
||||
}
|
||||
// 🔄 Refresh OAuth token after creation
|
||||
if refreshedCreds := s.refreshOAuthToken(ctx, account); refreshedCreds != nil {
|
||||
account.Credentials = refreshedCreds
|
||||
_ = s.accountRepo.Update(ctx, account)
|
||||
_ = persistAccountCredentials(ctx, s.accountRepo, account, refreshedCreds)
|
||||
}
|
||||
item.Action = "created"
|
||||
result.Created++
|
||||
@@ -652,8 +649,7 @@ func (s *CRSSyncService) SyncFromCRS(ctx context.Context, input SyncFromCRSInput
|
||||
|
||||
// 🔄 Refresh OAuth token after update
|
||||
if refreshedCreds := s.refreshOAuthToken(ctx, existing); refreshedCreds != nil {
|
||||
existing.Credentials = refreshedCreds
|
||||
_ = s.accountRepo.Update(ctx, existing)
|
||||
_ = persistAccountCredentials(ctx, s.accountRepo, existing, refreshedCreds)
|
||||
}
|
||||
|
||||
item.Action = "updated"
|
||||
@@ -862,8 +858,7 @@ func (s *CRSSyncService) SyncFromCRS(ctx context.Context, input SyncFromCRSInput
|
||||
continue
|
||||
}
|
||||
if refreshedCreds := s.refreshOAuthToken(ctx, account); refreshedCreds != nil {
|
||||
account.Credentials = refreshedCreds
|
||||
_ = s.accountRepo.Update(ctx, account)
|
||||
_ = persistAccountCredentials(ctx, s.accountRepo, account, refreshedCreds)
|
||||
}
|
||||
item.Action = "created"
|
||||
result.Created++
|
||||
@@ -893,8 +888,7 @@ func (s *CRSSyncService) SyncFromCRS(ctx context.Context, input SyncFromCRSInput
|
||||
}
|
||||
|
||||
if refreshedCreds := s.refreshOAuthToken(ctx, existing); refreshedCreds != nil {
|
||||
existing.Credentials = refreshedCreds
|
||||
_ = s.accountRepo.Update(ctx, existing)
|
||||
_ = persistAccountCredentials(ctx, s.accountRepo, existing, refreshedCreds)
|
||||
}
|
||||
|
||||
item.Action = "updated"
|
||||
|
||||
@@ -12,6 +12,7 @@ import (
|
||||
"net/smtp"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
@@ -111,7 +112,7 @@ func (s *EmailService) GetSMTPConfig(ctx context.Context) (*SMTPConfig, error) {
|
||||
return nil, fmt.Errorf("get smtp settings: %w", err)
|
||||
}
|
||||
|
||||
host := settings[SettingKeySMTPHost]
|
||||
host := strings.TrimSpace(settings[SettingKeySMTPHost])
|
||||
if host == "" {
|
||||
return nil, ErrEmailNotConfigured
|
||||
}
|
||||
@@ -128,10 +129,10 @@ func (s *EmailService) GetSMTPConfig(ctx context.Context) (*SMTPConfig, error) {
|
||||
return &SMTPConfig{
|
||||
Host: host,
|
||||
Port: port,
|
||||
Username: settings[SettingKeySMTPUsername],
|
||||
Password: settings[SettingKeySMTPPassword],
|
||||
From: settings[SettingKeySMTPFrom],
|
||||
FromName: settings[SettingKeySMTPFromName],
|
||||
Username: strings.TrimSpace(settings[SettingKeySMTPUsername]),
|
||||
Password: strings.TrimSpace(settings[SettingKeySMTPPassword]),
|
||||
From: strings.TrimSpace(settings[SettingKeySMTPFrom]),
|
||||
FromName: strings.TrimSpace(settings[SettingKeySMTPFromName]),
|
||||
UseTLS: useTLS,
|
||||
}, nil
|
||||
}
|
||||
|
||||
@@ -0,0 +1,485 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/apicompat"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/claude"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||||
"github.com/Wei-Shaw/sub2api/internal/util/responseheaders"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/tidwall/gjson"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// ForwardAsChatCompletions accepts an OpenAI Chat Completions API request body,
|
||||
// converts it to Anthropic Messages format (chained via Responses format),
|
||||
// forwards to the Anthropic upstream, and converts the response back to Chat
|
||||
// Completions format. This enables Chat Completions clients to access Anthropic
|
||||
// models through Anthropic platform groups.
|
||||
func (s *GatewayService) ForwardAsChatCompletions(
|
||||
ctx context.Context,
|
||||
c *gin.Context,
|
||||
account *Account,
|
||||
body []byte,
|
||||
parsed *ParsedRequest,
|
||||
) (*ForwardResult, error) {
|
||||
startTime := time.Now()
|
||||
|
||||
// 1. Parse Chat Completions request
|
||||
var ccReq apicompat.ChatCompletionsRequest
|
||||
if err := json.Unmarshal(body, &ccReq); err != nil {
|
||||
return nil, fmt.Errorf("parse chat completions request: %w", err)
|
||||
}
|
||||
originalModel := ccReq.Model
|
||||
clientStream := ccReq.Stream
|
||||
includeUsage := ccReq.StreamOptions != nil && ccReq.StreamOptions.IncludeUsage
|
||||
|
||||
// 2. Convert CC → Responses → Anthropic (chained conversion)
|
||||
responsesReq, err := apicompat.ChatCompletionsToResponses(&ccReq)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("convert chat completions to responses: %w", err)
|
||||
}
|
||||
|
||||
anthropicReq, err := apicompat.ResponsesToAnthropicRequest(responsesReq)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("convert responses to anthropic: %w", err)
|
||||
}
|
||||
|
||||
// 3. Force upstream streaming
|
||||
anthropicReq.Stream = true
|
||||
reqStream := true
|
||||
|
||||
// 4. Model mapping
|
||||
mappedModel := originalModel
|
||||
if account.Type == AccountTypeAPIKey {
|
||||
mappedModel = account.GetMappedModel(originalModel)
|
||||
}
|
||||
if mappedModel == originalModel && account.Platform == PlatformAnthropic && account.Type != AccountTypeAPIKey {
|
||||
normalized := claude.NormalizeModelID(originalModel)
|
||||
if normalized != originalModel {
|
||||
mappedModel = normalized
|
||||
}
|
||||
}
|
||||
anthropicReq.Model = mappedModel
|
||||
|
||||
logger.L().Debug("gateway forward_as_chat_completions: model mapping applied",
|
||||
zap.Int64("account_id", account.ID),
|
||||
zap.String("original_model", originalModel),
|
||||
zap.String("mapped_model", mappedModel),
|
||||
zap.Bool("client_stream", clientStream),
|
||||
)
|
||||
|
||||
// 5. Marshal Anthropic request body
|
||||
anthropicBody, err := json.Marshal(anthropicReq)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("marshal anthropic request: %w", err)
|
||||
}
|
||||
|
||||
// 6. Apply Claude Code mimicry for OAuth accounts
|
||||
isClaudeCode := false // CC API is never Claude Code
|
||||
shouldMimicClaudeCode := account.IsOAuth() && !isClaudeCode
|
||||
|
||||
if shouldMimicClaudeCode {
|
||||
if !strings.Contains(strings.ToLower(mappedModel), "haiku") &&
|
||||
!systemIncludesClaudeCodePrompt(anthropicReq.System) {
|
||||
anthropicBody = injectClaudeCodePrompt(anthropicBody, anthropicReq.System)
|
||||
}
|
||||
}
|
||||
|
||||
// 7. Enforce cache_control block limit
|
||||
anthropicBody = enforceCacheControlLimit(anthropicBody)
|
||||
|
||||
// 8. Get access token
|
||||
token, tokenType, err := s.GetAccessToken(ctx, account)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get access token: %w", err)
|
||||
}
|
||||
|
||||
// 9. Get proxy URL
|
||||
proxyURL := ""
|
||||
if account.ProxyID != nil && account.Proxy != nil {
|
||||
proxyURL = account.Proxy.URL()
|
||||
}
|
||||
|
||||
// 10. Build upstream request
|
||||
upstreamCtx, releaseUpstreamCtx := detachStreamUpstreamContext(ctx, reqStream)
|
||||
upstreamReq, err := s.buildUpstreamRequest(upstreamCtx, c, account, anthropicBody, token, tokenType, mappedModel, reqStream, shouldMimicClaudeCode)
|
||||
releaseUpstreamCtx()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("build upstream request: %w", err)
|
||||
}
|
||||
|
||||
// 11. Send request
|
||||
resp, err := s.httpUpstream.DoWithTLS(upstreamReq, proxyURL, account.ID, account.Concurrency, account.IsTLSFingerprintEnabled())
|
||||
if err != nil {
|
||||
if resp != nil && resp.Body != nil {
|
||||
_ = resp.Body.Close()
|
||||
}
|
||||
safeErr := sanitizeUpstreamErrorMessage(err.Error())
|
||||
setOpsUpstreamError(c, 0, safeErr, "")
|
||||
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
|
||||
Platform: account.Platform,
|
||||
AccountID: account.ID,
|
||||
AccountName: account.Name,
|
||||
UpstreamStatusCode: 0,
|
||||
Kind: "request_error",
|
||||
Message: safeErr,
|
||||
})
|
||||
writeGatewayCCError(c, http.StatusBadGateway, "server_error", "Upstream request failed")
|
||||
return nil, fmt.Errorf("upstream request failed: %s", safeErr)
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
// 12. Handle error response with failover
|
||||
if resp.StatusCode >= 400 {
|
||||
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
|
||||
_ = resp.Body.Close()
|
||||
resp.Body = io.NopCloser(bytes.NewReader(respBody))
|
||||
|
||||
upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody))
|
||||
upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
|
||||
|
||||
if s.shouldFailoverUpstreamError(resp.StatusCode) {
|
||||
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
|
||||
Platform: account.Platform,
|
||||
AccountID: account.ID,
|
||||
AccountName: account.Name,
|
||||
UpstreamStatusCode: resp.StatusCode,
|
||||
UpstreamRequestID: resp.Header.Get("x-request-id"),
|
||||
Kind: "failover",
|
||||
Message: upstreamMsg,
|
||||
})
|
||||
if s.rateLimitService != nil {
|
||||
s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
|
||||
}
|
||||
return nil, &UpstreamFailoverError{
|
||||
StatusCode: resp.StatusCode,
|
||||
ResponseBody: respBody,
|
||||
}
|
||||
}
|
||||
|
||||
writeGatewayCCError(c, mapUpstreamStatusCode(resp.StatusCode), "server_error", upstreamMsg)
|
||||
return nil, fmt.Errorf("upstream error: %d %s", resp.StatusCode, upstreamMsg)
|
||||
}
|
||||
|
||||
// 13. Extract reasoning effort from CC request body
|
||||
reasoningEffort := extractCCReasoningEffortFromBody(body)
|
||||
|
||||
// 14. Handle normal response
|
||||
// Read Anthropic SSE → convert to Responses events → convert to CC format
|
||||
var result *ForwardResult
|
||||
var handleErr error
|
||||
if clientStream {
|
||||
result, handleErr = s.handleCCStreamingFromAnthropic(resp, c, originalModel, mappedModel, reasoningEffort, startTime, includeUsage)
|
||||
} else {
|
||||
result, handleErr = s.handleCCBufferedFromAnthropic(resp, c, originalModel, mappedModel, reasoningEffort, startTime)
|
||||
}
|
||||
|
||||
return result, handleErr
|
||||
}
|
||||
|
||||
// extractCCReasoningEffortFromBody reads reasoning effort from a Chat Completions
|
||||
// request body. It checks both nested (reasoning.effort) and flat (reasoning_effort)
|
||||
// formats used by OpenAI-compatible clients.
|
||||
func extractCCReasoningEffortFromBody(body []byte) *string {
|
||||
raw := strings.TrimSpace(gjson.GetBytes(body, "reasoning.effort").String())
|
||||
if raw == "" {
|
||||
raw = strings.TrimSpace(gjson.GetBytes(body, "reasoning_effort").String())
|
||||
}
|
||||
if raw == "" {
|
||||
return nil
|
||||
}
|
||||
normalized := normalizeOpenAIReasoningEffort(raw)
|
||||
if normalized == "" {
|
||||
return nil
|
||||
}
|
||||
return &normalized
|
||||
}
|
||||
|
||||
// handleCCBufferedFromAnthropic reads Anthropic SSE events, assembles the full
|
||||
// response, then converts Anthropic → Responses → Chat Completions.
|
||||
func (s *GatewayService) handleCCBufferedFromAnthropic(
|
||||
resp *http.Response,
|
||||
c *gin.Context,
|
||||
originalModel string,
|
||||
mappedModel string,
|
||||
reasoningEffort *string,
|
||||
startTime time.Time,
|
||||
) (*ForwardResult, error) {
|
||||
requestID := resp.Header.Get("x-request-id")
|
||||
|
||||
scanner := bufio.NewScanner(resp.Body)
|
||||
maxLineSize := defaultMaxLineSize
|
||||
if s.cfg != nil && s.cfg.Gateway.MaxLineSize > 0 {
|
||||
maxLineSize = s.cfg.Gateway.MaxLineSize
|
||||
}
|
||||
scanner.Buffer(make([]byte, 0, 64*1024), maxLineSize)
|
||||
|
||||
var finalResp *apicompat.AnthropicResponse
|
||||
var usage ClaudeUsage
|
||||
|
||||
for scanner.Scan() {
|
||||
line := scanner.Text()
|
||||
if !strings.HasPrefix(line, "event: ") {
|
||||
continue
|
||||
}
|
||||
|
||||
if !scanner.Scan() {
|
||||
break
|
||||
}
|
||||
dataLine := scanner.Text()
|
||||
if !strings.HasPrefix(dataLine, "data: ") {
|
||||
continue
|
||||
}
|
||||
payload := dataLine[6:]
|
||||
|
||||
var event apicompat.AnthropicStreamEvent
|
||||
if err := json.Unmarshal([]byte(payload), &event); err != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
// message_start carries the initial response structure and cache usage
|
||||
if event.Type == "message_start" && event.Message != nil {
|
||||
finalResp = event.Message
|
||||
mergeAnthropicUsage(&usage, event.Message.Usage)
|
||||
}
|
||||
|
||||
// message_delta carries final usage and stop_reason
|
||||
if event.Type == "message_delta" {
|
||||
if event.Usage != nil {
|
||||
mergeAnthropicUsage(&usage, *event.Usage)
|
||||
}
|
||||
if event.Delta != nil && event.Delta.StopReason != "" && finalResp != nil {
|
||||
finalResp.StopReason = event.Delta.StopReason
|
||||
}
|
||||
}
|
||||
if event.Type == "content_block_start" && event.ContentBlock != nil && finalResp != nil {
|
||||
finalResp.Content = append(finalResp.Content, *event.ContentBlock)
|
||||
}
|
||||
if event.Type == "content_block_delta" && event.Delta != nil && finalResp != nil && event.Index != nil {
|
||||
idx := *event.Index
|
||||
if idx < len(finalResp.Content) {
|
||||
switch event.Delta.Type {
|
||||
case "text_delta":
|
||||
finalResp.Content[idx].Text += event.Delta.Text
|
||||
case "thinking_delta":
|
||||
finalResp.Content[idx].Thinking += event.Delta.Thinking
|
||||
case "input_json_delta":
|
||||
finalResp.Content[idx].Input = appendRawJSON(finalResp.Content[idx].Input, event.Delta.PartialJSON)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if err := scanner.Err(); err != nil {
|
||||
if !errors.Is(err, context.Canceled) && !errors.Is(err, context.DeadlineExceeded) {
|
||||
logger.L().Warn("forward_as_cc buffered: read error",
|
||||
zap.Error(err),
|
||||
zap.String("request_id", requestID),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
if finalResp == nil {
|
||||
writeGatewayCCError(c, http.StatusBadGateway, "server_error", "Upstream stream ended without a response")
|
||||
return nil, fmt.Errorf("upstream stream ended without response")
|
||||
}
|
||||
|
||||
// Update usage from accumulated delta
|
||||
if usage.InputTokens > 0 || usage.OutputTokens > 0 {
|
||||
finalResp.Usage = apicompat.AnthropicUsage{
|
||||
InputTokens: usage.InputTokens,
|
||||
OutputTokens: usage.OutputTokens,
|
||||
CacheCreationInputTokens: usage.CacheCreationInputTokens,
|
||||
CacheReadInputTokens: usage.CacheReadInputTokens,
|
||||
}
|
||||
}
|
||||
|
||||
// Chain: Anthropic → Responses → Chat Completions
|
||||
responsesResp := apicompat.AnthropicToResponsesResponse(finalResp)
|
||||
ccResp := apicompat.ResponsesToChatCompletions(responsesResp, originalModel)
|
||||
|
||||
if s.responseHeaderFilter != nil {
|
||||
responseheaders.WriteFilteredHeaders(c.Writer.Header(), resp.Header, s.responseHeaderFilter)
|
||||
}
|
||||
c.JSON(http.StatusOK, ccResp)
|
||||
|
||||
return &ForwardResult{
|
||||
RequestID: requestID,
|
||||
Usage: usage,
|
||||
Model: originalModel,
|
||||
UpstreamModel: mappedModel,
|
||||
ReasoningEffort: reasoningEffort,
|
||||
Stream: false,
|
||||
Duration: time.Since(startTime),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// handleCCStreamingFromAnthropic reads Anthropic SSE events, converts each
|
||||
// to Responses events, then to Chat Completions chunks, and writes them.
|
||||
func (s *GatewayService) handleCCStreamingFromAnthropic(
|
||||
resp *http.Response,
|
||||
c *gin.Context,
|
||||
originalModel string,
|
||||
mappedModel string,
|
||||
reasoningEffort *string,
|
||||
startTime time.Time,
|
||||
includeUsage bool,
|
||||
) (*ForwardResult, error) {
|
||||
requestID := resp.Header.Get("x-request-id")
|
||||
|
||||
if s.responseHeaderFilter != nil {
|
||||
responseheaders.WriteFilteredHeaders(c.Writer.Header(), resp.Header, s.responseHeaderFilter)
|
||||
}
|
||||
c.Writer.Header().Set("Content-Type", "text/event-stream")
|
||||
c.Writer.Header().Set("Cache-Control", "no-cache")
|
||||
c.Writer.Header().Set("Connection", "keep-alive")
|
||||
c.Writer.Header().Set("X-Accel-Buffering", "no")
|
||||
c.Writer.WriteHeader(http.StatusOK)
|
||||
|
||||
// Use Anthropic→Responses state machine, then convert Responses→CC
|
||||
anthState := apicompat.NewAnthropicEventToResponsesState()
|
||||
anthState.Model = originalModel
|
||||
ccState := apicompat.NewResponsesEventToChatState()
|
||||
ccState.Model = originalModel
|
||||
ccState.IncludeUsage = includeUsage
|
||||
|
||||
var usage ClaudeUsage
|
||||
var firstTokenMs *int
|
||||
firstChunk := true
|
||||
|
||||
scanner := bufio.NewScanner(resp.Body)
|
||||
maxLineSize := defaultMaxLineSize
|
||||
if s.cfg != nil && s.cfg.Gateway.MaxLineSize > 0 {
|
||||
maxLineSize = s.cfg.Gateway.MaxLineSize
|
||||
}
|
||||
scanner.Buffer(make([]byte, 0, 64*1024), maxLineSize)
|
||||
|
||||
resultWithUsage := func() *ForwardResult {
|
||||
return &ForwardResult{
|
||||
RequestID: requestID,
|
||||
Usage: usage,
|
||||
Model: originalModel,
|
||||
UpstreamModel: mappedModel,
|
||||
ReasoningEffort: reasoningEffort,
|
||||
Stream: true,
|
||||
Duration: time.Since(startTime),
|
||||
FirstTokenMs: firstTokenMs,
|
||||
}
|
||||
}
|
||||
|
||||
writeChunk := func(chunk apicompat.ChatCompletionsChunk) bool {
|
||||
sse, err := apicompat.ChatChunkToSSE(chunk)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
if _, err := fmt.Fprint(c.Writer, sse); err != nil {
|
||||
return true // client disconnected
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
processAnthropicEvent := func(event *apicompat.AnthropicStreamEvent) bool {
|
||||
if firstChunk {
|
||||
firstChunk = false
|
||||
ms := int(time.Since(startTime).Milliseconds())
|
||||
firstTokenMs = &ms
|
||||
}
|
||||
|
||||
// Extract usage from message_delta
|
||||
if event.Type == "message_delta" && event.Usage != nil {
|
||||
mergeAnthropicUsage(&usage, *event.Usage)
|
||||
}
|
||||
// Also capture usage from message_start (carries cache fields)
|
||||
if event.Type == "message_start" && event.Message != nil {
|
||||
mergeAnthropicUsage(&usage, event.Message.Usage)
|
||||
}
|
||||
|
||||
// Chain: Anthropic event → Responses events → CC chunks
|
||||
responsesEvents := apicompat.AnthropicEventToResponsesEvents(event, anthState)
|
||||
for _, resEvt := range responsesEvents {
|
||||
ccChunks := apicompat.ResponsesEventToChatChunks(&resEvt, ccState)
|
||||
for _, chunk := range ccChunks {
|
||||
if disconnected := writeChunk(chunk); disconnected {
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
c.Writer.Flush()
|
||||
return false
|
||||
}
|
||||
|
||||
for scanner.Scan() {
|
||||
line := scanner.Text()
|
||||
if !strings.HasPrefix(line, "event: ") {
|
||||
continue
|
||||
}
|
||||
|
||||
if !scanner.Scan() {
|
||||
break
|
||||
}
|
||||
dataLine := scanner.Text()
|
||||
if !strings.HasPrefix(dataLine, "data: ") {
|
||||
continue
|
||||
}
|
||||
payload := dataLine[6:]
|
||||
|
||||
var event apicompat.AnthropicStreamEvent
|
||||
if err := json.Unmarshal([]byte(payload), &event); err != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
if processAnthropicEvent(&event) {
|
||||
return resultWithUsage(), nil
|
||||
}
|
||||
}
|
||||
|
||||
if err := scanner.Err(); err != nil {
|
||||
if !errors.Is(err, context.Canceled) && !errors.Is(err, context.DeadlineExceeded) {
|
||||
logger.L().Warn("forward_as_cc stream: read error",
|
||||
zap.Error(err),
|
||||
zap.String("request_id", requestID),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
// Finalize both state machines
|
||||
finalResEvents := apicompat.FinalizeAnthropicResponsesStream(anthState)
|
||||
for _, resEvt := range finalResEvents {
|
||||
ccChunks := apicompat.ResponsesEventToChatChunks(&resEvt, ccState)
|
||||
for _, chunk := range ccChunks {
|
||||
writeChunk(chunk) //nolint:errcheck
|
||||
}
|
||||
}
|
||||
finalCCChunks := apicompat.FinalizeResponsesChatStream(ccState)
|
||||
for _, chunk := range finalCCChunks {
|
||||
writeChunk(chunk) //nolint:errcheck
|
||||
}
|
||||
|
||||
// Write [DONE] marker
|
||||
fmt.Fprint(c.Writer, "data: [DONE]\n\n") //nolint:errcheck
|
||||
c.Writer.Flush()
|
||||
|
||||
return resultWithUsage(), nil
|
||||
}
|
||||
|
||||
// writeGatewayCCError writes an error in OpenAI Chat Completions format for
|
||||
// the Anthropic-upstream CC forwarding path.
|
||||
func writeGatewayCCError(c *gin.Context, statusCode int, errType, message string) {
|
||||
c.JSON(statusCode, gin.H{
|
||||
"error": gin.H{
|
||||
"type": errType,
|
||||
"message": message,
|
||||
},
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,109 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestExtractCCReasoningEffortFromBody(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
t.Run("nested reasoning.effort", func(t *testing.T) {
|
||||
got := extractCCReasoningEffortFromBody([]byte(`{"reasoning":{"effort":"HIGH"}}`))
|
||||
require.NotNil(t, got)
|
||||
require.Equal(t, "high", *got)
|
||||
})
|
||||
|
||||
t.Run("flat reasoning_effort", func(t *testing.T) {
|
||||
got := extractCCReasoningEffortFromBody([]byte(`{"reasoning_effort":"x-high"}`))
|
||||
require.NotNil(t, got)
|
||||
require.Equal(t, "xhigh", *got)
|
||||
})
|
||||
|
||||
t.Run("missing effort", func(t *testing.T) {
|
||||
require.Nil(t, extractCCReasoningEffortFromBody([]byte(`{"model":"gpt-5"}`)))
|
||||
})
|
||||
}
|
||||
|
||||
func TestHandleCCBufferedFromAnthropic_PreservesMessageStartCacheUsageAndReasoning(t *testing.T) {
|
||||
t.Parallel()
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
|
||||
reasoningEffort := "high"
|
||||
resp := &http.Response{
|
||||
Header: http.Header{"x-request-id": []string{"rid_cc_buffered"}},
|
||||
Body: io.NopCloser(strings.NewReader(strings.Join([]string{
|
||||
`event: message_start`,
|
||||
`data: {"type":"message_start","message":{"id":"msg_1","type":"message","role":"assistant","content":[],"model":"claude-sonnet-4.5","stop_reason":"","usage":{"input_tokens":12,"cache_read_input_tokens":9,"cache_creation_input_tokens":3}}}`,
|
||||
``,
|
||||
`event: content_block_start`,
|
||||
`data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":"hello"}}`,
|
||||
``,
|
||||
`event: message_delta`,
|
||||
`data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":7}}`,
|
||||
``,
|
||||
}, "\n"))),
|
||||
}
|
||||
|
||||
svc := &GatewayService{}
|
||||
result, err := svc.handleCCBufferedFromAnthropic(resp, c, "gpt-5", "claude-sonnet-4.5", &reasoningEffort, time.Now())
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.Equal(t, 12, result.Usage.InputTokens)
|
||||
require.Equal(t, 7, result.Usage.OutputTokens)
|
||||
require.Equal(t, 9, result.Usage.CacheReadInputTokens)
|
||||
require.Equal(t, 3, result.Usage.CacheCreationInputTokens)
|
||||
require.NotNil(t, result.ReasoningEffort)
|
||||
require.Equal(t, "high", *result.ReasoningEffort)
|
||||
}
|
||||
|
||||
func TestHandleCCStreamingFromAnthropic_PreservesMessageStartCacheUsageAndReasoning(t *testing.T) {
|
||||
t.Parallel()
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
|
||||
reasoningEffort := "medium"
|
||||
resp := &http.Response{
|
||||
Header: http.Header{"x-request-id": []string{"rid_cc_stream"}},
|
||||
Body: io.NopCloser(strings.NewReader(strings.Join([]string{
|
||||
`event: message_start`,
|
||||
`data: {"type":"message_start","message":{"id":"msg_2","type":"message","role":"assistant","content":[],"model":"claude-sonnet-4.5","stop_reason":"","usage":{"input_tokens":20,"cache_read_input_tokens":11,"cache_creation_input_tokens":4}}}`,
|
||||
``,
|
||||
`event: content_block_start`,
|
||||
`data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":"hello"}}`,
|
||||
``,
|
||||
`event: message_delta`,
|
||||
`data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":8}}`,
|
||||
``,
|
||||
`event: message_stop`,
|
||||
`data: {"type":"message_stop"}`,
|
||||
``,
|
||||
}, "\n"))),
|
||||
}
|
||||
|
||||
svc := &GatewayService{}
|
||||
result, err := svc.handleCCStreamingFromAnthropic(resp, c, "gpt-5", "claude-sonnet-4.5", &reasoningEffort, time.Now(), true)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.Equal(t, 20, result.Usage.InputTokens)
|
||||
require.Equal(t, 8, result.Usage.OutputTokens)
|
||||
require.Equal(t, 11, result.Usage.CacheReadInputTokens)
|
||||
require.Equal(t, 4, result.Usage.CacheCreationInputTokens)
|
||||
require.NotNil(t, result.ReasoningEffort)
|
||||
require.Equal(t, "medium", *result.ReasoningEffort)
|
||||
require.Contains(t, rec.Body.String(), `[DONE]`)
|
||||
}
|
||||
@@ -0,0 +1,518 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/apicompat"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/claude"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||||
"github.com/Wei-Shaw/sub2api/internal/util/responseheaders"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/tidwall/gjson"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// ForwardAsResponses accepts an OpenAI Responses API request body, converts it
|
||||
// to Anthropic Messages format, forwards to the Anthropic upstream, and converts
|
||||
// the response back to Responses format. This enables OpenAI Responses API
|
||||
// clients to access Anthropic models through Anthropic platform groups.
|
||||
//
|
||||
// The method follows the same pattern as OpenAIGatewayService.ForwardAsAnthropic
|
||||
// but in reverse direction: Responses → Anthropic upstream → Responses.
|
||||
func (s *GatewayService) ForwardAsResponses(
|
||||
ctx context.Context,
|
||||
c *gin.Context,
|
||||
account *Account,
|
||||
body []byte,
|
||||
parsed *ParsedRequest,
|
||||
) (*ForwardResult, error) {
|
||||
startTime := time.Now()
|
||||
|
||||
// 1. Parse Responses request
|
||||
var responsesReq apicompat.ResponsesRequest
|
||||
if err := json.Unmarshal(body, &responsesReq); err != nil {
|
||||
return nil, fmt.Errorf("parse responses request: %w", err)
|
||||
}
|
||||
originalModel := responsesReq.Model
|
||||
clientStream := responsesReq.Stream
|
||||
|
||||
// 2. Convert Responses → Anthropic
|
||||
anthropicReq, err := apicompat.ResponsesToAnthropicRequest(&responsesReq)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("convert responses to anthropic: %w", err)
|
||||
}
|
||||
|
||||
// 3. Force upstream streaming (Anthropic works best with streaming)
|
||||
anthropicReq.Stream = true
|
||||
reqStream := true
|
||||
|
||||
// 4. Model mapping
|
||||
mappedModel := originalModel
|
||||
reasoningEffort := ExtractResponsesReasoningEffortFromBody(body)
|
||||
if account.Type == AccountTypeAPIKey {
|
||||
mappedModel = account.GetMappedModel(originalModel)
|
||||
}
|
||||
if mappedModel == originalModel && account.Platform == PlatformAnthropic && account.Type != AccountTypeAPIKey {
|
||||
normalized := claude.NormalizeModelID(originalModel)
|
||||
if normalized != originalModel {
|
||||
mappedModel = normalized
|
||||
}
|
||||
}
|
||||
anthropicReq.Model = mappedModel
|
||||
|
||||
logger.L().Debug("gateway forward_as_responses: model mapping applied",
|
||||
zap.Int64("account_id", account.ID),
|
||||
zap.String("original_model", originalModel),
|
||||
zap.String("mapped_model", mappedModel),
|
||||
zap.Bool("client_stream", clientStream),
|
||||
)
|
||||
|
||||
// 5. Marshal Anthropic request body
|
||||
anthropicBody, err := json.Marshal(anthropicReq)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("marshal anthropic request: %w", err)
|
||||
}
|
||||
|
||||
// 6. Apply Claude Code mimicry for OAuth accounts (non-Claude-Code endpoints)
|
||||
isClaudeCode := false // Responses API is never Claude Code
|
||||
shouldMimicClaudeCode := account.IsOAuth() && !isClaudeCode
|
||||
|
||||
if shouldMimicClaudeCode {
|
||||
if !strings.Contains(strings.ToLower(mappedModel), "haiku") &&
|
||||
!systemIncludesClaudeCodePrompt(anthropicReq.System) {
|
||||
anthropicBody = injectClaudeCodePrompt(anthropicBody, anthropicReq.System)
|
||||
}
|
||||
}
|
||||
|
||||
// 7. Enforce cache_control block limit
|
||||
anthropicBody = enforceCacheControlLimit(anthropicBody)
|
||||
|
||||
// 8. Get access token
|
||||
token, tokenType, err := s.GetAccessToken(ctx, account)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get access token: %w", err)
|
||||
}
|
||||
|
||||
// 9. Get proxy URL
|
||||
proxyURL := ""
|
||||
if account.ProxyID != nil && account.Proxy != nil {
|
||||
proxyURL = account.Proxy.URL()
|
||||
}
|
||||
|
||||
// 10. Build upstream request
|
||||
upstreamCtx, releaseUpstreamCtx := detachStreamUpstreamContext(ctx, reqStream)
|
||||
upstreamReq, err := s.buildUpstreamRequest(upstreamCtx, c, account, anthropicBody, token, tokenType, mappedModel, reqStream, shouldMimicClaudeCode)
|
||||
releaseUpstreamCtx()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("build upstream request: %w", err)
|
||||
}
|
||||
|
||||
// 11. Send request
|
||||
resp, err := s.httpUpstream.DoWithTLS(upstreamReq, proxyURL, account.ID, account.Concurrency, account.IsTLSFingerprintEnabled())
|
||||
if err != nil {
|
||||
if resp != nil && resp.Body != nil {
|
||||
_ = resp.Body.Close()
|
||||
}
|
||||
safeErr := sanitizeUpstreamErrorMessage(err.Error())
|
||||
setOpsUpstreamError(c, 0, safeErr, "")
|
||||
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
|
||||
Platform: account.Platform,
|
||||
AccountID: account.ID,
|
||||
AccountName: account.Name,
|
||||
UpstreamStatusCode: 0,
|
||||
Kind: "request_error",
|
||||
Message: safeErr,
|
||||
})
|
||||
writeResponsesError(c, http.StatusBadGateway, "server_error", "Upstream request failed")
|
||||
return nil, fmt.Errorf("upstream request failed: %s", safeErr)
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
// 12. Handle error response with failover
|
||||
if resp.StatusCode >= 400 {
|
||||
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
|
||||
_ = resp.Body.Close()
|
||||
resp.Body = io.NopCloser(bytes.NewReader(respBody))
|
||||
|
||||
upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody))
|
||||
upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
|
||||
|
||||
if s.shouldFailoverUpstreamError(resp.StatusCode) {
|
||||
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
|
||||
Platform: account.Platform,
|
||||
AccountID: account.ID,
|
||||
AccountName: account.Name,
|
||||
UpstreamStatusCode: resp.StatusCode,
|
||||
UpstreamRequestID: resp.Header.Get("x-request-id"),
|
||||
Kind: "failover",
|
||||
Message: upstreamMsg,
|
||||
})
|
||||
if s.rateLimitService != nil {
|
||||
s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
|
||||
}
|
||||
return nil, &UpstreamFailoverError{
|
||||
StatusCode: resp.StatusCode,
|
||||
ResponseBody: respBody,
|
||||
}
|
||||
}
|
||||
|
||||
// Non-failover error: return Responses-formatted error to client
|
||||
writeResponsesError(c, mapUpstreamStatusCode(resp.StatusCode), "server_error", upstreamMsg)
|
||||
return nil, fmt.Errorf("upstream error: %d %s", resp.StatusCode, upstreamMsg)
|
||||
}
|
||||
|
||||
// 13. Handle normal response (convert Anthropic → Responses)
|
||||
var result *ForwardResult
|
||||
var handleErr error
|
||||
if clientStream {
|
||||
result, handleErr = s.handleResponsesStreamingResponse(resp, c, originalModel, mappedModel, reasoningEffort, startTime)
|
||||
} else {
|
||||
result, handleErr = s.handleResponsesBufferedStreamingResponse(resp, c, originalModel, mappedModel, reasoningEffort, startTime)
|
||||
}
|
||||
|
||||
return result, handleErr
|
||||
}
|
||||
|
||||
// ExtractResponsesReasoningEffortFromBody reads Responses API reasoning.effort
|
||||
// and normalizes it for usage logging.
|
||||
func ExtractResponsesReasoningEffortFromBody(body []byte) *string {
|
||||
raw := strings.TrimSpace(gjson.GetBytes(body, "reasoning.effort").String())
|
||||
if raw == "" {
|
||||
return nil
|
||||
}
|
||||
normalized := normalizeOpenAIReasoningEffort(raw)
|
||||
if normalized == "" {
|
||||
return nil
|
||||
}
|
||||
return &normalized
|
||||
}
|
||||
|
||||
func mergeAnthropicUsage(dst *ClaudeUsage, src apicompat.AnthropicUsage) {
|
||||
if dst == nil {
|
||||
return
|
||||
}
|
||||
if src.InputTokens > 0 {
|
||||
dst.InputTokens = src.InputTokens
|
||||
}
|
||||
if src.OutputTokens > 0 {
|
||||
dst.OutputTokens = src.OutputTokens
|
||||
}
|
||||
if src.CacheReadInputTokens > 0 {
|
||||
dst.CacheReadInputTokens = src.CacheReadInputTokens
|
||||
}
|
||||
if src.CacheCreationInputTokens > 0 {
|
||||
dst.CacheCreationInputTokens = src.CacheCreationInputTokens
|
||||
}
|
||||
}
|
||||
|
||||
// handleResponsesBufferedStreamingResponse reads all Anthropic SSE events from
|
||||
// the upstream streaming response, assembles them into a complete Anthropic
|
||||
// response, converts to Responses API JSON format, and writes it to the client.
|
||||
func (s *GatewayService) handleResponsesBufferedStreamingResponse(
|
||||
resp *http.Response,
|
||||
c *gin.Context,
|
||||
originalModel string,
|
||||
mappedModel string,
|
||||
reasoningEffort *string,
|
||||
startTime time.Time,
|
||||
) (*ForwardResult, error) {
|
||||
requestID := resp.Header.Get("x-request-id")
|
||||
|
||||
scanner := bufio.NewScanner(resp.Body)
|
||||
maxLineSize := defaultMaxLineSize
|
||||
if s.cfg != nil && s.cfg.Gateway.MaxLineSize > 0 {
|
||||
maxLineSize = s.cfg.Gateway.MaxLineSize
|
||||
}
|
||||
scanner.Buffer(make([]byte, 0, 64*1024), maxLineSize)
|
||||
|
||||
// Accumulate the final Anthropic response from streaming events
|
||||
var finalResp *apicompat.AnthropicResponse
|
||||
var usage ClaudeUsage
|
||||
|
||||
for scanner.Scan() {
|
||||
line := scanner.Text()
|
||||
if !strings.HasPrefix(line, "event: ") {
|
||||
continue
|
||||
}
|
||||
eventType := strings.TrimPrefix(line, "event: ")
|
||||
|
||||
// Read the data line
|
||||
if !scanner.Scan() {
|
||||
break
|
||||
}
|
||||
dataLine := scanner.Text()
|
||||
if !strings.HasPrefix(dataLine, "data: ") {
|
||||
continue
|
||||
}
|
||||
payload := dataLine[6:]
|
||||
|
||||
var event apicompat.AnthropicStreamEvent
|
||||
if err := json.Unmarshal([]byte(payload), &event); err != nil {
|
||||
logger.L().Warn("forward_as_responses buffered: failed to parse event",
|
||||
zap.Error(err),
|
||||
zap.String("request_id", requestID),
|
||||
zap.String("event_type", eventType),
|
||||
)
|
||||
continue
|
||||
}
|
||||
|
||||
// message_start carries the initial response structure
|
||||
if event.Type == "message_start" && event.Message != nil {
|
||||
finalResp = event.Message
|
||||
mergeAnthropicUsage(&usage, event.Message.Usage)
|
||||
}
|
||||
|
||||
// message_delta carries final usage and stop_reason
|
||||
if event.Type == "message_delta" {
|
||||
if event.Usage != nil {
|
||||
mergeAnthropicUsage(&usage, *event.Usage)
|
||||
}
|
||||
if event.Delta != nil && event.Delta.StopReason != "" && finalResp != nil {
|
||||
finalResp.StopReason = event.Delta.StopReason
|
||||
}
|
||||
}
|
||||
|
||||
// Accumulate content blocks
|
||||
if event.Type == "content_block_start" && event.ContentBlock != nil && finalResp != nil {
|
||||
finalResp.Content = append(finalResp.Content, *event.ContentBlock)
|
||||
}
|
||||
if event.Type == "content_block_delta" && event.Delta != nil && finalResp != nil && event.Index != nil {
|
||||
idx := *event.Index
|
||||
if idx < len(finalResp.Content) {
|
||||
switch event.Delta.Type {
|
||||
case "text_delta":
|
||||
finalResp.Content[idx].Text += event.Delta.Text
|
||||
case "thinking_delta":
|
||||
finalResp.Content[idx].Thinking += event.Delta.Thinking
|
||||
case "input_json_delta":
|
||||
finalResp.Content[idx].Input = appendRawJSON(finalResp.Content[idx].Input, event.Delta.PartialJSON)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if err := scanner.Err(); err != nil {
|
||||
if !errors.Is(err, context.Canceled) && !errors.Is(err, context.DeadlineExceeded) {
|
||||
logger.L().Warn("forward_as_responses buffered: read error",
|
||||
zap.Error(err),
|
||||
zap.String("request_id", requestID),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
if finalResp == nil {
|
||||
writeResponsesError(c, http.StatusBadGateway, "server_error", "Upstream stream ended without a response")
|
||||
return nil, fmt.Errorf("upstream stream ended without response")
|
||||
}
|
||||
|
||||
// Update usage from accumulated delta
|
||||
if usage.InputTokens > 0 || usage.OutputTokens > 0 {
|
||||
finalResp.Usage = apicompat.AnthropicUsage{
|
||||
InputTokens: usage.InputTokens,
|
||||
OutputTokens: usage.OutputTokens,
|
||||
CacheCreationInputTokens: usage.CacheCreationInputTokens,
|
||||
CacheReadInputTokens: usage.CacheReadInputTokens,
|
||||
}
|
||||
}
|
||||
|
||||
// Convert to Responses format
|
||||
responsesResp := apicompat.AnthropicToResponsesResponse(finalResp)
|
||||
responsesResp.Model = originalModel // Use original model name
|
||||
|
||||
if s.responseHeaderFilter != nil {
|
||||
responseheaders.WriteFilteredHeaders(c.Writer.Header(), resp.Header, s.responseHeaderFilter)
|
||||
}
|
||||
c.JSON(http.StatusOK, responsesResp)
|
||||
|
||||
return &ForwardResult{
|
||||
RequestID: requestID,
|
||||
Usage: usage,
|
||||
Model: originalModel,
|
||||
UpstreamModel: mappedModel,
|
||||
ReasoningEffort: reasoningEffort,
|
||||
Stream: false,
|
||||
Duration: time.Since(startTime),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// handleResponsesStreamingResponse reads Anthropic SSE events from upstream,
|
||||
// converts each to Responses SSE events, and writes them to the client.
|
||||
func (s *GatewayService) handleResponsesStreamingResponse(
|
||||
resp *http.Response,
|
||||
c *gin.Context,
|
||||
originalModel string,
|
||||
mappedModel string,
|
||||
reasoningEffort *string,
|
||||
startTime time.Time,
|
||||
) (*ForwardResult, error) {
|
||||
requestID := resp.Header.Get("x-request-id")
|
||||
|
||||
if s.responseHeaderFilter != nil {
|
||||
responseheaders.WriteFilteredHeaders(c.Writer.Header(), resp.Header, s.responseHeaderFilter)
|
||||
}
|
||||
c.Writer.Header().Set("Content-Type", "text/event-stream")
|
||||
c.Writer.Header().Set("Cache-Control", "no-cache")
|
||||
c.Writer.Header().Set("Connection", "keep-alive")
|
||||
c.Writer.Header().Set("X-Accel-Buffering", "no")
|
||||
c.Writer.WriteHeader(http.StatusOK)
|
||||
|
||||
state := apicompat.NewAnthropicEventToResponsesState()
|
||||
state.Model = originalModel
|
||||
var usage ClaudeUsage
|
||||
var firstTokenMs *int
|
||||
firstChunk := true
|
||||
|
||||
scanner := bufio.NewScanner(resp.Body)
|
||||
maxLineSize := defaultMaxLineSize
|
||||
if s.cfg != nil && s.cfg.Gateway.MaxLineSize > 0 {
|
||||
maxLineSize = s.cfg.Gateway.MaxLineSize
|
||||
}
|
||||
scanner.Buffer(make([]byte, 0, 64*1024), maxLineSize)
|
||||
|
||||
resultWithUsage := func() *ForwardResult {
|
||||
return &ForwardResult{
|
||||
RequestID: requestID,
|
||||
Usage: usage,
|
||||
Model: originalModel,
|
||||
UpstreamModel: mappedModel,
|
||||
ReasoningEffort: reasoningEffort,
|
||||
Stream: true,
|
||||
Duration: time.Since(startTime),
|
||||
FirstTokenMs: firstTokenMs,
|
||||
}
|
||||
}
|
||||
|
||||
// processEvent handles a single parsed Anthropic SSE event.
|
||||
processEvent := func(event *apicompat.AnthropicStreamEvent) bool {
|
||||
if firstChunk {
|
||||
firstChunk = false
|
||||
ms := int(time.Since(startTime).Milliseconds())
|
||||
firstTokenMs = &ms
|
||||
}
|
||||
|
||||
// Extract usage from message_delta
|
||||
if event.Type == "message_delta" && event.Usage != nil {
|
||||
mergeAnthropicUsage(&usage, *event.Usage)
|
||||
}
|
||||
// Also capture usage from message_start
|
||||
if event.Type == "message_start" && event.Message != nil {
|
||||
mergeAnthropicUsage(&usage, event.Message.Usage)
|
||||
}
|
||||
|
||||
// Convert to Responses events
|
||||
events := apicompat.AnthropicEventToResponsesEvents(event, state)
|
||||
for _, evt := range events {
|
||||
sse, err := apicompat.ResponsesEventToSSE(evt)
|
||||
if err != nil {
|
||||
logger.L().Warn("forward_as_responses stream: failed to marshal event",
|
||||
zap.Error(err),
|
||||
zap.String("request_id", requestID),
|
||||
)
|
||||
continue
|
||||
}
|
||||
if _, err := fmt.Fprint(c.Writer, sse); err != nil {
|
||||
logger.L().Info("forward_as_responses stream: client disconnected",
|
||||
zap.String("request_id", requestID),
|
||||
)
|
||||
return true // client disconnected
|
||||
}
|
||||
}
|
||||
if len(events) > 0 {
|
||||
c.Writer.Flush()
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
finalizeStream := func() (*ForwardResult, error) {
|
||||
if finalEvents := apicompat.FinalizeAnthropicResponsesStream(state); len(finalEvents) > 0 {
|
||||
for _, evt := range finalEvents {
|
||||
sse, err := apicompat.ResponsesEventToSSE(evt)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
fmt.Fprint(c.Writer, sse) //nolint:errcheck
|
||||
}
|
||||
c.Writer.Flush()
|
||||
}
|
||||
return resultWithUsage(), nil
|
||||
}
|
||||
|
||||
// Read Anthropic SSE events
|
||||
for scanner.Scan() {
|
||||
line := scanner.Text()
|
||||
if !strings.HasPrefix(line, "event: ") {
|
||||
continue
|
||||
}
|
||||
eventType := strings.TrimPrefix(line, "event: ")
|
||||
|
||||
// Read data line
|
||||
if !scanner.Scan() {
|
||||
break
|
||||
}
|
||||
dataLine := scanner.Text()
|
||||
if !strings.HasPrefix(dataLine, "data: ") {
|
||||
continue
|
||||
}
|
||||
payload := dataLine[6:]
|
||||
|
||||
var event apicompat.AnthropicStreamEvent
|
||||
if err := json.Unmarshal([]byte(payload), &event); err != nil {
|
||||
logger.L().Warn("forward_as_responses stream: failed to parse event",
|
||||
zap.Error(err),
|
||||
zap.String("request_id", requestID),
|
||||
zap.String("event_type", eventType),
|
||||
)
|
||||
continue
|
||||
}
|
||||
|
||||
if processEvent(&event) {
|
||||
return resultWithUsage(), nil
|
||||
}
|
||||
}
|
||||
|
||||
if err := scanner.Err(); err != nil {
|
||||
if !errors.Is(err, context.Canceled) && !errors.Is(err, context.DeadlineExceeded) {
|
||||
logger.L().Warn("forward_as_responses stream: read error",
|
||||
zap.Error(err),
|
||||
zap.String("request_id", requestID),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
return finalizeStream()
|
||||
}
|
||||
|
||||
// appendRawJSON appends a JSON fragment string to existing raw JSON.
|
||||
func appendRawJSON(existing json.RawMessage, fragment string) json.RawMessage {
|
||||
if len(existing) == 0 {
|
||||
return json.RawMessage(fragment)
|
||||
}
|
||||
return json.RawMessage(string(existing) + fragment)
|
||||
}
|
||||
|
||||
// writeResponsesError writes an error response in OpenAI Responses API format.
|
||||
func writeResponsesError(c *gin.Context, statusCode int, code, message string) {
|
||||
c.JSON(statusCode, gin.H{
|
||||
"error": gin.H{
|
||||
"code": code,
|
||||
"message": message,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
// mapUpstreamStatusCode maps upstream HTTP status codes to appropriate client-facing codes.
|
||||
func mapUpstreamStatusCode(code int) int {
|
||||
if code >= 500 {
|
||||
return http.StatusBadGateway
|
||||
}
|
||||
return code
|
||||
}
|
||||
@@ -0,0 +1,94 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestExtractResponsesReasoningEffortFromBody(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
got := ExtractResponsesReasoningEffortFromBody([]byte(`{"model":"claude-sonnet-4.5","reasoning":{"effort":"HIGH"}}`))
|
||||
require.NotNil(t, got)
|
||||
require.Equal(t, "high", *got)
|
||||
|
||||
require.Nil(t, ExtractResponsesReasoningEffortFromBody([]byte(`{"model":"claude-sonnet-4.5"}`)))
|
||||
}
|
||||
|
||||
func TestHandleResponsesBufferedStreamingResponse_PreservesMessageStartCacheUsage(t *testing.T) {
|
||||
t.Parallel()
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
|
||||
resp := &http.Response{
|
||||
Header: http.Header{"x-request-id": []string{"rid_buffered"}},
|
||||
Body: io.NopCloser(strings.NewReader(strings.Join([]string{
|
||||
`event: message_start`,
|
||||
`data: {"type":"message_start","message":{"id":"msg_1","type":"message","role":"assistant","content":[],"model":"claude-sonnet-4.5","stop_reason":"","usage":{"input_tokens":12,"cache_read_input_tokens":9,"cache_creation_input_tokens":3}}}`,
|
||||
``,
|
||||
`event: content_block_start`,
|
||||
`data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":"hello"}}`,
|
||||
``,
|
||||
`event: message_delta`,
|
||||
`data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":7}}`,
|
||||
``,
|
||||
}, "\n"))),
|
||||
}
|
||||
|
||||
svc := &GatewayService{}
|
||||
result, err := svc.handleResponsesBufferedStreamingResponse(resp, c, "claude-sonnet-4.5", "claude-sonnet-4.5", nil, time.Now())
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.Equal(t, 12, result.Usage.InputTokens)
|
||||
require.Equal(t, 7, result.Usage.OutputTokens)
|
||||
require.Equal(t, 9, result.Usage.CacheReadInputTokens)
|
||||
require.Equal(t, 3, result.Usage.CacheCreationInputTokens)
|
||||
require.Contains(t, rec.Body.String(), `"cached_tokens":9`)
|
||||
}
|
||||
|
||||
func TestHandleResponsesStreamingResponse_PreservesMessageStartCacheUsage(t *testing.T) {
|
||||
t.Parallel()
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
|
||||
resp := &http.Response{
|
||||
Header: http.Header{"x-request-id": []string{"rid_stream"}},
|
||||
Body: io.NopCloser(strings.NewReader(strings.Join([]string{
|
||||
`event: message_start`,
|
||||
`data: {"type":"message_start","message":{"id":"msg_2","type":"message","role":"assistant","content":[],"model":"claude-sonnet-4.5","stop_reason":"","usage":{"input_tokens":20,"cache_read_input_tokens":11,"cache_creation_input_tokens":4}}}`,
|
||||
``,
|
||||
`event: content_block_start`,
|
||||
`data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":"hello"}}`,
|
||||
``,
|
||||
`event: message_delta`,
|
||||
`data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":8}}`,
|
||||
``,
|
||||
`event: message_stop`,
|
||||
`data: {"type":"message_stop"}`,
|
||||
``,
|
||||
}, "\n"))),
|
||||
}
|
||||
|
||||
svc := &GatewayService{}
|
||||
result, err := svc.handleResponsesStreamingResponse(resp, c, "claude-sonnet-4.5", "claude-sonnet-4.5", nil, time.Now())
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.Equal(t, 20, result.Usage.InputTokens)
|
||||
require.Equal(t, 8, result.Usage.OutputTokens)
|
||||
require.Equal(t, 11, result.Usage.CacheReadInputTokens)
|
||||
require.Equal(t, 4, result.Usage.CacheCreationInputTokens)
|
||||
require.Contains(t, rec.Body.String(), `response.completed`)
|
||||
}
|
||||
@@ -92,7 +92,7 @@ func (m *mockAccountRepoForPlatform) Delete(ctx context.Context, id int64) error
|
||||
func (m *mockAccountRepoForPlatform) List(ctx context.Context, params pagination.PaginationParams) ([]Account, *pagination.PaginationResult, error) {
|
||||
return nil, nil, nil
|
||||
}
|
||||
func (m *mockAccountRepoForPlatform) ListWithFilters(ctx context.Context, params pagination.PaginationParams, platform, accountType, status, search string, groupID int64) ([]Account, *pagination.PaginationResult, error) {
|
||||
func (m *mockAccountRepoForPlatform) ListWithFilters(ctx context.Context, params pagination.PaginationParams, platform, accountType, status, search string, groupID int64, privacyMode string) ([]Account, *pagination.PaginationResult, error) {
|
||||
return nil, nil, nil
|
||||
}
|
||||
func (m *mockAccountRepoForPlatform) ListByGroup(ctx context.Context, groupID int64) ([]Account, error) {
|
||||
|
||||
@@ -5,6 +5,8 @@ import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"math"
|
||||
"regexp"
|
||||
"sort"
|
||||
"strings"
|
||||
"unsafe"
|
||||
|
||||
@@ -34,6 +36,9 @@ var (
|
||||
patternEmptyTextSpaced = []byte(`"text": ""`)
|
||||
patternEmptyTextSp1 = []byte(`"text" : ""`)
|
||||
patternEmptyTextSp2 = []byte(`"text" :""`)
|
||||
|
||||
sessionUserAgentProductPattern = regexp.MustCompile(`([A-Za-z0-9._-]+)/[A-Za-z0-9._-]+`)
|
||||
sessionUserAgentVersionPattern = regexp.MustCompile(`\bv?\d+(?:\.\d+){1,3}\b`)
|
||||
)
|
||||
|
||||
// SessionContext 粘性会话上下文,用于区分不同来源的请求。
|
||||
@@ -75,6 +80,49 @@ type ParsedRequest struct {
|
||||
OnUpstreamAccepted func()
|
||||
}
|
||||
|
||||
// NormalizeSessionUserAgent reduces UA noise for sticky-session and digest hashing.
|
||||
// It preserves the set of product names from Product/Version tokens while
|
||||
// discarding version-only changes and incidental comments.
|
||||
func NormalizeSessionUserAgent(raw string) string {
|
||||
raw = strings.TrimSpace(raw)
|
||||
if raw == "" {
|
||||
return ""
|
||||
}
|
||||
|
||||
matches := sessionUserAgentProductPattern.FindAllStringSubmatch(raw, -1)
|
||||
if len(matches) == 0 {
|
||||
return normalizeSessionUserAgentFallback(raw)
|
||||
}
|
||||
|
||||
products := make([]string, 0, len(matches))
|
||||
seen := make(map[string]struct{}, len(matches))
|
||||
for _, match := range matches {
|
||||
if len(match) < 2 {
|
||||
continue
|
||||
}
|
||||
product := strings.ToLower(strings.TrimSpace(match[1]))
|
||||
if product == "" {
|
||||
continue
|
||||
}
|
||||
if _, exists := seen[product]; exists {
|
||||
continue
|
||||
}
|
||||
seen[product] = struct{}{}
|
||||
products = append(products, product)
|
||||
}
|
||||
if len(products) == 0 {
|
||||
return normalizeSessionUserAgentFallback(raw)
|
||||
}
|
||||
sort.Strings(products)
|
||||
return strings.Join(products, "+")
|
||||
}
|
||||
|
||||
func normalizeSessionUserAgentFallback(raw string) string {
|
||||
normalized := strings.ToLower(strings.Join(strings.Fields(raw), " "))
|
||||
normalized = sessionUserAgentVersionPattern.ReplaceAllString(normalized, "")
|
||||
return strings.Join(strings.Fields(normalized), " ")
|
||||
}
|
||||
|
||||
// ParseGatewayRequest 解析网关请求体并返回结构化结果。
|
||||
// protocol 指定请求协议格式(domain.PlatformAnthropic / domain.PlatformGemini),
|
||||
// 不同协议使用不同的 system/messages 字段名。
|
||||
@@ -205,6 +253,118 @@ func sliceRawFromBody(body []byte, r gjson.Result) []byte {
|
||||
return []byte(r.Raw)
|
||||
}
|
||||
|
||||
// stripEmptyTextBlocksFromSlice removes empty text blocks from a content slice (including nested tool_result content).
|
||||
// Returns (cleaned slice, true) if any blocks were removed, or (original, false) if unchanged.
|
||||
func stripEmptyTextBlocksFromSlice(blocks []any) ([]any, bool) {
|
||||
var result []any
|
||||
changed := false
|
||||
for i, block := range blocks {
|
||||
blockMap, ok := block.(map[string]any)
|
||||
if !ok {
|
||||
if result != nil {
|
||||
result = append(result, block)
|
||||
}
|
||||
continue
|
||||
}
|
||||
blockType, _ := blockMap["type"].(string)
|
||||
|
||||
// Strip empty text blocks
|
||||
if blockType == "text" {
|
||||
if txt, _ := blockMap["text"].(string); txt == "" {
|
||||
if result == nil {
|
||||
result = make([]any, 0, len(blocks))
|
||||
result = append(result, blocks[:i]...)
|
||||
}
|
||||
changed = true
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
// Recurse into tool_result nested content
|
||||
if blockType == "tool_result" {
|
||||
if nestedContent, ok := blockMap["content"].([]any); ok {
|
||||
if cleaned, nestedChanged := stripEmptyTextBlocksFromSlice(nestedContent); nestedChanged {
|
||||
if result == nil {
|
||||
result = make([]any, 0, len(blocks))
|
||||
result = append(result, blocks[:i]...)
|
||||
}
|
||||
changed = true
|
||||
blockCopy := make(map[string]any, len(blockMap))
|
||||
for k, v := range blockMap {
|
||||
blockCopy[k] = v
|
||||
}
|
||||
blockCopy["content"] = cleaned
|
||||
result = append(result, blockCopy)
|
||||
continue
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if result != nil {
|
||||
result = append(result, block)
|
||||
}
|
||||
}
|
||||
if !changed {
|
||||
return blocks, false
|
||||
}
|
||||
return result, true
|
||||
}
|
||||
|
||||
// StripEmptyTextBlocks removes empty text blocks from the request body (including nested tool_result content).
|
||||
// This is a lightweight pre-filter for the initial request path to prevent upstream 400 errors.
|
||||
// Returns the original body unchanged if no empty text blocks are found.
|
||||
func StripEmptyTextBlocks(body []byte) []byte {
|
||||
// Fast path: check if body contains empty text patterns
|
||||
hasEmptyTextBlock := bytes.Contains(body, patternEmptyText) ||
|
||||
bytes.Contains(body, patternEmptyTextSpaced) ||
|
||||
bytes.Contains(body, patternEmptyTextSp1) ||
|
||||
bytes.Contains(body, patternEmptyTextSp2)
|
||||
if !hasEmptyTextBlock {
|
||||
return body
|
||||
}
|
||||
|
||||
jsonStr := *(*string)(unsafe.Pointer(&body))
|
||||
msgsRes := gjson.Get(jsonStr, "messages")
|
||||
if !msgsRes.Exists() || !msgsRes.IsArray() {
|
||||
return body
|
||||
}
|
||||
|
||||
var messages []any
|
||||
if err := json.Unmarshal(sliceRawFromBody(body, msgsRes), &messages); err != nil {
|
||||
return body
|
||||
}
|
||||
|
||||
modified := false
|
||||
for _, msg := range messages {
|
||||
msgMap, ok := msg.(map[string]any)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
content, ok := msgMap["content"].([]any)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
if cleaned, changed := stripEmptyTextBlocksFromSlice(content); changed {
|
||||
modified = true
|
||||
msgMap["content"] = cleaned
|
||||
}
|
||||
}
|
||||
|
||||
if !modified {
|
||||
return body
|
||||
}
|
||||
|
||||
msgsBytes, err := json.Marshal(messages)
|
||||
if err != nil {
|
||||
return body
|
||||
}
|
||||
out, err := sjson.SetRawBytes(body, "messages", msgsBytes)
|
||||
if err != nil {
|
||||
return body
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// FilterThinkingBlocks removes thinking blocks from request body
|
||||
// Returns filtered body or original body if filtering fails (fail-safe)
|
||||
// This prevents 400 errors from invalid thinking block signatures
|
||||
@@ -378,6 +538,23 @@ func FilterThinkingBlocksForRetry(body []byte) []byte {
|
||||
}
|
||||
}
|
||||
|
||||
// Recursively strip empty text blocks from tool_result nested content.
|
||||
if blockType == "tool_result" {
|
||||
if nestedContent, ok := blockMap["content"].([]any); ok {
|
||||
if cleaned, changed := stripEmptyTextBlocksFromSlice(nestedContent); changed {
|
||||
modifiedThisMsg = true
|
||||
ensureNewContent(bi)
|
||||
blockCopy := make(map[string]any, len(blockMap))
|
||||
for k, v := range blockMap {
|
||||
blockCopy[k] = v
|
||||
}
|
||||
blockCopy["content"] = cleaned
|
||||
newContent = append(newContent, blockCopy)
|
||||
continue
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if newContent != nil {
|
||||
newContent = append(newContent, block)
|
||||
}
|
||||
|
||||
@@ -435,6 +435,122 @@ func TestFilterThinkingBlocksForRetry_StripsEmptyTextBlocks(t *testing.T) {
|
||||
require.NotEmpty(t, block1["text"])
|
||||
}
|
||||
|
||||
func TestFilterThinkingBlocksForRetry_StripsNestedEmptyTextInToolResult(t *testing.T) {
|
||||
// Empty text blocks nested inside tool_result content should also be stripped
|
||||
input := []byte(`{
|
||||
"messages":[
|
||||
{"role":"user","content":[
|
||||
{"type":"tool_result","tool_use_id":"t1","content":[
|
||||
{"type":"text","text":"valid result"},
|
||||
{"type":"text","text":""}
|
||||
]}
|
||||
]}
|
||||
]
|
||||
}`)
|
||||
|
||||
out := FilterThinkingBlocksForRetry(input)
|
||||
|
||||
var req map[string]any
|
||||
require.NoError(t, json.Unmarshal(out, &req))
|
||||
msgs := req["messages"].([]any)
|
||||
msg0 := msgs[0].(map[string]any)
|
||||
content0 := msg0["content"].([]any)
|
||||
require.Len(t, content0, 1)
|
||||
toolResult := content0[0].(map[string]any)
|
||||
require.Equal(t, "tool_result", toolResult["type"])
|
||||
nestedContent := toolResult["content"].([]any)
|
||||
require.Len(t, nestedContent, 1)
|
||||
require.Equal(t, "valid result", nestedContent[0].(map[string]any)["text"])
|
||||
}
|
||||
|
||||
func TestFilterThinkingBlocksForRetry_NestedAllEmptyGetsEmptySlice(t *testing.T) {
|
||||
// If all nested content blocks in tool_result are empty text, content becomes empty slice
|
||||
input := []byte(`{
|
||||
"messages":[
|
||||
{"role":"user","content":[
|
||||
{"type":"tool_result","tool_use_id":"t1","content":[
|
||||
{"type":"text","text":""}
|
||||
]},
|
||||
{"type":"text","text":"hello"}
|
||||
]}
|
||||
]
|
||||
}`)
|
||||
|
||||
out := FilterThinkingBlocksForRetry(input)
|
||||
|
||||
var req map[string]any
|
||||
require.NoError(t, json.Unmarshal(out, &req))
|
||||
msgs := req["messages"].([]any)
|
||||
msg0 := msgs[0].(map[string]any)
|
||||
content0 := msg0["content"].([]any)
|
||||
require.Len(t, content0, 2)
|
||||
toolResult := content0[0].(map[string]any)
|
||||
nestedContent := toolResult["content"].([]any)
|
||||
require.Len(t, nestedContent, 0)
|
||||
}
|
||||
|
||||
func TestStripEmptyTextBlocks(t *testing.T) {
|
||||
t.Run("strips top-level empty text", func(t *testing.T) {
|
||||
input := []byte(`{"messages":[{"role":"user","content":[{"type":"text","text":"hello"},{"type":"text","text":""}]}]}`)
|
||||
out := StripEmptyTextBlocks(input)
|
||||
var req map[string]any
|
||||
require.NoError(t, json.Unmarshal(out, &req))
|
||||
msgs := req["messages"].([]any)
|
||||
content := msgs[0].(map[string]any)["content"].([]any)
|
||||
require.Len(t, content, 1)
|
||||
require.Equal(t, "hello", content[0].(map[string]any)["text"])
|
||||
})
|
||||
|
||||
t.Run("strips nested empty text in tool_result", func(t *testing.T) {
|
||||
input := []byte(`{"messages":[{"role":"user","content":[{"type":"tool_result","tool_use_id":"t1","content":[{"type":"text","text":"ok"},{"type":"text","text":""}]}]}]}`)
|
||||
out := StripEmptyTextBlocks(input)
|
||||
var req map[string]any
|
||||
require.NoError(t, json.Unmarshal(out, &req))
|
||||
msgs := req["messages"].([]any)
|
||||
content := msgs[0].(map[string]any)["content"].([]any)
|
||||
toolResult := content[0].(map[string]any)
|
||||
nestedContent := toolResult["content"].([]any)
|
||||
require.Len(t, nestedContent, 1)
|
||||
require.Equal(t, "ok", nestedContent[0].(map[string]any)["text"])
|
||||
})
|
||||
|
||||
t.Run("no-op when no empty text", func(t *testing.T) {
|
||||
input := []byte(`{"messages":[{"role":"user","content":[{"type":"text","text":"hello"}]}]}`)
|
||||
out := StripEmptyTextBlocks(input)
|
||||
require.Equal(t, input, out)
|
||||
})
|
||||
|
||||
t.Run("preserves non-map blocks in content", func(t *testing.T) {
|
||||
// tool_result content can be a string; non-map blocks should pass through unchanged
|
||||
input := []byte(`{"messages":[{"role":"user","content":[{"type":"tool_result","tool_use_id":"t1","content":"string content"},{"type":"text","text":""}]}]}`)
|
||||
out := StripEmptyTextBlocks(input)
|
||||
var req map[string]any
|
||||
require.NoError(t, json.Unmarshal(out, &req))
|
||||
msgs := req["messages"].([]any)
|
||||
content := msgs[0].(map[string]any)["content"].([]any)
|
||||
require.Len(t, content, 1)
|
||||
toolResult := content[0].(map[string]any)
|
||||
require.Equal(t, "tool_result", toolResult["type"])
|
||||
require.Equal(t, "string content", toolResult["content"])
|
||||
})
|
||||
|
||||
t.Run("handles deeply nested tool_result", func(t *testing.T) {
|
||||
// Recursive: tool_result containing another tool_result with empty text
|
||||
input := []byte(`{"messages":[{"role":"user","content":[{"type":"tool_result","tool_use_id":"t1","content":[{"type":"tool_result","tool_use_id":"t2","content":[{"type":"text","text":""},{"type":"text","text":"deep"}]}]}]}]}`)
|
||||
out := StripEmptyTextBlocks(input)
|
||||
var req map[string]any
|
||||
require.NoError(t, json.Unmarshal(out, &req))
|
||||
msgs := req["messages"].([]any)
|
||||
content := msgs[0].(map[string]any)["content"].([]any)
|
||||
outer := content[0].(map[string]any)
|
||||
innerContent := outer["content"].([]any)
|
||||
inner := innerContent[0].(map[string]any)
|
||||
deepContent := inner["content"].([]any)
|
||||
require.Len(t, deepContent, 1)
|
||||
require.Equal(t, "deep", deepContent[0].(map[string]any)["text"])
|
||||
})
|
||||
}
|
||||
|
||||
func TestFilterThinkingBlocksForRetry_PreservesNonEmptyTextBlocks(t *testing.T) {
|
||||
// Non-empty text blocks should pass through unchanged
|
||||
input := []byte(`{
|
||||
|
||||
@@ -658,7 +658,7 @@ func (s *GatewayService) GenerateSessionHash(parsed *ParsedRequest) string {
|
||||
if parsed.SessionContext != nil {
|
||||
_, _ = combined.WriteString(parsed.SessionContext.ClientIP)
|
||||
_, _ = combined.WriteString(":")
|
||||
_, _ = combined.WriteString(parsed.SessionContext.UserAgent)
|
||||
_, _ = combined.WriteString(NormalizeSessionUserAgent(parsed.SessionContext.UserAgent))
|
||||
_, _ = combined.WriteString(":")
|
||||
_, _ = combined.WriteString(strconv.FormatInt(parsed.SessionContext.APIKeyID, 10))
|
||||
_, _ = combined.WriteString("|")
|
||||
@@ -4119,6 +4119,9 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
|
||||
// 调试日志:记录即将转发的账号信息
|
||||
logger.LegacyPrintf("service.gateway", "[Forward] Using account: ID=%d Name=%s Platform=%s Type=%s TLSFingerprint=%v Proxy=%s",
|
||||
account.ID, account.Name, account.Platform, account.Type, account.IsTLSFingerprintEnabled(), proxyURL)
|
||||
// Pre-filter: strip empty text blocks (including nested in tool_result) to prevent upstream 400.
|
||||
body = StripEmptyTextBlocks(body)
|
||||
|
||||
// 重试间复用同一请求体,避免每次 string(body) 产生额外分配。
|
||||
setOpsUpstreamRequestBody(c, body)
|
||||
|
||||
@@ -4148,6 +4151,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
|
||||
AccountID: account.ID,
|
||||
AccountName: account.Name,
|
||||
UpstreamStatusCode: 0,
|
||||
UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()),
|
||||
Kind: "request_error",
|
||||
Message: safeErr,
|
||||
})
|
||||
@@ -4174,6 +4178,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
|
||||
AccountName: account.Name,
|
||||
UpstreamStatusCode: resp.StatusCode,
|
||||
UpstreamRequestID: resp.Header.Get("x-request-id"),
|
||||
UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()),
|
||||
Kind: "signature_error",
|
||||
Message: extractUpstreamErrorMessage(respBody),
|
||||
Detail: func() string {
|
||||
@@ -4228,6 +4233,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
|
||||
AccountName: account.Name,
|
||||
UpstreamStatusCode: retryResp.StatusCode,
|
||||
UpstreamRequestID: retryResp.Header.Get("x-request-id"),
|
||||
UpstreamURL: safeUpstreamURL(retryReq.URL.String()),
|
||||
Kind: "signature_retry_thinking",
|
||||
Message: extractUpstreamErrorMessage(retryRespBody),
|
||||
Detail: func() string {
|
||||
@@ -4258,6 +4264,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
|
||||
AccountID: account.ID,
|
||||
AccountName: account.Name,
|
||||
UpstreamStatusCode: 0,
|
||||
UpstreamURL: safeUpstreamURL(retryReq2.URL.String()),
|
||||
Kind: "signature_retry_tools_request_error",
|
||||
Message: sanitizeUpstreamErrorMessage(retryErr2.Error()),
|
||||
})
|
||||
@@ -4297,6 +4304,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
|
||||
AccountName: account.Name,
|
||||
UpstreamStatusCode: resp.StatusCode,
|
||||
UpstreamRequestID: resp.Header.Get("x-request-id"),
|
||||
UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()),
|
||||
Kind: "budget_constraint_error",
|
||||
Message: errMsg,
|
||||
Detail: func() string {
|
||||
@@ -4358,6 +4366,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
|
||||
AccountName: account.Name,
|
||||
UpstreamStatusCode: resp.StatusCode,
|
||||
UpstreamRequestID: resp.Header.Get("x-request-id"),
|
||||
UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()),
|
||||
Kind: "retry",
|
||||
Message: extractUpstreamErrorMessage(respBody),
|
||||
Detail: func() string {
|
||||
@@ -4603,6 +4612,9 @@ func (s *GatewayService) forwardAnthropicAPIKeyPassthroughWithInput(
|
||||
if c != nil {
|
||||
c.Set("anthropic_passthrough", true)
|
||||
}
|
||||
// Pre-filter: strip empty text blocks (including nested in tool_result) to prevent upstream 400.
|
||||
input.Body = StripEmptyTextBlocks(input.Body)
|
||||
|
||||
// 重试间复用同一请求体,避免每次 string(body) 产生额外分配。
|
||||
setOpsUpstreamRequestBody(c, input.Body)
|
||||
|
||||
@@ -4628,6 +4640,7 @@ func (s *GatewayService) forwardAnthropicAPIKeyPassthroughWithInput(
|
||||
AccountID: account.ID,
|
||||
AccountName: account.Name,
|
||||
UpstreamStatusCode: 0,
|
||||
UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()),
|
||||
Passthrough: true,
|
||||
Kind: "request_error",
|
||||
Message: safeErr,
|
||||
@@ -4667,6 +4680,7 @@ func (s *GatewayService) forwardAnthropicAPIKeyPassthroughWithInput(
|
||||
AccountName: account.Name,
|
||||
UpstreamStatusCode: resp.StatusCode,
|
||||
UpstreamRequestID: resp.Header.Get("x-request-id"),
|
||||
UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()),
|
||||
Passthrough: true,
|
||||
Kind: "retry",
|
||||
Message: extractUpstreamErrorMessage(respBody),
|
||||
@@ -5344,6 +5358,7 @@ func (s *GatewayService) executeBedrockUpstream(
|
||||
AccountID: account.ID,
|
||||
AccountName: account.Name,
|
||||
UpstreamStatusCode: 0,
|
||||
UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()),
|
||||
Kind: "request_error",
|
||||
Message: safeErr,
|
||||
})
|
||||
@@ -5380,6 +5395,7 @@ func (s *GatewayService) executeBedrockUpstream(
|
||||
AccountID: account.ID,
|
||||
AccountName: account.Name,
|
||||
UpstreamStatusCode: resp.StatusCode,
|
||||
UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()),
|
||||
Kind: "retry",
|
||||
Message: extractUpstreamErrorMessage(respBody),
|
||||
Detail: func() string {
|
||||
@@ -7877,6 +7893,9 @@ func (s *GatewayService) ForwardCountTokens(ctx context.Context, c *gin.Context,
|
||||
body := parsed.Body
|
||||
reqModel := parsed.Model
|
||||
|
||||
// Pre-filter: strip empty text blocks to prevent upstream 400.
|
||||
body = StripEmptyTextBlocks(body)
|
||||
|
||||
isClaudeCode := isClaudeCodeRequest(ctx, c, parsed)
|
||||
shouldMimicClaudeCode := account.IsOAuth() && !isClaudeCode
|
||||
|
||||
@@ -8064,6 +8083,7 @@ func (s *GatewayService) forwardCountTokensAnthropicAPIKeyPassthrough(ctx contex
|
||||
AccountID: account.ID,
|
||||
AccountName: account.Name,
|
||||
UpstreamStatusCode: 0,
|
||||
UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()),
|
||||
Passthrough: true,
|
||||
Kind: "request_error",
|
||||
Message: sanitizeUpstreamErrorMessage(err.Error()),
|
||||
@@ -8119,6 +8139,7 @@ func (s *GatewayService) forwardCountTokensAnthropicAPIKeyPassthrough(ctx contex
|
||||
AccountName: account.Name,
|
||||
UpstreamStatusCode: resp.StatusCode,
|
||||
UpstreamRequestID: resp.Header.Get("x-request-id"),
|
||||
UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()),
|
||||
Passthrough: true,
|
||||
Kind: "http_error",
|
||||
Message: upstreamMsg,
|
||||
|
||||
@@ -79,7 +79,7 @@ func (m *mockAccountRepoForGemini) Delete(ctx context.Context, id int64) error
|
||||
func (m *mockAccountRepoForGemini) List(ctx context.Context, params pagination.PaginationParams) ([]Account, *pagination.PaginationResult, error) {
|
||||
return nil, nil, nil
|
||||
}
|
||||
func (m *mockAccountRepoForGemini) ListWithFilters(ctx context.Context, params pagination.PaginationParams, platform, accountType, status, search string, groupID int64) ([]Account, *pagination.PaginationResult, error) {
|
||||
func (m *mockAccountRepoForGemini) ListWithFilters(ctx context.Context, params pagination.PaginationParams, platform, accountType, status, search string, groupID int64, privacyMode string) ([]Account, *pagination.PaginationResult, error) {
|
||||
return nil, nil, nil
|
||||
}
|
||||
func (m *mockAccountRepoForGemini) ListByGroup(ctx context.Context, groupID int64) ([]Account, error) {
|
||||
|
||||
@@ -52,10 +52,11 @@ func BuildGeminiDigestChain(req *antigravity.GeminiRequest) string {
|
||||
// 返回 16 字符的 Base64 编码的 SHA256 前缀
|
||||
func GenerateGeminiPrefixHash(userID, apiKeyID int64, ip, userAgent, platform, model string) string {
|
||||
// 组合所有标识符
|
||||
normalizedUserAgent := NormalizeSessionUserAgent(userAgent)
|
||||
combined := strconv.FormatInt(userID, 10) + ":" +
|
||||
strconv.FormatInt(apiKeyID, 10) + ":" +
|
||||
ip + ":" +
|
||||
userAgent + ":" +
|
||||
normalizedUserAgent + ":" +
|
||||
platform + ":" +
|
||||
model
|
||||
|
||||
|
||||
@@ -152,6 +152,24 @@ func TestGenerateGeminiPrefixHash(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestGenerateGeminiPrefixHash_IgnoresUserAgentVersionNoise(t *testing.T) {
|
||||
hash1 := GenerateGeminiPrefixHash(1, 100, "192.168.1.1", "Mozilla/5.0 codex_cli_rs/0.1.0", "antigravity", "gemini-2.5-pro")
|
||||
hash2 := GenerateGeminiPrefixHash(1, 100, "192.168.1.1", "Mozilla/5.0 codex_cli_rs/0.1.1", "antigravity", "gemini-2.5-pro")
|
||||
|
||||
if hash1 != hash2 {
|
||||
t.Fatalf("version-only User-Agent changes should not perturb Gemini prefix hash: %s vs %s", hash1, hash2)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGenerateGeminiPrefixHash_IgnoresFreeformUserAgentVersionNoise(t *testing.T) {
|
||||
hash1 := GenerateGeminiPrefixHash(1, 100, "192.168.1.1", "Codex CLI 0.1.0", "antigravity", "gemini-2.5-pro")
|
||||
hash2 := GenerateGeminiPrefixHash(1, 100, "192.168.1.1", "Codex CLI 0.1.1", "antigravity", "gemini-2.5-pro")
|
||||
|
||||
if hash1 != hash2 {
|
||||
t.Fatalf("free-form version-only User-Agent changes should not perturb Gemini prefix hash: %s vs %s", hash1, hash2)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseGeminiSessionValue(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
|
||||
@@ -135,7 +135,7 @@ func (p *GeminiTokenProvider) GetAccessToken(ctx context.Context, account *Accou
|
||||
if tierID != "" {
|
||||
account.Credentials["tier_id"] = tierID
|
||||
}
|
||||
_ = p.accountRepo.Update(ctx, account)
|
||||
_ = persistAccountCredentials(ctx, p.accountRepo, account, account.Credentials)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -504,6 +504,48 @@ func TestGenerateSessionHash_SessionContext_UADifference(t *testing.T) {
|
||||
require.NotEqual(t, h1, h2, "different User-Agent should produce different hash")
|
||||
}
|
||||
|
||||
func TestGenerateSessionHash_SessionContext_UAVersionNoiseIgnored(t *testing.T) {
|
||||
svc := &GatewayService{}
|
||||
|
||||
base := func(ua string) *ParsedRequest {
|
||||
return &ParsedRequest{
|
||||
Messages: []any{
|
||||
map[string]any{"role": "user", "content": "test"},
|
||||
},
|
||||
SessionContext: &SessionContext{
|
||||
ClientIP: "1.1.1.1",
|
||||
UserAgent: ua,
|
||||
APIKeyID: 1,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
h1 := svc.GenerateSessionHash(base("Mozilla/5.0 codex_cli_rs/0.1.0"))
|
||||
h2 := svc.GenerateSessionHash(base("Mozilla/5.0 codex_cli_rs/0.1.1"))
|
||||
require.Equal(t, h1, h2, "version-only User-Agent changes should not perturb the sticky session hash")
|
||||
}
|
||||
|
||||
func TestGenerateSessionHash_SessionContext_FreeformUAVersionNoiseIgnored(t *testing.T) {
|
||||
svc := &GatewayService{}
|
||||
|
||||
base := func(ua string) *ParsedRequest {
|
||||
return &ParsedRequest{
|
||||
Messages: []any{
|
||||
map[string]any{"role": "user", "content": "test"},
|
||||
},
|
||||
SessionContext: &SessionContext{
|
||||
ClientIP: "1.1.1.1",
|
||||
UserAgent: ua,
|
||||
APIKeyID: 1,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
h1 := svc.GenerateSessionHash(base("Codex CLI 0.1.0"))
|
||||
h2 := svc.GenerateSessionHash(base("Codex CLI 0.1.1"))
|
||||
require.Equal(t, h1, h2, "free-form version-only User-Agent changes should not perturb the sticky session hash")
|
||||
}
|
||||
|
||||
func TestGenerateSessionHash_SessionContext_APIKeyIDDifference(t *testing.T) {
|
||||
svc := &GatewayService{}
|
||||
|
||||
|
||||
@@ -108,8 +108,7 @@ func (api *OAuthRefreshAPI) RefreshIfNeeded(
|
||||
// 5. 设置版本号 + 更新 DB
|
||||
if newCredentials != nil {
|
||||
newCredentials["_token_version"] = time.Now().UnixMilli()
|
||||
freshAccount.Credentials = newCredentials
|
||||
if updateErr := api.accountRepo.Update(ctx, freshAccount); updateErr != nil {
|
||||
if updateErr := persistAccountCredentials(ctx, api.accountRepo, freshAccount, newCredentials); updateErr != nil {
|
||||
slog.Error("oauth_refresh_update_failed",
|
||||
"account_id", freshAccount.ID,
|
||||
"error", updateErr,
|
||||
|
||||
@@ -16,10 +16,11 @@ import (
|
||||
// refreshAPIAccountRepo implements AccountRepository for OAuthRefreshAPI tests.
|
||||
type refreshAPIAccountRepo struct {
|
||||
mockAccountRepoForGemini
|
||||
account *Account // returned by GetByID
|
||||
getByIDErr error
|
||||
updateErr error
|
||||
updateCalls int
|
||||
account *Account // returned by GetByID
|
||||
getByIDErr error
|
||||
updateErr error
|
||||
updateCalls int
|
||||
updateCredentialsCalls int
|
||||
}
|
||||
|
||||
func (r *refreshAPIAccountRepo) GetByID(_ context.Context, _ int64) (*Account, error) {
|
||||
@@ -34,6 +35,19 @@ func (r *refreshAPIAccountRepo) Update(_ context.Context, _ *Account) error {
|
||||
return r.updateErr
|
||||
}
|
||||
|
||||
func (r *refreshAPIAccountRepo) UpdateCredentials(_ context.Context, id int64, credentials map[string]any) error {
|
||||
r.updateCalls++
|
||||
r.updateCredentialsCalls++
|
||||
if r.updateErr != nil {
|
||||
return r.updateErr
|
||||
}
|
||||
if r.account == nil || r.account.ID != id {
|
||||
r.account = &Account{ID: id}
|
||||
}
|
||||
r.account.Credentials = cloneCredentials(credentials)
|
||||
return nil
|
||||
}
|
||||
|
||||
// refreshAPIExecutorStub implements OAuthRefreshExecutor for tests.
|
||||
type refreshAPIExecutorStub struct {
|
||||
needsRefresh bool
|
||||
@@ -106,10 +120,36 @@ func TestRefreshIfNeeded_Success(t *testing.T) {
|
||||
require.Equal(t, "new-token", result.NewCredentials["access_token"])
|
||||
require.NotNil(t, result.NewCredentials["_token_version"]) // version stamp set
|
||||
require.Equal(t, 1, repo.updateCalls) // DB updated
|
||||
require.Equal(t, 1, cache.releaseCalls) // lock released
|
||||
require.Equal(t, 1, repo.updateCredentialsCalls)
|
||||
require.Equal(t, 1, cache.releaseCalls) // lock released
|
||||
require.Equal(t, 1, executor.refreshCalls)
|
||||
}
|
||||
|
||||
func TestRefreshIfNeeded_UpdateCredentialsPreservesRateLimitState(t *testing.T) {
|
||||
resetAt := time.Now().Add(45 * time.Minute)
|
||||
account := &Account{
|
||||
ID: 11,
|
||||
Platform: PlatformGemini,
|
||||
Type: AccountTypeOAuth,
|
||||
RateLimitResetAt: &resetAt,
|
||||
}
|
||||
repo := &refreshAPIAccountRepo{account: account}
|
||||
cache := &refreshAPICacheStub{lockResult: true}
|
||||
executor := &refreshAPIExecutorStub{
|
||||
needsRefresh: true,
|
||||
credentials: map[string]any{"access_token": "safe-token"},
|
||||
}
|
||||
|
||||
api := NewOAuthRefreshAPI(repo, cache)
|
||||
result, err := api.RefreshIfNeeded(context.Background(), account, executor, 3*time.Minute)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.True(t, result.Refreshed)
|
||||
require.Equal(t, 1, repo.updateCredentialsCalls)
|
||||
require.NotNil(t, repo.account.RateLimitResetAt)
|
||||
require.WithinDuration(t, resetAt, *repo.account.RateLimitResetAt, time.Second)
|
||||
}
|
||||
|
||||
func TestRefreshIfNeeded_LockHeld(t *testing.T) {
|
||||
account := &Account{ID: 2, Platform: PlatformAnthropic}
|
||||
repo := &refreshAPIAccountRepo{account: account}
|
||||
@@ -193,7 +233,7 @@ func TestRefreshIfNeeded_RefreshError(t *testing.T) {
|
||||
require.Error(t, err)
|
||||
require.Nil(t, result)
|
||||
require.Contains(t, err.Error(), "invalid_grant")
|
||||
require.Equal(t, 0, repo.updateCalls) // no DB update on refresh error
|
||||
require.Equal(t, 0, repo.updateCalls) // no DB update on refresh error
|
||||
require.Equal(t, 1, cache.releaseCalls) // lock still released via defer
|
||||
}
|
||||
|
||||
@@ -299,8 +339,8 @@ func TestMergeCredentials_NewOverridesOld(t *testing.T) {
|
||||
|
||||
result := MergeCredentials(old, new)
|
||||
|
||||
require.Equal(t, "new-token", result["access_token"]) // overridden
|
||||
require.Equal(t, "old-refresh", result["refresh_token"]) // preserved
|
||||
require.Equal(t, "new-token", result["access_token"]) // overridden
|
||||
require.Equal(t, "old-refresh", result["refresh_token"]) // preserved
|
||||
}
|
||||
|
||||
// ========== BuildClaudeAccountCredentials tests ==========
|
||||
|
||||
@@ -330,6 +330,11 @@ func (s *defaultOpenAIAccountScheduler) selectBySessionHash(
|
||||
_ = s.service.deleteStickySessionAccountID(ctx, req.GroupID, sessionHash)
|
||||
return nil, nil
|
||||
}
|
||||
account = s.service.recheckSelectedOpenAIAccountFromDB(ctx, account, req.RequestedModel)
|
||||
if account == nil {
|
||||
_ = s.service.deleteStickySessionAccountID(ctx, req.GroupID, sessionHash)
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
result, acquireErr := s.service.tryAcquireAccountSlot(ctx, accountID, account.Concurrency)
|
||||
if acquireErr == nil && result.Acquired {
|
||||
@@ -691,6 +696,10 @@ func (s *defaultOpenAIAccountScheduler) selectByLoadBalance(
|
||||
if fresh == nil || !s.isAccountTransportCompatible(fresh, req.RequiredTransport) {
|
||||
continue
|
||||
}
|
||||
fresh = s.service.recheckSelectedOpenAIAccountFromDB(ctx, fresh, req.RequestedModel)
|
||||
if fresh == nil || !s.isAccountTransportCompatible(fresh, req.RequiredTransport) {
|
||||
continue
|
||||
}
|
||||
result, acquireErr := s.service.tryAcquireAccountSlot(ctx, fresh.ID, fresh.Concurrency)
|
||||
if acquireErr != nil {
|
||||
return nil, len(candidates), topK, loadSkew, acquireErr
|
||||
|
||||
@@ -84,6 +84,61 @@ func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_SkipsFreshlyRa
|
||||
require.Equal(t, int64(32002), account.ID)
|
||||
}
|
||||
|
||||
func TestOpenAIGatewayService_SelectAccountWithScheduler_SessionStickyDBRuntimeRecheckSkipsStaleCachedAccount(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
groupID := int64(10103)
|
||||
rateLimitedUntil := time.Now().Add(30 * time.Minute)
|
||||
staleSticky := &Account{ID: 33001, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0}
|
||||
staleBackup := &Account{ID: 33002, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5}
|
||||
dbSticky := Account{ID: 33001, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, RateLimitResetAt: &rateLimitedUntil}
|
||||
dbBackup := Account{ID: 33002, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5}
|
||||
cache := &stubGatewayCache{sessionBindings: map[string]int64{"openai:session_hash_db_runtime_recheck": 33001}}
|
||||
snapshotCache := &openAISnapshotCacheStub{
|
||||
snapshotAccounts: []*Account{staleSticky, staleBackup},
|
||||
accountsByID: map[int64]*Account{33001: staleSticky, 33002: staleBackup},
|
||||
}
|
||||
snapshotService := &SchedulerSnapshotService{cache: snapshotCache}
|
||||
svc := &OpenAIGatewayService{
|
||||
accountRepo: stubOpenAIAccountRepo{accounts: []Account{dbSticky, dbBackup}},
|
||||
cache: cache,
|
||||
cfg: &config.Config{},
|
||||
schedulerSnapshot: snapshotService,
|
||||
concurrencyService: NewConcurrencyService(stubConcurrencyCache{}),
|
||||
}
|
||||
|
||||
selection, decision, err := svc.SelectAccountWithScheduler(ctx, &groupID, "", "session_hash_db_runtime_recheck", "gpt-5.1", nil, OpenAIUpstreamTransportAny)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, selection)
|
||||
require.NotNil(t, selection.Account)
|
||||
require.Equal(t, int64(33002), selection.Account.ID)
|
||||
require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
|
||||
}
|
||||
|
||||
func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_DBRuntimeRecheckSkipsStaleCachedCandidate(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
groupID := int64(10104)
|
||||
rateLimitedUntil := time.Now().Add(30 * time.Minute)
|
||||
stalePrimary := &Account{ID: 34001, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0}
|
||||
staleSecondary := &Account{ID: 34002, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5}
|
||||
dbPrimary := Account{ID: 34001, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, RateLimitResetAt: &rateLimitedUntil}
|
||||
dbSecondary := Account{ID: 34002, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5}
|
||||
snapshotCache := &openAISnapshotCacheStub{
|
||||
snapshotAccounts: []*Account{stalePrimary, staleSecondary},
|
||||
accountsByID: map[int64]*Account{34001: stalePrimary, 34002: staleSecondary},
|
||||
}
|
||||
snapshotService := &SchedulerSnapshotService{cache: snapshotCache}
|
||||
svc := &OpenAIGatewayService{
|
||||
accountRepo: stubOpenAIAccountRepo{accounts: []Account{dbPrimary, dbSecondary}},
|
||||
cfg: &config.Config{},
|
||||
schedulerSnapshot: snapshotService,
|
||||
}
|
||||
|
||||
account, err := svc.SelectAccountForModelWithExclusions(ctx, &groupID, "", "gpt-5.1", nil)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, account)
|
||||
require.Equal(t, int64(34002), account.ID)
|
||||
}
|
||||
|
||||
func TestOpenAIGatewayService_SelectAccountWithScheduler_PreviousResponseSticky(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
groupID := int64(9)
|
||||
|
||||
@@ -1201,6 +1201,11 @@ func (s *OpenAIGatewayService) tryStickySessionHit(ctx context.Context, groupID
|
||||
if requestedModel != "" && !account.IsModelSupported(requestedModel) {
|
||||
return nil
|
||||
}
|
||||
account = s.recheckSelectedOpenAIAccountFromDB(ctx, account, requestedModel)
|
||||
if account == nil {
|
||||
_ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash)
|
||||
return nil
|
||||
}
|
||||
|
||||
// 刷新会话 TTL 并返回账号
|
||||
// Refresh session TTL and return account
|
||||
@@ -1229,6 +1234,10 @@ func (s *OpenAIGatewayService) selectBestAccount(ctx context.Context, accounts [
|
||||
if fresh == nil {
|
||||
continue
|
||||
}
|
||||
fresh = s.recheckSelectedOpenAIAccountFromDB(ctx, fresh, requestedModel)
|
||||
if fresh == nil {
|
||||
continue
|
||||
}
|
||||
|
||||
// 选择优先级最高且最久未使用的账号
|
||||
// Select highest priority and least recently used
|
||||
@@ -1353,27 +1362,32 @@ func (s *OpenAIGatewayService) SelectAccountWithLoadAwareness(ctx context.Contex
|
||||
}
|
||||
if !clearSticky && account.IsSchedulable() && account.IsOpenAI() &&
|
||||
(requestedModel == "" || account.IsModelSupported(requestedModel)) {
|
||||
result, err := s.tryAcquireAccountSlot(ctx, accountID, account.Concurrency)
|
||||
if err == nil && result.Acquired {
|
||||
_ = s.refreshStickySessionTTL(ctx, groupID, sessionHash, openaiStickySessionTTL)
|
||||
return &AccountSelectionResult{
|
||||
Account: account,
|
||||
Acquired: true,
|
||||
ReleaseFunc: result.ReleaseFunc,
|
||||
}, nil
|
||||
}
|
||||
account = s.recheckSelectedOpenAIAccountFromDB(ctx, account, requestedModel)
|
||||
if account == nil {
|
||||
_ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash)
|
||||
} else {
|
||||
result, err := s.tryAcquireAccountSlot(ctx, accountID, account.Concurrency)
|
||||
if err == nil && result.Acquired {
|
||||
_ = s.refreshStickySessionTTL(ctx, groupID, sessionHash, openaiStickySessionTTL)
|
||||
return &AccountSelectionResult{
|
||||
Account: account,
|
||||
Acquired: true,
|
||||
ReleaseFunc: result.ReleaseFunc,
|
||||
}, nil
|
||||
}
|
||||
|
||||
waitingCount, _ := s.concurrencyService.GetAccountWaitingCount(ctx, accountID)
|
||||
if waitingCount < cfg.StickySessionMaxWaiting {
|
||||
return &AccountSelectionResult{
|
||||
Account: account,
|
||||
WaitPlan: &AccountWaitPlan{
|
||||
AccountID: accountID,
|
||||
MaxConcurrency: account.Concurrency,
|
||||
Timeout: cfg.StickySessionWaitTimeout,
|
||||
MaxWaiting: cfg.StickySessionMaxWaiting,
|
||||
},
|
||||
}, nil
|
||||
waitingCount, _ := s.concurrencyService.GetAccountWaitingCount(ctx, accountID)
|
||||
if waitingCount < cfg.StickySessionMaxWaiting {
|
||||
return &AccountSelectionResult{
|
||||
Account: account,
|
||||
WaitPlan: &AccountWaitPlan{
|
||||
AccountID: accountID,
|
||||
MaxConcurrency: account.Concurrency,
|
||||
Timeout: cfg.StickySessionWaitTimeout,
|
||||
MaxWaiting: cfg.StickySessionMaxWaiting,
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1560,6 +1574,28 @@ func (s *OpenAIGatewayService) resolveFreshSchedulableOpenAIAccount(ctx context.
|
||||
return fresh
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) recheckSelectedOpenAIAccountFromDB(ctx context.Context, account *Account, requestedModel string) *Account {
|
||||
if account == nil {
|
||||
return nil
|
||||
}
|
||||
if s.schedulerSnapshot == nil || s.accountRepo == nil {
|
||||
return account
|
||||
}
|
||||
|
||||
latest, err := s.accountRepo.GetByID(ctx, account.ID)
|
||||
if err != nil || latest == nil {
|
||||
return nil
|
||||
}
|
||||
syncOpenAICodexRateLimitFromExtra(ctx, s.accountRepo, latest, time.Now())
|
||||
if !latest.IsSchedulable() || !latest.IsOpenAI() {
|
||||
return nil
|
||||
}
|
||||
if requestedModel != "" && !latest.IsModelSupported(requestedModel) {
|
||||
return nil
|
||||
}
|
||||
return latest
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) getSchedulableAccount(ctx context.Context, accountID int64) (*Account, error) {
|
||||
var (
|
||||
account *Account
|
||||
@@ -2598,6 +2634,12 @@ func (s *OpenAIGatewayService) handleErrorResponsePassthrough(
|
||||
}
|
||||
setOpsUpstreamError(c, resp.StatusCode, upstreamMsg, upstreamDetail)
|
||||
logOpenAIInstructionsRequiredDebug(ctx, c, account, resp.StatusCode, upstreamMsg, requestBody, body)
|
||||
if s.rateLimitService != nil {
|
||||
// Passthrough mode preserves the raw upstream error response, but runtime
|
||||
// account state still needs to be updated so sticky routing can stop
|
||||
// reusing a freshly rate-limited account.
|
||||
_ = s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, body)
|
||||
}
|
||||
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
|
||||
Platform: account.Platform,
|
||||
AccountID: account.ID,
|
||||
|
||||
@@ -536,6 +536,55 @@ func TestOpenAIGatewayService_OAuthPassthrough_UpstreamErrorIncludesPassthroughF
|
||||
require.True(t, arr[len(arr)-1].Passthrough)
|
||||
}
|
||||
|
||||
func TestOpenAIGatewayService_OAuthPassthrough_429PersistsRateLimit(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(nil))
|
||||
c.Request.Header.Set("User-Agent", "codex_cli_rs/0.1.0")
|
||||
|
||||
originalBody := []byte(`{"model":"gpt-5.2","stream":false,"instructions":"local-test-instructions","input":[{"type":"text","text":"hi"}]}`)
|
||||
resetAt := time.Now().Add(7 * 24 * time.Hour).Unix()
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusTooManyRequests,
|
||||
Header: http.Header{
|
||||
"Content-Type": []string{"application/json"},
|
||||
"x-request-id": []string{"rid-rate-limit"},
|
||||
},
|
||||
Body: io.NopCloser(strings.NewReader(fmt.Sprintf(`{"error":{"message":"The usage limit has been reached","type":"usage_limit_reached","resets_at":%d}}`, resetAt))),
|
||||
}
|
||||
upstream := &httpUpstreamRecorder{resp: resp}
|
||||
repo := &openAIWSRateLimitSignalRepo{}
|
||||
rateSvc := &RateLimitService{accountRepo: repo}
|
||||
|
||||
svc := &OpenAIGatewayService{
|
||||
cfg: &config.Config{Gateway: config.GatewayConfig{ForceCodexCLI: false}},
|
||||
httpUpstream: upstream,
|
||||
rateLimitService: rateSvc,
|
||||
}
|
||||
|
||||
account := &Account{
|
||||
ID: 123,
|
||||
Name: "acc",
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeOAuth,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{"access_token": "oauth-token", "chatgpt_account_id": "chatgpt-acc"},
|
||||
Extra: map[string]any{"openai_passthrough": true},
|
||||
Status: StatusActive,
|
||||
Schedulable: true,
|
||||
RateMultiplier: f64p(1),
|
||||
}
|
||||
|
||||
_, err := svc.Forward(context.Background(), c, account, originalBody)
|
||||
require.Error(t, err)
|
||||
require.Equal(t, http.StatusTooManyRequests, rec.Code)
|
||||
require.Contains(t, rec.Body.String(), "usage_limit_reached")
|
||||
require.Len(t, repo.rateLimitCalls, 1)
|
||||
require.WithinDuration(t, time.Unix(resetAt, 0), repo.rateLimitCalls[0], 2*time.Second)
|
||||
}
|
||||
|
||||
func TestOpenAIGatewayService_OAuthPassthrough_NonCodexUAFallbackToCodexUA(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
|
||||
@@ -29,9 +29,10 @@ type soraSessionChunk struct {
|
||||
|
||||
// OpenAIOAuthService handles OpenAI OAuth authentication flows
|
||||
type OpenAIOAuthService struct {
|
||||
sessionStore *openai.SessionStore
|
||||
proxyRepo ProxyRepository
|
||||
oauthClient OpenAIOAuthClient
|
||||
sessionStore *openai.SessionStore
|
||||
proxyRepo ProxyRepository
|
||||
oauthClient OpenAIOAuthClient
|
||||
privacyClientFactory PrivacyClientFactory // 用于调用 chatgpt.com/backend-api(ImpersonateChrome)
|
||||
}
|
||||
|
||||
// NewOpenAIOAuthService creates a new OpenAI OAuth service
|
||||
@@ -43,6 +44,12 @@ func NewOpenAIOAuthService(proxyRepo ProxyRepository, oauthClient OpenAIOAuthCli
|
||||
}
|
||||
}
|
||||
|
||||
// SetPrivacyClientFactory 注入 ImpersonateChrome 客户端工厂,
|
||||
// 用于调用 chatgpt.com/backend-api 获取账号信息(plan_type 等)。
|
||||
func (s *OpenAIOAuthService) SetPrivacyClientFactory(factory PrivacyClientFactory) {
|
||||
s.privacyClientFactory = factory
|
||||
}
|
||||
|
||||
// OpenAIAuthURLResult contains the authorization URL and session info
|
||||
type OpenAIAuthURLResult struct {
|
||||
AuthURL string `json:"auth_url"`
|
||||
@@ -131,6 +138,7 @@ type OpenAITokenInfo struct {
|
||||
ChatGPTUserID string `json:"chatgpt_user_id,omitempty"`
|
||||
OrganizationID string `json:"organization_id,omitempty"`
|
||||
PlanType string `json:"plan_type,omitempty"`
|
||||
PrivacyMode string `json:"privacy_mode,omitempty"`
|
||||
}
|
||||
|
||||
// ExchangeCode exchanges authorization code for tokens
|
||||
@@ -251,6 +259,30 @@ func (s *OpenAIOAuthService) RefreshTokenWithClientID(ctx context.Context, refre
|
||||
tokenInfo.PlanType = userInfo.PlanType
|
||||
}
|
||||
|
||||
// id_token 中缺少 plan_type 时(如 Mobile RT),尝试通过 ChatGPT backend-api 补全
|
||||
if tokenInfo.PlanType == "" && tokenInfo.AccessToken != "" && s.privacyClientFactory != nil {
|
||||
// 从 access_token JWT 中提取 orgID(poid),用于匹配正确的账号
|
||||
orgID := tokenInfo.OrganizationID
|
||||
if orgID == "" {
|
||||
if atClaims, err := openai.DecodeIDToken(tokenInfo.AccessToken); err == nil && atClaims.OpenAIAuth != nil {
|
||||
orgID = atClaims.OpenAIAuth.POID
|
||||
}
|
||||
}
|
||||
if info := fetchChatGPTAccountInfo(ctx, s.privacyClientFactory, tokenInfo.AccessToken, proxyURL, orgID); info != nil {
|
||||
if tokenInfo.PlanType == "" && info.PlanType != "" {
|
||||
tokenInfo.PlanType = info.PlanType
|
||||
}
|
||||
if tokenInfo.Email == "" && info.Email != "" {
|
||||
tokenInfo.Email = info.Email
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 尝试设置隐私(关闭训练数据共享),best-effort
|
||||
if tokenInfo.AccessToken != "" && s.privacyClientFactory != nil {
|
||||
tokenInfo.PrivacyMode = disableOpenAITraining(ctx, s.privacyClientFactory, tokenInfo.AccessToken, proxyURL)
|
||||
}
|
||||
|
||||
return tokenInfo, nil
|
||||
}
|
||||
|
||||
|
||||
@@ -69,6 +69,139 @@ func disableOpenAITraining(ctx context.Context, clientFactory PrivacyClientFacto
|
||||
return PrivacyModeTrainingOff
|
||||
}
|
||||
|
||||
// ChatGPTAccountInfo 从 chatgpt.com/backend-api/accounts/check 获取的账号信息
|
||||
type ChatGPTAccountInfo struct {
|
||||
PlanType string
|
||||
Email string
|
||||
}
|
||||
|
||||
const chatGPTAccountsCheckURL = "https://chatgpt.com/backend-api/accounts/check/v4-2023-04-27"
|
||||
|
||||
// fetchChatGPTAccountInfo calls ChatGPT backend-api to get account info (plan_type, etc.).
|
||||
// Used as fallback when id_token doesn't contain these fields (e.g., Mobile RT).
|
||||
// orgID is used to match the correct account when multiple accounts exist (e.g., personal + team).
|
||||
// Returns nil on any failure (best-effort, non-blocking).
|
||||
func fetchChatGPTAccountInfo(ctx context.Context, clientFactory PrivacyClientFactory, accessToken, proxyURL, orgID string) *ChatGPTAccountInfo {
|
||||
if accessToken == "" || clientFactory == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(ctx, 15*time.Second)
|
||||
defer cancel()
|
||||
|
||||
client, err := clientFactory(proxyURL)
|
||||
if err != nil {
|
||||
slog.Debug("chatgpt_account_check_client_error", "error", err.Error())
|
||||
return nil
|
||||
}
|
||||
|
||||
var result map[string]any
|
||||
resp, err := client.R().
|
||||
SetContext(ctx).
|
||||
SetHeader("Authorization", "Bearer "+accessToken).
|
||||
SetHeader("Origin", "https://chatgpt.com").
|
||||
SetHeader("Referer", "https://chatgpt.com/").
|
||||
SetHeader("Accept", "application/json").
|
||||
SetSuccessResult(&result).
|
||||
Get(chatGPTAccountsCheckURL)
|
||||
|
||||
if err != nil {
|
||||
slog.Debug("chatgpt_account_check_request_error", "error", err.Error())
|
||||
return nil
|
||||
}
|
||||
|
||||
if !resp.IsSuccessState() {
|
||||
slog.Debug("chatgpt_account_check_failed", "status", resp.StatusCode, "body", truncate(resp.String(), 200))
|
||||
return nil
|
||||
}
|
||||
|
||||
info := &ChatGPTAccountInfo{}
|
||||
|
||||
accounts, ok := result["accounts"].(map[string]any)
|
||||
if !ok {
|
||||
slog.Debug("chatgpt_account_check_no_accounts", "body", truncate(resp.String(), 300))
|
||||
return nil
|
||||
}
|
||||
|
||||
// 优先匹配 orgID 对应的账号(access_token JWT 中的 poid)
|
||||
if orgID != "" {
|
||||
if matched := extractPlanFromAccount(accounts, orgID); matched != "" {
|
||||
info.PlanType = matched
|
||||
}
|
||||
}
|
||||
|
||||
// 未匹配到时,遍历所有账号:优先 is_default,次选非 free
|
||||
if info.PlanType == "" {
|
||||
var defaultPlan, paidPlan, anyPlan string
|
||||
for _, acctRaw := range accounts {
|
||||
acct, ok := acctRaw.(map[string]any)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
planType := extractPlanType(acct)
|
||||
if planType == "" {
|
||||
continue
|
||||
}
|
||||
if anyPlan == "" {
|
||||
anyPlan = planType
|
||||
}
|
||||
if account, ok := acct["account"].(map[string]any); ok {
|
||||
if isDefault, _ := account["is_default"].(bool); isDefault {
|
||||
defaultPlan = planType
|
||||
}
|
||||
}
|
||||
if !strings.EqualFold(planType, "free") && paidPlan == "" {
|
||||
paidPlan = planType
|
||||
}
|
||||
}
|
||||
// 优先级:default > 非 free > 任意
|
||||
switch {
|
||||
case defaultPlan != "":
|
||||
info.PlanType = defaultPlan
|
||||
case paidPlan != "":
|
||||
info.PlanType = paidPlan
|
||||
default:
|
||||
info.PlanType = anyPlan
|
||||
}
|
||||
}
|
||||
|
||||
if info.PlanType == "" {
|
||||
slog.Debug("chatgpt_account_check_no_plan_type", "body", truncate(resp.String(), 300))
|
||||
return nil
|
||||
}
|
||||
|
||||
slog.Info("chatgpt_account_check_success", "plan_type", info.PlanType, "org_id", orgID)
|
||||
return info
|
||||
}
|
||||
|
||||
// extractPlanFromAccount 从 accounts map 中按 key(account_id)精确匹配并提取 plan_type
|
||||
func extractPlanFromAccount(accounts map[string]any, accountKey string) string {
|
||||
acctRaw, ok := accounts[accountKey]
|
||||
if !ok {
|
||||
return ""
|
||||
}
|
||||
acct, ok := acctRaw.(map[string]any)
|
||||
if !ok {
|
||||
return ""
|
||||
}
|
||||
return extractPlanType(acct)
|
||||
}
|
||||
|
||||
// extractPlanType 从单个 account 对象中提取 plan_type
|
||||
func extractPlanType(acct map[string]any) string {
|
||||
if account, ok := acct["account"].(map[string]any); ok {
|
||||
if planType, ok := account["plan_type"].(string); ok && planType != "" {
|
||||
return planType
|
||||
}
|
||||
}
|
||||
if entitlement, ok := acct["entitlement"].(map[string]any); ok {
|
||||
if subPlan, ok := entitlement["subscription_plan"].(string); ok && subPlan != "" {
|
||||
return subPlan
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func truncate(s string, n int) string {
|
||||
if len(s) <= n {
|
||||
return s
|
||||
|
||||
@@ -85,6 +85,58 @@ func TestOpenAIGatewayService_SelectAccountByPreviousResponseID_RateLimitedMiss(
|
||||
require.Zero(t, boundAccountID)
|
||||
}
|
||||
|
||||
func TestOpenAIGatewayService_SelectAccountByPreviousResponseID_DBRuntimeRecheckRateLimitedMiss(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
groupID := int64(24)
|
||||
rateLimitedUntil := time.Now().Add(30 * time.Minute)
|
||||
staleAccount := &Account{
|
||||
ID: 13,
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeAPIKey,
|
||||
Status: StatusActive,
|
||||
Schedulable: true,
|
||||
Concurrency: 1,
|
||||
Extra: map[string]any{
|
||||
"openai_apikey_responses_websockets_v2_enabled": true,
|
||||
},
|
||||
}
|
||||
dbAccount := Account{
|
||||
ID: 13,
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeAPIKey,
|
||||
Status: StatusActive,
|
||||
Schedulable: true,
|
||||
Concurrency: 1,
|
||||
RateLimitResetAt: &rateLimitedUntil,
|
||||
Extra: map[string]any{
|
||||
"openai_apikey_responses_websockets_v2_enabled": true,
|
||||
},
|
||||
}
|
||||
cache := &stubGatewayCache{}
|
||||
store := NewOpenAIWSStateStore(cache)
|
||||
cfg := newOpenAIWSV2TestConfig()
|
||||
snapshotCache := &openAISnapshotCacheStub{
|
||||
accountsByID: map[int64]*Account{dbAccount.ID: staleAccount},
|
||||
}
|
||||
svc := &OpenAIGatewayService{
|
||||
accountRepo: stubOpenAIAccountRepo{accounts: []Account{dbAccount}},
|
||||
cache: cache,
|
||||
cfg: cfg,
|
||||
concurrencyService: NewConcurrencyService(stubConcurrencyCache{}),
|
||||
openaiWSStateStore: store,
|
||||
schedulerSnapshot: &SchedulerSnapshotService{cache: snapshotCache},
|
||||
}
|
||||
|
||||
require.NoError(t, store.BindResponseAccount(ctx, groupID, "resp_prev_db_rl", dbAccount.ID, time.Hour))
|
||||
|
||||
selection, err := svc.SelectAccountByPreviousResponseID(ctx, &groupID, "resp_prev_db_rl", "gpt-5.1", nil)
|
||||
require.NoError(t, err)
|
||||
require.Nil(t, selection, "DB 中已限流的账号不应继续命中 previous_response_id 粘连")
|
||||
boundAccountID, getErr := store.GetResponseAccount(ctx, groupID, "resp_prev_db_rl")
|
||||
require.NoError(t, getErr)
|
||||
require.Zero(t, boundAccountID)
|
||||
}
|
||||
|
||||
func TestOpenAIGatewayService_SelectAccountByPreviousResponseID_Excluded(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
groupID := int64(23)
|
||||
|
||||
@@ -3840,6 +3840,11 @@ func (s *OpenAIGatewayService) SelectAccountByPreviousResponseID(
|
||||
if requestedModel != "" && !account.IsModelSupported(requestedModel) {
|
||||
return nil, nil
|
||||
}
|
||||
account = s.recheckSelectedOpenAIAccountFromDB(ctx, account, requestedModel)
|
||||
if account == nil {
|
||||
_ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID)
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
result, acquireErr := s.tryAcquireAccountSlot(ctx, accountID, account.Concurrency)
|
||||
if acquireErr == nil && result.Acquired {
|
||||
|
||||
@@ -73,12 +73,13 @@ func (r *openAICodexExtraListRepo) SetRateLimited(_ context.Context, _ int64, re
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *openAICodexExtraListRepo) ListWithFilters(_ context.Context, params pagination.PaginationParams, platform, accountType, status, search string, groupID int64) ([]Account, *pagination.PaginationResult, error) {
|
||||
func (r *openAICodexExtraListRepo) ListWithFilters(_ context.Context, params pagination.PaginationParams, platform, accountType, status, search string, groupID int64, privacyMode string) ([]Account, *pagination.PaginationResult, error) {
|
||||
_ = platform
|
||||
_ = accountType
|
||||
_ = status
|
||||
_ = search
|
||||
_ = groupID
|
||||
_ = privacyMode
|
||||
return r.accounts, &pagination.PaginationResult{Total: int64(len(r.accounts)), Page: params.Page, PageSize: params.PageSize}, nil
|
||||
}
|
||||
|
||||
@@ -491,7 +492,7 @@ func TestAdminService_ListAccounts_ExhaustedCodexExtraReturnsRateLimitedAccount(
|
||||
}
|
||||
svc := &adminServiceImpl{accountRepo: repo}
|
||||
|
||||
accounts, total, err := svc.ListAccounts(context.Background(), 1, 20, PlatformOpenAI, AccountTypeOAuth, "", "", 0)
|
||||
accounts, total, err := svc.ListAccounts(context.Background(), 1, 20, PlatformOpenAI, AccountTypeOAuth, "", "", 0, "")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(1), total)
|
||||
require.Len(t, accounts, 1)
|
||||
|
||||
@@ -24,7 +24,7 @@ func (s *OpsService) listAllAccountsForOps(ctx context.Context, platformFilter s
|
||||
accounts, pageInfo, err := s.accountRepo.ListWithFilters(ctx, pagination.PaginationParams{
|
||||
Page: page,
|
||||
PageSize: opsAccountsPageSize,
|
||||
}, platformFilter, "", "", "", 0)
|
||||
}, platformFilter, "", "", "", 0, "")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
@@ -62,6 +62,12 @@ type OpsErrorLog struct {
|
||||
ClientIP *string `json:"client_ip"`
|
||||
RequestPath string `json:"request_path"`
|
||||
Stream bool `json:"stream"`
|
||||
|
||||
InboundEndpoint string `json:"inbound_endpoint"`
|
||||
UpstreamEndpoint string `json:"upstream_endpoint"`
|
||||
RequestedModel string `json:"requested_model"`
|
||||
UpstreamModel string `json:"upstream_model"`
|
||||
RequestType *int16 `json:"request_type"`
|
||||
}
|
||||
|
||||
type OpsErrorLogDetail struct {
|
||||
|
||||
@@ -79,6 +79,17 @@ type OpsInsertErrorLogInput struct {
|
||||
Model string
|
||||
RequestPath string
|
||||
Stream bool
|
||||
// InboundEndpoint is the normalized client-facing API endpoint path, e.g. /v1/chat/completions.
|
||||
InboundEndpoint string
|
||||
// UpstreamEndpoint is the normalized upstream endpoint path, e.g. /v1/responses.
|
||||
UpstreamEndpoint string
|
||||
// RequestedModel is the client-requested model name before mapping.
|
||||
RequestedModel string
|
||||
// UpstreamModel is the actual model sent to upstream after mapping. Empty means no mapping.
|
||||
UpstreamModel string
|
||||
// RequestType is the granular request type: 0=unknown, 1=sync, 2=stream, 3=ws_v2.
|
||||
// Matches service.RequestType enum semantics from usage_log.go.
|
||||
RequestType *int16
|
||||
UserAgent string
|
||||
|
||||
ErrorPhase string
|
||||
|
||||
@@ -93,6 +93,10 @@ type OpsUpstreamErrorEvent struct {
|
||||
UpstreamStatusCode int `json:"upstream_status_code,omitempty"`
|
||||
UpstreamRequestID string `json:"upstream_request_id,omitempty"`
|
||||
|
||||
// UpstreamURL is the actual upstream URL that was called (host + path, query/fragment stripped).
|
||||
// Helps debug 404/routing errors by showing which endpoint was targeted.
|
||||
UpstreamURL string `json:"upstream_url,omitempty"`
|
||||
|
||||
// Best-effort upstream request capture (sanitized+trimmed).
|
||||
// Required for retrying a specific upstream attempt.
|
||||
UpstreamRequestBody string `json:"upstream_request_body,omitempty"`
|
||||
@@ -119,6 +123,7 @@ func appendOpsUpstreamError(c *gin.Context, ev OpsUpstreamErrorEvent) {
|
||||
ev.UpstreamRequestBody = strings.TrimSpace(ev.UpstreamRequestBody)
|
||||
ev.UpstreamResponseBody = strings.TrimSpace(ev.UpstreamResponseBody)
|
||||
ev.Kind = strings.TrimSpace(ev.Kind)
|
||||
ev.UpstreamURL = strings.TrimSpace(ev.UpstreamURL)
|
||||
ev.Message = strings.TrimSpace(ev.Message)
|
||||
ev.Detail = strings.TrimSpace(ev.Detail)
|
||||
if ev.Message != "" {
|
||||
@@ -205,3 +210,19 @@ func ParseOpsUpstreamErrors(raw string) ([]*OpsUpstreamErrorEvent, error) {
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// safeUpstreamURL returns scheme + host + path from a URL, stripping query/fragment
|
||||
// to avoid leaking sensitive query parameters (e.g. OAuth tokens).
|
||||
func safeUpstreamURL(rawURL string) string {
|
||||
rawURL = strings.TrimSpace(rawURL)
|
||||
if rawURL == "" {
|
||||
return ""
|
||||
}
|
||||
if idx := strings.IndexByte(rawURL, '?'); idx >= 0 {
|
||||
rawURL = rawURL[:idx]
|
||||
}
|
||||
if idx := strings.IndexByte(rawURL, '#'); idx >= 0 {
|
||||
rawURL = rawURL[:idx]
|
||||
}
|
||||
return rawURL
|
||||
}
|
||||
|
||||
@@ -8,6 +8,27 @@ import (
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestSafeUpstreamURL(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
input string
|
||||
want string
|
||||
}{
|
||||
{"strips query", "https://api.anthropic.com/v1/messages?beta=true", "https://api.anthropic.com/v1/messages"},
|
||||
{"strips fragment", "https://api.openai.com/v1/responses#frag", "https://api.openai.com/v1/responses"},
|
||||
{"strips both", "https://host/path?token=secret#x", "https://host/path"},
|
||||
{"no query or fragment", "https://host/path", "https://host/path"},
|
||||
{"empty string", "", ""},
|
||||
{"whitespace only", " ", ""},
|
||||
{"query before fragment", "https://h/p?a=1#f", "https://h/p"},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
require.Equal(t, tt.want, safeUpstreamURL(tt.input))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAppendOpsUpstreamError_UsesRequestBodyBytesFromContext(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
rec := httptest.NewRecorder()
|
||||
|
||||
@@ -163,7 +163,7 @@ func (s *RateLimitService) HandleUpstreamError(ctx context.Context, account *Acc
|
||||
account.Credentials = make(map[string]any)
|
||||
}
|
||||
account.Credentials["expires_at"] = time.Now().Format(time.RFC3339)
|
||||
if err := s.accountRepo.Update(ctx, account); err != nil {
|
||||
if err := persistAccountCredentials(ctx, s.accountRepo, account, account.Credentials); err != nil {
|
||||
slog.Warn("oauth_401_force_refresh_update_failed", "account_id", account.ID, "error", err)
|
||||
} else {
|
||||
slog.Info("oauth_401_force_refresh_set", "account_id", account.ID, "platform", account.Platform)
|
||||
|
||||
@@ -15,9 +15,11 @@ import (
|
||||
|
||||
type rateLimitAccountRepoStub struct {
|
||||
mockAccountRepoForGemini
|
||||
setErrorCalls int
|
||||
tempCalls int
|
||||
lastErrorMsg string
|
||||
setErrorCalls int
|
||||
tempCalls int
|
||||
updateCredentialsCalls int
|
||||
lastCredentials map[string]any
|
||||
lastErrorMsg string
|
||||
}
|
||||
|
||||
func (r *rateLimitAccountRepoStub) SetError(ctx context.Context, id int64, errorMsg string) error {
|
||||
@@ -31,6 +33,12 @@ func (r *rateLimitAccountRepoStub) SetTempUnschedulable(ctx context.Context, id
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *rateLimitAccountRepoStub) UpdateCredentials(ctx context.Context, id int64, credentials map[string]any) error {
|
||||
r.updateCredentialsCalls++
|
||||
r.lastCredentials = cloneCredentials(credentials)
|
||||
return nil
|
||||
}
|
||||
|
||||
type tokenCacheInvalidatorRecorder struct {
|
||||
accounts []*Account
|
||||
err error
|
||||
@@ -110,6 +118,7 @@ func TestRateLimitService_HandleUpstreamError_OAuth401InvalidatorError(t *testin
|
||||
require.True(t, shouldDisable)
|
||||
require.Equal(t, 0, repo.setErrorCalls)
|
||||
require.Equal(t, 1, repo.tempCalls)
|
||||
require.Equal(t, 1, repo.updateCredentialsCalls)
|
||||
require.Len(t, invalidator.accounts, 1)
|
||||
}
|
||||
|
||||
@@ -130,3 +139,22 @@ func TestRateLimitService_HandleUpstreamError_NonOAuth401(t *testing.T) {
|
||||
require.Equal(t, 1, repo.setErrorCalls)
|
||||
require.Empty(t, invalidator.accounts)
|
||||
}
|
||||
|
||||
func TestRateLimitService_HandleUpstreamError_OAuth401UsesCredentialsUpdater(t *testing.T) {
|
||||
repo := &rateLimitAccountRepoStub{}
|
||||
service := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
|
||||
account := &Account{
|
||||
ID: 103,
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeOAuth,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "token",
|
||||
},
|
||||
}
|
||||
|
||||
shouldDisable := service.HandleUpstreamError(context.Background(), account, 401, http.Header{}, []byte("unauthorized"))
|
||||
|
||||
require.True(t, shouldDisable)
|
||||
require.Equal(t, 1, repo.updateCredentialsCalls)
|
||||
require.NotEmpty(t, repo.lastCredentials["expires_at"])
|
||||
}
|
||||
|
||||
@@ -81,7 +81,7 @@ func (m *sessionWindowMockRepo) Delete(context.Context, int64) error { panic(
|
||||
func (m *sessionWindowMockRepo) List(context.Context, pagination.PaginationParams) ([]Account, *pagination.PaginationResult, error) {
|
||||
panic("unexpected")
|
||||
}
|
||||
func (m *sessionWindowMockRepo) ListWithFilters(context.Context, pagination.PaginationParams, string, string, string, string, int64) ([]Account, *pagination.PaginationResult, error) {
|
||||
func (m *sessionWindowMockRepo) ListWithFilters(context.Context, pagination.PaginationParams, string, string, string, string, int64, string) ([]Account, *pagination.PaginationResult, error) {
|
||||
panic("unexpected")
|
||||
}
|
||||
func (m *sessionWindowMockRepo) ListByGroup(context.Context, int64) ([]Account, error) {
|
||||
|
||||
@@ -947,7 +947,7 @@ func (c *SoraSDKClient) applyRecoveredToken(ctx context.Context, account *Accoun
|
||||
}
|
||||
|
||||
if c.accountRepo != nil {
|
||||
if err := c.accountRepo.Update(ctx, account); err != nil && c.debugEnabled() {
|
||||
if err := persistAccountCredentials(ctx, c.accountRepo, account, account.Credentials); err != nil && c.debugEnabled() {
|
||||
c.debugLogf("persist_recovered_token_failed account_id=%d err=%s", account.ID, logredact.RedactText(err.Error()))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -280,8 +280,7 @@ func (s *TokenRefreshService) refreshWithRetry(ctx context.Context, account *Acc
|
||||
newCredentials, err = refresher.Refresh(ctx, account)
|
||||
if newCredentials != nil {
|
||||
newCredentials["_token_version"] = time.Now().UnixMilli()
|
||||
account.Credentials = newCredentials
|
||||
if saveErr := s.accountRepo.Update(ctx, account); saveErr != nil {
|
||||
if saveErr := persistAccountCredentials(ctx, s.accountRepo, account, newCredentials); saveErr != nil {
|
||||
return fmt.Errorf("failed to save credentials: %w", saveErr)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -14,19 +14,40 @@ import (
|
||||
|
||||
type tokenRefreshAccountRepo struct {
|
||||
mockAccountRepoForGemini
|
||||
updateCalls int
|
||||
setErrorCalls int
|
||||
clearTempCalls int
|
||||
lastAccount *Account
|
||||
updateErr error
|
||||
updateCalls int
|
||||
fullUpdateCalls int
|
||||
updateCredentialsCalls int
|
||||
setErrorCalls int
|
||||
clearTempCalls int
|
||||
lastAccount *Account
|
||||
updateErr error
|
||||
}
|
||||
|
||||
func (r *tokenRefreshAccountRepo) Update(ctx context.Context, account *Account) error {
|
||||
r.updateCalls++
|
||||
r.fullUpdateCalls++
|
||||
r.lastAccount = account
|
||||
return r.updateErr
|
||||
}
|
||||
|
||||
func (r *tokenRefreshAccountRepo) UpdateCredentials(ctx context.Context, id int64, credentials map[string]any) error {
|
||||
r.updateCalls++
|
||||
r.updateCredentialsCalls++
|
||||
if r.updateErr != nil {
|
||||
return r.updateErr
|
||||
}
|
||||
cloned := cloneCredentials(credentials)
|
||||
if r.accountsByID != nil {
|
||||
if acc, ok := r.accountsByID[id]; ok && acc != nil {
|
||||
acc.Credentials = cloned
|
||||
r.lastAccount = acc
|
||||
return nil
|
||||
}
|
||||
}
|
||||
r.lastAccount = &Account{ID: id, Credentials: cloned}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *tokenRefreshAccountRepo) SetError(ctx context.Context, id int64, errorMsg string) error {
|
||||
r.setErrorCalls++
|
||||
return nil
|
||||
@@ -112,6 +133,8 @@ func TestTokenRefreshService_RefreshWithRetry_InvalidatesCache(t *testing.T) {
|
||||
err := service.refreshWithRetry(context.Background(), account, refresher, refresher, time.Hour)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, repo.updateCalls)
|
||||
require.Equal(t, 1, repo.updateCredentialsCalls)
|
||||
require.Equal(t, 0, repo.fullUpdateCalls)
|
||||
require.Equal(t, 1, invalidator.calls)
|
||||
require.Equal(t, "new-token", account.GetCredential("access_token"))
|
||||
}
|
||||
@@ -249,9 +272,43 @@ func TestTokenRefreshService_RefreshWithRetry_OtherPlatformOAuth(t *testing.T) {
|
||||
err := service.refreshWithRetry(context.Background(), account, refresher, refresher, time.Hour)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, repo.updateCalls)
|
||||
require.Equal(t, 1, repo.updateCredentialsCalls)
|
||||
require.Equal(t, 1, invalidator.calls) // 所有 OAuth 账户刷新后触发缓存失效
|
||||
}
|
||||
|
||||
func TestTokenRefreshService_RefreshWithRetry_UsesCredentialsUpdater(t *testing.T) {
|
||||
repo := &tokenRefreshAccountRepo{}
|
||||
cfg := &config.Config{
|
||||
TokenRefresh: config.TokenRefreshConfig{
|
||||
MaxRetries: 1,
|
||||
RetryBackoffSeconds: 0,
|
||||
},
|
||||
}
|
||||
service := NewTokenRefreshService(repo, nil, nil, nil, nil, nil, nil, cfg, nil)
|
||||
resetAt := time.Now().Add(30 * time.Minute)
|
||||
account := &Account{
|
||||
ID: 17,
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeOAuth,
|
||||
RateLimitResetAt: &resetAt,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "old-token",
|
||||
},
|
||||
}
|
||||
refresher := &tokenRefresherStub{
|
||||
credentials: map[string]any{
|
||||
"access_token": "new-token",
|
||||
},
|
||||
}
|
||||
|
||||
err := service.refreshWithRetry(context.Background(), account, refresher, refresher, time.Hour)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, repo.updateCredentialsCalls)
|
||||
require.Equal(t, 0, repo.fullUpdateCalls)
|
||||
require.NotNil(t, account.RateLimitResetAt)
|
||||
require.WithinDuration(t, resetAt, *account.RateLimitResetAt, time.Second)
|
||||
}
|
||||
|
||||
// TestTokenRefreshService_RefreshWithRetry_UpdateFailed 测试更新失败的情况
|
||||
func TestTokenRefreshService_RefreshWithRetry_UpdateFailed(t *testing.T) {
|
||||
repo := &tokenRefreshAccountRepo{updateErr: errors.New("update failed")}
|
||||
@@ -390,7 +447,7 @@ func TestTokenRefreshService_RefreshWithRetry_ClearsTempUnschedulable(t *testing
|
||||
err := service.refreshWithRetry(context.Background(), account, refresher, refresher, time.Hour)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, repo.updateCalls)
|
||||
require.Equal(t, 1, repo.clearTempCalls) // DB 清除
|
||||
require.Equal(t, 1, repo.clearTempCalls) // DB 清除
|
||||
require.Equal(t, 1, tempCache.deleteCalls) // Redis 缓存也应清除
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user