mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-09-24 16:05:44 +08:00
Merge origin/main into fix/openai-passthrough-5xx-failover (resolve #4269 conflicts)
# Conflicts: # backend/internal/service/openai_gateway_passthrough.go
This commit is contained in:
@@ -96,9 +96,6 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
usageLogRepository := repository.NewUsageLogRepository(client, db)
|
||||
usageService := service.NewUsageService(usageLogRepository, userRepository, client, apiKeyAuthCacheInvalidator)
|
||||
opsRepository := repository.NewOpsRepository(db)
|
||||
batchImageRepository := repository.NewBatchImageRepository(db)
|
||||
batchImageQueue := repository.NewBatchImageQueue(redisClient, configConfig)
|
||||
batchImageDownloadLimiter := repository.NewBatchImageDownloadLimiter(redisClient, configConfig)
|
||||
usageBillingRepository := repository.NewUsageBillingRepository(client, db)
|
||||
gatewayCache := repository.NewGatewayCache(redisClient)
|
||||
schedulerOutboxRepository := repository.NewSchedulerOutboxRepository(db)
|
||||
@@ -137,11 +134,6 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
channelRepository := repository.NewChannelRepository(db)
|
||||
channelService := service.NewChannelService(channelRepository, groupRepository, apiKeyAuthCacheInvalidator, pricingService)
|
||||
modelPricingResolver := service.NewModelPricingResolver(channelService, billingService)
|
||||
batchImageModelPricingResolver := service.ProvideBatchImageModelPricingResolver(modelPricingResolver)
|
||||
batchImagePublicService := service.NewBatchImagePublicService(batchImageRepository, accountRepository, groupRepository, userGroupRateRepository, batchImageQueue, batchImageModelPricingResolver, usageBillingRepository, apiKeyAuthCacheInvalidator, configConfig)
|
||||
batchImageDownloadService := service.NewBatchImageDownloadService(batchImageRepository, accountRepository, batchImageDownloadLimiter, configConfig)
|
||||
batchImageCleanupService := service.ProvideBatchImageCleanupService(batchImageRepository, accountRepository, configConfig)
|
||||
batchImageWorkerRuntime := service.ProvideBatchImageWorkerRuntime(batchImageRepository, accountRepository, batchImageQueue, usageBillingRepository, usageLogRepository, batchImageModelPricingResolver, apiKeyAuthCacheInvalidator, configConfig)
|
||||
notificationEmailService := service.NewNotificationEmailService(settingRepository, emailService)
|
||||
balanceNotifyService := service.ProvideBalanceNotifyService(emailService, settingRepository, accountRepository, notificationEmailService)
|
||||
gatewayService := service.NewGatewayService(accountRepository, groupRepository, usageLogRepository, usageBillingRepository, userRepository, userSubscriptionRepository, userGroupRateRepository, gatewayCache, configConfig, schedulerSnapshotService, concurrencyService, billingService, rateLimitService, billingCacheService, identityService, httpUpstream, deferredService, claudeTokenProvider, sessionLimitCache, rpmCache, digestSessionStore, settingService, tlsFingerprintProfileService, channelService, modelPricingResolver, balanceNotifyService, serviceUserPlatformQuotaRepository)
|
||||
@@ -191,10 +183,10 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
antigravityQuotaFetcher := service.NewAntigravityQuotaFetcher(proxyRepository)
|
||||
grokQuotaFetcher := service.NewGrokQuotaFetcher()
|
||||
grokQuotaService := service.ProvideGrokQuotaService(accountRepository, proxyRepository, grokTokenProvider, httpUpstream, usageLogRepository)
|
||||
openAIQuotaService := service.ProvideOpenAIQuotaService(accountRepository, proxyRepository, openAITokenProvider, privacyClientFactory)
|
||||
openAIQuotaService := service.ProvideOpenAIQuotaService(accountRepository, proxyRepository, openAITokenProvider, privacyClientFactory, openAIGatewayService)
|
||||
usageCache := service.NewUsageCache()
|
||||
accountUsageService := service.NewAccountUsageService(accountRepository, usageLogRepository, claudeUsageFetcher, geminiQuotaService, antigravityQuotaFetcher, grokQuotaFetcher, grokQuotaService, openAIQuotaService, usageCache, identityCache, tlsFingerprintProfileService)
|
||||
accountTestService := service.NewAccountTestService(accountRepository, geminiTokenProvider, claudeTokenProvider, grokTokenProvider, antigravityGatewayService, httpUpstream, configConfig, tlsFingerprintProfileService)
|
||||
accountUsageService := service.ProvideAccountUsageService(accountRepository, usageLogRepository, claudeUsageFetcher, geminiQuotaService, antigravityQuotaFetcher, grokQuotaFetcher, grokQuotaService, openAIQuotaService, usageCache, identityCache, tlsFingerprintProfileService, openAIGatewayService)
|
||||
accountTestService := service.ProvideAccountTestService(accountRepository, geminiTokenProvider, claudeTokenProvider, grokTokenProvider, antigravityGatewayService, httpUpstream, configConfig, tlsFingerprintProfileService, openAIGatewayService)
|
||||
crsSyncService := service.NewCRSSyncService(accountRepository, proxyRepository, oAuthService, openAIOAuthService, geminiOAuthService, configConfig)
|
||||
accountHandler := admin.ProvideAccountHandler(adminService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, rateLimitService, accountUsageService, accountTestService, concurrencyService, crsSyncService, sessionLimitCache, rpmCache, compositeTokenCacheInvalidator, grokQuotaService)
|
||||
adminAnnouncementHandler := admin.NewAnnouncementHandler(announcementService)
|
||||
@@ -267,6 +259,13 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
handlerPaymentHandler := handler.NewPaymentHandler(paymentService, paymentConfigService)
|
||||
paymentWebhookHandler := handler.NewPaymentWebhookHandler(paymentService, registry)
|
||||
availableChannelHandler := handler.NewAvailableChannelHandler(channelService, apiKeyService, settingService)
|
||||
batchImageRepository := repository.NewBatchImageRepository(db)
|
||||
batchImageQueue := repository.NewBatchImageQueue(redisClient, configConfig)
|
||||
batchImageModelPricingResolver := service.ProvideBatchImageModelPricingResolver(modelPricingResolver)
|
||||
batchImagePublicService := service.NewBatchImagePublicService(batchImageRepository, accountRepository, groupRepository, userGroupRateRepository, batchImageQueue, batchImageModelPricingResolver, usageBillingRepository, apiKeyAuthCacheInvalidator, configConfig)
|
||||
batchImageDownloadLimiter := repository.NewBatchImageDownloadLimiter(redisClient, configConfig)
|
||||
batchImageDownloadService := service.NewBatchImageDownloadService(batchImageRepository, accountRepository, batchImageDownloadLimiter, configConfig)
|
||||
batchImageCleanupService := service.ProvideBatchImageCleanupService(batchImageRepository, accountRepository, configConfig)
|
||||
batchImageHandler := handler.NewBatchImageHandler(batchImagePublicService, batchImageDownloadService, batchImageCleanupService)
|
||||
idempotencyCoordinator := service.ProvideIdempotencyCoordinator(idempotencyRepository, configConfig)
|
||||
idempotencyCleanupService := service.ProvideIdempotencyCleanupService(idempotencyRepository, configConfig)
|
||||
@@ -285,6 +284,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
accountExpiryService := service.ProvideAccountExpiryService(accountRepository)
|
||||
proxyExpiryService := service.ProvideProxyExpiryService(proxyRepository)
|
||||
subscriptionExpiryService := service.ProvideSubscriptionExpiryService(userSubscriptionRepository, settingRepository, notificationEmailService, leaderLockCache, db)
|
||||
batchImageWorkerRuntime := service.ProvideBatchImageWorkerRuntime(batchImageRepository, accountRepository, batchImageQueue, usageBillingRepository, usageLogRepository, batchImageModelPricingResolver, apiKeyAuthCacheInvalidator, configConfig)
|
||||
scheduledTestRunnerService := service.ProvideScheduledTestRunnerService(scheduledTestPlanRepository, scheduledTestService, accountTestService, rateLimitService, configConfig)
|
||||
paymentOrderExpiryService := service.ProvidePaymentOrderExpiryService(paymentService, leaderLockCache, db)
|
||||
channelMonitorRunner := service.ProvideChannelMonitorRunner(channelMonitorService, settingService)
|
||||
|
||||
@@ -0,0 +1,47 @@
|
||||
package admin
|
||||
|
||||
import (
|
||||
"crypto/ed25519"
|
||||
"crypto/rand"
|
||||
"crypto/x509"
|
||||
"encoding/base64"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestNormalizeCodexImportEntryAcceptsAgentIdentityAuthJSON(t *testing.T) {
|
||||
_, privateKey, err := ed25519.GenerateKey(rand.Reader)
|
||||
require.NoError(t, err)
|
||||
der, err := x509.MarshalPKCS8PrivateKey(privateKey)
|
||||
require.NoError(t, err)
|
||||
privateKeyBase64 := base64.StdEncoding.EncodeToString(der)
|
||||
|
||||
item, err := normalizeCodexImportEntry(codexImportEntry{
|
||||
Index: 1,
|
||||
Value: map[string]any{
|
||||
"auth_mode": "agentIdentity",
|
||||
"agent_identity": map[string]any{
|
||||
"agent_runtime_id": "runtime-import",
|
||||
"agent_private_key": privateKeyBase64,
|
||||
"account_id": "account-import",
|
||||
"chatgpt_user_id": "user-import",
|
||||
"email": "agent@example.invalid",
|
||||
"plan_type": "pro",
|
||||
"chatgpt_account_is_fedramp": false,
|
||||
},
|
||||
},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, item)
|
||||
require.True(t, item.IsAgentIdentity)
|
||||
require.Equal(t, service.OpenAIAuthModeAgentIdentity, item.Credentials["auth_mode"])
|
||||
require.Equal(t, "runtime-import", item.Credentials["agent_runtime_id"])
|
||||
require.Equal(t, privateKeyBase64, item.Credentials["agent_private_key"])
|
||||
require.Equal(t, "account-import", item.Credentials["chatgpt_account_id"])
|
||||
require.Equal(t, "user-import", item.Credentials["chatgpt_user_id"])
|
||||
require.NotContains(t, item.Credentials, "access_token")
|
||||
require.NotContains(t, item.Credentials, "refresh_token")
|
||||
require.NotEmpty(t, item.WarningTexts)
|
||||
}
|
||||
@@ -72,20 +72,25 @@ type codexImportEntry struct {
|
||||
}
|
||||
|
||||
type codexImportAccount struct {
|
||||
Name string
|
||||
AccessToken string
|
||||
RefreshToken string
|
||||
IDToken string
|
||||
Email string
|
||||
AccountID string
|
||||
UserID string
|
||||
PlanType string
|
||||
Organization string
|
||||
Credentials map[string]any
|
||||
Extra map[string]any
|
||||
TokenExpiresAt *time.Time
|
||||
IdentityKeys []string
|
||||
WarningTexts []string
|
||||
Name string
|
||||
AccessToken string
|
||||
RefreshToken string
|
||||
IDToken string
|
||||
Email string
|
||||
AccountID string
|
||||
UserID string
|
||||
PlanType string
|
||||
Organization string
|
||||
AgentRuntimeID string
|
||||
AgentPrivateKey string
|
||||
AgentTaskID string
|
||||
AgentFedRAMP bool
|
||||
IsAgentIdentity bool
|
||||
Credentials map[string]any
|
||||
Extra map[string]any
|
||||
TokenExpiresAt *time.Time
|
||||
IdentityKeys []string
|
||||
WarningTexts []string
|
||||
}
|
||||
|
||||
type codexJWTClaims struct {
|
||||
@@ -496,6 +501,41 @@ func normalizeCodexImportEntry(entry codexImportEntry) (*codexImportAccount, err
|
||||
case string:
|
||||
item.AccessToken = strings.TrimSpace(raw)
|
||||
case map[string]any:
|
||||
if agentIdentity, ok := firstCodexMap(raw, []string{"agent_identity"}, []string{"agentIdentity"}); ok || strings.EqualFold(firstCodexString(raw, []string{"auth_mode"}, []string{"authMode"}), service.OpenAIAuthModeAgentIdentity) {
|
||||
if !ok {
|
||||
agentIdentity = raw
|
||||
}
|
||||
item.IsAgentIdentity = true
|
||||
item.AgentRuntimeID = firstCodexString(agentIdentity, []string{"agent_runtime_id"}, []string{"agentRuntimeId"})
|
||||
item.AgentPrivateKey = firstCodexString(agentIdentity, []string{"agent_private_key"}, []string{"agentPrivateKey"})
|
||||
item.AgentTaskID = firstCodexString(agentIdentity, []string{"task_id"}, []string{"taskId"})
|
||||
item.AccountID = firstCodexString(agentIdentity, []string{"account_id"}, []string{"accountId"})
|
||||
item.UserID = firstCodexString(agentIdentity, []string{"chatgpt_user_id"}, []string{"chatgptUserId"})
|
||||
item.Email = firstCodexString(agentIdentity, []string{"email"})
|
||||
item.PlanType = firstCodexString(agentIdentity, []string{"plan_type"}, []string{"planType"})
|
||||
item.AgentFedRAMP = firstCodexBool(agentIdentity, []string{"chatgpt_account_is_fedramp"}, []string{"chatgptAccountIsFedramp"})
|
||||
if item.AgentRuntimeID == "" || item.AgentPrivateKey == "" || item.AccountID == "" || item.UserID == "" {
|
||||
return nil, errors.New("agent identity 缺少必要字段")
|
||||
}
|
||||
if err := service.ValidateOpenAIAgentIdentityPrivateKey(item.AgentPrivateKey); err != nil {
|
||||
return nil, errors.New("agent identity private key 格式无效")
|
||||
}
|
||||
item.Credentials["auth_mode"] = service.OpenAIAuthModeAgentIdentity
|
||||
item.Credentials["agent_runtime_id"] = item.AgentRuntimeID
|
||||
item.Credentials["agent_private_key"] = item.AgentPrivateKey
|
||||
item.Credentials["chatgpt_account_id"] = item.AccountID
|
||||
item.Credentials["chatgpt_user_id"] = item.UserID
|
||||
item.Credentials["chatgpt_account_is_fedramp"] = item.AgentFedRAMP
|
||||
setCodexCredentialIfNotEmpty(item.Credentials, "task_id", item.AgentTaskID)
|
||||
setCodexCredentialIfNotEmpty(item.Credentials, "email", item.Email)
|
||||
setCodexCredentialIfNotEmpty(item.Credentials, "plan_type", item.PlanType)
|
||||
if item.AgentTaskID == "" {
|
||||
item.WarningTexts = append(item.WarningTexts, "未包含 task_id,首次请求会使用现有 runtime 注册新 task")
|
||||
}
|
||||
item.IdentityKeys = buildCodexAgentIdentityKeys(item.AccountID, item.UserID, item.Email, item.AgentRuntimeID)
|
||||
item.Name = buildCodexImportAccountName(item, entry.Index)
|
||||
return item, nil
|
||||
}
|
||||
item.AccessToken = firstCodexString(raw,
|
||||
[]string{"tokens", "access_token"},
|
||||
[]string{"tokens", "accessToken"},
|
||||
@@ -577,6 +617,9 @@ func normalizeCodexImportEntry(entry codexImportEntry) (*codexImportAccount, err
|
||||
return nil, fmt.Errorf("第 %d 条格式不支持", entry.Index)
|
||||
}
|
||||
|
||||
if item.IsAgentIdentity {
|
||||
return item, nil
|
||||
}
|
||||
if item.AccessToken == "" {
|
||||
return nil, errors.New("缺少 accessToken/access_token")
|
||||
}
|
||||
@@ -812,6 +855,9 @@ func sanitizeCodexImportCredentialExtras(input map[string]any) map[string]any {
|
||||
"openai_auth_mode": {},
|
||||
"token_type": {},
|
||||
"chatgpt_account_is_fedramp": {},
|
||||
"agent_runtime_id": {},
|
||||
"agent_private_key": {},
|
||||
"task_id": {},
|
||||
}
|
||||
out := make(map[string]any, len(input))
|
||||
for key, value := range input {
|
||||
@@ -842,6 +888,14 @@ func buildCodexImportIdentityKeys(accountID, userID, email, accessToken, refresh
|
||||
return buildCodexStoredIdentityKeys(accountID, userID, email, accessToken)
|
||||
}
|
||||
|
||||
func buildCodexAgentIdentityKeys(accountID, userID, email, runtimeID string) []string {
|
||||
keys := buildCodexStoredIdentityKeys(accountID, userID, email, "")
|
||||
if runtimeID = strings.TrimSpace(runtimeID); runtimeID != "" {
|
||||
keys = append([]string{"agent:" + runtimeID}, keys...)
|
||||
}
|
||||
return keys
|
||||
}
|
||||
|
||||
// buildCodexStoredIdentityKeys 生成存量账号索引键,保留 user/account 维度,
|
||||
// 让 accessToken-only 账号后续升级为完整 OAuth 时仍能命中并更新原账号。
|
||||
func buildCodexStoredIdentityKeys(accountID, userID, email, accessToken string) []string {
|
||||
@@ -891,6 +945,10 @@ func (i *codexAccountIndex) Add(account service.Account) {
|
||||
for _, key := range keys {
|
||||
i.accountsByKey[key] = upsertCodexAccount(i.accountsByKey[key], account)
|
||||
}
|
||||
if runtimeID := codexCredentialString(account.Credentials, "agent_runtime_id"); runtimeID != "" {
|
||||
key := "agent:" + runtimeID
|
||||
i.accountsByKey[key] = upsertCodexAccount(i.accountsByKey[key], account)
|
||||
}
|
||||
}
|
||||
|
||||
func (i *codexAccountIndex) remove(accountID int64) {
|
||||
@@ -1043,6 +1101,38 @@ func firstCodexString(obj map[string]any, paths ...[]string) string {
|
||||
return ""
|
||||
}
|
||||
|
||||
func firstCodexMap(obj map[string]any, paths ...[]string) (map[string]any, bool) {
|
||||
for _, path := range paths {
|
||||
value, ok := codexPathValue(obj, path)
|
||||
if !ok || value == nil {
|
||||
continue
|
||||
}
|
||||
if mapped, ok := value.(map[string]any); ok {
|
||||
return mapped, true
|
||||
}
|
||||
}
|
||||
return nil, false
|
||||
}
|
||||
|
||||
func firstCodexBool(obj map[string]any, paths ...[]string) bool {
|
||||
for _, path := range paths {
|
||||
value, ok := codexPathValue(obj, path)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
switch value := value.(type) {
|
||||
case bool:
|
||||
return value
|
||||
case string:
|
||||
parsed, err := strconv.ParseBool(strings.TrimSpace(value))
|
||||
if err == nil {
|
||||
return parsed
|
||||
}
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func copyCodexExtraString(obj map[string]any, extra map[string]any, key string, path []string) {
|
||||
value := firstCodexString(obj, path)
|
||||
if value != "" {
|
||||
|
||||
@@ -90,10 +90,13 @@ func TestGrokOAuthHandlerQueryQuotaProbesUpstream(t *testing.T) {
|
||||
ID: 42,
|
||||
Platform: service.PlatformGrok,
|
||||
Type: service.AccountTypeOAuth,
|
||||
Status: service.StatusActive,
|
||||
Schedulable: true,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "access-token",
|
||||
"expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
|
||||
"access_token": "access-token",
|
||||
"refresh_token": "refresh-token",
|
||||
"expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
|
||||
},
|
||||
}}
|
||||
upstream := &grokQuotaHandlerUpstream{}
|
||||
|
||||
@@ -386,8 +386,8 @@ func (h *OpsHandler) ListRequestErrorUpstreamErrors(c *gin.Context) {
|
||||
filter.EndTime = &endTime
|
||||
}
|
||||
filter.View = "all"
|
||||
filter.Phase = "upstream"
|
||||
// 上游错误列表需含 status<400 的 recovered 行,显式豁免客户端可见守卫。
|
||||
filter.ErrorPhasesAny = []string{"upstream", "account_auth"}
|
||||
// Provider-health list includes recovered inference and credential rows.
|
||||
filter.IncludeRecoveredUpstream = true
|
||||
filter.Owner = "provider"
|
||||
filter.Source = strings.TrimSpace(c.Query("error_source"))
|
||||
@@ -470,8 +470,8 @@ func (h *OpsHandler) ListUpstreamErrors(c *gin.Context) {
|
||||
}
|
||||
|
||||
filter.View = parseOpsViewParam(c)
|
||||
filter.Phase = "upstream"
|
||||
// 上游错误列表需含 status<400 的 recovered 行,显式豁免客户端可见守卫。
|
||||
filter.ErrorPhasesAny = []string{"upstream", "account_auth"}
|
||||
// Provider-health list includes recovered inference and credential rows.
|
||||
filter.IncludeRecoveredUpstream = true
|
||||
filter.Owner = "provider"
|
||||
filter.Source = strings.TrimSpace(c.Query("error_source"))
|
||||
|
||||
@@ -20,6 +20,7 @@ func TestRedactCredentials_StripsSensitiveKeysAndReportsStatus(t *testing.T) {
|
||||
"aws_secret_access_key": "aws-secret",
|
||||
"service_account_json": map[string]any{"private_key": "..."},
|
||||
"private_key": "raw-key",
|
||||
"agent_private_key": "agent-key-secret",
|
||||
// 非敏感
|
||||
"base_url": "https://api.example.com",
|
||||
"model_mapping": map[string]any{"foo": "bar"},
|
||||
@@ -35,6 +36,7 @@ func TestRedactCredentials_StripsSensitiveKeysAndReportsStatus(t *testing.T) {
|
||||
require.NotContains(t, out, "aws_secret_access_key")
|
||||
require.NotContains(t, out, "service_account_json")
|
||||
require.NotContains(t, out, "private_key")
|
||||
require.NotContains(t, out, "agent_private_key")
|
||||
|
||||
require.Equal(t, "https://api.example.com", out["base_url"])
|
||||
require.Equal(t, map[string]any{"foo": "bar"}, out["model_mapping"])
|
||||
@@ -47,6 +49,7 @@ func TestRedactCredentials_StripsSensitiveKeysAndReportsStatus(t *testing.T) {
|
||||
require.True(t, status["has_aws_secret_access_key"])
|
||||
require.True(t, status["has_service_account_json"])
|
||||
require.True(t, status["has_private_key"])
|
||||
require.True(t, status["has_agent_private_key"])
|
||||
|
||||
// 状态 map 不应携带非敏感键的 has_*
|
||||
require.NotContains(t, status, "has_base_url")
|
||||
@@ -84,6 +87,7 @@ func TestRedactCredentials_AllKnownSensitiveKeys(t *testing.T) {
|
||||
"api_key", "session_key", "cookie",
|
||||
"aws_secret_access_key", "aws_session_token",
|
||||
"service_account_json", "service_account", "private_key",
|
||||
"agent_private_key",
|
||||
}
|
||||
in := make(map[string]any, len(keys))
|
||||
for _, k := range keys {
|
||||
|
||||
@@ -5,6 +5,8 @@ import (
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"go.uber.org/zap"
|
||||
@@ -71,7 +73,15 @@ func (s *FailoverState) HandleFailoverError(
|
||||
retryLimit int,
|
||||
failoverErr *service.UpstreamFailoverError,
|
||||
) FailoverAction {
|
||||
// 客户端已断开:failover 只会用已取消的 context 重新选号并必然失败,
|
||||
// 不应再被当成账号耗尽处理(误报 502)。
|
||||
if ctx != nil && ctx.Err() != nil {
|
||||
return FailoverCanceled
|
||||
}
|
||||
s.LastFailoverErr = failoverErr
|
||||
if failoverErr == nil || !failoverErr.ShouldRetryNextAccount() {
|
||||
return FailoverExhausted
|
||||
}
|
||||
|
||||
// 缓存计费判断
|
||||
if needForceCacheBilling(s.hasBoundSession, failoverErr) {
|
||||
@@ -135,6 +145,12 @@ func (s *FailoverState) HandleFailoverError(
|
||||
// 返回 FailoverExhausted 时,调用方应返回错误响应。
|
||||
// 返回 FailoverCanceled 时,调用方应直接 return。
|
||||
func (s *FailoverState) HandleSelectionExhausted(ctx context.Context) FailoverAction {
|
||||
// 客户端已断开时选号失败是 context canceled 的必然结果,
|
||||
// 不代表账号耗尽,直接按取消终止。
|
||||
if ctx.Err() != nil {
|
||||
return FailoverCanceled
|
||||
}
|
||||
|
||||
if s.LastFailoverErr != nil &&
|
||||
s.LastFailoverErr.StatusCode == http.StatusServiceUnavailable &&
|
||||
s.SwitchCount <= s.MaxSwitches {
|
||||
@@ -163,6 +179,28 @@ func needForceCacheBilling(hasBoundSession bool, failoverErr *service.UpstreamFa
|
||||
return hasBoundSession || (failoverErr != nil && failoverErr.ForceCacheBilling)
|
||||
}
|
||||
|
||||
// failoverClientGone 判断下游客户端是否已断开(请求 context 已取消)。
|
||||
// 客户端断开后 failover 必须静默终止:用已取消的 context 重新选号只会得到
|
||||
// context.Canceled,并被误报成账号耗尽(通用 502);上游 detach 的在途请求
|
||||
// 照常完成计费,但不再为无人接收的响应启动新的上游尝试。
|
||||
// 响应尚未提交时把状态码标记为 499(client closed request),供访问日志归类。
|
||||
func failoverClientGone(c *gin.Context) bool {
|
||||
if c == nil || c.Request == nil || c.Request.Context().Err() == nil {
|
||||
return false
|
||||
}
|
||||
// 先停 compact 心跳(接管 ResponseWriter,建立 happens-before),与
|
||||
// handleStreamingAwareError/errorResponse 等终结路径对齐,避免心跳
|
||||
// goroutine 与下面的状态标记并发触碰同一 writer。心跳已提交 200 时
|
||||
// 状态码已固化,不再标 499。
|
||||
if service.StopOpenAICompactSSEKeepaliveCommitted(c) {
|
||||
return true
|
||||
}
|
||||
if !c.Writer.Written() {
|
||||
c.Status(statusClientClosedRequest)
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// sleepWithContext 等待指定时长,返回 false 表示 context 已取消。
|
||||
func sleepWithContext(ctx context.Context, d time.Duration) bool {
|
||||
if d <= 0 {
|
||||
|
||||
@@ -2,11 +2,14 @@ package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
@@ -128,6 +131,50 @@ func TestSleepWithContext(t *testing.T) {
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestHandleFailoverError_BasicSwitch(t *testing.T) {
|
||||
t.Run("显式停止不切换账号且旧错误默认仍切换", func(t *testing.T) {
|
||||
mock := &mockTempUnscheduler{}
|
||||
fs := NewFailoverState(3, false)
|
||||
stopErr := &service.UpstreamFailoverError{
|
||||
Stage: service.GatewayFailureStageAccountAuth,
|
||||
Scope: service.GatewayFailureScopeProvider,
|
||||
NextAccountAction: service.NextAccountStop,
|
||||
}
|
||||
|
||||
action := fs.HandleFailoverError(context.Background(), mock, 100, service.PlatformGrok, maxSameAccountRetries, stopErr)
|
||||
|
||||
require.Equal(t, FailoverExhausted, action)
|
||||
require.Zero(t, fs.SwitchCount)
|
||||
require.Empty(t, fs.FailedAccountIDs)
|
||||
require.Equal(t, stopErr, fs.LastFailoverErr)
|
||||
|
||||
legacyErr := newTestFailoverErr(http.StatusTooManyRequests, false, false)
|
||||
action = fs.HandleFailoverError(context.Background(), mock, 100, service.PlatformGrok, maxSameAccountRetries, legacyErr)
|
||||
|
||||
require.Equal(t, FailoverContinue, action)
|
||||
require.Equal(t, 1, fs.SwitchCount)
|
||||
require.Contains(t, fs.FailedAccountIDs, int64(100))
|
||||
})
|
||||
|
||||
t.Run("已取消的认证失败不改变切换状态", func(t *testing.T) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
mock := &mockTempUnscheduler{}
|
||||
fs := NewFailoverState(3, false)
|
||||
err := &service.UpstreamFailoverError{
|
||||
Stage: service.GatewayFailureStageAccountAuth,
|
||||
Scope: service.GatewayFailureScopeAccount,
|
||||
NextAccountAction: service.NextAccountRetry,
|
||||
}
|
||||
|
||||
action := fs.HandleFailoverError(ctx, mock, 101, service.PlatformGrok, maxSameAccountRetries, err)
|
||||
|
||||
require.Equal(t, FailoverCanceled, action)
|
||||
require.Zero(t, fs.SwitchCount)
|
||||
require.Empty(t, fs.FailedAccountIDs)
|
||||
require.Nil(t, fs.LastFailoverErr)
|
||||
require.Empty(t, mock.calls)
|
||||
})
|
||||
|
||||
t.Run("非重试错误_非Antigravity_直接切换", func(t *testing.T) {
|
||||
mock := &mockTempUnscheduler{}
|
||||
fs := NewFailoverState(3, false)
|
||||
@@ -444,7 +491,28 @@ func TestHandleFailoverError_ContextCanceled(t *testing.T) {
|
||||
err := newTestFailoverErr(400, true, false)
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel() // 立即取消
|
||||
go func() {
|
||||
time.Sleep(30 * time.Millisecond)
|
||||
cancel() // 通过入口检查后、sleep 期间取消
|
||||
}()
|
||||
|
||||
start := time.Now()
|
||||
action := fs.HandleFailoverError(ctx, mock, 100, "openai", maxSameAccountRetries, err)
|
||||
elapsed := time.Since(start)
|
||||
|
||||
require.Equal(t, FailoverCanceled, action)
|
||||
require.Less(t, elapsed, 400*time.Millisecond, "sleep 应被取消打断")
|
||||
// 进入重试分支后才取消:重试计数已递增
|
||||
require.Equal(t, 1, fs.SameAccountRetryCount[100])
|
||||
})
|
||||
|
||||
t.Run("入口即已取消_不改动任何failover状态", func(t *testing.T) {
|
||||
mock := &mockTempUnscheduler{}
|
||||
fs := NewFailoverState(3, false)
|
||||
err := newTestFailoverErr(520, false, false)
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel() // 调用前客户端已断开
|
||||
|
||||
start := time.Now()
|
||||
action := fs.HandleFailoverError(ctx, mock, 100, "openai", maxSameAccountRetries, err)
|
||||
@@ -452,8 +520,12 @@ func TestHandleFailoverError_ContextCanceled(t *testing.T) {
|
||||
|
||||
require.Equal(t, FailoverCanceled, action)
|
||||
require.Less(t, elapsed, 100*time.Millisecond, "应立即返回")
|
||||
// 重试计数仍应递增
|
||||
require.Equal(t, 1, fs.SameAccountRetryCount[100])
|
||||
// 入口已取消时不得改变任何 failover 状态。
|
||||
require.Equal(t, 0, fs.SwitchCount, "取消的请求不应计入切换")
|
||||
require.Equal(t, 0, fs.SameAccountRetryCount[100], "取消的请求不应改动重试计数")
|
||||
require.NotContains(t, fs.FailedAccountIDs, int64(100))
|
||||
require.Nil(t, fs.LastFailoverErr)
|
||||
require.Empty(t, mock.calls, "不应触发 TempUnschedule")
|
||||
})
|
||||
|
||||
t.Run("Antigravity延迟期间context取消", func(t *testing.T) {
|
||||
@@ -755,6 +827,29 @@ func TestHandleSelectionExhausted(t *testing.T) {
|
||||
require.Less(t, elapsed, 100*time.Millisecond, "应立即返回")
|
||||
})
|
||||
|
||||
t.Run("context已取消_非503也返回Canceled而非Exhausted", func(t *testing.T) {
|
||||
// #4257 核心场景:客户端断开后选号失败源于 context canceled,
|
||||
// 不应被当成账号耗尽转成 502。
|
||||
fs := NewFailoverState(3, false)
|
||||
fs.LastFailoverErr = newTestFailoverErr(520, false, false)
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
|
||||
action := fs.HandleSelectionExhausted(ctx)
|
||||
require.Equal(t, FailoverCanceled, action)
|
||||
})
|
||||
|
||||
t.Run("context已取消_无LastFailoverErr也返回Canceled", func(t *testing.T) {
|
||||
fs := NewFailoverState(3, false)
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
|
||||
action := fs.HandleSelectionExhausted(ctx)
|
||||
require.Equal(t, FailoverCanceled, action)
|
||||
})
|
||||
|
||||
t.Run("503且SwitchCount等于MaxSwitches_仍可重试", func(t *testing.T) {
|
||||
fs := NewFailoverState(2, false)
|
||||
fs.LastFailoverErr = newTestFailoverErr(503, false, false)
|
||||
@@ -764,3 +859,47 @@ func TestHandleSelectionExhausted(t *testing.T) {
|
||||
require.Equal(t, FailoverContinue, action)
|
||||
})
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// failoverClientGone 测试
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestFailoverClientGone(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
t.Run("活跃请求返回false", func(t *testing.T) {
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
|
||||
require.False(t, failoverClientGone(c))
|
||||
require.Equal(t, http.StatusOK, c.Writer.Status(), "不应改动状态码")
|
||||
})
|
||||
|
||||
t.Run("客户端已断开_返回true并标记499", func(t *testing.T) {
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil).WithContext(ctx)
|
||||
|
||||
require.True(t, failoverClientGone(c))
|
||||
require.Equal(t, statusClientClosedRequest, c.Writer.Status())
|
||||
})
|
||||
|
||||
t.Run("响应已提交_不改状态码", func(t *testing.T) {
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil).WithContext(ctx)
|
||||
c.String(http.StatusOK, "partial")
|
||||
|
||||
require.True(t, failoverClientGone(c))
|
||||
require.Equal(t, http.StatusOK, c.Writer.Status(), "已提交的状态码不应被覆盖")
|
||||
})
|
||||
|
||||
t.Run("nil安全", func(t *testing.T) {
|
||||
require.False(t, failoverClientGone(nil))
|
||||
})
|
||||
}
|
||||
|
||||
@@ -327,6 +327,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
|
||||
c.Request = c.Request.WithContext(ctx)
|
||||
continue
|
||||
case FailoverCanceled:
|
||||
failoverClientGone(c)
|
||||
return
|
||||
default: // FailoverExhausted
|
||||
if fs.LastFailoverErr != nil {
|
||||
@@ -456,6 +457,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
|
||||
h.handleFailoverExhausted(c, fs.LastFailoverErr, service.PlatformGemini, streamStarted)
|
||||
return
|
||||
case FailoverCanceled:
|
||||
failoverClientGone(c)
|
||||
return
|
||||
}
|
||||
}
|
||||
@@ -613,6 +615,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
|
||||
c.Request = c.Request.WithContext(ctx)
|
||||
continue
|
||||
case FailoverCanceled:
|
||||
failoverClientGone(c)
|
||||
return
|
||||
default: // FailoverExhausted
|
||||
if fs.LastFailoverErr != nil {
|
||||
@@ -876,6 +879,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
|
||||
h.handleFailoverExhausted(c, fs.LastFailoverErr, account.Platform, streamStarted)
|
||||
return
|
||||
case FailoverCanceled:
|
||||
failoverClientGone(c)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,98 @@
|
||||
//go:build unit
|
||||
|
||||
package handler
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
||||
middleware "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type countingGatewaySchedulerCache struct {
|
||||
*fakeSchedulerCache
|
||||
snapshotCalls atomic.Int64
|
||||
}
|
||||
|
||||
func (c *countingGatewaySchedulerCache) GetSnapshot(ctx context.Context, bucket service.SchedulerBucket) ([]*service.Account, bool, error) {
|
||||
c.snapshotCalls.Add(1)
|
||||
return c.fakeSchedulerCache.GetSnapshot(ctx, bucket)
|
||||
}
|
||||
|
||||
func TestGatewayHandlerPreCancelledCompatibleRequestsDoNotSelectAccount(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
groupID := int64(9100)
|
||||
group := &service.Group{ID: groupID, Hydrated: true, Platform: service.PlatformAnthropic, Status: service.StatusActive}
|
||||
account := &service.Account{
|
||||
ID: 9101, Platform: service.PlatformAnthropic, Type: service.AccountTypeAPIKey,
|
||||
Status: service.StatusActive, Schedulable: true, Concurrency: 1,
|
||||
AccountGroups: []service.AccountGroup{{AccountID: 9101, GroupID: groupID}},
|
||||
}
|
||||
schedulerCache := &countingGatewaySchedulerCache{fakeSchedulerCache: &fakeSchedulerCache{accounts: []*service.Account{account}}}
|
||||
schedulerSnapshot := service.NewSchedulerSnapshotService(schedulerCache, nil, nil, nil, nil)
|
||||
gatewayService := service.NewGatewayService(
|
||||
nil, &fakeGroupRepo{group: group}, nil, nil, nil, nil, nil, nil, nil,
|
||||
schedulerSnapshot, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil,
|
||||
)
|
||||
cfg := &config.Config{RunMode: config.RunModeSimple}
|
||||
billingCacheService := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil)
|
||||
t.Cleanup(billingCacheService.Stop)
|
||||
h := &GatewayHandler{
|
||||
gatewayService: gatewayService,
|
||||
billingCacheService: billingCacheService,
|
||||
concurrencyHelper: NewConcurrencyHelper(service.NewConcurrencyService(&fakeConcurrencyCache{}), SSEPingFormatClaude, 0),
|
||||
maxAccountSwitches: 1,
|
||||
cfg: cfg,
|
||||
}
|
||||
apiKey := &service.APIKey{
|
||||
ID: 9102, UserID: 9103, GroupID: &groupID, Group: group, Status: service.StatusActive,
|
||||
User: &service.User{ID: 9103, Concurrency: 10, Balance: 100},
|
||||
}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
path string
|
||||
body string
|
||||
call func(*gin.Context)
|
||||
}{
|
||||
{
|
||||
name: "responses", path: "/v1/responses", body: `{"model":"claude-test","input":"hello","stream":false}`,
|
||||
call: h.Responses,
|
||||
},
|
||||
{
|
||||
name: "chat completions", path: "/v1/chat/completions", body: `{"model":"claude-test","messages":[{"role":"user","content":"hello"}],"stream":false}`,
|
||||
call: h.ChatCompletions,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
schedulerCache.snapshotCalls.Store(0)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
ctx = context.WithValue(ctx, ctxkey.Group, group)
|
||||
req := httptest.NewRequest(http.MethodPost, tt.path, bytes.NewBufferString(tt.body)).WithContext(ctx)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
c.Request = req
|
||||
c.Set(string(middleware.ContextKeyAPIKey), apiKey)
|
||||
c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: apiKey.UserID, Concurrency: 10})
|
||||
|
||||
tt.call(c)
|
||||
|
||||
require.Zero(t, schedulerCache.snapshotCalls.Load(), "a cancelled request must stop before the account selector")
|
||||
_, selected := c.Get(opsAccountIDKey)
|
||||
require.False(t, selected)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -159,6 +159,9 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) {
|
||||
}
|
||||
|
||||
for {
|
||||
if c.Request.Context().Err() != nil {
|
||||
return
|
||||
}
|
||||
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, selectionSessionHash, reqModel, fs.FailedAccountIDs, "", int64(0))
|
||||
if err != nil {
|
||||
if len(fs.FailedAccountIDs) == 0 {
|
||||
@@ -178,6 +181,7 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) {
|
||||
case FailoverContinue:
|
||||
continue
|
||||
case FailoverCanceled:
|
||||
failoverClientGone(c)
|
||||
return
|
||||
default:
|
||||
if fs.LastFailoverErr != nil {
|
||||
@@ -262,6 +266,7 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) {
|
||||
h.handleCCFailoverExhausted(c, fs.LastFailoverErr, streamStarted)
|
||||
return
|
||||
case FailoverCanceled:
|
||||
failoverClientGone(c)
|
||||
return
|
||||
}
|
||||
}
|
||||
@@ -328,6 +333,14 @@ func (h *GatewayHandler) handleCCFailoverExhausted(c *gin.Context, lastErr *serv
|
||||
if streamStarted {
|
||||
return
|
||||
}
|
||||
if lastErr != nil {
|
||||
copyFailoverRetryAfter(c, lastErr.ResponseHeaders)
|
||||
}
|
||||
if lastErr != nil && lastErr.IsCredentialFailure() {
|
||||
status, message := credentialFailoverClientResponse(lastErr)
|
||||
h.chatCompletionsErrorResponse(c, status, "server_error", message)
|
||||
return
|
||||
}
|
||||
statusCode := http.StatusBadGateway
|
||||
if lastErr != nil && lastErr.StatusCode > 0 {
|
||||
statusCode = lastErr.StatusCode
|
||||
|
||||
@@ -157,6 +157,9 @@ func (h *GatewayHandler) Responses(c *gin.Context) {
|
||||
fs := NewFailoverState(h.maxAccountSwitches, false)
|
||||
|
||||
for {
|
||||
if requestCtx.Err() != nil {
|
||||
return
|
||||
}
|
||||
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(requestCtx, apiKey.GroupID, sessionHash, reqModel, fs.FailedAccountIDs, "", int64(0))
|
||||
if err != nil {
|
||||
if len(fs.FailedAccountIDs) == 0 {
|
||||
@@ -176,6 +179,7 @@ func (h *GatewayHandler) Responses(c *gin.Context) {
|
||||
case FailoverContinue:
|
||||
continue
|
||||
case FailoverCanceled:
|
||||
failoverClientGone(c)
|
||||
return
|
||||
default:
|
||||
if fs.LastFailoverErr != nil {
|
||||
@@ -241,6 +245,7 @@ func (h *GatewayHandler) Responses(c *gin.Context) {
|
||||
h.handleResponsesFailoverExhausted(c, fs.LastFailoverErr, streamStarted)
|
||||
return
|
||||
case FailoverCanceled:
|
||||
failoverClientGone(c)
|
||||
return
|
||||
}
|
||||
}
|
||||
@@ -307,6 +312,14 @@ func (h *GatewayHandler) handleResponsesFailoverExhausted(c *gin.Context, lastEr
|
||||
if streamStarted {
|
||||
return // Can't write error after stream started
|
||||
}
|
||||
if lastErr != nil {
|
||||
copyFailoverRetryAfter(c, lastErr.ResponseHeaders)
|
||||
}
|
||||
if lastErr != nil && lastErr.IsCredentialFailure() {
|
||||
status, message := credentialFailoverClientResponse(lastErr)
|
||||
h.responsesErrorResponse(c, status, "server_error", message)
|
||||
return
|
||||
}
|
||||
statusCode := http.StatusBadGateway
|
||||
if lastErr != nil && lastErr.StatusCode > 0 {
|
||||
statusCode = lastErr.StatusCode
|
||||
|
||||
@@ -163,20 +163,15 @@ func wrapReleaseOnDone(ctx context.Context, releaseFunc func()) func() {
|
||||
return nil
|
||||
}
|
||||
var once sync.Once
|
||||
var stop func() bool
|
||||
|
||||
release := func() {
|
||||
once.Do(func() {
|
||||
if stop != nil {
|
||||
_ = stop()
|
||||
}
|
||||
releaseFunc()
|
||||
})
|
||||
releaseOnce := func() {
|
||||
once.Do(releaseFunc)
|
||||
}
|
||||
stop := context.AfterFunc(ctx, releaseOnce)
|
||||
|
||||
stop = context.AfterFunc(ctx, release)
|
||||
|
||||
return release
|
||||
return func() {
|
||||
_ = stop()
|
||||
releaseOnce()
|
||||
}
|
||||
}
|
||||
|
||||
// IncrementWaitCount increments the wait count for a user
|
||||
|
||||
@@ -6,6 +6,8 @@ import (
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// TestWrapReleaseOnDone_NoGoroutineLeak 验证 wrapReleaseOnDone 修复后不会泄露 goroutine
|
||||
@@ -67,6 +69,21 @@ func TestWrapReleaseOnDone_ContextCancellation(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestWrapReleaseOnDone_AlreadyCancelledReleasesExactlyOnce(t *testing.T) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
|
||||
var releaseCount int32
|
||||
release := wrapReleaseOnDone(ctx, func() {
|
||||
atomic.AddInt32(&releaseCount, 1)
|
||||
})
|
||||
release()
|
||||
|
||||
require.Eventually(t, func() bool {
|
||||
return atomic.LoadInt32(&releaseCount) == 1
|
||||
}, time.Second, time.Millisecond)
|
||||
}
|
||||
|
||||
// TestWrapReleaseOnDone_MultipleCallsOnlyReleaseOnce 验证多次调用 release 只释放一次
|
||||
func TestWrapReleaseOnDone_MultipleCallsOnlyReleaseOnce(t *testing.T) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
|
||||
@@ -368,6 +368,7 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
|
||||
c.Request = c.Request.WithContext(ctx)
|
||||
continue
|
||||
case FailoverCanceled:
|
||||
failoverClientGone(c)
|
||||
return
|
||||
default: // FailoverExhausted
|
||||
h.handleGeminiFailoverExhausted(c, fs.LastFailoverErr)
|
||||
@@ -490,6 +491,7 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
|
||||
h.handleGeminiFailoverExhausted(c, fs.LastFailoverErr)
|
||||
return
|
||||
case FailoverCanceled:
|
||||
failoverClientGone(c)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
@@ -166,6 +166,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
failedAccountIDs := make(map[int64]struct{})
|
||||
sameAccountRetryCount := make(map[int64]int)
|
||||
var lastFailoverErr *service.UpstreamFailoverError
|
||||
var oauth429FailoverState service.OpenAIOAuth429FailoverState
|
||||
switchCount := 0
|
||||
maxAccountSwitches := h.maxAccountSwitches
|
||||
if maxAccountSwitches <= 0 {
|
||||
@@ -174,6 +175,9 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
routingStart := time.Now()
|
||||
|
||||
for {
|
||||
if failoverClientGone(c) {
|
||||
return
|
||||
}
|
||||
selection, scheduleDecision, err := h.gatewayService.SelectAccountWithSchedulerForCapability(
|
||||
requestCtx,
|
||||
apiKey.GroupID,
|
||||
@@ -188,6 +192,10 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
service.PlatformGrok,
|
||||
)
|
||||
if err != nil {
|
||||
if failoverClientGone(c) {
|
||||
reqLog.Info("grok_media.account_select_aborted_client_disconnected", zap.Error(err))
|
||||
return
|
||||
}
|
||||
reqLog.Warn("grok_media.account_select_failed",
|
||||
zap.Error(err),
|
||||
zap.Int("excluded_account_count", len(failedAccountIDs)),
|
||||
@@ -257,11 +265,24 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
if err != nil {
|
||||
var failoverErr *service.UpstreamFailoverError
|
||||
if errors.As(err, &failoverErr) {
|
||||
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil)
|
||||
if failoverClientGone(c) {
|
||||
reqLog.Info("grok_media.failover_aborted_client_disconnected",
|
||||
zap.Int64("account_id", account.ID),
|
||||
zap.Int("upstream_status", failoverErr.StatusCode),
|
||||
)
|
||||
return
|
||||
}
|
||||
if failoverErr.ShouldReportAccountScheduleFailure() {
|
||||
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil)
|
||||
}
|
||||
if c.Writer.Size() != writerSizeBeforeForward {
|
||||
h.handleFailoverExhausted(c, failoverErr, true)
|
||||
return
|
||||
}
|
||||
if !failoverErr.ShouldRetryNextAccount() {
|
||||
h.handleFailoverExhausted(c, failoverErr, false)
|
||||
return
|
||||
}
|
||||
if failoverErr.RetryableOnSameAccount {
|
||||
retryLimit := account.GetPoolModeRetryCount()
|
||||
if sameAccountRetryCount[account.ID] < retryLimit {
|
||||
@@ -288,6 +309,10 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
return
|
||||
}
|
||||
switchCount++
|
||||
if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount, &oauth429FailoverState) {
|
||||
h.handleFailoverExhausted(c, failoverErr, false)
|
||||
return
|
||||
}
|
||||
reqLog.Warn("grok_media.upstream_failover_switching",
|
||||
zap.Int64("account_id", account.ID),
|
||||
zap.Int("upstream_status", failoverErr.StatusCode),
|
||||
|
||||
@@ -106,6 +106,7 @@ func (h *OpenAIGatewayHandler) AlphaSearch(c *gin.Context) {
|
||||
failedAccountIDs := make(map[int64]struct{})
|
||||
var lastFailoverErr *service.UpstreamFailoverError
|
||||
switchCount := 0
|
||||
var oauth429FailoverState service.OpenAIOAuth429FailoverState
|
||||
routingStart := time.Now()
|
||||
|
||||
for {
|
||||
@@ -123,6 +124,10 @@ func (h *OpenAIGatewayHandler) AlphaSearch(c *gin.Context) {
|
||||
service.PlatformOpenAI,
|
||||
)
|
||||
if err != nil || selection == nil || selection.Account == nil {
|
||||
if failoverClientGone(c) {
|
||||
reqLog.Info("openai_alpha_search.account_select_aborted_client_disconnected", zap.Error(err))
|
||||
return
|
||||
}
|
||||
if len(failedAccountIDs) == 0 {
|
||||
cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, requestedModel, requestedModel, service.PlatformOpenAI)
|
||||
if !cls.ModelNotFound {
|
||||
@@ -180,6 +185,13 @@ func (h *OpenAIGatewayHandler) AlphaSearch(c *gin.Context) {
|
||||
h.handleFailoverExhausted(c, failoverErr, true)
|
||||
return
|
||||
}
|
||||
if failoverClientGone(c) {
|
||||
reqLog.Info("openai_alpha_search.failover_aborted_client_disconnected",
|
||||
zap.Int64("account_id", account.ID),
|
||||
zap.Int("upstream_status", failoverErr.StatusCode),
|
||||
)
|
||||
return
|
||||
}
|
||||
h.gatewayService.RecordOpenAIAccountSwitch()
|
||||
failedAccountIDs[account.ID] = struct{}{}
|
||||
lastFailoverErr = failoverErr
|
||||
@@ -188,7 +200,7 @@ func (h *OpenAIGatewayHandler) AlphaSearch(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
switchCount++
|
||||
if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount) {
|
||||
if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount, &oauth429FailoverState) {
|
||||
h.handleFailoverExhausted(c, failoverErr, false)
|
||||
return
|
||||
}
|
||||
|
||||
@@ -133,8 +133,12 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) {
|
||||
failedAccountIDs := make(map[int64]struct{})
|
||||
sameAccountRetryCount := make(map[int64]int)
|
||||
var lastFailoverErr *service.UpstreamFailoverError
|
||||
var oauth429FailoverState service.OpenAIOAuth429FailoverState
|
||||
|
||||
for {
|
||||
if failoverClientGone(c) {
|
||||
return
|
||||
}
|
||||
reqLog.Debug("openai_chat_completions.account_selecting", zap.Int("excluded_account_count", len(failedAccountIDs)))
|
||||
selection, scheduleDecision, err := h.gatewayService.SelectAccountWithSchedulerForCapability(
|
||||
c.Request.Context(),
|
||||
@@ -150,6 +154,10 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) {
|
||||
requestPlatform,
|
||||
)
|
||||
if err != nil {
|
||||
if failoverClientGone(c) {
|
||||
reqLog.Info("openai_chat_completions.account_select_aborted_client_disconnected", zap.Error(err))
|
||||
return
|
||||
}
|
||||
reqLog.Warn("openai_chat_completions.account_select_failed",
|
||||
zap.Error(openAICompatibleSelectionErrorForLog(err, requestPlatform)),
|
||||
zap.Int("excluded_account_count", len(failedAccountIDs)),
|
||||
@@ -231,11 +239,24 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) {
|
||||
} else {
|
||||
var failoverErr *service.UpstreamFailoverError
|
||||
if errors.As(err, &failoverErr) {
|
||||
if failoverClientGone(c) {
|
||||
reqLog.Info("openai_chat_completions.failover_aborted_client_disconnected",
|
||||
zap.Int64("account_id", account.ID),
|
||||
zap.Int("upstream_status", failoverErr.StatusCode),
|
||||
)
|
||||
return
|
||||
}
|
||||
if c.Writer.Size() != writerSizeBeforeForward {
|
||||
h.handleFailoverExhausted(c, failoverErr, true)
|
||||
return
|
||||
}
|
||||
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil)
|
||||
if failoverErr.ShouldReportAccountScheduleFailure() {
|
||||
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil)
|
||||
}
|
||||
if !failoverErr.ShouldRetryNextAccount() {
|
||||
h.handleFailoverExhausted(c, failoverErr, streamStarted)
|
||||
return
|
||||
}
|
||||
// Pool mode: retry on the same account
|
||||
if failoverErr.RetryableOnSameAccount {
|
||||
retryLimit := account.GetPoolModeRetryCount()
|
||||
@@ -263,7 +284,7 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
switchCount++
|
||||
if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount) {
|
||||
if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount, &oauth429FailoverState) {
|
||||
h.handleFailoverExhausted(c, failoverErr, streamStarted)
|
||||
return
|
||||
}
|
||||
|
||||
@@ -121,6 +121,10 @@ func (h *OpenAIGatewayHandler) Embeddings(c *gin.Context) {
|
||||
false,
|
||||
)
|
||||
if err != nil {
|
||||
if failoverClientGone(c) {
|
||||
reqLog.Info("openai_embeddings.account_select_aborted_client_disconnected", zap.Error(err))
|
||||
return
|
||||
}
|
||||
reqLog.Warn("openai_embeddings.account_select_failed",
|
||||
zap.Error(err),
|
||||
zap.Int("excluded_account_count", len(failedAccountIDs)),
|
||||
@@ -189,6 +193,13 @@ func (h *OpenAIGatewayHandler) Embeddings(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil)
|
||||
if failoverClientGone(c) {
|
||||
reqLog.Info("openai_embeddings.failover_aborted_client_disconnected",
|
||||
zap.Int64("account_id", account.ID),
|
||||
zap.Int("upstream_status", failoverErr.StatusCode),
|
||||
)
|
||||
return
|
||||
}
|
||||
h.gatewayService.RecordOpenAIAccountSwitch()
|
||||
failedAccountIDs[account.ID] = struct{}{}
|
||||
lastFailoverErr = failoverErr
|
||||
|
||||
@@ -0,0 +1,926 @@
|
||||
//go:build unit
|
||||
|
||||
package handler
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
|
||||
"github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
coderws "github.com/coder/websocket"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type grokCredentialHandlerRepo struct {
|
||||
service.AccountRepository
|
||||
mu sync.Mutex
|
||||
accounts []service.Account
|
||||
setErrorIDs []int64
|
||||
setTempIDs []int64
|
||||
rateLimitIDs []int64
|
||||
updateExtraIDs []int64
|
||||
selectionCalls int
|
||||
setErrorErr error
|
||||
setTempErr error
|
||||
missingOnGet map[int64]bool
|
||||
}
|
||||
|
||||
func (r *grokCredentialHandlerRepo) ListSchedulableByPlatform(_ context.Context, platform string) ([]service.Account, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.selectionCalls++
|
||||
out := make([]service.Account, 0, len(r.accounts))
|
||||
for _, account := range r.accounts {
|
||||
if account.Platform == platform && account.IsSchedulable() {
|
||||
out = append(out, account)
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (r *grokCredentialHandlerRepo) ListSchedulableByGroupIDAndPlatform(ctx context.Context, _ int64, platform string) ([]service.Account, error) {
|
||||
return r.ListSchedulableByPlatform(ctx, platform)
|
||||
}
|
||||
|
||||
func (r *grokCredentialHandlerRepo) ListSchedulableUngroupedByPlatform(ctx context.Context, platform string) ([]service.Account, error) {
|
||||
return r.ListSchedulableByPlatform(ctx, platform)
|
||||
}
|
||||
|
||||
func (r *grokCredentialHandlerRepo) GetByID(_ context.Context, id int64) (*service.Account, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
if r.missingOnGet[id] {
|
||||
return nil, nil
|
||||
}
|
||||
for _, account := range r.accounts {
|
||||
if account.ID == id {
|
||||
copy := account
|
||||
copy.Credentials = cloneCredentialMap(account.Credentials)
|
||||
return ©, nil
|
||||
}
|
||||
}
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (r *grokCredentialHandlerRepo) SetError(_ context.Context, id int64, message string) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.setErrorIDs = append(r.setErrorIDs, id)
|
||||
if r.setErrorErr != nil {
|
||||
return r.setErrorErr
|
||||
}
|
||||
for i := range r.accounts {
|
||||
if r.accounts[i].ID == id {
|
||||
r.accounts[i].Status = service.StatusError
|
||||
r.accounts[i].Schedulable = false
|
||||
r.accounts[i].ErrorMessage = message
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *grokCredentialHandlerRepo) SetTempUnschedulable(_ context.Context, id int64, until time.Time, _ string) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.setTempIDs = append(r.setTempIDs, id)
|
||||
if r.setTempErr != nil {
|
||||
return r.setTempErr
|
||||
}
|
||||
for i := range r.accounts {
|
||||
if r.accounts[i].ID == id {
|
||||
value := until
|
||||
r.accounts[i].TempUnschedulableUntil = &value
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *grokCredentialHandlerRepo) SetRateLimited(_ context.Context, id int64, resetAt time.Time) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.rateLimitIDs = append(r.rateLimitIDs, id)
|
||||
for i := range r.accounts {
|
||||
if r.accounts[i].ID != id {
|
||||
continue
|
||||
}
|
||||
now := time.Now()
|
||||
r.accounts[i].RateLimitedAt = &now
|
||||
value := resetAt
|
||||
r.accounts[i].RateLimitResetAt = &value
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *grokCredentialHandlerRepo) SetRateLimitedIfLater(ctx context.Context, id int64, resetAt time.Time) error {
|
||||
r.mu.Lock()
|
||||
for i := range r.accounts {
|
||||
if r.accounts[i].ID == id && r.accounts[i].RateLimitResetAt != nil && !resetAt.After(*r.accounts[i].RateLimitResetAt) {
|
||||
r.mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
}
|
||||
r.mu.Unlock()
|
||||
return r.SetRateLimited(ctx, id, resetAt)
|
||||
}
|
||||
|
||||
func (r *grokCredentialHandlerRepo) SetGrokCredentialErrorIfMatch(
|
||||
_ context.Context,
|
||||
id int64,
|
||||
snapshot service.GrokCredentialMutationSnapshot,
|
||||
message string,
|
||||
) (bool, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
for i := range r.accounts {
|
||||
account := &r.accounts[i]
|
||||
if account.ID != id || !handlerGrokCredentialSnapshotMatches(account, snapshot) {
|
||||
continue
|
||||
}
|
||||
r.setErrorIDs = append(r.setErrorIDs, id)
|
||||
if r.setErrorErr != nil {
|
||||
return false, r.setErrorErr
|
||||
}
|
||||
account.Status = service.StatusError
|
||||
account.Schedulable = false
|
||||
account.ErrorMessage = message
|
||||
return true, nil
|
||||
}
|
||||
return false, nil
|
||||
}
|
||||
|
||||
func (r *grokCredentialHandlerRepo) SetGrokCredentialTempUnschedulableIfMatch(
|
||||
_ context.Context,
|
||||
id int64,
|
||||
snapshot service.GrokCredentialMutationSnapshot,
|
||||
until time.Time,
|
||||
_ string,
|
||||
) (bool, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
for i := range r.accounts {
|
||||
account := &r.accounts[i]
|
||||
if account.ID != id || !handlerGrokCredentialSnapshotMatches(account, snapshot) {
|
||||
continue
|
||||
}
|
||||
r.setTempIDs = append(r.setTempIDs, id)
|
||||
if r.setTempErr != nil {
|
||||
return false, r.setTempErr
|
||||
}
|
||||
value := until
|
||||
account.TempUnschedulableUntil = &value
|
||||
return true, nil
|
||||
}
|
||||
return false, nil
|
||||
}
|
||||
|
||||
func handlerGrokCredentialSnapshotMatches(account *service.Account, snapshot service.GrokCredentialMutationSnapshot) bool {
|
||||
if account == nil {
|
||||
return false
|
||||
}
|
||||
credentialsJSON, err := json.Marshal(account.Credentials)
|
||||
return err == nil && account.IsGrokOAuth() && account.IsSchedulable() && string(credentialsJSON) == snapshot.CredentialsJSON &&
|
||||
handlerGrokCredentialProxyIDsEqual(account.ProxyID, snapshot.ProxyID)
|
||||
}
|
||||
|
||||
func handlerGrokCredentialProxyIDsEqual(left, right *int64) bool {
|
||||
if left == nil || right == nil {
|
||||
return left == nil && right == nil
|
||||
}
|
||||
return *left == *right
|
||||
}
|
||||
|
||||
func (r *grokCredentialHandlerRepo) UpdateExtra(_ context.Context, id int64, updates map[string]any) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.updateExtraIDs = append(r.updateExtraIDs, id)
|
||||
for i := range r.accounts {
|
||||
if r.accounts[i].ID != id {
|
||||
continue
|
||||
}
|
||||
if r.accounts[i].Extra == nil {
|
||||
r.accounts[i].Extra = map[string]any{}
|
||||
}
|
||||
for key, value := range updates {
|
||||
r.accounts[i].Extra[key] = value
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *grokCredentialHandlerRepo) errorIDs() []int64 {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
return append([]int64(nil), r.setErrorIDs...)
|
||||
}
|
||||
|
||||
func (r *grokCredentialHandlerRepo) selectorCalls() int {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
return r.selectionCalls
|
||||
}
|
||||
|
||||
func (r *grokCredentialHandlerRepo) rateLimitedAccountIDs() []int64 {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
return append([]int64(nil), r.rateLimitIDs...)
|
||||
}
|
||||
|
||||
type grokCredentialHandlerTokenCache struct {
|
||||
service.GrokTokenCache
|
||||
mu sync.Mutex
|
||||
deleteErr error
|
||||
}
|
||||
|
||||
func (c *grokCredentialHandlerTokenCache) GetAccessToken(context.Context, string) (string, error) {
|
||||
return "", errors.New("not cached")
|
||||
}
|
||||
|
||||
func (c *grokCredentialHandlerTokenCache) SetAccessToken(context.Context, string, string, time.Duration) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *grokCredentialHandlerTokenCache) DeleteAccessToken(context.Context, string) error {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
return c.deleteErr
|
||||
}
|
||||
|
||||
func (c *grokCredentialHandlerTokenCache) AcquireRefreshLock(context.Context, string, time.Duration) (bool, error) {
|
||||
return true, nil
|
||||
}
|
||||
|
||||
func (c *grokCredentialHandlerTokenCache) ReleaseRefreshLock(context.Context, string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func cloneCredentialMap(source map[string]any) map[string]any {
|
||||
cloned := make(map[string]any, len(source))
|
||||
for key, value := range source {
|
||||
cloned[key] = value
|
||||
}
|
||||
return cloned
|
||||
}
|
||||
|
||||
type grokCredentialHandlerRefresher struct {
|
||||
mode string
|
||||
started chan struct{}
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
func (r *grokCredentialHandlerRefresher) CacheKey(account *service.Account) string {
|
||||
return service.GrokTokenCacheKey(account)
|
||||
}
|
||||
|
||||
func (r *grokCredentialHandlerRefresher) CanRefresh(account *service.Account) bool {
|
||||
return account != nil && account.IsGrokOAuth()
|
||||
}
|
||||
|
||||
func (r *grokCredentialHandlerRefresher) NeedsRefresh(account *service.Account, _ time.Duration) bool {
|
||||
return account != nil && (account.ID == 801 || r.mode == "all_revoked")
|
||||
}
|
||||
|
||||
func (r *grokCredentialHandlerRefresher) Refresh(ctx context.Context, _ *service.Account) (map[string]any, error) {
|
||||
switch r.mode {
|
||||
case "revoked", "all_revoked", "mutation_set_error", "mutation_cache":
|
||||
return nil, infraerrors.New(http.StatusBadGateway, "GROK_OAUTH_TOKEN_REFRESH_FAILED", "invalid_grant")
|
||||
case "provider":
|
||||
return nil, infraerrors.New(http.StatusBadGateway, "GROK_OAUTH_TOKEN_REFRESH_FAILED", "invalid_client")
|
||||
case "cancel":
|
||||
r.once.Do(func() { close(r.started) })
|
||||
<-ctx.Done()
|
||||
return nil, ctx.Err()
|
||||
case "transient", "mutation_temp":
|
||||
return nil, errors.New("temporary refresh transport failure")
|
||||
default:
|
||||
return nil, nil
|
||||
}
|
||||
}
|
||||
|
||||
type grokCredentialHandlerUpstream struct {
|
||||
service.HTTPUpstream
|
||||
mu sync.Mutex
|
||||
hits []int64
|
||||
requestURLs []string
|
||||
authorization []string
|
||||
failAccountID int64
|
||||
rateLimitIDs map[int64]bool
|
||||
failureStatus map[int64]int
|
||||
cancelRequest context.CancelFunc
|
||||
}
|
||||
|
||||
func (u *grokCredentialHandlerUpstream) Do(req *http.Request, _ string, accountID int64, _ int) (*http.Response, error) {
|
||||
var requestBody []byte
|
||||
if req.Body != nil {
|
||||
requestBody, _ = io.ReadAll(req.Body)
|
||||
}
|
||||
u.mu.Lock()
|
||||
u.hits = append(u.hits, accountID)
|
||||
u.requestURLs = append(u.requestURLs, req.URL.String())
|
||||
u.authorization = append(u.authorization, req.Header.Get("Authorization"))
|
||||
failAccountID := u.failAccountID
|
||||
rateLimited := u.rateLimitIDs[accountID]
|
||||
failureStatus := u.failureStatus[accountID]
|
||||
cancelRequest := u.cancelRequest
|
||||
u.mu.Unlock()
|
||||
if rateLimited {
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusTooManyRequests,
|
||||
Header: http.Header{
|
||||
"Content-Type": []string{"application/json"},
|
||||
"Retry-After": []string{"60"},
|
||||
},
|
||||
Body: io.NopCloser(bytes.NewBufferString(`{"error":{"message":"rate limited"}}`)),
|
||||
}, nil
|
||||
}
|
||||
if failureStatus > 0 {
|
||||
return &http.Response{
|
||||
StatusCode: failureStatus,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(bytes.NewBufferString(`{"error":{"message":"upstream unavailable"}}`)),
|
||||
}, nil
|
||||
}
|
||||
if accountID == failAccountID {
|
||||
if cancelRequest != nil {
|
||||
cancelRequest()
|
||||
}
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusPaymentRequired,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(bytes.NewBufferString(`{"error":{"message":"payment required"}}`)),
|
||||
}, nil
|
||||
}
|
||||
if bytes.Contains(requestBody, []byte(`"stream":true`)) {
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
||||
Body: io.NopCloser(bytes.NewBufferString(
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_healthy\",\"model\":\"grok-4.5\",\"usage\":{\"input_tokens\":1,\"output_tokens\":1}}}\n\n",
|
||||
)),
|
||||
}, nil
|
||||
}
|
||||
if strings.Contains(req.URL.Path, "/chat/completions") {
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(bytes.NewBufferString(
|
||||
`{"id":"chatcmpl_healthy","object":"chat.completion","model":"grok-4.5","choices":[{"index":0,"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}`,
|
||||
)),
|
||||
}, nil
|
||||
}
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(bytes.NewBufferString(
|
||||
`{"id":"resp_healthy","object":"response","model":"grok-4.5","status":"completed","output":[],"usage":{"input_tokens":1,"output_tokens":1}}`,
|
||||
)),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (u *grokCredentialHandlerUpstream) accountHits() []int64 {
|
||||
u.mu.Lock()
|
||||
defer u.mu.Unlock()
|
||||
return append([]int64(nil), u.hits...)
|
||||
}
|
||||
|
||||
func (u *grokCredentialHandlerUpstream) requests() ([]string, []string) {
|
||||
u.mu.Lock()
|
||||
defer u.mu.Unlock()
|
||||
return append([]string(nil), u.requestURLs...), append([]string(nil), u.authorization...)
|
||||
}
|
||||
|
||||
func TestResponsesCredentialFailoverLoop(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
t.Run("revoked account selects healthy account", func(t *testing.T) {
|
||||
h, repo, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "revoked")
|
||||
defer cleanup()
|
||||
_ = h
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/openai/v1/responses", bytes.NewBufferString(`{"model":"grok","input":"hello","stream":false}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
router.ServeHTTP(recorder, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, recorder.Code, recorder.Body.String())
|
||||
require.Contains(t, recorder.Body.String(), "resp_healthy")
|
||||
require.Equal(t, []int64{801}, repo.errorIDs())
|
||||
require.Equal(t, []int64{802}, upstream.accountHits())
|
||||
requestURLs, authorization := upstream.requests()
|
||||
require.Equal(t, []string{xai.DefaultCLIBaseURL + "/responses"}, requestURLs)
|
||||
require.Equal(t, []string{"Bearer healthy-access"}, authorization)
|
||||
})
|
||||
|
||||
t.Run("provider configuration stops before healthy account", func(t *testing.T) {
|
||||
h, repo, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "provider")
|
||||
defer cleanup()
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/openai/v1/responses", bytes.NewBufferString(`{"model":"grok","input":"hello","stream":false}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
router.ServeHTTP(recorder, req)
|
||||
|
||||
require.Equal(t, http.StatusServiceUnavailable, recorder.Code)
|
||||
require.Contains(t, recorder.Body.String(), service.GrokCredentialUnavailableClientMessage)
|
||||
require.Empty(t, repo.errorIDs())
|
||||
require.Empty(t, upstream.accountHits())
|
||||
require.Equal(t, 1, repo.selectorCalls())
|
||||
require.Zero(t, h.gatewayService.SnapshotOpenAIAccountSchedulerMetrics().RuntimeStatsAccountCount,
|
||||
"provider-scoped auth failure must not penalize the selected account")
|
||||
})
|
||||
|
||||
t.Run("parent cancellation stops before healthy account", func(t *testing.T) {
|
||||
_, repo, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "cancel")
|
||||
defer cleanup()
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
req := httptest.NewRequest(http.MethodPost, "/openai/v1/responses", bytes.NewBufferString(`{"model":"grok","input":"hello","stream":false}`)).WithContext(ctx)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
recorder := httptest.NewRecorder()
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
defer close(done)
|
||||
router.ServeHTTP(recorder, req)
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("credential refresh did not start")
|
||||
case <-findHandlerRefresherStarted(router):
|
||||
cancel()
|
||||
}
|
||||
select {
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("handler did not stop after cancellation")
|
||||
case <-done:
|
||||
}
|
||||
|
||||
require.Empty(t, repo.errorIDs())
|
||||
require.Empty(t, upstream.accountHits())
|
||||
})
|
||||
|
||||
t.Run("post-mapping cancellation stops before scheduler mutation or reselection", func(t *testing.T) {
|
||||
h, repo, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "postmap_cancel")
|
||||
defer cleanup()
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
upstream.mu.Lock()
|
||||
upstream.failAccountID = 801
|
||||
upstream.cancelRequest = cancel
|
||||
upstream.mu.Unlock()
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/openai/v1/responses", bytes.NewBufferString(`{"model":"grok","input":"hello","stream":false}`)).WithContext(ctx)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
router.ServeHTTP(recorder, req)
|
||||
|
||||
require.Equal(t, []int64{801}, upstream.accountHits())
|
||||
require.Empty(t, repo.errorIDs())
|
||||
require.Equal(t, 1, repo.selectorCalls())
|
||||
require.Zero(t, h.gatewayService.SnapshotOpenAIAccountSchedulerMetrics().RuntimeStatsAccountCount)
|
||||
})
|
||||
|
||||
t.Run("pre-cancelled request never invokes an account selector", func(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
method string
|
||||
path string
|
||||
body string
|
||||
}{
|
||||
{name: "responses", method: http.MethodPost, path: "/openai/v1/responses", body: `{"model":"grok","input":"hello","stream":false}`},
|
||||
{name: "messages", method: http.MethodPost, path: "/openai/v1/messages", body: `{"model":"grok","max_tokens":16,"messages":[{"role":"user","content":"hello"}]}`},
|
||||
{name: "chat completions", method: http.MethodPost, path: "/openai/v1/chat/completions", body: `{"model":"grok","messages":[{"role":"user","content":"hello"}],"stream":false}`},
|
||||
{name: "grok media", method: http.MethodGet, path: "/openai/v1/videos/request-1"},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
_, repo, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "revoked")
|
||||
defer cleanup()
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
recorder := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(tt.method, tt.path, bytes.NewBufferString(tt.body)).WithContext(ctx)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
router.ServeHTTP(recorder, req)
|
||||
|
||||
require.Zero(t, repo.selectorCalls())
|
||||
require.Empty(t, upstream.accountHits())
|
||||
})
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("credential state mutation failures stop before reselection", func(t *testing.T) {
|
||||
for _, mode := range []string{"mutation_set_error", "mutation_temp", "mutation_cache"} {
|
||||
t.Run(mode, func(t *testing.T) {
|
||||
_, repo, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, mode)
|
||||
defer cleanup()
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/openai/v1/responses", bytes.NewBufferString(`{"model":"grok","input":"hello","stream":false}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
router.ServeHTTP(recorder, req)
|
||||
|
||||
require.Equal(t, http.StatusServiceUnavailable, recorder.Code, recorder.Body.String())
|
||||
require.Contains(t, recorder.Body.String(), service.GrokCredentialUnavailableClientMessage)
|
||||
require.Empty(t, upstream.accountHits())
|
||||
require.Equal(t, 1, repo.selectorCalls())
|
||||
})
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("missing credential provider stops before upstream or reselection", func(t *testing.T) {
|
||||
_, repo, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "nil_provider")
|
||||
defer cleanup()
|
||||
recorder := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/openai/v1/responses", bytes.NewBufferString(`{"model":"grok","input":"hello","stream":false}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
router.ServeHTTP(recorder, req)
|
||||
|
||||
require.Equal(t, http.StatusServiceUnavailable, recorder.Code, recorder.Body.String())
|
||||
require.Contains(t, recorder.Body.String(), service.GrokCredentialUnavailableClientMessage)
|
||||
require.Equal(t, 1, repo.selectorCalls())
|
||||
require.Empty(t, upstream.accountHits())
|
||||
require.Empty(t, repo.errorIDs())
|
||||
})
|
||||
}
|
||||
|
||||
func TestResponsesGrok429FailoverIsBounded(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
t.Run("first rate limited account selects healthy account", func(t *testing.T) {
|
||||
_, repo, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "first_429")
|
||||
defer cleanup()
|
||||
recorder := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/openai/v1/responses", bytes.NewBufferString(`{"model":"grok","input":"hello","stream":false}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
router.ServeHTTP(recorder, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, recorder.Code, recorder.Body.String())
|
||||
require.Contains(t, recorder.Body.String(), "resp_healthy")
|
||||
require.Equal(t, []int64{801, 802}, upstream.accountHits())
|
||||
require.Equal(t, []int64{801}, repo.rateLimitedAccountIDs())
|
||||
})
|
||||
|
||||
t.Run("two rate limited accounts stop without sweeping the pool", func(t *testing.T) {
|
||||
_, repo, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "all_429")
|
||||
defer cleanup()
|
||||
recorder := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/openai/v1/responses", bytes.NewBufferString(`{"model":"grok","input":"hello","stream":false}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
router.ServeHTTP(recorder, req)
|
||||
|
||||
require.Equal(t, http.StatusTooManyRequests, recorder.Code, recorder.Body.String())
|
||||
require.Equal(t, []int64{801, 802}, upstream.accountHits())
|
||||
require.Equal(t, []int64{801, 802}, repo.rateLimitedAccountIDs())
|
||||
require.NotContains(t, recorder.Body.String(), "expired")
|
||||
require.NotContains(t, recorder.Body.String(), "healthy-access")
|
||||
require.NotContains(t, recorder.Body.String(), "rate limited")
|
||||
})
|
||||
}
|
||||
|
||||
func TestResponsesGrok429FailoverHandlesMixedStatuses(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
t.Run("429 then 500 stops after the bounded followup", func(t *testing.T) {
|
||||
_, _, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "mixed_429_500")
|
||||
defer cleanup()
|
||||
recorder := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/openai/v1/responses", bytes.NewBufferString(`{"model":"grok","input":"hello","stream":false}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
router.ServeHTTP(recorder, req)
|
||||
|
||||
require.Equal(t, http.StatusBadGateway, recorder.Code, recorder.Body.String())
|
||||
require.Equal(t, []int64{801, 802}, upstream.accountHits())
|
||||
require.NotContains(t, recorder.Body.String(), "upstream unavailable")
|
||||
})
|
||||
|
||||
t.Run("500 then 429 permits one healthy followup", func(t *testing.T) {
|
||||
_, _, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "mixed_500_429")
|
||||
defer cleanup()
|
||||
recorder := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/openai/v1/responses", bytes.NewBufferString(`{"model":"grok","input":"hello","stream":false}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
router.ServeHTTP(recorder, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, recorder.Code, recorder.Body.String())
|
||||
require.Equal(t, []int64{801, 802, 803}, upstream.accountHits())
|
||||
})
|
||||
|
||||
t.Run("OAuth 429 then API-key failure cannot bypass the bound", func(t *testing.T) {
|
||||
_, _, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "oauth_429_apikey_500")
|
||||
defer cleanup()
|
||||
recorder := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/openai/v1/responses", bytes.NewBufferString(`{"model":"grok","input":"hello","stream":false}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
router.ServeHTTP(recorder, req)
|
||||
|
||||
require.Equal(t, http.StatusBadGateway, recorder.Code, recorder.Body.String())
|
||||
require.Equal(t, []int64{801, 802}, upstream.accountHits())
|
||||
})
|
||||
}
|
||||
|
||||
func TestGrokMedia429FailoverIsBounded(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
t.Run("first 429 selects one healthy followup", func(t *testing.T) {
|
||||
_, _, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "first_429")
|
||||
defer cleanup()
|
||||
recorder := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/openai/v1/videos/request-1", nil)
|
||||
|
||||
router.ServeHTTP(recorder, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, recorder.Code, recorder.Body.String())
|
||||
require.Equal(t, []int64{801, 802}, upstream.accountHits())
|
||||
})
|
||||
|
||||
t.Run("second 429 stops without sweeping a third account", func(t *testing.T) {
|
||||
_, _, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "all_429")
|
||||
defer cleanup()
|
||||
recorder := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/openai/v1/videos/request-1", nil)
|
||||
|
||||
router.ServeHTTP(recorder, req)
|
||||
|
||||
require.Equal(t, http.StatusTooManyRequests, recorder.Code, recorder.Body.String())
|
||||
require.Equal(t, []int64{801, 802}, upstream.accountHits())
|
||||
require.NotContains(t, recorder.Body.String(), "rate limited")
|
||||
})
|
||||
}
|
||||
|
||||
func TestGrokOAuthCredentialFailoverAcrossHTTPHandlers(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
endpoints := []struct {
|
||||
name string
|
||||
method string
|
||||
path string
|
||||
body string
|
||||
}{
|
||||
{name: "messages", method: http.MethodPost, path: "/openai/v1/messages", body: `{"model":"grok","max_tokens":16,"messages":[{"role":"user","content":"hello"}]}`},
|
||||
{name: "chat completions", method: http.MethodPost, path: "/openai/v1/chat/completions", body: `{"model":"grok","messages":[{"role":"user","content":"hello"}],"stream":false}`},
|
||||
{name: "chat completions raw fallback", method: http.MethodPost, path: "/openai/v1/chat/completions", body: `{"model":"grok","messages":[{"role":"user","content":"hello"}],"stop":["END"],"stream":false}`},
|
||||
{name: "grok media", method: http.MethodGet, path: "/openai/v1/videos/request-1"},
|
||||
}
|
||||
|
||||
for _, endpoint := range endpoints {
|
||||
t.Run(endpoint.name+" revoked selects healthy", func(t *testing.T) {
|
||||
_, repo, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "revoked")
|
||||
defer cleanup()
|
||||
recorder := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(endpoint.method, endpoint.path, bytes.NewBufferString(endpoint.body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
router.ServeHTTP(recorder, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, recorder.Code, recorder.Body.String())
|
||||
require.Equal(t, []int64{801}, repo.errorIDs())
|
||||
require.Equal(t, []int64{802}, upstream.accountHits())
|
||||
})
|
||||
|
||||
t.Run(endpoint.name+" all accounts exhausted safely", func(t *testing.T) {
|
||||
_, repo, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "all_revoked")
|
||||
defer cleanup()
|
||||
recorder := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(endpoint.method, endpoint.path, bytes.NewBufferString(endpoint.body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
router.ServeHTTP(recorder, req)
|
||||
|
||||
require.Equal(t, http.StatusServiceUnavailable, recorder.Code, recorder.Body.String())
|
||||
require.Contains(t, recorder.Body.String(), service.GrokCredentialUnavailableClientMessage)
|
||||
require.NotContains(t, recorder.Body.String(), "revoked-refresh")
|
||||
require.NotContains(t, recorder.Body.String(), "healthy-refresh")
|
||||
require.Equal(t, []int64{801, 802}, repo.errorIDs())
|
||||
require.Empty(t, upstream.accountHits())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGrokOAuthMissingSelectedRowRetriesHealthyAccountWithoutMutation(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
_, repo, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "missing_row")
|
||||
defer cleanup()
|
||||
recorder := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/openai/v1/responses", bytes.NewBufferString(`{"model":"grok","input":"hello","stream":false}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
router.ServeHTTP(recorder, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, recorder.Code, recorder.Body.String())
|
||||
require.Equal(t, []int64{802}, upstream.accountHits())
|
||||
require.Empty(t, repo.errorIDs())
|
||||
require.Empty(t, repo.setTempIDs)
|
||||
}
|
||||
|
||||
func TestResponsesWebSocketCredentialFailoverLoop(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
dial := func(t *testing.T, router *gin.Engine) (*coderws.Conn, func()) {
|
||||
t.Helper()
|
||||
server := httptest.NewServer(router)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
|
||||
conn, _, err := coderws.Dial(ctx, "ws"+strings.TrimPrefix(server.URL, "http")+"/openai/v1/responses", nil)
|
||||
cancel()
|
||||
require.NoError(t, err)
|
||||
return conn, func() {
|
||||
_ = conn.CloseNow()
|
||||
server.Close()
|
||||
}
|
||||
}
|
||||
writeFirst := func(t *testing.T, conn *coderws.Conn) {
|
||||
t.Helper()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
|
||||
defer cancel()
|
||||
require.NoError(t, conn.Write(ctx, coderws.MessageText, []byte(`{"type":"response.create","model":"grok","input":"hello","stream":false}`)))
|
||||
}
|
||||
|
||||
t.Run("revoked account selects healthy account", func(t *testing.T) {
|
||||
_, repo, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "revoked")
|
||||
defer cleanup()
|
||||
conn, closeConn := dial(t, router)
|
||||
defer closeConn()
|
||||
writeFirst(t, conn)
|
||||
|
||||
readCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
_, payload, err := conn.Read(readCtx)
|
||||
cancel()
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, string(payload), "resp_healthy")
|
||||
require.Equal(t, []int64{801}, repo.errorIDs())
|
||||
require.Equal(t, 2, repo.selectorCalls())
|
||||
require.Equal(t, []int64{802}, upstream.accountHits())
|
||||
})
|
||||
|
||||
t.Run("provider configuration stops", func(t *testing.T) {
|
||||
_, repo, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "provider")
|
||||
defer cleanup()
|
||||
conn, closeConn := dial(t, router)
|
||||
defer closeConn()
|
||||
writeFirst(t, conn)
|
||||
|
||||
readCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
|
||||
_, _, err := conn.Read(readCtx)
|
||||
cancel()
|
||||
var closeErr coderws.CloseError
|
||||
require.ErrorAs(t, err, &closeErr)
|
||||
require.Contains(t, closeErr.Reason, service.GrokCredentialUnavailableClientMessage)
|
||||
require.Equal(t, 1, repo.selectorCalls())
|
||||
require.Empty(t, upstream.accountHits())
|
||||
})
|
||||
|
||||
t.Run("parent cancellation prevents reselection", func(t *testing.T) {
|
||||
_, repo, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "cancel")
|
||||
defer cleanup()
|
||||
conn, closeConn := dial(t, router)
|
||||
writeFirst(t, conn)
|
||||
select {
|
||||
case <-findHandlerRefresherStarted(router):
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("credential refresh did not start")
|
||||
}
|
||||
closeConn()
|
||||
|
||||
require.Eventually(t, func() bool { return repo.selectorCalls() == 1 }, 2*time.Second, 20*time.Millisecond)
|
||||
require.Empty(t, repo.errorIDs())
|
||||
require.Empty(t, upstream.accountHits())
|
||||
})
|
||||
}
|
||||
|
||||
var handlerRefresherStarted sync.Map
|
||||
|
||||
func findHandlerRefresherStarted(router *gin.Engine) <-chan struct{} {
|
||||
value, _ := handlerRefresherStarted.Load(router)
|
||||
return value.(chan struct{})
|
||||
}
|
||||
|
||||
func newGrokCredentialFailoverHandler(t *testing.T, mode string) (*OpenAIGatewayHandler, *grokCredentialHandlerRepo, *grokCredentialHandlerUpstream, *gin.Engine, func()) {
|
||||
t.Helper()
|
||||
groupID := int64(901)
|
||||
accounts := []service.Account{
|
||||
{
|
||||
ID: 801, Name: "revoked", Platform: service.PlatformGrok, Type: service.AccountTypeOAuth,
|
||||
Status: service.StatusActive, Schedulable: true, Concurrency: 1, Priority: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "expired", "refresh_token": "revoked-refresh",
|
||||
"expires_at": time.Now().Add(-time.Minute).UTC().Format(time.RFC3339),
|
||||
},
|
||||
},
|
||||
{
|
||||
ID: 802, Name: "healthy", Platform: service.PlatformGrok, Type: service.AccountTypeOAuth,
|
||||
Status: service.StatusActive, Schedulable: true, Concurrency: 1, Priority: 2,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "healthy-access", "refresh_token": "healthy-refresh",
|
||||
"expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
|
||||
},
|
||||
},
|
||||
}
|
||||
if mode == "postmap_cancel" || mode == "first_429" || mode == "all_429" || mode == "mixed_429_500" || mode == "mixed_500_429" || mode == "oauth_429_apikey_500" {
|
||||
accounts[0].Credentials["expires_at"] = time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339)
|
||||
}
|
||||
if mode == "all_429" || mode == "mixed_429_500" || mode == "mixed_500_429" || mode == "oauth_429_apikey_500" {
|
||||
accounts = append(accounts, service.Account{
|
||||
ID: 803, Name: "untried-healthy", Platform: service.PlatformGrok, Type: service.AccountTypeOAuth,
|
||||
Status: service.StatusActive, Schedulable: true, Concurrency: 1, Priority: 3,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "untried-healthy-access", "refresh_token": "untried-healthy-refresh",
|
||||
"expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
|
||||
},
|
||||
})
|
||||
}
|
||||
if mode == "oauth_429_apikey_500" {
|
||||
accounts[1].Type = service.AccountTypeAPIKey
|
||||
accounts[1].Credentials = map[string]any{"api_key": "third-party-key"}
|
||||
}
|
||||
if mode == "all_revoked" {
|
||||
accounts[1].Credentials["expires_at"] = time.Now().Add(-time.Minute).UTC().Format(time.RFC3339)
|
||||
}
|
||||
repo := &grokCredentialHandlerRepo{accounts: accounts, missingOnGet: map[int64]bool{}}
|
||||
if mode == "missing_row" {
|
||||
repo.missingOnGet[801] = true
|
||||
}
|
||||
if mode == "mutation_set_error" {
|
||||
repo.setErrorErr = errors.New("database write failed")
|
||||
}
|
||||
if mode == "mutation_temp" {
|
||||
repo.setTempErr = errors.New("database write failed")
|
||||
}
|
||||
refresher := &grokCredentialHandlerRefresher{mode: mode, started: make(chan struct{})}
|
||||
tokenCache := &grokCredentialHandlerTokenCache{}
|
||||
if mode == "mutation_cache" {
|
||||
tokenCache.deleteErr = errors.New("cache delete failed")
|
||||
}
|
||||
var provider *service.GrokTokenProvider
|
||||
if mode != "nil_provider" {
|
||||
provider = service.NewGrokTokenProvider(repo, tokenCache)
|
||||
provider.SetRefreshAPI(service.NewOAuthRefreshAPI(repo, tokenCache), refresher)
|
||||
}
|
||||
upstream := &grokCredentialHandlerUpstream{}
|
||||
switch mode {
|
||||
case "first_429":
|
||||
upstream.rateLimitIDs = map[int64]bool{801: true}
|
||||
case "all_429":
|
||||
upstream.rateLimitIDs = map[int64]bool{801: true, 802: true}
|
||||
case "mixed_429_500":
|
||||
upstream.rateLimitIDs = map[int64]bool{801: true}
|
||||
upstream.failureStatus = map[int64]int{802: http.StatusInternalServerError}
|
||||
case "mixed_500_429":
|
||||
upstream.failureStatus = map[int64]int{801: http.StatusInternalServerError}
|
||||
upstream.rateLimitIDs = map[int64]bool{802: true}
|
||||
case "oauth_429_apikey_500":
|
||||
upstream.rateLimitIDs = map[int64]bool{801: true}
|
||||
upstream.failureStatus = map[int64]int{802: http.StatusInternalServerError}
|
||||
}
|
||||
cfg := &config.Config{RunMode: config.RunModeSimple}
|
||||
cfg.Gateway.MaxAccountSwitches = 3
|
||||
billingCache := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil)
|
||||
gateway := service.NewOpenAIGatewayService(
|
||||
repo, nil, nil, nil, nil, nil, nil, cfg, nil, nil,
|
||||
service.NewBillingService(cfg, nil), nil, billingCache, upstream,
|
||||
&service.DeferredService{}, nil, provider, nil, nil, nil, nil, nil,
|
||||
)
|
||||
cache := &concurrencyCacheMock{
|
||||
acquireUserSlotFn: func(context.Context, int64, int, string) (bool, error) { return true, nil },
|
||||
acquireAccountSlotFn: func(context.Context, int64, int, string) (bool, error) { return true, nil },
|
||||
}
|
||||
h := NewOpenAIGatewayHandler(gateway, service.NewConcurrencyService(cache), billingCache, &service.APIKeyService{}, nil, nil, nil, nil, cfg)
|
||||
apiKey := &service.APIKey{
|
||||
ID: 902, GroupID: &groupID,
|
||||
User: &service.User{ID: 903, Status: service.StatusActive},
|
||||
Group: &service.Group{ID: groupID, Platform: service.PlatformGrok, Status: service.StatusActive},
|
||||
}
|
||||
router := gin.New()
|
||||
router.Use(func(c *gin.Context) {
|
||||
c.Set(string(middleware.ContextKeyAPIKey), apiKey)
|
||||
c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: apiKey.User.ID, Concurrency: 1})
|
||||
c.Next()
|
||||
})
|
||||
router.POST("/openai/v1/responses", h.Responses)
|
||||
router.GET("/openai/v1/responses", h.ResponsesWebSocket)
|
||||
router.POST("/openai/v1/messages", h.Messages)
|
||||
router.POST("/openai/v1/chat/completions", h.ChatCompletions)
|
||||
router.GET("/openai/v1/videos/:request_id", h.GrokVideoStatus)
|
||||
handlerRefresherStarted.Store(router, refresher.started)
|
||||
cleanup := func() {
|
||||
handlerRefresherStarted.Delete(router)
|
||||
billingCache.Stop()
|
||||
}
|
||||
return h, repo, upstream, router, cleanup
|
||||
}
|
||||
@@ -0,0 +1,211 @@
|
||||
//go:build unit
|
||||
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestGatewayChatCredentialStopDoesNotSelectAnotherAccountAndReturnsSafe503(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
stopErr := &service.UpstreamFailoverError{
|
||||
Stage: service.GatewayFailureStageAccountAuth,
|
||||
Scope: service.GatewayFailureScopeProvider,
|
||||
Reason: service.GrokCredentialReasonProviderConfig,
|
||||
NextAccountAction: service.NextAccountStop,
|
||||
ClientStatusCode: http.StatusTeapot,
|
||||
ClientMessage: "invalid_client client_secret=must-not-leak",
|
||||
}
|
||||
state := NewFailoverState(3, false)
|
||||
action := state.HandleFailoverError(context.Background(), &mockTempUnscheduler{}, 71, service.PlatformGrok, 0, stopErr)
|
||||
|
||||
require.Equal(t, FailoverExhausted, action)
|
||||
require.Zero(t, state.SwitchCount)
|
||||
require.Empty(t, state.FailedAccountIDs)
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
(&GatewayHandler{}).handleCCFailoverExhausted(c, state.LastFailoverErr, false)
|
||||
|
||||
require.Equal(t, http.StatusServiceUnavailable, recorder.Code)
|
||||
require.Contains(t, recorder.Body.String(), service.GrokCredentialUnavailableClientMessage)
|
||||
require.NotContains(t, recorder.Body.String(), "invalid_client")
|
||||
require.NotContains(t, recorder.Body.String(), "client_secret")
|
||||
}
|
||||
|
||||
func TestGatewayChatInferenceExhaustionRestoresRetryAfter(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
|
||||
(&GatewayHandler{}).handleCCFailoverExhausted(c, &service.UpstreamFailoverError{
|
||||
StatusCode: http.StatusTooManyRequests,
|
||||
ResponseHeaders: http.Header{"Retry-After": []string{"45"}},
|
||||
}, false)
|
||||
|
||||
require.Equal(t, http.StatusTooManyRequests, recorder.Code)
|
||||
require.Equal(t, "45", recorder.Header().Get("Retry-After"))
|
||||
}
|
||||
|
||||
func TestCredentialFailoverExhaustionReturnsFixedSafe503(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
h := &OpenAIGatewayHandler{}
|
||||
|
||||
h.handleFailoverExhausted(c, &service.UpstreamFailoverError{
|
||||
Stage: service.GatewayFailureStageAccountAuth,
|
||||
Scope: service.GatewayFailureScopeAccount,
|
||||
Reason: service.GrokCredentialReasonRevoked,
|
||||
NextAccountAction: service.NextAccountRetry,
|
||||
ClientStatusCode: http.StatusTeapot,
|
||||
ClientMessage: "invalid_grant refresh_token=must-not-leak",
|
||||
}, false)
|
||||
|
||||
require.Equal(t, http.StatusServiceUnavailable, recorder.Code)
|
||||
require.Contains(t, recorder.Body.String(), service.GrokCredentialUnavailableClientMessage)
|
||||
require.NotContains(t, strings.ToLower(recorder.Body.String()), "invalid_grant")
|
||||
require.NotContains(t, strings.ToLower(recorder.Body.String()), "refresh_token")
|
||||
require.NotContains(t, recorder.Body.String(), "must-not-leak")
|
||||
}
|
||||
|
||||
func TestInferenceFailoverExhaustionRestoresRetryAfter(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
h := &OpenAIGatewayHandler{}
|
||||
|
||||
h.handleFailoverExhausted(c, &service.UpstreamFailoverError{
|
||||
StatusCode: http.StatusTooManyRequests,
|
||||
ResponseHeaders: http.Header{"Retry-After": []string{"17"}},
|
||||
}, false)
|
||||
|
||||
require.Equal(t, http.StatusTooManyRequests, recorder.Code)
|
||||
require.Equal(t, "17", recorder.Header().Get("Retry-After"))
|
||||
}
|
||||
|
||||
func TestFailoverExhaustionRejectsSecretBearingRetryAfter(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
h := &OpenAIGatewayHandler{}
|
||||
|
||||
h.handleFailoverExhausted(c, &service.UpstreamFailoverError{
|
||||
StatusCode: http.StatusTooManyRequests,
|
||||
ResponseHeaders: http.Header{"Retry-After": []string{"refresh_token=must-not-leak"}},
|
||||
}, false)
|
||||
|
||||
require.Equal(t, http.StatusTooManyRequests, recorder.Code)
|
||||
require.Empty(t, recorder.Header().Get("Retry-After"))
|
||||
require.NotContains(t, recorder.Body.String(), "must-not-leak")
|
||||
}
|
||||
|
||||
func TestFailoverExhaustionRejectsFarFutureRetryAfterDate(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
h := &OpenAIGatewayHandler{}
|
||||
|
||||
h.handleFailoverExhausted(c, &service.UpstreamFailoverError{
|
||||
StatusCode: http.StatusTooManyRequests,
|
||||
ResponseHeaders: http.Header{
|
||||
"Retry-After": []string{time.Now().Add(30 * 24 * time.Hour).UTC().Format(http.TimeFormat)},
|
||||
},
|
||||
}, false)
|
||||
|
||||
require.Equal(t, http.StatusTooManyRequests, recorder.Code)
|
||||
require.Empty(t, recorder.Header().Get("Retry-After"))
|
||||
}
|
||||
|
||||
func TestFailoverExhaustionAllowsBoundedRetryAfterDate(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
h := &OpenAIGatewayHandler{}
|
||||
retryAfter := time.Now().Add(time.Hour).UTC().Format(http.TimeFormat)
|
||||
|
||||
h.handleFailoverExhausted(c, &service.UpstreamFailoverError{
|
||||
StatusCode: http.StatusTooManyRequests,
|
||||
ResponseHeaders: http.Header{"Retry-After": []string{retryAfter}},
|
||||
}, false)
|
||||
|
||||
require.Equal(t, http.StatusTooManyRequests, recorder.Code)
|
||||
require.Equal(t, retryAfter, recorder.Header().Get("Retry-After"))
|
||||
}
|
||||
|
||||
func TestOpsClassificationTreatsCredentialFailureAsAuthNotInference(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Set(service.OpsUpstreamStatusCodeKey, http.StatusForbidden)
|
||||
c.Set(service.OpsUpstreamErrorMessageKey, "stale inference message")
|
||||
c.Set(service.OpsUpstreamErrorDetailKey, "stale inference detail")
|
||||
c.Set(service.OpsUpstreamErrorsKey, []*service.OpsUpstreamErrorEvent{
|
||||
{Stage: string(service.GatewayFailureStageInference), UpstreamStatusCode: http.StatusForbidden, Message: "stale inference message", Detail: "stale inference detail"},
|
||||
{
|
||||
Stage: string(service.GatewayFailureStageAccountAuth),
|
||||
Scope: string(service.GatewayFailureScopeAccount),
|
||||
Reason: string(service.GrokCredentialReasonRevoked),
|
||||
UpstreamStatusCode: 0,
|
||||
Message: "Grok OAuth credentials require account action",
|
||||
},
|
||||
})
|
||||
|
||||
phase, _, owner, source := classifyOpsErrorLog(c, "upstream_error", service.GrokCredentialUnavailableClientMessage, "", http.StatusServiceUnavailable)
|
||||
require.Equal(t, "account_auth", phase)
|
||||
require.Equal(t, "provider", owner)
|
||||
require.Equal(t, "gateway", source)
|
||||
|
||||
entry := &service.OpsInsertErrorLogInput{}
|
||||
applyOpsUpstreamFieldsFromContext(c, entry)
|
||||
require.NotNil(t, entry.UpstreamStatusCode)
|
||||
require.Zero(t, *entry.UpstreamStatusCode)
|
||||
require.NotNil(t, entry.UpstreamErrorMessage)
|
||||
require.Equal(t, "Grok OAuth credentials require account action", *entry.UpstreamErrorMessage)
|
||||
require.Nil(t, entry.UpstreamErrorDetail)
|
||||
require.Len(t, entry.UpstreamErrors, 2)
|
||||
require.Equal(t, http.StatusForbidden, entry.UpstreamErrors[0].UpstreamStatusCode)
|
||||
}
|
||||
|
||||
func TestOpsRecoveredCredentialFailoverUsesAccountAuthAttribution(t *testing.T) {
|
||||
setupOpsErrorLogTestQueue(t, 2)
|
||||
gin.SetMode(gin.TestMode)
|
||||
ops := service.NewOpsService(nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
|
||||
router := gin.New()
|
||||
router.Use(OpsErrorLoggerMiddleware(ops))
|
||||
router.GET("/openai/v1/responses", func(c *gin.Context) {
|
||||
c.Set(service.OpsUpstreamErrorsKey, []*service.OpsUpstreamErrorEvent{
|
||||
{Stage: string(service.GatewayFailureStageInference), UpstreamStatusCode: http.StatusForbidden, Message: "earlier inference failure"},
|
||||
{
|
||||
Stage: string(service.GatewayFailureStageAccountAuth), Scope: string(service.GatewayFailureScopeAccount),
|
||||
Reason: string(service.GrokCredentialReasonRevoked), Message: "Grok OAuth credentials require account action",
|
||||
},
|
||||
})
|
||||
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||||
})
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
router.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/openai/v1/responses", nil))
|
||||
|
||||
require.Equal(t, http.StatusOK, recorder.Code)
|
||||
require.Equal(t, int64(1), OpsErrorLogQueueLength())
|
||||
job := <-opsErrorLogQueue
|
||||
require.Equal(t, "account_auth", job.entry.ErrorPhase)
|
||||
require.Equal(t, "provider", job.entry.ErrorOwner)
|
||||
require.Equal(t, "gateway", job.entry.ErrorSource)
|
||||
require.Contains(t, job.entry.ErrorMessage, "Recovered account authentication failure")
|
||||
require.NotContains(t, job.entry.ErrorMessage, "403")
|
||||
require.NotContains(t, job.entry.ErrorMessage, "earlier inference failure")
|
||||
require.NotNil(t, job.entry.UpstreamStatusCode)
|
||||
require.Zero(t, *job.entry.UpstreamStatusCode)
|
||||
require.Len(t, job.entry.UpstreamErrors, 2)
|
||||
require.Equal(t, http.StatusForbidden, job.entry.UpstreamErrors[0].UpstreamStatusCode)
|
||||
}
|
||||
@@ -332,8 +332,12 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
|
||||
failedAccountIDs := make(map[int64]struct{})
|
||||
sameAccountRetryCount := make(map[int64]int)
|
||||
var lastFailoverErr *service.UpstreamFailoverError
|
||||
var oauth429FailoverState service.OpenAIOAuth429FailoverState
|
||||
|
||||
for {
|
||||
if failoverClientGone(c) {
|
||||
return
|
||||
}
|
||||
// Select account supporting the requested model
|
||||
reqLog.Debug("openai.account_selecting", zap.Int("excluded_account_count", len(failedAccountIDs)))
|
||||
selection, scheduleDecision, err := h.gatewayService.SelectAccountWithSchedulerForCapability(
|
||||
@@ -350,6 +354,10 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
|
||||
requestPlatform,
|
||||
)
|
||||
if err != nil {
|
||||
if failoverClientGone(c) {
|
||||
reqLog.Info("openai.account_select_aborted_client_disconnected", zap.Error(err))
|
||||
return
|
||||
}
|
||||
reqLog.Warn("openai.account_select_failed",
|
||||
zap.Error(openAICompatibleSelectionErrorForLog(err, requestPlatform)),
|
||||
zap.Int("excluded_account_count", len(failedAccountIDs)),
|
||||
@@ -443,11 +451,24 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
|
||||
} else {
|
||||
var failoverErr *service.UpstreamFailoverError
|
||||
if errors.As(err, &failoverErr) {
|
||||
if failoverClientGone(c) {
|
||||
reqLog.Info("openai.failover_aborted_client_disconnected",
|
||||
zap.Int64("account_id", account.ID),
|
||||
zap.Int("upstream_status", failoverErr.StatusCode),
|
||||
)
|
||||
return
|
||||
}
|
||||
if service.OpenAICompactKeepaliveAdjustedWrittenSize(c) != writerSizeBeforeForward {
|
||||
h.handleFailoverExhausted(c, failoverErr, true)
|
||||
return
|
||||
}
|
||||
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil)
|
||||
if failoverErr.ShouldReportAccountScheduleFailure() {
|
||||
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil)
|
||||
}
|
||||
if !failoverErr.ShouldRetryNextAccount() {
|
||||
h.handleFailoverExhausted(c, failoverErr, streamStarted)
|
||||
return
|
||||
}
|
||||
// 池模式:同账号重试
|
||||
if failoverErr.RetryableOnSameAccount {
|
||||
retryLimit := account.GetPoolModeRetryCount()
|
||||
@@ -475,16 +496,27 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
switchCount++
|
||||
if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount) {
|
||||
if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount, &oauth429FailoverState) {
|
||||
h.handleFailoverExhausted(c, failoverErr, streamStarted)
|
||||
return
|
||||
}
|
||||
reqLog.Warn("openai.upstream_failover_switching",
|
||||
failoverSwitchFields := []zap.Field{
|
||||
zap.Int64("account_id", account.ID),
|
||||
zap.Int("upstream_status", failoverErr.StatusCode),
|
||||
zap.Int("switch_count", switchCount),
|
||||
zap.Int("max_switches", maxAccountSwitches),
|
||||
)
|
||||
}
|
||||
if account.Proxy != nil {
|
||||
failoverSwitchFields = append(failoverSwitchFields,
|
||||
zap.Int64("proxy_id", account.Proxy.ID),
|
||||
zap.String("proxy_name", account.Proxy.Name),
|
||||
zap.String("proxy_host", account.Proxy.Host),
|
||||
zap.Int("proxy_port", account.Proxy.Port),
|
||||
)
|
||||
} else if account.ProxyID != nil {
|
||||
failoverSwitchFields = append(failoverSwitchFields, zap.Int64p("proxy_id", account.ProxyID))
|
||||
}
|
||||
reqLog.Warn("openai.upstream_failover_switching", failoverSwitchFields...)
|
||||
continue
|
||||
}
|
||||
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil)
|
||||
@@ -832,9 +864,13 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) {
|
||||
failedAccountIDs := make(map[int64]struct{})
|
||||
sameAccountRetryCount := make(map[int64]int)
|
||||
var lastFailoverErr *service.UpstreamFailoverError
|
||||
var oauth429FailoverState service.OpenAIOAuth429FailoverState
|
||||
effectiveMappedModel := preferredMappedModel
|
||||
|
||||
for {
|
||||
if failoverClientGone(c) {
|
||||
return
|
||||
}
|
||||
currentRoutingModel := routingModel
|
||||
if effectiveMappedModel != "" {
|
||||
currentRoutingModel = effectiveMappedModel
|
||||
@@ -854,6 +890,10 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) {
|
||||
requestPlatform,
|
||||
)
|
||||
if err != nil {
|
||||
if failoverClientGone(c) {
|
||||
reqLog.Info("openai_messages.account_select_aborted_client_disconnected", zap.Error(err))
|
||||
return
|
||||
}
|
||||
reqLog.Warn("openai_messages.account_select_failed",
|
||||
zap.Error(openAICompatibleSelectionErrorForLog(err, requestPlatform)),
|
||||
zap.Int("excluded_account_count", len(failedAccountIDs)),
|
||||
@@ -935,11 +975,24 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) {
|
||||
} else {
|
||||
var failoverErr *service.UpstreamFailoverError
|
||||
if errors.As(err, &failoverErr) {
|
||||
if failoverClientGone(c) {
|
||||
reqLog.Info("openai_messages.failover_aborted_client_disconnected",
|
||||
zap.Int64("account_id", account.ID),
|
||||
zap.Int("upstream_status", failoverErr.StatusCode),
|
||||
)
|
||||
return
|
||||
}
|
||||
if c.Writer.Size() != writerSizeBeforeForward {
|
||||
h.handleAnthropicFailoverExhausted(c, failoverErr, true)
|
||||
return
|
||||
}
|
||||
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil)
|
||||
if failoverErr.ShouldReportAccountScheduleFailure() {
|
||||
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil)
|
||||
}
|
||||
if !failoverErr.ShouldRetryNextAccount() {
|
||||
h.handleAnthropicFailoverExhausted(c, failoverErr, streamStarted)
|
||||
return
|
||||
}
|
||||
// 池模式:同账号重试
|
||||
if failoverErr.RetryableOnSameAccount {
|
||||
retryLimit := account.GetPoolModeRetryCount()
|
||||
@@ -967,7 +1020,7 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
switchCount++
|
||||
if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount) {
|
||||
if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount, &oauth429FailoverState) {
|
||||
h.handleAnthropicFailoverExhausted(c, failoverErr, streamStarted)
|
||||
return
|
||||
}
|
||||
@@ -1093,6 +1146,14 @@ func (h *OpenAIGatewayHandler) anthropicStreamingAwareError(c *gin.Context, stat
|
||||
|
||||
// handleAnthropicFailoverExhausted maps upstream failover errors to Anthropic format.
|
||||
func (h *OpenAIGatewayHandler) handleAnthropicFailoverExhausted(c *gin.Context, failoverErr *service.UpstreamFailoverError, streamStarted bool) {
|
||||
if failoverErr != nil {
|
||||
copyFailoverRetryAfter(c, failoverErr.ResponseHeaders)
|
||||
}
|
||||
if failoverErr != nil && failoverErr.IsCredentialFailure() {
|
||||
status, message := credentialFailoverClientResponse(failoverErr)
|
||||
h.anthropicStreamingAwareError(c, status, "api_error", message, streamStarted)
|
||||
return
|
||||
}
|
||||
status, errType, errMsg := h.mapUpstreamError(failoverErr.StatusCode)
|
||||
h.anthropicStreamingAwareError(c, status, errType, errMsg, streamStarted)
|
||||
}
|
||||
@@ -1464,8 +1525,50 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
switchCount := 0
|
||||
failedAccountIDs := make(map[int64]struct{})
|
||||
var lastFailoverErr *service.UpstreamFailoverError
|
||||
var oauth429FailoverState service.OpenAIOAuth429FailoverState
|
||||
handleWSFailover := func(account *service.Account, failoverErr *service.UpstreamFailoverError) bool {
|
||||
if ctx.Err() != nil {
|
||||
return false
|
||||
}
|
||||
if failoverErr.ShouldReportAccountScheduleFailure() {
|
||||
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil)
|
||||
}
|
||||
releaseAccountSlot()
|
||||
if !failoverErr.ShouldRetryNextAccount() {
|
||||
closeOpenAIWSFailoverExhausted(wsConn, failoverErr)
|
||||
return false
|
||||
}
|
||||
if ctx.Err() != nil {
|
||||
return false
|
||||
}
|
||||
h.gatewayService.RecordOpenAIAccountSwitch()
|
||||
failedAccountIDs[account.ID] = struct{}{}
|
||||
lastFailoverErr = failoverErr
|
||||
if switchCount >= maxAccountSwitches {
|
||||
closeOpenAIWSFailoverExhausted(wsConn, failoverErr)
|
||||
return false
|
||||
}
|
||||
switchCount++
|
||||
if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount, &oauth429FailoverState) {
|
||||
closeOpenAIWSFailoverExhausted(wsConn, failoverErr)
|
||||
return false
|
||||
}
|
||||
reqLog.Warn("openai.websocket_upstream_failover_switching",
|
||||
zap.Int64("account_id", account.ID),
|
||||
zap.Int("upstream_status", failoverErr.StatusCode),
|
||||
zap.Int("switch_count", switchCount),
|
||||
zap.Int("max_switches", maxAccountSwitches),
|
||||
)
|
||||
if ctx.Err() != nil {
|
||||
return false
|
||||
}
|
||||
return ensureUserSlotHeld()
|
||||
}
|
||||
|
||||
for {
|
||||
if ctx.Err() != nil {
|
||||
return
|
||||
}
|
||||
reqLog.Debug("openai.websocket_account_selecting", zap.Int("excluded_account_count", len(failedAccountIDs)))
|
||||
selection, scheduleDecision, err := h.gatewayService.SelectAccountWithSchedulerForCapability(
|
||||
ctx,
|
||||
@@ -1533,9 +1636,19 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
reqLog.Warn("openai.websocket_bind_sticky_session_failed", zap.Int64("account_id", account.ID), zap.Error(err))
|
||||
}
|
||||
|
||||
token, _, err := h.gatewayService.GetAccessToken(ctx, account)
|
||||
token, _, err := h.gatewayService.GetRequestCredential(ctx, c, account)
|
||||
if err != nil {
|
||||
reqLog.Warn("openai.websocket_get_access_token_failed", zap.Int64("account_id", account.ID), zap.Error(err))
|
||||
if ctx.Err() != nil {
|
||||
return
|
||||
}
|
||||
var failoverErr *service.UpstreamFailoverError
|
||||
if errors.As(err, &failoverErr) {
|
||||
if handleWSFailover(account, failoverErr) {
|
||||
continue
|
||||
}
|
||||
return
|
||||
}
|
||||
closeOpenAIClientWS(wsConn, coderws.StatusInternalError, "failed to get access token")
|
||||
return
|
||||
}
|
||||
@@ -1692,30 +1805,10 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
if err := h.gatewayService.ProxyResponsesWebSocketFromClient(ctx, c, wsConn, account, token, wsFirstMessage, hooks); err != nil {
|
||||
var failoverErr *service.UpstreamFailoverError
|
||||
if errors.As(err, &failoverErr) {
|
||||
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil)
|
||||
releaseAccountSlot()
|
||||
failedAccountIDs[account.ID] = struct{}{}
|
||||
lastFailoverErr = failoverErr
|
||||
if switchCount >= maxAccountSwitches {
|
||||
closeOpenAIWSFailoverExhausted(wsConn, failoverErr)
|
||||
return
|
||||
if handleWSFailover(account, failoverErr) {
|
||||
continue
|
||||
}
|
||||
switchCount++
|
||||
if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount) {
|
||||
closeOpenAIWSFailoverExhausted(wsConn, failoverErr)
|
||||
return
|
||||
}
|
||||
h.gatewayService.RecordOpenAIAccountSwitch()
|
||||
reqLog.Warn("openai.websocket_upstream_failover_switching",
|
||||
zap.Int64("account_id", account.ID),
|
||||
zap.Int("upstream_status", failoverErr.StatusCode),
|
||||
zap.Int("switch_count", switchCount),
|
||||
zap.Int("max_switches", maxAccountSwitches),
|
||||
)
|
||||
if !ensureUserSlotHeld() {
|
||||
return
|
||||
}
|
||||
continue
|
||||
return
|
||||
}
|
||||
|
||||
if errors.Is(context.Cause(ctx), service.ErrOpenAIWSIngressLeaseLost) {
|
||||
@@ -1739,12 +1832,23 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
|
||||
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil)
|
||||
closeStatus, closeReason := summarizeWSCloseErrorForLog(err)
|
||||
reqLog.Warn("openai.websocket_proxy_failed",
|
||||
proxyFailedFields := []zap.Field{
|
||||
zap.Int64("account_id", account.ID),
|
||||
zap.Error(err),
|
||||
zap.String("close_status", closeStatus),
|
||||
zap.String("close_reason", closeReason),
|
||||
)
|
||||
}
|
||||
if account.Proxy != nil {
|
||||
proxyFailedFields = append(proxyFailedFields,
|
||||
zap.Int64("proxy_id", account.Proxy.ID),
|
||||
zap.String("proxy_name", account.Proxy.Name),
|
||||
zap.String("proxy_host", account.Proxy.Host),
|
||||
zap.Int("proxy_port", account.Proxy.Port),
|
||||
)
|
||||
} else if account.ProxyID != nil {
|
||||
proxyFailedFields = append(proxyFailedFields, zap.Int64p("proxy_id", account.ProxyID))
|
||||
}
|
||||
reqLog.Warn("openai.websocket_proxy_failed", proxyFailedFields...)
|
||||
if errors.As(err, &closeErr) {
|
||||
closeOpenAIClientWS(wsConn, closeErr.StatusCode(), closeErr.Reason())
|
||||
return
|
||||
@@ -1946,6 +2050,16 @@ func (h *OpenAIGatewayHandler) handleConcurrencyError(c *gin.Context, err error,
|
||||
}
|
||||
|
||||
func (h *OpenAIGatewayHandler) handleFailoverExhausted(c *gin.Context, failoverErr *service.UpstreamFailoverError, streamStarted bool) {
|
||||
if failoverErr == nil {
|
||||
h.handleFailoverExhaustedSimple(c, http.StatusBadGateway, streamStarted)
|
||||
return
|
||||
}
|
||||
copyFailoverRetryAfter(c, failoverErr.ResponseHeaders)
|
||||
if failoverErr.IsCredentialFailure() {
|
||||
status, message := credentialFailoverClientResponse(failoverErr)
|
||||
h.handleStreamingAwareError(c, status, "upstream_error", message, streamStarted)
|
||||
return
|
||||
}
|
||||
statusCode := failoverErr.StatusCode
|
||||
responseBody := failoverErr.ResponseBody
|
||||
if service.IsOpenAISilentRefusalErrorBody(responseBody) {
|
||||
@@ -1987,6 +2101,41 @@ func (h *OpenAIGatewayHandler) handleFailoverExhausted(c *gin.Context, failoverE
|
||||
h.handleStreamingAwareError(c, status, errType, errMsg, streamStarted)
|
||||
}
|
||||
|
||||
func credentialFailoverClientResponse(failoverErr *service.UpstreamFailoverError) (int, string) {
|
||||
_ = failoverErr
|
||||
return http.StatusServiceUnavailable, service.GrokCredentialUnavailableClientMessage
|
||||
}
|
||||
|
||||
func copyFailoverRetryAfter(c *gin.Context, headers http.Header) {
|
||||
if c == nil || headers == nil {
|
||||
return
|
||||
}
|
||||
retryAfter := strings.TrimSpace(headers.Get("Retry-After"))
|
||||
if retryAfter == "" || len(retryAfter) > 128 || strings.ContainsAny(retryAfter, "\r\n") || !isSafeRetryAfter(retryAfter) {
|
||||
return
|
||||
}
|
||||
c.Header("Retry-After", retryAfter)
|
||||
}
|
||||
|
||||
func isSafeRetryAfter(value string) bool {
|
||||
digitsOnly := true
|
||||
for _, char := range value {
|
||||
if char < '0' || char > '9' {
|
||||
digitsOnly = false
|
||||
break
|
||||
}
|
||||
}
|
||||
if digitsOnly {
|
||||
seconds, err := strconv.ParseUint(value, 10, 32)
|
||||
return err == nil && seconds <= uint64((7*24*time.Hour)/time.Second)
|
||||
}
|
||||
retryAt, err := http.ParseTime(value)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
return !retryAt.After(time.Now().Add(7 * 24 * time.Hour))
|
||||
}
|
||||
|
||||
// handleFailoverExhaustedSimple 简化版本,用于没有响应体的情况
|
||||
func (h *OpenAIGatewayHandler) handleFailoverExhaustedSimple(c *gin.Context, statusCode int, streamStarted bool) {
|
||||
status, errType, errMsg := h.mapUpstreamError(statusCode)
|
||||
@@ -2194,6 +2343,10 @@ func closeOpenAIWSFailoverExhausted(conn *coderws.Conn, failoverErr *service.Ups
|
||||
closeOpenAIClientWS(conn, coderws.StatusInternalError, "upstream websocket proxy failed")
|
||||
return
|
||||
}
|
||||
if failoverErr.Stage == service.GatewayFailureStageAccountAuth {
|
||||
closeOpenAIClientWS(conn, coderws.StatusTryAgainLater, service.GrokCredentialUnavailableClientMessage)
|
||||
return
|
||||
}
|
||||
switch failoverErr.StatusCode {
|
||||
case http.StatusTooManyRequests:
|
||||
closeOpenAIClientWS(conn, coderws.StatusTryAgainLater, "upstream rate limit exceeded, please retry later")
|
||||
|
||||
@@ -145,6 +145,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) {
|
||||
stopJSONKeepalive := func() {}
|
||||
jsonKeepaliveStarted := false
|
||||
defer func() { stopJSONKeepalive() }()
|
||||
var oauth429FailoverState service.OpenAIOAuth429FailoverState
|
||||
|
||||
for {
|
||||
reqLog.Debug("openai.images.account_selecting", zap.Int("excluded_account_count", len(failedAccountIDs)))
|
||||
@@ -157,6 +158,10 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) {
|
||||
parsed.RequiredCapability,
|
||||
)
|
||||
if err != nil {
|
||||
if failoverClientGone(c) {
|
||||
reqLog.Info("openai.images.account_select_aborted_client_disconnected", zap.Error(err))
|
||||
return
|
||||
}
|
||||
reqLog.Warn("openai.images.account_select_failed",
|
||||
zap.Error(err),
|
||||
zap.Int("excluded_account_count", len(failedAccountIDs)),
|
||||
@@ -273,6 +278,13 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) {
|
||||
h.handleFailoverExhausted(c, failoverErr, true)
|
||||
return
|
||||
}
|
||||
if failoverClientGone(c) {
|
||||
reqLog.Info("openai.images.failover_aborted_client_disconnected",
|
||||
zap.Int64("account_id", account.ID),
|
||||
zap.Int("upstream_status", failoverErr.StatusCode),
|
||||
)
|
||||
return
|
||||
}
|
||||
if failoverErr.RetryableOnSameAccount {
|
||||
retryLimit := account.GetPoolModeRetryCount()
|
||||
if sameAccountRetryCount[account.ID] < retryLimit {
|
||||
@@ -299,7 +311,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
switchCount++
|
||||
if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount) {
|
||||
if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount, &oauth429FailoverState) {
|
||||
h.handleFailoverExhausted(c, failoverErr, streamStarted)
|
||||
return
|
||||
}
|
||||
|
||||
@@ -0,0 +1,194 @@
|
||||
//go:build unit
|
||||
|
||||
package handler
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
// openAIResponsesFailoverCancelUpstream 固定返回 HTTP 520,可在首次上游调用时
|
||||
// 触发回调(用于模拟“上游在途期间客户端断开”)。
|
||||
type openAIResponsesFailoverCancelUpstream struct {
|
||||
service.HTTPUpstream
|
||||
mu sync.Mutex
|
||||
accountIDs []int64
|
||||
onFirstDo func()
|
||||
}
|
||||
|
||||
func (u *openAIResponsesFailoverCancelUpstream) Do(_ *http.Request, _ string, accountID int64, _ int) (*http.Response, error) {
|
||||
u.mu.Lock()
|
||||
u.accountIDs = append(u.accountIDs, accountID)
|
||||
first := len(u.accountIDs) == 1
|
||||
u.mu.Unlock()
|
||||
if first && u.onFirstDo != nil {
|
||||
u.onFirstDo()
|
||||
}
|
||||
return &http.Response{
|
||||
StatusCode: 520,
|
||||
Header: http.Header{"Content-Type": []string{"text/html"}},
|
||||
Body: io.NopCloser(bytes.NewBufferString("<html>520: unknown error</html>")),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (u *openAIResponsesFailoverCancelUpstream) calls() []int64 {
|
||||
u.mu.Lock()
|
||||
defer u.mu.Unlock()
|
||||
return append([]int64(nil), u.accountIDs...)
|
||||
}
|
||||
|
||||
func newOpenAIResponsesFailoverTestHandler(t *testing.T, upstream service.HTTPUpstream) *OpenAIGatewayHandler {
|
||||
t.Helper()
|
||||
accounts := []service.Account{
|
||||
{
|
||||
ID: 1,
|
||||
Name: "responses-account-1",
|
||||
Platform: service.PlatformOpenAI,
|
||||
Type: service.AccountTypeOAuth,
|
||||
Status: service.StatusActive,
|
||||
Schedulable: true,
|
||||
Concurrency: 0,
|
||||
Priority: 0,
|
||||
Credentials: map[string]any{"access_token": "token-1"},
|
||||
},
|
||||
{
|
||||
ID: 2,
|
||||
Name: "responses-account-2",
|
||||
Platform: service.PlatformOpenAI,
|
||||
Type: service.AccountTypeOAuth,
|
||||
Status: service.StatusActive,
|
||||
Schedulable: true,
|
||||
Concurrency: 0,
|
||||
Priority: 1,
|
||||
Credentials: map[string]any{"access_token": "token-2"},
|
||||
},
|
||||
}
|
||||
accountRepo := openAIImagesFailoverAccountRepo{accounts: accounts}
|
||||
cfg := &config.Config{RunMode: config.RunModeSimple}
|
||||
gatewayService := service.NewOpenAIGatewayService(
|
||||
accountRepo,
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
cfg,
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
upstream,
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
)
|
||||
billingService := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil)
|
||||
t.Cleanup(billingService.Stop)
|
||||
concurrencyService := service.NewConcurrencyService(nil)
|
||||
handler := NewOpenAIGatewayHandler(
|
||||
gatewayService,
|
||||
concurrencyService,
|
||||
billingService,
|
||||
service.NewAPIKeyService(nil, nil, nil, nil, nil, nil, cfg),
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
cfg,
|
||||
)
|
||||
handler.maxAccountSwitches = 10
|
||||
return handler
|
||||
}
|
||||
|
||||
func newOpenAIResponsesFailoverTestContext(t *testing.T, ctx context.Context) (*gin.Context, *httptest.ResponseRecorder) {
|
||||
t.Helper()
|
||||
groupID := int64(3131)
|
||||
body := []byte(`{"model":"gpt-5.1","stream":false,"input":"hello"}`)
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body))
|
||||
if ctx != nil {
|
||||
req = req.WithContext(ctx)
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = req
|
||||
c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{
|
||||
ID: 99,
|
||||
GroupID: &groupID,
|
||||
Group: &service.Group{
|
||||
ID: groupID,
|
||||
Platform: service.PlatformOpenAI,
|
||||
},
|
||||
User: &service.User{ID: 100},
|
||||
})
|
||||
c.Set(string(middleware2.ContextKeyUser), middleware2.AuthSubject{UserID: 100, Concurrency: 0})
|
||||
return c, rec
|
||||
}
|
||||
|
||||
// TestOpenAIGatewayHandlerResponses_FailoverAbortsWhenClientDisconnected 复现
|
||||
// #4257:客户端在上游请求在途期间断开,上游随后返回可 failover 的 520。
|
||||
// 期望:不再用已取消的 context 重新选号(不触达账号 2)、不把取消误报成
|
||||
// 502 账号耗尽、请求按 499 归类。
|
||||
func TestOpenAIGatewayHandlerResponses_FailoverAbortsWhenClientDisconnected(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
upstream := &openAIResponsesFailoverCancelUpstream{onFirstDo: cancel}
|
||||
handler := newOpenAIResponsesFailoverTestHandler(t, upstream)
|
||||
c, rec := newOpenAIResponsesFailoverTestContext(t, ctx)
|
||||
|
||||
handler.Responses(c)
|
||||
|
||||
require.Equal(t, []int64{1}, upstream.calls(), "客户端断开后不应再切换到账号 2")
|
||||
require.Equal(t, statusClientClosedRequest, c.Writer.Status(), "应按 499 归类")
|
||||
require.Zero(t, rec.Body.Len(), "不应写入 502 错误响应体")
|
||||
|
||||
_, hasFinalUpstreamErr := c.Get(service.OpsUpstreamStatusCodeKey)
|
||||
require.False(t, hasFinalUpstreamErr, "不应记录 failover 耗尽的上游错误终态")
|
||||
|
||||
// 真实发生过的 520 应保留 failover 事件(service 层在返回 failover 错误前记录)
|
||||
rawEvents, ok := c.Get(service.OpsUpstreamErrorsKey)
|
||||
require.True(t, ok)
|
||||
events, ok := rawEvents.([]*service.OpsUpstreamErrorEvent)
|
||||
require.True(t, ok)
|
||||
require.Len(t, events, 1)
|
||||
require.Equal(t, "failover", events[0].Kind)
|
||||
require.Equal(t, 520, events[0].UpstreamStatusCode)
|
||||
}
|
||||
|
||||
// TestOpenAIGatewayHandlerResponses_FailoverContinuesForConnectedClient 回归
|
||||
// 守卫:客户端在线时 failover 行为不变——切换到账号 2,两个账号都 520 后按
|
||||
// 耗尽返回 502。
|
||||
func TestOpenAIGatewayHandlerResponses_FailoverContinuesForConnectedClient(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
upstream := &openAIResponsesFailoverCancelUpstream{}
|
||||
handler := newOpenAIResponsesFailoverTestHandler(t, upstream)
|
||||
c, rec := newOpenAIResponsesFailoverTestContext(t, nil)
|
||||
|
||||
handler.Responses(c)
|
||||
|
||||
require.Equal(t, []int64{1, 2}, upstream.calls(), "在线客户端应正常切换账号")
|
||||
require.Equal(t, http.StatusBadGateway, rec.Code)
|
||||
require.Equal(t, "upstream_error", gjson.GetBytes(rec.Body.Bytes(), "error.type").String())
|
||||
}
|
||||
@@ -721,10 +721,15 @@ func OpsErrorLoggerMiddleware(ops *service.OpsService) gin.HandlerFunc {
|
||||
var upstreamStatusCode *int
|
||||
var upstreamErrorMessage *string
|
||||
var upstreamErrorDetail *string
|
||||
finalAccountAuth := false
|
||||
if len(events) > 0 {
|
||||
last := events[len(events)-1]
|
||||
if last != nil {
|
||||
if last.UpstreamStatusCode > 0 {
|
||||
finalAccountAuth = last.Stage == string(service.GatewayFailureStageAccountAuth)
|
||||
if finalAccountAuth {
|
||||
code := 0
|
||||
upstreamStatusCode = &code
|
||||
} else if last.UpstreamStatusCode > 0 {
|
||||
code := last.UpstreamStatusCode
|
||||
upstreamStatusCode = &code
|
||||
}
|
||||
@@ -737,7 +742,7 @@ func OpsErrorLoggerMiddleware(ops *service.OpsService) gin.HandlerFunc {
|
||||
}
|
||||
}
|
||||
|
||||
if upstreamStatusCode == nil {
|
||||
if !finalAccountAuth && upstreamStatusCode == nil {
|
||||
if v, ok := c.Get(service.OpsUpstreamStatusCodeKey); ok {
|
||||
switch t := v.(type) {
|
||||
case int:
|
||||
@@ -753,7 +758,7 @@ func OpsErrorLoggerMiddleware(ops *service.OpsService) gin.HandlerFunc {
|
||||
}
|
||||
}
|
||||
}
|
||||
if upstreamErrorMessage == nil {
|
||||
if !finalAccountAuth && upstreamErrorMessage == nil {
|
||||
if v, ok := c.Get(service.OpsUpstreamErrorMessageKey); ok {
|
||||
if s, ok := v.(string); ok && strings.TrimSpace(s) != "" {
|
||||
msg := strings.TrimSpace(s)
|
||||
@@ -761,7 +766,7 @@ func OpsErrorLoggerMiddleware(ops *service.OpsService) gin.HandlerFunc {
|
||||
}
|
||||
}
|
||||
}
|
||||
if upstreamErrorDetail == nil {
|
||||
if !finalAccountAuth && upstreamErrorDetail == nil {
|
||||
if v, ok := c.Get(service.OpsUpstreamErrorDetailKey); ok {
|
||||
if s, ok := v.(string); ok && strings.TrimSpace(s) != "" {
|
||||
detail := strings.TrimSpace(s)
|
||||
@@ -781,13 +786,18 @@ func OpsErrorLoggerMiddleware(ops *service.OpsService) gin.HandlerFunc {
|
||||
}
|
||||
|
||||
recoveredMsg := "Recovered upstream error"
|
||||
if effectiveUpstreamStatus > 0 {
|
||||
if finalAccountAuth {
|
||||
recoveredMsg = "Recovered account authentication failure"
|
||||
} else if effectiveUpstreamStatus > 0 {
|
||||
recoveredMsg += " " + strconvItoa(effectiveUpstreamStatus)
|
||||
}
|
||||
if upstreamErrorMessage != nil && strings.TrimSpace(*upstreamErrorMessage) != "" {
|
||||
recoveredMsg += ": " + strings.TrimSpace(*upstreamErrorMessage)
|
||||
}
|
||||
recoveredMsg = truncateString(recoveredMsg, 2048)
|
||||
recoveredPhase, recoveredBusinessLimited, recoveredOwner, recoveredSource := classifyOpsErrorLog(
|
||||
c, "upstream_error", recoveredMsg, "", effectiveUpstreamStatus,
|
||||
)
|
||||
|
||||
entry := &service.OpsInsertErrorLogInput{
|
||||
RequestID: requestID,
|
||||
@@ -828,19 +838,19 @@ func OpsErrorLoggerMiddleware(ops *service.OpsService) gin.HandlerFunc {
|
||||
}(),
|
||||
UserAgent: c.GetHeader("User-Agent"),
|
||||
|
||||
ErrorPhase: "upstream",
|
||||
ErrorPhase: recoveredPhase,
|
||||
ErrorType: "upstream_error",
|
||||
// Severity should reflect the upstream failure, not the final client status (200).
|
||||
Severity: classifyOpsSeverity("upstream_error", effectiveUpstreamStatus),
|
||||
StatusCode: status,
|
||||
IsBusinessLimited: false,
|
||||
IsBusinessLimited: recoveredBusinessLimited,
|
||||
IsCountTokens: isCountTokensRequest(c),
|
||||
|
||||
ErrorMessage: recoveredMsg,
|
||||
ErrorBody: "",
|
||||
|
||||
ErrorSource: "upstream_http",
|
||||
ErrorOwner: "provider",
|
||||
ErrorSource: recoveredSource,
|
||||
ErrorOwner: recoveredOwner,
|
||||
|
||||
UpstreamStatusCode: upstreamStatusCode,
|
||||
UpstreamErrorMessage: upstreamErrorMessage,
|
||||
@@ -850,6 +860,7 @@ func OpsErrorLoggerMiddleware(ops *service.OpsService) gin.HandlerFunc {
|
||||
CreatedAt: time.Now(),
|
||||
}
|
||||
applyOpsLatencyFieldsFromContext(c, entry)
|
||||
applyOpsUpstreamFieldsFromContext(c, entry)
|
||||
|
||||
if apiKey != nil {
|
||||
entry.APIKeyID = &apiKey.ID
|
||||
@@ -987,60 +998,7 @@ func OpsErrorLoggerMiddleware(ops *service.OpsService) gin.HandlerFunc {
|
||||
CreatedAt: time.Now(),
|
||||
}
|
||||
applyOpsLatencyFieldsFromContext(c, entry)
|
||||
|
||||
// Capture upstream error context set by gateway services (if present).
|
||||
// This does NOT affect the client response; it enriches Ops troubleshooting data.
|
||||
{
|
||||
if v, ok := c.Get(service.OpsUpstreamStatusCodeKey); ok {
|
||||
switch t := v.(type) {
|
||||
case int:
|
||||
if t > 0 {
|
||||
code := t
|
||||
entry.UpstreamStatusCode = &code
|
||||
}
|
||||
case int64:
|
||||
if t > 0 {
|
||||
code := int(t)
|
||||
entry.UpstreamStatusCode = &code
|
||||
}
|
||||
}
|
||||
}
|
||||
if v, ok := c.Get(service.OpsUpstreamErrorMessageKey); ok {
|
||||
if s, ok := v.(string); ok {
|
||||
if msg := strings.TrimSpace(s); msg != "" {
|
||||
entry.UpstreamErrorMessage = &msg
|
||||
}
|
||||
}
|
||||
}
|
||||
if v, ok := c.Get(service.OpsUpstreamErrorDetailKey); ok {
|
||||
if s, ok := v.(string); ok {
|
||||
if detail := strings.TrimSpace(s); detail != "" {
|
||||
entry.UpstreamErrorDetail = &detail
|
||||
}
|
||||
}
|
||||
}
|
||||
if v, ok := c.Get(service.OpsUpstreamErrorsKey); ok {
|
||||
if events, ok := v.([]*service.OpsUpstreamErrorEvent); ok && len(events) > 0 {
|
||||
entry.UpstreamErrors = events
|
||||
// Best-effort backfill the single upstream fields from the last event when missing.
|
||||
last := events[len(events)-1]
|
||||
if last != nil {
|
||||
if entry.UpstreamStatusCode == nil && last.UpstreamStatusCode > 0 {
|
||||
code := last.UpstreamStatusCode
|
||||
entry.UpstreamStatusCode = &code
|
||||
}
|
||||
if entry.UpstreamErrorMessage == nil && strings.TrimSpace(last.Message) != "" {
|
||||
msg := strings.TrimSpace(last.Message)
|
||||
entry.UpstreamErrorMessage = &msg
|
||||
}
|
||||
if entry.UpstreamErrorDetail == nil && strings.TrimSpace(last.Detail) != "" {
|
||||
detail := strings.TrimSpace(last.Detail)
|
||||
entry.UpstreamErrorDetail = &detail
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
applyOpsUpstreamFieldsFromContext(c, entry)
|
||||
|
||||
if apiKey != nil {
|
||||
entry.APIKeyID = &apiKey.ID
|
||||
@@ -1235,6 +1193,77 @@ func applyOpsLatencyFieldsFromContext(c *gin.Context, entry *service.OpsInsertEr
|
||||
entry.TimeToFirstTokenMs = getContextLatencyMs(c, service.OpsTimeToFirstTokenMsKey)
|
||||
}
|
||||
|
||||
// applyOpsUpstreamFieldsFromContext captures attempt-level upstream context.
|
||||
// A final account_auth event owns the top-level status and forces it to zero;
|
||||
// prior inference statuses remain available in UpstreamErrors.
|
||||
func applyOpsUpstreamFieldsFromContext(c *gin.Context, entry *service.OpsInsertErrorLogInput) {
|
||||
if c == nil || entry == nil {
|
||||
return
|
||||
}
|
||||
if v, ok := c.Get(service.OpsUpstreamStatusCodeKey); ok {
|
||||
switch t := v.(type) {
|
||||
case int:
|
||||
if t > 0 {
|
||||
code := t
|
||||
entry.UpstreamStatusCode = &code
|
||||
}
|
||||
case int64:
|
||||
if t > 0 {
|
||||
code := int(t)
|
||||
entry.UpstreamStatusCode = &code
|
||||
}
|
||||
}
|
||||
}
|
||||
if v, ok := c.Get(service.OpsUpstreamErrorMessageKey); ok {
|
||||
if value, ok := v.(string); ok {
|
||||
if message := strings.TrimSpace(value); message != "" {
|
||||
entry.UpstreamErrorMessage = &message
|
||||
}
|
||||
}
|
||||
}
|
||||
if v, ok := c.Get(service.OpsUpstreamErrorDetailKey); ok {
|
||||
if value, ok := v.(string); ok {
|
||||
if detail := strings.TrimSpace(value); detail != "" {
|
||||
entry.UpstreamErrorDetail = &detail
|
||||
}
|
||||
}
|
||||
}
|
||||
if v, ok := c.Get(service.OpsUpstreamErrorsKey); ok {
|
||||
if events, ok := v.([]*service.OpsUpstreamErrorEvent); ok && len(events) > 0 {
|
||||
entry.UpstreamErrors = events
|
||||
last := events[len(events)-1]
|
||||
if last == nil {
|
||||
return
|
||||
}
|
||||
if last.Stage == string(service.GatewayFailureStageAccountAuth) {
|
||||
code := 0
|
||||
entry.UpstreamStatusCode = &code
|
||||
entry.UpstreamErrorMessage = nil
|
||||
if message := strings.TrimSpace(last.Message); message != "" {
|
||||
entry.UpstreamErrorMessage = &message
|
||||
}
|
||||
entry.UpstreamErrorDetail = nil
|
||||
if detail := strings.TrimSpace(last.Detail); detail != "" {
|
||||
entry.UpstreamErrorDetail = &detail
|
||||
}
|
||||
} else {
|
||||
if entry.UpstreamStatusCode == nil && last.UpstreamStatusCode > 0 {
|
||||
code := last.UpstreamStatusCode
|
||||
entry.UpstreamStatusCode = &code
|
||||
}
|
||||
if entry.UpstreamErrorMessage == nil && strings.TrimSpace(last.Message) != "" {
|
||||
message := strings.TrimSpace(last.Message)
|
||||
entry.UpstreamErrorMessage = &message
|
||||
}
|
||||
if entry.UpstreamErrorDetail == nil && strings.TrimSpace(last.Detail) != "" {
|
||||
detail := strings.TrimSpace(last.Detail)
|
||||
entry.UpstreamErrorDetail = &detail
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func getContextLatencyMs(c *gin.Context, key string) *int64 {
|
||||
if c == nil || strings.TrimSpace(key) == "" {
|
||||
return nil
|
||||
@@ -1390,7 +1419,7 @@ func normalizeOpsErrorType(errType string, code string) string {
|
||||
|
||||
func classifyOpsPhase(errType, message, code string) string {
|
||||
msg := strings.ToLower(message)
|
||||
// Standardized phases: request|auth|routing|upstream|network|internal
|
||||
// Standardized phases: request|auth|account_auth|routing|upstream|network|internal
|
||||
// Map billing/concurrency/response => request; scheduling => routing.
|
||||
if isOpsClientAuthError(code, msg) {
|
||||
return "auth"
|
||||
@@ -1445,7 +1474,10 @@ func classifyOpsErrorLog(c *gin.Context, errType, message, code string, status i
|
||||
routingCapacityLimited := isOpsRoutingCapacityLimited(c)
|
||||
clientBusinessLimited := service.HasOpsClientBusinessLimited(c)
|
||||
upstreamError := hasOpsUpstreamErrorContext(c)
|
||||
if upstreamError && !routingCapacityLimited {
|
||||
accountAuthFailure := hasOpsAccountAuthFailure(c)
|
||||
if accountAuthFailure && !routingCapacityLimited {
|
||||
phase = "account_auth"
|
||||
} else if upstreamError && !routingCapacityLimited {
|
||||
phase = "upstream"
|
||||
}
|
||||
if clientBusinessLimited && !upstreamError && !routingCapacityLimited {
|
||||
@@ -1570,6 +1602,22 @@ func hasOpsUpstreamErrorContext(c *gin.Context) bool {
|
||||
return false
|
||||
}
|
||||
|
||||
func hasOpsAccountAuthFailure(c *gin.Context) bool {
|
||||
if c == nil {
|
||||
return false
|
||||
}
|
||||
if v, ok := c.Get(service.OpsUpstreamErrorsKey); ok {
|
||||
if events, ok := v.([]*service.OpsUpstreamErrorEvent); ok {
|
||||
for i := len(events) - 1; i >= 0; i-- {
|
||||
if events[i] != nil {
|
||||
return events[i].Stage == string(service.GatewayFailureStageAccountAuth)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func isOpsNoAvailableAccountMessage(message string) bool {
|
||||
msg := strings.ToLower(message)
|
||||
return strings.Contains(msg, opsErrNoAvailableAccounts) ||
|
||||
@@ -1584,6 +1632,8 @@ func classifyOpsErrorOwner(phase string, message string) string {
|
||||
switch phase {
|
||||
case "upstream", "network":
|
||||
return "provider"
|
||||
case "account_auth":
|
||||
return "provider"
|
||||
case "request", "auth":
|
||||
return "client"
|
||||
case "routing", "internal":
|
||||
@@ -1601,6 +1651,8 @@ func classifyOpsErrorSource(phase string, message string) string {
|
||||
switch phase {
|
||||
case "upstream":
|
||||
return "upstream_http"
|
||||
case "account_auth":
|
||||
return "gateway"
|
||||
case "network":
|
||||
return "gateway"
|
||||
case "request", "auth":
|
||||
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
"crypto/sha256"
|
||||
"encoding/base64"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/url"
|
||||
"os"
|
||||
@@ -164,6 +165,24 @@ func ValidatedBaseURL(override string) (string, error) {
|
||||
return ValidateBaseURL(EffectiveBaseURL(override))
|
||||
}
|
||||
|
||||
// BaseURLValidator applies the caller's outbound URL trust policy before xAI
|
||||
// endpoint paths are appended. The service layer uses this for API-key accounts
|
||||
// so the global security.url_allowlist policy remains the single source of
|
||||
// truth; OAuth callers keep using the strict trusted-host validator.
|
||||
type BaseURLValidator func(string) (string, error)
|
||||
|
||||
func validatedBaseURLWithValidator(override string, validator BaseURLValidator) (string, error) {
|
||||
if validator == nil {
|
||||
return ValidatedBaseURL(override)
|
||||
}
|
||||
raw := EffectiveBaseURL(override)
|
||||
validated, err := validator(raw)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return normalizeKnownBaseURLPath(validated)
|
||||
}
|
||||
|
||||
type RuntimeSanityCheck struct {
|
||||
Value string `json:"value"`
|
||||
Valid bool `json:"valid"`
|
||||
@@ -282,7 +301,16 @@ func ValidateTrustedBaseURL(raw string) (string, error) {
|
||||
func normalizeKnownBaseURLPath(raw string) (string, error) {
|
||||
parsed, err := url.Parse(raw)
|
||||
if err != nil || parsed.Scheme == "" || parsed.Host == "" {
|
||||
return "", fmt.Errorf("invalid url: %s", raw)
|
||||
return "", errors.New("invalid base URL")
|
||||
}
|
||||
if parsed.User != nil {
|
||||
return "", errors.New("base URL must not include userinfo")
|
||||
}
|
||||
if parsed.RawQuery != "" {
|
||||
return "", errors.New("base URL must not include a query")
|
||||
}
|
||||
if parsed.Fragment != "" {
|
||||
return "", errors.New("base URL must not include a fragment")
|
||||
}
|
||||
path := strings.TrimRight(parsed.Path, "/")
|
||||
if path == "" {
|
||||
@@ -435,7 +463,11 @@ func ParseAuthorizationInput(raw string) AuthorizationInput {
|
||||
}
|
||||
|
||||
func BuildResponsesURL(baseURL string) (string, error) {
|
||||
validatedBaseURL, err := ValidatedBaseURL(baseURL)
|
||||
return BuildResponsesURLWithValidator(baseURL, nil)
|
||||
}
|
||||
|
||||
func BuildResponsesURLWithValidator(baseURL string, validator BaseURLValidator) (string, error) {
|
||||
validatedBaseURL, err := validatedBaseURLWithValidator(baseURL, validator)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("invalid base url: %w", err)
|
||||
}
|
||||
@@ -443,7 +475,11 @@ func BuildResponsesURL(baseURL string) (string, error) {
|
||||
}
|
||||
|
||||
func BuildChatCompletionsURL(baseURL string) (string, error) {
|
||||
validatedBaseURL, err := ValidatedBaseURL(baseURL)
|
||||
return BuildChatCompletionsURLWithValidator(baseURL, nil)
|
||||
}
|
||||
|
||||
func BuildChatCompletionsURLWithValidator(baseURL string, validator BaseURLValidator) (string, error) {
|
||||
validatedBaseURL, err := validatedBaseURLWithValidator(baseURL, validator)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("invalid base url: %w", err)
|
||||
}
|
||||
@@ -451,7 +487,11 @@ func BuildChatCompletionsURL(baseURL string) (string, error) {
|
||||
}
|
||||
|
||||
func BuildImagesGenerationsURL(baseURL string) (string, error) {
|
||||
validatedBaseURL, err := ValidatedBaseURL(baseURL)
|
||||
return BuildImagesGenerationsURLWithValidator(baseURL, nil)
|
||||
}
|
||||
|
||||
func BuildImagesGenerationsURLWithValidator(baseURL string, validator BaseURLValidator) (string, error) {
|
||||
validatedBaseURL, err := validatedBaseURLWithValidator(baseURL, validator)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("invalid base url: %w", err)
|
||||
}
|
||||
@@ -459,7 +499,11 @@ func BuildImagesGenerationsURL(baseURL string) (string, error) {
|
||||
}
|
||||
|
||||
func BuildImagesEditsURL(baseURL string) (string, error) {
|
||||
validatedBaseURL, err := ValidatedBaseURL(baseURL)
|
||||
return BuildImagesEditsURLWithValidator(baseURL, nil)
|
||||
}
|
||||
|
||||
func BuildImagesEditsURLWithValidator(baseURL string, validator BaseURLValidator) (string, error) {
|
||||
validatedBaseURL, err := validatedBaseURLWithValidator(baseURL, validator)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("invalid base url: %w", err)
|
||||
}
|
||||
@@ -467,7 +511,11 @@ func BuildImagesEditsURL(baseURL string) (string, error) {
|
||||
}
|
||||
|
||||
func BuildVideosGenerationsURL(baseURL string) (string, error) {
|
||||
validatedBaseURL, err := ValidatedBaseURL(baseURL)
|
||||
return BuildVideosGenerationsURLWithValidator(baseURL, nil)
|
||||
}
|
||||
|
||||
func BuildVideosGenerationsURLWithValidator(baseURL string, validator BaseURLValidator) (string, error) {
|
||||
validatedBaseURL, err := validatedBaseURLWithValidator(baseURL, validator)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("invalid base url: %w", err)
|
||||
}
|
||||
@@ -475,7 +523,11 @@ func BuildVideosGenerationsURL(baseURL string) (string, error) {
|
||||
}
|
||||
|
||||
func BuildVideosEditsURL(baseURL string) (string, error) {
|
||||
validatedBaseURL, err := ValidatedBaseURL(baseURL)
|
||||
return BuildVideosEditsURLWithValidator(baseURL, nil)
|
||||
}
|
||||
|
||||
func BuildVideosEditsURLWithValidator(baseURL string, validator BaseURLValidator) (string, error) {
|
||||
validatedBaseURL, err := validatedBaseURLWithValidator(baseURL, validator)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("invalid base url: %w", err)
|
||||
}
|
||||
@@ -483,7 +535,11 @@ func BuildVideosEditsURL(baseURL string) (string, error) {
|
||||
}
|
||||
|
||||
func BuildVideosExtensionsURL(baseURL string) (string, error) {
|
||||
validatedBaseURL, err := ValidatedBaseURL(baseURL)
|
||||
return BuildVideosExtensionsURLWithValidator(baseURL, nil)
|
||||
}
|
||||
|
||||
func BuildVideosExtensionsURLWithValidator(baseURL string, validator BaseURLValidator) (string, error) {
|
||||
validatedBaseURL, err := validatedBaseURLWithValidator(baseURL, validator)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("invalid base url: %w", err)
|
||||
}
|
||||
@@ -491,7 +547,11 @@ func BuildVideosExtensionsURL(baseURL string) (string, error) {
|
||||
}
|
||||
|
||||
func BuildVideoURL(baseURL, requestID string) (string, error) {
|
||||
validatedBaseURL, err := ValidatedBaseURL(baseURL)
|
||||
return BuildVideoURLWithValidator(baseURL, requestID, nil)
|
||||
}
|
||||
|
||||
func BuildVideoURLWithValidator(baseURL, requestID string, validator BaseURLValidator) (string, error) {
|
||||
validatedBaseURL, err := validatedBaseURLWithValidator(baseURL, validator)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("invalid base url: %w", err)
|
||||
}
|
||||
|
||||
@@ -6,6 +6,7 @@ import (
|
||||
"net/url"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/util/urlvalidator"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
@@ -165,6 +166,44 @@ func TestValidateBaseURLAllowsPublicThirdPartyGrokAPI(t *testing.T) {
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func TestBuildResponsesURLWithValidatorUsesCallerPolicy(t *testing.T) {
|
||||
validator := func(raw string) (string, error) {
|
||||
return urlvalidator.ValidateURLFormat(raw, true)
|
||||
}
|
||||
|
||||
target, err := BuildResponsesURLWithValidator("http://grok.example.test/v1/", validator)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "http://grok.example.test/v1/responses", target)
|
||||
}
|
||||
|
||||
func TestBuildResponsesURLPreservesUnsafeOverrideCustomPath(t *testing.T) {
|
||||
t.Setenv(EnvAllowUnsafeURLOverrides, "true")
|
||||
|
||||
target, err := BuildResponsesURL("http://localhost:8080/custom")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "http://localhost:8080/custom/responses", target)
|
||||
}
|
||||
|
||||
func TestBuildResponsesURLWithValidatorRejectsBaseURLComponents(t *testing.T) {
|
||||
permissive := func(raw string) (string, error) { return raw, nil }
|
||||
tests := []struct {
|
||||
name string
|
||||
raw string
|
||||
}{
|
||||
{name: "userinfo", raw: "https://user:secret@grok.example.test/v1"},
|
||||
{name: "query", raw: "https://grok.example.test/v1?token=secret"},
|
||||
{name: "fragment", raw: "https://grok.example.test/v1#secret"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
_, err := BuildResponsesURLWithValidator(tt.raw, permissive)
|
||||
require.Error(t, err)
|
||||
require.NotContains(t, err.Error(), "secret")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateXAIURLsAllowUnsafeDevOverride(t *testing.T) {
|
||||
t.Setenv(EnvAllowUnsafeURLOverrides, "true")
|
||||
|
||||
|
||||
@@ -859,6 +859,54 @@ func (r *accountRepository) SetError(ctx context.Context, id int64, errorMsg str
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *accountRepository) SetGrokCredentialErrorIfMatch(
|
||||
ctx context.Context,
|
||||
id int64,
|
||||
snapshot service.GrokCredentialMutationSnapshot,
|
||||
errorMsg string,
|
||||
) (bool, error) {
|
||||
result, err := r.sql.ExecContext(ctx, `
|
||||
WITH updated AS (
|
||||
UPDATE accounts AS a
|
||||
SET status = $1,
|
||||
error_message = $2,
|
||||
schedulable = false,
|
||||
updated_at = NOW()
|
||||
WHERE a.id = $3
|
||||
AND a.deleted_at IS NULL
|
||||
AND a.status = $4
|
||||
AND a.platform = $5
|
||||
AND a.type = $6
|
||||
AND a.schedulable IS TRUE
|
||||
AND (a.temp_unschedulable_until IS NULL OR a.temp_unschedulable_until <= NOW())
|
||||
AND (a.rate_limit_reset_at IS NULL OR a.rate_limit_reset_at <= NOW())
|
||||
AND (a.overload_until IS NULL OR a.overload_until <= NOW())
|
||||
AND (a.auto_pause_on_expired IS NOT TRUE OR a.expires_at IS NULL OR a.expires_at > NOW())
|
||||
AND a.credentials = $7::jsonb
|
||||
AND a.proxy_id IS NOT DISTINCT FROM $8
|
||||
AND ($2 <> $9 OR (
|
||||
a.proxy_id IS NOT NULL AND NOT EXISTS (
|
||||
SELECT 1 FROM proxies p WHERE p.id = a.proxy_id AND p.deleted_at IS NULL
|
||||
)
|
||||
))
|
||||
RETURNING a.id
|
||||
)
|
||||
INSERT INTO scheduler_outbox (event_type, account_id, group_id, payload)
|
||||
SELECT $10, updated.id, NULL, NULL FROM updated
|
||||
`, service.StatusError, errorMsg, id, service.StatusActive, service.PlatformGrok, service.AccountTypeOAuth,
|
||||
snapshot.CredentialsJSON, snapshot.ProxyID, string(service.GrokCredentialReasonProxyInvalid),
|
||||
service.SchedulerOutboxEventAccountChanged)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
affected, err := result.RowsAffected()
|
||||
if err != nil || affected == 0 {
|
||||
return false, err
|
||||
}
|
||||
r.syncSchedulerAccountSnapshotDetached(ctx, id)
|
||||
return true, nil
|
||||
}
|
||||
|
||||
// syncSchedulerAccountSnapshot 在账号状态变更时主动同步快照到调度器缓存。
|
||||
// 当账号被设置为错误、禁用、不可调度或临时不可调度时调用,
|
||||
// 确保调度器和粘性会话逻辑能及时感知账号的最新状态,避免继续使用不可用账号。
|
||||
@@ -881,6 +929,16 @@ func (r *accountRepository) syncSchedulerAccountSnapshot(ctx context.Context, ac
|
||||
}
|
||||
}
|
||||
|
||||
func (r *accountRepository) syncSchedulerAccountSnapshotDetached(ctx context.Context, accountID int64) {
|
||||
base := context.Background()
|
||||
if ctx != nil {
|
||||
base = context.WithoutCancel(ctx)
|
||||
}
|
||||
propagationCtx, cancel := context.WithTimeout(base, 2*time.Second)
|
||||
defer cancel()
|
||||
r.syncSchedulerAccountSnapshot(propagationCtx, accountID)
|
||||
}
|
||||
|
||||
func (r *accountRepository) deleteSchedulerAccountSnapshot(ctx context.Context, accountID int64) {
|
||||
if r == nil || r.schedulerCache == nil || accountID <= 0 {
|
||||
return
|
||||
@@ -1050,8 +1108,42 @@ func (r *accountRepository) BindGroups(ctx context.Context, accountID int64, gro
|
||||
}
|
||||
|
||||
func (r *accountRepository) ListSchedulable(ctx context.Context) ([]service.Account, error) {
|
||||
now := time.Now()
|
||||
accounts, err := r.client.Account.Query().
|
||||
accounts, err := r.schedulableAccountsQuery(time.Now()).All(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return r.accountsToService(ctx, accounts)
|
||||
}
|
||||
|
||||
func (r *accountRepository) ListSchedulableAccountLoads(ctx context.Context) ([]service.AccountWithConcurrency, error) {
|
||||
accounts, err := r.schedulableAccountsQuery(time.Now()).
|
||||
Select(
|
||||
dbaccount.FieldID,
|
||||
dbaccount.FieldConcurrency,
|
||||
dbaccount.FieldLoadFactor,
|
||||
).
|
||||
All(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
loads := make([]service.AccountWithConcurrency, 0, len(accounts))
|
||||
for _, account := range accounts {
|
||||
projection := service.Account{
|
||||
ID: account.ID,
|
||||
Concurrency: account.Concurrency,
|
||||
LoadFactor: account.LoadFactor,
|
||||
}
|
||||
loads = append(loads, service.AccountWithConcurrency{
|
||||
ID: account.ID,
|
||||
MaxConcurrency: projection.EffectiveLoadFactor(),
|
||||
})
|
||||
}
|
||||
return loads, nil
|
||||
}
|
||||
|
||||
func (r *accountRepository) schedulableAccountsQuery(now time.Time) *dbent.AccountQuery {
|
||||
return r.client.Account.Query().
|
||||
Where(
|
||||
dbaccount.StatusEQ(service.StatusActive),
|
||||
dbaccount.SchedulableEQ(true),
|
||||
@@ -1060,12 +1152,7 @@ func (r *accountRepository) ListSchedulable(ctx context.Context) ([]service.Acco
|
||||
dbaccount.Or(dbaccount.OverloadUntilIsNil(), dbaccount.OverloadUntilLTE(now)),
|
||||
dbaccount.Or(dbaccount.RateLimitResetAtIsNil(), dbaccount.RateLimitResetAtLTE(now)),
|
||||
).
|
||||
Order(dbent.Asc(dbaccount.FieldPriority)).
|
||||
All(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return r.accountsToService(ctx, accounts)
|
||||
Order(dbent.Asc(dbaccount.FieldPriority))
|
||||
}
|
||||
|
||||
func (r *accountRepository) ListSchedulableByGroupID(ctx context.Context, groupID int64) ([]service.Account, error) {
|
||||
@@ -1319,6 +1406,35 @@ func (r *accountRepository) SetRateLimitedIfLater(ctx context.Context, id int64,
|
||||
return nil
|
||||
}
|
||||
|
||||
// ClearRateLimitIfObserved clears exactly the Grok rate-limit generation seen
|
||||
// by a successful request. Matching both timestamps prevents a stale success
|
||||
// from erasing a later clear/re-arm generation with an equal or shorter reset.
|
||||
func (r *accountRepository) ClearRateLimitIfObserved(ctx context.Context, id int64, observedLimitedAt, observedResetAt time.Time) (bool, error) {
|
||||
updated, err := r.client.Account.Update().
|
||||
Where(
|
||||
dbaccount.IDEQ(id),
|
||||
dbaccount.PlatformEQ(service.PlatformGrok),
|
||||
dbaccount.TypeEQ(service.AccountTypeOAuth),
|
||||
dbaccount.RateLimitedAtEQ(observedLimitedAt),
|
||||
dbaccount.RateLimitResetAtEQ(observedResetAt),
|
||||
).
|
||||
ClearRateLimitedAt().
|
||||
ClearRateLimitResetAt().
|
||||
Save(ctx)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
if updated == 0 {
|
||||
r.syncSchedulerAccountSnapshot(ctx, id)
|
||||
return false, nil
|
||||
}
|
||||
if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
|
||||
logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue observed rate-limit clear failed: account=%d err=%v", id, err)
|
||||
}
|
||||
r.syncSchedulerAccountSnapshot(ctx, id)
|
||||
return true, nil
|
||||
}
|
||||
|
||||
func (r *accountRepository) SetModelRateLimit(ctx context.Context, id int64, scope string, resetAt time.Time, reason ...string) error {
|
||||
if scope == "" {
|
||||
return nil
|
||||
@@ -1414,6 +1530,51 @@ func (r *accountRepository) SetTempUnschedulable(ctx context.Context, id int64,
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *accountRepository) SetGrokCredentialTempUnschedulableIfMatch(
|
||||
ctx context.Context,
|
||||
id int64,
|
||||
snapshot service.GrokCredentialMutationSnapshot,
|
||||
until time.Time,
|
||||
reason string,
|
||||
) (bool, error) {
|
||||
result, err := r.sql.ExecContext(ctx, `
|
||||
WITH updated AS (
|
||||
UPDATE accounts AS a
|
||||
SET temp_unschedulable_until = CASE
|
||||
WHEN a.temp_unschedulable_until IS NULL OR a.temp_unschedulable_until < $1 THEN $1
|
||||
ELSE a.temp_unschedulable_until
|
||||
END,
|
||||
temp_unschedulable_reason = $2,
|
||||
updated_at = NOW()
|
||||
WHERE a.id = $3
|
||||
AND a.deleted_at IS NULL
|
||||
AND a.status = $4
|
||||
AND a.platform = $5
|
||||
AND a.type = $6
|
||||
AND a.schedulable IS TRUE
|
||||
AND (a.temp_unschedulable_until IS NULL OR a.temp_unschedulable_until <= NOW())
|
||||
AND (a.rate_limit_reset_at IS NULL OR a.rate_limit_reset_at <= NOW())
|
||||
AND (a.overload_until IS NULL OR a.overload_until <= NOW())
|
||||
AND (a.auto_pause_on_expired IS NOT TRUE OR a.expires_at IS NULL OR a.expires_at > NOW())
|
||||
AND a.credentials = $7::jsonb
|
||||
AND a.proxy_id IS NOT DISTINCT FROM $8
|
||||
RETURNING a.id
|
||||
)
|
||||
INSERT INTO scheduler_outbox (event_type, account_id, group_id, payload)
|
||||
SELECT $9, updated.id, NULL, NULL FROM updated
|
||||
`, until, reason, id, service.StatusActive, service.PlatformGrok, service.AccountTypeOAuth,
|
||||
snapshot.CredentialsJSON, snapshot.ProxyID, service.SchedulerOutboxEventAccountChanged)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
affected, err := result.RowsAffected()
|
||||
if err != nil || affected == 0 {
|
||||
return false, err
|
||||
}
|
||||
r.syncSchedulerAccountSnapshotDetached(ctx, id)
|
||||
return true, nil
|
||||
}
|
||||
|
||||
func (r *accountRepository) ClearTempUnschedulable(ctx context.Context, id int64) error {
|
||||
_, err := r.sql.ExecContext(ctx, `
|
||||
UPDATE accounts
|
||||
|
||||
@@ -722,6 +722,56 @@ func (s *AccountRepoSuite) TestSetRateLimitedIfLaterDoesNotShortenReset() {
|
||||
s.Require().WithinDuration(later, *cacheRecorder.setAccounts[1].RateLimitResetAt, time.Second)
|
||||
}
|
||||
|
||||
func (s *AccountRepoSuite) TestClearRateLimitIfObservedProtectsRearmed429Generation() {
|
||||
account := mustCreateAccount(s.T(), s.client, &service.Account{
|
||||
Name: "acc-rl-conditional-clear",
|
||||
Platform: service.PlatformGrok,
|
||||
Type: service.AccountTypeOAuth,
|
||||
})
|
||||
firstReset := time.Now().Add(30 * time.Minute).UTC().Truncate(time.Second)
|
||||
rearmedReset := time.Now().Add(5 * time.Minute).UTC().Truncate(time.Second)
|
||||
|
||||
s.Require().NoError(s.repo.SetRateLimitedIfLater(s.ctx, account.ID, firstReset))
|
||||
staleGeneration, err := s.repo.GetByID(s.ctx, account.ID)
|
||||
s.Require().NoError(err)
|
||||
s.Require().NotNil(staleGeneration.RateLimitedAt)
|
||||
s.Require().NotNil(staleGeneration.RateLimitResetAt)
|
||||
cleared, err := s.repo.ClearRateLimitIfObserved(s.ctx, account.ID, *staleGeneration.RateLimitedAt, *staleGeneration.RateLimitResetAt)
|
||||
s.Require().NoError(err)
|
||||
s.Require().True(cleared)
|
||||
|
||||
// A newer generation may legitimately re-arm a shorter boundary after the
|
||||
// first generation was cleared. The stale success must not erase it.
|
||||
s.Require().NoError(s.repo.SetRateLimitedIfLater(s.ctx, account.ID, rearmedReset))
|
||||
cleared, err = s.repo.ClearRateLimitIfObserved(s.ctx, account.ID, *staleGeneration.RateLimitedAt, *staleGeneration.RateLimitResetAt)
|
||||
s.Require().NoError(err)
|
||||
s.Require().False(cleared)
|
||||
|
||||
got, err := s.repo.GetByID(s.ctx, account.ID)
|
||||
s.Require().NoError(err)
|
||||
s.Require().NotNil(got.RateLimitedAt)
|
||||
s.Require().NotNil(got.RateLimitResetAt)
|
||||
s.Require().WithinDuration(rearmedReset, *got.RateLimitResetAt, time.Second)
|
||||
|
||||
// An admin can retype the row while the successful OAuth request is still
|
||||
// in flight. The stale OAuth recovery must not cross into API-key state even
|
||||
// when both observed timestamps still match.
|
||||
_, err = s.client.Account.UpdateOneID(account.ID).
|
||||
SetType(service.AccountTypeAPIKey).
|
||||
Save(s.ctx)
|
||||
s.Require().NoError(err)
|
||||
cleared, err = s.repo.ClearRateLimitIfObserved(s.ctx, account.ID, *got.RateLimitedAt, *got.RateLimitResetAt)
|
||||
s.Require().NoError(err)
|
||||
s.Require().False(cleared)
|
||||
|
||||
retyped, err := s.repo.GetByID(s.ctx, account.ID)
|
||||
s.Require().NoError(err)
|
||||
s.Require().Equal(service.AccountTypeAPIKey, retyped.Type)
|
||||
s.Require().NotNil(retyped.RateLimitedAt)
|
||||
s.Require().NotNil(retyped.RateLimitResetAt)
|
||||
s.Require().WithinDuration(rearmedReset, *retyped.RateLimitResetAt, time.Second)
|
||||
}
|
||||
|
||||
func (s *AccountRepoSuite) TestClearRateLimit() {
|
||||
account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-clear"})
|
||||
until := time.Now().Add(1 * time.Hour)
|
||||
|
||||
@@ -0,0 +1,107 @@
|
||||
//go:build integration
|
||||
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestListSchedulableAccountLoadsMatchesListSchedulable(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
tx := testEntTx(t)
|
||||
client := tx.Client()
|
||||
repo := newAccountRepositoryWithSQL(client, tx, nil)
|
||||
now := time.Now()
|
||||
past := now.Add(-time.Hour)
|
||||
future := now.Add(time.Hour)
|
||||
|
||||
create := func(name string) *service.Account {
|
||||
return mustCreateAccount(t, client, &service.Account{Name: name, Schedulable: true})
|
||||
}
|
||||
|
||||
positiveLoad := create("projection-positive-load")
|
||||
_, err := client.Account.UpdateOneID(positiveLoad.ID).SetConcurrency(2).SetLoadFactor(9).SetPriority(30).Save(ctx)
|
||||
require.NoError(t, err)
|
||||
concurrencyFallback := create("projection-concurrency-fallback")
|
||||
_, err = client.Account.UpdateOneID(concurrencyFallback.ID).SetConcurrency(4).SetPriority(10).Save(ctx)
|
||||
require.NoError(t, err)
|
||||
zeroFallback := create("projection-zero-fallback")
|
||||
_, err = client.Account.UpdateOneID(zeroFallback.ID).SetConcurrency(0).SetLoadFactor(0).SetPriority(20).Save(ctx)
|
||||
require.NoError(t, err)
|
||||
|
||||
disabled := create("projection-disabled")
|
||||
_, err = client.Account.UpdateOneID(disabled.ID).SetStatus(service.StatusDisabled).Save(ctx)
|
||||
require.NoError(t, err)
|
||||
unschedulable := create("projection-unschedulable")
|
||||
_, err = client.Account.UpdateOneID(unschedulable.ID).SetSchedulable(false).Save(ctx)
|
||||
require.NoError(t, err)
|
||||
expired := create("projection-expired")
|
||||
_, err = client.Account.UpdateOneID(expired.ID).SetExpiresAt(past).SetAutoPauseOnExpired(true).Save(ctx)
|
||||
require.NoError(t, err)
|
||||
expiredAllowed := create("projection-expired-allowed")
|
||||
_, err = client.Account.UpdateOneID(expiredAllowed.ID).SetExpiresAt(past).SetAutoPauseOnExpired(false).Save(ctx)
|
||||
require.NoError(t, err)
|
||||
overloaded := create("projection-overloaded")
|
||||
_, err = client.Account.UpdateOneID(overloaded.ID).SetOverloadUntil(future).Save(ctx)
|
||||
require.NoError(t, err)
|
||||
overloadCleared := create("projection-overload-cleared")
|
||||
_, err = client.Account.UpdateOneID(overloadCleared.ID).SetOverloadUntil(past).Save(ctx)
|
||||
require.NoError(t, err)
|
||||
rateLimited := create("projection-rate-limited")
|
||||
_, err = client.Account.UpdateOneID(rateLimited.ID).SetRateLimitResetAt(future).Save(ctx)
|
||||
require.NoError(t, err)
|
||||
rateLimitCleared := create("projection-rate-limit-cleared")
|
||||
_, err = client.Account.UpdateOneID(rateLimitCleared.ID).SetRateLimitResetAt(past).Save(ctx)
|
||||
require.NoError(t, err)
|
||||
tempBlocked := create("projection-temp-blocked")
|
||||
_, err = client.Account.UpdateOneID(tempBlocked.ID).SetTempUnschedulableUntil(future).Save(ctx)
|
||||
require.NoError(t, err)
|
||||
tempCleared := create("projection-temp-cleared")
|
||||
_, err = client.Account.UpdateOneID(tempCleared.ID).SetTempUnschedulableUntil(past).Save(ctx)
|
||||
require.NoError(t, err)
|
||||
|
||||
accounts, err := repo.ListSchedulable(ctx)
|
||||
require.NoError(t, err)
|
||||
loads, err := repo.ListSchedulableAccountLoads(ctx)
|
||||
require.NoError(t, err)
|
||||
|
||||
accountIDs := make([]int64, 0, len(accounts))
|
||||
wantByID := make(map[int64]int, len(accounts))
|
||||
for i := range accounts {
|
||||
accountIDs = append(accountIDs, accounts[i].ID)
|
||||
wantByID[accounts[i].ID] = accounts[i].EffectiveLoadFactor()
|
||||
}
|
||||
|
||||
loadIDs := make([]int64, 0, len(loads))
|
||||
byID := make(map[int64]int, len(loads))
|
||||
for _, load := range loads {
|
||||
loadIDs = append(loadIDs, load.ID)
|
||||
byID[load.ID] = load.MaxConcurrency
|
||||
}
|
||||
require.Equal(t, accountIDs, loadIDs)
|
||||
targetIDs := map[int64]struct{}{
|
||||
positiveLoad.ID: {}, concurrencyFallback.ID: {}, zeroFallback.ID: {},
|
||||
}
|
||||
targetOrder := make([]int64, 0, len(targetIDs))
|
||||
for _, id := range loadIDs {
|
||||
if _, ok := targetIDs[id]; ok {
|
||||
targetOrder = append(targetOrder, id)
|
||||
}
|
||||
}
|
||||
require.Equal(t, []int64{concurrencyFallback.ID, zeroFallback.ID, positiveLoad.ID}, targetOrder)
|
||||
require.Equal(t, wantByID, byID)
|
||||
require.Equal(t, 9, byID[positiveLoad.ID])
|
||||
require.Equal(t, 4, byID[concurrencyFallback.ID])
|
||||
require.Equal(t, 1, byID[zeroFallback.ID])
|
||||
for _, included := range []*service.Account{expiredAllowed, overloadCleared, rateLimitCleared, tempCleared} {
|
||||
require.Contains(t, byID, included.ID)
|
||||
}
|
||||
for _, excluded := range []*service.Account{disabled, unschedulable, expired, overloaded, rateLimited, tempBlocked} {
|
||||
require.NotContains(t, byID, excluded.ID)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,82 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/DATA-DOG/go-sqlmock"
|
||||
dbent "github.com/Wei-Shaw/sub2api/ent"
|
||||
_ "github.com/Wei-Shaw/sub2api/ent/runtime"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"entgo.io/ent/dialect"
|
||||
entsql "entgo.io/ent/dialect/sql"
|
||||
)
|
||||
|
||||
type captureEntQueryMatcher struct {
|
||||
actual *string
|
||||
}
|
||||
|
||||
func (m captureEntQueryMatcher) Match(_, actual string) error {
|
||||
if m.actual == nil {
|
||||
return fmt.Errorf("query capture target is nil")
|
||||
}
|
||||
*m.actual = actual
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestListSchedulableAccountLoadsUsesSingleProjectionQuery(t *testing.T) {
|
||||
var capturedSQL string
|
||||
db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(captureEntQueryMatcher{actual: &capturedSQL}))
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
|
||||
driver := entsql.OpenDB(dialect.Postgres, db)
|
||||
client := dbent.NewClient(dbent.Driver(driver))
|
||||
t.Cleanup(func() { _ = client.Close() })
|
||||
repo := newAccountRepositoryWithSQL(client, db, nil)
|
||||
|
||||
mock.ExpectQuery("schedulable account load projection").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id", "concurrency", "load_factor"}).
|
||||
AddRow(int64(11), 3, nil).
|
||||
AddRow(int64(12), 2, 7))
|
||||
|
||||
loads, err := repo.ListSchedulableAccountLoads(context.Background())
|
||||
require.NoError(t, err)
|
||||
require.Len(t, loads, 2)
|
||||
require.Equal(t, int64(11), loads[0].ID)
|
||||
require.Equal(t, 3, loads[0].MaxConcurrency)
|
||||
require.Equal(t, int64(12), loads[1].ID)
|
||||
require.Equal(t, 7, loads[1].MaxConcurrency)
|
||||
require.NoError(t, mock.ExpectationsWereMet(), "projection path must execute exactly one query")
|
||||
|
||||
normalized := normalizeSQLWhitespace(capturedSQL)
|
||||
selectClause, _, found := strings.Cut(normalized, " FROM ")
|
||||
require.True(t, found, "unexpected projection SQL: %s", normalized)
|
||||
require.Equal(t, 2, strings.Count(selectClause, ","), "projection must select exactly three columns: %s", selectClause)
|
||||
require.Contains(t, selectClause, `"id"`)
|
||||
require.Contains(t, selectClause, `"concurrency"`)
|
||||
require.Contains(t, selectClause, `"load_factor"`)
|
||||
require.NotContains(t, selectClause, "credentials")
|
||||
require.NotContains(t, selectClause, "extra")
|
||||
require.NotContains(t, selectClause, "proxy_id")
|
||||
require.NotContains(t, normalized, "account_groups")
|
||||
require.NotContains(t, normalized, "proxies")
|
||||
for _, predicateColumn := range []string{
|
||||
"status",
|
||||
"schedulable",
|
||||
"temp_unschedulable_until",
|
||||
"expires_at",
|
||||
"auto_pause_on_expired",
|
||||
"overload_until",
|
||||
"rate_limit_reset_at",
|
||||
"deleted_at",
|
||||
} {
|
||||
require.Contains(t, normalized, predicateColumn)
|
||||
}
|
||||
_, orderClause, hasOrder := strings.Cut(normalized, " ORDER BY ")
|
||||
require.True(t, hasOrder, "projection query must preserve schedulable account order: %s", normalized)
|
||||
require.Contains(t, orderClause, `"priority" ASC`)
|
||||
}
|
||||
@@ -9,6 +9,7 @@ import (
|
||||
"time"
|
||||
|
||||
sqlmock "github.com/DATA-DOG/go-sqlmock"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
@@ -24,6 +25,103 @@ func TestAccountRepository_SetTempUnschedulable_NoRowsAffectedDoesNotWriteOutbox
|
||||
require.NotContains(t, strings.Join(exec.execQueries, "\n"), "scheduler_outbox")
|
||||
}
|
||||
|
||||
func TestAccountRepository_GrokCredentialConditionalMutationsAreEligibleAndAtomicallyPropagated(t *testing.T) {
|
||||
proxyID := int64(77)
|
||||
snapshot := service.GrokCredentialMutationSnapshot{
|
||||
CredentialsJSON: `{"access_token":"access","refresh_token":"refresh","_token_version":123}`,
|
||||
ProxyID: &proxyID,
|
||||
}
|
||||
|
||||
t.Run("permanent", func(t *testing.T) {
|
||||
exec := &recordingSQLExecutor{result: rowsAffectedResult(0)}
|
||||
repo := newAccountRepositoryWithSQL(nil, exec, nil)
|
||||
|
||||
updated, err := repo.SetGrokCredentialErrorIfMatch(context.Background(), 42, snapshot, "revoked")
|
||||
|
||||
require.NoError(t, err)
|
||||
require.False(t, updated)
|
||||
require.Len(t, exec.execQueries, 1)
|
||||
normalized := normalizeSQLWhitespace(exec.execQueries[0])
|
||||
require.Contains(t, normalized, "WITH updated AS ( UPDATE accounts AS a")
|
||||
require.Contains(t, normalized, "a.schedulable IS TRUE")
|
||||
require.Contains(t, normalized, "a.temp_unschedulable_until IS NULL OR a.temp_unschedulable_until <= NOW()")
|
||||
require.Contains(t, normalized, "a.rate_limit_reset_at IS NULL OR a.rate_limit_reset_at <= NOW()")
|
||||
require.Contains(t, normalized, "a.overload_until IS NULL OR a.overload_until <= NOW()")
|
||||
require.Contains(t, normalized, "a.credentials = $7::jsonb")
|
||||
require.Contains(t, normalized, "a.proxy_id IS NOT DISTINCT FROM $8")
|
||||
require.Contains(t, normalized, "NOT EXISTS ( SELECT 1 FROM proxies p")
|
||||
require.Contains(t, normalized, "INSERT INTO scheduler_outbox")
|
||||
require.Len(t, exec.execArgs[0], 10)
|
||||
require.Equal(t, snapshot.CredentialsJSON, exec.execArgs[0][6])
|
||||
require.Equal(t, &proxyID, exec.execArgs[0][7])
|
||||
require.Equal(t, string(service.GrokCredentialReasonProxyInvalid), exec.execArgs[0][8])
|
||||
require.Equal(t, service.SchedulerOutboxEventAccountChanged, exec.execArgs[0][9])
|
||||
})
|
||||
|
||||
t.Run("transient", func(t *testing.T) {
|
||||
exec := &recordingSQLExecutor{result: rowsAffectedResult(0)}
|
||||
repo := newAccountRepositoryWithSQL(nil, exec, nil)
|
||||
|
||||
updated, err := repo.SetGrokCredentialTempUnschedulableIfMatch(
|
||||
context.Background(), 42, snapshot, time.Now().Add(time.Minute), "temporary",
|
||||
)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.False(t, updated)
|
||||
require.Len(t, exec.execQueries, 1)
|
||||
normalized := normalizeSQLWhitespace(exec.execQueries[0])
|
||||
require.Contains(t, normalized, "WITH updated AS ( UPDATE accounts AS a")
|
||||
require.Contains(t, normalized, "a.schedulable IS TRUE")
|
||||
require.Contains(t, normalized, "a.temp_unschedulable_until IS NULL OR a.temp_unschedulable_until <= NOW()")
|
||||
require.Contains(t, normalized, "a.rate_limit_reset_at IS NULL OR a.rate_limit_reset_at <= NOW()")
|
||||
require.Contains(t, normalized, "a.overload_until IS NULL OR a.overload_until <= NOW()")
|
||||
require.Contains(t, normalized, "a.credentials = $7::jsonb")
|
||||
require.Contains(t, normalized, "a.proxy_id IS NOT DISTINCT FROM $8")
|
||||
require.Contains(t, normalized, "INSERT INTO scheduler_outbox")
|
||||
require.Len(t, exec.execArgs[0], 9)
|
||||
require.Equal(t, snapshot.CredentialsJSON, exec.execArgs[0][6])
|
||||
require.Equal(t, &proxyID, exec.execArgs[0][7])
|
||||
require.Equal(t, service.SchedulerOutboxEventAccountChanged, exec.execArgs[0][8])
|
||||
})
|
||||
}
|
||||
|
||||
func TestAccountRepository_GrokCredentialCommitCarriesOutboxAcrossCallerCancellation(t *testing.T) {
|
||||
snapshot := service.GrokCredentialMutationSnapshot{CredentialsJSON: `{"access_token":"access","refresh_token":"refresh"}`}
|
||||
tests := []struct {
|
||||
name string
|
||||
mutate func(context.Context, *accountRepository) (bool, error)
|
||||
}{
|
||||
{
|
||||
name: "permanent",
|
||||
mutate: func(ctx context.Context, repo *accountRepository) (bool, error) {
|
||||
return repo.SetGrokCredentialErrorIfMatch(ctx, 42, snapshot, string(service.GrokCredentialReasonRevoked))
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "transient",
|
||||
mutate: func(ctx context.Context, repo *accountRepository) (bool, error) {
|
||||
return repo.SetGrokCredentialTempUnschedulableIfMatch(ctx, 42, snapshot, time.Now().Add(time.Minute), string(service.GrokCredentialReasonRefreshTransient))
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
exec := &recordingSQLExecutor{result: rowsAffectedResult(1), afterExec: cancel}
|
||||
repo := newAccountRepositoryWithSQL(nil, exec, nil)
|
||||
|
||||
updated, err := tt.mutate(ctx, repo)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.True(t, updated)
|
||||
require.ErrorIs(t, ctx.Err(), context.Canceled)
|
||||
require.Len(t, exec.execQueries, 1, "state update and scheduler outbox must share one atomic SQL statement")
|
||||
require.Contains(t, normalizeSQLWhitespace(exec.execQueries[0]), "INSERT INTO scheduler_outbox")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAccountRepository_ListOAuthRefreshCandidates_SQLFilter(t *testing.T) {
|
||||
db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherRegexp))
|
||||
require.NoError(t, err)
|
||||
@@ -87,14 +185,20 @@ func (r rowsAffectedResult) RowsAffected() (int64, error) { return int64(r), nil
|
||||
type recordingSQLExecutor struct {
|
||||
result sql.Result
|
||||
err error
|
||||
afterExec func()
|
||||
execQueries []string
|
||||
execArgs [][]any
|
||||
}
|
||||
|
||||
func (e *recordingSQLExecutor) ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) {
|
||||
e.execQueries = append(e.execQueries, query)
|
||||
e.execArgs = append(e.execArgs, append([]any(nil), args...))
|
||||
if e.err != nil {
|
||||
return nil, e.err
|
||||
}
|
||||
if e.afterExec != nil {
|
||||
e.afterExec()
|
||||
}
|
||||
return e.result, nil
|
||||
}
|
||||
|
||||
|
||||
@@ -153,14 +153,29 @@ func grokOAuthStatusError(code, message string, resp *req.Response) error {
|
||||
statusCode := http.StatusBadGateway
|
||||
errorCode := code
|
||||
upstreamStatus := 0
|
||||
if resp != nil && resp.StatusCode == http.StatusForbidden {
|
||||
statusCode = http.StatusForbidden
|
||||
errorCode = "GROK_OAUTH_ENTITLEMENT_DENIED"
|
||||
}
|
||||
body := ""
|
||||
if resp != nil {
|
||||
upstreamStatus = resp.StatusCode
|
||||
body = logredact.RedactText(resp.String())
|
||||
if resp.StatusCode == http.StatusForbidden && grokOAuthHasExplicitEntitlementDenial(body) {
|
||||
statusCode = http.StatusForbidden
|
||||
errorCode = "GROK_OAUTH_ENTITLEMENT_DENIED"
|
||||
}
|
||||
}
|
||||
return infraerrors.Newf(statusCode, errorCode, "%s: status %d, body: %s", message, upstreamStatus, body)
|
||||
}
|
||||
|
||||
func grokOAuthHasExplicitEntitlementDenial(body string) bool {
|
||||
lower := strings.ToLower(body)
|
||||
compact := strings.NewReplacer(" ", "", "\n", "", "\r", "", "\t", "").Replace(lower)
|
||||
for _, field := range []string{"error", "code", "reason"} {
|
||||
for _, value := range []string{"access_denied", "entitlement_denied", "subscription_required", "no_active_subscription"} {
|
||||
if strings.Contains(compact, `"`+field+`":"`+value+`"`) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
return strings.Contains(lower, "entitlement denied") ||
|
||||
strings.Contains(lower, "subscription required") ||
|
||||
strings.Contains(lower, "no active grok subscription")
|
||||
}
|
||||
|
||||
@@ -50,15 +50,7 @@ func TestGrokOAuthClientExchangeAndRefreshUseFormFields(t *testing.T) {
|
||||
t.Setenv(xai.EnvTokenURL, server.URL)
|
||||
|
||||
client := NewGrokOAuthClient()
|
||||
|
||||
exchanged, err := client.ExchangeCode(
|
||||
context.Background(),
|
||||
"auth-code",
|
||||
"verifier",
|
||||
"http://127.0.0.1:56121/callback",
|
||||
"",
|
||||
"client-id",
|
||||
)
|
||||
exchanged, err := client.ExchangeCode(context.Background(), "auth-code", "verifier", "http://127.0.0.1:56121/callback", "", "client-id")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "exchange-access", exchanged.AccessToken)
|
||||
require.Equal(t, "exchange-refresh", exchanged.RefreshToken)
|
||||
@@ -72,18 +64,30 @@ func TestGrokOAuthClientExchangeAndRefreshUseFormFields(t *testing.T) {
|
||||
require.Equal(t, int64(7200), refreshed.ExpiresIn)
|
||||
}
|
||||
|
||||
func TestGrokOAuthClientRefreshForbiddenClassifiesEntitlement(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusForbidden)
|
||||
_, _ = w.Write([]byte(`{"error":"subscription required"}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
t.Setenv(xai.EnvTokenURL, server.URL)
|
||||
func TestGrokOAuthClientRefreshForbiddenClassifiesOnlyExplicitEntitlement(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
body string
|
||||
wantReason string
|
||||
}{
|
||||
{name: "explicit entitlement", body: `{"error":"access_denied"}`, wantReason: "GROK_OAUTH_ENTITLEMENT_DENIED"},
|
||||
{name: "generic forbidden", body: `{"error":"forbidden"}`, wantReason: "GROK_OAUTH_TOKEN_REFRESH_FAILED"},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusForbidden)
|
||||
_, _ = w.Write([]byte(tt.body))
|
||||
}))
|
||||
defer server.Close()
|
||||
t.Setenv(xai.EnvTokenURL, server.URL)
|
||||
|
||||
client := NewGrokOAuthClient()
|
||||
_, err := client.RefreshToken(context.Background(), "refresh-token", "", "client-id")
|
||||
require.Error(t, err)
|
||||
require.Contains(t, strings.ToUpper(err.Error()), "GROK_OAUTH_ENTITLEMENT_DENIED")
|
||||
client := NewGrokOAuthClient()
|
||||
_, err := client.RefreshToken(context.Background(), "refresh-token", "", "client-id")
|
||||
require.Error(t, err)
|
||||
require.Contains(t, strings.ToUpper(err.Error()), tt.wantReason)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGrokOAuthClientStatusErrorRedactsSensitiveResponseBody(t *testing.T) {
|
||||
@@ -105,3 +109,13 @@ func TestGrokOAuthClientStatusErrorRedactsSensitiveResponseBody(t *testing.T) {
|
||||
require.NotContains(t, errText, "refresh-secret")
|
||||
require.NotContains(t, errText, "verifier-secret")
|
||||
}
|
||||
|
||||
func TestGrokOAuthEntitlementDenialRequiresExplicitEvidence(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
require.True(t, grokOAuthHasExplicitEntitlementDenial(`{"error":"access_denied"}`))
|
||||
require.True(t, grokOAuthHasExplicitEntitlementDenial(`{"code":"entitlement_denied"}`))
|
||||
require.True(t, grokOAuthHasExplicitEntitlementDenial(`{"message":"no active Grok subscription"}`))
|
||||
require.False(t, grokOAuthHasExplicitEntitlementDenial(`{"error":"forbidden","message":"request forbidden"}`))
|
||||
require.False(t, grokOAuthHasExplicitEntitlementDenial(`<html>403 Forbidden</html>`))
|
||||
}
|
||||
|
||||
@@ -101,6 +101,40 @@ func TestBuildOpsErrorLogsWhere_CyberPolicyStatusExemption(t *testing.T) {
|
||||
if strings.Contains(whereRecovered, "status_code") {
|
||||
t.Fatalf("upstream phase with IncludeRecoveredUpstream must not add any status_code clause\nfull: %s", whereRecovered)
|
||||
}
|
||||
|
||||
// account_auth uses the same explicit provider-health opt-in but remains a
|
||||
// distinct phase from inference upstream errors.
|
||||
whereAccountAuth, _ := buildOpsErrorLogsWhere(&service.OpsErrorLogFilter{Phase: "account_auth", IncludeRecoveredUpstream: true})
|
||||
if strings.Contains(whereAccountAuth, "status_code") {
|
||||
t.Fatalf("account_auth phase with IncludeRecoveredUpstream must expose recovered rows\nfull: %s", whereAccountAuth)
|
||||
}
|
||||
if !strings.Contains(whereAccountAuth, "e.error_phase = $") {
|
||||
t.Fatalf("account_auth recovered filter must retain its explicit phase\nfull: %s", whereAccountAuth)
|
||||
}
|
||||
|
||||
whereProviderHealth, _ := buildOpsErrorLogsWhere(&service.OpsErrorLogFilter{
|
||||
ErrorPhasesAny: []string{"upstream", "account_auth"},
|
||||
IncludeRecoveredUpstream: true,
|
||||
})
|
||||
if strings.Contains(whereProviderHealth, "status_code") {
|
||||
t.Fatalf("provider-health ANY filter must expose recovered inference and credential rows\nfull: %s", whereProviderHealth)
|
||||
}
|
||||
if !strings.Contains(whereProviderHealth, "e.error_phase = ANY($") {
|
||||
t.Fatalf("provider-health filter must preserve distinct phase values\nfull: %s", whereProviderHealth)
|
||||
}
|
||||
|
||||
whereUserAccountAuth, _ := buildOpsErrorLogsWhere(&service.OpsErrorLogFilter{ErrorPhasesAny: []string{"account_auth"}})
|
||||
if !strings.Contains(whereUserAccountAuth, "COALESCE(e.status_code, 0) >= 400") {
|
||||
t.Fatalf("request-error account_auth filters must exclude recovered successes\nfull: %s", whereUserAccountAuth)
|
||||
}
|
||||
|
||||
whereMixed, _ := buildOpsErrorLogsWhere(&service.OpsErrorLogFilter{
|
||||
ErrorPhasesAny: []string{"account_auth", "request"},
|
||||
IncludeRecoveredUpstream: true,
|
||||
})
|
||||
if !strings.Contains(whereMixed, "COALESCE(e.status_code, 0) >= 400") {
|
||||
t.Fatalf("recovered opt-in must not bypass the guard for non-provider phases\nfull: %s", whereMixed)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildOpsErrorLogsWhere_MatchDeletedKeyOwner(t *testing.T) {
|
||||
|
||||
@@ -160,7 +160,7 @@ func opsInsertErrorLogArgs(input *service.OpsInsertErrorLogInput) []any {
|
||||
opsNullString(input.ErrorBody),
|
||||
opsNullString(input.ErrorSource),
|
||||
opsNullString(input.ErrorOwner),
|
||||
opsNullInt(input.UpstreamStatusCode),
|
||||
opsNullableIntPointer(input.UpstreamStatusCode),
|
||||
opsNullString(input.UpstreamErrorMessage),
|
||||
opsNullString(input.UpstreamErrorDetail),
|
||||
opsNullString(input.UpstreamErrorsJSON),
|
||||
@@ -582,7 +582,7 @@ LIMIT 1`
|
||||
s := clientIP.String
|
||||
out.ClientIP = &s
|
||||
}
|
||||
if upstreamStatusCode.Valid && upstreamStatusCode.Int64 > 0 {
|
||||
if upstreamStatusCode.Valid {
|
||||
v := int(upstreamStatusCode.Int64)
|
||||
out.UpstreamStatusCode = &v
|
||||
}
|
||||
@@ -978,13 +978,13 @@ func buildOpsErrorLogsWhere(filter *service.OpsErrorLogFilter) (string, []any) {
|
||||
resolvedFilter = filter.Resolved
|
||||
}
|
||||
// Keep list endpoints scoped to client errors unless the caller explicitly opts
|
||||
// into recovered upstream rows (Phase=="upstream" + IncludeRecoveredUpstream,
|
||||
// ops 专用上游列表)。请求错误语义的端点即便过滤 phase=upstream 也保留该守卫。
|
||||
// into recovered provider-health rows (upstream/account_auth). Request-error
|
||||
// endpoints never set the opt-in and retain this guard.
|
||||
// cyber_policy is exempt from the status >= 400 guard: streaming cyber hits arrive with
|
||||
// status 200 (the SSE stream opened successfully before upstream returned response.failed),
|
||||
// but they are always client-visible blocked requests that belong in admin + user error
|
||||
// lists. Without the exemption the entire streaming-path cyber sink would be invisible.
|
||||
if phaseFilter != "upstream" || filter == nil || !filter.IncludeRecoveredUpstream {
|
||||
if !opsFilterIncludesRecoveredProviderRows(filter, phaseFilter) {
|
||||
clauses = append(clauses, "(COALESCE(e.status_code, 0) >= 400 OR e.error_type = 'cyber_policy')")
|
||||
}
|
||||
|
||||
@@ -1117,6 +1117,28 @@ func buildOpsErrorLogsWhere(filter *service.OpsErrorLogFilter) (string, []any) {
|
||||
return "WHERE " + strings.Join(clauses, " AND "), args
|
||||
}
|
||||
|
||||
func opsFilterIncludesRecoveredProviderRows(filter *service.OpsErrorLogFilter, phaseFilter string) bool {
|
||||
if filter == nil || !filter.IncludeRecoveredUpstream {
|
||||
return false
|
||||
}
|
||||
if phaseFilter != "" {
|
||||
return phaseFilter == "upstream" || phaseFilter == "account_auth"
|
||||
}
|
||||
if len(filter.ErrorPhasesAny) == 0 {
|
||||
return false
|
||||
}
|
||||
sawProviderPhase := false
|
||||
for _, rawPhase := range filter.ErrorPhasesAny {
|
||||
switch strings.TrimSpace(strings.ToLower(rawPhase)) {
|
||||
case "upstream", "account_auth":
|
||||
sawProviderPhase = true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
return sawProviderPhase
|
||||
}
|
||||
|
||||
func buildOpsSystemLogsWhere(filter *service.OpsSystemLogFilter) (string, []any, bool) {
|
||||
clauses := make([]string, 0, 10)
|
||||
args := make([]any, 0, 10)
|
||||
@@ -1269,6 +1291,16 @@ func opsNullInt(v any) any {
|
||||
}
|
||||
}
|
||||
|
||||
// opsNullableIntPointer distinguishes an absent value from an explicitly
|
||||
// observed zero. Credential-stage failures intentionally persist upstream
|
||||
// status 0 because no inference request was sent.
|
||||
func opsNullableIntPointer(v *int) any {
|
||||
if v == nil {
|
||||
return sql.NullInt64{}
|
||||
}
|
||||
return sql.NullInt64{Int64: int64(*v), Valid: true}
|
||||
}
|
||||
|
||||
func opsNullInt16(v *int16) any {
|
||||
if v == nil {
|
||||
return sql.NullInt64{}
|
||||
|
||||
@@ -0,0 +1,37 @@
|
||||
//go:build unit
|
||||
|
||||
package repository
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestOpsInsertErrorLogArgsPreservesExplicitZeroUpstreamStatus(t *testing.T) {
|
||||
zero := 0
|
||||
args := opsInsertErrorLogArgs(&service.OpsInsertErrorLogInput{UpstreamStatusCode: &zero})
|
||||
|
||||
require.Len(t, args, 41)
|
||||
encoded, ok := args[27].(sql.NullInt64)
|
||||
require.True(t, ok)
|
||||
require.True(t, encoded.Valid)
|
||||
require.Zero(t, encoded.Int64)
|
||||
}
|
||||
|
||||
func TestOpsNullableIntPointerDistinguishesNilZeroAndStatus(t *testing.T) {
|
||||
missing := opsNullableIntPointer(nil).(sql.NullInt64)
|
||||
require.False(t, missing.Valid)
|
||||
|
||||
zeroValue := 0
|
||||
zero := opsNullableIntPointer(&zeroValue).(sql.NullInt64)
|
||||
require.True(t, zero.Valid)
|
||||
require.Zero(t, zero.Int64)
|
||||
|
||||
statusValue := 503
|
||||
status := opsNullableIntPointer(&statusValue).(sql.NullInt64)
|
||||
require.True(t, status.Valid)
|
||||
require.EqualValues(t, 503, status.Int64)
|
||||
}
|
||||
@@ -91,4 +91,21 @@ func TestGetErrorLogByID_DeletedKeyOwner(t *testing.T) {
|
||||
require.Equal(t, "sk-valid", valid.APIKeyPrefix)
|
||||
require.Empty(t, valid.AttemptedKeyPrefix, "attempted prefix and api key prefix are mutually exclusive")
|
||||
require.Nil(t, valid.DeletedKeyOwnerUserID, "valid key error has no deleted owner")
|
||||
|
||||
// ── Case 4: account_auth with no inference attempt preserves explicit 0 ──
|
||||
zero := 0
|
||||
credentialFailureID, err := repo.InsertErrorLog(ctx, &service.OpsInsertErrorLogInput{
|
||||
ErrorPhase: "account_auth",
|
||||
ErrorType: "upstream_error",
|
||||
Severity: "error",
|
||||
StatusCode: 503,
|
||||
UpstreamStatusCode: &zero,
|
||||
CreatedAt: time.Now(),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
credentialFailure, err := repo.GetErrorLogByID(ctx, credentialFailureID)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, credentialFailure.UpstreamStatusCode)
|
||||
require.Zero(t, *credentialFailure.UpstreamStatusCode)
|
||||
}
|
||||
|
||||
@@ -6,7 +6,6 @@ import (
|
||||
"errors"
|
||||
"hash/fnv"
|
||||
"log/slog"
|
||||
"net/url"
|
||||
"reflect"
|
||||
"sort"
|
||||
"strconv"
|
||||
@@ -1268,16 +1267,13 @@ func (a *Account) GetGrokBaseURL() string {
|
||||
if !a.IsGrok() {
|
||||
return ""
|
||||
}
|
||||
baseURL := a.GetCredential("base_url")
|
||||
if a.IsGrokOAuth() {
|
||||
if strings.TrimSpace(baseURL) == "" || isOfficialGrokAPIBaseURL(baseURL) {
|
||||
return xai.DefaultCLIBaseURL
|
||||
}
|
||||
if _, err := xai.ValidateTrustedBaseURL(baseURL); err == nil {
|
||||
return baseURL
|
||||
}
|
||||
// OAuth bearer credentials are subscription credentials and may only be
|
||||
// sent to the supported CLI gateway. Stored base_url values and unsafe
|
||||
// development overrides apply exclusively to API-key accounts.
|
||||
return xai.DefaultCLIBaseURL
|
||||
}
|
||||
baseURL := a.GetCredential("base_url")
|
||||
if baseURL != "" {
|
||||
return baseURL
|
||||
}
|
||||
@@ -1286,57 +1282,17 @@ func (a *Account) GetGrokBaseURL() string {
|
||||
|
||||
// GetGrokMediaBaseURL selects the upstream used by Grok Imagine APIs.
|
||||
//
|
||||
// OAuth text requests need the CLI subscription proxy, but that proxy has a
|
||||
// smaller request-body limit than the official Imagine API. Media requests can
|
||||
// contain large base64 inputs, so default OAuth accounts must use api.x.ai.
|
||||
// API-key accounts and explicit unsafe development overrides retain their
|
||||
// configured base URL.
|
||||
// OAuth media credentials have the same trust boundary as OAuth text traffic:
|
||||
// they are pinned to the supported CLI gateway even for large request bodies.
|
||||
// API-key accounts retain their configured public/custom upstream behavior.
|
||||
func (a *Account) GetGrokMediaBaseURL() string {
|
||||
if !a.IsGrok() {
|
||||
return ""
|
||||
}
|
||||
if !a.IsGrokOAuth() {
|
||||
return a.GetGrokBaseURL()
|
||||
if a.IsGrokOAuth() {
|
||||
return xai.DefaultCLIBaseURL
|
||||
}
|
||||
|
||||
baseURL := a.GetCredential("base_url")
|
||||
if strings.TrimSpace(baseURL) == "" || isOfficialGrokAPIBaseURL(baseURL) || isOfficialGrokCLIBaseURL(baseURL) {
|
||||
return xai.DefaultBaseURL
|
||||
}
|
||||
if _, err := xai.ValidateTrustedBaseURL(baseURL); err == nil {
|
||||
return baseURL
|
||||
}
|
||||
return xai.DefaultBaseURL
|
||||
}
|
||||
|
||||
func isOfficialGrokAPIBaseURL(raw string) bool {
|
||||
return isOfficialGrokBaseURL(raw, xai.DefaultBaseURL)
|
||||
}
|
||||
|
||||
func isOfficialGrokCLIBaseURL(raw string) bool {
|
||||
return isOfficialGrokBaseURL(raw, xai.DefaultCLIBaseURL)
|
||||
}
|
||||
|
||||
func isOfficialGrokBaseURL(raw, expected string) bool {
|
||||
parsed, err := url.Parse(strings.TrimSpace(raw))
|
||||
if err != nil || parsed == nil || parsed.Opaque != "" || parsed.User != nil || parsed.RawQuery != "" || parsed.Fragment != "" {
|
||||
return false
|
||||
}
|
||||
defaultURL, err := url.Parse(expected)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
if !strings.EqualFold(parsed.Scheme, defaultURL.Scheme) || !strings.EqualFold(parsed.Hostname(), defaultURL.Hostname()) {
|
||||
return false
|
||||
}
|
||||
if port := parsed.Port(); port != "" {
|
||||
portNumber, err := strconv.Atoi(port)
|
||||
if err != nil || portNumber != 443 {
|
||||
return false
|
||||
}
|
||||
}
|
||||
path := strings.TrimRight(parsed.Path, "/")
|
||||
return path == "" || path == strings.TrimRight(defaultURL.Path, "/")
|
||||
return a.GetGrokBaseURL()
|
||||
}
|
||||
|
||||
func (a *Account) GetGrokAccessToken() string {
|
||||
|
||||
@@ -255,7 +255,7 @@ func TestGetGrokBaseURLUsesSubscriptionProxyForOAuth(t *testing.T) {
|
||||
expected: xai.DefaultCLIBaseURL,
|
||||
},
|
||||
{
|
||||
name: "oauth non-default API port remains an explicit override",
|
||||
name: "oauth non-default API port remains pinned to CLI proxy",
|
||||
account: Account{
|
||||
Type: AccountTypeOAuth,
|
||||
Platform: PlatformGrok,
|
||||
@@ -263,7 +263,7 @@ func TestGetGrokBaseURLUsesSubscriptionProxyForOAuth(t *testing.T) {
|
||||
"base_url": "https://api.x.ai:8443/v1",
|
||||
},
|
||||
},
|
||||
expected: "https://api.x.ai:8443/v1",
|
||||
expected: xai.DefaultCLIBaseURL,
|
||||
},
|
||||
{
|
||||
name: "oauth explicit custom base_url stays pinned to CLI proxy by default",
|
||||
@@ -294,7 +294,7 @@ func TestGetGrokBaseURLUsesSubscriptionProxyForOAuth(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetGrokBaseURLAllowsExplicitOAuthOverrideWhenUnsafeOverridesEnabled(t *testing.T) {
|
||||
func TestGetGrokBaseURLPinsOAuthWhenUnsafeOverridesEnabled(t *testing.T) {
|
||||
t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
|
||||
account := Account{
|
||||
Type: AccountTypeOAuth,
|
||||
@@ -304,26 +304,26 @@ func TestGetGrokBaseURLAllowsExplicitOAuthOverrideWhenUnsafeOverridesEnabled(t *
|
||||
},
|
||||
}
|
||||
|
||||
require.Equal(t, "https://custom.example.com/v1", account.GetGrokBaseURL())
|
||||
require.Equal(t, xai.DefaultCLIBaseURL, account.GetGrokBaseURL())
|
||||
}
|
||||
|
||||
func TestGetGrokMediaBaseURLSeparatesOAuthMediaFromCLIProxy(t *testing.T) {
|
||||
func TestGetGrokMediaBaseURLPinsOAuthMediaToCLIProxy(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
account Account
|
||||
expected string
|
||||
}{
|
||||
{
|
||||
name: "oauth without base_url uses official media API",
|
||||
name: "oauth without base_url uses CLI subscription proxy",
|
||||
account: Account{
|
||||
Type: AccountTypeOAuth,
|
||||
Platform: PlatformGrok,
|
||||
Credentials: map[string]any{},
|
||||
},
|
||||
expected: xai.DefaultBaseURL,
|
||||
expected: xai.DefaultCLIBaseURL,
|
||||
},
|
||||
{
|
||||
name: "oauth stored CLI proxy uses official media API",
|
||||
name: "oauth stored CLI proxy stays on CLI subscription proxy",
|
||||
account: Account{
|
||||
Type: AccountTypeOAuth,
|
||||
Platform: PlatformGrok,
|
||||
@@ -331,10 +331,10 @@ func TestGetGrokMediaBaseURLSeparatesOAuthMediaFromCLIProxy(t *testing.T) {
|
||||
"base_url": xai.DefaultCLIBaseURL,
|
||||
},
|
||||
},
|
||||
expected: xai.DefaultBaseURL,
|
||||
expected: xai.DefaultCLIBaseURL,
|
||||
},
|
||||
{
|
||||
name: "oauth stored CLI proxy variant uses official media API",
|
||||
name: "oauth stored CLI proxy variant is canonicalized to CLI proxy",
|
||||
account: Account{
|
||||
Type: AccountTypeOAuth,
|
||||
Platform: PlatformGrok,
|
||||
@@ -342,10 +342,10 @@ func TestGetGrokMediaBaseURLSeparatesOAuthMediaFromCLIProxy(t *testing.T) {
|
||||
"base_url": "HTTPS://CLI-CHAT-PROXY.GROK.COM:443/%76%31/",
|
||||
},
|
||||
},
|
||||
expected: xai.DefaultBaseURL,
|
||||
expected: xai.DefaultCLIBaseURL,
|
||||
},
|
||||
{
|
||||
name: "oauth legacy official API remains on official media API",
|
||||
name: "oauth legacy official API is pinned to CLI proxy",
|
||||
account: Account{
|
||||
Type: AccountTypeOAuth,
|
||||
Platform: PlatformGrok,
|
||||
@@ -353,10 +353,10 @@ func TestGetGrokMediaBaseURLSeparatesOAuthMediaFromCLIProxy(t *testing.T) {
|
||||
"base_url": xai.DefaultBaseURL,
|
||||
},
|
||||
},
|
||||
expected: xai.DefaultBaseURL,
|
||||
expected: xai.DefaultCLIBaseURL,
|
||||
},
|
||||
{
|
||||
name: "oauth untrusted custom base_url is pinned to official media API",
|
||||
name: "oauth untrusted custom base_url is pinned to CLI proxy",
|
||||
account: Account{
|
||||
Type: AccountTypeOAuth,
|
||||
Platform: PlatformGrok,
|
||||
@@ -364,7 +364,7 @@ func TestGetGrokMediaBaseURLSeparatesOAuthMediaFromCLIProxy(t *testing.T) {
|
||||
"base_url": "https://custom.example.com/v1",
|
||||
},
|
||||
},
|
||||
expected: xai.DefaultBaseURL,
|
||||
expected: xai.DefaultCLIBaseURL,
|
||||
},
|
||||
{
|
||||
name: "API key retains its configured media API",
|
||||
@@ -395,7 +395,7 @@ func TestGetGrokMediaBaseURLSeparatesOAuthMediaFromCLIProxy(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetGrokMediaBaseURLAllowsExplicitOAuthOverrideWhenUnsafeOverridesEnabled(t *testing.T) {
|
||||
func TestGetGrokMediaBaseURLPinsOAuthWhenUnsafeOverridesEnabled(t *testing.T) {
|
||||
t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
|
||||
account := Account{
|
||||
Type: AccountTypeOAuth,
|
||||
@@ -405,5 +405,5 @@ func TestGetGrokMediaBaseURLAllowsExplicitOAuthOverrideWhenUnsafeOverridesEnable
|
||||
},
|
||||
}
|
||||
|
||||
require.Equal(t, "https://custom.example.com/v1", account.GetGrokMediaBaseURL())
|
||||
require.Equal(t, xai.DefaultCLIBaseURL, account.GetGrokMediaBaseURL())
|
||||
}
|
||||
|
||||
@@ -4,7 +4,7 @@ package service
|
||||
// dto 层做响应脱敏、service 层做更新合并都引用此清单——新增凭证类型时务必同步。
|
||||
var SensitiveCredentialKeys = []string{
|
||||
// OAuth
|
||||
"access_token", "refresh_token", "id_token",
|
||||
"access_token", "refresh_token", "id_token", "agent_private_key",
|
||||
// API Key 类
|
||||
"api_key", "session_key", "cookie",
|
||||
// 云服务凭据
|
||||
|
||||
@@ -65,6 +65,13 @@ func tryModelFilePricing(billingService *BillingService, model string, tokens Us
|
||||
if err != nil || pricing == nil {
|
||||
return nil
|
||||
}
|
||||
if billingService.shouldApplySessionLongContextPricing(tokens, pricing) {
|
||||
breakdown, err := billingService.CalculateCost(model, tokens, 1)
|
||||
if err != nil || breakdown == nil || breakdown.TotalCost <= 0 {
|
||||
return nil
|
||||
}
|
||||
return &breakdown.TotalCost
|
||||
}
|
||||
cost := float64(tokens.InputTokens)*pricing.InputPricePerToken +
|
||||
float64(tokens.OutputTokens)*pricing.OutputPricePerToken +
|
||||
float64(tokens.CacheCreationTokens)*pricing.CacheCreationPricePerToken +
|
||||
|
||||
@@ -459,6 +459,26 @@ func TestTryModelFilePricing_Success(t *testing.T) {
|
||||
require.InDelta(t, 0.2, *result, 1e-12)
|
||||
}
|
||||
|
||||
func TestTryModelFilePricing_AppliesLongContextPricing(t *testing.T) {
|
||||
bs := newTestBillingServiceWithPrices(map[string]*ModelPricing{
|
||||
"gpt-5.6-sol": {
|
||||
InputPricePerToken: 0.001,
|
||||
OutputPricePerToken: 0.002,
|
||||
CacheReadPricePerToken: 0.0001,
|
||||
LongContextInputThreshold: 100,
|
||||
LongContextInputMultiplier: 2,
|
||||
LongContextOutputMultiplier: 1.5,
|
||||
},
|
||||
})
|
||||
tokens := UsageTokens{InputTokens: 101, OutputTokens: 10, CacheReadTokens: 5}
|
||||
|
||||
result := tryModelFilePricing(bs, "gpt-5.6-sol", tokens)
|
||||
|
||||
require.NotNil(t, result)
|
||||
// Input and cache-read use the 2x input tier; output uses the 1.5x tier.
|
||||
require.InDelta(t, 0.233, *result, 1e-12)
|
||||
}
|
||||
|
||||
func TestTryModelFilePricing_PricingNotFound(t *testing.T) {
|
||||
// "nonexistent-model" does not match any fallback pattern
|
||||
bs := newTestBillingServiceWithPrices(map[string]*ModelPricing{})
|
||||
|
||||
@@ -15,6 +15,7 @@ import (
|
||||
"net/http/httptest"
|
||||
"regexp"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
@@ -72,6 +73,8 @@ type AccountTestService struct {
|
||||
httpUpstream HTTPUpstream
|
||||
cfg *config.Config
|
||||
tlsFPProfileService *TLSFingerprintProfileService
|
||||
agentIdentityTaskMu sync.Mutex
|
||||
agentIdentityWS agentIdentityWSConnectionInvalidator
|
||||
}
|
||||
|
||||
// NewAccountTestService creates a new AccountTestService
|
||||
@@ -544,9 +547,11 @@ func (s *AccountTestService) testOpenAIAccountConnection(c *gin.Context, account
|
||||
|
||||
if credentialAccount.IsOAuth() {
|
||||
isOAuth = true
|
||||
// OAuth - use Bearer token with ChatGPT internal API
|
||||
authToken = credentialAccount.GetOpenAIAccessToken()
|
||||
if authToken == "" {
|
||||
// Agent Identity signs each request and does not retain the OAuth token.
|
||||
if !credentialAccount.IsOpenAIAgentIdentity() {
|
||||
authToken = credentialAccount.GetOpenAIAccessToken()
|
||||
}
|
||||
if authToken == "" && !credentialAccount.IsOpenAIAgentIdentity() {
|
||||
return s.sendErrorAndEnd(c, "No access token available")
|
||||
}
|
||||
|
||||
@@ -591,8 +596,11 @@ func (s *AccountTestService) testOpenAIAccountConnection(c *gin.Context, account
|
||||
payload := createOpenAITestPayload(upstreamTestModelID, isOAuth)
|
||||
payloadBytes, _ := json.Marshal(payload)
|
||||
|
||||
// Send test_start event
|
||||
s.sendEvent(c, TestEvent{Type: "test_start", Model: testModelID})
|
||||
// Send test_start event once. A task-invalid Agent Identity response may
|
||||
// restart this probe after registering a replacement task.
|
||||
if !agentIdentityTaskRecoveryWasTried(ctx) {
|
||||
s.sendEvent(c, TestEvent{Type: "test_start", Model: testModelID})
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, "POST", apiURL, bytes.NewReader(payloadBytes))
|
||||
if err != nil {
|
||||
@@ -602,7 +610,19 @@ func (s *AccountTestService) testOpenAIAccountConnection(c *gin.Context, account
|
||||
|
||||
// Set common headers
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Authorization", "Bearer "+authToken)
|
||||
if credentialAccount.IsOpenAIAgentIdentity() {
|
||||
authHeaders, authErr := buildAgentIdentityAuthenticationHeaders(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, credentialAccount)
|
||||
if authErr != nil {
|
||||
return s.sendErrorAndEnd(c, "Failed to build Agent Identity authentication")
|
||||
}
|
||||
for key, values := range authHeaders {
|
||||
for _, value := range values {
|
||||
req.Header.Add(key, value)
|
||||
}
|
||||
}
|
||||
} else {
|
||||
req.Header.Set("Authorization", "Bearer "+authToken)
|
||||
}
|
||||
|
||||
// Set OAuth-specific headers for ChatGPT internal API
|
||||
if isOAuth {
|
||||
@@ -644,6 +664,15 @@ func (s *AccountTestService) testOpenAIAccountConnection(c *gin.Context, account
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
body = redactAgentIdentitySensitiveBodyForAccount(ctx, s.accountRepo, credentialAccount, body)
|
||||
if !agentIdentityTaskRecoveryWasTried(ctx) && credentialAccount.IsOpenAIAgentIdentity() && isAgentIdentityTaskInvalidHTTPResponse(resp.StatusCode, body) {
|
||||
expectedTaskID := credentialAccount.GetCredential("task_id")
|
||||
if err := ensureAgentIdentityTaskForAccount(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, credentialAccount, expectedTaskID); err != nil {
|
||||
return s.sendErrorAndEnd(c, fmt.Sprintf("Agent Identity task recovery failed: %s", err.Error()))
|
||||
}
|
||||
c.Request = c.Request.WithContext(markAgentIdentityTaskRecoveryTried(ctx))
|
||||
return s.testOpenAIAccountConnection(c, account, modelID, prompt, mode)
|
||||
}
|
||||
if resp.StatusCode == http.StatusTooManyRequests {
|
||||
s.reconcileOpenAI429State(ctx, account, resp.Header, body)
|
||||
}
|
||||
@@ -695,7 +724,7 @@ func (s *AccountTestService) testGrokAccountConnection(c *gin.Context, account *
|
||||
return s.sendErrorAndEnd(c, fmt.Sprintf("Unsupported Grok account type: %s", account.Type))
|
||||
}
|
||||
|
||||
apiURL, err := xai.BuildResponsesURL(account.GetGrokBaseURL())
|
||||
apiURL, err := buildGrokResponsesURL(account, s.cfg)
|
||||
if err != nil {
|
||||
return s.sendErrorAndEnd(c, fmt.Sprintf("Invalid Grok base URL: %s", err.Error()))
|
||||
}
|
||||
@@ -715,7 +744,9 @@ func (s *AccountTestService) testGrokAccountConnection(c *gin.Context, account *
|
||||
return s.sendErrorAndEnd(c, "Failed to create Grok test payload")
|
||||
}
|
||||
|
||||
s.sendEvent(c, TestEvent{Type: "test_start", Model: testModelID})
|
||||
if !agentIdentityTaskRecoveryWasTried(ctx) {
|
||||
s.sendEvent(c, TestEvent{Type: "test_start", Model: testModelID})
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, bytes.NewReader(payloadBytes))
|
||||
if err != nil {
|
||||
@@ -724,7 +755,9 @@ func (s *AccountTestService) testGrokAccountConnection(c *gin.Context, account *
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Accept", "application/json, text/event-stream")
|
||||
req.Header.Set("Authorization", "Bearer "+authToken)
|
||||
applyGrokCLIHeaders(req.Header)
|
||||
if account.IsGrokOAuth() {
|
||||
applyGrokCLIHeaders(req.Header)
|
||||
}
|
||||
|
||||
proxyURL := ""
|
||||
if account.ProxyID != nil && account.Proxy != nil {
|
||||
@@ -740,7 +773,7 @@ func (s *AccountTestService) testGrokAccountConnection(c *gin.Context, account *
|
||||
now := time.Now()
|
||||
snapshot := parseGrokQuotaSnapshot(resp.Header, resp.StatusCode, now)
|
||||
if snapshot != nil && s.accountRepo != nil {
|
||||
resetAt, limited := grokRateLimitResetAt(snapshot, now)
|
||||
resetAt, limited := grokRateLimitResetAtForAccount(account, snapshot, now)
|
||||
if limited {
|
||||
normalizeGrokExhaustedWindowResets(snapshot, resetAt, now)
|
||||
}
|
||||
@@ -749,7 +782,11 @@ func (s *AccountTestService) testGrokAccountConnection(c *gin.Context, account *
|
||||
})
|
||||
if limited {
|
||||
persistGrokRateLimit(ctx, s.accountRepo, account, resetAt)
|
||||
} else if isSuccessfulGrokRateLimitRecovery(account, snapshot) {
|
||||
clearGrokRateLimitAfterRecovery(ctx, s.accountRepo, account)
|
||||
}
|
||||
} else if s.accountRepo != nil && isSuccessfulGrokRateLimitRecovery(account, &xai.QuotaSnapshot{StatusCode: resp.StatusCode}) {
|
||||
clearGrokRateLimitAfterRecovery(ctx, s.accountRepo, account)
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
@@ -827,16 +864,26 @@ func (s *AccountTestService) testOpenAIChatCompletionsConnection(
|
||||
// resulting capability state on the account.
|
||||
func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account *Account, testModelID string) error {
|
||||
ctx := c.Request.Context()
|
||||
credentialAccount := account
|
||||
if account.IsShadow() {
|
||||
resolved, err := resolveCredentialAccount(ctx, s.accountRepo, account)
|
||||
if err != nil {
|
||||
return s.sendErrorAndEnd(c, "Failed to resolve account credentials")
|
||||
}
|
||||
credentialAccount = resolved
|
||||
}
|
||||
|
||||
authToken := ""
|
||||
apiURL := ""
|
||||
isOAuth := false
|
||||
|
||||
switch {
|
||||
case account.IsOAuth():
|
||||
case credentialAccount.IsOAuth():
|
||||
isOAuth = true
|
||||
authToken = account.GetOpenAIAccessToken()
|
||||
if authToken == "" {
|
||||
if !credentialAccount.IsOpenAIAgentIdentity() {
|
||||
authToken = credentialAccount.GetOpenAIAccessToken()
|
||||
}
|
||||
if authToken == "" && !credentialAccount.IsOpenAIAgentIdentity() {
|
||||
return s.sendErrorAndEnd(c, "No access token available")
|
||||
}
|
||||
apiURL = chatgptCodexAPIURL + "/compact"
|
||||
@@ -865,7 +912,9 @@ func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account
|
||||
c.Writer.Flush()
|
||||
|
||||
payloadBytes, _ := json.Marshal(createOpenAICompactProbePayload(testModelID))
|
||||
s.sendEvent(c, TestEvent{Type: "test_start", Model: testModelID})
|
||||
if !agentIdentityTaskRecoveryWasTried(ctx) {
|
||||
s.sendEvent(c, TestEvent{Type: "test_start", Model: testModelID})
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, "POST", apiURL, bytes.NewReader(payloadBytes))
|
||||
if err != nil {
|
||||
@@ -875,7 +924,19 @@ func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account
|
||||
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Accept", "application/json")
|
||||
req.Header.Set("Authorization", "Bearer "+authToken)
|
||||
if credentialAccount.IsOpenAIAgentIdentity() {
|
||||
authHeaders, authErr := buildAgentIdentityAuthenticationHeaders(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, credentialAccount)
|
||||
if authErr != nil {
|
||||
return s.sendErrorAndEnd(c, "Failed to build Agent Identity authentication")
|
||||
}
|
||||
for key, values := range authHeaders {
|
||||
for _, value := range values {
|
||||
req.Header.Add(key, value)
|
||||
}
|
||||
}
|
||||
} else {
|
||||
req.Header.Set("Authorization", "Bearer "+authToken)
|
||||
}
|
||||
req.Header.Set("OpenAI-Beta", "responses=experimental")
|
||||
req.Header.Set("Originator", "codex_cli_rs")
|
||||
req.Header.Set("User-Agent", codexCLIUserAgent)
|
||||
@@ -886,7 +947,7 @@ func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account
|
||||
|
||||
if isOAuth {
|
||||
req.Host = "chatgpt.com"
|
||||
setOpenAIChatGPTAccountHeaders(req.Header, account)
|
||||
setOpenAIChatGPTAccountHeaders(req.Header, credentialAccount)
|
||||
}
|
||||
|
||||
// 账号级请求头覆写:测试请求与真实转发保持一致的最终头
|
||||
@@ -909,6 +970,15 @@ func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
body, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
|
||||
body = redactAgentIdentitySensitiveBodyForAccount(ctx, s.accountRepo, credentialAccount, body)
|
||||
if !agentIdentityTaskRecoveryWasTried(ctx) && credentialAccount.IsOpenAIAgentIdentity() && isAgentIdentityTaskInvalidHTTPResponse(resp.StatusCode, body) {
|
||||
expectedTaskID := credentialAccount.GetCredential("task_id")
|
||||
if err := ensureAgentIdentityTaskForAccount(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, credentialAccount, expectedTaskID); err != nil {
|
||||
return s.sendErrorAndEnd(c, fmt.Sprintf("Agent Identity task recovery failed: %s", err.Error()))
|
||||
}
|
||||
c.Request = c.Request.WithContext(markAgentIdentityTaskRecoveryTried(ctx))
|
||||
return s.testOpenAICompactConnection(c, account, testModelID)
|
||||
}
|
||||
|
||||
if s.accountRepo != nil {
|
||||
updates := buildOpenAICompactProbeExtraUpdates(resp, body, nil, time.Now())
|
||||
@@ -1700,8 +1770,19 @@ func (s *AccountTestService) testOpenAIImageAPIKey(c *gin.Context, ctx context.C
|
||||
|
||||
// testOpenAIImageOAuth tests OpenAI image generation using an OAuth account via Codex /responses API.
|
||||
func (s *AccountTestService) testOpenAIImageOAuth(c *gin.Context, ctx context.Context, account *Account, modelID, prompt string) error {
|
||||
authToken := account.GetOpenAIAccessToken()
|
||||
if authToken == "" {
|
||||
credentialAccount := account
|
||||
if account.IsShadow() {
|
||||
resolved, err := resolveCredentialAccount(ctx, s.accountRepo, account)
|
||||
if err != nil {
|
||||
return s.sendErrorAndEnd(c, "Failed to resolve account credentials")
|
||||
}
|
||||
credentialAccount = resolved
|
||||
}
|
||||
authToken := ""
|
||||
if !credentialAccount.IsOpenAIAgentIdentity() {
|
||||
authToken = credentialAccount.GetOpenAIAccessToken()
|
||||
}
|
||||
if authToken == "" && !credentialAccount.IsOpenAIAgentIdentity() {
|
||||
return s.sendErrorAndEnd(c, "No access token available")
|
||||
}
|
||||
|
||||
@@ -1733,17 +1814,29 @@ func (s *AccountTestService) testOpenAIImageOAuth(c *gin.Context, ctx context.Co
|
||||
}
|
||||
req = req.WithContext(WithHTTPUpstreamProfile(req.Context(), HTTPUpstreamProfileOpenAI))
|
||||
req.Host = "chatgpt.com"
|
||||
req.Header.Set("Authorization", "Bearer "+authToken)
|
||||
if credentialAccount.IsOpenAIAgentIdentity() {
|
||||
authHeaders, authErr := buildAgentIdentityAuthenticationHeaders(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, credentialAccount)
|
||||
if authErr != nil {
|
||||
return s.sendErrorAndEnd(c, "Failed to build Agent Identity authentication")
|
||||
}
|
||||
for key, values := range authHeaders {
|
||||
for _, value := range values {
|
||||
req.Header.Add(key, value)
|
||||
}
|
||||
}
|
||||
} else {
|
||||
req.Header.Set("Authorization", "Bearer "+authToken)
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Accept", "text/event-stream")
|
||||
req.Header.Set("OpenAI-Beta", "responses=experimental")
|
||||
req.Header.Set("originator", "codex_cli_rs")
|
||||
if customUA := strings.TrimSpace(account.GetOpenAIUserAgent()); customUA != "" {
|
||||
if customUA := strings.TrimSpace(credentialAccount.GetOpenAIUserAgent()); customUA != "" {
|
||||
req.Header.Set("User-Agent", customUA)
|
||||
} else {
|
||||
req.Header.Set("User-Agent", codexCLIUserAgent)
|
||||
}
|
||||
setOpenAIChatGPTAccountHeaders(req.Header, account)
|
||||
setOpenAIChatGPTAccountHeaders(req.Header, credentialAccount)
|
||||
// 与真实转发一致:originator 与最终 User-Agent 首段配套(原 opencode 与 Codex UA 错配会 404,issue #3901)。
|
||||
enforceCodexIdentityHeaders(req.Header)
|
||||
|
||||
@@ -1762,6 +1855,7 @@ func (s *AccountTestService) testOpenAIImageOAuth(c *gin.Context, ctx context.Co
|
||||
}()
|
||||
if resp.StatusCode >= 400 {
|
||||
body, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
|
||||
body = redactAgentIdentitySensitiveBodyForAccount(ctx, s.accountRepo, credentialAccount, body)
|
||||
message := strings.TrimSpace(extractUpstreamErrorMessage(body))
|
||||
if message == "" {
|
||||
message = fmt.Sprintf("Responses API returned %d", resp.StatusCode)
|
||||
@@ -1773,6 +1867,7 @@ func (s *AccountTestService) testOpenAIImageOAuth(c *gin.Context, ctx context.Co
|
||||
if err != nil {
|
||||
return s.sendErrorAndEnd(c, fmt.Sprintf("Failed to read image response: %s", err.Error()))
|
||||
}
|
||||
body = redactAgentIdentitySensitiveBodyForAccount(ctx, s.accountRepo, credentialAccount, body)
|
||||
|
||||
results, _, _, _, _, err := collectOpenAIImagesFromResponsesBody(body)
|
||||
if err != nil {
|
||||
|
||||
@@ -40,8 +40,9 @@ func TestAccountTestService_TestAccountConnection_GrokUsesXAIResponses(t *testin
|
||||
Schedulable: true,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "grok-access-token",
|
||||
"expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
|
||||
"access_token": "grok-access-token",
|
||||
"refresh_token": "grok-refresh-token",
|
||||
"expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
|
||||
"model_mapping": map[string]any{
|
||||
"grok": "grok-4.3",
|
||||
},
|
||||
@@ -92,8 +93,9 @@ func TestAccountTestService_TestAccountConnection_GrokDefaultsEmptyModelTo45(t *
|
||||
Schedulable: true,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "grok-access-token",
|
||||
"expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
|
||||
"access_token": "grok-access-token",
|
||||
"refresh_token": "grok-refresh-token",
|
||||
"expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
|
||||
},
|
||||
}
|
||||
repo := &mockAccountRepoForGemini{accountsByID: map[int64]*Account{account.ID: account}}
|
||||
@@ -133,8 +135,9 @@ func TestAccountTestService_Grok429PersistsRateLimitReset(t *testing.T) {
|
||||
Schedulable: true,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "grok-access-token",
|
||||
"expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
|
||||
"access_token": "grok-access-token",
|
||||
"refresh_token": "grok-refresh-token",
|
||||
"expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
|
||||
},
|
||||
}
|
||||
baseRepo := &mockAccountRepoForGemini{accountsByID: map[int64]*Account{account.ID: account}}
|
||||
@@ -166,8 +169,9 @@ func TestAccountTestService_Grok429WithoutQuotaHeadersUsesFallback(t *testing.T)
|
||||
ID: 15, Name: "grok-oauth-limited-no-headers", Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "grok-access-token",
|
||||
"expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
|
||||
"access_token": "grok-access-token",
|
||||
"refresh_token": "grok-refresh-token",
|
||||
"expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
|
||||
},
|
||||
}
|
||||
baseRepo := &mockAccountRepoForGemini{accountsByID: map[int64]*Account{account.ID: account}}
|
||||
|
||||
@@ -299,6 +299,8 @@ type AccountUsageService struct {
|
||||
cache *UsageCache
|
||||
identityCache IdentityCache
|
||||
tlsFPProfileService *TLSFingerprintProfileService
|
||||
agentIdentityTaskMu sync.Mutex
|
||||
agentIdentityWS agentIdentityWSConnectionInvalidator
|
||||
}
|
||||
|
||||
// NewAccountUsageService 创建AccountUsageService实例
|
||||
@@ -692,8 +694,11 @@ func (s *AccountUsageService) probeOpenAICodexSnapshot(ctx context.Context, acco
|
||||
if account == nil || !account.IsOAuth() {
|
||||
return nil, nil
|
||||
}
|
||||
accessToken := account.GetOpenAIAccessToken()
|
||||
if accessToken == "" {
|
||||
accessToken := ""
|
||||
if !account.IsOpenAIAgentIdentity() {
|
||||
accessToken = account.GetOpenAIAccessToken()
|
||||
}
|
||||
if accessToken == "" && !account.IsOpenAIAgentIdentity() {
|
||||
return nil, fmt.Errorf("no access token available")
|
||||
}
|
||||
modelID := openaipkg.DefaultTestModel
|
||||
@@ -711,7 +716,19 @@ func (s *AccountUsageService) probeOpenAICodexSnapshot(ctx context.Context, acco
|
||||
}
|
||||
req.Host = "chatgpt.com"
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Authorization", "Bearer "+accessToken)
|
||||
if account.IsOpenAIAgentIdentity() {
|
||||
authHeaders, authErr := buildAgentIdentityAuthenticationHeaders(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, account)
|
||||
if authErr != nil {
|
||||
return nil, fmt.Errorf("build Agent Identity authentication: %w", authErr)
|
||||
}
|
||||
for key, values := range authHeaders {
|
||||
for _, value := range values {
|
||||
req.Header.Add(key, value)
|
||||
}
|
||||
}
|
||||
} else {
|
||||
req.Header.Set("Authorization", "Bearer "+accessToken)
|
||||
}
|
||||
req.Header.Set("Accept", "text/event-stream")
|
||||
req.Header.Set("OpenAI-Beta", "responses=experimental")
|
||||
req.Header.Set("Originator", "codex_cli_rs")
|
||||
|
||||
@@ -92,6 +92,9 @@ const (
|
||||
contentModerationCleanupInterval = 24 * time.Hour
|
||||
contentModerationCleanupTimeout = 30 * time.Minute
|
||||
contentModerationCleanupDelay = 5 * time.Minute
|
||||
|
||||
contentModerationRuntimeCacheTTL = time.Second
|
||||
contentModerationRuntimeRefreshTimeout = 5 * time.Second
|
||||
)
|
||||
|
||||
var contentModerationCategoryOrder = []string{
|
||||
@@ -512,10 +515,22 @@ type ContentModerationService struct {
|
||||
lastCleanupUnix atomic.Int64
|
||||
lastCleanupDeletedHit atomic.Int64
|
||||
lastCleanupDeletedNonHit atomic.Int64
|
||||
runtimeSnapshot atomic.Pointer[contentModerationRuntimeSnapshot]
|
||||
runtimeRefreshMu sync.Mutex
|
||||
runtimeCacheTTL time.Duration
|
||||
runtimeRefreshRetryAt atomic.Int64
|
||||
keyHealthMu sync.Mutex
|
||||
keyHealth map[string]*contentModerationKeyHealth
|
||||
}
|
||||
|
||||
type contentModerationRuntimeSnapshot struct {
|
||||
riskControlEnabled bool
|
||||
config *ContentModerationConfig
|
||||
keywordMatcher *contentModerationKeywordMatcher
|
||||
configDigest [sha256.Size]byte
|
||||
loadedAt time.Time
|
||||
}
|
||||
|
||||
type contentModerationTask struct {
|
||||
input ContentModerationCheckInput
|
||||
content ContentModerationInput
|
||||
@@ -700,6 +715,7 @@ func (s *ContentModerationService) UpdateConfig(ctx context.Context, input Updat
|
||||
if err := s.settingRepo.Set(ctx, SettingKeyContentModerationConfig, string(raw)); err != nil {
|
||||
return nil, fmt.Errorf("save content moderation config: %w", err)
|
||||
}
|
||||
s.replaceRuntimeConfig(cfg, raw)
|
||||
return s.configView(cfg), nil
|
||||
}
|
||||
|
||||
@@ -776,16 +792,7 @@ func (s *ContentModerationService) Check(ctx context.Context, input ContentModer
|
||||
"protocol", input.Protocol)
|
||||
return allow, nil
|
||||
}
|
||||
if !s.isRiskControlEnabled(ctx) {
|
||||
slog.Info("content_moderation.skip_feature_disabled",
|
||||
"user_id", input.UserID,
|
||||
"api_key_id", input.APIKeyID,
|
||||
"group_id", contentModerationLogGroupID(input.GroupID),
|
||||
"endpoint", input.Endpoint,
|
||||
"protocol", input.Protocol)
|
||||
return allow, nil
|
||||
}
|
||||
cfg, err := s.loadConfig(ctx)
|
||||
runtimeSnapshot, err := s.loadRuntimeSnapshot(ctx)
|
||||
if err != nil {
|
||||
slog.Warn("content_moderation.skip_config_load_failed",
|
||||
"user_id", input.UserID,
|
||||
@@ -796,6 +803,16 @@ func (s *ContentModerationService) Check(ctx context.Context, input ContentModer
|
||||
"error", err)
|
||||
return allow, nil
|
||||
}
|
||||
if !runtimeSnapshot.riskControlEnabled {
|
||||
slog.Info("content_moderation.skip_feature_disabled",
|
||||
"user_id", input.UserID,
|
||||
"api_key_id", input.APIKeyID,
|
||||
"group_id", contentModerationLogGroupID(input.GroupID),
|
||||
"endpoint", input.Endpoint,
|
||||
"protocol", input.Protocol)
|
||||
return allow, nil
|
||||
}
|
||||
cfg := runtimeSnapshot.config
|
||||
inGroupScope := cfg.includesGroup(input.GroupID)
|
||||
inModelScope := cfg.includesModel(input.Model)
|
||||
slog.Info("content_moderation.config_loaded",
|
||||
@@ -885,7 +902,7 @@ func (s *ContentModerationService) Check(ctx context.Context, input ContentModer
|
||||
hashText := content.Hash()
|
||||
if cfg.Mode == ContentModerationModePreBlock {
|
||||
if cfg.KeywordBlockingMode != ContentModerationKeywordModeAPIOnly && len(cfg.BlockedKeywords) > 0 {
|
||||
if keyword, hit := matchBlockedKeyword(content.Text, cfg.BlockedKeywords); hit {
|
||||
if keyword, hit := runtimeSnapshot.matchBlockedKeyword(content.Text); hit {
|
||||
s.recordPreBlockSyncMetric(0, ContentModerationActionKeywordBlock)
|
||||
slog.Info("content_moderation.keyword_block",
|
||||
"user_id", input.UserID,
|
||||
@@ -1178,12 +1195,13 @@ func (s *ContentModerationService) enqueueRecord(input ContentModerationCheckInp
|
||||
func (s *ContentModerationService) worker(id int) {
|
||||
for {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), maxContentModerationTimeoutMS*time.Millisecond+10*time.Second)
|
||||
cfg, err := s.loadConfig(ctx)
|
||||
if err != nil || id >= cfg.WorkerCount {
|
||||
runtimeSnapshot, err := s.loadRuntimeSnapshot(ctx)
|
||||
if err != nil || runtimeSnapshot == nil || runtimeSnapshot.config == nil || id >= runtimeSnapshot.config.WorkerCount {
|
||||
cancel()
|
||||
time.Sleep(time.Second)
|
||||
continue
|
||||
}
|
||||
cfg := runtimeSnapshot.config
|
||||
task, ok := s.dequeueAsyncTask(ctx, time.Second)
|
||||
if !ok {
|
||||
cancel()
|
||||
@@ -1438,15 +1456,18 @@ func (s *ContentModerationService) runCleanupOnce() {
|
||||
}
|
||||
|
||||
func (s *ContentModerationService) loadConfig(ctx context.Context) (*ContentModerationConfig, error) {
|
||||
cfg := defaultContentModerationConfig()
|
||||
raw, err := s.settingRepo.GetValue(ctx, SettingKeyContentModerationConfig)
|
||||
if err != nil {
|
||||
if errors.Is(err, ErrSettingNotFound) {
|
||||
cfg.normalize()
|
||||
return cfg, nil
|
||||
return parseContentModerationConfig("")
|
||||
}
|
||||
return nil, fmt.Errorf("get content moderation config: %w", err)
|
||||
}
|
||||
return parseContentModerationConfig(raw)
|
||||
}
|
||||
|
||||
func parseContentModerationConfig(raw string) (*ContentModerationConfig, error) {
|
||||
cfg := defaultContentModerationConfig()
|
||||
if strings.TrimSpace(raw) == "" {
|
||||
cfg.normalize()
|
||||
return cfg, nil
|
||||
@@ -1458,6 +1479,137 @@ func (s *ContentModerationService) loadConfig(ctx context.Context) (*ContentMode
|
||||
return cfg, nil
|
||||
}
|
||||
|
||||
func (s *ContentModerationService) loadRuntimeSnapshot(ctx context.Context) (*contentModerationRuntimeSnapshot, error) {
|
||||
if s == nil || s.settingRepo == nil {
|
||||
return nil, errors.New("content moderation setting repository unavailable")
|
||||
}
|
||||
now := time.Now()
|
||||
if snapshot := s.runtimeSnapshot.Load(); snapshot != nil {
|
||||
if now.Sub(snapshot.loadedAt) < s.runtimeSnapshotTTL() {
|
||||
return snapshot, nil
|
||||
}
|
||||
s.triggerRuntimeSnapshotRefresh()
|
||||
return snapshot, nil
|
||||
}
|
||||
|
||||
s.runtimeRefreshMu.Lock()
|
||||
defer s.runtimeRefreshMu.Unlock()
|
||||
if snapshot := s.runtimeSnapshot.Load(); snapshot != nil {
|
||||
return snapshot, nil
|
||||
}
|
||||
return s.refreshRuntimeSnapshot(ctx)
|
||||
}
|
||||
|
||||
func (s *ContentModerationService) runtimeSnapshotTTL() time.Duration {
|
||||
if s != nil && s.runtimeCacheTTL > 0 {
|
||||
return s.runtimeCacheTTL
|
||||
}
|
||||
return contentModerationRuntimeCacheTTL
|
||||
}
|
||||
|
||||
func (s *ContentModerationService) triggerRuntimeSnapshotRefresh() {
|
||||
if s == nil || s.runtimeRefreshDeferred() || !s.runtimeRefreshMu.TryLock() {
|
||||
return
|
||||
}
|
||||
if s.runtimeRefreshDeferred() {
|
||||
s.runtimeRefreshMu.Unlock()
|
||||
return
|
||||
}
|
||||
go func() {
|
||||
defer s.runtimeRefreshMu.Unlock()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), contentModerationRuntimeRefreshTimeout)
|
||||
defer cancel()
|
||||
if _, err := s.refreshRuntimeSnapshot(ctx); err != nil {
|
||||
s.runtimeRefreshRetryAt.Store(time.Now().Add(s.runtimeSnapshotTTL()).UnixNano())
|
||||
slog.Warn("content_moderation.runtime_snapshot_refresh_failed", "error", err)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
func (s *ContentModerationService) runtimeRefreshDeferred() bool {
|
||||
if s == nil {
|
||||
return false
|
||||
}
|
||||
return time.Now().UnixNano() < s.runtimeRefreshRetryAt.Load()
|
||||
}
|
||||
|
||||
func (s *ContentModerationService) refreshRuntimeSnapshot(ctx context.Context) (*contentModerationRuntimeSnapshot, error) {
|
||||
values, err := s.settingRepo.GetMultiple(ctx, []string{
|
||||
SettingKeyRiskControlEnabled,
|
||||
SettingKeyContentModerationConfig,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get content moderation runtime settings: %w", err)
|
||||
}
|
||||
rawConfig := values[SettingKeyContentModerationConfig]
|
||||
configDigest := sha256.Sum256([]byte(rawConfig))
|
||||
if current := s.runtimeSnapshot.Load(); current != nil && current.configDigest == configDigest {
|
||||
snapshot := &contentModerationRuntimeSnapshot{
|
||||
riskControlEnabled: values[SettingKeyRiskControlEnabled] == "true",
|
||||
config: current.config,
|
||||
keywordMatcher: current.keywordMatcher,
|
||||
configDigest: configDigest,
|
||||
loadedAt: time.Now(),
|
||||
}
|
||||
s.runtimeSnapshot.Store(snapshot)
|
||||
s.runtimeRefreshRetryAt.Store(0)
|
||||
return snapshot, nil
|
||||
}
|
||||
cfg, err := parseContentModerationConfig(rawConfig)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
snapshot := &contentModerationRuntimeSnapshot{
|
||||
riskControlEnabled: values[SettingKeyRiskControlEnabled] == "true",
|
||||
config: cfg,
|
||||
keywordMatcher: newContentModerationKeywordMatcher(cfg.BlockedKeywords),
|
||||
configDigest: configDigest,
|
||||
loadedAt: time.Now(),
|
||||
}
|
||||
s.runtimeSnapshot.Store(snapshot)
|
||||
s.runtimeRefreshRetryAt.Store(0)
|
||||
return snapshot, nil
|
||||
}
|
||||
|
||||
func (s *ContentModerationService) replaceRuntimeConfig(cfg *ContentModerationConfig, raw []byte) {
|
||||
if s == nil || cfg == nil {
|
||||
return
|
||||
}
|
||||
s.runtimeRefreshMu.Lock()
|
||||
hasSnapshot := s.runtimeSnapshot.Load() != nil
|
||||
s.runtimeRefreshMu.Unlock()
|
||||
if !hasSnapshot {
|
||||
return
|
||||
}
|
||||
config := cloneContentModerationConfig(cfg)
|
||||
keywordMatcher := newContentModerationKeywordMatcher(cfg.BlockedKeywords)
|
||||
configDigest := sha256.Sum256(raw)
|
||||
|
||||
s.runtimeRefreshMu.Lock()
|
||||
defer s.runtimeRefreshMu.Unlock()
|
||||
current := s.runtimeSnapshot.Load()
|
||||
if current == nil {
|
||||
return
|
||||
}
|
||||
s.runtimeSnapshot.Store(&contentModerationRuntimeSnapshot{
|
||||
riskControlEnabled: current.riskControlEnabled,
|
||||
config: config,
|
||||
keywordMatcher: keywordMatcher,
|
||||
configDigest: configDigest,
|
||||
loadedAt: time.Now(),
|
||||
})
|
||||
}
|
||||
|
||||
func (s *contentModerationRuntimeSnapshot) matchBlockedKeyword(text string) (string, bool) {
|
||||
if s == nil || s.config == nil {
|
||||
return "", false
|
||||
}
|
||||
if s.keywordMatcher != nil {
|
||||
return s.keywordMatcher.Match(text)
|
||||
}
|
||||
return matchBlockedKeyword(text, s.config.BlockedKeywords)
|
||||
}
|
||||
|
||||
func (s *ContentModerationService) isRiskControlEnabled(ctx context.Context) bool {
|
||||
raw, err := s.settingRepo.GetValue(ctx, SettingKeyRiskControlEnabled)
|
||||
if err != nil {
|
||||
|
||||
@@ -0,0 +1,222 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"strings"
|
||||
)
|
||||
|
||||
type contentModerationKeywordMatcher struct {
|
||||
nodes []contentModerationKeywordNode
|
||||
edges []contentModerationKeywordEdge
|
||||
rootTransitions [256]int32
|
||||
keywords []string
|
||||
}
|
||||
|
||||
type contentModerationKeywordNode struct {
|
||||
failure int32
|
||||
bestKeyword int32
|
||||
edgeStart uint32
|
||||
edgeCount uint16
|
||||
}
|
||||
|
||||
type contentModerationKeywordEdge struct {
|
||||
target int32
|
||||
label byte
|
||||
}
|
||||
|
||||
type contentModerationKeywordBuildEdge struct {
|
||||
target int32
|
||||
nextSibling int32
|
||||
label byte
|
||||
}
|
||||
|
||||
func newContentModerationKeywordMatcher(keywords []string) *contentModerationKeywordMatcher {
|
||||
if len(keywords) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
buildNodes := []contentModerationKeywordNode{newContentModerationKeywordNode()}
|
||||
buildEdges := make([]contentModerationKeywordBuildEdge, 0)
|
||||
originalKeywords := append([]string(nil), keywords...)
|
||||
|
||||
for keywordIndex, keyword := range keywords {
|
||||
if keyword == "" {
|
||||
continue
|
||||
}
|
||||
state := int32(0)
|
||||
for _, label := range []byte(strings.ToLower(keyword)) {
|
||||
next := contentModerationKeywordBuildTransition(buildNodes, buildEdges, state, label)
|
||||
if next < 0 {
|
||||
next = int32(len(buildNodes))
|
||||
buildNodes = append(buildNodes, newContentModerationKeywordNode())
|
||||
buildEdges = append(buildEdges, contentModerationKeywordBuildEdge{
|
||||
target: next,
|
||||
nextSibling: contentModerationKeywordBuildFirstEdge(buildNodes[state]),
|
||||
label: label,
|
||||
})
|
||||
buildNodes[state].edgeStart = uint32(len(buildEdges))
|
||||
}
|
||||
state = next
|
||||
}
|
||||
if current := buildNodes[state].bestKeyword; current < 0 || int32(keywordIndex) < current {
|
||||
buildNodes[state].bestKeyword = int32(keywordIndex)
|
||||
}
|
||||
}
|
||||
|
||||
if len(buildNodes) == 1 {
|
||||
return nil
|
||||
}
|
||||
|
||||
queue := make([]int32, 0, len(buildNodes)-1)
|
||||
var rootTransitions [256]int32
|
||||
for edgeIndex := contentModerationKeywordBuildFirstEdge(buildNodes[0]); edgeIndex >= 0; edgeIndex = buildEdges[edgeIndex].nextSibling {
|
||||
edge := buildEdges[edgeIndex]
|
||||
rootTransitions[edge.label] = edge.target
|
||||
queue = append(queue, edge.target)
|
||||
}
|
||||
|
||||
for queueIndex := 0; queueIndex < len(queue); queueIndex++ {
|
||||
state := queue[queueIndex]
|
||||
for edgeIndex := contentModerationKeywordBuildFirstEdge(buildNodes[state]); edgeIndex >= 0; edgeIndex = buildEdges[edgeIndex].nextSibling {
|
||||
edge := buildEdges[edgeIndex]
|
||||
failure := buildNodes[state].failure
|
||||
fallback := contentModerationKeywordBuildTransition(buildNodes, buildEdges, failure, edge.label)
|
||||
for fallback < 0 && failure != 0 {
|
||||
failure = buildNodes[failure].failure
|
||||
fallback = contentModerationKeywordBuildTransition(buildNodes, buildEdges, failure, edge.label)
|
||||
}
|
||||
if fallback >= 0 {
|
||||
buildNodes[edge.target].failure = fallback
|
||||
}
|
||||
buildNodes[edge.target].bestKeyword = minKeywordIndex(
|
||||
buildNodes[edge.target].bestKeyword,
|
||||
buildNodes[buildNodes[edge.target].failure].bestKeyword,
|
||||
)
|
||||
queue = append(queue, edge.target)
|
||||
}
|
||||
}
|
||||
|
||||
edges := make([]contentModerationKeywordEdge, 0, len(buildEdges))
|
||||
var outgoing [256]contentModerationKeywordEdge
|
||||
for nodeIndex := range buildNodes {
|
||||
count := 0
|
||||
for edgeIndex := contentModerationKeywordBuildFirstEdge(buildNodes[nodeIndex]); edgeIndex >= 0; edgeIndex = buildEdges[edgeIndex].nextSibling {
|
||||
edge := buildEdges[edgeIndex]
|
||||
outgoing[count] = contentModerationKeywordEdge{target: edge.target, label: edge.label}
|
||||
count++
|
||||
}
|
||||
for index := 1; index < count; index++ {
|
||||
current := outgoing[index]
|
||||
insertAt := index
|
||||
for insertAt > 0 && current.label < outgoing[insertAt-1].label {
|
||||
outgoing[insertAt] = outgoing[insertAt-1]
|
||||
insertAt--
|
||||
}
|
||||
outgoing[insertAt] = current
|
||||
}
|
||||
buildNodes[nodeIndex].edgeStart = uint32(len(edges))
|
||||
buildNodes[nodeIndex].edgeCount = uint16(count)
|
||||
edges = append(edges, outgoing[:count]...)
|
||||
}
|
||||
|
||||
return &contentModerationKeywordMatcher{
|
||||
nodes: buildNodes,
|
||||
edges: edges,
|
||||
rootTransitions: rootTransitions,
|
||||
keywords: originalKeywords,
|
||||
}
|
||||
}
|
||||
|
||||
func newContentModerationKeywordNode() contentModerationKeywordNode {
|
||||
return contentModerationKeywordNode{bestKeyword: -1}
|
||||
}
|
||||
|
||||
func contentModerationKeywordBuildFirstEdge(node contentModerationKeywordNode) int32 {
|
||||
if node.edgeStart == 0 {
|
||||
return -1
|
||||
}
|
||||
return int32(node.edgeStart - 1)
|
||||
}
|
||||
|
||||
func contentModerationKeywordBuildTransition(
|
||||
nodes []contentModerationKeywordNode,
|
||||
edges []contentModerationKeywordBuildEdge,
|
||||
state int32,
|
||||
label byte,
|
||||
) int32 {
|
||||
if state < 0 || int(state) >= len(nodes) {
|
||||
return -1
|
||||
}
|
||||
for edgeIndex := contentModerationKeywordBuildFirstEdge(nodes[state]); edgeIndex >= 0; edgeIndex = edges[edgeIndex].nextSibling {
|
||||
if edges[edgeIndex].label == label {
|
||||
return edges[edgeIndex].target
|
||||
}
|
||||
}
|
||||
return -1
|
||||
}
|
||||
|
||||
func minKeywordIndex(left, right int32) int32 {
|
||||
if left < 0 {
|
||||
return right
|
||||
}
|
||||
if right < 0 || left < right {
|
||||
return left
|
||||
}
|
||||
return right
|
||||
}
|
||||
|
||||
func (m *contentModerationKeywordMatcher) Match(text string) (string, bool) {
|
||||
if m == nil || text == "" || len(m.nodes) == 0 || len(m.keywords) == 0 {
|
||||
return "", false
|
||||
}
|
||||
lower := strings.ToLower(text)
|
||||
state := int32(0)
|
||||
bestKeyword := int32(-1)
|
||||
for index := 0; index < len(lower); index++ {
|
||||
label := lower[index]
|
||||
for {
|
||||
next := m.next(state, label)
|
||||
if next != 0 {
|
||||
state = next
|
||||
break
|
||||
}
|
||||
if state == 0 {
|
||||
break
|
||||
}
|
||||
state = m.nodes[state].failure
|
||||
}
|
||||
bestKeyword = minKeywordIndex(bestKeyword, m.nodes[state].bestKeyword)
|
||||
if bestKeyword == 0 {
|
||||
return m.keywords[0], true
|
||||
}
|
||||
}
|
||||
if bestKeyword < 0 || int(bestKeyword) >= len(m.keywords) {
|
||||
return "", false
|
||||
}
|
||||
return m.keywords[bestKeyword], true
|
||||
}
|
||||
|
||||
func (m *contentModerationKeywordMatcher) next(state int32, label byte) int32 {
|
||||
if state == 0 {
|
||||
return m.rootTransitions[label]
|
||||
}
|
||||
if state < 0 || int(state) >= len(m.nodes) {
|
||||
return 0
|
||||
}
|
||||
node := m.nodes[state]
|
||||
left := int(node.edgeStart)
|
||||
right := left + int(node.edgeCount)
|
||||
for left < right {
|
||||
middle := left + (right-left)/2
|
||||
edge := m.edges[middle]
|
||||
if edge.label < label {
|
||||
left = middle + 1
|
||||
continue
|
||||
}
|
||||
right = middle
|
||||
}
|
||||
end := int(node.edgeStart) + int(node.edgeCount)
|
||||
if left < end && m.edges[left].label == label {
|
||||
return m.edges[left].target
|
||||
}
|
||||
return 0
|
||||
}
|
||||
@@ -0,0 +1,59 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"math/rand"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestContentModerationKeywordMatcherMatchesLegacyBehavior(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
text string
|
||||
keywords []string
|
||||
}{
|
||||
{name: "miss", text: "clean prompt", keywords: []string{"blocked", "secret"}},
|
||||
{name: "case insensitive", text: "contains SECRET value", keywords: []string{"secret"}},
|
||||
{name: "configured order wins", text: "early appears before later", keywords: []string{"later", "early"}},
|
||||
{name: "overlap uses configured order", text: "abc", keywords: []string{"bc", "abc"}},
|
||||
{name: "unicode", text: "这里包含敏感词和世界", keywords: []string{"世界", "敏感词"}},
|
||||
{name: "duplicates", text: "duplicate", keywords: []string{"duplicate", "DUPLICATE"}},
|
||||
{name: "empty entries", text: "blocked", keywords: []string{"", "blocked"}},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
wantKeyword, wantHit := matchBlockedKeyword(tt.text, tt.keywords)
|
||||
gotKeyword, gotHit := newContentModerationKeywordMatcher(tt.keywords).Match(tt.text)
|
||||
require.Equal(t, wantHit, gotHit)
|
||||
require.Equal(t, wantKeyword, gotKeyword)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestContentModerationKeywordMatcherRandomizedParity(t *testing.T) {
|
||||
rng := rand.New(rand.NewSource(20260714))
|
||||
const alphabet = "abcXYZ"
|
||||
for iteration := 0; iteration < 1000; iteration++ {
|
||||
keywords := make([]string, 1+rng.Intn(30))
|
||||
for index := range keywords {
|
||||
length := 1 + rng.Intn(8)
|
||||
var value strings.Builder
|
||||
for range length {
|
||||
_ = value.WriteByte(alphabet[rng.Intn(len(alphabet))])
|
||||
}
|
||||
keywords[index] = value.String()
|
||||
}
|
||||
var text strings.Builder
|
||||
for range 20 + rng.Intn(100) {
|
||||
_ = text.WriteByte(alphabet[rng.Intn(len(alphabet))])
|
||||
}
|
||||
|
||||
wantKeyword, wantHit := matchBlockedKeyword(text.String(), keywords)
|
||||
gotKeyword, gotHit := newContentModerationKeywordMatcher(keywords).Match(text.String())
|
||||
require.Equal(t, wantHit, gotHit, "iteration %d", iteration)
|
||||
require.Equal(t, wantKeyword, gotKeyword, "iteration %d", iteration)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,451 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type contentModerationRuntimeSettingRepo struct {
|
||||
mu sync.Mutex
|
||||
values map[string]string
|
||||
getValueCalls int
|
||||
getMultipleCalls int
|
||||
getMultipleErr error
|
||||
getMultipleStart chan<- struct{}
|
||||
getMultipleWait <-chan struct{}
|
||||
}
|
||||
|
||||
func (r *contentModerationRuntimeSettingRepo) Get(_ context.Context, key string) (*Setting, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
value, ok := r.values[key]
|
||||
if !ok {
|
||||
return nil, ErrSettingNotFound
|
||||
}
|
||||
return &Setting{Key: key, Value: value}, nil
|
||||
}
|
||||
|
||||
func (r *contentModerationRuntimeSettingRepo) GetValue(_ context.Context, key string) (string, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.getValueCalls++
|
||||
value, ok := r.values[key]
|
||||
if !ok {
|
||||
return "", ErrSettingNotFound
|
||||
}
|
||||
return value, nil
|
||||
}
|
||||
|
||||
func (r *contentModerationRuntimeSettingRepo) Set(_ context.Context, key, value string) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
if r.values == nil {
|
||||
r.values = make(map[string]string)
|
||||
}
|
||||
r.values[key] = value
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *contentModerationRuntimeSettingRepo) GetMultiple(_ context.Context, keys []string) (map[string]string, error) {
|
||||
r.mu.Lock()
|
||||
r.getMultipleCalls++
|
||||
if err := r.getMultipleErr; err != nil {
|
||||
r.mu.Unlock()
|
||||
return nil, err
|
||||
}
|
||||
out := make(map[string]string, len(keys))
|
||||
for _, key := range keys {
|
||||
if value, ok := r.values[key]; ok {
|
||||
out[key] = value
|
||||
}
|
||||
}
|
||||
start := r.getMultipleStart
|
||||
wait := r.getMultipleWait
|
||||
r.getMultipleStart = nil
|
||||
r.getMultipleWait = nil
|
||||
r.mu.Unlock()
|
||||
if start != nil {
|
||||
start <- struct{}{}
|
||||
}
|
||||
if wait != nil {
|
||||
<-wait
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (r *contentModerationRuntimeSettingRepo) SetMultiple(_ context.Context, values map[string]string) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
if r.values == nil {
|
||||
r.values = make(map[string]string)
|
||||
}
|
||||
for key, value := range values {
|
||||
r.values[key] = value
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *contentModerationRuntimeSettingRepo) GetAll(_ context.Context) (map[string]string, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
out := make(map[string]string, len(r.values))
|
||||
for key, value := range r.values {
|
||||
out[key] = value
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (r *contentModerationRuntimeSettingRepo) Delete(_ context.Context, key string) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
delete(r.values, key)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *contentModerationRuntimeSettingRepo) calls() (getValue, getMultiple int) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
return r.getValueCalls, r.getMultipleCalls
|
||||
}
|
||||
|
||||
func (r *contentModerationRuntimeSettingRepo) failMultiple(err error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.getMultipleErr = err
|
||||
}
|
||||
|
||||
func (r *contentModerationRuntimeSettingRepo) blockNextMultiple(start chan<- struct{}, wait <-chan struct{}) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.getMultipleStart = start
|
||||
r.getMultipleWait = wait
|
||||
}
|
||||
|
||||
func runtimeCacheTestConfig(t *testing.T, keywords ...string) string {
|
||||
t.Helper()
|
||||
cfg := defaultContentModerationConfig()
|
||||
cfg.Enabled = true
|
||||
cfg.Mode = ContentModerationModePreBlock
|
||||
cfg.KeywordBlockingMode = ContentModerationKeywordModeKeywordOnly
|
||||
cfg.BlockedKeywords = keywords
|
||||
raw, err := json.Marshal(cfg)
|
||||
require.NoError(t, err)
|
||||
return string(raw)
|
||||
}
|
||||
|
||||
func runtimeCacheTestService(repo *contentModerationRuntimeSettingRepo, ttl time.Duration) *ContentModerationService {
|
||||
return &ContentModerationService{
|
||||
settingRepo: repo,
|
||||
repo: &contentModerationTestRepo{},
|
||||
runtimeCacheTTL: ttl,
|
||||
}
|
||||
}
|
||||
|
||||
func runtimeCacheTestInput(text string) ContentModerationCheckInput {
|
||||
return ContentModerationCheckInput{
|
||||
Protocol: ContentModerationProtocolOpenAIChat,
|
||||
Model: "risk-cache-test",
|
||||
Body: []byte(`{"messages":[{"role":"user","content":"` + text + `"}]}`),
|
||||
}
|
||||
}
|
||||
|
||||
func TestContentModerationRuntimeSnapshotCachesSettings(t *testing.T) {
|
||||
repo := &contentModerationRuntimeSettingRepo{values: map[string]string{
|
||||
SettingKeyRiskControlEnabled: "true",
|
||||
SettingKeyContentModerationConfig: runtimeCacheTestConfig(t, "blocked"),
|
||||
}}
|
||||
svc := runtimeCacheTestService(repo, time.Hour)
|
||||
|
||||
for range 20 {
|
||||
decision, err := svc.Check(context.Background(), runtimeCacheTestInput("clean prompt"))
|
||||
require.NoError(t, err)
|
||||
require.True(t, decision.Allowed)
|
||||
}
|
||||
|
||||
getValue, getMultiple := repo.calls()
|
||||
require.Zero(t, getValue)
|
||||
require.Equal(t, 1, getMultiple)
|
||||
}
|
||||
|
||||
func TestContentModerationRuntimeSnapshotUpdateConfigIsImmediate(t *testing.T) {
|
||||
repo := &contentModerationRuntimeSettingRepo{values: map[string]string{
|
||||
SettingKeyRiskControlEnabled: "true",
|
||||
SettingKeyContentModerationConfig: runtimeCacheTestConfig(t, "old-keyword"),
|
||||
}}
|
||||
svc := runtimeCacheTestService(repo, time.Hour)
|
||||
|
||||
decision, err := svc.Check(context.Background(), runtimeCacheTestInput("new-keyword"))
|
||||
require.NoError(t, err)
|
||||
require.True(t, decision.Allowed)
|
||||
|
||||
keywords := []string{"new-keyword"}
|
||||
_, err = svc.UpdateConfig(context.Background(), UpdateContentModerationConfigInput{
|
||||
BlockedKeywords: &keywords,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
decision, err = svc.Check(context.Background(), runtimeCacheTestInput("new-keyword"))
|
||||
require.NoError(t, err)
|
||||
require.True(t, decision.Blocked)
|
||||
|
||||
_, getMultiple := repo.calls()
|
||||
require.Equal(t, 1, getMultiple)
|
||||
}
|
||||
|
||||
func TestContentModerationRuntimeSnapshotUpdateWinsOverInitialLoad(t *testing.T) {
|
||||
repo := &contentModerationRuntimeSettingRepo{values: map[string]string{
|
||||
SettingKeyRiskControlEnabled: "true",
|
||||
SettingKeyContentModerationConfig: runtimeCacheTestConfig(t, "old-keyword"),
|
||||
}}
|
||||
svc := runtimeCacheTestService(repo, time.Hour)
|
||||
|
||||
refreshStarted := make(chan struct{}, 1)
|
||||
releaseRefresh := make(chan struct{})
|
||||
released := false
|
||||
defer func() {
|
||||
if !released {
|
||||
close(releaseRefresh)
|
||||
}
|
||||
}()
|
||||
repo.blockNextMultiple(refreshStarted, releaseRefresh)
|
||||
|
||||
initialCheckDone := make(chan error, 1)
|
||||
go func() {
|
||||
decision, err := svc.Check(context.Background(), runtimeCacheTestInput("clean prompt"))
|
||||
if err == nil && (decision == nil || !decision.Allowed) {
|
||||
err = errors.New("unexpected initial moderation decision")
|
||||
}
|
||||
initialCheckDone <- err
|
||||
}()
|
||||
require.Eventually(t, func() bool {
|
||||
select {
|
||||
case <-refreshStarted:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}, time.Second, time.Millisecond)
|
||||
|
||||
updateDone := make(chan error, 1)
|
||||
go func() {
|
||||
keywords := []string{"new-keyword"}
|
||||
_, updateErr := svc.UpdateConfig(context.Background(), UpdateContentModerationConfigInput{
|
||||
BlockedKeywords: &keywords,
|
||||
})
|
||||
updateDone <- updateErr
|
||||
}()
|
||||
select {
|
||||
case updateErr := <-updateDone:
|
||||
require.NoError(t, updateErr)
|
||||
t.Fatal("configuration update completed before the initial load released its lock")
|
||||
case <-time.After(10 * time.Millisecond):
|
||||
}
|
||||
|
||||
close(releaseRefresh)
|
||||
released = true
|
||||
require.NoError(t, <-initialCheckDone)
|
||||
require.NoError(t, <-updateDone)
|
||||
|
||||
decision, err := svc.Check(context.Background(), runtimeCacheTestInput("new-keyword"))
|
||||
require.NoError(t, err)
|
||||
require.True(t, decision.Blocked)
|
||||
decision, err = svc.Check(context.Background(), runtimeCacheTestInput("old-keyword"))
|
||||
require.NoError(t, err)
|
||||
require.True(t, decision.Allowed)
|
||||
}
|
||||
|
||||
func TestContentModerationRuntimeSnapshotRefreshFailureKeepsStaleConfig(t *testing.T) {
|
||||
repo := &contentModerationRuntimeSettingRepo{values: map[string]string{
|
||||
SettingKeyRiskControlEnabled: "true",
|
||||
SettingKeyContentModerationConfig: runtimeCacheTestConfig(t, "blocked"),
|
||||
}}
|
||||
svc := runtimeCacheTestService(repo, time.Nanosecond)
|
||||
input := runtimeCacheTestInput("blocked")
|
||||
|
||||
decision, err := svc.Check(context.Background(), input)
|
||||
require.NoError(t, err)
|
||||
require.True(t, decision.Blocked)
|
||||
|
||||
repo.failMultiple(errors.New("database unavailable"))
|
||||
decision, err = svc.Check(context.Background(), input)
|
||||
require.NoError(t, err)
|
||||
require.True(t, decision.Blocked)
|
||||
require.Eventually(t, func() bool {
|
||||
_, calls := repo.calls()
|
||||
return calls >= 2
|
||||
}, time.Second, time.Millisecond)
|
||||
}
|
||||
|
||||
func TestContentModerationRuntimeSnapshotRefreshFailureBacksOff(t *testing.T) {
|
||||
repo := &contentModerationRuntimeSettingRepo{values: map[string]string{
|
||||
SettingKeyRiskControlEnabled: "true",
|
||||
SettingKeyContentModerationConfig: runtimeCacheTestConfig(t, "blocked"),
|
||||
}}
|
||||
svc := runtimeCacheTestService(repo, time.Minute)
|
||||
input := runtimeCacheTestInput("blocked")
|
||||
|
||||
decision, err := svc.Check(context.Background(), input)
|
||||
require.NoError(t, err)
|
||||
require.True(t, decision.Blocked)
|
||||
|
||||
current := svc.runtimeSnapshot.Load()
|
||||
require.NotNil(t, current)
|
||||
expired := *current
|
||||
expired.loadedAt = time.Now().Add(-2 * time.Minute)
|
||||
svc.runtimeSnapshot.Store(&expired)
|
||||
repo.failMultiple(errors.New("database unavailable"))
|
||||
|
||||
decision, err = svc.Check(context.Background(), input)
|
||||
require.NoError(t, err)
|
||||
require.True(t, decision.Blocked)
|
||||
require.Eventually(t, func() bool {
|
||||
_, calls := repo.calls()
|
||||
return calls == 2
|
||||
}, time.Second, time.Millisecond)
|
||||
|
||||
for range 100 {
|
||||
decision, err = svc.Check(context.Background(), input)
|
||||
require.NoError(t, err)
|
||||
require.True(t, decision.Blocked)
|
||||
}
|
||||
_, calls := repo.calls()
|
||||
require.Equal(t, 2, calls)
|
||||
}
|
||||
|
||||
func TestContentModerationRuntimeSnapshotRefreshReusesUnchangedMatcher(t *testing.T) {
|
||||
repo := &contentModerationRuntimeSettingRepo{values: map[string]string{
|
||||
SettingKeyRiskControlEnabled: "true",
|
||||
SettingKeyContentModerationConfig: runtimeCacheTestConfig(t, "blocked"),
|
||||
}}
|
||||
svc := runtimeCacheTestService(repo, time.Minute)
|
||||
input := runtimeCacheTestInput("blocked")
|
||||
|
||||
decision, err := svc.Check(context.Background(), input)
|
||||
require.NoError(t, err)
|
||||
require.True(t, decision.Blocked)
|
||||
|
||||
current := svc.runtimeSnapshot.Load()
|
||||
require.NotNil(t, current)
|
||||
expired := *current
|
||||
expired.loadedAt = time.Now().Add(-2 * time.Minute)
|
||||
svc.runtimeSnapshot.Store(&expired)
|
||||
|
||||
decision, err = svc.Check(context.Background(), input)
|
||||
require.NoError(t, err)
|
||||
require.True(t, decision.Blocked)
|
||||
require.Eventually(t, func() bool {
|
||||
refreshed := svc.runtimeSnapshot.Load()
|
||||
return refreshed != nil && refreshed.loadedAt.After(expired.loadedAt)
|
||||
}, time.Second, time.Millisecond)
|
||||
|
||||
refreshed := svc.runtimeSnapshot.Load()
|
||||
require.Same(t, current.config, refreshed.config)
|
||||
require.Same(t, current.keywordMatcher, refreshed.keywordMatcher)
|
||||
_, calls := repo.calls()
|
||||
require.Equal(t, 2, calls)
|
||||
}
|
||||
|
||||
func TestContentModerationRuntimeSnapshotUpdateWinsOverInFlightRefresh(t *testing.T) {
|
||||
repo := &contentModerationRuntimeSettingRepo{values: map[string]string{
|
||||
SettingKeyRiskControlEnabled: "true",
|
||||
SettingKeyContentModerationConfig: runtimeCacheTestConfig(t, "old-keyword"),
|
||||
}}
|
||||
svc := runtimeCacheTestService(repo, time.Minute)
|
||||
|
||||
decision, err := svc.Check(context.Background(), runtimeCacheTestInput("old-keyword"))
|
||||
require.NoError(t, err)
|
||||
require.True(t, decision.Blocked)
|
||||
|
||||
current := svc.runtimeSnapshot.Load()
|
||||
require.NotNil(t, current)
|
||||
expired := *current
|
||||
expired.loadedAt = time.Now().Add(-2 * time.Minute)
|
||||
svc.runtimeSnapshot.Store(&expired)
|
||||
|
||||
refreshStarted := make(chan struct{}, 1)
|
||||
releaseRefresh := make(chan struct{})
|
||||
repo.blockNextMultiple(refreshStarted, releaseRefresh)
|
||||
decision, err = svc.Check(context.Background(), runtimeCacheTestInput("clean prompt"))
|
||||
require.NoError(t, err)
|
||||
require.True(t, decision.Allowed)
|
||||
require.Eventually(t, func() bool {
|
||||
select {
|
||||
case <-refreshStarted:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}, time.Second, time.Millisecond)
|
||||
|
||||
updateDone := make(chan error, 1)
|
||||
go func() {
|
||||
keywords := []string{"new-keyword"}
|
||||
_, updateErr := svc.UpdateConfig(context.Background(), UpdateContentModerationConfigInput{
|
||||
BlockedKeywords: &keywords,
|
||||
})
|
||||
updateDone <- updateErr
|
||||
}()
|
||||
select {
|
||||
case updateErr := <-updateDone:
|
||||
require.NoError(t, updateErr)
|
||||
t.Fatal("configuration update completed before the in-flight refresh released its lock")
|
||||
case <-time.After(10 * time.Millisecond):
|
||||
}
|
||||
|
||||
close(releaseRefresh)
|
||||
require.NoError(t, <-updateDone)
|
||||
decision, err = svc.Check(context.Background(), runtimeCacheTestInput("new-keyword"))
|
||||
require.NoError(t, err)
|
||||
require.True(t, decision.Blocked)
|
||||
decision, err = svc.Check(context.Background(), runtimeCacheTestInput("old-keyword"))
|
||||
require.NoError(t, err)
|
||||
require.True(t, decision.Allowed)
|
||||
}
|
||||
|
||||
func TestContentModerationRuntimeSnapshotConcurrentReadAndReplace(t *testing.T) {
|
||||
repo := &contentModerationRuntimeSettingRepo{values: map[string]string{
|
||||
SettingKeyRiskControlEnabled: "true",
|
||||
SettingKeyContentModerationConfig: runtimeCacheTestConfig(t, "blocked-0"),
|
||||
}}
|
||||
svc := runtimeCacheTestService(repo, time.Hour)
|
||||
_, err := svc.Check(context.Background(), runtimeCacheTestInput("clean prompt"))
|
||||
require.NoError(t, err)
|
||||
|
||||
var wg sync.WaitGroup
|
||||
errs := make(chan error, 8)
|
||||
for range 8 {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for range 100 {
|
||||
decision, checkErr := svc.Check(context.Background(), runtimeCacheTestInput("clean prompt"))
|
||||
if checkErr != nil {
|
||||
errs <- checkErr
|
||||
return
|
||||
}
|
||||
if decision == nil || !decision.Allowed {
|
||||
errs <- errors.New("unexpected moderation decision")
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
for i := 1; i <= 20; i++ {
|
||||
keywords := []string{"blocked-" + time.Duration(i).String()}
|
||||
_, err := svc.UpdateConfig(context.Background(), UpdateContentModerationConfigInput{
|
||||
BlockedKeywords: &keywords,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
}
|
||||
wg.Wait()
|
||||
close(errs)
|
||||
for err := range errs {
|
||||
require.NoError(t, err)
|
||||
}
|
||||
}
|
||||
@@ -558,19 +558,79 @@ type ForwardResult struct {
|
||||
ImageSizeBreakdown map[string]int
|
||||
}
|
||||
|
||||
// UpstreamFailoverError indicates an upstream error that should trigger account failover.
|
||||
// GatewayFailureStage identifies which request stage failed. The zero value is
|
||||
// intentionally treated as inference so existing UpstreamFailoverError callers
|
||||
// retain their current behavior.
|
||||
type GatewayFailureStage string
|
||||
|
||||
const (
|
||||
GatewayFailureStageInference GatewayFailureStage = "inference"
|
||||
GatewayFailureStageAccountAuth GatewayFailureStage = "account_auth"
|
||||
)
|
||||
|
||||
// GatewayFailureScope identifies whether selecting another account can help.
|
||||
type GatewayFailureScope string
|
||||
|
||||
const (
|
||||
GatewayFailureScopeAccount GatewayFailureScope = "account"
|
||||
GatewayFailureScopeProvider GatewayFailureScope = "provider"
|
||||
GatewayFailureScopeRequest GatewayFailureScope = "request"
|
||||
)
|
||||
|
||||
// NextAccountAction is tri-state for backwards compatibility. The zero value
|
||||
// means legacy retry behavior; only NextAccountStop explicitly short-circuits.
|
||||
type NextAccountAction uint8
|
||||
|
||||
const (
|
||||
NextAccountLegacyRetry NextAccountAction = iota
|
||||
NextAccountRetry
|
||||
NextAccountStop
|
||||
)
|
||||
|
||||
type GatewayFailureReason string
|
||||
|
||||
// UpstreamFailoverError indicates an upstream or credential error that may
|
||||
// trigger account failover. Additive metadata keeps existing composite literals
|
||||
// source-compatible and preserves their legacy retry-next-account behavior.
|
||||
type UpstreamFailoverError struct {
|
||||
StatusCode int
|
||||
ResponseBody []byte // 上游响应体,用于错误透传规则匹配
|
||||
ResponseHeaders http.Header // 上游响应头,用于透传 cf-ray/cf-mitigated/content-type 等诊断信息
|
||||
ForceCacheBilling bool // Antigravity 粘性会话切换时设为 true
|
||||
RetryableOnSameAccount bool // 临时性错误(如 Google 间歇性 400、空响应),应在同一账号上重试 N 次再切换
|
||||
Stage GatewayFailureStage
|
||||
Scope GatewayFailureScope
|
||||
Reason GatewayFailureReason
|
||||
NextAccountAction NextAccountAction
|
||||
ClientStatusCode int
|
||||
ClientMessage string
|
||||
}
|
||||
|
||||
func (e *UpstreamFailoverError) Error() string {
|
||||
if e != nil && e.Stage == GatewayFailureStageAccountAuth {
|
||||
return fmt.Sprintf("credential failure: %s (failover)", e.Reason)
|
||||
}
|
||||
return fmt.Sprintf("upstream error: %d (failover)", e.StatusCode)
|
||||
}
|
||||
|
||||
func (e *UpstreamFailoverError) ShouldRetryNextAccount() bool {
|
||||
return e != nil && e.NextAccountAction != NextAccountStop
|
||||
}
|
||||
|
||||
func (e *UpstreamFailoverError) IsCredentialFailure() bool {
|
||||
return e != nil && e.Stage == GatewayFailureStageAccountAuth
|
||||
}
|
||||
|
||||
// ShouldReportAccountScheduleFailure prevents provider- and request-scoped
|
||||
// credential failures from being misattributed to the selected account. Legacy
|
||||
// and inference failures retain their existing scheduler-health behavior.
|
||||
func (e *UpstreamFailoverError) ShouldReportAccountScheduleFailure() bool {
|
||||
if e == nil {
|
||||
return false
|
||||
}
|
||||
return !e.IsCredentialFailure() || e.Scope == GatewayFailureScopeAccount
|
||||
}
|
||||
|
||||
// sseStreamErrorEventError 表示上游 SSE 流体内出现 event:error 帧。
|
||||
// RawData 是该事件 data: 行的原始 JSON 字符串
|
||||
// (Anthropic 标准结构 {"type":"error","error":{"type":"...","message":"..."}})。
|
||||
|
||||
@@ -0,0 +1,655 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
const (
|
||||
grokCredentialFailoverDeadlineKey = "grok_credential_failover_deadline"
|
||||
grokCredentialFailoverBudget = 15 * time.Second
|
||||
grokCredentialMutationTimeout = 5 * time.Second
|
||||
grokCredentialMutationConfirmWait = 250 * time.Millisecond
|
||||
grokCredentialCacheCleanupTimeout = 500 * time.Millisecond
|
||||
|
||||
GrokCredentialUnavailableClientMessage = "No healthy Grok OAuth account is currently available"
|
||||
|
||||
GrokCredentialReasonRevoked GatewayFailureReason = "grok_oauth_credential_revoked"
|
||||
GrokCredentialReasonMissing GatewayFailureReason = "grok_oauth_credentials_missing"
|
||||
GrokCredentialReasonEntitlement GatewayFailureReason = "grok_oauth_entitlement_action_required"
|
||||
GrokCredentialReasonProxyInvalid GatewayFailureReason = "grok_oauth_proxy_invalid"
|
||||
GrokCredentialReasonRefreshTransient GatewayFailureReason = "grok_oauth_refresh_transient"
|
||||
GrokCredentialReasonProviderConfig GatewayFailureReason = "grok_oauth_provider_config"
|
||||
GrokCredentialReasonProviderDown GatewayFailureReason = "grok_oauth_provider_unavailable"
|
||||
GrokCredentialReasonAccountChanged GatewayFailureReason = "grok_oauth_account_state_changed"
|
||||
GrokCredentialReasonStateUpdate GatewayFailureReason = "grok_oauth_account_state_update_failed"
|
||||
GrokCredentialReasonFailoverTimeout GatewayFailureReason = "grok_oauth_failover_timeout"
|
||||
)
|
||||
|
||||
var errGrokCredentialStateUpdateFailed = errors.New("grok oauth account state update failed")
|
||||
|
||||
type grokCredentialFailureClass struct {
|
||||
scope GatewayFailureScope
|
||||
reason GatewayFailureReason
|
||||
action NextAccountAction
|
||||
permanent bool
|
||||
transient bool
|
||||
message string
|
||||
snapshot *GrokCredentialMutationSnapshot
|
||||
}
|
||||
|
||||
// GrokCredentialMutationSnapshot is the credential identity observed when the
|
||||
// request selected an account. Repository mutations compare all fields before
|
||||
// quarantining that account so a concurrent refresh cannot be overwritten.
|
||||
type GrokCredentialMutationSnapshot struct {
|
||||
CredentialsJSON string
|
||||
AccessToken string
|
||||
RefreshToken string
|
||||
TokenVersion int64
|
||||
ProxyID *int64
|
||||
}
|
||||
|
||||
type grokCredentialFailureSnapshotError struct {
|
||||
cause error
|
||||
snapshot GrokCredentialMutationSnapshot
|
||||
}
|
||||
|
||||
func (e *grokCredentialFailureSnapshotError) Error() string { return e.cause.Error() }
|
||||
func (e *grokCredentialFailureSnapshotError) Unwrap() error { return e.cause }
|
||||
|
||||
func withGrokCredentialFailureSnapshot(err error, account *Account) error {
|
||||
if err == nil || account == nil || !account.IsGrokOAuth() {
|
||||
return err
|
||||
}
|
||||
var existing *grokCredentialFailureSnapshotError
|
||||
if errors.As(err, &existing) {
|
||||
return err
|
||||
}
|
||||
return &grokCredentialFailureSnapshotError{cause: err, snapshot: grokCredentialMutationSnapshot(account)}
|
||||
}
|
||||
|
||||
func grokCredentialFailureSnapshot(err error) (GrokCredentialMutationSnapshot, bool) {
|
||||
var snapshotErr *grokCredentialFailureSnapshotError
|
||||
if !errors.As(err, &snapshotErr) || snapshotErr == nil {
|
||||
return GrokCredentialMutationSnapshot{}, false
|
||||
}
|
||||
return snapshotErr.snapshot, true
|
||||
}
|
||||
|
||||
type grokCredentialConditionalStateRepository interface {
|
||||
SetGrokCredentialErrorIfMatch(context.Context, int64, GrokCredentialMutationSnapshot, string) (bool, error)
|
||||
SetGrokCredentialTempUnschedulableIfMatch(context.Context, int64, GrokCredentialMutationSnapshot, time.Time, string) (bool, error)
|
||||
}
|
||||
|
||||
// GetRequestCredential applies the request-path credential and failover contract
|
||||
// before any upstream transport is opened.
|
||||
func (s *OpenAIGatewayService) GetRequestCredential(ctx context.Context, c *gin.Context, account *Account) (string, string, error) {
|
||||
return s.getRequestCredential(ctx, c, account)
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) getRequestCredential(ctx context.Context, c *gin.Context, account *Account) (string, string, error) {
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
if account == nil {
|
||||
return "", "", errors.New("account is nil")
|
||||
}
|
||||
if !account.IsGrokOAuth() {
|
||||
return s.GetAccessToken(ctx, account)
|
||||
}
|
||||
if err := ctx.Err(); err != nil {
|
||||
return "", "", err
|
||||
}
|
||||
if s == nil || s.grokTokenProvider == nil {
|
||||
return "", "", s.newGrokCredentialFailover(c, account, grokCredentialFailureClass{
|
||||
scope: GatewayFailureScopeProvider,
|
||||
reason: GrokCredentialReasonProviderConfig,
|
||||
action: NextAccountStop,
|
||||
message: "Grok OAuth credential provider is unavailable",
|
||||
})
|
||||
}
|
||||
if s.isOpenAIAccountRuntimeBlocked(account) {
|
||||
return "", "", s.newGrokCredentialFailover(c, account, grokCredentialFailureClass{
|
||||
scope: GatewayFailureScopeAccount,
|
||||
reason: GrokCredentialReasonAccountChanged,
|
||||
action: NextAccountRetry,
|
||||
message: "Grok OAuth account is not currently schedulable",
|
||||
})
|
||||
}
|
||||
|
||||
credentialCtx, cancel, budgetExpired := grokCredentialAcquisitionContext(ctx, c)
|
||||
if cancel != nil {
|
||||
defer cancel()
|
||||
}
|
||||
if budgetExpired {
|
||||
return "", "", s.newGrokCredentialFailover(c, account, grokCredentialFailureClass{
|
||||
scope: GatewayFailureScopeRequest,
|
||||
reason: GrokCredentialReasonFailoverTimeout,
|
||||
action: NextAccountStop,
|
||||
message: "Grok OAuth credential failover budget exhausted",
|
||||
})
|
||||
}
|
||||
|
||||
token, kind, err := s.GetAccessToken(credentialCtx, account)
|
||||
if err == nil {
|
||||
if s.isOpenAIAccountRuntimeBlocked(account) {
|
||||
return "", "", s.newGrokCredentialFailover(c, account, grokCredentialFailureClass{
|
||||
scope: GatewayFailureScopeAccount,
|
||||
reason: GrokCredentialReasonAccountChanged,
|
||||
action: NextAccountRetry,
|
||||
message: "Grok OAuth account is not currently schedulable",
|
||||
})
|
||||
}
|
||||
return token, kind, nil
|
||||
}
|
||||
if parentErr := ctx.Err(); parentErr != nil {
|
||||
return "", "", parentErr
|
||||
}
|
||||
if credentialCtx.Err() != nil {
|
||||
return "", "", s.newGrokCredentialFailover(c, account, grokCredentialFailureClass{
|
||||
scope: GatewayFailureScopeRequest,
|
||||
reason: GrokCredentialReasonFailoverTimeout,
|
||||
action: NextAccountStop,
|
||||
message: "Grok OAuth credential failover budget exhausted",
|
||||
})
|
||||
}
|
||||
|
||||
class := classifyGrokCredentialFailure(account, err)
|
||||
if snapshot, ok := grokCredentialFailureSnapshot(err); ok {
|
||||
class.snapshot = &snapshot
|
||||
}
|
||||
if ctx.Err() != nil {
|
||||
return "", "", ctx.Err()
|
||||
}
|
||||
if class.permanent || class.transient {
|
||||
freshToken, mutationErr := s.applyGrokCredentialAccountFailure(credentialCtx, account, class)
|
||||
if freshToken != "" {
|
||||
return freshToken, "oauth", nil
|
||||
}
|
||||
if mutationErr != nil {
|
||||
if ctx.Err() != nil {
|
||||
return "", "", ctx.Err()
|
||||
}
|
||||
if credentialCtx.Err() != nil {
|
||||
return "", "", s.newGrokCredentialFailover(c, account, grokCredentialFailureClass{
|
||||
scope: GatewayFailureScopeRequest,
|
||||
reason: GrokCredentialReasonFailoverTimeout,
|
||||
action: NextAccountStop,
|
||||
message: "Grok OAuth credential failover budget exhausted",
|
||||
})
|
||||
}
|
||||
if errors.Is(mutationErr, errOAuthRefreshAccountStateChanged) {
|
||||
class = grokCredentialFailureClass{
|
||||
scope: GatewayFailureScopeAccount,
|
||||
reason: GrokCredentialReasonAccountChanged,
|
||||
action: NextAccountRetry,
|
||||
message: "Grok OAuth account eligibility changed",
|
||||
}
|
||||
} else if errors.Is(mutationErr, errOAuthRefreshAccountRereadFailed) {
|
||||
class = grokCredentialFailureClass{
|
||||
scope: GatewayFailureScopeProvider,
|
||||
reason: GrokCredentialReasonProviderDown,
|
||||
action: NextAccountStop,
|
||||
message: "Grok OAuth account state is temporarily unavailable",
|
||||
}
|
||||
} else {
|
||||
class = grokCredentialFailureClass{
|
||||
scope: GatewayFailureScopeProvider,
|
||||
reason: GrokCredentialReasonStateUpdate,
|
||||
action: NextAccountStop,
|
||||
message: "Grok OAuth account state could not be updated safely",
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return "", "", s.newGrokCredentialFailover(c, account, class)
|
||||
}
|
||||
|
||||
func grokCredentialAcquisitionContext(ctx context.Context, c *gin.Context) (context.Context, context.CancelFunc, bool) {
|
||||
if c == nil {
|
||||
return ctx, nil, false
|
||||
}
|
||||
deadline := time.Time{}
|
||||
if raw, ok := c.Get(grokCredentialFailoverDeadlineKey); ok {
|
||||
deadline, _ = raw.(time.Time)
|
||||
}
|
||||
if deadline.IsZero() {
|
||||
deadline = time.Now().Add(grokCredentialFailoverBudget)
|
||||
c.Set(grokCredentialFailoverDeadlineKey, deadline)
|
||||
}
|
||||
if !time.Now().Before(deadline) {
|
||||
return ctx, nil, true
|
||||
}
|
||||
acquireCtx, cancel := context.WithDeadline(ctx, deadline)
|
||||
return acquireCtx, cancel, false
|
||||
}
|
||||
|
||||
func classifyGrokCredentialFailure(account *Account, err error) grokCredentialFailureClass {
|
||||
stableReason := strings.ToLower(strings.TrimSpace(infraerrors.Reason(err)))
|
||||
message := ""
|
||||
if err != nil {
|
||||
message = strings.ToLower(err.Error())
|
||||
}
|
||||
contains := func(values ...string) bool {
|
||||
for _, value := range values {
|
||||
if strings.Contains(stableReason, value) || strings.Contains(message, value) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
switch {
|
||||
case errors.Is(err, errGrokOAuthRefreshTokenMissing), errors.Is(err, errGrokOAuthAccessTokenMissing), errors.Is(err, errGrokOAuthAccessTokenExpired):
|
||||
return grokCredentialFailureClass{scope: GatewayFailureScopeAccount, reason: GrokCredentialReasonMissing, action: NextAccountRetry, permanent: true, message: "Grok OAuth credentials are missing or expired"}
|
||||
case contains("invalid_grant", "invalid_refresh_token", "token_expired", "refresh_token_reused", "refresh_token_invalidated", "app_session_terminated"):
|
||||
return grokCredentialFailureClass{scope: GatewayFailureScopeAccount, reason: GrokCredentialReasonRevoked, action: NextAccountRetry, permanent: true, message: "Grok OAuth credentials require account action"}
|
||||
case contains("grok_oauth_entitlement_denied", "entitlement_denied", "access_denied", "subscription required", "no active grok subscription"):
|
||||
return grokCredentialFailureClass{scope: GatewayFailureScopeAccount, reason: GrokCredentialReasonEntitlement, action: NextAccountRetry, permanent: true, message: "Grok OAuth entitlement requires account action"}
|
||||
case errors.Is(err, errGrokOAuthConfiguredProxyMiss), contains("grok_oauth_proxy_not_found"):
|
||||
return grokCredentialFailureClass{scope: GatewayFailureScopeAccount, reason: GrokCredentialReasonProxyInvalid, action: NextAccountRetry, permanent: true, message: "Grok OAuth account proxy configuration is invalid"}
|
||||
case errors.Is(err, errOAuthRefreshAccountRereadFailed):
|
||||
return grokCredentialFailureClass{scope: GatewayFailureScopeProvider, reason: GrokCredentialReasonProviderDown, action: NextAccountStop, message: "Grok OAuth account state is temporarily unavailable"}
|
||||
case errors.Is(err, errOAuthRefreshCredentialPersist):
|
||||
return grokCredentialFailureClass{scope: GatewayFailureScopeProvider, reason: GrokCredentialReasonProviderDown, action: NextAccountStop, message: "Grok OAuth shared credential state is temporarily unavailable"}
|
||||
case errors.Is(err, errOAuthRefreshAccountStateChanged):
|
||||
return grokCredentialFailureClass{scope: GatewayFailureScopeAccount, reason: GrokCredentialReasonAccountChanged, action: NextAccountRetry, message: "Grok OAuth account eligibility changed"}
|
||||
case errors.Is(err, errGrokOAuthRefreshNotConfigured), contains("invalid_client", "unauthorized_client", "invalid_scope", "unknown scope", "grok oauth service is not configured", "grok_oauth_proxy_not_available"):
|
||||
return grokCredentialFailureClass{scope: GatewayFailureScopeProvider, reason: GrokCredentialReasonProviderConfig, action: NextAccountStop, message: "Grok OAuth provider configuration is unavailable"}
|
||||
case contains("grok_oauth_proxy_lookup_failed"),
|
||||
contains("grok_oauth_token_refresh_failed") && contains("status 403") && (account == nil || account.ProxyID == nil):
|
||||
return grokCredentialFailureClass{scope: GatewayFailureScopeProvider, reason: GrokCredentialReasonProviderDown, action: NextAccountStop, message: "Grok OAuth provider is temporarily unavailable"}
|
||||
case contains("grok_oauth_client_init_failed") && (account == nil || account.ProxyID == nil):
|
||||
return grokCredentialFailureClass{scope: GatewayFailureScopeProvider, reason: GrokCredentialReasonProviderConfig, action: NextAccountStop, message: "Grok OAuth provider configuration is unavailable"}
|
||||
case contains("grok_oauth_request_failed") && (account == nil || account.ProxyID == nil):
|
||||
return grokCredentialFailureClass{scope: GatewayFailureScopeProvider, reason: GrokCredentialReasonProviderDown, action: NextAccountStop, message: "Grok OAuth provider is temporarily unavailable"}
|
||||
case contains("status 429", "status 500", "status 502", "status 503", "status 504") && (account == nil || account.ProxyID == nil):
|
||||
return grokCredentialFailureClass{scope: GatewayFailureScopeProvider, reason: GrokCredentialReasonProviderDown, action: NextAccountStop, message: "Grok OAuth provider is temporarily unavailable"}
|
||||
default:
|
||||
return grokCredentialFailureClass{scope: GatewayFailureScopeAccount, reason: GrokCredentialReasonRefreshTransient, action: NextAccountRetry, transient: true, message: "Grok OAuth credential refresh is temporarily unavailable"}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) applyGrokCredentialAccountFailure(ctx context.Context, account *Account, class grokCredentialFailureClass) (string, error) {
|
||||
if s == nil || account == nil || ctx == nil || ctx.Err() != nil {
|
||||
if ctx != nil {
|
||||
return "", ctx.Err()
|
||||
}
|
||||
return "", context.Canceled
|
||||
}
|
||||
mutationMu := s.grokCredentialMutationLock(account.ID)
|
||||
if err := mutationMu.Lock(ctx); err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer mutationMu.Unlock()
|
||||
stateRepo, hasConditionalStateRepo := s.accountRepo.(grokCredentialConditionalStateRepository)
|
||||
snapshot := grokCredentialMutationSnapshot(account)
|
||||
if class.snapshot != nil {
|
||||
snapshot = *class.snapshot
|
||||
}
|
||||
if token, err := s.validateCurrentGrokCredentialFailure(ctx, account.ID, snapshot, class); err != nil || token != "" {
|
||||
return token, err
|
||||
}
|
||||
|
||||
if class.permanent {
|
||||
if token, ok := s.grokCredentialConcurrentlyRefreshedToken(ctx, account.ID, snapshot); ok {
|
||||
return token, nil
|
||||
}
|
||||
if ctx.Err() != nil {
|
||||
return "", ctx.Err()
|
||||
}
|
||||
rollbackRuntime := s.blockGrokCredentialRuntime(account, time.Time{}, string(class.reason))
|
||||
keepRuntimeBlock := false
|
||||
runtimeRollbackDone := false
|
||||
defer func() {
|
||||
if !keepRuntimeBlock && !runtimeRollbackDone {
|
||||
rollbackRuntime()
|
||||
}
|
||||
}()
|
||||
if s.accountRepo == nil {
|
||||
keepRuntimeBlock = true
|
||||
return "", fmt.Errorf("%w: account repository is not configured", errGrokCredentialStateUpdateFailed)
|
||||
}
|
||||
if !hasConditionalStateRepo {
|
||||
keepRuntimeBlock = true
|
||||
return "", fmt.Errorf("%w: conditional account repository is not configured", errGrokCredentialStateUpdateFailed)
|
||||
}
|
||||
stateCtx, cancel := context.WithTimeout(ctx, grokCredentialMutationTimeout)
|
||||
if err := ctx.Err(); err != nil {
|
||||
cancel()
|
||||
return "", err
|
||||
}
|
||||
updated, err := stateRepo.SetGrokCredentialErrorIfMatch(stateCtx, account.ID, snapshot, string(class.reason))
|
||||
requestErr := ctx.Err()
|
||||
cancel()
|
||||
if err != nil {
|
||||
slog.Warn("grok_credential_failure.set_error_failed", "account_id", account.ID, "reason", class.reason, "error", err)
|
||||
if requestErr != nil || errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
|
||||
if s.grokCredentialMutationCommitted(account.ID, class, time.Time{}) {
|
||||
updated = true
|
||||
} else if requestErr != nil {
|
||||
return "", requestErr
|
||||
} else {
|
||||
keepRuntimeBlock = true
|
||||
return "", fmt.Errorf("%w: permanent state commit could not be confirmed: %v", errGrokCredentialStateUpdateFailed, err)
|
||||
}
|
||||
} else {
|
||||
keepRuntimeBlock = true
|
||||
return "", fmt.Errorf("%w: persist permanent state: %v", errGrokCredentialStateUpdateFailed, err)
|
||||
}
|
||||
}
|
||||
if !updated {
|
||||
rollbackRuntime()
|
||||
runtimeRollbackDone = true
|
||||
return s.resolveGrokCredentialCASMiss(ctx, account.ID, snapshot)
|
||||
}
|
||||
// SetError is the linearization point: the durable quarantine is now
|
||||
// committed and this node's runtime block must not be rolled back.
|
||||
keepRuntimeBlock = true
|
||||
if s.grokTokenProvider == nil {
|
||||
if ctx.Err() != nil {
|
||||
return "", ctx.Err()
|
||||
}
|
||||
return "", fmt.Errorf("%w: token provider is not configured", errGrokCredentialStateUpdateFailed)
|
||||
}
|
||||
invalidateCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), grokCredentialCacheCleanupTimeout)
|
||||
err = s.grokTokenProvider.InvalidateToken(invalidateCtx, account)
|
||||
cancel()
|
||||
if err != nil {
|
||||
slog.Warn("grok_credential_failure.invalidate_token_failed", "account_id", account.ID, "reason", class.reason, "error", err)
|
||||
if ctx.Err() != nil {
|
||||
return "", ctx.Err()
|
||||
}
|
||||
return "", fmt.Errorf("%w: invalidate cached credential: %v", errGrokCredentialStateUpdateFailed, err)
|
||||
}
|
||||
if ctx.Err() != nil {
|
||||
return "", ctx.Err()
|
||||
}
|
||||
return "", nil
|
||||
}
|
||||
|
||||
if class.transient {
|
||||
until := time.Now().Add(tokenRefreshTempUnschedDuration)
|
||||
if ctx.Err() != nil {
|
||||
return "", ctx.Err()
|
||||
}
|
||||
rollbackRuntime := s.blockGrokCredentialRuntime(account, until, string(class.reason))
|
||||
keepRuntimeBlock := false
|
||||
runtimeRollbackDone := false
|
||||
defer func() {
|
||||
if !keepRuntimeBlock && !runtimeRollbackDone {
|
||||
rollbackRuntime()
|
||||
}
|
||||
}()
|
||||
stateCtx, cancel := context.WithTimeout(ctx, grokCredentialMutationTimeout)
|
||||
if s.accountRepo == nil {
|
||||
cancel()
|
||||
keepRuntimeBlock = true
|
||||
return "", fmt.Errorf("%w: account repository is not configured", errGrokCredentialStateUpdateFailed)
|
||||
}
|
||||
if !hasConditionalStateRepo {
|
||||
cancel()
|
||||
keepRuntimeBlock = true
|
||||
return "", fmt.Errorf("%w: conditional account repository is not configured", errGrokCredentialStateUpdateFailed)
|
||||
}
|
||||
if err := ctx.Err(); err != nil {
|
||||
cancel()
|
||||
return "", err
|
||||
}
|
||||
updated, err := stateRepo.SetGrokCredentialTempUnschedulableIfMatch(stateCtx, account.ID, snapshot, until, string(class.reason))
|
||||
requestErr := ctx.Err()
|
||||
cancel()
|
||||
if err != nil {
|
||||
slog.Warn("grok_credential_failure.set_temp_unschedulable_failed", "account_id", account.ID, "reason", class.reason, "error", err)
|
||||
if requestErr != nil || errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
|
||||
if s.grokCredentialMutationCommitted(account.ID, class, until) {
|
||||
updated = true
|
||||
} else if requestErr != nil {
|
||||
return "", requestErr
|
||||
} else {
|
||||
keepRuntimeBlock = true
|
||||
return "", fmt.Errorf("%w: transient state commit could not be confirmed: %v", errGrokCredentialStateUpdateFailed, err)
|
||||
}
|
||||
} else {
|
||||
keepRuntimeBlock = true
|
||||
return "", fmt.Errorf("%w: persist transient state: %v", errGrokCredentialStateUpdateFailed, err)
|
||||
}
|
||||
}
|
||||
if !updated {
|
||||
rollbackRuntime()
|
||||
runtimeRollbackDone = true
|
||||
return s.resolveGrokCredentialCASMiss(ctx, account.ID, snapshot)
|
||||
}
|
||||
// The temporary quarantine is durable after SetTempUnschedulable succeeds.
|
||||
keepRuntimeBlock = true
|
||||
if ctx.Err() != nil {
|
||||
return "", ctx.Err()
|
||||
}
|
||||
return "", nil
|
||||
}
|
||||
|
||||
return "", nil
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) validateCurrentGrokCredentialFailure(
|
||||
ctx context.Context,
|
||||
accountID int64,
|
||||
snapshot GrokCredentialMutationSnapshot,
|
||||
class grokCredentialFailureClass,
|
||||
) (string, error) {
|
||||
if s == nil || s.accountRepo == nil || accountID <= 0 || ctx == nil || ctx.Err() != nil {
|
||||
if ctx != nil {
|
||||
return "", ctx.Err()
|
||||
}
|
||||
return "", context.Canceled
|
||||
}
|
||||
checkCtx, cancel := context.WithTimeout(ctx, time.Second)
|
||||
defer cancel()
|
||||
latest, err := s.accountRepo.GetByID(checkCtx, accountID)
|
||||
if err != nil {
|
||||
if errors.Is(err, ErrAccountNotFound) {
|
||||
return "", errOAuthRefreshAccountStateChanged
|
||||
}
|
||||
return "", fmt.Errorf("%w: %v", errOAuthRefreshAccountRereadFailed, err)
|
||||
}
|
||||
if latest == nil || !latest.IsGrokOAuth() || !latest.IsSchedulable() || s.isOpenAIAccountRuntimeBlocked(latest) {
|
||||
return "", errOAuthRefreshAccountStateChanged
|
||||
}
|
||||
latestSnapshot := grokCredentialMutationSnapshot(latest)
|
||||
if latestSnapshot.CredentialsJSON != snapshot.CredentialsJSON ||
|
||||
!grokCredentialProxyIDsEqual(latestSnapshot.ProxyID, snapshot.ProxyID) {
|
||||
if token, ok := s.grokCredentialConcurrentlyRefreshedToken(ctx, accountID, snapshot); ok {
|
||||
return token, nil
|
||||
}
|
||||
return "", errOAuthRefreshAccountStateChanged
|
||||
}
|
||||
|
||||
// A configured proxy is external to the account-row CAS identity. Recheck
|
||||
// the hydrated proxy object so restoring a deleted row under the same ID
|
||||
// wins over a stale proxy-invalid failure.
|
||||
if class.reason == GrokCredentialReasonProxyInvalid {
|
||||
if latest.ProxyID == nil || latest.Proxy != nil {
|
||||
return "", errOAuthRefreshAccountStateChanged
|
||||
}
|
||||
} else if latest.ProxyID != nil && latest.Proxy == nil {
|
||||
return "", errOAuthRefreshAccountStateChanged
|
||||
}
|
||||
|
||||
if class.reason == GrokCredentialReasonMissing {
|
||||
expiresAt := latest.GetCredentialAsTime("expires_at")
|
||||
credentialsStillMissing := strings.TrimSpace(latest.GetGrokAccessToken()) == "" ||
|
||||
strings.TrimSpace(latest.GetGrokRefreshToken()) == "" || expiresAt == nil || !time.Now().Before(*expiresAt)
|
||||
if !credentialsStillMissing {
|
||||
return "", errOAuthRefreshAccountStateChanged
|
||||
}
|
||||
}
|
||||
return "", nil
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) grokCredentialMutationLock(accountID int64) *oauthRefreshLocalLock {
|
||||
actual, _ := s.grokCredentialMutationLocks.LoadOrStore(accountID, newOAuthRefreshLocalLock())
|
||||
mu, ok := actual.(*oauthRefreshLocalLock)
|
||||
if !ok {
|
||||
mu = newOAuthRefreshLocalLock()
|
||||
s.grokCredentialMutationLocks.Store(accountID, mu)
|
||||
}
|
||||
return mu
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) grokCredentialMutationCommitted(accountID int64, class grokCredentialFailureClass, until time.Time) bool {
|
||||
if s == nil || s.accountRepo == nil || accountID <= 0 {
|
||||
return false
|
||||
}
|
||||
confirmCtx, cancel := context.WithTimeout(context.Background(), grokCredentialMutationConfirmWait)
|
||||
defer cancel()
|
||||
latest, err := s.accountRepo.GetByID(confirmCtx, accountID)
|
||||
if err != nil || latest == nil {
|
||||
return false
|
||||
}
|
||||
if class.permanent {
|
||||
return latest.Status == StatusError && !latest.Schedulable && latest.ErrorMessage == string(class.reason)
|
||||
}
|
||||
if class.transient {
|
||||
return latest.TempUnschedulableUntil != nil && !latest.TempUnschedulableUntil.Before(until) &&
|
||||
latest.TempUnschedulableReason == string(class.reason)
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func grokCredentialMutationSnapshot(account *Account) GrokCredentialMutationSnapshot {
|
||||
if account == nil {
|
||||
return GrokCredentialMutationSnapshot{}
|
||||
}
|
||||
credentialsJSON := "null"
|
||||
if encoded, err := json.Marshal(account.Credentials); err == nil {
|
||||
credentialsJSON = string(encoded)
|
||||
}
|
||||
snapshot := GrokCredentialMutationSnapshot{
|
||||
CredentialsJSON: credentialsJSON,
|
||||
AccessToken: strings.TrimSpace(account.GetGrokAccessToken()),
|
||||
RefreshToken: strings.TrimSpace(account.GetGrokRefreshToken()),
|
||||
TokenVersion: account.GetCredentialAsInt64("_token_version"),
|
||||
}
|
||||
if account.ProxyID != nil {
|
||||
proxyID := *account.ProxyID
|
||||
snapshot.ProxyID = &proxyID
|
||||
}
|
||||
return snapshot
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) resolveGrokCredentialCASMiss(ctx context.Context, accountID int64, snapshot GrokCredentialMutationSnapshot) (string, error) {
|
||||
if ctx.Err() != nil {
|
||||
return "", ctx.Err()
|
||||
}
|
||||
if token, ok := s.grokCredentialConcurrentlyRefreshedToken(ctx, accountID, snapshot); ok {
|
||||
return token, nil
|
||||
}
|
||||
return "", errOAuthRefreshAccountStateChanged
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) blockGrokCredentialRuntime(account *Account, until time.Time, reason string) func() {
|
||||
if s == nil || account == nil {
|
||||
return func() {}
|
||||
}
|
||||
mu := s.openAIAccountRuntimeBlockLock(account.ID)
|
||||
mu.Lock()
|
||||
before, hadBefore := s.openaiAccountRuntimeBlockUntil.Load(account.ID)
|
||||
installedGeneration, changed := s.blockAccountSchedulingLocked(account, until, reason)
|
||||
installed, installedOK := s.openaiAccountRuntimeBlockUntil.Load(account.ID)
|
||||
installedUntil, isTime := installed.(time.Time)
|
||||
mu.Unlock()
|
||||
if !changed || !installedOK || !isTime {
|
||||
return func() {}
|
||||
}
|
||||
if hadBefore {
|
||||
if beforeUntil, ok := before.(time.Time); ok && beforeUntil.Equal(installedUntil) {
|
||||
return func() {}
|
||||
}
|
||||
}
|
||||
return func() {
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
generation, ok := s.openaiAccountRuntimeBlockGeneration.Load(account.ID)
|
||||
if !ok || generation != installedGeneration {
|
||||
return
|
||||
}
|
||||
current, ok := s.openaiAccountRuntimeBlockUntil.Load(account.ID)
|
||||
currentUntil, isTime := current.(time.Time)
|
||||
if !ok || !isTime || !currentUntil.Equal(installedUntil) {
|
||||
return
|
||||
}
|
||||
if hadBefore {
|
||||
s.openaiAccountRuntimeBlockUntil.Store(account.ID, before)
|
||||
s.openaiAccountRuntimeBlockGeneration.Store(account.ID, s.openaiAccountRuntimeBlockSequence.Add(1))
|
||||
return
|
||||
}
|
||||
s.openaiAccountRuntimeBlockUntil.Delete(account.ID)
|
||||
s.openaiAccountRuntimeBlockGeneration.Store(account.ID, s.openaiAccountRuntimeBlockSequence.Add(1))
|
||||
}
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) grokCredentialConcurrentlyRefreshedToken(ctx context.Context, accountID int64, baseline GrokCredentialMutationSnapshot) (string, bool) {
|
||||
if s == nil || s.accountRepo == nil || accountID <= 0 || ctx == nil || ctx.Err() != nil {
|
||||
return "", false
|
||||
}
|
||||
checkCtx, cancel := context.WithTimeout(ctx, time.Second)
|
||||
defer cancel()
|
||||
latest, err := s.accountRepo.GetByID(checkCtx, accountID)
|
||||
if err != nil || latest == nil {
|
||||
return "", false
|
||||
}
|
||||
latestSnapshot := grokCredentialMutationSnapshot(latest)
|
||||
if !grokCredentialProxyIDsEqual(latestSnapshot.ProxyID, baseline.ProxyID) ||
|
||||
latestSnapshot.CredentialsJSON == baseline.CredentialsJSON || !latest.IsSchedulable() ||
|
||||
(latest.ProxyID != nil && latest.Proxy == nil) || s.isOpenAIAccountRuntimeBlocked(latest) {
|
||||
return "", false
|
||||
}
|
||||
latestToken := strings.TrimSpace(latest.GetGrokAccessToken())
|
||||
if latestToken == "" || strings.TrimSpace(latest.GetGrokRefreshToken()) == "" {
|
||||
return "", false
|
||||
}
|
||||
expiresAt := latest.GetCredentialAsTime("expires_at")
|
||||
if expiresAt == nil || !time.Now().Before(*expiresAt) {
|
||||
return "", false
|
||||
}
|
||||
return latestToken, true
|
||||
}
|
||||
|
||||
func grokCredentialProxyIDsEqual(left, right *int64) bool {
|
||||
if left == nil || right == nil {
|
||||
return left == nil && right == nil
|
||||
}
|
||||
return *left == *right
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) newGrokCredentialFailover(c *gin.Context, account *Account, class grokCredentialFailureClass) error {
|
||||
if strings.TrimSpace(class.message) == "" {
|
||||
class.message = "Grok OAuth credentials are unavailable"
|
||||
}
|
||||
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
|
||||
Platform: PlatformGrok,
|
||||
AccountID: account.ID,
|
||||
Stage: string(GatewayFailureStageAccountAuth),
|
||||
Scope: string(class.scope),
|
||||
Reason: string(class.reason),
|
||||
Kind: "credential_failover",
|
||||
Message: class.message,
|
||||
})
|
||||
return &UpstreamFailoverError{
|
||||
Stage: GatewayFailureStageAccountAuth,
|
||||
Scope: class.scope,
|
||||
Reason: class.reason,
|
||||
NextAccountAction: class.action,
|
||||
ClientStatusCode: http.StatusServiceUnavailable,
|
||||
ClientMessage: GrokCredentialUnavailableClientMessage,
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -13,7 +13,6 @@ import (
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
|
||||
"github.com/Wei-Shaw/sub2api/internal/util/responseheaders"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/tidwall/gjson"
|
||||
@@ -268,25 +267,6 @@ func (s *OpenAIGatewayService) BindGrokMediaVideoRequestAccount(ctx context.Cont
|
||||
return s.BindStickySession(ctx, groupID, GrokMediaVideoRequestSessionHash(requestID), accountID)
|
||||
}
|
||||
|
||||
func (e GrokMediaEndpoint) upstreamURL(baseURL, requestID string) (string, error) {
|
||||
switch e {
|
||||
case GrokMediaEndpointImagesGenerations:
|
||||
return xai.BuildImagesGenerationsURL(baseURL)
|
||||
case GrokMediaEndpointImagesEdits:
|
||||
return xai.BuildImagesEditsURL(baseURL)
|
||||
case GrokMediaEndpointVideosGenerations:
|
||||
return xai.BuildVideosGenerationsURL(baseURL)
|
||||
case GrokMediaEndpointVideosEdits:
|
||||
return xai.BuildVideosEditsURL(baseURL)
|
||||
case GrokMediaEndpointVideosExtensions:
|
||||
return xai.BuildVideosExtensionsURL(baseURL)
|
||||
case GrokMediaEndpointVideoStatus:
|
||||
return xai.BuildVideoURL(baseURL, requestID)
|
||||
default:
|
||||
return "", fmt.Errorf("unsupported grok media endpoint: %s", e)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) ForwardGrokMedia(
|
||||
ctx context.Context,
|
||||
c *gin.Context,
|
||||
@@ -304,11 +284,11 @@ func (s *OpenAIGatewayService) ForwardGrokMedia(
|
||||
return nil, fmt.Errorf("account platform %s is not supported for grok media", account.Platform)
|
||||
}
|
||||
|
||||
token, _, err := s.GetAccessToken(ctx, account)
|
||||
token, _, err := s.getRequestCredential(ctx, c, account)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
targetURL, err := endpoint.upstreamURL(account.GetGrokMediaBaseURL(), requestID)
|
||||
targetURL, err := buildGrokMediaURL(account, s.cfg, endpoint, requestID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -339,7 +319,9 @@ func (s *OpenAIGatewayService) ForwardGrokMedia(
|
||||
}
|
||||
upstreamReq.Header.Set("Authorization", "Bearer "+token)
|
||||
upstreamReq.Header.Set("Accept", "application/json")
|
||||
applyGrokCLIHeaders(upstreamReq.Header)
|
||||
if account.IsGrokOAuth() {
|
||||
applyGrokCLIHeaders(upstreamReq.Header)
|
||||
}
|
||||
if endpoint.RequiresRequestBody() {
|
||||
contentType = strings.TrimSpace(contentType)
|
||||
if contentType == "" {
|
||||
@@ -366,7 +348,7 @@ func (s *OpenAIGatewayService) ForwardGrokMedia(
|
||||
return s.handleGrokMediaErrorResponse(ctx, resp, c, account, requestIDHeader, requestModel)
|
||||
}
|
||||
|
||||
s.updateGrokUsageSnapshot(ctx, account, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode))
|
||||
s.updateGrokUsageFromResponse(ctx, account, resp.Header, resp.StatusCode)
|
||||
respBody, err := ReadUpstreamResponseBody(resp.Body, s.cfg, c, openAITooLargeError)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -635,6 +617,7 @@ func (s *OpenAIGatewayService) handleGrokMediaErrorResponse(
|
||||
return nil, &UpstreamFailoverError{
|
||||
StatusCode: resp.StatusCode,
|
||||
ResponseBody: body,
|
||||
ResponseHeaders: resp.Header.Clone(),
|
||||
RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3,6 +3,7 @@ package service
|
||||
import (
|
||||
"context"
|
||||
"crypto/subtle"
|
||||
"errors"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
@@ -314,10 +315,13 @@ func (s *GrokOAuthService) proxyURL(ctx context.Context, proxyID *int64) (string
|
||||
}
|
||||
proxy, err := s.proxyRepo.GetByID(ctx, *proxyID)
|
||||
if err != nil {
|
||||
return "", infraerrors.Newf(http.StatusBadRequest, "GROK_OAUTH_PROXY_NOT_FOUND", "proxy not found: %v", err)
|
||||
if errors.Is(err, ErrProxyNotFound) {
|
||||
return "", infraerrors.New(http.StatusBadRequest, "GROK_OAUTH_PROXY_NOT_FOUND", "configured proxy was not found")
|
||||
}
|
||||
return "", infraerrors.New(http.StatusServiceUnavailable, "GROK_OAUTH_PROXY_LOOKUP_FAILED", "proxy lookup is temporarily unavailable")
|
||||
}
|
||||
if proxy == nil {
|
||||
return "", nil
|
||||
return "", infraerrors.New(http.StatusBadRequest, "GROK_OAUTH_PROXY_NOT_FOUND", "configured proxy was not found")
|
||||
}
|
||||
return proxy.URL(), nil
|
||||
}
|
||||
|
||||
@@ -135,7 +135,7 @@ func (s *GrokQuotaService) probeUsage(ctx context.Context, accountID int64) (*Gr
|
||||
if err != nil {
|
||||
return nil, infraerrors.Newf(http.StatusBadRequest, "GROK_QUOTA_PROBE_BODY_ERROR", "failed to build probe body: %v", err)
|
||||
}
|
||||
targetURL, err := xai.BuildResponsesURL(account.GetGrokBaseURL())
|
||||
targetURL, err := buildGrokResponsesURL(account, nil)
|
||||
if err != nil {
|
||||
return nil, infraerrors.Newf(http.StatusBadRequest, "GROK_QUOTA_BASE_URL_INVALID", "invalid Grok base_url: %v", err)
|
||||
}
|
||||
@@ -149,7 +149,9 @@ func (s *GrokQuotaService) probeUsage(ctx context.Context, accountID int64) (*Gr
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Accept", "application/json")
|
||||
applyGrokCLIHeaders(req.Header)
|
||||
if account.IsGrokOAuth() {
|
||||
applyGrokCLIHeaders(req.Header)
|
||||
}
|
||||
|
||||
resp, err := s.httpUpstream.Do(req, proxyURL, account.ID, maxInt(account.Concurrency, 1))
|
||||
if err != nil {
|
||||
@@ -158,7 +160,7 @@ func (s *GrokQuotaService) probeUsage(ctx context.Context, accountID int64) (*Gr
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
snapshot := xai.ObserveQuotaHeaders(resp.Header, resp.StatusCode, "active_probe")
|
||||
resetAt, limited := grokRateLimitResetAt(snapshot, time.Now())
|
||||
resetAt, limited := grokRateLimitResetAtForAccount(account, snapshot, time.Now())
|
||||
if limited {
|
||||
normalizeGrokExhaustedWindowResets(snapshot, resetAt, time.Now())
|
||||
}
|
||||
@@ -167,6 +169,8 @@ func (s *GrokQuotaService) probeUsage(ctx context.Context, accountID int64) (*Gr
|
||||
})
|
||||
if limited {
|
||||
persistGrokRateLimit(ctx, s.accountRepo, account, resetAt)
|
||||
} else if isSuccessfulGrokRateLimitRecovery(account, snapshot) {
|
||||
clearGrokRateLimitAfterRecovery(ctx, s.accountRepo, account)
|
||||
}
|
||||
|
||||
result := &GrokQuotaProbeResult{
|
||||
@@ -387,6 +391,7 @@ func (s *GrokQuotaService) prepareProbe(ctx context.Context, accountID int64) (*
|
||||
if err != nil {
|
||||
return nil, "", "", err
|
||||
}
|
||||
proxyURL := s.resolveProxyURL(ctx, account)
|
||||
|
||||
token, err := s.tokenProvider.GetAccessToken(ctx, account)
|
||||
if err != nil {
|
||||
@@ -396,7 +401,7 @@ func (s *GrokQuotaService) prepareProbe(ctx context.Context, accountID int64) (*
|
||||
return nil, "", "", infraerrors.New(http.StatusBadGateway, "GROK_QUOTA_TOKEN_UNAVAILABLE", "access token is empty")
|
||||
}
|
||||
|
||||
return account, token, s.resolveProxyURL(ctx, account), nil
|
||||
return account, token, proxyURL, nil
|
||||
}
|
||||
|
||||
func (s *GrokQuotaService) resolveProxyURL(ctx context.Context, account *Account) string {
|
||||
@@ -408,6 +413,7 @@ func (s *GrokQuotaService) resolveProxyURL(ctx context.Context, account *Account
|
||||
return account.Proxy.URL()
|
||||
case s != nil && s.proxyRepo != nil:
|
||||
if proxy, err := s.proxyRepo.GetByID(ctx, *account.ProxyID); err == nil && proxy != nil {
|
||||
account.Proxy = proxy
|
||||
return proxy.URL()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -32,6 +32,10 @@ type grokQuotaAccountRepo struct {
|
||||
lastTempUnschedID int64
|
||||
lastTempUnschedUntil time.Time
|
||||
lastTempUnschedReason string
|
||||
recoveryClearCalls int
|
||||
recoveryObservedAt time.Time
|
||||
recoveryObservedReset time.Time
|
||||
recoveryClearResult bool
|
||||
}
|
||||
|
||||
func (r *grokQuotaAccountRepo) UpdateExtra(_ context.Context, id int64, updates map[string]any) error {
|
||||
@@ -54,6 +58,13 @@ func (r *grokQuotaAccountRepo) SetRateLimitedIfLater(ctx context.Context, id int
|
||||
return r.SetRateLimited(ctx, id, resetAt)
|
||||
}
|
||||
|
||||
func (r *grokQuotaAccountRepo) ClearRateLimitIfObserved(_ context.Context, _ int64, observedLimitedAt, observedResetAt time.Time) (bool, error) {
|
||||
r.recoveryClearCalls++
|
||||
r.recoveryObservedAt = observedLimitedAt
|
||||
r.recoveryObservedReset = observedResetAt
|
||||
return r.recoveryClearResult, nil
|
||||
}
|
||||
|
||||
func (r *grokQuotaAccountRepo) SetTempUnschedulable(_ context.Context, id int64, until time.Time, reason string) error {
|
||||
r.tempUnschedCalls++
|
||||
r.lastTempUnschedID = id
|
||||
@@ -180,19 +191,26 @@ func (r *grokQuotaProxyRepo) GetByID(_ context.Context, id int64) (*Proxy, error
|
||||
return r.proxies[id], nil
|
||||
}
|
||||
|
||||
func healthyGrokQuotaOAuthAccount(id int64) *Account {
|
||||
return &Account{
|
||||
ID: id,
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Status: StatusActive,
|
||||
Schedulable: true,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "access-token",
|
||||
"refresh_token": "refresh-token",
|
||||
"expires_at": time.Now().Add(2 * grokTokenRefreshSkew).UTC().Format(time.RFC3339),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func TestGrokQuotaServiceProbeUsageStoresHeaders(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
account := &Account{
|
||||
ID: 42,
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "access-token",
|
||||
"expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
|
||||
},
|
||||
}
|
||||
account := healthyGrokQuotaOAuthAccount(42)
|
||||
repo := &grokQuotaAccountRepo{
|
||||
mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
|
||||
accountsByID: map[int64]*Account{42: account},
|
||||
@@ -236,19 +254,10 @@ func TestGrokQuotaServiceProbeUsageStoresHeaders(t *testing.T) {
|
||||
func TestGrokQuotaServiceProbeUsageIgnoresAccountGrokMapping(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
account := &Account{
|
||||
ID: 47,
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "access-token",
|
||||
"expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
|
||||
"model_mapping": map[string]any{
|
||||
"grok": "grok-composer",
|
||||
"grok-composer": "grok-composer-2.5-fast",
|
||||
},
|
||||
},
|
||||
account := healthyGrokQuotaOAuthAccount(47)
|
||||
account.Credentials["model_mapping"] = map[string]any{
|
||||
"grok": "grok-composer",
|
||||
"grok-composer": "grok-composer-2.5-fast",
|
||||
}
|
||||
repo := &grokQuotaAccountRepo{
|
||||
mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
|
||||
@@ -272,16 +281,7 @@ func TestGrokQuotaServiceProbeUsageIgnoresAccountGrokMapping(t *testing.T) {
|
||||
func TestGrokQuotaServiceProbeUsageReportsProbeModelOnUpstreamError(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
account := &Account{
|
||||
ID: 48,
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "access-token",
|
||||
"expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
|
||||
},
|
||||
}
|
||||
account := healthyGrokQuotaOAuthAccount(48)
|
||||
repo := &grokQuotaAccountRepo{
|
||||
mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
|
||||
accountsByID: map[int64]*Account{48: account},
|
||||
@@ -302,16 +302,7 @@ func TestGrokQuotaServiceProbeUsageReportsProbeModelOnUpstreamError(t *testing.T
|
||||
|
||||
func TestGrokQuotaServiceProbeUsageRedactsUpstreamErrorBodyFromErrorAndLogs(t *testing.T) {
|
||||
const upstreamSecret = "upstream-secret-refresh-token"
|
||||
account := &Account{
|
||||
ID: 49,
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "access-token",
|
||||
"expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
|
||||
},
|
||||
}
|
||||
account := healthyGrokQuotaOAuthAccount(49)
|
||||
repo := &grokQuotaAccountRepo{
|
||||
mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
|
||||
accountsByID: map[int64]*Account{49: account},
|
||||
@@ -352,17 +343,8 @@ func TestGrokQuotaServiceProbeUsageLoadsProxyWhenAccountEdgeMissing(t *testing.T
|
||||
t.Parallel()
|
||||
|
||||
proxyID := int64(7)
|
||||
account := &Account{
|
||||
ID: 46,
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Concurrency: 1,
|
||||
ProxyID: &proxyID,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "access-token",
|
||||
"expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
|
||||
},
|
||||
}
|
||||
account := healthyGrokQuotaOAuthAccount(46)
|
||||
account.ProxyID = &proxyID
|
||||
repo := &grokQuotaAccountRepo{
|
||||
mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
|
||||
accountsByID: map[int64]*Account{46: account},
|
||||
@@ -394,20 +376,16 @@ func TestGrokQuotaServiceProbeUsageLoadsProxyWhenAccountEdgeMissing(t *testing.T
|
||||
func TestGrokQuotaServiceProbeUsageStoresNoHeadersState(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
account := &Account{
|
||||
ID: 45,
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "access-token",
|
||||
"expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
|
||||
},
|
||||
}
|
||||
account := healthyGrokQuotaOAuthAccount(45)
|
||||
observedResetAt := time.Now().Add(-time.Second).UTC().Truncate(time.Second)
|
||||
observedLimitedAt := observedResetAt.Add(-grokRateLimitRepeatCooldown)
|
||||
account.RateLimitedAt = &observedLimitedAt
|
||||
account.RateLimitResetAt = &observedResetAt
|
||||
repo := &grokQuotaAccountRepo{
|
||||
mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
|
||||
accountsByID: map[int64]*Account{45: account},
|
||||
},
|
||||
recoveryClearResult: true,
|
||||
}
|
||||
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
@@ -430,20 +408,15 @@ func TestGrokQuotaServiceProbeUsageStoresNoHeadersState(t *testing.T) {
|
||||
require.True(t, ok)
|
||||
require.False(t, stored.HeadersObserved)
|
||||
require.Equal(t, http.StatusOK, stored.StatusCode)
|
||||
require.Equal(t, 1, repo.recoveryClearCalls)
|
||||
require.Equal(t, observedLimitedAt, repo.recoveryObservedAt)
|
||||
require.Equal(t, observedResetAt, repo.recoveryObservedReset)
|
||||
}
|
||||
|
||||
func TestGrokQuotaServiceProbeUsageReturnsRateLimitedSnapshot(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
account := &Account{
|
||||
ID: 43,
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "access-token",
|
||||
"expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
|
||||
},
|
||||
}
|
||||
account := healthyGrokQuotaOAuthAccount(43)
|
||||
repo := &grokQuotaAccountRepo{
|
||||
mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
|
||||
accountsByID: map[int64]*Account{43: account},
|
||||
@@ -471,13 +444,7 @@ func TestGrokQuotaServiceProbeUsageReturnsRateLimitedSnapshot(t *testing.T) {
|
||||
func TestGrokQuotaServiceQueryQuotaFreeFallsBackToGrok45(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
account := &Account{
|
||||
ID: 51, Platform: PlatformGrok, Type: AccountTypeOAuth, Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "access-token",
|
||||
"expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
|
||||
},
|
||||
}
|
||||
account := healthyGrokQuotaOAuthAccount(51)
|
||||
repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
|
||||
accountsByID: map[int64]*Account{account.ID: account},
|
||||
}}
|
||||
@@ -518,13 +485,7 @@ func TestGrokQuotaServiceQueryQuotaFreeFallsBackToGrok45(t *testing.T) {
|
||||
func TestGrokQuotaServiceQueryQuotaPaidBillingSkipsActiveProbe(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
account := &Account{
|
||||
ID: 52, Platform: PlatformGrok, Type: AccountTypeOAuth, Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "access-token",
|
||||
"expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
|
||||
},
|
||||
}
|
||||
account := healthyGrokQuotaOAuthAccount(52)
|
||||
repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
|
||||
accountsByID: map[int64]*Account{account.ID: account},
|
||||
}}
|
||||
@@ -552,13 +513,7 @@ func TestGrokQuotaServiceQueryQuotaPaidBillingSkipsActiveProbe(t *testing.T) {
|
||||
func TestGrokQuotaServiceQueryQuotaCustomPaidMonthlyLimitSkipsActiveProbe(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
account := &Account{
|
||||
ID: 57, Platform: PlatformGrok, Type: AccountTypeOAuth, Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "access-token",
|
||||
"expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
|
||||
},
|
||||
}
|
||||
account := healthyGrokQuotaOAuthAccount(57)
|
||||
repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
|
||||
accountsByID: map[int64]*Account{account.ID: account},
|
||||
}}
|
||||
@@ -691,13 +646,7 @@ func TestGrokLocalUsageForBillingOnlyReturnsAvailableWindows(t *testing.T) {
|
||||
func TestAccountUsageServiceGrokRefreshUsesBillingOnly(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
account := &Account{
|
||||
ID: 54, Platform: PlatformGrok, Type: AccountTypeOAuth, Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "access-token",
|
||||
"expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
|
||||
},
|
||||
}
|
||||
account := healthyGrokQuotaOAuthAccount(54)
|
||||
repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
|
||||
accountsByID: map[int64]*Account{account.ID: account},
|
||||
}}
|
||||
@@ -732,13 +681,7 @@ func TestAccountUsageServiceGrokRefreshUsesBillingOnly(t *testing.T) {
|
||||
func TestGrokQuotaServiceProbeFlightsDeduplicateBillingAndSeparateActive(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
account := &Account{
|
||||
ID: 55, Platform: PlatformGrok, Type: AccountTypeOAuth, Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "access-token",
|
||||
"expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
|
||||
},
|
||||
}
|
||||
account := healthyGrokQuotaOAuthAccount(55)
|
||||
repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
|
||||
accountsByID: map[int64]*Account{account.ID: account},
|
||||
}}
|
||||
@@ -794,13 +737,7 @@ func TestGrokQuotaServiceProbeFlightsDeduplicateBillingAndSeparateActive(t *test
|
||||
func TestGrokQuotaServiceBilling429DoesNotPauseModelScheduling(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
account := &Account{
|
||||
ID: 56, Platform: PlatformGrok, Type: AccountTypeOAuth, Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "access-token",
|
||||
"expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
|
||||
},
|
||||
}
|
||||
account := healthyGrokQuotaOAuthAccount(56)
|
||||
repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
|
||||
accountsByID: map[int64]*Account{account.ID: account},
|
||||
}}
|
||||
@@ -820,13 +757,7 @@ func TestGrokQuotaServiceBilling429DoesNotPauseModelScheduling(t *testing.T) {
|
||||
func TestGrokQuotaServiceQueryQuotaFree429PersistsLimitAndKeepsBilling(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
account := &Account{
|
||||
ID: 53, Platform: PlatformGrok, Type: AccountTypeOAuth, Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "access-token",
|
||||
"expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
|
||||
},
|
||||
}
|
||||
account := healthyGrokQuotaOAuthAccount(53)
|
||||
repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
|
||||
accountsByID: map[int64]*Account{account.ID: account},
|
||||
}}
|
||||
|
||||
@@ -3,19 +3,25 @@ package service
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"log/slog"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/util/logredact"
|
||||
)
|
||||
|
||||
const (
|
||||
grokTokenCacheSkew = 5 * time.Minute
|
||||
grokRequestRefreshTimeout = 8 * time.Second
|
||||
grokTokenProviderLogComponent = "grok_token_provider"
|
||||
grokTempUnschedulableErrorCode = "token_refresh_failed"
|
||||
grokTokenCacheSkew = 5 * time.Minute
|
||||
grokRequestRefreshTimeout = 8 * time.Second
|
||||
grokRefreshLockWaitTimeout = 2 * time.Second
|
||||
grokRefreshLockPollInterval = 25 * time.Millisecond
|
||||
)
|
||||
|
||||
var (
|
||||
errGrokOAuthRefreshNotConfigured = errors.New("grok oauth refresh is not configured")
|
||||
errGrokOAuthRefreshTokenMissing = errors.New("grok oauth refresh token is missing")
|
||||
errGrokOAuthAccessTokenMissing = errors.New("grok oauth access token is missing")
|
||||
errGrokOAuthAccessTokenExpired = errors.New("grok oauth access token is expired")
|
||||
errGrokOAuthConfiguredProxyMiss = errors.New("grok oauth configured proxy is missing")
|
||||
)
|
||||
|
||||
type GrokTokenCache = GeminiTokenCache
|
||||
@@ -36,7 +42,7 @@ func NewGrokTokenProvider(
|
||||
return &GrokTokenProvider{
|
||||
accountRepo: accountRepo,
|
||||
tokenCache: tokenCache,
|
||||
refreshPolicy: AntigravityProviderRefreshPolicy(),
|
||||
refreshPolicy: GrokProviderRefreshPolicy(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -60,32 +66,57 @@ func (p *GrokTokenProvider) GetAccessToken(ctx context.Context, account *Account
|
||||
if account.Platform != PlatformGrok || account.Type != AccountTypeOAuth {
|
||||
return "", errors.New("not a grok oauth account")
|
||||
}
|
||||
|
||||
cacheKey := GrokTokenCacheKey(account)
|
||||
if p.tokenCache != nil {
|
||||
if token, err := p.tokenCache.GetAccessToken(ctx, cacheKey); err == nil && strings.TrimSpace(token) != "" {
|
||||
return token, nil
|
||||
}
|
||||
selectedProxyID := cloneGrokProxyID(account.ProxyID)
|
||||
if eligibilityErr := grokOAuthRequestAccountEligibilityError(account); eligibilityErr != nil {
|
||||
return "", withGrokCredentialFailureSnapshot(eligibilityErr, account)
|
||||
}
|
||||
|
||||
expiresAt := account.GetCredentialAsTime("expires_at")
|
||||
needsRefresh := expiresAt == nil || time.Until(*expiresAt) <= grokTokenRefreshSkew
|
||||
if needsRefresh && strings.TrimSpace(account.GetGrokRefreshToken()) == "" {
|
||||
if expiresAt == nil || !time.Now().Before(*expiresAt) {
|
||||
return "", errors.New("grok access_token expired and refresh_token is missing")
|
||||
}
|
||||
needsRefresh = false
|
||||
accountAccessToken := strings.TrimSpace(account.GetGrokAccessToken())
|
||||
if accountAccessToken == "" {
|
||||
return "", withGrokCredentialFailureSnapshot(errGrokOAuthAccessTokenMissing, account)
|
||||
}
|
||||
if needsRefresh && p.refreshAPI != nil && p.executor != nil {
|
||||
if strings.TrimSpace(account.GetGrokRefreshToken()) == "" {
|
||||
return "", withGrokCredentialFailureSnapshot(errGrokOAuthRefreshTokenMissing, account)
|
||||
}
|
||||
cacheKey := GrokTokenCacheKey(account)
|
||||
if p.tokenCache != nil {
|
||||
if token, err := p.tokenCache.GetAccessToken(ctx, cacheKey); err == nil {
|
||||
cachedToken := strings.TrimSpace(token)
|
||||
if cachedToken != "" && accountAccessToken != "" && cachedToken == accountAccessToken &&
|
||||
expiresAt != nil && time.Until(*expiresAt) > grokTokenRefreshSkew {
|
||||
return cachedToken, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
needsRefresh := expiresAt == nil || time.Until(*expiresAt) <= grokTokenRefreshSkew
|
||||
if needsRefresh {
|
||||
if p.refreshAPI == nil || p.executor == nil {
|
||||
return "", errGrokOAuthRefreshNotConfigured
|
||||
}
|
||||
refreshCtx, cancel := context.WithTimeout(ctx, grokRequestRefreshTimeout)
|
||||
defer cancel()
|
||||
result, err := p.refreshAPI.RefreshIfNeeded(refreshCtx, account, p.executor, grokTokenRefreshSkew)
|
||||
result, err := p.refreshAPI.RefreshIfNeeded(withOAuthRefreshRequestPath(refreshCtx), account, p.executor, grokTokenRefreshSkew)
|
||||
if err != nil {
|
||||
p.markTempUnschedulable(account, err)
|
||||
if p.refreshPolicy.OnRefreshError == ProviderRefreshErrorReturn {
|
||||
return "", err
|
||||
}
|
||||
} else if !result.LockHeld && result.Account != nil {
|
||||
} else if result != nil && result.LockHeld {
|
||||
if p.refreshPolicy.OnLockHeld == ProviderLockHeldWaitForCache {
|
||||
token, waitErr := p.waitForRefreshedToken(refreshCtx, account, cacheKey)
|
||||
return token, withGrokCredentialFailureSnapshot(waitErr, account)
|
||||
}
|
||||
if expiresAt == nil || !time.Now().Before(*expiresAt) {
|
||||
return "", withGrokCredentialFailureSnapshot(errGrokOAuthAccessTokenExpired, account)
|
||||
}
|
||||
} else if result != nil && result.Account != nil {
|
||||
if eligibilityErr := grokOAuthRequestAccountEligibilityError(result.Account); eligibilityErr != nil {
|
||||
return "", withGrokCredentialFailureSnapshot(eligibilityErr, result.Account)
|
||||
}
|
||||
if !grokCredentialProxyIDsEqual(result.Account.ProxyID, selectedProxyID) {
|
||||
return "", withGrokCredentialFailureSnapshot(errOAuthRefreshAccountStateChanged, result.Account)
|
||||
}
|
||||
account = result.Account
|
||||
expiresAt = account.GetCredentialAsTime("expires_at")
|
||||
}
|
||||
@@ -93,15 +124,28 @@ func (p *GrokTokenProvider) GetAccessToken(ctx context.Context, account *Account
|
||||
|
||||
accessToken := account.GetGrokAccessToken()
|
||||
if strings.TrimSpace(accessToken) == "" {
|
||||
return "", errors.New("access_token not found in credentials")
|
||||
return "", withGrokCredentialFailureSnapshot(errGrokOAuthAccessTokenMissing, account)
|
||||
}
|
||||
if expiresAt != nil && !time.Now().Before(*expiresAt) {
|
||||
return "", withGrokCredentialFailureSnapshot(errGrokOAuthAccessTokenExpired, account)
|
||||
}
|
||||
|
||||
if p.tokenCache != nil {
|
||||
latestAccount, isStale := CheckTokenVersion(ctx, account, p.accountRepo)
|
||||
if isStale && latestAccount != nil {
|
||||
if eligibilityErr := grokOAuthRequestAccountEligibilityError(latestAccount); eligibilityErr != nil {
|
||||
return "", withGrokCredentialFailureSnapshot(eligibilityErr, latestAccount)
|
||||
}
|
||||
if !grokCredentialProxyIDsEqual(latestAccount.ProxyID, selectedProxyID) {
|
||||
return "", withGrokCredentialFailureSnapshot(errOAuthRefreshAccountStateChanged, latestAccount)
|
||||
}
|
||||
accessToken = latestAccount.GetGrokAccessToken()
|
||||
if strings.TrimSpace(accessToken) == "" {
|
||||
return "", errors.New("access_token not found after version check")
|
||||
return "", withGrokCredentialFailureSnapshot(errGrokOAuthAccessTokenMissing, latestAccount)
|
||||
}
|
||||
latestExpiry := latestAccount.GetCredentialAsTime("expires_at")
|
||||
if latestExpiry == nil || !time.Now().Before(*latestExpiry) {
|
||||
return "", withGrokCredentialFailureSnapshot(errGrokOAuthAccessTokenExpired, latestAccount)
|
||||
}
|
||||
} else {
|
||||
ttl := 30 * time.Minute
|
||||
@@ -123,40 +167,105 @@ func (p *GrokTokenProvider) GetAccessToken(ctx context.Context, account *Account
|
||||
return accessToken, nil
|
||||
}
|
||||
|
||||
func (p *GrokTokenProvider) markTempUnschedulable(account *Account, refreshErr error) {
|
||||
if p == nil || p.accountRepo == nil || account == nil {
|
||||
return
|
||||
}
|
||||
now := time.Now()
|
||||
until := now.Add(tokenRefreshTempUnschedDuration)
|
||||
redactedErr := "unknown error"
|
||||
if refreshErr != nil {
|
||||
redactedErr = logredact.RedactText(refreshErr.Error())
|
||||
}
|
||||
if isNonRetryableRefreshError(refreshErr) {
|
||||
if err := p.accountRepo.SetError(context.Background(), account.ID, "grok token refresh failed (non-retryable): "+redactedErr); err != nil {
|
||||
slog.Warn(grokTokenProviderLogComponent+".set_error_status_failed", "account_id", account.ID, "error", err)
|
||||
func (p *GrokTokenProvider) waitForRefreshedToken(ctx context.Context, account *Account, cacheKey string) (string, error) {
|
||||
waitCtx, cancel := context.WithTimeout(ctx, grokRefreshLockWaitTimeout)
|
||||
defer cancel()
|
||||
|
||||
initialToken := strings.TrimSpace(account.GetGrokAccessToken())
|
||||
initialVersion := account.GetCredentialAsInt64("_token_version")
|
||||
selectedProxyID := cloneGrokProxyID(account.ProxyID)
|
||||
sawAuthoritativeState := false
|
||||
var lastAccountReadErr error
|
||||
ticker := time.NewTicker(grokRefreshLockPollInterval)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
cachedToken := ""
|
||||
if p.tokenCache != nil {
|
||||
if token, err := p.tokenCache.GetAccessToken(waitCtx, cacheKey); err == nil {
|
||||
cachedToken = strings.TrimSpace(token)
|
||||
}
|
||||
}
|
||||
return
|
||||
}
|
||||
reason := "grok token refresh failed on request path: " + redactedErr
|
||||
bgCtx := context.Background()
|
||||
if err := p.accountRepo.SetTempUnschedulable(bgCtx, account.ID, until, reason); err != nil {
|
||||
slog.Warn(grokTokenProviderLogComponent+".set_temp_unschedulable_failed", "account_id", account.ID, "error", err)
|
||||
return
|
||||
}
|
||||
if p.tempUnschedCache != nil {
|
||||
state := &TempUnschedState{
|
||||
UntilUnix: until.Unix(),
|
||||
TriggeredAtUnix: now.Unix(),
|
||||
ErrorMessage: grokTempUnschedulableErrorCode + ": " + reason,
|
||||
|
||||
if p.accountRepo != nil {
|
||||
latest, err := p.accountRepo.GetByID(waitCtx, account.ID)
|
||||
if err != nil {
|
||||
lastAccountReadErr = err
|
||||
} else if latest == nil {
|
||||
return "", errOAuthRefreshAccountStateChanged
|
||||
} else {
|
||||
sawAuthoritativeState = true
|
||||
if eligibilityErr := grokOAuthRequestAccountEligibilityError(latest); eligibilityErr != nil {
|
||||
return "", withGrokCredentialFailureSnapshot(eligibilityErr, latest)
|
||||
}
|
||||
if !grokCredentialProxyIDsEqual(latest.ProxyID, selectedProxyID) {
|
||||
return "", withGrokCredentialFailureSnapshot(errOAuthRefreshAccountStateChanged, latest)
|
||||
}
|
||||
token := strings.TrimSpace(latest.GetGrokAccessToken())
|
||||
version := latest.GetCredentialAsInt64("_token_version")
|
||||
expiresAt := latest.GetCredentialAsTime("expires_at")
|
||||
changed := token != initialToken || (version > 0 && version > initialVersion)
|
||||
valid := expiresAt != nil && time.Now().Before(*expiresAt)
|
||||
if token != "" && changed && valid {
|
||||
// The versioned DB credential is authoritative. A stale cache must
|
||||
// not hold the request on the old expired token; repair it best-effort.
|
||||
if cachedToken != "" && cachedToken != token {
|
||||
ttl := time.Until(*expiresAt)
|
||||
if ttl > grokTokenCacheSkew {
|
||||
ttl -= grokTokenCacheSkew
|
||||
}
|
||||
_ = p.tokenCache.SetAccessToken(waitCtx, cacheKey, token, ttl)
|
||||
}
|
||||
return token, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
if err := p.tempUnschedCache.SetTempUnsched(bgCtx, account.ID, state); err != nil {
|
||||
slog.Warn(grokTokenProviderLogComponent+".temp_unsched_cache_set_failed", "account_id", account.ID, "error", err)
|
||||
|
||||
select {
|
||||
case <-waitCtx.Done():
|
||||
if ctx.Err() != nil {
|
||||
return "", ctx.Err()
|
||||
}
|
||||
if !sawAuthoritativeState {
|
||||
if lastAccountReadErr == nil {
|
||||
lastAccountReadErr = waitCtx.Err()
|
||||
}
|
||||
return "", fmt.Errorf("%w: %v", errOAuthRefreshAccountRereadFailed, lastAccountReadErr)
|
||||
}
|
||||
// Another worker still owns the refresh and the authoritative row is
|
||||
// unchanged. Do not quarantine the old credential: its refresh may
|
||||
// commit immediately after this bounded wait.
|
||||
return "", errOAuthRefreshAccountStateChanged
|
||||
case <-ticker.C:
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func grokOAuthRequestAccountEligibilityError(account *Account) error {
|
||||
if account == nil || !account.IsGrokOAuth() || !account.IsSchedulable() {
|
||||
return errOAuthRefreshAccountStateChanged
|
||||
}
|
||||
if account.ProxyID != nil && account.Proxy == nil {
|
||||
return errGrokOAuthConfiguredProxyMiss
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func cloneGrokProxyID(proxyID *int64) *int64 {
|
||||
if proxyID == nil {
|
||||
return nil
|
||||
}
|
||||
value := *proxyID
|
||||
return &value
|
||||
}
|
||||
|
||||
func (p *GrokTokenProvider) InvalidateToken(ctx context.Context, account *Account) error {
|
||||
if p == nil || p.tokenCache == nil || account == nil {
|
||||
return nil
|
||||
}
|
||||
return p.tokenCache.DeleteAccessToken(ctx, GrokTokenCacheKey(account))
|
||||
}
|
||||
|
||||
func GrokTokenCacheKey(account *Account) string {
|
||||
if account == nil {
|
||||
return "grok:account:0"
|
||||
|
||||
@@ -5,6 +5,7 @@ package service
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
@@ -19,9 +20,33 @@ type grokTokenCacheForProviderTest struct {
|
||||
setTTL time.Duration
|
||||
lockResult bool
|
||||
releaseCalls int
|
||||
deletedKeys []string
|
||||
deleteErr error
|
||||
getCalls int
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
type grokCredentialRaceRepo struct {
|
||||
*tokenRefreshAccountRepo
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
func (r *grokCredentialRaceRepo) GetByID(ctx context.Context, id int64) (*Account, error) {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
return r.tokenRefreshAccountRepo.GetByID(ctx, id)
|
||||
}
|
||||
|
||||
func (r *grokCredentialRaceRepo) setAccount(account *Account) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.accountsByID[account.ID] = account
|
||||
}
|
||||
|
||||
func (c *grokTokenCacheForProviderTest) GetAccessToken(context.Context, string) (string, error) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
c.getCalls++
|
||||
if c.token == "" {
|
||||
return "", errors.New("not cached")
|
||||
}
|
||||
@@ -29,14 +54,19 @@ func (c *grokTokenCacheForProviderTest) GetAccessToken(context.Context, string)
|
||||
}
|
||||
|
||||
func (c *grokTokenCacheForProviderTest) SetAccessToken(_ context.Context, key string, token string, ttl time.Duration) error {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
c.setKey = key
|
||||
c.setToken = token
|
||||
c.setTTL = ttl
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *grokTokenCacheForProviderTest) DeleteAccessToken(context.Context, string) error {
|
||||
return nil
|
||||
func (c *grokTokenCacheForProviderTest) DeleteAccessToken(_ context.Context, key string) error {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
c.deletedKeys = append(c.deletedKeys, key)
|
||||
return c.deleteErr
|
||||
}
|
||||
|
||||
func (c *grokTokenCacheForProviderTest) AcquireRefreshLock(context.Context, string, time.Duration) (bool, error) {
|
||||
@@ -53,9 +83,11 @@ func TestGrokTokenProviderRefreshesExpiredTokenOnRequestPath(t *testing.T) {
|
||||
|
||||
expiredAt := time.Now().Add(-time.Minute).UTC().Format(time.RFC3339)
|
||||
account := &Account{
|
||||
ID: 54,
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
ID: 54,
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Status: StatusActive,
|
||||
Schedulable: true,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "expired-access-token",
|
||||
"refresh_token": "refresh-token",
|
||||
@@ -95,9 +127,11 @@ func TestGrokTokenProviderRefreshesExpiredTokenOnRequestPath(t *testing.T) {
|
||||
func TestGrokTokenProviderRefreshFailureUnschedulesWithRedactedReason(t *testing.T) {
|
||||
expiredAt := time.Now().Add(-time.Minute).UTC().Format(time.RFC3339)
|
||||
account := &Account{
|
||||
ID: 55,
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
ID: 55,
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Status: StatusActive,
|
||||
Schedulable: true,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "expired-access-token",
|
||||
"refresh_token": "refresh-token",
|
||||
@@ -108,24 +142,178 @@ func TestGrokTokenProviderRefreshFailureUnschedulesWithRedactedReason(t *testing
|
||||
repo := &tokenRefreshAccountRepo{}
|
||||
repo.accountsByID = map[int64]*Account{55: account}
|
||||
cache := &grokTokenCacheForProviderTest{lockResult: true}
|
||||
tempCache := &tempUnschedCacheStub{}
|
||||
provider := NewGrokTokenProvider(repo, cache)
|
||||
provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{
|
||||
err: errors.New("temporary refresh failure access_token=leaked-access refresh_token=leaked-refresh"),
|
||||
})
|
||||
provider.SetTempUnschedCache(tempCache)
|
||||
|
||||
token, err := provider.GetAccessToken(context.Background(), account)
|
||||
require.Error(t, err)
|
||||
require.Empty(t, token)
|
||||
require.Equal(t, 1, repo.setTempUnschedCalls)
|
||||
require.Equal(t, 0, repo.setTempUnschedCalls)
|
||||
require.Equal(t, 0, repo.setErrorCalls)
|
||||
require.Contains(t, repo.lastTempUnschedReason, "access_token=***")
|
||||
require.Contains(t, repo.lastTempUnschedReason, "refresh_token=***")
|
||||
require.NotContains(t, repo.lastTempUnschedReason, "leaked-access")
|
||||
require.NotContains(t, repo.lastTempUnschedReason, "leaked-refresh")
|
||||
require.Equal(t, 1, tempCache.setCalls)
|
||||
require.NotNil(t, tempCache.lastState)
|
||||
require.NotContains(t, tempCache.lastState.ErrorMessage, "leaked-access")
|
||||
require.NotContains(t, tempCache.lastState.ErrorMessage, "leaked-refresh")
|
||||
}
|
||||
|
||||
func TestGrokTokenProviderLockHeldWaitsForRefreshedCacheAndNeverUsesExpiredToken(t *testing.T) {
|
||||
account := expiredGrokOAuthAccountForCredentialTest(56)
|
||||
baseRepo := &tokenRefreshAccountRepo{}
|
||||
baseRepo.accountsByID = map[int64]*Account{account.ID: account}
|
||||
repo := &grokCredentialRaceRepo{tokenRefreshAccountRepo: baseRepo}
|
||||
cache := &grokTokenCacheForProviderTest{lockResult: false, token: "expired-access-token"}
|
||||
provider := NewGrokTokenProvider(repo, cache)
|
||||
provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{})
|
||||
|
||||
go func() {
|
||||
time.Sleep(40 * time.Millisecond)
|
||||
refreshed := *account
|
||||
refreshed.Credentials = shallowCopyMap(account.Credentials)
|
||||
refreshed.Credentials["access_token"] = "refreshed-after-lock"
|
||||
refreshed.Credentials["expires_at"] = time.Now().Add(time.Hour).UTC().Format(time.RFC3339)
|
||||
refreshed.Credentials["_token_version"] = time.Now().UnixMilli()
|
||||
repo.setAccount(&refreshed)
|
||||
cache.mu.Lock()
|
||||
cache.token = "refreshed-after-lock"
|
||||
cache.mu.Unlock()
|
||||
}()
|
||||
|
||||
startedAt := time.Now()
|
||||
token, err := provider.GetAccessToken(context.Background(), account)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "refreshed-after-lock", token)
|
||||
require.NotEqual(t, "expired-access-token", token)
|
||||
require.GreaterOrEqual(t, time.Since(startedAt), 25*time.Millisecond,
|
||||
"expired account metadata must prevent returning the old cached token")
|
||||
}
|
||||
|
||||
func TestGrokTokenProviderLockHeldTimeoutDoesNotReturnExpiredToken(t *testing.T) {
|
||||
account := expiredGrokOAuthAccountForCredentialTest(57)
|
||||
repo := &tokenRefreshAccountRepo{}
|
||||
repo.accountsByID = map[int64]*Account{account.ID: account}
|
||||
cache := &grokTokenCacheForProviderTest{lockResult: false}
|
||||
provider := NewGrokTokenProvider(repo, cache)
|
||||
provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{})
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 80*time.Millisecond)
|
||||
defer cancel()
|
||||
|
||||
token, err := provider.GetAccessToken(ctx, account)
|
||||
require.Error(t, err)
|
||||
require.Empty(t, token)
|
||||
}
|
||||
|
||||
func TestGrokTokenProviderLockHeldRejectsChangedTokenWithoutExpiry(t *testing.T) {
|
||||
account := expiredGrokOAuthAccountForCredentialTest(58)
|
||||
baseRepo := &tokenRefreshAccountRepo{}
|
||||
baseRepo.accountsByID = map[int64]*Account{account.ID: account}
|
||||
repo := &grokCredentialRaceRepo{tokenRefreshAccountRepo: baseRepo}
|
||||
cache := &grokTokenCacheForProviderTest{lockResult: false, token: "expired-access-token"}
|
||||
provider := NewGrokTokenProvider(repo, cache)
|
||||
provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{})
|
||||
|
||||
go func() {
|
||||
time.Sleep(30 * time.Millisecond)
|
||||
refreshed := *account
|
||||
refreshed.Credentials = shallowCopyMap(account.Credentials)
|
||||
refreshed.Credentials["access_token"] = "changed-without-expiry"
|
||||
delete(refreshed.Credentials, "expires_at")
|
||||
refreshed.Credentials["_token_version"] = time.Now().UnixMilli()
|
||||
repo.setAccount(&refreshed)
|
||||
cache.mu.Lock()
|
||||
cache.token = "changed-without-expiry"
|
||||
cache.mu.Unlock()
|
||||
}()
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
|
||||
defer cancel()
|
||||
token, err := provider.GetAccessToken(ctx, account)
|
||||
|
||||
require.Error(t, err)
|
||||
require.Empty(t, token, "an unbounded credential must not win the lock-held race")
|
||||
}
|
||||
|
||||
func TestGrokTokenProviderLockHeldUsesVersionedDBTokenAndRepairsStaleCache(t *testing.T) {
|
||||
account := expiredGrokOAuthAccountForCredentialTest(60)
|
||||
baseRepo := &tokenRefreshAccountRepo{}
|
||||
baseRepo.accountsByID = map[int64]*Account{account.ID: account}
|
||||
repo := &grokCredentialRaceRepo{tokenRefreshAccountRepo: baseRepo}
|
||||
cache := &grokTokenCacheForProviderTest{lockResult: false, token: "expired-access-token"}
|
||||
provider := NewGrokTokenProvider(repo, cache)
|
||||
provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{})
|
||||
|
||||
go func() {
|
||||
time.Sleep(30 * time.Millisecond)
|
||||
refreshed := *account
|
||||
refreshed.Credentials = shallowCopyMap(account.Credentials)
|
||||
refreshed.Credentials["access_token"] = "db-authoritative-token"
|
||||
refreshed.Credentials["expires_at"] = time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339)
|
||||
refreshed.Credentials["_token_version"] = time.Now().UnixMilli()
|
||||
repo.setAccount(&refreshed)
|
||||
}()
|
||||
|
||||
token, err := provider.GetAccessToken(context.Background(), account)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "db-authoritative-token", token)
|
||||
require.Equal(t, "db-authoritative-token", cache.setToken)
|
||||
require.Greater(t, cache.setTTL, time.Duration(0))
|
||||
}
|
||||
|
||||
func TestGrokTokenProviderRejectsStaleDBTokenWithoutExpiry(t *testing.T) {
|
||||
expiresAt := time.Now().Add(2 * grokTokenRefreshSkew).UTC().Format(time.RFC3339)
|
||||
account := &Account{
|
||||
ID: 59,
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Status: StatusActive,
|
||||
Schedulable: true,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "old-access-token",
|
||||
"refresh_token": "refresh-token",
|
||||
"expires_at": expiresAt,
|
||||
},
|
||||
}
|
||||
latest := *account
|
||||
latest.Credentials = shallowCopyMap(account.Credentials)
|
||||
latest.Credentials["access_token"] = "new-access-token-without-expiry"
|
||||
latest.Credentials["_token_version"] = time.Now().UnixMilli()
|
||||
delete(latest.Credentials, "expires_at")
|
||||
repo := &tokenRefreshAccountRepo{}
|
||||
repo.accountsByID = map[int64]*Account{account.ID: &latest}
|
||||
cache := &grokTokenCacheForProviderTest{}
|
||||
provider := NewGrokTokenProvider(repo, cache)
|
||||
|
||||
token, err := provider.GetAccessToken(context.Background(), account)
|
||||
|
||||
require.ErrorIs(t, err, errGrokOAuthAccessTokenExpired)
|
||||
require.Empty(t, token)
|
||||
}
|
||||
|
||||
func TestGrokTokenProviderRejectsIneligibleSelectedAccountBeforeWarmCache(t *testing.T) {
|
||||
future := time.Now().Add(time.Hour)
|
||||
tests := []struct {
|
||||
name string
|
||||
mutate func(*Account)
|
||||
}{
|
||||
{name: "disabled", mutate: func(account *Account) { account.Status = StatusDisabled }},
|
||||
{name: "not schedulable", mutate: func(account *Account) { account.Schedulable = false }},
|
||||
{name: "temporarily unschedulable", mutate: func(account *Account) { account.TempUnschedulableUntil = &future }},
|
||||
{name: "rate limited", mutate: func(account *Account) { account.RateLimitResetAt = &future }},
|
||||
{name: "overloaded", mutate: func(account *Account) { account.OverloadUntil = &future }},
|
||||
}
|
||||
|
||||
for index, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
account := expiredGrokOAuthAccountForCredentialTest(int64(90 + index))
|
||||
account.Credentials["access_token"] = "warm-cache-token"
|
||||
account.Credentials["expires_at"] = time.Now().Add(2 * grokTokenRefreshSkew).UTC().Format(time.RFC3339)
|
||||
tt.mutate(account)
|
||||
cache := &grokTokenCacheForProviderTest{token: "warm-cache-token"}
|
||||
provider := NewGrokTokenProvider(&tokenRefreshAccountRepo{}, cache)
|
||||
|
||||
token, err := provider.GetAccessToken(context.Background(), account)
|
||||
|
||||
require.ErrorIs(t, err, errOAuthRefreshAccountStateChanged)
|
||||
require.Empty(t, token)
|
||||
require.Zero(t, cache.getCalls, "an ineligible selected account must be rejected before cache lookup")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -22,7 +22,8 @@ func (r *GrokTokenRefresher) CacheKey(account *Account) string {
|
||||
}
|
||||
|
||||
func (r *GrokTokenRefresher) CanRefresh(account *Account) bool {
|
||||
return account != nil && account.Platform == PlatformGrok && account.Type == AccountTypeOAuth
|
||||
return account != nil && account.Platform == PlatformGrok && account.Type == AccountTypeOAuth &&
|
||||
strings.TrimSpace(account.GetGrokRefreshToken()) != ""
|
||||
}
|
||||
|
||||
func (r *GrokTokenRefresher) NeedsRefresh(account *Account, refreshWindow time.Duration) bool {
|
||||
|
||||
@@ -0,0 +1,90 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
|
||||
"github.com/Wei-Shaw/sub2api/internal/util/urlvalidator"
|
||||
)
|
||||
|
||||
func grokBaseURLValidator(account *Account, cfg *config.Config) (xai.BaseURLValidator, error) {
|
||||
if account == nil || !account.IsGrok() {
|
||||
return nil, fmt.Errorf("grok account is required")
|
||||
}
|
||||
switch account.Type {
|
||||
case AccountTypeOAuth:
|
||||
// Subscription credentials are never governed by the operator's API-key
|
||||
// URL policy. They stay pinned to the supported CLI gateway.
|
||||
return redactedGrokBaseURLValidator(xai.ValidateTrustedBaseURL), nil
|
||||
case AccountTypeAPIKey:
|
||||
if cfg == nil {
|
||||
return redactedGrokBaseURLValidator(xai.ValidateBaseURL), nil
|
||||
}
|
||||
if !cfg.Security.URLAllowlist.Enabled {
|
||||
return redactedGrokBaseURLValidator(func(raw string) (string, error) {
|
||||
return urlvalidator.ValidateURLFormat(raw, cfg.Security.URLAllowlist.AllowInsecureHTTP)
|
||||
}), nil
|
||||
}
|
||||
return redactedGrokBaseURLValidator(func(raw string) (string, error) {
|
||||
return urlvalidator.ValidateHTTPSURL(raw, urlvalidator.ValidationOptions{
|
||||
AllowedHosts: cfg.Security.URLAllowlist.UpstreamHosts,
|
||||
RequireAllowlist: true,
|
||||
AllowPrivate: cfg.Security.URLAllowlist.AllowPrivateHosts,
|
||||
})
|
||||
}), nil
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported grok account type: %s", account.Type)
|
||||
}
|
||||
}
|
||||
|
||||
func redactedGrokBaseURLValidator(validator xai.BaseURLValidator) xai.BaseURLValidator {
|
||||
return func(raw string) (string, error) {
|
||||
validated, err := validator(raw)
|
||||
if err != nil {
|
||||
return "", errors.New("base URL rejected by URL security policy")
|
||||
}
|
||||
return validated, nil
|
||||
}
|
||||
}
|
||||
|
||||
func buildGrokResponsesURL(account *Account, cfg *config.Config) (string, error) {
|
||||
validator, err := grokBaseURLValidator(account, cfg)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return xai.BuildResponsesURLWithValidator(account.GetGrokBaseURL(), validator)
|
||||
}
|
||||
|
||||
func buildGrokChatCompletionsURL(account *Account, cfg *config.Config) (string, error) {
|
||||
validator, err := grokBaseURLValidator(account, cfg)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return xai.BuildChatCompletionsURLWithValidator(account.GetGrokBaseURL(), validator)
|
||||
}
|
||||
|
||||
func buildGrokMediaURL(account *Account, cfg *config.Config, endpoint GrokMediaEndpoint, requestID string) (string, error) {
|
||||
validator, err := grokBaseURLValidator(account, cfg)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
baseURL := account.GetGrokMediaBaseURL()
|
||||
switch endpoint {
|
||||
case GrokMediaEndpointImagesGenerations:
|
||||
return xai.BuildImagesGenerationsURLWithValidator(baseURL, validator)
|
||||
case GrokMediaEndpointImagesEdits:
|
||||
return xai.BuildImagesEditsURLWithValidator(baseURL, validator)
|
||||
case GrokMediaEndpointVideosGenerations:
|
||||
return xai.BuildVideosGenerationsURLWithValidator(baseURL, validator)
|
||||
case GrokMediaEndpointVideosEdits:
|
||||
return xai.BuildVideosEditsURLWithValidator(baseURL, validator)
|
||||
case GrokMediaEndpointVideosExtensions:
|
||||
return xai.BuildVideosExtensionsURLWithValidator(baseURL, validator)
|
||||
case GrokMediaEndpointVideoStatus:
|
||||
return xai.BuildVideoURLWithValidator(baseURL, requestID, validator)
|
||||
default:
|
||||
return "", fmt.Errorf("unsupported grok media endpoint: %s", endpoint)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,122 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestGrokAPIKeyURLPolicyFollowsGlobalSecurityConfig(t *testing.T) {
|
||||
account := &Account{
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeAPIKey,
|
||||
Credentials: map[string]any{
|
||||
"base_url": "http://grok.example.test/v1",
|
||||
},
|
||||
}
|
||||
|
||||
t.Run("insecure HTTP enabled with allowlist disabled", func(t *testing.T) {
|
||||
cfg := &config.Config{}
|
||||
cfg.Security.URLAllowlist.Enabled = false
|
||||
cfg.Security.URLAllowlist.AllowInsecureHTTP = true
|
||||
|
||||
responsesURL, err := buildGrokResponsesURL(account, cfg)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "http://grok.example.test/v1/responses", responsesURL)
|
||||
|
||||
chatURL, err := buildGrokChatCompletionsURL(account, cfg)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "http://grok.example.test/v1/chat/completions", chatURL)
|
||||
|
||||
mediaURL, err := buildGrokMediaURL(account, cfg, GrokMediaEndpointImagesGenerations, "")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "http://grok.example.test/v1/images/generations", mediaURL)
|
||||
})
|
||||
|
||||
t.Run("insecure HTTP disabled", func(t *testing.T) {
|
||||
cfg := &config.Config{}
|
||||
cfg.Security.URLAllowlist.Enabled = false
|
||||
cfg.Security.URLAllowlist.AllowInsecureHTTP = false
|
||||
|
||||
_, err := buildGrokResponsesURL(account, cfg)
|
||||
require.EqualError(t, err, "invalid base url: base URL rejected by URL security policy")
|
||||
})
|
||||
|
||||
t.Run("enabled allowlist remains HTTPS only", func(t *testing.T) {
|
||||
cfg := &config.Config{}
|
||||
cfg.Security.URLAllowlist.Enabled = true
|
||||
cfg.Security.URLAllowlist.AllowInsecureHTTP = true
|
||||
cfg.Security.URLAllowlist.UpstreamHosts = []string{"grok.example.test"}
|
||||
|
||||
_, err := buildGrokResponsesURL(account, cfg)
|
||||
require.EqualError(t, err, "invalid base url: base URL rejected by URL security policy")
|
||||
})
|
||||
}
|
||||
|
||||
func TestGrokAPIKeyURLPolicyAppliesAllowlistAndPrivateHostControls(t *testing.T) {
|
||||
account := &Account{
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeAPIKey,
|
||||
Credentials: map[string]any{
|
||||
"base_url": "https://grok.example.test/v1",
|
||||
},
|
||||
}
|
||||
cfg := &config.Config{}
|
||||
cfg.Security.URLAllowlist.Enabled = true
|
||||
cfg.Security.URLAllowlist.UpstreamHosts = []string{"grok.example.test"}
|
||||
|
||||
target, err := buildGrokResponsesURL(account, cfg)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "https://grok.example.test/v1/responses", target)
|
||||
|
||||
cfg.Security.URLAllowlist.UpstreamHosts = []string{"other.example.test"}
|
||||
_, err = buildGrokResponsesURL(account, cfg)
|
||||
require.EqualError(t, err, "invalid base url: base URL rejected by URL security policy")
|
||||
|
||||
account.Credentials["base_url"] = "https://127.0.0.1/v1"
|
||||
cfg.Security.URLAllowlist.UpstreamHosts = []string{"127.0.0.1"}
|
||||
_, err = buildGrokResponsesURL(account, cfg)
|
||||
require.EqualError(t, err, "invalid base url: base URL rejected by URL security policy")
|
||||
|
||||
cfg.Security.URLAllowlist.AllowPrivateHosts = true
|
||||
target, err = buildGrokResponsesURL(account, cfg)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "https://127.0.0.1/v1/responses", target)
|
||||
}
|
||||
|
||||
func TestGrokAPIKeyURLPolicyRedactsMalformedConfiguredURL(t *testing.T) {
|
||||
account := &Account{
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeAPIKey,
|
||||
Credentials: map[string]any{
|
||||
"base_url": "https://%zz:secret@grok.example.test/v1",
|
||||
},
|
||||
}
|
||||
cfg := &config.Config{}
|
||||
cfg.Security.URLAllowlist.AllowInsecureHTTP = true
|
||||
|
||||
_, err := buildGrokResponsesURL(account, cfg)
|
||||
require.EqualError(t, err, "invalid base url: base URL rejected by URL security policy")
|
||||
require.NotContains(t, err.Error(), "secret")
|
||||
}
|
||||
|
||||
func TestGrokOAuthURLPolicyIgnoresAPIKeyOverrides(t *testing.T) {
|
||||
account := &Account{
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Credentials: map[string]any{
|
||||
"base_url": "http://attacker.example.test/v1",
|
||||
},
|
||||
}
|
||||
cfg := &config.Config{}
|
||||
cfg.Security.URLAllowlist.Enabled = false
|
||||
cfg.Security.URLAllowlist.AllowInsecureHTTP = true
|
||||
|
||||
target, err := buildGrokResponsesURL(account, cfg)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, xai.DefaultCLIBaseURL+"/responses", target)
|
||||
}
|
||||
@@ -2,6 +2,7 @@ package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"strconv"
|
||||
@@ -20,6 +21,45 @@ type OAuthRefreshExecutor interface {
|
||||
}
|
||||
|
||||
const defaultRefreshLockTTL = 60 * time.Second
|
||||
const oauthRefreshLockCleanupTimeout = 2 * time.Second
|
||||
|
||||
var (
|
||||
errOAuthRefreshAccountRereadFailed = errors.New("oauth refresh account reread failed")
|
||||
errOAuthRefreshAccountStateChanged = errors.New("oauth refresh account state changed")
|
||||
errOAuthRefreshCredentialPersist = errors.New("oauth refresh credential persistence failed")
|
||||
)
|
||||
|
||||
type oauthRefreshRequestPathKey struct{}
|
||||
|
||||
func withOAuthRefreshRequestPath(ctx context.Context) context.Context {
|
||||
return context.WithValue(ctx, oauthRefreshRequestPathKey{}, true)
|
||||
}
|
||||
|
||||
func isOAuthRefreshRequestPath(ctx context.Context) bool {
|
||||
requestPath, _ := ctx.Value(oauthRefreshRequestPathKey{}).(bool)
|
||||
return requestPath
|
||||
}
|
||||
|
||||
type oauthRefreshLocalLock struct {
|
||||
semaphore chan struct{}
|
||||
}
|
||||
|
||||
func newOAuthRefreshLocalLock() *oauthRefreshLocalLock {
|
||||
return &oauthRefreshLocalLock{semaphore: make(chan struct{}, 1)}
|
||||
}
|
||||
|
||||
func (l *oauthRefreshLocalLock) Lock(ctx context.Context) error {
|
||||
select {
|
||||
case l.semaphore <- struct{}{}:
|
||||
return nil
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
func (l *oauthRefreshLocalLock) Unlock() {
|
||||
<-l.semaphore
|
||||
}
|
||||
|
||||
// OAuthRefreshResult 统一刷新结果
|
||||
type OAuthRefreshResult struct {
|
||||
@@ -35,7 +75,7 @@ type OAuthRefreshAPI struct {
|
||||
accountRepo AccountRepository
|
||||
tokenCache GeminiTokenCache // 可选,nil = 无分布式锁
|
||||
lockTTL time.Duration
|
||||
localLocks sync.Map // key: cacheKey string -> value: *sync.Mutex
|
||||
localLocks sync.Map // key: cacheKey string -> value: *oauthRefreshLocalLock
|
||||
}
|
||||
|
||||
// NewOAuthRefreshAPI 创建统一刷新 API
|
||||
@@ -53,11 +93,11 @@ func NewOAuthRefreshAPI(accountRepo AccountRepository, tokenCache GeminiTokenCac
|
||||
}
|
||||
|
||||
// getLocalLock 返回指定 cacheKey 的进程内互斥锁
|
||||
func (api *OAuthRefreshAPI) getLocalLock(cacheKey string) *sync.Mutex {
|
||||
actual, _ := api.localLocks.LoadOrStore(cacheKey, &sync.Mutex{})
|
||||
mu, ok := actual.(*sync.Mutex)
|
||||
func (api *OAuthRefreshAPI) getLocalLock(cacheKey string) *oauthRefreshLocalLock {
|
||||
actual, _ := api.localLocks.LoadOrStore(cacheKey, newOAuthRefreshLocalLock())
|
||||
mu, ok := actual.(*oauthRefreshLocalLock)
|
||||
if !ok {
|
||||
mu = &sync.Mutex{}
|
||||
mu = newOAuthRefreshLocalLock()
|
||||
api.localLocks.Store(cacheKey, mu)
|
||||
}
|
||||
return mu
|
||||
@@ -78,15 +118,25 @@ func (api *OAuthRefreshAPI) RefreshIfNeeded(
|
||||
executor OAuthRefreshExecutor,
|
||||
refreshWindow time.Duration,
|
||||
) (*OAuthRefreshResult, error) {
|
||||
if api == nil || api.accountRepo == nil {
|
||||
return nil, errors.New("oauth refresh account repository is not configured")
|
||||
}
|
||||
if account == nil {
|
||||
return nil, errors.New("oauth refresh account is nil")
|
||||
}
|
||||
if executor == nil {
|
||||
return nil, errors.New("oauth refresh executor is nil")
|
||||
}
|
||||
cacheKey := executor.CacheKey(account)
|
||||
|
||||
// 0. 获取进程内互斥锁(防止同一进程内的并发刷新竞争)
|
||||
localMu := api.getLocalLock(cacheKey)
|
||||
localMu.Lock()
|
||||
if err := localMu.Lock(ctx); err != nil {
|
||||
return nil, fmt.Errorf("oauth refresh local lock: %w", err)
|
||||
}
|
||||
defer localMu.Unlock()
|
||||
|
||||
// 1. 获取分布式锁
|
||||
lockAcquired := false
|
||||
if api.tokenCache != nil {
|
||||
acquired, lockErr := api.tokenCache.AcquireRefreshLock(ctx, cacheKey, api.lockTTL)
|
||||
if lockErr != nil {
|
||||
@@ -100,22 +150,38 @@ func (api *OAuthRefreshAPI) RefreshIfNeeded(
|
||||
// 锁被其他 worker 持有
|
||||
return &OAuthRefreshResult{LockHeld: true}, nil
|
||||
} else {
|
||||
lockAcquired = true
|
||||
defer func() { _ = api.tokenCache.ReleaseRefreshLock(ctx, cacheKey) }()
|
||||
defer func() {
|
||||
cleanupCtx, cancel := context.WithTimeout(context.Background(), oauthRefreshLockCleanupTimeout)
|
||||
defer cancel()
|
||||
_ = api.tokenCache.ReleaseRefreshLock(cleanupCtx, cacheKey)
|
||||
}()
|
||||
}
|
||||
}
|
||||
|
||||
// 2. 从 DB 重读最新 account(锁保护下,确保使用最新的 refresh_token)
|
||||
freshAccount, err := api.accountRepo.GetByID(ctx, account.ID)
|
||||
if err != nil {
|
||||
slog.Warn("oauth_refresh_db_reread_failed",
|
||||
"account_id", account.ID,
|
||||
"error", err,
|
||||
)
|
||||
// 降级使用传入的 account
|
||||
freshAccount = account
|
||||
} else if freshAccount == nil {
|
||||
freshAccount = account
|
||||
return nil, fmt.Errorf("%w: %v", errOAuthRefreshAccountRereadFailed, err)
|
||||
}
|
||||
if freshAccount == nil {
|
||||
return nil, fmt.Errorf("%w: account not found", errOAuthRefreshAccountStateChanged)
|
||||
}
|
||||
if freshAccount.ID != account.ID {
|
||||
return nil, fmt.Errorf("%w: account identity mismatch", errOAuthRefreshAccountRereadFailed)
|
||||
}
|
||||
if !freshAccount.IsActive() {
|
||||
return nil, fmt.Errorf("%w: account is not active", errOAuthRefreshAccountStateChanged)
|
||||
}
|
||||
if isOAuthRefreshRequestPath(ctx) && freshAccount.Platform == PlatformGrok {
|
||||
if eligibilityErr := grokOAuthRequestAccountEligibilityError(freshAccount); eligibilityErr != nil {
|
||||
return nil, withGrokCredentialFailureSnapshot(eligibilityErr, freshAccount)
|
||||
}
|
||||
}
|
||||
if !executor.CanRefresh(freshAccount) {
|
||||
if freshAccount.IsGrokOAuth() && strings.TrimSpace(freshAccount.GetGrokRefreshToken()) == "" {
|
||||
return nil, withGrokCredentialFailureSnapshot(errGrokOAuthRefreshTokenMissing, freshAccount)
|
||||
}
|
||||
return nil, fmt.Errorf("%w: account is no longer refreshable", errOAuthRefreshAccountStateChanged)
|
||||
}
|
||||
|
||||
// 3. 二次检查是否仍需刷新(另一条路径可能已刷新)
|
||||
@@ -127,11 +193,19 @@ func (api *OAuthRefreshAPI) RefreshIfNeeded(
|
||||
|
||||
// 4. 执行平台特定刷新逻辑
|
||||
newCredentials, refreshErr := executor.Refresh(ctx, freshAccount)
|
||||
if err := ctx.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if refreshErr != nil {
|
||||
// 竞争恢复:invalid_grant 可能是另一个 worker 已消费了旧 refresh_token
|
||||
// 重新读取 DB,如果 refresh_token 已更新则说明是竞争,返回成功
|
||||
if isInvalidGrantError(refreshErr) {
|
||||
if recoveredAccount, recovered := api.tryRecoverFromRefreshRace(ctx, freshAccount); recovered {
|
||||
if isOAuthRefreshRequestPath(ctx) && recoveredAccount.Platform == PlatformGrok {
|
||||
if eligibilityErr := grokOAuthRequestAccountEligibilityError(recoveredAccount); eligibilityErr != nil {
|
||||
return nil, withGrokCredentialFailureSnapshot(eligibilityErr, recoveredAccount)
|
||||
}
|
||||
}
|
||||
slog.Info("oauth_refresh_race_recovered",
|
||||
"account_id", freshAccount.ID,
|
||||
"platform", freshAccount.Platform,
|
||||
@@ -141,7 +215,7 @@ func (api *OAuthRefreshAPI) RefreshIfNeeded(
|
||||
}, nil
|
||||
}
|
||||
}
|
||||
return nil, refreshErr
|
||||
return nil, withGrokCredentialFailureSnapshot(refreshErr, freshAccount)
|
||||
}
|
||||
|
||||
// 5. 设置版本号 + 更新 DB
|
||||
@@ -152,16 +226,30 @@ func (api *OAuthRefreshAPI) RefreshIfNeeded(
|
||||
"account_id", freshAccount.ID,
|
||||
"error", updateErr,
|
||||
)
|
||||
return nil, fmt.Errorf("oauth refresh succeeded but DB update failed: %w", updateErr)
|
||||
return nil, withGrokCredentialFailureSnapshot(
|
||||
fmt.Errorf("%w: %v", errOAuthRefreshCredentialPersist, updateErr), freshAccount,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
_ = lockAcquired // suppress unused warning when tokenCache is nil
|
||||
resultAccount := freshAccount
|
||||
if isOAuthRefreshRequestPath(ctx) && freshAccount.Platform == PlatformGrok {
|
||||
latestAccount, rereadErr := api.accountRepo.GetByID(ctx, freshAccount.ID)
|
||||
if rereadErr != nil {
|
||||
return nil, fmt.Errorf("%w: %v", errOAuthRefreshAccountRereadFailed, rereadErr)
|
||||
}
|
||||
if latestAccount == nil {
|
||||
return nil, fmt.Errorf("%w: account not found after refresh", errOAuthRefreshAccountStateChanged)
|
||||
}
|
||||
if eligibilityErr := grokOAuthRequestAccountEligibilityError(latestAccount); eligibilityErr != nil {
|
||||
return nil, withGrokCredentialFailureSnapshot(eligibilityErr, latestAccount)
|
||||
}
|
||||
resultAccount = latestAccount
|
||||
}
|
||||
|
||||
return &OAuthRefreshResult{
|
||||
Refreshed: true,
|
||||
NewCredentials: newCredentials,
|
||||
Account: freshAccount,
|
||||
Account: resultAccount,
|
||||
}, nil
|
||||
}
|
||||
|
||||
|
||||
@@ -51,13 +51,14 @@ func (r *refreshAPIAccountRepo) UpdateCredentials(_ context.Context, id int64, c
|
||||
|
||||
// refreshAPIExecutorStub implements OAuthRefreshExecutor for tests.
|
||||
type refreshAPIExecutorStub struct {
|
||||
needsRefresh bool
|
||||
credentials map[string]any
|
||||
err error
|
||||
refreshCalls int
|
||||
needsRefresh bool
|
||||
cannotRefresh bool
|
||||
credentials map[string]any
|
||||
err error
|
||||
refreshCalls int
|
||||
}
|
||||
|
||||
func (e *refreshAPIExecutorStub) CanRefresh(_ *Account) bool { return true }
|
||||
func (e *refreshAPIExecutorStub) CanRefresh(_ *Account) bool { return !e.cannotRefresh }
|
||||
|
||||
func (e *refreshAPIExecutorStub) NeedsRefresh(_ *Account, _ time.Duration) bool {
|
||||
return e.needsRefresh
|
||||
@@ -77,9 +78,10 @@ func (e *refreshAPIExecutorStub) CacheKey(account *Account) string {
|
||||
|
||||
// refreshAPICacheStub implements GeminiTokenCache for OAuthRefreshAPI tests.
|
||||
type refreshAPICacheStub struct {
|
||||
lockResult bool
|
||||
lockErr error
|
||||
releaseCalls int
|
||||
lockResult bool
|
||||
lockErr error
|
||||
releaseCalls int
|
||||
releaseCtxErr error
|
||||
}
|
||||
|
||||
func (c *refreshAPICacheStub) GetAccessToken(context.Context, string) (string, error) {
|
||||
@@ -96,15 +98,16 @@ func (c *refreshAPICacheStub) AcquireRefreshLock(context.Context, string, time.D
|
||||
return c.lockResult, c.lockErr
|
||||
}
|
||||
|
||||
func (c *refreshAPICacheStub) ReleaseRefreshLock(context.Context, string) error {
|
||||
func (c *refreshAPICacheStub) ReleaseRefreshLock(ctx context.Context, _ string) error {
|
||||
c.releaseCalls++
|
||||
c.releaseCtxErr = ctx.Err()
|
||||
return nil
|
||||
}
|
||||
|
||||
// ========== RefreshIfNeeded tests ==========
|
||||
|
||||
func TestRefreshIfNeeded_Success(t *testing.T) {
|
||||
account := &Account{ID: 1, Platform: PlatformAnthropic, Type: AccountTypeOAuth}
|
||||
account := &Account{ID: 1, Platform: PlatformAnthropic, Type: AccountTypeOAuth, Status: StatusActive}
|
||||
repo := &refreshAPIAccountRepo{account: account}
|
||||
cache := &refreshAPICacheStub{lockResult: true}
|
||||
executor := &refreshAPIExecutorStub{
|
||||
@@ -132,6 +135,7 @@ func TestRefreshIfNeeded_UpdateCredentialsPreservesRateLimitState(t *testing.T)
|
||||
ID: 11,
|
||||
Platform: PlatformGemini,
|
||||
Type: AccountTypeOAuth,
|
||||
Status: StatusActive,
|
||||
RateLimitResetAt: &resetAt,
|
||||
}
|
||||
repo := &refreshAPIAccountRepo{account: account}
|
||||
@@ -152,7 +156,7 @@ func TestRefreshIfNeeded_UpdateCredentialsPreservesRateLimitState(t *testing.T)
|
||||
}
|
||||
|
||||
func TestRefreshIfNeeded_LockHeld(t *testing.T) {
|
||||
account := &Account{ID: 2, Platform: PlatformAnthropic}
|
||||
account := &Account{ID: 2, Platform: PlatformAnthropic, Status: StatusActive}
|
||||
repo := &refreshAPIAccountRepo{account: account}
|
||||
cache := &refreshAPICacheStub{lockResult: false} // lock not acquired
|
||||
executor := &refreshAPIExecutorStub{needsRefresh: true}
|
||||
@@ -168,7 +172,7 @@ func TestRefreshIfNeeded_LockHeld(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestRefreshIfNeeded_LockErrorDegrades(t *testing.T) {
|
||||
account := &Account{ID: 3, Platform: PlatformGemini, Type: AccountTypeOAuth}
|
||||
account := &Account{ID: 3, Platform: PlatformGemini, Type: AccountTypeOAuth, Status: StatusActive}
|
||||
repo := &refreshAPIAccountRepo{account: account}
|
||||
cache := &refreshAPICacheStub{lockErr: errors.New("redis down")} // lock error
|
||||
executor := &refreshAPIExecutorStub{
|
||||
@@ -187,7 +191,7 @@ func TestRefreshIfNeeded_LockErrorDegrades(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestRefreshIfNeeded_NoCacheNoLock(t *testing.T) {
|
||||
account := &Account{ID: 4, Platform: PlatformGemini, Type: AccountTypeOAuth}
|
||||
account := &Account{ID: 4, Platform: PlatformGemini, Type: AccountTypeOAuth, Status: StatusActive}
|
||||
repo := &refreshAPIAccountRepo{account: account}
|
||||
executor := &refreshAPIExecutorStub{
|
||||
needsRefresh: true,
|
||||
@@ -203,7 +207,7 @@ func TestRefreshIfNeeded_NoCacheNoLock(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestRefreshIfNeeded_AlreadyRefreshed(t *testing.T) {
|
||||
account := &Account{ID: 5, Platform: PlatformAnthropic}
|
||||
account := &Account{ID: 5, Platform: PlatformAnthropic, Status: StatusActive}
|
||||
repo := &refreshAPIAccountRepo{account: account}
|
||||
cache := &refreshAPICacheStub{lockResult: true}
|
||||
executor := &refreshAPIExecutorStub{needsRefresh: false} // already refreshed
|
||||
@@ -220,7 +224,7 @@ func TestRefreshIfNeeded_AlreadyRefreshed(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestRefreshIfNeeded_RefreshError(t *testing.T) {
|
||||
account := &Account{ID: 6, Platform: PlatformAnthropic}
|
||||
account := &Account{ID: 6, Platform: PlatformAnthropic, Status: StatusActive}
|
||||
repo := &refreshAPIAccountRepo{account: account}
|
||||
cache := &refreshAPICacheStub{lockResult: true}
|
||||
executor := &refreshAPIExecutorStub{
|
||||
@@ -239,7 +243,7 @@ func TestRefreshIfNeeded_RefreshError(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestRefreshIfNeeded_DBUpdateError(t *testing.T) {
|
||||
account := &Account{ID: 7, Platform: PlatformGemini, Type: AccountTypeOAuth}
|
||||
account := &Account{ID: 7, Platform: PlatformGemini, Type: AccountTypeOAuth, Status: StatusActive}
|
||||
repo := &refreshAPIAccountRepo{
|
||||
account: account,
|
||||
updateErr: errors.New("db connection lost"),
|
||||
@@ -255,12 +259,12 @@ func TestRefreshIfNeeded_DBUpdateError(t *testing.T) {
|
||||
|
||||
require.Error(t, err)
|
||||
require.Nil(t, result)
|
||||
require.Contains(t, err.Error(), "DB update failed")
|
||||
require.ErrorIs(t, err, errOAuthRefreshCredentialPersist)
|
||||
require.Equal(t, 1, repo.updateCalls) // attempted
|
||||
}
|
||||
|
||||
func TestRefreshIfNeeded_DBRereadFails(t *testing.T) {
|
||||
account := &Account{ID: 8, Platform: PlatformAnthropic, Type: AccountTypeOAuth}
|
||||
account := &Account{ID: 8, Platform: PlatformAnthropic, Type: AccountTypeOAuth, Status: StatusActive}
|
||||
repo := &refreshAPIAccountRepo{
|
||||
account: nil, // GetByID returns nil
|
||||
getByIDErr: errors.New("db timeout"),
|
||||
@@ -274,13 +278,95 @@ func TestRefreshIfNeeded_DBRereadFails(t *testing.T) {
|
||||
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, executor.refreshCalls) // still refreshes using passed-in account
|
||||
require.ErrorContains(t, err, "oauth refresh account reread")
|
||||
require.Nil(t, result)
|
||||
require.Zero(t, executor.refreshCalls, "must not refresh with the stale caller snapshot")
|
||||
require.Zero(t, repo.updateCalls)
|
||||
require.Equal(t, 1, cache.releaseCalls)
|
||||
}
|
||||
|
||||
func TestRefreshIfNeeded_DBRereadNilFailsClosed(t *testing.T) {
|
||||
account := &Account{ID: 81, Platform: PlatformGrok, Type: AccountTypeOAuth, Status: StatusActive}
|
||||
repo := &refreshAPIAccountRepo{}
|
||||
cache := &refreshAPICacheStub{lockResult: true}
|
||||
executor := &refreshAPIExecutorStub{needsRefresh: true}
|
||||
|
||||
api := NewOAuthRefreshAPI(repo, cache)
|
||||
result, err := api.RefreshIfNeeded(context.Background(), account, executor, 3*time.Minute)
|
||||
|
||||
require.ErrorIs(t, err, errOAuthRefreshAccountStateChanged)
|
||||
require.Nil(t, result)
|
||||
require.Zero(t, executor.refreshCalls)
|
||||
require.Zero(t, repo.updateCalls)
|
||||
require.Equal(t, 1, cache.releaseCalls)
|
||||
}
|
||||
|
||||
func TestRefreshIfNeeded_DBRereadInactiveFailsClosed(t *testing.T) {
|
||||
account := &Account{ID: 82, Platform: PlatformGrok, Type: AccountTypeOAuth, Status: StatusActive}
|
||||
freshAccount := &Account{ID: account.ID, Platform: PlatformGrok, Type: AccountTypeOAuth, Status: StatusDisabled}
|
||||
repo := &refreshAPIAccountRepo{account: freshAccount}
|
||||
executor := &refreshAPIExecutorStub{needsRefresh: true}
|
||||
|
||||
api := NewOAuthRefreshAPI(repo, nil)
|
||||
result, err := api.RefreshIfNeeded(context.Background(), account, executor, 3*time.Minute)
|
||||
|
||||
require.ErrorContains(t, err, "account is not active")
|
||||
require.Nil(t, result)
|
||||
require.Zero(t, executor.refreshCalls)
|
||||
require.Zero(t, repo.updateCalls)
|
||||
}
|
||||
|
||||
func TestRefreshIfNeeded_DBRereadRevalidatesExecutorContract(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
freshPlatform string
|
||||
freshType string
|
||||
}{
|
||||
{name: "platform changed", freshPlatform: PlatformAnthropic, freshType: AccountTypeOAuth},
|
||||
{name: "type changed", freshPlatform: PlatformGrok, freshType: AccountTypeUpstream},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
account := &Account{ID: 83, Platform: PlatformGrok, Type: AccountTypeOAuth, Status: StatusActive}
|
||||
freshAccount := &Account{ID: account.ID, Platform: tt.freshPlatform, Type: tt.freshType, Status: StatusActive}
|
||||
repo := &refreshAPIAccountRepo{account: freshAccount}
|
||||
executor := NewGrokTokenRefresher(nil)
|
||||
|
||||
api := NewOAuthRefreshAPI(repo, nil)
|
||||
result, err := api.RefreshIfNeeded(context.Background(), account, executor, 3*time.Minute)
|
||||
|
||||
require.ErrorContains(t, err, "no longer refreshable")
|
||||
require.Nil(t, result)
|
||||
require.Zero(t, repo.updateCalls)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestRefreshIfNeeded_DBRereadMissingGrokRefreshCredentialReturnsPermanentSignal(t *testing.T) {
|
||||
account := &Account{
|
||||
ID: 84,
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Status: StatusActive,
|
||||
Credentials: map[string]any{
|
||||
"refresh_token": "caller-snapshot-refresh-token",
|
||||
},
|
||||
}
|
||||
freshAccount := &Account{ID: account.ID, Platform: PlatformGrok, Type: AccountTypeOAuth, Status: StatusActive}
|
||||
repo := &refreshAPIAccountRepo{account: freshAccount}
|
||||
executor := NewGrokTokenRefresher(nil)
|
||||
|
||||
api := NewOAuthRefreshAPI(repo, nil)
|
||||
result, err := api.RefreshIfNeeded(context.Background(), account, executor, 3*time.Minute)
|
||||
|
||||
require.ErrorIs(t, err, errGrokOAuthRefreshTokenMissing)
|
||||
require.Nil(t, result)
|
||||
require.Zero(t, repo.updateCalls)
|
||||
}
|
||||
|
||||
func TestRefreshIfNeeded_NilCredentials(t *testing.T) {
|
||||
account := &Account{ID: 9, Platform: PlatformGemini, Type: AccountTypeOAuth}
|
||||
account := &Account{ID: 9, Platform: PlatformGemini, Type: AccountTypeOAuth, Status: StatusActive}
|
||||
repo := &refreshAPIAccountRepo{account: account}
|
||||
cache := &refreshAPICacheStub{lockResult: true}
|
||||
executor := &refreshAPIExecutorStub{
|
||||
@@ -413,6 +499,7 @@ func TestRefreshIfNeeded_InvalidGrantRaceRecovered(t *testing.T) {
|
||||
ID: 10,
|
||||
Platform: PlatformAnthropic,
|
||||
Type: AccountTypeOAuth,
|
||||
Status: StatusActive,
|
||||
Credentials: map[string]any{"refresh_token": "old-rt", "access_token": "old-at"},
|
||||
}
|
||||
// After race, DB has new refresh token from another worker
|
||||
@@ -420,6 +507,7 @@ func TestRefreshIfNeeded_InvalidGrantRaceRecovered(t *testing.T) {
|
||||
ID: 10,
|
||||
Platform: PlatformAnthropic,
|
||||
Type: AccountTypeOAuth,
|
||||
Status: StatusActive,
|
||||
Credentials: map[string]any{"refresh_token": "new-rt", "access_token": "new-at"},
|
||||
}
|
||||
repo := &refreshAPIAccountRepoWithRace{
|
||||
@@ -449,6 +537,7 @@ func TestRefreshIfNeeded_InvalidGrantGenuine(t *testing.T) {
|
||||
ID: 11,
|
||||
Platform: PlatformAnthropic,
|
||||
Type: AccountTypeOAuth,
|
||||
Status: StatusActive,
|
||||
Credentials: map[string]any{"refresh_token": "revoked-rt", "access_token": "old-at"},
|
||||
}
|
||||
repo := &refreshAPIAccountRepoWithRace{
|
||||
@@ -474,6 +563,7 @@ func TestRefreshIfNeeded_InvalidGrantDBRereadFailsOnRecovery(t *testing.T) {
|
||||
ID: 12,
|
||||
Platform: PlatformAnthropic,
|
||||
Type: AccountTypeOAuth,
|
||||
Status: StatusActive,
|
||||
Credentials: map[string]any{"refresh_token": "old-rt"},
|
||||
}
|
||||
repo := &refreshAPIAccountRepoWithRace{
|
||||
@@ -500,6 +590,7 @@ func TestRefreshIfNeeded_LocalMutexSerializesConcurrent(t *testing.T) {
|
||||
ID: 20,
|
||||
Platform: PlatformAnthropic,
|
||||
Type: AccountTypeOAuth,
|
||||
Status: StatusActive,
|
||||
Credentials: map[string]any{"refresh_token": "new-rt", "access_token": "new-at"},
|
||||
}
|
||||
callCount := 0
|
||||
@@ -507,6 +598,7 @@ func TestRefreshIfNeeded_LocalMutexSerializesConcurrent(t *testing.T) {
|
||||
ID: 20,
|
||||
Platform: PlatformAnthropic,
|
||||
Type: AccountTypeOAuth,
|
||||
Status: StatusActive,
|
||||
Credentials: map[string]any{"refresh_token": "old-rt"},
|
||||
}}
|
||||
|
||||
@@ -556,6 +648,78 @@ func TestRefreshIfNeeded_LocalMutexSerializesConcurrent(t *testing.T) {
|
||||
mu.Unlock()
|
||||
}
|
||||
|
||||
func TestRefreshIfNeeded_LocalLockWaitHonorsContextCancellation(t *testing.T) {
|
||||
account := &Account{ID: 21, Platform: PlatformGrok, Type: AccountTypeOAuth, Status: StatusActive}
|
||||
repo := &refreshAPIAccountRepo{account: account}
|
||||
refreshStarted := make(chan struct{})
|
||||
releaseRefresh := make(chan struct{})
|
||||
var once sync.Once
|
||||
executor := &dynamicRefreshExecutor{
|
||||
canRefresh: true,
|
||||
cacheKey: "test:context-lock:grok",
|
||||
needsRefreshFunc: func() bool { return true },
|
||||
refreshFunc: func(context.Context, *Account) (map[string]any, error) {
|
||||
once.Do(func() { close(refreshStarted) })
|
||||
<-releaseRefresh
|
||||
return map[string]any{"access_token": "new-at"}, nil
|
||||
},
|
||||
}
|
||||
api := NewOAuthRefreshAPI(repo, nil)
|
||||
firstDone := make(chan error, 1)
|
||||
go func() {
|
||||
_, err := api.RefreshIfNeeded(context.Background(), account, executor, 3*time.Minute)
|
||||
firstDone <- err
|
||||
}()
|
||||
<-refreshStarted
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond)
|
||||
defer cancel()
|
||||
startedAt := time.Now()
|
||||
result, err := api.RefreshIfNeeded(ctx, account, executor, 3*time.Minute)
|
||||
|
||||
require.Nil(t, result)
|
||||
require.ErrorIs(t, err, context.DeadlineExceeded)
|
||||
require.Less(t, time.Since(startedAt), 500*time.Millisecond)
|
||||
close(releaseRefresh)
|
||||
require.NoError(t, <-firstDone)
|
||||
}
|
||||
|
||||
func TestRefreshIfNeeded_ReleasesDistributedLockWithCleanupContext(t *testing.T) {
|
||||
account := &Account{
|
||||
ID: 22,
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Status: StatusActive,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "old-access",
|
||||
"refresh_token": "old-refresh",
|
||||
},
|
||||
}
|
||||
repo := &refreshAPIAccountRepo{account: account}
|
||||
cache := &refreshAPICacheStub{lockResult: true}
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
executor := &dynamicRefreshExecutor{
|
||||
canRefresh: true,
|
||||
cacheKey: "test:cleanup:grok",
|
||||
needsRefreshFunc: func() bool { return true },
|
||||
refreshFunc: func(context.Context, *Account) (map[string]any, error) {
|
||||
cancel()
|
||||
return map[string]any{"access_token": "new-at"}, nil
|
||||
},
|
||||
}
|
||||
api := NewOAuthRefreshAPI(repo, cache)
|
||||
|
||||
result, err := api.RefreshIfNeeded(ctx, account, executor, 3*time.Minute)
|
||||
|
||||
require.ErrorIs(t, err, context.Canceled)
|
||||
require.Nil(t, result)
|
||||
require.Zero(t, repo.updateCalls)
|
||||
require.Equal(t, "old-access", account.GetGrokAccessToken())
|
||||
require.Zero(t, account.GetCredentialAsInt64("_token_version"))
|
||||
require.Equal(t, 1, cache.releaseCalls)
|
||||
require.NoError(t, cache.releaseCtxErr)
|
||||
}
|
||||
|
||||
// dynamicRefreshExecutor is a test helper with function-based NeedsRefresh and Refresh.
|
||||
type dynamicRefreshExecutor struct {
|
||||
canRefresh bool
|
||||
|
||||
@@ -3,6 +3,7 @@ package service
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
@@ -15,6 +16,13 @@ const (
|
||||
openAIOAuth429StormMaxAccountSwitches = 1
|
||||
)
|
||||
|
||||
// OpenAIOAuth429FailoverState tracks the request-local follow-up budget after
|
||||
// the first Grok OAuth 429. Once that 429 occurs, exactly one different account
|
||||
// may be attempted; any failure from that follow-up account ends failover.
|
||||
type OpenAIOAuth429FailoverState struct {
|
||||
grokOAuth429FollowupPending bool
|
||||
}
|
||||
|
||||
func openAIAccountStateContext(ctx context.Context) (context.Context, context.CancelFunc) {
|
||||
base := context.Background()
|
||||
if ctx != nil {
|
||||
@@ -96,6 +104,25 @@ func (s *OpenAIGatewayService) BlockAccountScheduling(account *Account, until ti
|
||||
if s == nil || !isOpenAIAccount(account) {
|
||||
return
|
||||
}
|
||||
mu := s.openAIAccountRuntimeBlockLock(account.ID)
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
_, _ = s.blockAccountSchedulingLocked(account, until, reason)
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) openAIAccountRuntimeBlockLock(accountID int64) *sync.Mutex {
|
||||
actual, _ := s.openaiAccountRuntimeBlockLocks.LoadOrStore(accountID, &sync.Mutex{})
|
||||
mu, ok := actual.(*sync.Mutex)
|
||||
if !ok {
|
||||
mu = &sync.Mutex{}
|
||||
s.openaiAccountRuntimeBlockLocks.Store(accountID, mu)
|
||||
}
|
||||
return mu
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) blockAccountSchedulingLocked(account *Account, until time.Time, _ string) (uint64, bool) {
|
||||
generation := s.openaiAccountRuntimeBlockSequence.Add(1)
|
||||
s.openaiAccountRuntimeBlockGeneration.Store(account.ID, generation)
|
||||
now := time.Now()
|
||||
blockUntil := until
|
||||
if blockUntil.IsZero() || !blockUntil.After(now) {
|
||||
@@ -107,7 +134,7 @@ func (s *OpenAIGatewayService) BlockAccountScheduling(account *Account, until ti
|
||||
if !loaded {
|
||||
actual, stored := s.openaiAccountRuntimeBlockUntil.LoadOrStore(account.ID, blockUntil)
|
||||
if !stored {
|
||||
return
|
||||
return generation, true
|
||||
}
|
||||
current = actual
|
||||
}
|
||||
@@ -115,15 +142,15 @@ func (s *OpenAIGatewayService) BlockAccountScheduling(account *Account, until ti
|
||||
currentUntil, ok := current.(time.Time)
|
||||
if !ok || currentUntil.IsZero() {
|
||||
if s.openaiAccountRuntimeBlockUntil.CompareAndSwap(account.ID, current, blockUntil) {
|
||||
return
|
||||
return generation, true
|
||||
}
|
||||
continue
|
||||
}
|
||||
if currentUntil.After(blockUntil) {
|
||||
return
|
||||
if !blockUntil.After(currentUntil) {
|
||||
return generation, false
|
||||
}
|
||||
if s.openaiAccountRuntimeBlockUntil.CompareAndSwap(account.ID, current, blockUntil) {
|
||||
return
|
||||
return generation, true
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -132,13 +159,20 @@ func (s *OpenAIGatewayService) ClearAccountSchedulingBlock(accountID int64) {
|
||||
if s == nil || accountID <= 0 {
|
||||
return
|
||||
}
|
||||
mu := s.openAIAccountRuntimeBlockLock(accountID)
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
s.openaiAccountRuntimeBlockUntil.Delete(accountID)
|
||||
s.openaiAccountRuntimeBlockGeneration.Store(accountID, s.openaiAccountRuntimeBlockSequence.Add(1))
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) isOpenAIAccountRuntimeBlocked(account *Account) bool {
|
||||
if s == nil || !isOpenAIAccount(account) {
|
||||
return false
|
||||
}
|
||||
mu := s.openAIAccountRuntimeBlockLock(account.ID)
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
value, ok := s.openaiAccountRuntimeBlockUntil.Load(account.ID)
|
||||
if !ok {
|
||||
return false
|
||||
@@ -146,12 +180,14 @@ func (s *OpenAIGatewayService) isOpenAIAccountRuntimeBlocked(account *Account) b
|
||||
cooldownUntil, ok := value.(time.Time)
|
||||
if !ok || cooldownUntil.IsZero() {
|
||||
s.openaiAccountRuntimeBlockUntil.Delete(account.ID)
|
||||
s.openaiAccountRuntimeBlockGeneration.Store(account.ID, s.openaiAccountRuntimeBlockSequence.Add(1))
|
||||
return false
|
||||
}
|
||||
if time.Now().Before(cooldownUntil) {
|
||||
return true
|
||||
}
|
||||
s.openaiAccountRuntimeBlockUntil.Delete(account.ID)
|
||||
s.openaiAccountRuntimeBlockGeneration.Store(account.ID, s.openaiAccountRuntimeBlockSequence.Add(1))
|
||||
return false
|
||||
}
|
||||
|
||||
@@ -181,14 +217,28 @@ func (s *OpenAIGatewayService) isOpenAIOAuth429Storm() bool {
|
||||
return s.openaiOAuth429WindowCount.Load() >= openAIOAuth429StormThreshold
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) ShouldStopOpenAIOAuth429Failover(account *Account, statusCode int, failedSwitches int) bool {
|
||||
if statusCode != http.StatusTooManyRequests || failedSwitches < openAIOAuth429StormMaxAccountSwitches {
|
||||
func (s *OpenAIGatewayService) ShouldStopOpenAIOAuth429Failover(account *Account, statusCode int, failedSwitches int, state *OpenAIOAuth429FailoverState) bool {
|
||||
if failedSwitches < openAIOAuth429StormMaxAccountSwitches {
|
||||
return false
|
||||
}
|
||||
if isGrokOAuthAccount(account) {
|
||||
if state != nil && state.grokOAuth429FollowupPending {
|
||||
// The follow-up budget was armed by a Grok OAuth 429. Consume it on
|
||||
// any failing follow-up account, even if a mixed pool selected an API-key
|
||||
// account next.
|
||||
return true
|
||||
}
|
||||
if !isOpenAIOAuthAccount(account) {
|
||||
if isGrokOAuthAccount(account) {
|
||||
if state == nil {
|
||||
// Preserve the old threshold for callers that have not adopted the
|
||||
// request-local state contract yet.
|
||||
return statusCode == http.StatusTooManyRequests && failedSwitches >= 2
|
||||
}
|
||||
if statusCode == http.StatusTooManyRequests {
|
||||
state.grokOAuth429FollowupPending = true
|
||||
}
|
||||
return false
|
||||
}
|
||||
if statusCode != http.StatusTooManyRequests || !isOpenAIOAuthAccount(account) {
|
||||
return false
|
||||
}
|
||||
return s.isOpenAIOAuth429Storm()
|
||||
|
||||
@@ -135,26 +135,45 @@ func TestShouldStopOpenAIOAuth429Failover_OnlyDuringStorm(t *testing.T) {
|
||||
svc := &OpenAIGatewayService{}
|
||||
account := &Account{ID: 42, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
|
||||
apiKeyAccount := &Account{ID: 43, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
|
||||
var state OpenAIOAuth429FailoverState
|
||||
|
||||
require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 1))
|
||||
require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 1, &state))
|
||||
|
||||
for i := 0; i < openAIOAuth429StormThreshold; i++ {
|
||||
svc.recordOpenAIOAuth429()
|
||||
}
|
||||
|
||||
require.True(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 1))
|
||||
require.False(t, svc.ShouldStopOpenAIOAuth429Failover(apiKeyAccount, http.StatusTooManyRequests, 1))
|
||||
require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusInternalServerError, 1))
|
||||
require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 0))
|
||||
require.True(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 1, &state))
|
||||
require.False(t, svc.ShouldStopOpenAIOAuth429Failover(apiKeyAccount, http.StatusTooManyRequests, 1, &state))
|
||||
require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusInternalServerError, 1, &state))
|
||||
require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 0, &state))
|
||||
}
|
||||
|
||||
func TestShouldStopOpenAIOAuth429Failover_StopsGrokAfterFirst429Switch(t *testing.T) {
|
||||
func TestShouldStopOpenAIOAuth429Failover_TracksOneGrokFollowupAttempt(t *testing.T) {
|
||||
svc := &OpenAIGatewayService{}
|
||||
account := &Account{ID: 44, Platform: PlatformGrok, Type: AccountTypeOAuth}
|
||||
apiKeyAccount := &Account{ID: 45, Platform: PlatformGrok, Type: AccountTypeAPIKey}
|
||||
|
||||
require.True(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 1))
|
||||
require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 0))
|
||||
require.False(t, svc.ShouldStopOpenAIOAuth429Failover(apiKeyAccount, http.StatusTooManyRequests, 1))
|
||||
require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusInternalServerError, 1))
|
||||
t.Run("429 then 500 stops after one followup", func(t *testing.T) {
|
||||
var state OpenAIOAuth429FailoverState
|
||||
require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 1, &state))
|
||||
require.True(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusInternalServerError, 2, &state))
|
||||
})
|
||||
|
||||
t.Run("500 then 429 still allows one followup", func(t *testing.T) {
|
||||
var state OpenAIOAuth429FailoverState
|
||||
require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusInternalServerError, 1, &state))
|
||||
require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 2, &state))
|
||||
require.True(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusBadGateway, 3, &state))
|
||||
})
|
||||
|
||||
t.Run("OAuth 429 then API-key failure consumes the same followup", func(t *testing.T) {
|
||||
var state OpenAIOAuth429FailoverState
|
||||
require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 1, &state))
|
||||
require.True(t, svc.ShouldStopOpenAIOAuth429Failover(apiKeyAccount, http.StatusInternalServerError, 2, &state))
|
||||
})
|
||||
|
||||
var state OpenAIOAuth429FailoverState
|
||||
require.False(t, svc.ShouldStopOpenAIOAuth429Failover(account, http.StatusTooManyRequests, 0, &state))
|
||||
require.False(t, svc.ShouldStopOpenAIOAuth429Failover(apiKeyAccount, http.StatusTooManyRequests, 2, &state))
|
||||
}
|
||||
|
||||
@@ -0,0 +1,522 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto"
|
||||
"crypto/ed25519"
|
||||
"crypto/sha512"
|
||||
"crypto/x509"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/httpclient"
|
||||
"golang.org/x/crypto/curve25519"
|
||||
"golang.org/x/crypto/nacl/box"
|
||||
)
|
||||
|
||||
const (
|
||||
OpenAIAuthModeAgentIdentity = "agentIdentity"
|
||||
agentIdentityAuthAPIBaseURL = "https://auth.openai.com/api/accounts"
|
||||
agentIdentityTaskRegistrationTimeout = 30 * time.Second
|
||||
)
|
||||
|
||||
var openAIAgentIdentityAuthAPIBaseURL = agentIdentityAuthAPIBaseURL
|
||||
|
||||
var agentIdentityTaskLocks sync.Map // map[int64]*sync.Mutex
|
||||
|
||||
type agentIdentityWSConnectionInvalidator interface {
|
||||
InvalidateAgentIdentityWSConnections(accountID int64)
|
||||
}
|
||||
|
||||
type agentIdentityKey struct {
|
||||
runtimeID string
|
||||
privateKey ed25519.PrivateKey
|
||||
taskID string
|
||||
}
|
||||
|
||||
type agentIdentityTaskRegistrationResponse struct {
|
||||
TaskID string `json:"task_id"`
|
||||
TaskIDCamel string `json:"taskId"`
|
||||
EncryptedTaskID string `json:"encrypted_task_id"`
|
||||
EncryptedTaskIDCamel string `json:"encryptedTaskId"`
|
||||
}
|
||||
|
||||
type agentIdentityTaskRecoveredError struct{}
|
||||
|
||||
func (e *agentIdentityTaskRecoveredError) Error() string {
|
||||
return "agent identity task recovered"
|
||||
}
|
||||
|
||||
func (a *Account) IsOpenAIAgentIdentity() bool {
|
||||
if a == nil || !a.IsOpenAIOAuth() {
|
||||
return false
|
||||
}
|
||||
return strings.EqualFold(strings.TrimSpace(a.GetCredential(openAIAuthModeCredentialKey)), OpenAIAuthModeAgentIdentity)
|
||||
}
|
||||
|
||||
func agentIdentityPrivateKey(account *Account) (ed25519.PrivateKey, error) {
|
||||
if account == nil {
|
||||
return nil, errors.New("agent identity account is nil")
|
||||
}
|
||||
raw := strings.TrimSpace(account.GetCredential("agent_private_key"))
|
||||
if raw == "" {
|
||||
return nil, errors.New("agent identity private key is missing")
|
||||
}
|
||||
der, err := base64.StdEncoding.DecodeString(raw)
|
||||
if err != nil {
|
||||
return nil, errors.New("agent identity private key is not valid base64")
|
||||
}
|
||||
key, err := x509.ParsePKCS8PrivateKey(der)
|
||||
if err != nil {
|
||||
return nil, errors.New("agent identity private key is not valid PKCS#8")
|
||||
}
|
||||
privateKey, ok := key.(ed25519.PrivateKey)
|
||||
if !ok || len(privateKey) != ed25519.PrivateKeySize {
|
||||
return nil, errors.New("agent identity private key is not Ed25519")
|
||||
}
|
||||
return privateKey, nil
|
||||
}
|
||||
|
||||
// ValidateOpenAIAgentIdentityPrivateKey validates the stored PKCS#8 Ed25519
|
||||
// form without returning or logging the key material.
|
||||
func ValidateOpenAIAgentIdentityPrivateKey(encoded string) error {
|
||||
account := &Account{Credentials: map[string]any{"agent_private_key": encoded}}
|
||||
_, err := agentIdentityPrivateKey(account)
|
||||
return err
|
||||
}
|
||||
|
||||
func agentIdentityKeyFromAccount(account *Account) (agentIdentityKey, error) {
|
||||
privateKey, err := agentIdentityPrivateKey(account)
|
||||
if err != nil {
|
||||
return agentIdentityKey{}, err
|
||||
}
|
||||
runtimeID := strings.TrimSpace(account.GetCredential("agent_runtime_id"))
|
||||
if runtimeID == "" {
|
||||
return agentIdentityKey{}, errors.New("agent identity runtime id is missing")
|
||||
}
|
||||
return agentIdentityKey{
|
||||
runtimeID: runtimeID,
|
||||
privateKey: privateKey,
|
||||
taskID: strings.TrimSpace(account.GetCredential("task_id")),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func buildAgentAssertion(key agentIdentityKey, now time.Time) (string, error) {
|
||||
if key.runtimeID == "" || key.taskID == "" {
|
||||
return "", errors.New("agent identity runtime or task id is missing")
|
||||
}
|
||||
timestamp := now.UTC().Format(time.RFC3339)
|
||||
payload := []byte(key.runtimeID + ":" + key.taskID + ":" + timestamp)
|
||||
signature, err := key.privateKey.Sign(nil, payload, crypto.Hash(0))
|
||||
if err != nil {
|
||||
return "", errors.New("failed to sign agent assertion")
|
||||
}
|
||||
envelope := map[string]string{
|
||||
"agent_runtime_id": key.runtimeID,
|
||||
"task_id": key.taskID,
|
||||
"timestamp": timestamp,
|
||||
"signature": base64.StdEncoding.EncodeToString(signature),
|
||||
}
|
||||
encoded, err := json.Marshal(envelope)
|
||||
if err != nil {
|
||||
return "", errors.New("failed to serialize agent assertion")
|
||||
}
|
||||
return "AgentAssertion " + base64.RawURLEncoding.EncodeToString(encoded), nil
|
||||
}
|
||||
|
||||
func signAgentTaskRegistration(key agentIdentityKey, timestamp time.Time) (string, string, error) {
|
||||
if key.runtimeID == "" {
|
||||
return "", "", errors.New("agent identity runtime id is missing")
|
||||
}
|
||||
formatted := timestamp.UTC().Format(time.RFC3339)
|
||||
signature, err := key.privateKey.Sign(nil, []byte(key.runtimeID+":"+formatted), crypto.Hash(0))
|
||||
if err != nil {
|
||||
return "", "", errors.New("failed to sign agent task registration")
|
||||
}
|
||||
return formatted, base64.StdEncoding.EncodeToString(signature), nil
|
||||
}
|
||||
|
||||
func decryptAgentTaskID(key agentIdentityKey, encoded string) (string, error) {
|
||||
ciphertext, err := base64.StdEncoding.DecodeString(strings.TrimSpace(encoded))
|
||||
if err != nil {
|
||||
return "", errors.New("encrypted agent task id is not valid base64")
|
||||
}
|
||||
seed := key.privateKey.Seed()
|
||||
digest := sha512.Sum512(seed)
|
||||
var curvePrivate [32]byte
|
||||
copy(curvePrivate[:], digest[:32])
|
||||
curvePrivate[0] &= 248
|
||||
curvePrivate[31] &= 127
|
||||
curvePrivate[31] |= 64
|
||||
curvePublicBytes, err := curve25519.X25519(curvePrivate[:], curve25519.Basepoint)
|
||||
if err != nil {
|
||||
return "", errors.New("failed to derive agent identity decryption key")
|
||||
}
|
||||
var curvePublic [32]byte
|
||||
copy(curvePublic[:], curvePublicBytes)
|
||||
plaintext, ok := box.OpenAnonymous(nil, ciphertext, &curvePublic, &curvePrivate)
|
||||
if !ok {
|
||||
return "", errors.New("failed to decrypt encrypted agent task id")
|
||||
}
|
||||
taskID := strings.TrimSpace(string(plaintext))
|
||||
if taskID == "" {
|
||||
return "", errors.New("decrypted agent task id is empty")
|
||||
}
|
||||
return taskID, nil
|
||||
}
|
||||
|
||||
func registerAgentIdentityTask(ctx context.Context, account *Account) (string, error) {
|
||||
key, err := agentIdentityKeyFromAccount(account)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
timestamp, signature, err := signAgentTaskRegistration(key, time.Now())
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
proxyURL := ""
|
||||
if account.ProxyID != nil && account.Proxy != nil {
|
||||
proxyURL = account.Proxy.URL()
|
||||
}
|
||||
client, err := httpclient.GetClient(httpclient.Options{
|
||||
ProxyURL: proxyURL,
|
||||
Timeout: agentIdentityTaskRegistrationTimeout,
|
||||
ResponseHeaderTimeout: 15 * time.Second,
|
||||
})
|
||||
if err != nil {
|
||||
return "", errors.New("invalid proxy configuration for agent task registration")
|
||||
}
|
||||
body, err := json.Marshal(map[string]string{
|
||||
"timestamp": timestamp,
|
||||
"signature": signature,
|
||||
})
|
||||
if err != nil {
|
||||
return "", errors.New("failed to serialize agent task registration")
|
||||
}
|
||||
url := strings.TrimRight(strings.TrimSpace(openAIAgentIdentityAuthAPIBaseURL), "/") + "/v1/agent/" + key.runtimeID + "/task/register"
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, strings.NewReader(string(body)))
|
||||
if err != nil {
|
||||
return "", errors.New("failed to build agent task registration request")
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Accept", "application/json")
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return "", errors.New("agent task registration request failed")
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
|
||||
return "", fmt.Errorf("agent task registration returned status %d", resp.StatusCode)
|
||||
}
|
||||
var result agentIdentityTaskRegistrationResponse
|
||||
if err := json.NewDecoder(io.LimitReader(resp.Body, 64*1024)).Decode(&result); err != nil {
|
||||
return "", errors.New("agent task registration response is invalid")
|
||||
}
|
||||
if taskID := strings.TrimSpace(result.TaskID); taskID != "" {
|
||||
return taskID, nil
|
||||
}
|
||||
if taskID := strings.TrimSpace(result.TaskIDCamel); taskID != "" {
|
||||
return taskID, nil
|
||||
}
|
||||
encrypted := strings.TrimSpace(result.EncryptedTaskID)
|
||||
if encrypted == "" {
|
||||
encrypted = strings.TrimSpace(result.EncryptedTaskIDCamel)
|
||||
}
|
||||
if encrypted == "" {
|
||||
return "", errors.New("agent task registration response omitted task id")
|
||||
}
|
||||
return decryptAgentTaskID(key, encrypted)
|
||||
}
|
||||
|
||||
func ensureAgentIdentityTaskForAccount(ctx context.Context, repo AccountRepository, wsInvalidator agentIdentityWSConnectionInvalidator, taskMu *sync.Mutex, account *Account, expectedTaskID string) error {
|
||||
if account == nil || !account.IsOpenAIAgentIdentity() {
|
||||
return nil
|
||||
}
|
||||
credAccount := account
|
||||
if account.IsShadow() {
|
||||
resolved, err := resolveCredentialAccount(ctx, repo, account)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
credAccount = resolved
|
||||
}
|
||||
if credAccount == nil || !credAccount.IsOpenAIAgentIdentity() {
|
||||
return errors.New("agent identity credentials are unavailable")
|
||||
}
|
||||
currentTaskID := strings.TrimSpace(credAccount.GetCredential("task_id"))
|
||||
if currentTaskID != "" && (expectedTaskID == "" || currentTaskID != expectedTaskID) {
|
||||
return nil
|
||||
}
|
||||
if taskMu == nil {
|
||||
return errors.New("agent identity task lock is unavailable")
|
||||
}
|
||||
sharedTaskMu := taskMu
|
||||
if credAccount.ID > 0 {
|
||||
candidate := &sync.Mutex{}
|
||||
actual, _ := agentIdentityTaskLocks.LoadOrStore(credAccount.ID, candidate)
|
||||
loadedTaskMu, ok := actual.(*sync.Mutex)
|
||||
if !ok {
|
||||
return errors.New("agent identity task lock has invalid type")
|
||||
}
|
||||
sharedTaskMu = loadedTaskMu
|
||||
}
|
||||
sharedTaskMu.Lock()
|
||||
defer sharedTaskMu.Unlock()
|
||||
// Re-read inside the shared lock. Different request paths often receive
|
||||
// independent repository snapshots; checking only the caller's snapshot
|
||||
// would allow sequential duplicate registrations after the first writer
|
||||
// has already persisted a new task.
|
||||
if repo != nil && credAccount.ID > 0 {
|
||||
if refreshed, refreshErr := repo.GetByID(ctx, credAccount.ID); refreshErr == nil && refreshed != nil {
|
||||
if refreshed.IsShadow() {
|
||||
if resolved, resolveErr := resolveCredentialAccount(ctx, repo, refreshed); resolveErr == nil && resolved != nil {
|
||||
refreshed = resolved
|
||||
}
|
||||
}
|
||||
if refreshed.IsOpenAIAgentIdentity() {
|
||||
credAccount = refreshed
|
||||
if !account.IsShadow() {
|
||||
account.Credentials = shallowCopyMap(credAccount.Credentials)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
currentTaskID = strings.TrimSpace(credAccount.GetCredential("task_id"))
|
||||
if currentTaskID != "" && (expectedTaskID == "" || currentTaskID != expectedTaskID) {
|
||||
return nil
|
||||
}
|
||||
newTaskID, err := registerAgentIdentityTask(ctx, credAccount)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
credentials := make(map[string]any, len(credAccount.Credentials)+1)
|
||||
for key, value := range credAccount.Credentials {
|
||||
credentials[key] = value
|
||||
}
|
||||
credentials["task_id"] = newTaskID
|
||||
if err := persistAccountCredentials(ctx, repo, credAccount, credentials); err != nil {
|
||||
return err
|
||||
}
|
||||
if !account.IsShadow() && account != credAccount {
|
||||
account.Credentials = shallowCopyMap(credAccount.Credentials)
|
||||
}
|
||||
if wsInvalidator != nil {
|
||||
wsInvalidator.InvalidateAgentIdentityWSConnections(credAccount.ID)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) ensureAgentIdentityTask(ctx context.Context, account *Account, expectedTaskID string) error {
|
||||
if s == nil {
|
||||
return errors.New("openai gateway service is nil")
|
||||
}
|
||||
return ensureAgentIdentityTaskForAccount(ctx, s.accountRepo, s, &s.agentIdentityTaskMu, account, expectedTaskID)
|
||||
}
|
||||
|
||||
func isAgentIdentityTaskInvalidHTTPResponse(statusCode int, body []byte) bool {
|
||||
if statusCode != http.StatusUnauthorized {
|
||||
return false
|
||||
}
|
||||
lower := strings.ToLower(string(body))
|
||||
compact := strings.NewReplacer(" ", "", "\t", "", "\r", "", "\n", "").Replace(lower)
|
||||
for _, marker := range []string{
|
||||
`"code":"invalid_task_id"`,
|
||||
`"code":"task_not_found"`,
|
||||
`"code":"task_expired"`,
|
||||
`"error":"invalid_task_id"`,
|
||||
} {
|
||||
if strings.Contains(compact, marker) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
for _, marker := range []string{
|
||||
"invalid task_id",
|
||||
"invalid task id",
|
||||
"task_id is invalid",
|
||||
"task id is invalid",
|
||||
"task not found",
|
||||
"task expired",
|
||||
"unknown task_id",
|
||||
"unknown task id",
|
||||
} {
|
||||
if strings.Contains(lower, marker) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
type agentIdentityTaskRecoveryContextKey struct{}
|
||||
|
||||
func markAgentIdentityTaskRecoveryTried(ctx context.Context) context.Context {
|
||||
return context.WithValue(ctx, agentIdentityTaskRecoveryContextKey{}, true)
|
||||
}
|
||||
|
||||
func agentIdentityTaskRecoveryWasTried(ctx context.Context) bool {
|
||||
tried, _ := ctx.Value(agentIdentityTaskRecoveryContextKey{}).(bool)
|
||||
return tried
|
||||
}
|
||||
|
||||
func isAgentIdentityTaskInvalidWSDialError(err *openAIWSDialError) bool {
|
||||
return err != nil && isAgentIdentityTaskInvalidHTTPResponse(err.StatusCode, err.ResponseBody)
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) buildOpenAIAuthenticationHeaders(ctx context.Context, account *Account, token string) (http.Header, error) {
|
||||
if account == nil {
|
||||
return nil, errors.New("account is nil")
|
||||
}
|
||||
credAccount := account
|
||||
if account.IsShadow() {
|
||||
resolved, err := resolveCredentialAccount(ctx, s.accountRepo, account)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
credAccount = resolved
|
||||
}
|
||||
headers := make(http.Header)
|
||||
if credAccount != nil && credAccount.IsOpenAIAgentIdentity() {
|
||||
agentHeaders, err := buildAgentIdentityAuthenticationHeaders(ctx, s.accountRepo, s, &s.agentIdentityTaskMu, credAccount)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return agentHeaders, nil
|
||||
}
|
||||
headers.Set("Authorization", "Bearer "+token)
|
||||
return headers, nil
|
||||
}
|
||||
|
||||
func buildAgentIdentityAuthenticationHeaders(ctx context.Context, repo AccountRepository, wsInvalidator agentIdentityWSConnectionInvalidator, taskMu *sync.Mutex, account *Account) (http.Header, error) {
|
||||
if account == nil || !account.IsOpenAIAgentIdentity() {
|
||||
return nil, errors.New("agent identity account is required")
|
||||
}
|
||||
if err := ensureAgentIdentityTaskForAccount(ctx, repo, wsInvalidator, taskMu, account, ""); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
key, err := agentIdentityKeyFromAccount(account)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
assertion, err := buildAgentAssertion(key, time.Now())
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
headers := make(http.Header)
|
||||
headers.Set("Authorization", assertion)
|
||||
return headers, nil
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) refreshOpenAIAgentIdentityHeaders(ctx context.Context, account *Account, headers http.Header) (http.Header, error) {
|
||||
if account == nil {
|
||||
return cloneHeader(headers), nil
|
||||
}
|
||||
credAccount := account
|
||||
if account.IsShadow() {
|
||||
resolved, err := resolveCredentialAccount(ctx, s.accountRepo, account)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
credAccount = resolved
|
||||
}
|
||||
if !credAccount.IsOpenAIAgentIdentity() {
|
||||
return cloneHeader(headers), nil
|
||||
}
|
||||
refreshed := cloneHeader(headers)
|
||||
if refreshed == nil {
|
||||
refreshed = make(http.Header)
|
||||
}
|
||||
authHeaders, err := buildAgentIdentityAuthenticationHeaders(ctx, s.accountRepo, s, &s.agentIdentityTaskMu, credAccount)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
refreshed.Set("Authorization", authHeaders.Get("Authorization"))
|
||||
return refreshed, nil
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) recoverAgentIdentityTask(ctx context.Context, account *Account, expectedTaskID string) error {
|
||||
if account != nil && account.IsShadow() {
|
||||
if resolved, err := resolveCredentialAccount(ctx, s.accountRepo, account); err == nil && resolved != nil && strings.TrimSpace(expectedTaskID) == "" {
|
||||
expectedTaskID = strings.TrimSpace(resolved.GetCredential("task_id"))
|
||||
}
|
||||
}
|
||||
return s.ensureAgentIdentityTask(ctx, account, expectedTaskID)
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) isAgentIdentityAccount(ctx context.Context, account *Account) bool {
|
||||
if account == nil {
|
||||
return false
|
||||
}
|
||||
credAccount := account
|
||||
if account.IsShadow() {
|
||||
resolved, err := resolveCredentialAccount(ctx, s.accountRepo, account)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
credAccount = resolved
|
||||
}
|
||||
return credAccount != nil && credAccount.IsOpenAIAgentIdentity()
|
||||
}
|
||||
|
||||
// redactAgentIdentitySensitiveBody removes credential values before an
|
||||
// upstream error can reach logs, ops events, or returned error text. Agent
|
||||
// Identity responses should not echo these values, but keeping this boundary
|
||||
// defensive prevents accidental disclosure if an upstream error does.
|
||||
func redactAgentIdentitySensitiveBodyForAccount(ctx context.Context, repo AccountRepository, account *Account, body []byte) []byte {
|
||||
if account == nil || len(body) == 0 {
|
||||
return body
|
||||
}
|
||||
credAccount := account
|
||||
if account != nil && account.IsShadow() {
|
||||
if resolved, err := resolveCredentialAccount(ctx, repo, account); err == nil && resolved != nil {
|
||||
credAccount = resolved
|
||||
}
|
||||
}
|
||||
if credAccount == nil || !credAccount.IsOpenAIAgentIdentity() {
|
||||
return body
|
||||
}
|
||||
redacted := string(body)
|
||||
for _, key := range []string{
|
||||
"agent_private_key",
|
||||
"agent_runtime_id",
|
||||
"task_id",
|
||||
"access_token",
|
||||
"refresh_token",
|
||||
"id_token",
|
||||
"api_key",
|
||||
"session_key",
|
||||
"cookie",
|
||||
} {
|
||||
if value := strings.TrimSpace(credAccount.GetCredential(key)); value != "" {
|
||||
redacted = strings.ReplaceAll(redacted, value, "[redacted]")
|
||||
}
|
||||
}
|
||||
const assertionPrefix = "AgentAssertion "
|
||||
for offset := 0; offset < len(redacted); {
|
||||
relativeStart := strings.Index(redacted[offset:], assertionPrefix)
|
||||
if relativeStart < 0 {
|
||||
break
|
||||
}
|
||||
start := offset + relativeStart
|
||||
valueStart := start + len(assertionPrefix)
|
||||
end := valueStart
|
||||
for end < len(redacted) && !strings.ContainsRune(" \t\r\n\"',}", rune(redacted[end])) {
|
||||
end++
|
||||
}
|
||||
redacted = redacted[:valueStart] + "[redacted]" + redacted[end:]
|
||||
offset = valueStart + len("[redacted]")
|
||||
}
|
||||
return []byte(redacted)
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) redactAgentIdentitySensitiveBody(ctx context.Context, account *Account, body []byte) []byte {
|
||||
if !s.isAgentIdentityAccount(ctx, account) {
|
||||
return body
|
||||
}
|
||||
return redactAgentIdentitySensitiveBodyForAccount(ctx, s.accountRepo, account, body)
|
||||
}
|
||||
@@ -0,0 +1,500 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestAccountTestServiceOpenAICompactAgentIdentityUsesFreshAssertion(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
key, privateKey := newTestAgentIdentityKey(t)
|
||||
account := Account{
|
||||
ID: 21,
|
||||
Name: "agent-identity",
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeOAuth,
|
||||
Status: StatusActive,
|
||||
Schedulable: true,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"auth_mode": OpenAIAuthModeAgentIdentity,
|
||||
"agent_runtime_id": key.runtimeID,
|
||||
"agent_private_key": privateKey,
|
||||
"task_id": key.taskID,
|
||||
"chatgpt_account_id": "account-agent-test",
|
||||
"chatgpt_account_is_fedramp": true,
|
||||
},
|
||||
}
|
||||
repo := &snapshotUpdateAccountRepo{stubOpenAIAccountRepo: stubOpenAIAccountRepo{accounts: []Account{account}}}
|
||||
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(`{"id":"compact-agent","status":"completed"}`)),
|
||||
}}
|
||||
svc := &AccountTestService{accountRepo: repo, httpUpstream: upstream}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/21/test", bytes.NewReader(nil))
|
||||
|
||||
require.NoError(t, svc.TestAccountConnection(c, account.ID, "gpt-5.4", "", AccountTestModeCompact))
|
||||
require.Equal(t, "AgentAssertion", strings.SplitN(upstream.lastReq.Header.Get("Authorization"), " ", 2)[0])
|
||||
require.Equal(t, "account-agent-test", upstream.lastReq.Header.Get("chatgpt-account-id"))
|
||||
require.Equal(t, "true", upstream.lastReq.Header.Get("x-openai-fedramp"))
|
||||
require.NotContains(t, upstream.lastReq.Header.Get("Authorization"), privateKey)
|
||||
}
|
||||
|
||||
func TestAccountTestServiceOpenAICompactAgentIdentityRecoversInvalidTaskOnce(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
key, privateKey := newTestAgentIdentityKey(t)
|
||||
account := &Account{
|
||||
ID: 22,
|
||||
Name: "agent-identity-recovery",
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeOAuth,
|
||||
Status: StatusActive,
|
||||
Schedulable: true,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"auth_mode": OpenAIAuthModeAgentIdentity,
|
||||
"agent_runtime_id": key.runtimeID,
|
||||
"agent_private_key": privateKey,
|
||||
"task_id": "task-compact-old",
|
||||
"chatgpt_account_id": "account-agent-compact-recovery",
|
||||
},
|
||||
}
|
||||
repo := &accountTestAgentIdentityRepo{account: account}
|
||||
registerCalls := 0
|
||||
registerServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
registerCalls++
|
||||
_, _ = io.WriteString(w, `{"task_id":"task-compact-new"}`)
|
||||
}))
|
||||
defer registerServer.Close()
|
||||
oldBase := openAIAgentIdentityAuthAPIBaseURL
|
||||
openAIAgentIdentityAuthAPIBaseURL = registerServer.URL
|
||||
t.Cleanup(func() { openAIAgentIdentityAuthAPIBaseURL = oldBase })
|
||||
|
||||
upstream := &httpUpstreamRecorder{responses: []*http.Response{
|
||||
{StatusCode: http.StatusUnauthorized, Header: http.Header{"Content-Type": []string{"application/json"}}, Body: io.NopCloser(strings.NewReader(`{"error":{"code":"invalid_task_id"}}`))},
|
||||
{StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"application/json"}}, Body: io.NopCloser(strings.NewReader(`{"id":"compact-agent","status":"completed"}`))},
|
||||
}}
|
||||
invalidator := &agentIdentityWSInvalidationRecorder{}
|
||||
svc := &AccountTestService{accountRepo: repo, httpUpstream: upstream, agentIdentityWS: invalidator}
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/22/test", bytes.NewReader(nil))
|
||||
|
||||
require.NoError(t, svc.TestAccountConnection(c, account.ID, "gpt-5.4", "", AccountTestModeCompact))
|
||||
require.Equal(t, 1, registerCalls)
|
||||
require.Len(t, upstream.requests, 2)
|
||||
require.Equal(t, "task-compact-new", account.GetCredential("task_id"))
|
||||
require.Equal(t, 0, repo.setErrorCalls)
|
||||
require.Equal(t, []int64{account.ID}, invalidator.accountIDs)
|
||||
}
|
||||
|
||||
func TestOpenAIAgentIdentityPassthroughKeepsSessionAndPromptCacheHeaders(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
key, privateKey := newTestAgentIdentityKey(t)
|
||||
account := &Account{
|
||||
ID: 24,
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeOAuth,
|
||||
Credentials: map[string]any{
|
||||
"auth_mode": OpenAIAuthModeAgentIdentity,
|
||||
"agent_runtime_id": key.runtimeID,
|
||||
"agent_private_key": privateKey,
|
||||
"task_id": key.taskID,
|
||||
"chatgpt_account_id": "account-agent-passthrough",
|
||||
},
|
||||
}
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
body := []byte(`{"model":"gpt-5.4","instructions":"Reply OK","input":[],"stream":true,"prompt_cache_key":"cache-agent"}`)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body))
|
||||
c.Request.Header.Set("session_id", "client-session")
|
||||
c.Request.Header.Set("conversation_id", "client-conversation")
|
||||
c.Request.Header.Set("Authorization", "Bearer inbound-must-not-forward")
|
||||
|
||||
svc := &OpenAIGatewayService{}
|
||||
req, err := svc.buildUpstreamRequestOpenAIPassthrough(context.Background(), c, account, body, "")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "AgentAssertion", strings.SplitN(req.Header.Get("Authorization"), " ", 2)[0])
|
||||
require.Equal(t, "account-agent-passthrough", req.Header.Get("chatgpt-account-id"))
|
||||
require.NotEqual(t, "client-session", req.Header.Get("session_id"))
|
||||
require.NotEqual(t, "client-conversation", req.Header.Get("conversation_id"))
|
||||
require.Equal(t, isolateOpenAISessionID(0, "client-session"), req.Header.Get("session_id"))
|
||||
require.Equal(t, isolateOpenAISessionID(0, "client-conversation"), req.Header.Get("conversation_id"))
|
||||
requestBody, err := io.ReadAll(req.Body)
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, string(requestBody), `"prompt_cache_key":"cache-agent"`)
|
||||
|
||||
// Authentication mode must not affect session isolation or prompt-cache
|
||||
// behavior. Compare the same request with the existing OAuth path instead
|
||||
// of pinning this test to an implementation-specific hash.
|
||||
oauthAccount := &Account{
|
||||
ID: 26,
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeOAuth,
|
||||
Credentials: map[string]any{
|
||||
"chatgpt_account_id": "account-oauth-passthrough",
|
||||
},
|
||||
}
|
||||
oauthRecorder := httptest.NewRecorder()
|
||||
oauthContext, _ := gin.CreateTestContext(oauthRecorder)
|
||||
oauthContext.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body))
|
||||
oauthContext.Request.Header.Set("session_id", "client-session")
|
||||
oauthContext.Request.Header.Set("conversation_id", "client-conversation")
|
||||
oauthReq, err := svc.buildUpstreamRequestOpenAIPassthrough(context.Background(), oauthContext, oauthAccount, body, "oauth-token")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, oauthReq.Header.Get("session_id"), req.Header.Get("session_id"))
|
||||
require.Equal(t, oauthReq.Header.Get("conversation_id"), req.Header.Get("conversation_id"))
|
||||
}
|
||||
|
||||
func TestOpenAIAgentIdentityErrorRedactionDoesNotLeakCredentialValues(t *testing.T) {
|
||||
key, privateKey := newTestAgentIdentityKey(t)
|
||||
account := &Account{
|
||||
ID: 25,
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeOAuth,
|
||||
Credentials: map[string]any{
|
||||
"auth_mode": OpenAIAuthModeAgentIdentity,
|
||||
"agent_runtime_id": key.runtimeID,
|
||||
"agent_private_key": privateKey,
|
||||
"task_id": key.taskID,
|
||||
"access_token": key.runtimeID + "-oauth-value",
|
||||
},
|
||||
}
|
||||
svc := &OpenAIGatewayService{}
|
||||
oauthValue := account.GetCredential("access_token")
|
||||
redacted := svc.redactAgentIdentitySensitiveBody(context.Background(), account, []byte(`{"message":"runtime-test task-test `+oauthValue+` AgentAssertion abc123"}`))
|
||||
require.NotContains(t, string(redacted), key.runtimeID)
|
||||
require.NotContains(t, string(redacted), key.taskID)
|
||||
require.NotContains(t, string(redacted), oauthValue)
|
||||
require.NotContains(t, string(redacted), "AgentAssertion abc123")
|
||||
require.Contains(t, string(redacted), "[redacted]")
|
||||
}
|
||||
|
||||
func TestOpenAIAuthenticationHeadersPreserveOAuthPATAndAPIKeyBearerModes(t *testing.T) {
|
||||
svc := &OpenAIGatewayService{}
|
||||
tests := []struct {
|
||||
name string
|
||||
account *Account
|
||||
token string
|
||||
}{
|
||||
{name: "oauth", account: &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth}, token: "oauth-runtime-token"},
|
||||
{name: "personal access token", account: &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth, Credentials: map[string]any{"auth_mode": OpenAIAuthModePersonalAccessToken}}, token: "pat-runtime-token"},
|
||||
{name: "api key", account: &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}, token: "api-key-runtime-token"},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
headers, err := svc.buildOpenAIAuthenticationHeaders(context.Background(), tt.account, tt.token)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "Bearer "+tt.token, headers.Get("Authorization"))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIWSAgentIdentityRecoveryRequiresTaskInvalidBody(t *testing.T) {
|
||||
require.False(t, isAgentIdentityTaskInvalidWSDialError(&openAIWSDialError{
|
||||
StatusCode: http.StatusUnauthorized,
|
||||
ResponseBody: []byte(`{"error":{"code":"invalid_signature"}}`),
|
||||
}))
|
||||
require.True(t, isAgentIdentityTaskInvalidWSDialError(&openAIWSDialError{
|
||||
StatusCode: http.StatusUnauthorized,
|
||||
ResponseBody: []byte(`{"error":{"code":"invalid_task_id"}}`),
|
||||
}))
|
||||
}
|
||||
|
||||
func TestValidateOpenAIWSBearerTokenAllowsAgentIdentityWithoutStoredToken(t *testing.T) {
|
||||
t.Run("Given Agent Identity When a WS path receives no bearer token Then dial-time assertion auth is allowed", func(t *testing.T) {
|
||||
account := &Account{
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeOAuth,
|
||||
Credentials: map[string]any{
|
||||
"auth_mode": OpenAIAuthModeAgentIdentity,
|
||||
},
|
||||
}
|
||||
|
||||
require.NoError(t, validateOpenAIWSBearerToken(account, ""))
|
||||
})
|
||||
|
||||
t.Run("Given bearer credentials When a WS path receives no token Then the request is rejected", func(t *testing.T) {
|
||||
accounts := []*Account{
|
||||
{Platform: PlatformOpenAI, Type: AccountTypeOAuth},
|
||||
{Platform: PlatformOpenAI, Type: AccountTypeOAuth, Credentials: map[string]any{"auth_mode": OpenAIAuthModePersonalAccessToken}},
|
||||
{Platform: PlatformOpenAI, Type: AccountTypeAPIKey},
|
||||
}
|
||||
|
||||
for _, account := range accounts {
|
||||
require.EqualError(t, validateOpenAIWSBearerToken(account, ""), "token is empty")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestOpenAIWSConnPoolHeadersFactoryRunsAtDialAndStalePrewarmIsDiscarded(t *testing.T) {
|
||||
cfg := &config.Config{}
|
||||
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
|
||||
pool := newOpenAIWSConnPool(cfg)
|
||||
defer pool.Close()
|
||||
pool.setClientDialerForTest(&openAIWSFakeDialer{})
|
||||
|
||||
accountID := int64(22)
|
||||
ap := pool.getOrCreateAccountPool(accountID)
|
||||
factoryCalls := 0
|
||||
latestHeader := ""
|
||||
req := openAIWSAcquireRequest{
|
||||
Account: &Account{ID: accountID, Platform: PlatformOpenAI, Type: AccountTypeOAuth},
|
||||
WSURL: "wss://example.com/v1/responses",
|
||||
HeadersFactory: func(_ context.Context, headers http.Header) (http.Header, error) {
|
||||
factoryCalls++
|
||||
latestHeader = "AgentAssertion dial-" + string(rune('0'+factoryCalls))
|
||||
if headers == nil {
|
||||
headers = make(http.Header)
|
||||
}
|
||||
headers.Set("Authorization", latestHeader)
|
||||
return headers, nil
|
||||
},
|
||||
}
|
||||
ap.mu.Lock()
|
||||
ap.lastAcquire = &req
|
||||
generation := ap.generation
|
||||
ap.mu.Unlock()
|
||||
|
||||
pool.prewarmConns(accountID, req, 1, generation)
|
||||
require.Equal(t, 1, factoryCalls, "prewarm must generate authorization inside the actual dial")
|
||||
require.Equal(t, "AgentAssertion dial-1", latestHeader)
|
||||
|
||||
pool.ClearAccount(accountID)
|
||||
ap.mu.Lock()
|
||||
require.Empty(t, ap.conns, "credential recovery must remove pooled connections")
|
||||
require.Nil(t, ap.lastAcquire, "credential recovery must discard delayed acquire state")
|
||||
require.Equal(t, generation+1, ap.generation)
|
||||
ap.mu.Unlock()
|
||||
|
||||
// A prewarm captured before ClearAccount must not be admitted after recovery.
|
||||
pool.prewarmConns(accountID, req, 1, generation)
|
||||
ap.mu.Lock()
|
||||
require.Empty(t, ap.conns)
|
||||
ap.mu.Unlock()
|
||||
}
|
||||
|
||||
func TestOpenAIAgentIdentityTaskInvalidRetriesExactlyOnce(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
key, privateKey := newTestAgentIdentityKey(t)
|
||||
account := &Account{
|
||||
ID: 23,
|
||||
Name: "agent-identity",
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeOAuth,
|
||||
Status: StatusActive,
|
||||
Schedulable: true,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"auth_mode": OpenAIAuthModeAgentIdentity,
|
||||
"agent_runtime_id": key.runtimeID,
|
||||
"agent_private_key": privateKey,
|
||||
"task_id": "task-old",
|
||||
"chatgpt_account_id": "account-agent-retry",
|
||||
},
|
||||
}
|
||||
repo := &agentIdentityForwardRepo{account: account}
|
||||
registerCalls := 0
|
||||
registerServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
registerCalls++
|
||||
_, _ = io.WriteString(w, `{"task_id":"task-new"}`)
|
||||
}))
|
||||
defer registerServer.Close()
|
||||
oldBase := openAIAgentIdentityAuthAPIBaseURL
|
||||
openAIAgentIdentityAuthAPIBaseURL = registerServer.URL
|
||||
t.Cleanup(func() { openAIAgentIdentityAuthAPIBaseURL = oldBase })
|
||||
|
||||
successBody := `{"id":"resp-agent-retry","object":"response","model":"gpt-5.4","output":[],"usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}`
|
||||
upstream := &httpUpstreamRecorder{responses: []*http.Response{
|
||||
{StatusCode: http.StatusUnauthorized, Header: http.Header{"Content-Type": []string{"application/json"}}, Body: io.NopCloser(strings.NewReader(`{"error":{"code":"invalid_task_id"}}`))},
|
||||
{StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"application/json"}}, Body: io.NopCloser(strings.NewReader(successBody))},
|
||||
}}
|
||||
require.True(t, isAgentIdentityTaskInvalidHTTPResponse(http.StatusUnauthorized, []byte(`{"error":{"code":"invalid_task_id"}}`)))
|
||||
svc := &OpenAIGatewayService{cfg: &config.Config{}, accountRepo: repo, httpUpstream: upstream}
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"gpt-5.4","instructions":"Reply OK","input":[],"stream":false}`))
|
||||
|
||||
_, err := svc.Forward(context.Background(), c, account, []byte(`{"model":"gpt-5.4","instructions":"Reply OK","input":[],"stream":false}`))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, registerCalls)
|
||||
require.Len(t, upstream.requests, 2)
|
||||
require.NotEqual(t, upstream.requests[0].Header.Get("Authorization"), upstream.requests[1].Header.Get("Authorization"))
|
||||
require.Equal(t, "task-new", decodeAgentAssertionTask(t, upstream.requests[1].Header.Get("Authorization")))
|
||||
|
||||
// Two consecutive invalid responses still produce only one retry for this
|
||||
// request; the recovery path must not loop indefinitely.
|
||||
upstream.responses = []*http.Response{
|
||||
{StatusCode: http.StatusUnauthorized, Header: http.Header{"Content-Type": []string{"application/json"}}, Body: io.NopCloser(strings.NewReader(`{"error":{"code":"invalid_task_id"}}`))},
|
||||
{StatusCode: http.StatusUnauthorized, Header: http.Header{"Content-Type": []string{"application/json"}}, Body: io.NopCloser(strings.NewReader(`{"error":{"code":"invalid_task_id"}}`))},
|
||||
}
|
||||
rec2 := httptest.NewRecorder()
|
||||
c2, _ := gin.CreateTestContext(rec2)
|
||||
c2.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"gpt-5.4","instructions":"Reply OK","input":[],"stream":false}`))
|
||||
_, err = svc.Forward(context.Background(), c2, account, []byte(`{"model":"gpt-5.4","instructions":"Reply OK","input":[],"stream":false}`))
|
||||
require.Error(t, err)
|
||||
require.Equal(t, 2, registerCalls)
|
||||
require.Len(t, upstream.requests, 4)
|
||||
|
||||
// Passthrough uses the same one-shot task recovery contract.
|
||||
account.Extra = map[string]any{"openai_passthrough": true}
|
||||
account.Credentials["task_id"] = "task-old-passthrough"
|
||||
upstream.responses = []*http.Response{
|
||||
{StatusCode: http.StatusUnauthorized, Header: http.Header{"Content-Type": []string{"application/json"}}, Body: io.NopCloser(strings.NewReader(`{"error":{"code":"invalid_task_id"}}`))},
|
||||
{StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"text/event-stream"}}, Body: io.NopCloser(strings.NewReader("data: {\"type\":\"response.completed\",\"response\":{\"usage\":{\"input_tokens\":1,\"output_tokens\":1,\"total_tokens\":2}}}\n\ndata: [DONE]\n\n"))},
|
||||
}
|
||||
rec3 := httptest.NewRecorder()
|
||||
c3, _ := gin.CreateTestContext(rec3)
|
||||
c3.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"gpt-5.4","instructions":"Reply OK","input":[],"stream":false}`))
|
||||
_, err = svc.Forward(context.Background(), c3, account, []byte(`{"model":"gpt-5.4","instructions":"Reply OK","input":[],"stream":false}`))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 3, registerCalls)
|
||||
require.Len(t, upstream.requests, 6)
|
||||
}
|
||||
|
||||
func TestOpenAIAgentIdentityCompatRoutesRecoverInvalidTaskOnce(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
tests := []struct {
|
||||
name string
|
||||
path string
|
||||
body []byte
|
||||
call func(*OpenAIGatewayService, context.Context, *gin.Context, *Account, []byte) (*OpenAIForwardResult, error)
|
||||
}{
|
||||
{
|
||||
name: "chat completions",
|
||||
path: "/v1/chat/completions",
|
||||
body: []byte(`{"model":"gpt-5.4","stream":false,"messages":[{"role":"user","content":"hi"}]}`),
|
||||
call: func(s *OpenAIGatewayService, ctx context.Context, c *gin.Context, account *Account, body []byte) (*OpenAIForwardResult, error) {
|
||||
return s.ForwardAsChatCompletions(ctx, c, account, body, "", "gpt-5.4")
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "anthropic messages",
|
||||
path: "/v1/messages",
|
||||
body: []byte(`{"model":"gpt-5.4","stream":false,"max_tokens":32,"messages":[{"role":"user","content":"hi"}]}`),
|
||||
call: func(s *OpenAIGatewayService, ctx context.Context, c *gin.Context, account *Account, body []byte) (*OpenAIForwardResult, error) {
|
||||
return s.ForwardAsAnthropic(ctx, c, account, body, "", "gpt-5.4")
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for index, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
key, privateKey := newTestAgentIdentityKey(t)
|
||||
account := &Account{
|
||||
ID: int64(40 + index),
|
||||
Name: "agent-identity-compat",
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeOAuth,
|
||||
Status: StatusActive,
|
||||
Schedulable: true,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"auth_mode": OpenAIAuthModeAgentIdentity,
|
||||
"agent_runtime_id": key.runtimeID,
|
||||
"agent_private_key": privateKey,
|
||||
"task_id": "task-compat-old",
|
||||
"chatgpt_account_id": "account-compat-recovery",
|
||||
},
|
||||
}
|
||||
repo := &agentIdentityForwardRepo{account: account}
|
||||
registerCalls := 0
|
||||
registerServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
registerCalls++
|
||||
_, _ = io.WriteString(w, `{"task_id":"task-compat-new"}`)
|
||||
}))
|
||||
defer registerServer.Close()
|
||||
oldBase := openAIAgentIdentityAuthAPIBaseURL
|
||||
openAIAgentIdentityAuthAPIBaseURL = registerServer.URL
|
||||
t.Cleanup(func() { openAIAgentIdentityAuthAPIBaseURL = oldBase })
|
||||
|
||||
upstream := &httpUpstreamRecorder{responses: []*http.Response{
|
||||
{StatusCode: http.StatusUnauthorized, Header: http.Header{"Content-Type": []string{"application/json"}}, Body: io.NopCloser(strings.NewReader(`{"error":{"code":"invalid_task_id"}}`))},
|
||||
{StatusCode: http.StatusUnauthorized, Header: http.Header{"Content-Type": []string{"application/json"}}, Body: io.NopCloser(strings.NewReader(`{"error":{"code":"invalid_task_id"}}`))},
|
||||
}}
|
||||
svc := &OpenAIGatewayService{cfg: &config.Config{}, accountRepo: repo, httpUpstream: upstream}
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, tt.path, bytes.NewReader(tt.body))
|
||||
|
||||
_, err := tt.call(svc, context.Background(), c, account, tt.body)
|
||||
require.Error(t, err)
|
||||
require.Equal(t, 1, registerCalls)
|
||||
require.Len(t, upstream.requests, 2)
|
||||
require.Equal(t, "task-compat-new", account.GetCredential("task_id"))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func decodeAgentAssertionTask(t *testing.T, header string) string {
|
||||
t.Helper()
|
||||
encoded := strings.TrimPrefix(header, "AgentAssertion ")
|
||||
decoded, err := base64.RawURLEncoding.DecodeString(encoded)
|
||||
require.NoError(t, err)
|
||||
var envelope struct {
|
||||
TaskID string `json:"task_id"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal(decoded, &envelope))
|
||||
return envelope.TaskID
|
||||
}
|
||||
|
||||
type agentIdentityForwardRepo struct {
|
||||
AccountRepository
|
||||
account *Account
|
||||
}
|
||||
|
||||
type agentIdentityWSInvalidationRecorder struct {
|
||||
accountIDs []int64
|
||||
}
|
||||
|
||||
func (r *agentIdentityWSInvalidationRecorder) InvalidateAgentIdentityWSConnections(accountID int64) {
|
||||
r.accountIDs = append(r.accountIDs, accountID)
|
||||
}
|
||||
|
||||
type accountTestAgentIdentityRepo struct {
|
||||
AccountRepository
|
||||
account *Account
|
||||
setErrorCalls int
|
||||
}
|
||||
|
||||
func (r *accountTestAgentIdentityRepo) GetByID(_ context.Context, _ int64) (*Account, error) {
|
||||
return r.account, nil
|
||||
}
|
||||
|
||||
func (r *accountTestAgentIdentityRepo) UpdateCredentials(_ context.Context, _ int64, credentials map[string]any) error {
|
||||
r.account.Credentials = credentials
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *accountTestAgentIdentityRepo) UpdateExtra(_ context.Context, _ int64, _ map[string]any) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *accountTestAgentIdentityRepo) SetError(_ context.Context, _ int64, _ string) error {
|
||||
r.setErrorCalls++
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *agentIdentityForwardRepo) GetByID(_ context.Context, _ int64) (*Account, error) {
|
||||
return r.account, nil
|
||||
}
|
||||
|
||||
func (r *agentIdentityForwardRepo) UpdateCredentials(_ context.Context, _ int64, credentials map[string]any) error {
|
||||
r.account.Credentials = credentials
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,229 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/ed25519"
|
||||
"crypto/rand"
|
||||
"crypto/sha512"
|
||||
"crypto/x509"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
"golang.org/x/crypto/curve25519"
|
||||
"golang.org/x/crypto/nacl/box"
|
||||
)
|
||||
|
||||
func newTestAgentIdentityKey(t *testing.T) (agentIdentityKey, string) {
|
||||
t.Helper()
|
||||
_, privateKey, err := ed25519.GenerateKey(rand.Reader)
|
||||
require.NoError(t, err)
|
||||
der, err := x509.MarshalPKCS8PrivateKey(privateKey)
|
||||
require.NoError(t, err)
|
||||
return agentIdentityKey{
|
||||
runtimeID: "runtime-test",
|
||||
privateKey: privateKey,
|
||||
taskID: "task-test",
|
||||
}, base64.StdEncoding.EncodeToString(der)
|
||||
}
|
||||
|
||||
func TestBuildAgentAssertionMatchesCodexEnvelopeAndSignature(t *testing.T) {
|
||||
key, _ := newTestAgentIdentityKey(t)
|
||||
now := time.Date(2026, 7, 14, 8, 9, 10, 0, time.FixedZone("UTC+8", 8*60*60))
|
||||
assertion, err := buildAgentAssertion(key, now)
|
||||
require.NoError(t, err)
|
||||
require.True(t, strings.HasPrefix(assertion, "AgentAssertion "))
|
||||
|
||||
encoded := strings.TrimPrefix(assertion, "AgentAssertion ")
|
||||
decoded, err := base64.RawURLEncoding.DecodeString(encoded)
|
||||
require.NoError(t, err)
|
||||
var envelope struct {
|
||||
AgentRuntimeID string `json:"agent_runtime_id"`
|
||||
TaskID string `json:"task_id"`
|
||||
Timestamp string `json:"timestamp"`
|
||||
Signature string `json:"signature"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal(decoded, &envelope))
|
||||
require.Equal(t, "runtime-test", envelope.AgentRuntimeID)
|
||||
require.Equal(t, "task-test", envelope.TaskID)
|
||||
require.Equal(t, "2026-07-14T00:09:10Z", envelope.Timestamp)
|
||||
signature, err := base64.StdEncoding.DecodeString(envelope.Signature)
|
||||
require.NoError(t, err)
|
||||
publicKey, ok := key.privateKey.Public().(ed25519.PublicKey)
|
||||
require.True(t, ok)
|
||||
require.True(t, ed25519.Verify(publicKey, []byte("runtime-test:task-test:2026-07-14T00:09:10Z"), signature))
|
||||
}
|
||||
|
||||
func TestDecryptAgentTaskIDSupportsCodexSealedBoxResponse(t *testing.T) {
|
||||
key, _ := newTestAgentIdentityKey(t)
|
||||
digest := sha512.Sum512(key.privateKey.Seed())
|
||||
var curvePrivate [32]byte
|
||||
copy(curvePrivate[:], digest[:32])
|
||||
curvePrivate[0] &= 248
|
||||
curvePrivate[31] &= 127
|
||||
curvePrivate[31] |= 64
|
||||
curvePublicBytes, err := curve25519.X25519(curvePrivate[:], curve25519.Basepoint)
|
||||
require.NoError(t, err)
|
||||
var curvePublic [32]byte
|
||||
copy(curvePublic[:], curvePublicBytes)
|
||||
ciphertext, err := box.SealAnonymous(nil, []byte("task-sealed"), &curvePublic, rand.Reader)
|
||||
require.NoError(t, err)
|
||||
got, err := decryptAgentTaskID(key, base64.StdEncoding.EncodeToString(ciphertext))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "task-sealed", got)
|
||||
}
|
||||
|
||||
func TestRegisterAgentIdentityTaskAcceptsPlaintextAndEncryptedResponses(t *testing.T) {
|
||||
key, privateKey := newTestAgentIdentityKey(t)
|
||||
requestCount := 0
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
require.Equal(t, http.MethodPost, r.Method)
|
||||
require.Equal(t, "/v1/agent/runtime-test/task/register", r.URL.Path)
|
||||
var request map[string]string
|
||||
require.NoError(t, json.NewDecoder(r.Body).Decode(&request))
|
||||
require.NotEmpty(t, request["timestamp"])
|
||||
require.NotEmpty(t, request["signature"])
|
||||
requestCount++
|
||||
if requestCount == 2 {
|
||||
digest := sha512.Sum512(key.privateKey.Seed())
|
||||
var curvePrivate [32]byte
|
||||
copy(curvePrivate[:], digest[:32])
|
||||
curvePrivate[0] &= 248
|
||||
curvePrivate[31] &= 127
|
||||
curvePrivate[31] |= 64
|
||||
curvePublicBytes, curveErr := curve25519.X25519(curvePrivate[:], curve25519.Basepoint)
|
||||
require.NoError(t, curveErr)
|
||||
var curvePublic [32]byte
|
||||
copy(curvePublic[:], curvePublicBytes)
|
||||
ciphertext, sealErr := box.SealAnonymous(nil, []byte("task-encrypted"), &curvePublic, rand.Reader)
|
||||
require.NoError(t, sealErr)
|
||||
_, _ = fmt.Fprintf(w, `{"encrypted_task_id":%q}`, base64.StdEncoding.EncodeToString(ciphertext))
|
||||
return
|
||||
}
|
||||
_, _ = w.Write([]byte(`{"task_id":"task-plain"}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
oldBase := openAIAgentIdentityAuthAPIBaseURL
|
||||
openAIAgentIdentityAuthAPIBaseURL = server.URL
|
||||
t.Cleanup(func() { openAIAgentIdentityAuthAPIBaseURL = oldBase })
|
||||
|
||||
account := &Account{ID: 1, Type: AccountTypeOAuth, Platform: PlatformOpenAI, Credentials: map[string]any{
|
||||
"auth_mode": OpenAIAuthModeAgentIdentity,
|
||||
"agent_runtime_id": key.runtimeID,
|
||||
"agent_private_key": privateKey,
|
||||
}}
|
||||
taskID, err := registerAgentIdentityTask(context.Background(), account)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "task-plain", taskID)
|
||||
taskID, err = registerAgentIdentityTask(context.Background(), account)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "task-encrypted", taskID)
|
||||
}
|
||||
|
||||
func TestEnsureAgentIdentityTaskPersistsAndRedactsCredentials(t *testing.T) {
|
||||
key, privateKey := newTestAgentIdentityKey(t)
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
_, _ = w.Write([]byte(`{"task_id":"task-persisted"}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
oldBase := openAIAgentIdentityAuthAPIBaseURL
|
||||
openAIAgentIdentityAuthAPIBaseURL = server.URL
|
||||
t.Cleanup(func() { openAIAgentIdentityAuthAPIBaseURL = oldBase })
|
||||
|
||||
repo := &agentIdentityCredentialsRepo{}
|
||||
account := &Account{ID: 7, Type: AccountTypeOAuth, Platform: PlatformOpenAI, Credentials: map[string]any{
|
||||
"auth_mode": OpenAIAuthModeAgentIdentity,
|
||||
"agent_runtime_id": key.runtimeID,
|
||||
"agent_private_key": privateKey,
|
||||
"chatgpt_account_id": "account-test",
|
||||
}}
|
||||
service := &OpenAIGatewayService{accountRepo: repo}
|
||||
require.NoError(t, service.ensureAgentIdentityTask(context.Background(), account, ""))
|
||||
require.Equal(t, "task-persisted", account.GetCredential("task_id"))
|
||||
require.Equal(t, "task-persisted", repo.credentials["task_id"])
|
||||
require.True(t, IsSensitiveCredentialKey("agent_private_key"))
|
||||
redacted := make(map[string]any)
|
||||
for key, value := range account.Credentials {
|
||||
if !IsSensitiveCredentialKey(key) {
|
||||
redacted[key] = value
|
||||
}
|
||||
}
|
||||
require.NotContains(t, string(mustAgentIdentityJSON(t, redacted)), privateKey)
|
||||
}
|
||||
|
||||
func TestEnsureAgentIdentityTaskSharesLockAcrossServicesForSameAccount(t *testing.T) {
|
||||
key, privateKey := newTestAgentIdentityKey(t)
|
||||
account := &Account{ID: 9001, Type: AccountTypeOAuth, Platform: PlatformOpenAI, Credentials: map[string]any{
|
||||
"auth_mode": OpenAIAuthModeAgentIdentity,
|
||||
"agent_runtime_id": key.runtimeID,
|
||||
"agent_private_key": privateKey,
|
||||
}}
|
||||
repo := &agentIdentityCredentialsRepo{account: account}
|
||||
registerCalls := 0
|
||||
var registerMu sync.Mutex
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
registerMu.Lock()
|
||||
registerCalls++
|
||||
registerMu.Unlock()
|
||||
_, _ = w.Write([]byte(`{"task_id":"task-shared"}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
oldBase := openAIAgentIdentityAuthAPIBaseURL
|
||||
openAIAgentIdentityAuthAPIBaseURL = server.URL
|
||||
t.Cleanup(func() { openAIAgentIdentityAuthAPIBaseURL = oldBase })
|
||||
|
||||
start := make(chan struct{})
|
||||
errors := make(chan error, 2)
|
||||
requests := []*Account{cloneAgentIdentityTestAccount(account), cloneAgentIdentityTestAccount(account)}
|
||||
for _, request := range requests {
|
||||
go func() {
|
||||
<-start
|
||||
errors <- ensureAgentIdentityTaskForAccount(context.Background(), repo, nil, &sync.Mutex{}, request, "")
|
||||
}()
|
||||
}
|
||||
close(start)
|
||||
require.NoError(t, <-errors)
|
||||
require.NoError(t, <-errors)
|
||||
registerMu.Lock()
|
||||
defer registerMu.Unlock()
|
||||
require.Equal(t, 1, registerCalls)
|
||||
require.Equal(t, "task-shared", repo.account.GetCredential("task_id"))
|
||||
}
|
||||
|
||||
func cloneAgentIdentityTestAccount(account *Account) *Account {
|
||||
copy := *account
|
||||
copy.Credentials = shallowCopyMap(account.Credentials)
|
||||
return ©
|
||||
}
|
||||
|
||||
type agentIdentityCredentialsRepo struct {
|
||||
AccountRepository
|
||||
credentials map[string]any
|
||||
account *Account
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
func (r *agentIdentityCredentialsRepo) GetByID(_ context.Context, _ int64) (*Account, error) {
|
||||
return r.account, nil
|
||||
}
|
||||
|
||||
func (r *agentIdentityCredentialsRepo) UpdateCredentials(_ context.Context, _ int64, credentials map[string]any) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.credentials = credentials
|
||||
return nil
|
||||
}
|
||||
|
||||
func mustAgentIdentityJSON(t *testing.T, value any) []byte {
|
||||
t.Helper()
|
||||
encoded, err := json.Marshal(value)
|
||||
require.NoError(t, err)
|
||||
return encoded
|
||||
}
|
||||
@@ -241,7 +241,7 @@ func (s *OpenAIGatewayService) FetchCodexModelsManifest(ctx context.Context, acc
|
||||
switch {
|
||||
case credAccount.IsOpenAIOAuth():
|
||||
authToken = strings.TrimSpace(credAccount.GetOpenAIAccessToken())
|
||||
if authToken == "" {
|
||||
if authToken == "" && !credAccount.IsOpenAIAgentIdentity() {
|
||||
return nil, infraerrors.New(http.StatusBadGateway, "OPENAI_CODEX_MODELS_TOKEN_MISSING", "account has no Codex backend access token")
|
||||
}
|
||||
case credAccount.IsOpenAIApiKey():
|
||||
@@ -277,16 +277,25 @@ func (s *OpenAIGatewayService) FetchCodexModelsManifest(ctx context.Context, acc
|
||||
}
|
||||
|
||||
headers := make(http.Header)
|
||||
headers.Set("Authorization", "Bearer "+authToken)
|
||||
if useAPIKeyUpstream {
|
||||
headers.Set("Authorization", "Bearer "+authToken)
|
||||
credAccount.ApplyHeaderOverrides(headers)
|
||||
} else {
|
||||
authHeaders, authErr := s.buildOpenAIAuthenticationHeaders(ctx, credAccount, authToken)
|
||||
if authErr != nil {
|
||||
return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_CODEX_MODELS_AUTH_FAILED", "build Codex models authentication: %v", authErr)
|
||||
}
|
||||
for key, values := range authHeaders {
|
||||
for _, value := range values {
|
||||
headers.Add(key, value)
|
||||
}
|
||||
}
|
||||
setOpenAIChatGPTAccountHeaders(headers, credAccount)
|
||||
}
|
||||
headers.Set("Accept", "application/json")
|
||||
headers.Set("Originator", "codex_cli_rs")
|
||||
headers.Set("Version", clientVersion)
|
||||
headers.Set("User-Agent", codexCLIUserAgent)
|
||||
if useAPIKeyUpstream {
|
||||
credAccount.ApplyHeaderOverrides(headers)
|
||||
} else {
|
||||
setOpenAIChatGPTAccountHeaders(headers, credAccount)
|
||||
}
|
||||
|
||||
proxyURL := ""
|
||||
if account.ProxyID != nil && account.Proxy != nil {
|
||||
|
||||
@@ -199,6 +199,49 @@ func TestFetchCodexModelsManifestPassthrough(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestFetchCodexModelsManifestAgentIdentityUsesAssertionWithoutOAuthToken(t *testing.T) {
|
||||
key, privateKey := newTestAgentIdentityKey(t)
|
||||
account := &Account{
|
||||
ID: 3,
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeOAuth,
|
||||
Credentials: map[string]any{
|
||||
"auth_mode": OpenAIAuthModeAgentIdentity,
|
||||
"agent_runtime_id": key.runtimeID,
|
||||
"agent_private_key": privateKey,
|
||||
"task_id": key.taskID,
|
||||
"chatgpt_account_id": "acc-agent",
|
||||
},
|
||||
}
|
||||
|
||||
var gotAuth, gotAccountID string
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
gotAuth = r.Header.Get("Authorization")
|
||||
gotAccountID = r.Header.Get("chatgpt-account-id")
|
||||
_, _ = w.Write([]byte(`{"models":[]}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
original := chatgptCodexModelsURL
|
||||
chatgptCodexModelsURL = server.URL
|
||||
defer func() { chatgptCodexModelsURL = original }()
|
||||
|
||||
s := &OpenAIGatewayService{}
|
||||
manifest, err := s.FetchCodexModelsManifest(context.Background(), account, "0.137.0", "")
|
||||
if err != nil {
|
||||
t.Fatalf("FetchCodexModelsManifest returned error: %v", err)
|
||||
}
|
||||
if string(manifest.Body) != `{"models":[]}` {
|
||||
t.Fatalf("unexpected manifest body: %q", manifest.Body)
|
||||
}
|
||||
if !strings.HasPrefix(gotAuth, "AgentAssertion ") {
|
||||
t.Fatalf("authorization scheme: got %q", strings.SplitN(gotAuth, " ", 2)[0])
|
||||
}
|
||||
if gotAccountID != "acc-agent" {
|
||||
t.Fatalf("chatgpt-account-id header: got %q", gotAccountID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFetchCodexModelsManifestDefaultClientVersion(t *testing.T) {
|
||||
var gotClientVersion string
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
|
||||
@@ -25,31 +26,96 @@ func deriveOpenAIContentSessionSeed(body []byte) string {
|
||||
return ""
|
||||
}
|
||||
|
||||
const (
|
||||
modelField = iota
|
||||
toolsField
|
||||
functionsField
|
||||
instructionsField
|
||||
messagesField
|
||||
inputField
|
||||
contentSessionSeedFieldCount
|
||||
allContentSessionSeedFields = 1<<contentSessionSeedFieldCount - 1
|
||||
)
|
||||
var fields [contentSessionSeedFieldCount]gjson.Result
|
||||
var seen uint8
|
||||
// Match gjson.GetBytes by starting at the first root container, even when
|
||||
// malformed input has a non-JSON prefix.
|
||||
root := body
|
||||
for i := 0; i < len(body); i++ {
|
||||
switch body[i] {
|
||||
case '{':
|
||||
root = body[i:]
|
||||
goto scanRoot
|
||||
case '[':
|
||||
return ""
|
||||
}
|
||||
}
|
||||
return ""
|
||||
|
||||
scanRoot:
|
||||
nextKeyOffset := 1
|
||||
parseRawJSONView(root).ForEach(func(key, value gjson.Result) bool {
|
||||
if key.Index < nextKeyOffset || key.Index > len(root) {
|
||||
return false
|
||||
}
|
||||
// Result.ForEach can continue after the root '}' on malformed input.
|
||||
// The separator range excludes braces inside the preceding parsed value.
|
||||
if bytes.IndexByte(root[nextKeyOffset:key.Index], '}') >= 0 {
|
||||
return false
|
||||
}
|
||||
nextKeyOffset = value.Index + len(value.Raw)
|
||||
|
||||
field := -1
|
||||
switch key.Str {
|
||||
case "model":
|
||||
field = modelField
|
||||
case "tools":
|
||||
field = toolsField
|
||||
case "functions":
|
||||
field = functionsField
|
||||
case "instructions":
|
||||
field = instructionsField
|
||||
case "messages":
|
||||
field = messagesField
|
||||
case "input":
|
||||
field = inputField
|
||||
}
|
||||
if field < 0 {
|
||||
return true
|
||||
}
|
||||
mask := uint8(1 << field)
|
||||
if seen&mask == 0 {
|
||||
fields[field] = value
|
||||
seen |= mask
|
||||
}
|
||||
return seen != allContentSessionSeedFields
|
||||
})
|
||||
|
||||
var b strings.Builder
|
||||
|
||||
if model := gjson.GetBytes(body, "model").String(); model != "" {
|
||||
if model := fields[modelField].String(); model != "" {
|
||||
_, _ = b.WriteString("model=")
|
||||
_, _ = b.WriteString(model)
|
||||
}
|
||||
|
||||
if tools := gjson.GetBytes(body, "tools"); tools.Exists() && tools.IsArray() && tools.Raw != "[]" {
|
||||
if tools := fields[toolsField]; tools.Exists() && tools.IsArray() && tools.Raw != "[]" {
|
||||
_, _ = b.WriteString("|tools=")
|
||||
_, _ = b.WriteString(normalizeCompatSeedJSON(json.RawMessage(tools.Raw)))
|
||||
}
|
||||
|
||||
if funcs := gjson.GetBytes(body, "functions"); funcs.Exists() && funcs.IsArray() && funcs.Raw != "[]" {
|
||||
if funcs := fields[functionsField]; funcs.Exists() && funcs.IsArray() && funcs.Raw != "[]" {
|
||||
_, _ = b.WriteString("|functions=")
|
||||
_, _ = b.WriteString(normalizeCompatSeedJSON(json.RawMessage(funcs.Raw)))
|
||||
}
|
||||
|
||||
if instr := gjson.GetBytes(body, "instructions").String(); instr != "" {
|
||||
if instr := fields[instructionsField].String(); instr != "" {
|
||||
_, _ = b.WriteString("|instructions=")
|
||||
_, _ = b.WriteString(instr)
|
||||
}
|
||||
|
||||
firstUserCaptured := false
|
||||
|
||||
msgs := gjson.GetBytes(body, "messages")
|
||||
msgs := fields[messagesField]
|
||||
if msgs.Exists() && msgs.IsArray() {
|
||||
msgs.ForEach(func(_, msg gjson.Result) bool {
|
||||
role := msg.Get("role").String()
|
||||
@@ -70,7 +136,7 @@ func deriveOpenAIContentSessionSeed(body []byte) string {
|
||||
}
|
||||
return true
|
||||
})
|
||||
} else if inp := gjson.GetBytes(body, "input"); inp.Exists() {
|
||||
} else if inp := fields[inputField]; inp.Exists() {
|
||||
if inp.Type == gjson.String {
|
||||
_, _ = b.WriteString("|input=")
|
||||
_, _ = b.WriteString(inp.String())
|
||||
|
||||
@@ -0,0 +1,29 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
var benchmarkOpenAIContentSessionSeed string
|
||||
|
||||
func BenchmarkDeriveOpenAIContentSessionSeedLargeBody(b *testing.B) {
|
||||
largeHistory := strings.Repeat("payload", 1<<17)
|
||||
tests := []struct {
|
||||
name string
|
||||
body []byte
|
||||
}{
|
||||
{name: "ChatCompletions", body: []byte(`{"model":"gpt-5.4","tools":[{"type":"function","function":{"name":"lookup"}}],"messages":[{"role":"system","content":"Be concise."},{"role":"user","content":"Hello"},{"role":"assistant","content":"` + largeHistory + `"},{"role":"user","content":"Follow-up"}]}`)},
|
||||
{name: "Responses", body: []byte(`{"model":"gpt-5.4","instructions":"Be concise.","tools":[{"type":"function","name":"lookup"}],"input":[{"role":"system","content":"System prompt"},{"role":"user","content":"Hello"},{"role":"assistant","content":"` + largeHistory + `"}]}`)},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
b.Run(test.name, func(b *testing.B) {
|
||||
b.ReportAllocs()
|
||||
b.SetBytes(int64(len(test.body)))
|
||||
for range b.N {
|
||||
benchmarkOpenAIContentSessionSeed = deriveOpenAIContentSessionSeed(test.body)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -1,9 +1,12 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
func TestDeriveOpenAIContentSessionSeed_EmptyInputs(t *testing.T) {
|
||||
@@ -197,6 +200,184 @@ func TestDeriveOpenAIContentSessionSeed_JSONCanonicalisation(t *testing.T) {
|
||||
require.Equal(t, s1, s2, "different formatting of identical JSON should produce the same seed")
|
||||
}
|
||||
|
||||
func TestDeriveOpenAIContentSessionSeed_SingleScanMatchesLegacyBytes(t *testing.T) {
|
||||
largeValue := strings.Repeat("payload", 1<<17)
|
||||
tests := []struct {
|
||||
name string
|
||||
body []byte
|
||||
}{
|
||||
{
|
||||
name: "large chat completions",
|
||||
body: []byte(`{"metadata":"` + largeValue + `","model":"gpt-5.4","tools":[{"type":"function","function":{"name":"lookup"}}],"functions":[{"name":"legacy_lookup"}],"messages":[{"role":"system","content":"System prompt"},{"role":"developer","content":[{"type":"text","text":"Developer prompt"}]},{"role":"user","content":"Hello"}]}`),
|
||||
},
|
||||
{
|
||||
name: "large responses",
|
||||
body: []byte(`{"metadata":"` + largeValue + `","model":"gpt-5.4","instructions":"Be concise.","tools":[{"type":"function","name":"lookup"}],"input":[{"role":"system","content":"System prompt"},{"role":"user","content":[{"type":"input_text","text":"Hello"}]}]}`),
|
||||
},
|
||||
{
|
||||
name: "fields in reverse order",
|
||||
body: []byte(`{"input":"fallback input","messages":[{"role":"user","content":"chat wins"}],"instructions":"Be concise.","functions":[{"name":"lookup"}],"tools":[{"type":"function","name":"lookup"}],"model":"gpt-5.4"}`),
|
||||
},
|
||||
{
|
||||
name: "missing and wrong type fields",
|
||||
body: []byte(`{"tools":[],"functions":null,"instructions":0,"messages":{},"input":[{"type":"input_text","text":"fallback"}]}`),
|
||||
},
|
||||
{
|
||||
name: "duplicate fields keep first value",
|
||||
body: []byte(`{"model":"first","model":"second","tools":[{"name":"first"}],"tools":[{"name":"second"}],"functions":[],"functions":[{"name":"second"}],"instructions":"first","instructions":"second","messages":null,"messages":[{"role":"user","content":"second"}],"input":"first input","input":"second input"}`),
|
||||
},
|
||||
{
|
||||
name: "escaped field names",
|
||||
body: []byte(`{"mo\u0064el":"gpt-5.4","mess\u0061ges":[{"role":"user","content":"Hello"}]}`),
|
||||
},
|
||||
{
|
||||
name: "trailing object fields are outside the root",
|
||||
body: []byte(`{"foo":1}{"model":"trailing","input":"trailing input"}`),
|
||||
},
|
||||
{
|
||||
name: "trailing quoted fields are outside the root",
|
||||
body: []byte(`{"model":"root"}"input":"trailing input"`),
|
||||
},
|
||||
{
|
||||
name: "leading garbage before the root",
|
||||
body: []byte(`garbage{"model":"gpt-5.4","input":"Hello"}`),
|
||||
},
|
||||
{
|
||||
name: "escaped braces remain inside string values",
|
||||
body: []byte(`{"metadata":"escaped } and [ and \" quote","model":"root","input":"Hello"}{"model":"trailing"}`),
|
||||
},
|
||||
{
|
||||
name: "nested braces do not end the root",
|
||||
body: []byte(`{"metadata":{"nested":"} ]"},"model":"root","input":"Hello"}{"model":"trailing"}`),
|
||||
},
|
||||
{
|
||||
name: "root array does not expose nested or trailing object fields",
|
||||
body: []byte(`[{"model":"nested"}]{"model":"trailing","input":"trailing input"}`),
|
||||
},
|
||||
{
|
||||
name: "trailing messages do not override root input",
|
||||
body: []byte(`{"input":"root input"}{"messages":[{"role":"user","content":"trailing"}]}`),
|
||||
},
|
||||
{
|
||||
name: "truncated string containing a closing brace",
|
||||
body: []byte(`{"model":"root","metadata":"still } inside`),
|
||||
},
|
||||
{
|
||||
name: "lenient truncated body",
|
||||
body: []byte(`{"model":"gpt-5.4","messages":[{"role":"user","content":"Hello"}]`),
|
||||
},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
require.Equal(t, legacyDeriveOpenAIContentSessionSeed(test.body), deriveOpenAIContentSessionSeed(test.body))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeriveOpenAIContentSessionSeed_AllTruncationOffsetsMatchLegacyBytes(t *testing.T) {
|
||||
bodies := []string{
|
||||
`{"model":"gpt-5.4","tools":[{"type":"function","function":{"name":"lookup"}}],"functions":[{"name":"legacy"}],"instructions":"escaped \" } text","messages":[{"role":"system","content":"System"},{"role":"user","content":[{"type":"text","text":"Hello"}]}],"input":"fallback"}`,
|
||||
`{"model":"gpt-5.4","instructions":"Be concise.","tools":[{"type":"function","name":"lookup"}],"input":[{"role":"system","content":"System"},{"role":"user","content":[{"type":"input_text","text":"Hello"}]}]}`,
|
||||
}
|
||||
for bodyIndex, body := range bodies {
|
||||
for end := 1; end < len(body); end++ {
|
||||
truncated := []byte(body[:end])
|
||||
require.Equalf(t, legacyDeriveOpenAIContentSessionSeed(truncated), deriveOpenAIContentSessionSeed(truncated), "body %d truncated at byte %d", bodyIndex, end)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func legacyDeriveOpenAIContentSessionSeed(body []byte) string {
|
||||
if len(body) == 0 {
|
||||
return ""
|
||||
}
|
||||
|
||||
var b strings.Builder
|
||||
|
||||
if model := gjson.GetBytes(body, "model").String(); model != "" {
|
||||
_, _ = b.WriteString("model=")
|
||||
_, _ = b.WriteString(model)
|
||||
}
|
||||
|
||||
if tools := gjson.GetBytes(body, "tools"); tools.Exists() && tools.IsArray() && tools.Raw != "[]" {
|
||||
_, _ = b.WriteString("|tools=")
|
||||
_, _ = b.WriteString(normalizeCompatSeedJSON(json.RawMessage(tools.Raw)))
|
||||
}
|
||||
|
||||
if funcs := gjson.GetBytes(body, "functions"); funcs.Exists() && funcs.IsArray() && funcs.Raw != "[]" {
|
||||
_, _ = b.WriteString("|functions=")
|
||||
_, _ = b.WriteString(normalizeCompatSeedJSON(json.RawMessage(funcs.Raw)))
|
||||
}
|
||||
|
||||
if instr := gjson.GetBytes(body, "instructions").String(); instr != "" {
|
||||
_, _ = b.WriteString("|instructions=")
|
||||
_, _ = b.WriteString(instr)
|
||||
}
|
||||
|
||||
firstUserCaptured := false
|
||||
|
||||
msgs := gjson.GetBytes(body, "messages")
|
||||
if msgs.Exists() && msgs.IsArray() {
|
||||
msgs.ForEach(func(_, msg gjson.Result) bool {
|
||||
role := msg.Get("role").String()
|
||||
switch role {
|
||||
case "system", "developer":
|
||||
_, _ = b.WriteString("|system=")
|
||||
if c := msg.Get("content"); c.Exists() {
|
||||
_, _ = b.WriteString(normalizeCompatSeedJSON(json.RawMessage(c.Raw)))
|
||||
}
|
||||
case "user":
|
||||
if !firstUserCaptured {
|
||||
_, _ = b.WriteString("|first_user=")
|
||||
if c := msg.Get("content"); c.Exists() {
|
||||
_, _ = b.WriteString(normalizeCompatSeedJSON(json.RawMessage(c.Raw)))
|
||||
}
|
||||
firstUserCaptured = true
|
||||
}
|
||||
}
|
||||
return true
|
||||
})
|
||||
} else if inp := gjson.GetBytes(body, "input"); inp.Exists() {
|
||||
if inp.Type == gjson.String {
|
||||
_, _ = b.WriteString("|input=")
|
||||
_, _ = b.WriteString(inp.String())
|
||||
} else if inp.IsArray() {
|
||||
inp.ForEach(func(_, item gjson.Result) bool {
|
||||
role := item.Get("role").String()
|
||||
switch role {
|
||||
case "system", "developer":
|
||||
_, _ = b.WriteString("|system=")
|
||||
if c := item.Get("content"); c.Exists() {
|
||||
_, _ = b.WriteString(normalizeCompatSeedJSON(json.RawMessage(c.Raw)))
|
||||
}
|
||||
case "user":
|
||||
if !firstUserCaptured {
|
||||
_, _ = b.WriteString("|first_user=")
|
||||
if c := item.Get("content"); c.Exists() {
|
||||
_, _ = b.WriteString(normalizeCompatSeedJSON(json.RawMessage(c.Raw)))
|
||||
}
|
||||
firstUserCaptured = true
|
||||
}
|
||||
}
|
||||
if !firstUserCaptured && item.Get("type").String() == "input_text" {
|
||||
_, _ = b.WriteString("|first_user=")
|
||||
if text := item.Get("text").String(); text != "" {
|
||||
_, _ = b.WriteString(text)
|
||||
}
|
||||
firstUserCaptured = true
|
||||
}
|
||||
return true
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
if b.Len() == 0 {
|
||||
return ""
|
||||
}
|
||||
return contentSessionSeedPrefix + b.String()
|
||||
}
|
||||
|
||||
func TestDeriveOpenAIContentSessionSeed_ResponsesAPI_InputTextTypedItem(t *testing.T) {
|
||||
body := []byte(`{
|
||||
"model": "gpt-5.4",
|
||||
|
||||
@@ -118,6 +118,7 @@ func (s *OpenAIGatewayService) failoverOpenAIUpstreamHTTPError(
|
||||
return &UpstreamFailoverError{
|
||||
StatusCode: resp.StatusCode,
|
||||
ResponseBody: respBody,
|
||||
ResponseHeaders: resp.Header.Clone(),
|
||||
RetryableOnSameAccount: account.IsPoolMode() && (account.IsPoolModeRetryableStatus(resp.StatusCode) || isOpenAITransientProcessingError(resp.StatusCode, upstreamMsg, respBody)),
|
||||
}
|
||||
}
|
||||
@@ -197,7 +198,9 @@ func (s *OpenAIGatewayService) sendCCUpstreamRequest(
|
||||
// 账号级请求头覆写(仅 openai api_key 账号启用时生效)
|
||||
account.ApplyHeaderOverrides(upstreamReq.Header)
|
||||
if account.Platform == PlatformGrok {
|
||||
applyGrokCLIHeaders(upstreamReq.Header)
|
||||
if account.IsGrokOAuth() {
|
||||
applyGrokCLIHeaders(upstreamReq.Header)
|
||||
}
|
||||
applyGrokCacheHeaders(upstreamReq.Header, grokCacheIdentity)
|
||||
}
|
||||
|
||||
|
||||
@@ -277,6 +277,13 @@ func (s *OpenAIGatewayService) ForwardAsChatCompletions(
|
||||
// 8. Handle error response with failover
|
||||
if resp.StatusCode >= 400 {
|
||||
respBody, upstreamMsg := s.readOpenAIUpstreamError(resp)
|
||||
if !agentIdentityTaskRecoveryWasTried(ctx) && s.isAgentIdentityAccount(ctx, account) && isAgentIdentityTaskInvalidHTTPResponse(resp.StatusCode, respBody) {
|
||||
expectedTaskID := account.GetCredential("task_id")
|
||||
if err := s.recoverAgentIdentityTask(ctx, account, expectedTaskID); err != nil {
|
||||
return nil, fmt.Errorf("agent identity task recovery failed: %w", err)
|
||||
}
|
||||
return s.ForwardAsChatCompletions(markAgentIdentityTaskRecoveryTried(ctx), c, account, body, promptCacheKey, defaultMappedModel)
|
||||
}
|
||||
if account.Type == AccountTypeAPIKey &&
|
||||
openai_compat.ResolveResponsesSupport(account.Extra) == openai_compat.ResponsesSupportUnknown &&
|
||||
!isResponsesEndpointSupportedByStatus(resp.StatusCode) {
|
||||
|
||||
@@ -9,7 +9,6 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
|
||||
"github.com/Wei-Shaw/sub2api/internal/util/responseheaders"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/tidwall/gjson"
|
||||
@@ -109,7 +108,7 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions(
|
||||
|
||||
// Grok Composer does not accept image_url parts directly, but Grok Build
|
||||
// can describe the images first. Bridge only this exact failure mode.
|
||||
token, tokenKind, err := s.GetAccessToken(ctx, account)
|
||||
token, tokenKind, err := s.getRequestCredential(ctx, c, account)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -162,7 +161,7 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions(
|
||||
}
|
||||
SetActualOpenAIUpstreamEndpoint(c, grokChatRawEndpoint)
|
||||
customUA := account.GetOpenAIUserAgent()
|
||||
if customUA == "" && account.Platform == PlatformGrok {
|
||||
if customUA == "" && account.IsGrokOAuth() {
|
||||
customUA = "sub2api-grok/1.0"
|
||||
}
|
||||
resp, err := s.sendCCUpstreamRequest(ctx, c, account, targetURL, upstreamBody, clientStream, token, customUA, grokCacheIdentity)
|
||||
@@ -189,6 +188,7 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions(
|
||||
return nil, &UpstreamFailoverError{
|
||||
StatusCode: resp.StatusCode,
|
||||
ResponseBody: respBody,
|
||||
ResponseHeaders: resp.Header.Clone(),
|
||||
RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode),
|
||||
}
|
||||
}
|
||||
@@ -201,7 +201,7 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions(
|
||||
}
|
||||
|
||||
if account.Platform == PlatformGrok {
|
||||
s.updateGrokUsageSnapshot(ctx, account, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode))
|
||||
s.updateGrokUsageFromResponse(ctx, account, resp.Header, resp.StatusCode)
|
||||
}
|
||||
|
||||
// 8. Forward response
|
||||
@@ -221,7 +221,7 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions(
|
||||
|
||||
func (s *OpenAIGatewayService) rawChatCompletionsURL(account *Account) (string, error) {
|
||||
if account.Platform == PlatformGrok {
|
||||
targetURL, err := xai.BuildChatCompletionsURL(account.GetGrokBaseURL())
|
||||
targetURL, err := buildGrokChatCompletionsURL(account, s.cfg)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("invalid grok base_url: %w", err)
|
||||
}
|
||||
|
||||
@@ -215,7 +215,15 @@ func (s *OpenAIGatewayService) buildInputTokensUpstreamRequest(
|
||||
return nil, err
|
||||
}
|
||||
req = req.WithContext(WithHTTPUpstreamProfile(req.Context(), HTTPUpstreamProfileOpenAI))
|
||||
req.Header.Set("authorization", "Bearer "+token)
|
||||
authHeaders, err := s.buildOpenAIAuthenticationHeaders(ctx, account, token)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for key, values := range authHeaders {
|
||||
for _, value := range values {
|
||||
req.Header.Add(key, value)
|
||||
}
|
||||
}
|
||||
req.Header.Set("content-type", "application/json")
|
||||
req.Header.Set("accept", "application/json")
|
||||
|
||||
|
||||
@@ -499,6 +499,7 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco
|
||||
var wsResult *OpenAIForwardResult
|
||||
var wsErr error
|
||||
wsLastFailureReason := ""
|
||||
agentTaskRecoveryTried := false
|
||||
wsPrevResponseRecoveryTried := false
|
||||
wsInvalidEncryptedContentRecoveryTried := false
|
||||
recoverPrevResponseNotFound := func(attempt int) bool {
|
||||
@@ -583,6 +584,7 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco
|
||||
startTime,
|
||||
attempt,
|
||||
wsLastFailureReason,
|
||||
&agentTaskRecoveryTried,
|
||||
)
|
||||
if wsErr == nil {
|
||||
break
|
||||
@@ -590,6 +592,10 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco
|
||||
if c != nil && c.Writer != nil && c.Writer.Written() {
|
||||
break
|
||||
}
|
||||
var taskRecoveredErr *agentIdentityTaskRecoveredError
|
||||
if errors.As(wsErr, &taskRecoveredErr) {
|
||||
continue
|
||||
}
|
||||
|
||||
reason, retryable := classifyOpenAIWSReconnectReason(wsErr)
|
||||
if reason != "" {
|
||||
@@ -696,6 +702,7 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco
|
||||
}
|
||||
|
||||
httpInvalidEncryptedContentRetryTried := false
|
||||
agentTaskRecoveryTried := false
|
||||
for {
|
||||
// Build upstream request
|
||||
upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx)
|
||||
@@ -731,6 +738,16 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco
|
||||
upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody))
|
||||
upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
|
||||
upstreamCode := extractUpstreamErrorCode(respBody)
|
||||
if !agentTaskRecoveryTried && s.isAgentIdentityAccount(ctx, account) && isAgentIdentityTaskInvalidHTTPResponse(resp.StatusCode, respBody) {
|
||||
agentTaskRecoveryTried = true
|
||||
expectedTaskID := account.GetCredential("task_id")
|
||||
if err := s.recoverAgentIdentityTask(ctx, account, expectedTaskID); err != nil {
|
||||
return nil, fmt.Errorf("agent identity task recovery failed: %w", err)
|
||||
}
|
||||
continue
|
||||
}
|
||||
respBody = s.redactAgentIdentitySensitiveBody(ctx, account, respBody)
|
||||
resp.Body = io.NopCloser(bytes.NewReader(respBody))
|
||||
if !httpInvalidEncryptedContentRetryTried && resp.StatusCode == http.StatusBadRequest && upstreamCode == "invalid_encrypted_content" {
|
||||
decoded, decodeErr := ensureReqBody()
|
||||
if decodeErr != nil {
|
||||
@@ -881,8 +898,17 @@ func (s *OpenAIGatewayService) buildUpstreamRequest(ctx context.Context, c *gin.
|
||||
}
|
||||
req = req.WithContext(WithHTTPUpstreamProfile(req.Context(), HTTPUpstreamProfileOpenAI))
|
||||
|
||||
// Set authentication header
|
||||
req.Header.Set("authorization", "Bearer "+token)
|
||||
// Build authentication for this request. Agent Identity signs a fresh
|
||||
// assertion here; OAuth/PAT/API-key keep their existing Bearer behavior.
|
||||
authHeaders, err := s.buildOpenAIAuthenticationHeaders(ctx, account, token)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("build openai authentication headers: %w", err)
|
||||
}
|
||||
for key, values := range authHeaders {
|
||||
for _, value := range values {
|
||||
req.Header.Add(key, value)
|
||||
}
|
||||
}
|
||||
|
||||
// Set headers specific to OAuth accounts (ChatGPT internal API)
|
||||
if account.Type == AccountTypeOAuth {
|
||||
|
||||
@@ -11,6 +11,7 @@ import (
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/apicompat"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
|
||||
"github.com/gin-gonic/gin"
|
||||
@@ -25,6 +26,10 @@ const (
|
||||
grokCLIVersion = "0.2.93"
|
||||
grokDefaultResponsesModel = "grok-4.5"
|
||||
grokRateLimitFallbackCooldown = 2 * time.Minute
|
||||
grokRateLimitRepeatCooldown = 10 * time.Minute
|
||||
grokRateLimitSustainedCooldown = 30 * time.Minute
|
||||
grokRateLimitMaxAdaptiveCooldown = time.Hour
|
||||
grokRateLimitBackoffQuietPeriod = time.Hour
|
||||
)
|
||||
|
||||
func (s *OpenAIGatewayService) forwardGrokResponses(
|
||||
@@ -54,14 +59,14 @@ func (s *OpenAIGatewayService) forwardGrokResponses(
|
||||
return nil, fmt.Errorf("apply grok prompt cache identity: %w", err)
|
||||
}
|
||||
|
||||
token, _, err := s.GetAccessToken(ctx, account)
|
||||
token, _, err := s.getRequestCredential(ctx, c, account)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx)
|
||||
defer releaseUpstreamCtx()
|
||||
upstreamReq, err := buildGrokResponsesRequest(upstreamCtx, c, account, patchedBody, token, cacheIdentity)
|
||||
upstreamReq, err := buildGrokResponsesRequest(upstreamCtx, c, account, patchedBody, token, cacheIdentity, s.cfg)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -100,13 +105,14 @@ func (s *OpenAIGatewayService) forwardGrokResponses(
|
||||
return nil, &UpstreamFailoverError{
|
||||
StatusCode: resp.StatusCode,
|
||||
ResponseBody: respBody,
|
||||
ResponseHeaders: resp.Header.Clone(),
|
||||
RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode),
|
||||
}
|
||||
}
|
||||
return s.handleErrorResponse(ctx, resp, c, account, patchedBody, upstreamModel)
|
||||
}
|
||||
|
||||
s.updateGrokUsageSnapshot(ctx, account, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode))
|
||||
s.updateGrokUsageFromResponse(ctx, account, resp.Header, resp.StatusCode)
|
||||
|
||||
var usage *OpenAIUsage
|
||||
var firstTokenMs *int
|
||||
@@ -573,7 +579,7 @@ func (s *OpenAIGatewayService) describeGrokComposerImage(
|
||||
upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx)
|
||||
// Image-description probes are auxiliary requests, not conversation turns.
|
||||
// Do not bind them to the caller's Grok prompt-cache identity.
|
||||
upstreamReq, err := buildGrokResponsesRequest(upstreamCtx, c, account, body, token, "")
|
||||
upstreamReq, err := buildGrokResponsesRequest(upstreamCtx, c, account, body, token, "", s.cfg)
|
||||
releaseUpstreamCtx()
|
||||
if err != nil {
|
||||
return "", OpenAIUsage{}, fmt.Errorf("build grok composer image bridge request: %w", err)
|
||||
@@ -610,13 +616,14 @@ func (s *OpenAIGatewayService) describeGrokComposerImage(
|
||||
return "", OpenAIUsage{}, &UpstreamFailoverError{
|
||||
StatusCode: resp.StatusCode,
|
||||
ResponseBody: respBody,
|
||||
ResponseHeaders: resp.Header.Clone(),
|
||||
RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode),
|
||||
}
|
||||
}
|
||||
return "", OpenAIUsage{}, fmt.Errorf("grok composer image bridge upstream error: %s", upstreamMsg)
|
||||
}
|
||||
|
||||
s.updateGrokUsageSnapshot(ctx, account, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode))
|
||||
s.updateGrokUsageFromResponse(ctx, account, resp.Header, resp.StatusCode)
|
||||
respBody, err := ReadUpstreamResponseBody(resp.Body, s.cfg, c, nil)
|
||||
if err != nil {
|
||||
return "", OpenAIUsage{}, fmt.Errorf("read grok composer image bridge response: %w", err)
|
||||
@@ -738,8 +745,8 @@ func addOpenAIUsage(dst *OpenAIUsage, usage OpenAIUsage) {
|
||||
dst.ImageOutputTokens += usage.ImageOutputTokens
|
||||
}
|
||||
|
||||
func buildGrokResponsesRequest(ctx context.Context, c *gin.Context, account *Account, body []byte, token, cacheIdentity string) (*http.Request, error) {
|
||||
targetURL, err := xai.BuildResponsesURL(account.GetGrokBaseURL())
|
||||
func buildGrokResponsesRequest(ctx context.Context, c *gin.Context, account *Account, body []byte, token, cacheIdentity string, cfg *config.Config) (*http.Request, error) {
|
||||
targetURL, err := buildGrokResponsesURL(account, cfg)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -750,7 +757,9 @@ func buildGrokResponsesRequest(ctx context.Context, c *gin.Context, account *Acc
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Accept", "application/json, text/event-stream")
|
||||
applyGrokCLIHeaders(req.Header)
|
||||
if account.IsGrokOAuth() {
|
||||
applyGrokCLIHeaders(req.Header)
|
||||
}
|
||||
applyGrokCacheHeaders(req.Header, cacheIdentity)
|
||||
if c != nil {
|
||||
if v := c.GetHeader("OpenAI-Beta"); strings.TrimSpace(v) != "" {
|
||||
@@ -776,11 +785,12 @@ func (s *OpenAIGatewayService) updateGrokUsageSnapshot(ctx context.Context, acco
|
||||
}
|
||||
accountID := account.ID
|
||||
now := time.Now()
|
||||
resetAt, hasActiveLimit := grokRateLimitResetAt(snapshot, now)
|
||||
resetAt, hasActiveLimit := grokRateLimitResetAtForAccount(account, snapshot, now)
|
||||
if hasActiveLimit {
|
||||
normalizeGrokExhaustedWindowResets(snapshot, resetAt, now)
|
||||
}
|
||||
critical := snapshot.StatusCode == http.StatusTooManyRequests || hasActiveLimit
|
||||
recovery := isSuccessfulGrokRateLimitRecovery(account, snapshot)
|
||||
critical := snapshot.StatusCode == http.StatusTooManyRequests || hasActiveLimit || recovery
|
||||
if s.codexSnapshotThrottle != nil {
|
||||
allowed := s.codexSnapshotThrottle.Allow(accountID, now)
|
||||
if !critical && !allowed {
|
||||
@@ -806,6 +816,23 @@ func (s *OpenAIGatewayService) updateGrokUsageSnapshot(ctx context.Context, acco
|
||||
// on the passive snapshot scheduler check.
|
||||
if hasActiveLimit {
|
||||
s.rateLimitGrok(stateCtx, account, resetAt)
|
||||
} else if recovery {
|
||||
clearGrokRateLimitAfterRecovery(stateCtx, s.accountRepo, account)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) updateGrokUsageFromResponse(ctx context.Context, account *Account, headers http.Header, statusCode int) {
|
||||
snapshot := parseGrokQuotaSnapshot(headers, statusCode, time.Now())
|
||||
if snapshot != nil {
|
||||
s.updateGrokUsageSnapshot(ctx, account, snapshot)
|
||||
return
|
||||
}
|
||||
// Successful responses are recovery evidence even when the upstream omits
|
||||
// optional quota headers. Do not replace an informative stored snapshot with
|
||||
// an empty one; only clear the exact observed cooldown generation.
|
||||
recoverySnapshot := &xai.QuotaSnapshot{StatusCode: statusCode}
|
||||
if isSuccessfulGrokRateLimitRecovery(account, recoverySnapshot) {
|
||||
clearGrokRateLimitAfterRecovery(ctx, s.accountRepo, account)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -896,6 +923,37 @@ func grokRateLimitResetAt(snapshot *xai.QuotaSnapshot, now time.Time) (time.Time
|
||||
return time.Time{}, false
|
||||
}
|
||||
|
||||
func grokRateLimitResetAtForAccount(account *Account, snapshot *xai.QuotaSnapshot, now time.Time) (time.Time, bool) {
|
||||
resetAt, limited := grokRateLimitResetAt(snapshot, now)
|
||||
if !limited || !isGrokOAuthAccount(account) || snapshot == nil || snapshot.StatusCode != http.StatusTooManyRequests {
|
||||
return resetAt, limited
|
||||
}
|
||||
if account.RateLimitedAt == nil || account.RateLimitResetAt == nil {
|
||||
return resetAt, true
|
||||
}
|
||||
previousResetAt := *account.RateLimitResetAt
|
||||
if previousResetAt.After(now) || now.Sub(previousResetAt) > grokRateLimitBackoffQuietPeriod {
|
||||
return resetAt, true
|
||||
}
|
||||
previousCooldown := previousResetAt.Sub(*account.RateLimitedAt)
|
||||
if previousCooldown <= 0 {
|
||||
return resetAt, true
|
||||
}
|
||||
|
||||
adaptiveCooldown := grokRateLimitRepeatCooldown
|
||||
switch {
|
||||
case previousCooldown >= grokRateLimitSustainedCooldown:
|
||||
adaptiveCooldown = grokRateLimitMaxAdaptiveCooldown
|
||||
case previousCooldown >= grokRateLimitRepeatCooldown:
|
||||
adaptiveCooldown = grokRateLimitSustainedCooldown
|
||||
}
|
||||
adaptiveResetAt := now.Add(adaptiveCooldown)
|
||||
if adaptiveResetAt.After(resetAt) {
|
||||
resetAt = adaptiveResetAt
|
||||
}
|
||||
return resetAt, true
|
||||
}
|
||||
|
||||
func normalizeGrokRateLimitResetAt(account *Account, resetAt, now time.Time) time.Time {
|
||||
if !resetAt.After(now) {
|
||||
resetAt = now.Add(grokRateLimitFallbackCooldown)
|
||||
@@ -910,6 +968,33 @@ type grokRateLimitExtendingRepository interface {
|
||||
SetRateLimitedIfLater(ctx context.Context, id int64, resetAt time.Time) error
|
||||
}
|
||||
|
||||
type grokRateLimitRecoveryRepository interface {
|
||||
ClearRateLimitIfObserved(ctx context.Context, id int64, observedLimitedAt, observedResetAt time.Time) (bool, error)
|
||||
}
|
||||
|
||||
func isSuccessfulGrokRateLimitRecovery(account *Account, snapshot *xai.QuotaSnapshot) bool {
|
||||
return isGrokOAuthAccount(account) &&
|
||||
account.RateLimitedAt != nil &&
|
||||
account.RateLimitResetAt != nil &&
|
||||
snapshot != nil &&
|
||||
snapshot.StatusCode >= http.StatusOK &&
|
||||
snapshot.StatusCode < http.StatusMultipleChoices
|
||||
}
|
||||
|
||||
func clearGrokRateLimitAfterRecovery(ctx context.Context, repo AccountRepository, account *Account) {
|
||||
if repo == nil || account == nil || account.RateLimitedAt == nil || account.RateLimitResetAt == nil || ctx.Err() != nil {
|
||||
return
|
||||
}
|
||||
recoveryRepo, ok := repo.(grokRateLimitRecoveryRepository)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
_, err := recoveryRepo.ClearRateLimitIfObserved(ctx, account.ID, *account.RateLimitedAt, *account.RateLimitResetAt)
|
||||
if err != nil {
|
||||
slog.Warn("grok_rate_limit_recovery_clear_failed", "account_id", account.ID, "error", err)
|
||||
}
|
||||
}
|
||||
|
||||
func persistGrokRateLimit(ctx context.Context, repo AccountRepository, account *Account, resetAt time.Time) {
|
||||
if repo == nil || account == nil || account.ID <= 0 {
|
||||
return
|
||||
|
||||
@@ -10,7 +10,6 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/apicompat"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
@@ -253,12 +252,12 @@ func (s *OpenAIGatewayService) forwardGrokChatCompletionsViaResponses(
|
||||
}
|
||||
responsesBody = updatedBody
|
||||
|
||||
token, _, err := s.GetAccessToken(ctx, account)
|
||||
token, _, err := s.getRequestCredential(ctx, c, account)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get grok access token: %w", err)
|
||||
}
|
||||
upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx)
|
||||
upstreamReq, err := buildGrokResponsesRequest(upstreamCtx, c, account, responsesBody, token, cacheIdentity)
|
||||
upstreamReq, err := buildGrokResponsesRequest(upstreamCtx, c, account, responsesBody, token, cacheIdentity, s.cfg)
|
||||
releaseUpstreamCtx()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("build grok responses bridge request: %w", err)
|
||||
@@ -294,13 +293,14 @@ func (s *OpenAIGatewayService) forwardGrokChatCompletionsViaResponses(
|
||||
return nil, &UpstreamFailoverError{
|
||||
StatusCode: resp.StatusCode,
|
||||
ResponseBody: respBody,
|
||||
ResponseHeaders: resp.Header.Clone(),
|
||||
RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode),
|
||||
}
|
||||
}
|
||||
return s.handleChatCompletionsErrorResponse(resp, c, account, billingModel)
|
||||
}
|
||||
|
||||
s.updateGrokUsageSnapshot(ctx, account, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode))
|
||||
s.updateGrokUsageFromResponse(ctx, account, resp.Header, resp.StatusCode)
|
||||
|
||||
var result *OpenAIForwardResult
|
||||
if clientStream {
|
||||
|
||||
@@ -284,6 +284,7 @@ func TestForwardGrokChatViaResponses429UsesGrokRateLimitPolicy(t *testing.T) {
|
||||
var failoverErr *UpstreamFailoverError
|
||||
require.True(t, errors.As(err, &failoverErr))
|
||||
require.Equal(t, http.StatusTooManyRequests, failoverErr.StatusCode)
|
||||
require.Equal(t, "45", failoverErr.ResponseHeaders.Get("Retry-After"))
|
||||
require.Equal(t, xai.DefaultCLIBaseURL+"/responses", upstream.lastReq.URL.String())
|
||||
require.Equal(t, grokChatResponsesEndpoint, GetActualOpenAIUpstreamEndpoint(c))
|
||||
require.Equal(t, 1, repo.rateLimitedCalls)
|
||||
@@ -292,6 +293,45 @@ func TestForwardGrokChatViaResponses429UsesGrokRateLimitPolicy(t *testing.T) {
|
||||
require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
|
||||
}
|
||||
|
||||
func TestForwardGrokRawChat429PreservesRetryAfter(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
body := []byte(`{"model":"grok","messages":[{"role":"user","content":"hi"}],"stream":false,"stop":"done"}`)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, grokChatRawEndpoint, bytes.NewReader(body))
|
||||
c.Set("api_key", &APIKey{ID: 7551})
|
||||
|
||||
account := grokChatBridgeTestAccount(755)
|
||||
account.Credentials["expires_at"] = time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339)
|
||||
repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
|
||||
accountsByID: map[int64]*Account{account.ID: account},
|
||||
}}
|
||||
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
||||
StatusCode: http.StatusTooManyRequests,
|
||||
Header: http.Header{
|
||||
"Content-Type": []string{"application/json"},
|
||||
"Retry-After": []string{"45"},
|
||||
},
|
||||
Body: io.NopCloser(strings.NewReader(`{"error":{"message":"rate limited"}}`)),
|
||||
}}
|
||||
svc := &OpenAIGatewayService{
|
||||
httpUpstream: upstream,
|
||||
grokTokenProvider: NewGrokTokenProvider(repo, nil),
|
||||
accountRepo: repo,
|
||||
}
|
||||
|
||||
result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "")
|
||||
|
||||
require.Error(t, err)
|
||||
require.Nil(t, result)
|
||||
var failoverErr *UpstreamFailoverError
|
||||
require.ErrorAs(t, err, &failoverErr)
|
||||
require.Equal(t, http.StatusTooManyRequests, failoverErr.StatusCode)
|
||||
require.Equal(t, "45", failoverErr.ResponseHeaders.Get("Retry-After"))
|
||||
require.Equal(t, xai.DefaultCLIBaseURL+"/chat/completions", upstream.lastReq.URL.String())
|
||||
}
|
||||
|
||||
func TestForwardGrokRawChatErrorRecordsActualEndpoint(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
@@ -329,11 +369,14 @@ func grokChatBridgeTestAccount(id int64) *Account {
|
||||
Name: "grok-cache-bridge",
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Status: StatusActive,
|
||||
Schedulable: true,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "access-token",
|
||||
"expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
|
||||
"base_url": xai.DefaultCLIBaseURL,
|
||||
"access_token": "access-token",
|
||||
"refresh_token": "refresh-token",
|
||||
"expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
|
||||
"base_url": xai.DefaultCLIBaseURL,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
@@ -256,10 +256,10 @@ func TestBuildGrokResponsesRequestUsesAccountBaseURLAndBearerToken(t *testing.T)
|
||||
},
|
||||
}
|
||||
|
||||
req, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "access-token", "isolated-cache-id")
|
||||
req, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "access-token", "isolated-cache-id", nil)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, http.MethodPost, req.Method)
|
||||
require.Equal(t, "https://xai.test/v1/responses", req.URL.String())
|
||||
require.Equal(t, xai.DefaultCLIBaseURL+"/responses", req.URL.String())
|
||||
require.Equal(t, "Bearer access-token", req.Header.Get("Authorization"))
|
||||
require.Equal(t, "application/json", req.Header.Get("Content-Type"))
|
||||
require.Contains(t, req.Header.Get("Accept"), "text/event-stream")
|
||||
@@ -280,10 +280,12 @@ func TestBuildGrokResponsesRequestAllowsPublicAPIKeyBaseURLByDefault(t *testing.
|
||||
},
|
||||
}
|
||||
|
||||
req, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "api-key", "")
|
||||
req, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "api-key", "", nil)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "https://grok.example.test/v1/responses", req.URL.String())
|
||||
require.Equal(t, "Bearer api-key", req.Header.Get("Authorization"))
|
||||
require.Empty(t, req.Header.Get("X-Grok-Client-Version"))
|
||||
require.NotEqual(t, grokUpstreamUserAgent, req.Header.Get("User-Agent"))
|
||||
}
|
||||
|
||||
func TestBuildGrokResponsesRequestPinsOAuthCustomBaseURLByDefault(t *testing.T) {
|
||||
@@ -297,7 +299,7 @@ func TestBuildGrokResponsesRequestPinsOAuthCustomBaseURLByDefault(t *testing.T)
|
||||
},
|
||||
}
|
||||
|
||||
req, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "access-token", "")
|
||||
req, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "access-token", "", nil)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, xai.DefaultCLIBaseURL+"/responses", req.URL.String())
|
||||
}
|
||||
@@ -425,7 +427,8 @@ func TestForwardGrokMediaImagesGenerationNormalizesImagineAlias(t *testing.T) {
|
||||
require.Equal(t, http.MethodPost, upstream.lastReq.Method)
|
||||
require.Equal(t, "Bearer api-key", upstream.lastReq.Header.Get("Authorization"))
|
||||
require.Equal(t, "application/json", upstream.lastReq.Header.Get("Content-Type"))
|
||||
require.Equal(t, grokCLIVersion, upstream.lastReq.Header.Get("X-Grok-Client-Version"))
|
||||
require.Empty(t, upstream.lastReq.Header.Get("X-Grok-Client-Version"))
|
||||
require.NotEqual(t, grokUpstreamUserAgent, upstream.lastReq.Header.Get("User-Agent"))
|
||||
require.JSONEq(t, `{"model":"grok-imagine-image-quality","prompt":"draw a cat"}`, string(upstream.lastBody))
|
||||
require.Equal(t, http.StatusOK, recorder.Code)
|
||||
require.JSONEq(t, `{"data":[]}`, recorder.Body.String())
|
||||
@@ -611,7 +614,7 @@ func TestForwardGrokMediaVideoGenerationPreservesImageToVideoModel(t *testing.T)
|
||||
require.Equal(t, VideoBillingDefaultDurationSeconds, result.VideoDurationSeconds)
|
||||
}
|
||||
|
||||
func TestForwardGrokMediaOAuthImageToVideoUsesOfficialAPIForLargeBody(t *testing.T) {
|
||||
func TestForwardGrokMediaOAuthImageToVideoKeepsCLIGatewayForLargeBody(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
@@ -626,10 +629,14 @@ func TestForwardGrokMediaOAuthImageToVideoUsesOfficialAPIForLargeBody(t *testing
|
||||
Name: "grok-oauth",
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Status: StatusActive,
|
||||
Schedulable: true,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "oauth-access-token",
|
||||
"base_url": xai.DefaultCLIBaseURL,
|
||||
"access_token": "oauth-access-token",
|
||||
"refresh_token": "oauth-refresh-token",
|
||||
"expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
|
||||
"base_url": xai.DefaultCLIBaseURL,
|
||||
},
|
||||
}
|
||||
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
||||
@@ -639,11 +646,11 @@ func TestForwardGrokMediaOAuthImageToVideoUsesOfficialAPIForLargeBody(t *testing
|
||||
},
|
||||
Body: io.NopCloser(strings.NewReader(`{"request_id":"video-request-oauth"}`)),
|
||||
}}
|
||||
svc := &OpenAIGatewayService{httpUpstream: upstream}
|
||||
svc := &OpenAIGatewayService{httpUpstream: upstream, grokTokenProvider: NewGrokTokenProvider(nil, nil)}
|
||||
|
||||
_, err := svc.ForwardGrokMedia(context.Background(), c, account, GrokMediaEndpointVideosGenerations, "", body, "application/json")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, xai.DefaultBaseURL+"/videos/generations", upstream.lastReq.URL.String())
|
||||
require.Equal(t, xai.DefaultCLIBaseURL+"/videos/generations", upstream.lastReq.URL.String())
|
||||
require.Equal(t, "data:image/png;base64,"+imageData, gjson.GetBytes(upstream.lastBody, "image.image_url").String())
|
||||
}
|
||||
|
||||
@@ -681,6 +688,8 @@ func TestForwardGrokMediaVideoStatusUsesGETWithoutBody(t *testing.T) {
|
||||
require.Equal(t, "https://xai.test/v1/videos/request-123", upstream.lastReq.URL.String())
|
||||
require.Equal(t, http.MethodGet, upstream.lastReq.Method)
|
||||
require.Equal(t, "Bearer api-key", upstream.lastReq.Header.Get("Authorization"))
|
||||
require.Empty(t, upstream.lastReq.Header.Get("X-Grok-Client-Version"))
|
||||
require.NotEqual(t, grokUpstreamUserAgent, upstream.lastReq.Header.Get("User-Agent"))
|
||||
require.Empty(t, upstream.lastReq.Header.Get("Content-Type"))
|
||||
require.Empty(t, upstream.lastBody)
|
||||
require.Equal(t, http.StatusOK, recorder.Code)
|
||||
@@ -792,6 +801,54 @@ func TestForwardGrokMedia429ReconcilesRateLimitBeforeCustomErrorBypass(t *testin
|
||||
require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
|
||||
}
|
||||
|
||||
func TestGrokMedia429FailoverPreservesRetryAfter(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil)
|
||||
account := &Account{
|
||||
ID: 641, Name: "grok-oauth", Platform: PlatformGrok, Type: AccountTypeOAuth,
|
||||
Status: StatusActive, Schedulable: true,
|
||||
Credentials: map[string]any{
|
||||
"custom_error_codes_enabled": true,
|
||||
"custom_error_codes": []any{float64(http.StatusTooManyRequests)},
|
||||
},
|
||||
}
|
||||
repo := &grokQuotaAccountRepo{}
|
||||
svc := &OpenAIGatewayService{accountRepo: repo}
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusTooManyRequests,
|
||||
Header: http.Header{"Retry-After": []string{"45"}},
|
||||
Body: io.NopCloser(strings.NewReader(`{"error":{"message":"rate limited"}}`)),
|
||||
}
|
||||
|
||||
result, err := svc.handleGrokMediaErrorResponse(context.Background(), resp, c, account, "request-id", "grok-imagine")
|
||||
|
||||
require.Nil(t, result)
|
||||
var failoverErr *UpstreamFailoverError
|
||||
require.ErrorAs(t, err, &failoverErr)
|
||||
require.Equal(t, http.StatusTooManyRequests, failoverErr.StatusCode)
|
||||
require.Equal(t, "45", failoverErr.ResponseHeaders.Get("Retry-After"))
|
||||
}
|
||||
|
||||
func healthyGrokOAuthGatewayTestAccount(id int64, token string) *Account {
|
||||
return &Account{
|
||||
ID: id,
|
||||
Name: "grok",
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Status: StatusActive,
|
||||
Schedulable: true,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": token,
|
||||
"refresh_token": "refresh-token",
|
||||
"expires_at": time.Now().Add(2 * grokTokenRefreshSkew).UTC().Format(time.RFC3339),
|
||||
"base_url": xai.DefaultCLIBaseURL,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func TestForwardAsChatCompletionsForGrokStopFallsBackToXAIChatCompletions(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
@@ -801,18 +858,7 @@ func TestForwardAsChatCompletionsForGrokStopFallsBackToXAIChatCompletions(t *tes
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
|
||||
c.Set("api_key", &APIKey{ID: 5101})
|
||||
|
||||
account := &Account{
|
||||
ID: 51,
|
||||
Name: "grok",
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "access-token",
|
||||
"expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
|
||||
"base_url": xai.DefaultCLIBaseURL,
|
||||
},
|
||||
}
|
||||
account := healthyGrokOAuthGatewayTestAccount(51, "access-token")
|
||||
repo := &grokQuotaAccountRepo{
|
||||
mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
|
||||
accountsByID: map[int64]*Account{51: account},
|
||||
@@ -864,18 +910,7 @@ func TestForwardGrokResponsesStreamingDefaultsEmptyModelTo45AndSnapshots(t *test
|
||||
c.Request.Header.Set("OpenAI-Beta", "responses=experimental")
|
||||
c.Set("api_key", &APIKey{ID: 5201})
|
||||
|
||||
account := &Account{
|
||||
ID: 52,
|
||||
Name: "grok",
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "access-token",
|
||||
"expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
|
||||
"base_url": xai.DefaultCLIBaseURL,
|
||||
},
|
||||
}
|
||||
account := healthyGrokOAuthGatewayTestAccount(52, "access-token")
|
||||
repo := &grokQuotaAccountRepo{
|
||||
mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
|
||||
accountsByID: map[int64]*Account{52: account},
|
||||
@@ -968,12 +1003,46 @@ func TestForwardGrokResponsesAPIKeyUsesXAIResponses(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "https://api.x.ai/v1/responses", upstream.lastReq.URL.String())
|
||||
require.Equal(t, "Bearer xai-test-key", upstream.lastReq.Header.Get("Authorization"))
|
||||
require.Empty(t, upstream.lastReq.Header.Get("X-Grok-Client-Version"))
|
||||
require.NotEqual(t, grokUpstreamUserAgent, upstream.lastReq.Header.Get("User-Agent"))
|
||||
require.Equal(t, "grok-4.5", gjson.GetBytes(upstream.lastBody, "model").String())
|
||||
require.Equal(t, "resp_grok_api_key", result.ResponseID)
|
||||
require.Equal(t, 2, result.Usage.InputTokens)
|
||||
require.Equal(t, 1, result.Usage.OutputTokens)
|
||||
}
|
||||
|
||||
func TestForwardAsChatCompletionsForGrokAPIKeyUsesConfiguredRawEndpointWithoutOAuthIdentity(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
body := []byte(`{"model":"grok","messages":[{"role":"user","content":"hi"}],"stream":false}`)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
account := &Account{
|
||||
ID: 706,
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeAPIKey,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"api_key": "third-party-key",
|
||||
"base_url": "https://grok.example.test/v1",
|
||||
},
|
||||
}
|
||||
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(`{"id":"chatcmpl","object":"chat.completion","model":"grok-4.5","choices":[],"usage":{"prompt_tokens":1,"completion_tokens":1}}`)),
|
||||
}}
|
||||
svc := &OpenAIGatewayService{cfg: rawChatCompletionsTestConfig(), httpUpstream: upstream}
|
||||
|
||||
_, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "https://grok.example.test/v1/chat/completions", upstream.lastReq.URL.String())
|
||||
require.Equal(t, "Bearer third-party-key", upstream.lastReq.Header.Get("Authorization"))
|
||||
require.Empty(t, upstream.lastReq.Header.Get("X-Grok-Client-Version"))
|
||||
require.NotEqual(t, grokUpstreamUserAgent, upstream.lastReq.Header.Get("User-Agent"))
|
||||
}
|
||||
|
||||
func TestAccountTestServiceGrokAPIKeyUsesXAIResponses(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
@@ -1008,6 +1077,41 @@ func TestAccountTestServiceGrokAPIKeyUsesXAIResponses(t *testing.T) {
|
||||
require.Contains(t, recorder.Body.String(), `"type":"test_complete"`)
|
||||
}
|
||||
|
||||
func TestAccountTestServiceGrokAPIKeyAllowsConfiguredHTTPWhenGlobalPolicyDoes(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
account := &Account{
|
||||
ID: 55,
|
||||
Name: "grok-api-key-http",
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeAPIKey,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"api_key": "third-party-key",
|
||||
"base_url": "http://grok.example.test/v1",
|
||||
},
|
||||
}
|
||||
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
||||
Body: io.NopCloser(strings.NewReader(
|
||||
"data: {\"type\":\"response.output_text.delta\",\"delta\":\"ok\"}\n\n" +
|
||||
"data: {\"type\":\"response.completed\"}\n\n",
|
||||
)),
|
||||
}}
|
||||
svc := &AccountTestService{cfg: rawChatCompletionsTestConfig(), httpUpstream: upstream}
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/55/test", nil)
|
||||
|
||||
err := svc.testGrokAccountConnection(c, account, "grok")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "http://grok.example.test/v1/responses", upstream.lastReq.URL.String())
|
||||
require.Equal(t, "Bearer third-party-key", upstream.lastReq.Header.Get("Authorization"))
|
||||
require.Empty(t, upstream.lastReq.Header.Get("X-Grok-Client-Version"))
|
||||
require.Contains(t, recorder.Body.String(), `"type":"test_complete"`)
|
||||
}
|
||||
|
||||
func TestForwardAsChatCompletionsForGrokStreamingUsesRawXAIChatCompletions(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
@@ -1017,18 +1121,7 @@ func TestForwardAsChatCompletionsForGrokStreamingUsesRawXAIChatCompletions(t *te
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
|
||||
account := &Account{
|
||||
ID: 53,
|
||||
Name: "grok",
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "access-token",
|
||||
"expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
|
||||
"base_url": xai.DefaultCLIBaseURL,
|
||||
},
|
||||
}
|
||||
account := healthyGrokOAuthGatewayTestAccount(53, "access-token")
|
||||
repo := &grokQuotaAccountRepo{
|
||||
mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
|
||||
accountsByID: map[int64]*Account{53: account},
|
||||
@@ -1085,22 +1178,16 @@ func TestForwardGrokResponsesNonStreamingUsesCacheIdentityAndCachedUsage(t *test
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
c.Set("api_key", &APIKey{ID: 5202})
|
||||
|
||||
account := &Account{
|
||||
ID: 56,
|
||||
Name: "grok",
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "access-token",
|
||||
"expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
|
||||
"base_url": xai.DefaultCLIBaseURL,
|
||||
},
|
||||
}
|
||||
account := healthyGrokOAuthGatewayTestAccount(56, "access-token")
|
||||
observedResetAt := time.Now().Add(-time.Second).UTC().Truncate(time.Second)
|
||||
observedLimitedAt := observedResetAt.Add(-grokRateLimitRepeatCooldown)
|
||||
account.RateLimitedAt = &observedLimitedAt
|
||||
account.RateLimitResetAt = &observedResetAt
|
||||
repo := &grokQuotaAccountRepo{
|
||||
mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
|
||||
accountsByID: map[int64]*Account{56: account},
|
||||
},
|
||||
recoveryClearResult: true,
|
||||
}
|
||||
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
@@ -1133,6 +1220,9 @@ func TestForwardGrokResponsesNonStreamingUsesCacheIdentityAndCachedUsage(t *test
|
||||
require.False(t, gjson.GetBytes(upstream.lastBody, "tools").Exists())
|
||||
require.False(t, gjson.GetBytes(upstream.lastBody, "tool_choice").Exists())
|
||||
require.Equal(t, "resp_grok_non_stream", gjson.Get(recorder.Body.String(), "id").String())
|
||||
require.Equal(t, 1, repo.recoveryClearCalls)
|
||||
require.Equal(t, observedLimitedAt, repo.recoveryObservedAt)
|
||||
require.Equal(t, observedResetAt, repo.recoveryObservedReset)
|
||||
}
|
||||
|
||||
func TestForwardGrokResponsesFailoverKeepsCacheIdentityAcrossAccounts(t *testing.T) {
|
||||
@@ -1145,18 +1235,9 @@ func TestForwardGrokResponsesFailoverKeepsCacheIdentityAcrossAccounts(t *testing
|
||||
c.Set("api_key", &APIKey{ID: 5203})
|
||||
|
||||
newAccount := func(id int64, token string) *Account {
|
||||
return &Account{
|
||||
ID: id,
|
||||
Name: fmt.Sprintf("grok-%d", id),
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": token,
|
||||
"expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
|
||||
"base_url": xai.DefaultCLIBaseURL,
|
||||
},
|
||||
}
|
||||
account := healthyGrokOAuthGatewayTestAccount(id, token)
|
||||
account.Name = fmt.Sprintf("grok-%d", id)
|
||||
return account
|
||||
}
|
||||
firstAccount := newAccount(58, "access-token-a")
|
||||
secondAccount := newAccount(59, "access-token-b")
|
||||
@@ -1213,18 +1294,7 @@ func TestForwardAsChatCompletionsForGrokStreamingStopFallsBackToRawXAIChatComple
|
||||
c.Request.Header.Set(grokConversationIDHeader, "native-client-conversation")
|
||||
c.Set("api_key", &APIKey{ID: 5301})
|
||||
|
||||
account := &Account{
|
||||
ID: 53,
|
||||
Name: "grok",
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "access-token",
|
||||
"expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
|
||||
"base_url": xai.DefaultCLIBaseURL,
|
||||
},
|
||||
}
|
||||
account := healthyGrokOAuthGatewayTestAccount(53, "access-token")
|
||||
repo := &grokQuotaAccountRepo{
|
||||
mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
|
||||
accountsByID: map[int64]*Account{53: account},
|
||||
@@ -1284,18 +1354,7 @@ func TestForwardAsChatCompletionsForGrokComposerBridgesImageInput(t *testing.T)
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
c.Set("api_key", &APIKey{ID: 5501})
|
||||
|
||||
account := &Account{
|
||||
ID: 55,
|
||||
Name: "grok",
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "access-token",
|
||||
"expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
|
||||
"base_url": xai.DefaultCLIBaseURL,
|
||||
},
|
||||
}
|
||||
account := healthyGrokOAuthGatewayTestAccount(55, "access-token")
|
||||
repo := &grokQuotaAccountRepo{
|
||||
mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
|
||||
accountsByID: map[int64]*Account{55: account},
|
||||
@@ -1358,18 +1417,7 @@ func TestForwardAsAnthropicForGrokUsesXAIResponses(t *testing.T) {
|
||||
c.Request.Header.Set("OpenAI-Beta", "grok-experimental")
|
||||
c.Request.Header.Set("originator", "opencode")
|
||||
|
||||
account := &Account{
|
||||
ID: 54,
|
||||
Name: "grok",
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "access-token",
|
||||
"expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
|
||||
"base_url": xai.DefaultCLIBaseURL,
|
||||
},
|
||||
}
|
||||
account := healthyGrokOAuthGatewayTestAccount(54, "access-token")
|
||||
repo := &grokQuotaAccountRepo{
|
||||
mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
|
||||
accountsByID: map[int64]*Account{54: account},
|
||||
@@ -1419,18 +1467,7 @@ func TestForwardAsAnthropicForGrokStreamingPreservesCacheUsage(t *testing.T) {
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(body))
|
||||
c.Set("api_key", &APIKey{ID: 5402})
|
||||
|
||||
account := &Account{
|
||||
ID: 57,
|
||||
Name: "grok",
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "access-token",
|
||||
"expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
|
||||
"base_url": xai.DefaultCLIBaseURL,
|
||||
},
|
||||
}
|
||||
account := healthyGrokOAuthGatewayTestAccount(57, "access-token")
|
||||
repo := &grokQuotaAccountRepo{
|
||||
mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
|
||||
accountsByID: map[int64]*Account{57: account},
|
||||
@@ -1571,6 +1608,102 @@ func TestHandleGrokAccountUpstreamError429UsesFallbackReset(t *testing.T) {
|
||||
require.Zero(t, repo.tempUnschedCalls)
|
||||
}
|
||||
|
||||
func TestGrokRateLimitResetAtForAccountEscalatesRepeated429s(t *testing.T) {
|
||||
now := time.Now().UTC().Truncate(time.Second)
|
||||
retryAfter := 45
|
||||
snapshot := &xai.QuotaSnapshot{
|
||||
StatusCode: http.StatusTooManyRequests,
|
||||
RetryAfterSeconds: &retryAfter,
|
||||
UpdatedAt: now.Format(time.RFC3339),
|
||||
}
|
||||
tests := []struct {
|
||||
name string
|
||||
previousCooldown time.Duration
|
||||
wantCooldown time.Duration
|
||||
}{
|
||||
{name: "repeat after short boundary", previousCooldown: 45 * time.Second, wantCooldown: grokRateLimitRepeatCooldown},
|
||||
{name: "sustained repeat", previousCooldown: grokRateLimitRepeatCooldown, wantCooldown: grokRateLimitSustainedCooldown},
|
||||
{name: "capped repeat", previousCooldown: grokRateLimitSustainedCooldown, wantCooldown: grokRateLimitMaxAdaptiveCooldown},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
previousReset := now.Add(-time.Second)
|
||||
previousLimited := previousReset.Add(-tt.previousCooldown)
|
||||
account := &Account{
|
||||
ID: 630,
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
RateLimitedAt: &previousLimited,
|
||||
RateLimitResetAt: &previousReset,
|
||||
}
|
||||
|
||||
resetAt, limited := grokRateLimitResetAtForAccount(account, snapshot, now)
|
||||
|
||||
require.True(t, limited)
|
||||
require.WithinDuration(t, now.Add(tt.wantCooldown), resetAt, time.Second)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGrokRateLimitResetAtForAccountPreservesAuthoritativeAndQuietRecovery(t *testing.T) {
|
||||
now := time.Now().UTC().Truncate(time.Second)
|
||||
retryAfter := 45
|
||||
previousReset := now.Add(-grokRateLimitBackoffQuietPeriod - time.Second)
|
||||
previousLimited := previousReset.Add(-grokRateLimitSustainedCooldown)
|
||||
account := &Account{
|
||||
ID: 631,
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
RateLimitedAt: &previousLimited,
|
||||
RateLimitResetAt: &previousReset,
|
||||
}
|
||||
snapshot := &xai.QuotaSnapshot{
|
||||
StatusCode: http.StatusTooManyRequests,
|
||||
RetryAfterSeconds: &retryAfter,
|
||||
UpdatedAt: now.Format(time.RFC3339),
|
||||
}
|
||||
|
||||
resetAt, limited := grokRateLimitResetAtForAccount(account, snapshot, now)
|
||||
require.True(t, limited)
|
||||
require.WithinDuration(t, now.Add(45*time.Second), resetAt, time.Second)
|
||||
|
||||
authoritativeReset := now.Add(2 * time.Hour)
|
||||
remaining := int64(0)
|
||||
snapshot.Requests = &xai.QuotaWindow{Remaining: &remaining, ResetUnix: grokInt64PtrForTest(authoritativeReset.Unix())}
|
||||
recentReset := now.Add(-time.Second)
|
||||
recentLimited := recentReset.Add(-grokRateLimitSustainedCooldown)
|
||||
account.RateLimitResetAt = &recentReset
|
||||
account.RateLimitedAt = &recentLimited
|
||||
|
||||
resetAt, limited = grokRateLimitResetAtForAccount(account, snapshot, now)
|
||||
require.True(t, limited)
|
||||
require.WithinDuration(t, authoritativeReset, resetAt, time.Second)
|
||||
}
|
||||
|
||||
func TestGrokRateLimitResetAtForAccountLeavesAPIKey429PolicyUnchanged(t *testing.T) {
|
||||
now := time.Now().UTC().Truncate(time.Second)
|
||||
retryAfter := 45
|
||||
previousReset := now.Add(-time.Second)
|
||||
previousLimited := previousReset.Add(-grokRateLimitSustainedCooldown)
|
||||
account := &Account{
|
||||
ID: 632,
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeAPIKey,
|
||||
RateLimitedAt: &previousLimited,
|
||||
RateLimitResetAt: &previousReset,
|
||||
}
|
||||
snapshot := &xai.QuotaSnapshot{
|
||||
StatusCode: http.StatusTooManyRequests,
|
||||
RetryAfterSeconds: &retryAfter,
|
||||
UpdatedAt: now.Format(time.RFC3339),
|
||||
}
|
||||
|
||||
resetAt, limited := grokRateLimitResetAtForAccount(account, snapshot, now)
|
||||
require.True(t, limited)
|
||||
require.WithinDuration(t, now.Add(45*time.Second), resetAt, time.Second)
|
||||
}
|
||||
|
||||
func TestGrokRateLimitResetAtUsesFutureWindowAfterRetryAfterExpires(t *testing.T) {
|
||||
now := time.Now().UTC().Truncate(time.Second)
|
||||
observedAt := now.Add(-2 * time.Minute)
|
||||
@@ -1672,6 +1805,72 @@ func TestUpdateGrokUsageSnapshotAvailableSuccessDoesNotSetRateLimited(t *testing
|
||||
require.Zero(t, repo.rateLimitedCalls)
|
||||
}
|
||||
|
||||
func TestUpdateGrokUsageFromResponseHeaderlessSuccessClearsObservedCooldown(t *testing.T) {
|
||||
now := time.Now().UTC().Truncate(time.Second)
|
||||
limitedAt := now.Add(-grokRateLimitRepeatCooldown)
|
||||
observedResetAt := now.Add(-time.Second)
|
||||
account := &Account{
|
||||
ID: 660,
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
RateLimitedAt: &limitedAt,
|
||||
RateLimitResetAt: &observedResetAt,
|
||||
}
|
||||
repo := &grokQuotaAccountRepo{recoveryClearResult: true}
|
||||
svc := &OpenAIGatewayService{
|
||||
accountRepo: repo,
|
||||
codexSnapshotThrottle: newAccountWriteThrottle(time.Hour),
|
||||
}
|
||||
|
||||
svc.updateGrokUsageFromResponse(context.Background(), account, nil, http.StatusOK)
|
||||
|
||||
require.Zero(t, repo.updateCalls, "headerless success must not overwrite an informative quota snapshot")
|
||||
require.Equal(t, 1, repo.recoveryClearCalls)
|
||||
require.Equal(t, limitedAt, repo.recoveryObservedAt)
|
||||
require.Equal(t, observedResetAt, repo.recoveryObservedReset)
|
||||
require.Same(t, &observedResetAt, account.RateLimitResetAt, "shared account snapshots must not be mutated in place")
|
||||
}
|
||||
|
||||
func TestUpdateGrokUsageFromResponseRecoveryRespectsCancellationAndAPIKeyBoundary(t *testing.T) {
|
||||
now := time.Now().UTC().Truncate(time.Second)
|
||||
observedResetAt := now.Add(-time.Second)
|
||||
observedLimitedAt := observedResetAt.Add(-grokRateLimitRepeatCooldown)
|
||||
|
||||
t.Run("parent cancellation does not mutate account state", func(t *testing.T) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
account := &Account{
|
||||
ID: 661,
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
RateLimitedAt: &observedLimitedAt,
|
||||
RateLimitResetAt: &observedResetAt,
|
||||
}
|
||||
repo := &grokQuotaAccountRepo{recoveryClearResult: true}
|
||||
svc := &OpenAIGatewayService{accountRepo: repo}
|
||||
|
||||
svc.updateGrokUsageFromResponse(ctx, account, nil, http.StatusOK)
|
||||
|
||||
require.Zero(t, repo.recoveryClearCalls)
|
||||
})
|
||||
|
||||
t.Run("API key success does not alter OAuth cooldown state", func(t *testing.T) {
|
||||
account := &Account{
|
||||
ID: 662,
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeAPIKey,
|
||||
RateLimitedAt: &observedLimitedAt,
|
||||
RateLimitResetAt: &observedResetAt,
|
||||
}
|
||||
repo := &grokQuotaAccountRepo{recoveryClearResult: true}
|
||||
svc := &OpenAIGatewayService{accountRepo: repo}
|
||||
|
||||
svc.updateGrokUsageFromResponse(context.Background(), account, nil, http.StatusOK)
|
||||
|
||||
require.Zero(t, repo.recoveryClearCalls)
|
||||
})
|
||||
}
|
||||
|
||||
func TestUpdateGrokUsageSnapshotExhaustedSuccessWithoutResetUsesFallback(t *testing.T) {
|
||||
repo := &grokQuotaAccountRepo{}
|
||||
svc := &OpenAIGatewayService{accountRepo: repo}
|
||||
|
||||
@@ -14,7 +14,6 @@ import (
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/claude"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/openai_compat"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
|
||||
"github.com/Wei-Shaw/sub2api/internal/util/responseheaders"
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
@@ -258,7 +257,7 @@ func (s *OpenAIGatewayService) ForwardAsAnthropic(
|
||||
}
|
||||
|
||||
// 5. Get access token
|
||||
token, _, err := s.GetAccessToken(ctx, account)
|
||||
token, _, err := s.getRequestCredential(ctx, c, account)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get access token: %w", err)
|
||||
}
|
||||
@@ -273,7 +272,7 @@ func (s *OpenAIGatewayService) ForwardAsAnthropic(
|
||||
upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx)
|
||||
var upstreamReq *http.Request
|
||||
if account.Platform == PlatformGrok {
|
||||
upstreamReq, err = buildGrokResponsesRequest(upstreamCtx, c, account, responsesBody, token, grokCacheIdentity)
|
||||
upstreamReq, err = buildGrokResponsesRequest(upstreamCtx, c, account, responsesBody, token, grokCacheIdentity, s.cfg)
|
||||
} else {
|
||||
upstreamReq, err = s.buildUpstreamRequest(upstreamCtx, c, account, responsesBody, token, isStream, promptCacheKey, false)
|
||||
}
|
||||
@@ -324,6 +323,13 @@ func (s *OpenAIGatewayService) ForwardAsAnthropic(
|
||||
// 8. Handle error response with failover
|
||||
if resp.StatusCode >= 400 {
|
||||
respBody, upstreamMsg := s.readOpenAIUpstreamError(resp)
|
||||
if !agentIdentityTaskRecoveryWasTried(ctx) && s.isAgentIdentityAccount(ctx, account) && isAgentIdentityTaskInvalidHTTPResponse(resp.StatusCode, respBody) {
|
||||
expectedTaskID := account.GetCredential("task_id")
|
||||
if err := s.recoverAgentIdentityTask(ctx, account, expectedTaskID); err != nil {
|
||||
return nil, fmt.Errorf("agent identity task recovery failed: %w", err)
|
||||
}
|
||||
return s.ForwardAsAnthropic(markAgentIdentityTaskRecoveryTried(ctx), c, account, body, promptCacheKey, defaultMappedModel)
|
||||
}
|
||||
if previousResponseID != "" && (isOpenAICompatPreviousResponseNotFound(resp.StatusCode, upstreamMsg, respBody) || isOpenAICompatPreviousResponseUnsupported(resp.StatusCode, upstreamMsg, respBody)) {
|
||||
if isOpenAICompatPreviousResponseUnsupported(resp.StatusCode, upstreamMsg, respBody) {
|
||||
s.disableOpenAICompatSessionContinuation(ctx, c, account, promptCacheKey)
|
||||
@@ -344,7 +350,7 @@ func (s *OpenAIGatewayService) ForwardAsAnthropic(
|
||||
return s.handleAnthropicErrorResponse(resp, c, account, billingModel)
|
||||
}
|
||||
if account.Platform == PlatformGrok && account.Type == AccountTypeOAuth && !account.IsShadow() {
|
||||
s.updateGrokUsageSnapshot(ctx, account, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode))
|
||||
s.updateGrokUsageFromResponse(ctx, account, resp.Header, resp.StatusCode)
|
||||
}
|
||||
|
||||
if account.Type == AccountTypeOAuth && promptCacheKey != "" {
|
||||
|
||||
@@ -11,6 +11,7 @@ import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"sort"
|
||||
"strings"
|
||||
@@ -161,13 +162,6 @@ func (s *OpenAIGatewayService) forwardOpenAIPassthrough(
|
||||
return nil, err
|
||||
}
|
||||
|
||||
upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx)
|
||||
upstreamReq, err := s.buildUpstreamRequestOpenAIPassthrough(upstreamCtx, c, account, body, token)
|
||||
releaseUpstreamCtx()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
proxyURL := ""
|
||||
if account.ProxyID != nil && account.Proxy != nil {
|
||||
proxyURL = account.Proxy.URL()
|
||||
@@ -177,26 +171,51 @@ func (s *OpenAIGatewayService) forwardOpenAIPassthrough(
|
||||
c.Set("openai_passthrough", true)
|
||||
}
|
||||
|
||||
upstreamStart := time.Now()
|
||||
resp, err := s.httpUpstream.Do(upstreamReq, proxyURL, account.ID, account.Concurrency)
|
||||
SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds())
|
||||
if err != nil {
|
||||
// Transport-level failure (proxy/DNS/TCP/TLS — no HTTP response). Convert to
|
||||
// a failover so the handler switches to a healthy account, and temporarily
|
||||
// unschedule the account on durable faults (e.g. rejected proxy credentials).
|
||||
return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, err, true)
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
agentTaskRecoveryTried := false
|
||||
var resp *http.Response
|
||||
for {
|
||||
upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx)
|
||||
upstreamReq, buildErr := s.buildUpstreamRequestOpenAIPassthrough(upstreamCtx, c, account, body, token)
|
||||
releaseUpstreamCtx()
|
||||
if buildErr != nil {
|
||||
return nil, buildErr
|
||||
}
|
||||
|
||||
upstreamStart := time.Now()
|
||||
resp, err = s.httpUpstream.Do(upstreamReq, proxyURL, account.ID, account.Concurrency)
|
||||
SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds())
|
||||
if err != nil {
|
||||
// Transport-level failure (proxy/DNS/TCP/TLS — no HTTP response). Convert to
|
||||
// a failover so the handler switches to a healthy account.
|
||||
return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, err, true)
|
||||
}
|
||||
if resp.StatusCode < 400 {
|
||||
break
|
||||
}
|
||||
|
||||
// Peek only to identify an invalid task. Restore the body so the existing
|
||||
// passthrough error handling sees the same response after recovery fails.
|
||||
probeBody := s.readUpstreamErrorBody(resp)
|
||||
_ = resp.Body.Close()
|
||||
resp.Body = io.NopCloser(bytes.NewReader(probeBody))
|
||||
if !agentTaskRecoveryTried && s.isAgentIdentityAccount(ctx, account) && isAgentIdentityTaskInvalidHTTPResponse(resp.StatusCode, probeBody) {
|
||||
agentTaskRecoveryTried = true
|
||||
expectedTaskID := account.GetCredential("task_id")
|
||||
if recoveryErr := s.recoverAgentIdentityTask(ctx, account, expectedTaskID); recoveryErr != nil {
|
||||
return nil, fmt.Errorf("agent identity task recovery failed: %w", recoveryErr)
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
if resp.StatusCode >= 400 {
|
||||
responseBody := s.readUpstreamErrorBody(resp)
|
||||
// 透传模式默认保持原样代理;容量错误以及 API-key 上游的瞬时
|
||||
// 5xx 应先触发多账号 failover,且此时尚未写入下游响应。
|
||||
if shouldFailoverOpenAIPassthroughResponse(account, resp.StatusCode, responseBody) {
|
||||
return nil, s.handleFailoverErrorResponsePassthrough(ctx, resp, c, account, body, responseBody)
|
||||
// probeBody 已在上方任务探测时读取过一次,直接复用避免重复读取。
|
||||
if shouldFailoverOpenAIPassthroughResponse(account, resp.StatusCode, probeBody) {
|
||||
return nil, s.handleFailoverErrorResponsePassthrough(ctx, resp, c, account, body, probeBody)
|
||||
}
|
||||
return nil, s.handleErrorResponsePassthrough(ctx, resp, c, account, body, responseBody)
|
||||
return nil, s.handleErrorResponsePassthrough(ctx, resp, c, account, body, probeBody)
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
serviceTier := extractOpenAIServiceTierFromBody(body)
|
||||
|
||||
@@ -339,7 +358,15 @@ func (s *OpenAIGatewayService) buildUpstreamRequestOpenAIPassthrough(
|
||||
req.Header.Del("authorization")
|
||||
req.Header.Del("x-api-key")
|
||||
req.Header.Del("x-goog-api-key")
|
||||
req.Header.Set("authorization", "Bearer "+token)
|
||||
authHeaders, err := s.buildOpenAIAuthenticationHeaders(ctx, account, token)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("build openai authentication headers: %w", err)
|
||||
}
|
||||
for key, values := range authHeaders {
|
||||
for _, value := range values {
|
||||
req.Header.Add(key, value)
|
||||
}
|
||||
}
|
||||
|
||||
// OAuth 透传到 ChatGPT internal API 时补齐必要头。
|
||||
if account.Type == AccountTypeOAuth {
|
||||
@@ -449,7 +476,7 @@ func (s *OpenAIGatewayService) handleFailoverErrorResponsePassthrough(
|
||||
requestBody []byte,
|
||||
responseBody []byte,
|
||||
) error {
|
||||
body := responseBody
|
||||
body := s.redactAgentIdentitySensitiveBody(ctx, account, responseBody)
|
||||
|
||||
upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(body))
|
||||
upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
|
||||
@@ -494,7 +521,7 @@ func (s *OpenAIGatewayService) handleErrorResponsePassthrough(
|
||||
responseBody []byte,
|
||||
) error {
|
||||
MarkResponseCommitted(c)
|
||||
body := responseBody
|
||||
body := s.redactAgentIdentitySensitiveBody(ctx, account, responseBody)
|
||||
|
||||
// cyber_policy:透传账号本就把原始 body 回给客户端(下方 c.Data),此处仅打标记,
|
||||
// 供 handler 事后写风控/邮件。cyber 是上游网络安全策略拦截,不冷却账号,
|
||||
|
||||
@@ -393,6 +393,7 @@ type OpenAIGatewayService struct {
|
||||
openaiWSStateStoreOnce sync.Once
|
||||
openaiSchedulerOnce sync.Once
|
||||
openaiWSPassthroughDialerOnce sync.Once
|
||||
agentIdentityTaskMu sync.Mutex
|
||||
openaiWSPool *openAIWSConnPool
|
||||
openaiWSStateStore OpenAIWSStateStore
|
||||
openaiScheduler OpenAIAccountScheduler
|
||||
@@ -401,6 +402,10 @@ type OpenAIGatewayService struct {
|
||||
|
||||
openaiWSFallbackUntil sync.Map // key: int64(accountID), value: time.Time
|
||||
openaiAccountRuntimeBlockUntil sync.Map // key: int64(accountID), value: time.Time
|
||||
openaiAccountRuntimeBlockLocks sync.Map // key: int64(accountID), value: *sync.Mutex
|
||||
openaiAccountRuntimeBlockGeneration sync.Map // key: int64(accountID), value: uint64
|
||||
openaiAccountRuntimeBlockSequence atomic.Uint64
|
||||
grokCredentialMutationLocks sync.Map // key: int64(accountID), value: *sync.Mutex
|
||||
openaiOAuth429WindowStartUnixNano atomic.Int64
|
||||
openaiOAuth429WindowCount atomic.Int64
|
||||
openaiWSRetryMetrics openAIWSRetryMetrics
|
||||
@@ -591,6 +596,12 @@ func (s *OpenAIGatewayService) CloseOpenAIWSPool() {
|
||||
}
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) InvalidateAgentIdentityWSConnections(accountID int64) {
|
||||
if pool := s.getOpenAIWSConnPool(); pool != nil {
|
||||
pool.ClearAccount(accountID)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) logOpenAIWSModeBootstrap() {
|
||||
if s == nil || s.cfg == nil {
|
||||
return
|
||||
@@ -1090,6 +1101,9 @@ func (s *OpenAIGatewayService) GetAccessToken(ctx context.Context, account *Acco
|
||||
}
|
||||
switch account.Type {
|
||||
case AccountTypeOAuth:
|
||||
if account.IsOpenAIAgentIdentity() {
|
||||
return "", OpenAIAuthModeAgentIdentity, nil
|
||||
}
|
||||
if account.Platform == PlatformGrok {
|
||||
if s.grokTokenProvider != nil {
|
||||
accessToken, err := s.grokTokenProvider.GetAccessToken(ctx, account)
|
||||
|
||||
@@ -279,6 +279,7 @@ func (s *OpenAIGatewayService) handleErrorResponse(
|
||||
requestedModel ...string,
|
||||
) (*OpenAIForwardResult, error) {
|
||||
body := s.readUpstreamErrorBody(resp)
|
||||
body = s.redactAgentIdentitySensitiveBody(ctx, account, body)
|
||||
|
||||
// cyber_policy 硬阻断:透传上游原始错误体给客户端(不重包成通用 502),不冷却账号。
|
||||
// 当前请求恒透传(需求1);标记供 handler 事后写风控/邮件。400 cyber 不可 failover
|
||||
@@ -470,6 +471,7 @@ func (s *OpenAIGatewayService) handleCompatErrorResponse(
|
||||
requestedModel ...string,
|
||||
) (*OpenAIForwardResult, error) {
|
||||
body := s.readUpstreamErrorBody(resp)
|
||||
body = s.redactAgentIdentitySensitiveBody(context.Background(), account, body)
|
||||
|
||||
// cyber_policy:兼容路径(Chat Completions / Anthropic)以各自格式回写错误,
|
||||
// 不原样透传 responses 格式的 cyber body(否则对下游格式不合法)。cyber 是上游网络
|
||||
|
||||
@@ -0,0 +1,92 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"encoding/base64"
|
||||
"encoding/binary"
|
||||
"fmt"
|
||||
"image"
|
||||
_ "image/jpeg"
|
||||
_ "image/png"
|
||||
"io"
|
||||
"strings"
|
||||
)
|
||||
|
||||
const maxOpenAIImageDimensionProbeBytes int64 = 1 << 20
|
||||
|
||||
func detectOpenAIImageResultSize(encoded string) string {
|
||||
payload := strings.TrimSpace(encoded)
|
||||
if strings.HasPrefix(strings.ToLower(payload), "data:") {
|
||||
comma := strings.IndexByte(payload, ',')
|
||||
if comma < 0 || comma+1 >= len(payload) {
|
||||
return ""
|
||||
}
|
||||
payload = strings.TrimSpace(payload[comma+1:])
|
||||
}
|
||||
if payload == "" {
|
||||
return ""
|
||||
}
|
||||
|
||||
for _, encoding := range []*base64.Encoding{base64.StdEncoding, base64.RawStdEncoding} {
|
||||
decoded := base64.NewDecoder(encoding, strings.NewReader(payload))
|
||||
buffered := bufio.NewReader(io.LimitReader(decoded, maxOpenAIImageDimensionProbeBytes))
|
||||
prefix, _ := buffered.Peek(30)
|
||||
if width, height, ok := detectOpenAIWebPDimensions(prefix); ok {
|
||||
return fmt.Sprintf("%dx%d", width, height)
|
||||
}
|
||||
cfg, _, err := image.DecodeConfig(buffered)
|
||||
if err != nil || cfg.Width <= 0 || cfg.Height <= 0 {
|
||||
continue
|
||||
}
|
||||
return fmt.Sprintf("%dx%d", cfg.Width, cfg.Height)
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func detectOpenAIWebPDimensions(header []byte) (int, int, bool) {
|
||||
if len(header) < 16 || string(header[:4]) != "RIFF" || string(header[8:12]) != "WEBP" {
|
||||
return 0, 0, false
|
||||
}
|
||||
|
||||
switch string(header[12:16]) {
|
||||
case "VP8X":
|
||||
if len(header) < 30 {
|
||||
return 0, 0, false
|
||||
}
|
||||
width := 1 + int(header[24]) + int(header[25])<<8 + int(header[26])<<16
|
||||
height := 1 + int(header[27]) + int(header[28])<<8 + int(header[29])<<16
|
||||
return width, height, width > 0 && height > 0
|
||||
case "VP8 ":
|
||||
if len(header) < 30 || string(header[23:26]) != "\x9d\x01\x2a" {
|
||||
return 0, 0, false
|
||||
}
|
||||
width := int(binary.LittleEndian.Uint16(header[26:28]) & 0x3fff)
|
||||
height := int(binary.LittleEndian.Uint16(header[28:30]) & 0x3fff)
|
||||
return width, height, width > 0 && height > 0
|
||||
case "VP8L":
|
||||
if len(header) < 25 || header[20] != 0x2f {
|
||||
return 0, 0, false
|
||||
}
|
||||
width := 1 + int(header[21]) + int(header[22]&0x3f)<<8
|
||||
height := 1 + int(header[22]>>6) + int(header[23])<<2 + int(header[24]&0x0f)<<10
|
||||
return width, height, width > 0 && height > 0
|
||||
default:
|
||||
return 0, 0, false
|
||||
}
|
||||
}
|
||||
|
||||
func reconcileOpenAIResponsesImageResultSizes(results []openAIResponsesImageResult, firstMeta *openAIResponsesImageResult) {
|
||||
for i := range results {
|
||||
// ChatGPT OAuth can normalize requested controls to "auto". The final
|
||||
// image bytes are authoritative for response metadata and tier billing.
|
||||
if actualSize := detectOpenAIImageResultSize(results[i].Result); actualSize != "" {
|
||||
results[i].Size = actualSize
|
||||
}
|
||||
}
|
||||
if firstMeta == nil || len(results) == 0 {
|
||||
return
|
||||
}
|
||||
if size := strings.TrimSpace(results[0].Size); size != "" {
|
||||
firstMeta.Size = size
|
||||
}
|
||||
}
|
||||
@@ -632,6 +632,7 @@ func (s *OpenAIGatewayService) forwardOpenAIImagesAPIKey(
|
||||
if resp.StatusCode >= 400 {
|
||||
respBody := s.readUpstreamErrorBody(resp)
|
||||
_ = resp.Body.Close()
|
||||
respBody = s.redactAgentIdentitySensitiveBody(upstreamCtx, account, respBody)
|
||||
resp.Body = io.NopCloser(bytes.NewReader(respBody))
|
||||
upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody))
|
||||
upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
|
||||
@@ -752,7 +753,15 @@ func (s *OpenAIGatewayService) buildOpenAIImagesRequest(
|
||||
return nil, err
|
||||
}
|
||||
req = req.WithContext(WithHTTPUpstreamProfile(req.Context(), HTTPUpstreamProfileOpenAI))
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
authHeaders, err := s.buildOpenAIAuthenticationHeaders(ctx, account, token)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("build openai authentication headers: %w", err)
|
||||
}
|
||||
for key, values := range authHeaders {
|
||||
for _, value := range values {
|
||||
req.Header.Add(key, value)
|
||||
}
|
||||
}
|
||||
for key, values := range c.Request.Header {
|
||||
if !openaiPassthroughAllowedHeaders[strings.ToLower(key)] {
|
||||
continue
|
||||
|
||||
@@ -0,0 +1,174 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/binary"
|
||||
"fmt"
|
||||
"image"
|
||||
"image/color"
|
||||
"image/jpeg"
|
||||
"image/png"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
func TestDetectOpenAIImageResultSize(t *testing.T) {
|
||||
pngEncoded := encodeOpenAIImageTestPNG(t, 1672, 941)
|
||||
jpegEncoded := encodeOpenAIImageTestJPEG(t, 640, 360)
|
||||
webpVP8XEncoded := encodeOpenAIImageTestWebPVP8X(1920, 1080)
|
||||
webpVP8Encoded := encodeOpenAIImageTestWebPVP8(1280, 720)
|
||||
webpVP8LEncoded := encodeOpenAIImageTestWebPVP8L(640, 480)
|
||||
|
||||
require.Equal(t, "1672x941", detectOpenAIImageResultSize(pngEncoded))
|
||||
require.Equal(t, "1672x941", detectOpenAIImageResultSize(strings.TrimRight(pngEncoded, "=")))
|
||||
require.Equal(t, "1672x941", detectOpenAIImageResultSize("data:image/png;base64,"+pngEncoded))
|
||||
require.Equal(t, "640x360", detectOpenAIImageResultSize(jpegEncoded))
|
||||
require.Equal(t, "1920x1080", detectOpenAIImageResultSize(webpVP8XEncoded))
|
||||
require.Equal(t, "1280x720", detectOpenAIImageResultSize(webpVP8Encoded))
|
||||
require.Equal(t, "640x480", detectOpenAIImageResultSize(webpVP8LEncoded))
|
||||
require.Empty(t, detectOpenAIImageResultSize("data:image/png;base64"))
|
||||
require.Empty(t, detectOpenAIImageResultSize("not-image-data"))
|
||||
}
|
||||
|
||||
func TestOpenAIGatewayServiceForwardImages_OAuthUsesDecodedOutputDimensions(t *testing.T) {
|
||||
run := runOpenAIOAuthImageActualSizeTest(t, false)
|
||||
|
||||
require.Equal(t, "3840x2160", gjson.GetBytes(run.upstream.lastBody, "tools.0.size").String())
|
||||
require.Equal(t, "low", gjson.GetBytes(run.upstream.lastBody, "tools.0.quality").String())
|
||||
require.Equal(t, "1672x941", gjson.Get(run.recorder.Body.String(), "size").String())
|
||||
require.Equal(t, "auto", gjson.Get(run.recorder.Body.String(), "quality").String())
|
||||
require.Equal(t, []string{"1672x941"}, run.result.ImageOutputSizes)
|
||||
|
||||
ApplyOpenAIImageBillingResolution(run.result)
|
||||
require.Equal(t, ImageBillingSize2K, run.result.ImageSize)
|
||||
require.Equal(t, "1672x941", run.result.ImageOutputSize)
|
||||
require.Equal(t, ImageSizeSourceOutput, run.result.ImageSizeSource)
|
||||
}
|
||||
|
||||
func TestOpenAIGatewayServiceForwardImages_OAuthStreamingUsesDecodedOutputDimensions(t *testing.T) {
|
||||
run := runOpenAIOAuthImageActualSizeTest(t, true)
|
||||
|
||||
events := parseOpenAIImageTestSSEEvents(run.recorder.Body.String())
|
||||
completed, ok := findOpenAIImageTestSSEEvent(events, "image_generation.completed")
|
||||
require.True(t, ok)
|
||||
require.Equal(t, "1672x941", gjson.Get(completed.Data, "size").String())
|
||||
require.Equal(t, "auto", gjson.Get(completed.Data, "quality").String())
|
||||
require.Equal(t, []string{"1672x941"}, run.result.ImageOutputSizes)
|
||||
}
|
||||
|
||||
type openAIOAuthImageActualSizeTestRun struct {
|
||||
result *OpenAIForwardResult
|
||||
recorder *httptest.ResponseRecorder
|
||||
upstream *httpUpstreamRecorder
|
||||
}
|
||||
|
||||
func runOpenAIOAuthImageActualSizeTest(t *testing.T, stream bool) openAIOAuthImageActualSizeTestRun {
|
||||
t.Helper()
|
||||
gin.SetMode(gin.TestMode)
|
||||
body := []byte(fmt.Sprintf(`{"model":"gpt-image-2","prompt":"draw a test chart","size":"3840x2160","quality":"low","output_format":"png","stream":%t}`, stream))
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = req
|
||||
c.Set("api_key", &APIKey{ID: 42})
|
||||
|
||||
encoded := encodeOpenAIImageTestPNG(t, 1672, 941)
|
||||
upstreamBody := fmt.Sprintf(
|
||||
"data: {\"type\":\"response.created\",\"response\":{\"created_at\":1710000000,\"tools\":[{\"type\":\"image_generation\",\"model\":\"gpt-image-2\",\"size\":\"auto\",\"quality\":\"auto\",\"output_format\":\"png\"}]}}\n\n"+
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"created_at\":1710000000,\"tools\":[{\"type\":\"image_generation\",\"model\":\"gpt-image-2\",\"size\":\"auto\",\"quality\":\"auto\",\"output_format\":\"png\"}],\"output\":[{\"id\":\"ig_actual_size\",\"type\":\"image_generation_call\",\"result\":%q}]}}\n\n"+
|
||||
"data: [DONE]\n\n",
|
||||
encoded,
|
||||
)
|
||||
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{
|
||||
"Content-Type": []string{"text/event-stream"},
|
||||
"X-Request-Id": []string{"req_img_actual_size"},
|
||||
},
|
||||
Body: io.NopCloser(strings.NewReader(upstreamBody)),
|
||||
}}
|
||||
svc := &OpenAIGatewayService{httpUpstream: upstream}
|
||||
parsed, err := svc.ParseOpenAIImagesRequest(c, body)
|
||||
require.NoError(t, err)
|
||||
|
||||
account := &Account{
|
||||
ID: 1,
|
||||
Name: "openai-oauth",
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeOAuth,
|
||||
Credentials: map[string]any{
|
||||
"access_token": "token-123",
|
||||
"chatgpt_account_id": "acct-123",
|
||||
},
|
||||
}
|
||||
result, err := svc.ForwardImages(context.Background(), c, account, body, parsed, "")
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
return openAIOAuthImageActualSizeTestRun{result: result, recorder: rec, upstream: upstream}
|
||||
}
|
||||
|
||||
func encodeOpenAIImageTestPNG(t *testing.T, width, height int) string {
|
||||
t.Helper()
|
||||
img := image.NewNRGBA(image.Rect(0, 0, width, height))
|
||||
img.SetNRGBA(0, 0, color.NRGBA{R: 0xff, A: 0xff})
|
||||
var buf bytes.Buffer
|
||||
require.NoError(t, png.Encode(&buf, img))
|
||||
return base64.StdEncoding.EncodeToString(buf.Bytes())
|
||||
}
|
||||
|
||||
func encodeOpenAIImageTestJPEG(t *testing.T, width, height int) string {
|
||||
t.Helper()
|
||||
img := image.NewNRGBA(image.Rect(0, 0, width, height))
|
||||
img.SetNRGBA(0, 0, color.NRGBA{G: 0xff, A: 0xff})
|
||||
var buf bytes.Buffer
|
||||
require.NoError(t, jpeg.Encode(&buf, img, nil))
|
||||
return base64.StdEncoding.EncodeToString(buf.Bytes())
|
||||
}
|
||||
|
||||
func encodeOpenAIImageTestWebPVP8X(width, height int) string {
|
||||
header := make([]byte, 30)
|
||||
copy(header[0:4], "RIFF")
|
||||
copy(header[8:12], "WEBP")
|
||||
copy(header[12:16], "VP8X")
|
||||
width--
|
||||
height--
|
||||
header[24], header[25], header[26] = byte(width), byte(width>>8), byte(width>>16)
|
||||
header[27], header[28], header[29] = byte(height), byte(height>>8), byte(height>>16)
|
||||
return base64.StdEncoding.EncodeToString(header)
|
||||
}
|
||||
|
||||
func encodeOpenAIImageTestWebPVP8(width, height int) string {
|
||||
header := make([]byte, 30)
|
||||
copy(header[0:4], "RIFF")
|
||||
copy(header[8:12], "WEBP")
|
||||
copy(header[12:16], "VP8 ")
|
||||
copy(header[23:26], "\x9d\x01\x2a")
|
||||
binary.LittleEndian.PutUint16(header[26:28], uint16(width))
|
||||
binary.LittleEndian.PutUint16(header[28:30], uint16(height))
|
||||
return base64.StdEncoding.EncodeToString(header)
|
||||
}
|
||||
|
||||
func encodeOpenAIImageTestWebPVP8L(width, height int) string {
|
||||
header := make([]byte, 25)
|
||||
copy(header[0:4], "RIFF")
|
||||
copy(header[8:12], "WEBP")
|
||||
copy(header[12:16], "VP8L")
|
||||
header[20] = 0x2f
|
||||
width--
|
||||
height--
|
||||
header[21] = byte(width)
|
||||
header[22] = byte(width>>8)&0x3f | byte(height&0x03)<<6
|
||||
header[23] = byte(height >> 2)
|
||||
header[24] = byte(height>>10) & 0x0f
|
||||
return base64.StdEncoding.EncodeToString(header)
|
||||
}
|
||||
@@ -565,12 +565,14 @@ func collectOpenAIImagesFromResponsesBody(body []byte) ([]openAIResponsesImageRe
|
||||
return nil, 0, nil, openAIResponsesImageResult{}, false, collectErr
|
||||
}
|
||||
if len(finalResults) > 0 {
|
||||
reconcileOpenAIResponsesImageResultSizes(finalResults, &finalMeta)
|
||||
return finalResults, createdAt, usageRaw, finalMeta, true, nil
|
||||
}
|
||||
|
||||
if len(fallbackResults) > 0 {
|
||||
firstMeta := fallbackResults[0]
|
||||
mergeOpenAIResponsesImageMeta(&firstMeta, responseMeta)
|
||||
reconcileOpenAIResponsesImageResultSizes(fallbackResults, &firstMeta)
|
||||
return fallbackResults, createdAt, usageRaw, firstMeta, foundFinal, nil
|
||||
}
|
||||
return nil, createdAt, usageRaw, openAIResponsesImageResult{}, foundFinal, nil
|
||||
@@ -1262,6 +1264,7 @@ func (s *OpenAIGatewayService) handleOpenAIImagesOAuthStreamingResponse(
|
||||
mergeOpenAIResponsesImageMeta(&img, streamMeta)
|
||||
appendOpenAIResponsesImageResultDedup(&finalResults, finalSeen, "", img)
|
||||
}
|
||||
reconcileOpenAIResponsesImageResultSizes(finalResults, nil)
|
||||
if len(finalResults) == 0 {
|
||||
outputErr := fmt.Errorf("upstream did not return image output")
|
||||
// 软失败:response.completed 事件里没有图片。记录上游诊断摘要到 ops,
|
||||
@@ -1324,8 +1327,12 @@ func (s *OpenAIGatewayService) handleOpenAIImagesOAuthStreamingResponse(
|
||||
}
|
||||
if len(pendingResults) > 0 {
|
||||
eventName := streamPrefix + ".completed"
|
||||
for _, img := range pendingResults {
|
||||
mergeOpenAIResponsesImageMeta(&img, streamMeta)
|
||||
finalResults := append([]openAIResponsesImageResult(nil), pendingResults...)
|
||||
for i := range finalResults {
|
||||
mergeOpenAIResponsesImageMeta(&finalResults[i], streamMeta)
|
||||
}
|
||||
reconcileOpenAIResponsesImageResultSizes(finalResults, nil)
|
||||
for _, img := range finalResults {
|
||||
key := openAIResponsesImageResultKey("", img)
|
||||
if _, exists := emitted[key]; exists {
|
||||
continue
|
||||
@@ -1335,7 +1342,7 @@ func (s *OpenAIGatewayService) handleOpenAIImagesOAuthStreamingResponse(
|
||||
s.tryWriteOpenAIImagesStreamEvent(c, flusher, &clientDisconnected, &lastDownstreamWriteAt, eventName, payload)
|
||||
}
|
||||
imageCount = len(emitted)
|
||||
imageOutputSizes = openAIResponsesImageResultSizes(pendingResults)
|
||||
imageOutputSizes = openAIResponsesImageResultSizes(finalResults)
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -1564,6 +1571,14 @@ func (s *OpenAIGatewayService) forwardOpenAIImagesOAuth(
|
||||
if resp.StatusCode >= 400 {
|
||||
respBody := s.readUpstreamErrorBody(resp)
|
||||
_ = resp.Body.Close()
|
||||
respBody = s.redactAgentIdentitySensitiveBody(upstreamCtx, account, respBody)
|
||||
if !agentIdentityTaskRecoveryWasTried(ctx) && s.isAgentIdentityAccount(ctx, account) && isAgentIdentityTaskInvalidHTTPResponse(resp.StatusCode, respBody) {
|
||||
expectedTaskID := account.GetCredential("task_id")
|
||||
if err := s.recoverAgentIdentityTask(ctx, account, expectedTaskID); err != nil {
|
||||
return nil, fmt.Errorf("agent identity task recovery failed: %w", err)
|
||||
}
|
||||
return s.forwardOpenAIImagesOAuth(markAgentIdentityTaskRecoveryTried(ctx), c, account, parsed, channelMappedModel)
|
||||
}
|
||||
resp.Body = io.NopCloser(bytes.NewReader(respBody))
|
||||
upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody))
|
||||
upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
|
||||
|
||||
@@ -8,6 +8,7 @@ import (
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
@@ -21,6 +22,8 @@ import (
|
||||
// errors.Is still matches it by identity since ResetCredit returns this var.
|
||||
var ErrSparkShadowResetNotSupported = infraerrors.New(http.StatusConflict, "SPARK_SHADOW_RESET_NOT_SUPPORTED", "spark shadow account does not support credit reset; reset the parent account")
|
||||
|
||||
var ErrAgentIdentityResetNotSupported = infraerrors.New(http.StatusConflict, "AGENT_IDENTITY_RESET_NOT_SUPPORTED", "agent identity does not support rate-limit reset credit consumption")
|
||||
|
||||
// Endpoints used by the OpenAI/ChatGPT/Codex quota query and reset feature.
|
||||
const (
|
||||
chatGPTUsageURL = "https://chatgpt.com/backend-api/wham/usage"
|
||||
@@ -116,6 +119,8 @@ type OpenAIQuotaService struct {
|
||||
proxyRepo ProxyRepository
|
||||
tokenProvider *OpenAITokenProvider
|
||||
privacyClientFactory PrivacyClientFactory
|
||||
agentIdentityTaskMu sync.Mutex
|
||||
agentIdentityWS agentIdentityWSConnectionInvalidator
|
||||
}
|
||||
|
||||
// NewOpenAIQuotaService constructs a quota service. token provider is required —
|
||||
@@ -151,21 +156,36 @@ func (s *OpenAIQuotaService) QueryUsage(ctx context.Context, accountID int64) (*
|
||||
|
||||
callCtx, cancel := context.WithTimeout(ctx, openaiQuotaUpstreamTimeout)
|
||||
defer cancel()
|
||||
agentIdentity := s.isAgentIdentityAccount(ctx, accountID)
|
||||
|
||||
var payload OpenAIQuotaUsage
|
||||
resp, err := client.R().
|
||||
SetContext(callCtx).
|
||||
SetHeaders(buildCodexCommonHeaders(accessToken, chatGPTAccountID, fedRAMP)).
|
||||
SetSuccessResult(&payload).
|
||||
Get(chatGPTUsageURL)
|
||||
if err != nil {
|
||||
return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_QUOTA_REQUEST_FAILED", "upstream request failed: %v", err)
|
||||
}
|
||||
if !resp.IsSuccessState() {
|
||||
status := resp.StatusCode
|
||||
body := truncate(resp.String(), 240)
|
||||
slog.Warn("openai_quota_query_failed", "account_id", accountID, "status", status, "body", body)
|
||||
return nil, infraerrors.Newf(mapUpstreamStatus(status), "OPENAI_QUOTA_UPSTREAM_ERROR", "upstream returned %d: %s", status, body)
|
||||
for recovered := false; ; {
|
||||
quotaHeaders, headerErr := s.buildCodexQuotaHeaders(callCtx, accountID, accessToken, chatGPTAccountID, fedRAMP)
|
||||
if headerErr != nil {
|
||||
return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_QUOTA_AUTH_FAILED", "failed to build upstream authentication: %v", headerErr)
|
||||
}
|
||||
resp, err := client.R().
|
||||
SetContext(callCtx).
|
||||
SetHeaders(quotaHeaders).
|
||||
SetSuccessResult(&payload).
|
||||
Get(chatGPTUsageURL)
|
||||
if err != nil {
|
||||
return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_QUOTA_REQUEST_FAILED", "upstream request failed: %v", err)
|
||||
}
|
||||
if !resp.IsSuccessState() {
|
||||
if agentIdentity && !recovered && isAgentIdentityTaskInvalidHTTPResponse(resp.StatusCode, []byte(resp.String())) {
|
||||
recovered = true
|
||||
if err := s.recoverAgentIdentityTask(ctx, accountID); err != nil {
|
||||
return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_QUOTA_AUTH_FAILED", "agent identity task recovery failed: %v", err)
|
||||
}
|
||||
continue
|
||||
}
|
||||
status := resp.StatusCode
|
||||
body := truncate(s.redactQuotaErrorBody(ctx, accountID, resp.String()), 240)
|
||||
slog.Warn("openai_quota_query_failed", "account_id", accountID, "status", status, "body", body)
|
||||
return nil, infraerrors.Newf(mapUpstreamStatus(status), "OPENAI_QUOTA_UPSTREAM_ERROR", "upstream returned %d: %s", status, body)
|
||||
}
|
||||
break
|
||||
}
|
||||
|
||||
payload.FetchedAt = time.Now().Unix()
|
||||
@@ -189,9 +209,14 @@ func (s *OpenAIQuotaService) QueryUsage(ctx context.Context, accountID int64) (*
|
||||
}
|
||||
|
||||
func (s *OpenAIQuotaService) queryResetCreditDetails(ctx context.Context, client *req.Client, accessToken, chatGPTAccountID string, fedRAMP bool, accountID int64) *openAIRateLimitResetCreditDetails {
|
||||
quotaHeaders, headerErr := s.buildCodexQuotaHeaders(ctx, accountID, accessToken, chatGPTAccountID, fedRAMP)
|
||||
if headerErr != nil {
|
||||
slog.Warn("openai_quota_reset_credit_details_auth_failed", "account_id", accountID, "error", headerErr)
|
||||
return nil
|
||||
}
|
||||
resp, err := client.R().
|
||||
SetContext(ctx).
|
||||
SetHeaders(buildCodexCommonHeaders(accessToken, chatGPTAccountID, fedRAMP)).
|
||||
SetHeaders(quotaHeaders).
|
||||
Get(chatGPTRateLimitCreditsURL)
|
||||
if err != nil {
|
||||
slog.Warn("openai_quota_reset_credit_details_failed", "account_id", accountID, "error", err)
|
||||
@@ -233,6 +258,9 @@ func (s *OpenAIQuotaService) ResetCredit(ctx context.Context, accountID int64) (
|
||||
if acc.IsShadow() {
|
||||
return nil, ErrSparkShadowResetNotSupported
|
||||
}
|
||||
if acc.IsOpenAIAgentIdentity() {
|
||||
return nil, ErrAgentIdentityResetNotSupported
|
||||
}
|
||||
}
|
||||
|
||||
accessToken, chatGPTAccountID, proxyURL, fedRAMP, err := s.prepareUpstreamCall(ctx, accountID)
|
||||
@@ -252,11 +280,12 @@ func (s *OpenAIQuotaService) ResetCredit(ctx context.Context, accountID int64) (
|
||||
|
||||
callCtx, cancel := context.WithTimeout(ctx, openaiQuotaUpstreamTimeout)
|
||||
defer cancel()
|
||||
|
||||
headers := buildCodexCommonHeaders(accessToken, chatGPTAccountID, fedRAMP)
|
||||
headers["content-type"] = "application/json"
|
||||
|
||||
var payload OpenAIQuotaResetResult
|
||||
headers, headerErr := s.buildCodexQuotaHeaders(callCtx, accountID, accessToken, chatGPTAccountID, fedRAMP)
|
||||
if headerErr != nil {
|
||||
return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_QUOTA_AUTH_FAILED", "failed to build upstream authentication: %v", headerErr)
|
||||
}
|
||||
headers["content-type"] = "application/json"
|
||||
resp, err := client.R().
|
||||
SetContext(callCtx).
|
||||
SetHeaders(headers).
|
||||
@@ -268,7 +297,7 @@ func (s *OpenAIQuotaService) ResetCredit(ctx context.Context, accountID int64) (
|
||||
}
|
||||
if !resp.IsSuccessState() {
|
||||
status := resp.StatusCode
|
||||
body := truncate(resp.String(), 240)
|
||||
body := truncate(s.redactQuotaErrorBody(callCtx, accountID, resp.String()), 240)
|
||||
slog.Warn("openai_quota_reset_failed", "account_id", accountID, "status", status, "body", body)
|
||||
return nil, infraerrors.Newf(mapUpstreamStatus(status), "OPENAI_QUOTA_RESET_UPSTREAM_ERROR", "upstream returned %d: %s", status, body)
|
||||
}
|
||||
@@ -285,7 +314,7 @@ func (s *OpenAIQuotaService) ResetCredit(ctx context.Context, accountID int64) (
|
||||
// token via the shared TokenProvider, and resolves the chatgpt-account-id and
|
||||
// proxy URL. Centralized so QueryUsage / ResetCredit share validation.
|
||||
func (s *OpenAIQuotaService) prepareUpstreamCall(ctx context.Context, accountID int64) (accessToken, chatGPTAccountID, proxyURL string, fedRAMP bool, err error) {
|
||||
if s == nil || s.accountRepo == nil || s.tokenProvider == nil || s.privacyClientFactory == nil {
|
||||
if s == nil || s.accountRepo == nil || s.privacyClientFactory == nil {
|
||||
return "", "", "", false, infraerrors.New(http.StatusInternalServerError, "OPENAI_QUOTA_NOT_CONFIGURED", "openai quota service is not configured")
|
||||
}
|
||||
|
||||
@@ -323,12 +352,17 @@ func (s *OpenAIQuotaService) prepareUpstreamCall(ctx context.Context, accountID
|
||||
return "", "", "", false, infraerrors.New(http.StatusBadRequest, "OPENAI_QUOTA_MISSING_ACCOUNT_ID", "chatgpt_account_id is missing; please re-authorize this account")
|
||||
}
|
||||
|
||||
accessToken, err = s.tokenProvider.GetAccessToken(ctx, account)
|
||||
if err != nil {
|
||||
return "", "", "", false, infraerrors.Newf(http.StatusBadGateway, "OPENAI_QUOTA_TOKEN_UNAVAILABLE", "failed to acquire access token: %v", err)
|
||||
}
|
||||
if strings.TrimSpace(accessToken) == "" {
|
||||
return "", "", "", false, infraerrors.New(http.StatusBadGateway, "OPENAI_QUOTA_TOKEN_UNAVAILABLE", "access token is empty")
|
||||
if !account.IsOpenAIAgentIdentity() {
|
||||
if s.tokenProvider == nil {
|
||||
return "", "", "", false, infraerrors.New(http.StatusInternalServerError, "OPENAI_QUOTA_NOT_CONFIGURED", "openai quota token provider is not configured")
|
||||
}
|
||||
accessToken, err = s.tokenProvider.GetAccessToken(ctx, account)
|
||||
if err != nil {
|
||||
return "", "", "", false, infraerrors.Newf(http.StatusBadGateway, "OPENAI_QUOTA_TOKEN_UNAVAILABLE", "failed to acquire access token: %v", err)
|
||||
}
|
||||
if strings.TrimSpace(accessToken) == "" {
|
||||
return "", "", "", false, infraerrors.New(http.StatusBadGateway, "OPENAI_QUOTA_TOKEN_UNAVAILABLE", "access token is empty")
|
||||
}
|
||||
}
|
||||
fedRAMP = account.IsChatGPTAccountFedRAMP()
|
||||
|
||||
@@ -351,6 +385,91 @@ func (s *OpenAIQuotaService) prepareUpstreamCall(ctx context.Context, accountID
|
||||
return accessToken, chatGPTAccountID, proxyURL, fedRAMP, nil
|
||||
}
|
||||
|
||||
func (s *OpenAIQuotaService) recoverAgentIdentityTask(ctx context.Context, accountID int64) error {
|
||||
if s == nil || s.accountRepo == nil {
|
||||
return fmt.Errorf("account repository is unavailable")
|
||||
}
|
||||
account, err := s.accountRepo.GetByID(ctx, accountID)
|
||||
if err != nil || account == nil {
|
||||
return fmt.Errorf("account is unavailable")
|
||||
}
|
||||
if account.IsShadow() {
|
||||
account, err = resolveCredentialAccount(ctx, s.accountRepo, account)
|
||||
if err != nil || account == nil {
|
||||
return fmt.Errorf("credential account is unavailable")
|
||||
}
|
||||
}
|
||||
if !account.IsOpenAIAgentIdentity() {
|
||||
return nil
|
||||
}
|
||||
return ensureAgentIdentityTaskForAccount(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, account, account.GetCredential("task_id"))
|
||||
}
|
||||
|
||||
func (s *OpenAIQuotaService) isAgentIdentityAccount(ctx context.Context, accountID int64) bool {
|
||||
if s == nil || s.accountRepo == nil {
|
||||
return false
|
||||
}
|
||||
account, err := s.accountRepo.GetByID(ctx, accountID)
|
||||
if err != nil || account == nil {
|
||||
return false
|
||||
}
|
||||
if account.IsShadow() {
|
||||
account, err = resolveCredentialAccount(ctx, s.accountRepo, account)
|
||||
if err != nil || account == nil {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return account.IsOpenAIAgentIdentity()
|
||||
}
|
||||
|
||||
func (s *OpenAIQuotaService) buildCodexQuotaHeaders(ctx context.Context, accountID int64, accessToken, chatGPTAccountID string, fedRAMP bool) (map[string]string, error) {
|
||||
headers := buildCodexCommonHeaders(accessToken, chatGPTAccountID, fedRAMP)
|
||||
if s == nil || s.accountRepo == nil {
|
||||
return headers, nil
|
||||
}
|
||||
account, err := s.accountRepo.GetByID(ctx, accountID)
|
||||
if err != nil || account == nil {
|
||||
if strings.TrimSpace(accessToken) == "" {
|
||||
return nil, fmt.Errorf("agent identity account credentials are unavailable")
|
||||
}
|
||||
return headers, nil
|
||||
}
|
||||
if account.IsShadow() {
|
||||
if resolved, resolveErr := resolveCredentialAccount(ctx, s.accountRepo, account); resolveErr == nil && resolved != nil {
|
||||
account = resolved
|
||||
} else if strings.TrimSpace(accessToken) == "" {
|
||||
return nil, fmt.Errorf("agent identity shadow credentials are unavailable")
|
||||
}
|
||||
}
|
||||
if !account.IsOpenAIAgentIdentity() {
|
||||
return headers, nil
|
||||
}
|
||||
if err := ensureAgentIdentityTaskForAccount(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, account, ""); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
key, err := agentIdentityKeyFromAccount(account)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
assertion, err := buildAgentAssertion(key, time.Now())
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
headers["authorization"] = assertion
|
||||
return headers, nil
|
||||
}
|
||||
|
||||
func (s *OpenAIQuotaService) redactQuotaErrorBody(ctx context.Context, accountID int64, body string) string {
|
||||
if s == nil || s.accountRepo == nil {
|
||||
return body
|
||||
}
|
||||
account, err := s.accountRepo.GetByID(ctx, accountID)
|
||||
if err != nil || account == nil {
|
||||
return body
|
||||
}
|
||||
return string(redactAgentIdentitySensitiveBodyForAccount(ctx, s.accountRepo, account, []byte(body)))
|
||||
}
|
||||
|
||||
// buildCodexCommonHeaders sets the request headers expected by the chatgpt.com
|
||||
// backend so calls succeed past Cloudflare/WASM checks.
|
||||
func buildCodexCommonHeaders(accessToken, chatGPTAccountID string, fedRAMP bool) map[string]string {
|
||||
|
||||
@@ -2,6 +2,10 @@ package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/ed25519"
|
||||
"crypto/rand"
|
||||
"crypto/x509"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
@@ -34,6 +38,15 @@ func (r *stubQuotaAccountRepo) GetByID(_ context.Context, id int64) (*Account, e
|
||||
return acc, nil
|
||||
}
|
||||
|
||||
func (r *stubQuotaAccountRepo) UpdateCredentials(_ context.Context, id int64, credentials map[string]any) error {
|
||||
acc, ok := r.accounts[id]
|
||||
if !ok {
|
||||
return fmt.Errorf("account %d not found", id)
|
||||
}
|
||||
acc.Credentials = credentials
|
||||
return nil
|
||||
}
|
||||
|
||||
// stubQuotaTokenCache 实现 OpenAITokenCache,返回预设静态 token。
|
||||
type stubQuotaTokenCache struct {
|
||||
tokens map[string]string
|
||||
@@ -163,6 +176,23 @@ func TestResetCreditShadowRejected(t *testing.T) {
|
||||
"shadow ResetCredit 应映射为 409 Conflict 而非 500")
|
||||
}
|
||||
|
||||
func TestResetCreditAgentIdentityRejectedBeforeUpstream(t *testing.T) {
|
||||
account := &Account{
|
||||
ID: 201,
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeOAuth,
|
||||
Credentials: map[string]any{
|
||||
"auth_mode": OpenAIAuthModeAgentIdentity,
|
||||
},
|
||||
}
|
||||
repo := &stubQuotaAccountRepo{accounts: map[int64]*Account{account.ID: account}}
|
||||
svc := &OpenAIQuotaService{accountRepo: repo}
|
||||
|
||||
_, err := svc.ResetCredit(context.Background(), account.ID)
|
||||
require.ErrorIs(t, err, ErrAgentIdentityResetNotSupported)
|
||||
require.Equal(t, http.StatusConflict, infraerrors.Code(err))
|
||||
}
|
||||
|
||||
// ── Part B: prepareUpstreamCall 影子 resolve ──────────────────────────────
|
||||
|
||||
// TestPrepareUpstreamCallShadowResolve 验证影子账号(200)QueryUsage 时:
|
||||
@@ -214,6 +244,101 @@ func TestPrepareUpstreamCallShadowResolve(t *testing.T) {
|
||||
"prepareUpstreamCall should use parent's chatgpt_account_id after shadow resolve")
|
||||
}
|
||||
|
||||
func TestQueryUsageAgentIdentityUsesAssertionWithoutOAuthToken(t *testing.T) {
|
||||
_, privateKey, err := ed25519.GenerateKey(rand.Reader)
|
||||
require.NoError(t, err)
|
||||
der, err := x509.MarshalPKCS8PrivateKey(privateKey)
|
||||
require.NoError(t, err)
|
||||
account := &Account{
|
||||
ID: 300,
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeOAuth,
|
||||
Credentials: map[string]any{
|
||||
"auth_mode": OpenAIAuthModeAgentIdentity,
|
||||
"agent_runtime_id": "runtime-quota",
|
||||
"agent_private_key": base64.StdEncoding.EncodeToString(der),
|
||||
"task_id": "task-quota",
|
||||
"chatgpt_account_id": "account-quota",
|
||||
"chatgpt_account_is_fedramp": true,
|
||||
},
|
||||
}
|
||||
repo := &stubQuotaAccountRepo{accounts: map[int64]*Account{account.ID: account}}
|
||||
var authorization string
|
||||
var accountHeader string
|
||||
var fedrampHeader string
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
authorization = r.Header.Get("authorization")
|
||||
accountHeader = r.Header.Get("chatgpt-account-id")
|
||||
fedrampHeader = r.Header.Get("x-openai-fedramp")
|
||||
w.Header().Set("content-type", "application/json")
|
||||
_, _ = w.Write([]byte(`{"plan_type":"pro","rate_limit":{"allowed":true}}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
svc := NewOpenAIQuotaService(repo, nil, nil, newQuotaRedirectingFactory(srv))
|
||||
usage, err := svc.QueryUsage(context.Background(), account.ID)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, usage)
|
||||
require.True(t, strings.HasPrefix(authorization, "AgentAssertion "))
|
||||
require.Equal(t, "account-quota", accountHeader)
|
||||
require.Equal(t, "true", fedrampHeader)
|
||||
}
|
||||
|
||||
func TestQueryUsageAgentIdentityRecoversInvalidTaskOnce(t *testing.T) {
|
||||
_, privateKey, err := ed25519.GenerateKey(rand.Reader)
|
||||
require.NoError(t, err)
|
||||
der, err := x509.MarshalPKCS8PrivateKey(privateKey)
|
||||
require.NoError(t, err)
|
||||
account := &Account{
|
||||
ID: 301,
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeOAuth,
|
||||
Credentials: map[string]any{
|
||||
"auth_mode": OpenAIAuthModeAgentIdentity,
|
||||
"agent_runtime_id": "runtime-quota-recovery",
|
||||
"agent_private_key": base64.StdEncoding.EncodeToString(der),
|
||||
"task_id": "task-quota-old",
|
||||
"chatgpt_account_id": "account-quota-recovery",
|
||||
},
|
||||
}
|
||||
repo := &stubQuotaAccountRepo{accounts: map[int64]*Account{account.ID: account}}
|
||||
usageCalls := 0
|
||||
registerCalls := 0
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("content-type", "application/json")
|
||||
if strings.Contains(r.URL.Path, "/task/register") {
|
||||
registerCalls++
|
||||
_, _ = w.Write([]byte(`{"task_id":"task-quota-new"}`))
|
||||
return
|
||||
}
|
||||
if strings.Contains(r.URL.Path, "rate-limit-reset-credits") {
|
||||
_, _ = w.Write([]byte(`{}`))
|
||||
return
|
||||
}
|
||||
usageCalls++
|
||||
if usageCalls == 1 {
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
_, _ = w.Write([]byte(`{"error":{"code":"invalid_task_id"}}`))
|
||||
return
|
||||
}
|
||||
_, _ = w.Write([]byte(`{"plan_type":"pro","rate_limit":{"allowed":true}}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
oldBase := openAIAgentIdentityAuthAPIBaseURL
|
||||
openAIAgentIdentityAuthAPIBaseURL = srv.URL
|
||||
t.Cleanup(func() { openAIAgentIdentityAuthAPIBaseURL = oldBase })
|
||||
|
||||
invalidator := &agentIdentityWSInvalidationRecorder{}
|
||||
svc := NewOpenAIQuotaService(repo, nil, nil, newQuotaRedirectingFactory(srv))
|
||||
svc.agentIdentityWS = invalidator
|
||||
usage, err := svc.QueryUsage(context.Background(), account.ID)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, usage)
|
||||
require.Equal(t, 2, usageCalls)
|
||||
require.Equal(t, 1, registerCalls)
|
||||
require.Equal(t, "task-quota-new", account.GetCredential("task_id"))
|
||||
require.Equal(t, []int64{account.ID}, invalidator.accountIDs)
|
||||
}
|
||||
|
||||
func TestParseOpenAIRateLimitResetCreditDetails_CompatibleContainers(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
|
||||
@@ -4,6 +4,7 @@ import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
@@ -68,6 +69,28 @@ type coderOpenAIWSClientDialer struct {
|
||||
proxyMisses atomic.Int64
|
||||
}
|
||||
|
||||
// openAIWSHandshakeError keeps a bounded, non-logged HTTP error body so the
|
||||
// Agent Identity recovery path can distinguish an invalid task from other
|
||||
// 401 handshake failures.
|
||||
type openAIWSHandshakeError struct {
|
||||
Body []byte
|
||||
Err error
|
||||
}
|
||||
|
||||
func (e *openAIWSHandshakeError) Error() string {
|
||||
if e == nil || e.Err == nil {
|
||||
return "openai ws handshake failed"
|
||||
}
|
||||
return e.Err.Error()
|
||||
}
|
||||
|
||||
func (e *openAIWSHandshakeError) Unwrap() error {
|
||||
if e == nil {
|
||||
return nil
|
||||
}
|
||||
return e.Err
|
||||
}
|
||||
|
||||
type openAIWSProxyClientEntry struct {
|
||||
client *http.Client
|
||||
lastUsedUnixNano int64
|
||||
@@ -104,7 +127,12 @@ func (d *coderOpenAIWSClientDialer) Dial(
|
||||
status = resp.StatusCode
|
||||
respHeaders = cloneHeader(resp.Header)
|
||||
}
|
||||
return nil, status, respHeaders, err
|
||||
var body []byte
|
||||
if resp != nil && resp.Body != nil {
|
||||
body, _ = io.ReadAll(io.LimitReader(resp.Body, 8<<10))
|
||||
_ = resp.Body.Close()
|
||||
}
|
||||
return nil, status, respHeaders, &openAIWSHandshakeError{Body: body, Err: err}
|
||||
}
|
||||
// coder/websocket 默认单消息读取上限为 32KB,Codex WS 事件(如 rate_limits/大 delta)
|
||||
// 可能超过该阈值,需显式提高上限,避免本地 read_fail(message too big)。
|
||||
|
||||
@@ -46,8 +46,8 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
|
||||
if account == nil {
|
||||
return errors.New("account is nil")
|
||||
}
|
||||
if strings.TrimSpace(token) == "" {
|
||||
return errors.New("token is empty")
|
||||
if err := validateOpenAIWSBearerToken(account, token); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// 预取一次 OpenAI Fast Policy settings,绑定到 ctx,让该 WS session
|
||||
@@ -577,6 +577,9 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
|
||||
Account: account,
|
||||
WSURL: wsURL,
|
||||
Headers: wsHeaders,
|
||||
HeadersFactory: func(factoryCtx context.Context, headers http.Header) (http.Header, error) {
|
||||
return s.refreshOpenAIAgentIdentityHeaders(factoryCtx, account, headers)
|
||||
},
|
||||
ProxyURL: func() string {
|
||||
if account.ProxyID != nil && account.Proxy != nil {
|
||||
return account.Proxy.URL()
|
||||
@@ -640,7 +643,9 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
|
||||
acquireTimeout = 30 * time.Second
|
||||
}
|
||||
|
||||
acquireTurnLease := func(turn int, preferred string, forcePreferredConn bool) (*openAIWSConnLease, error) {
|
||||
agentTaskRecoveryTried := false
|
||||
var acquireTurnLease func(int, string, bool) (*openAIWSConnLease, error)
|
||||
acquireTurnLease = func(turn int, preferred string, forcePreferredConn bool) (*openAIWSConnLease, error) {
|
||||
req := cloneOpenAIWSAcquireRequest(baseAcquireReq)
|
||||
req.PreferredConnID = strings.TrimSpace(preferred)
|
||||
req.ForcePreferredConn = forcePreferredConn
|
||||
@@ -649,6 +654,14 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
|
||||
acquireCtx, acquireCancel := context.WithTimeout(ctx, acquireTimeout)
|
||||
lease, acquireErr := pool.Acquire(acquireCtx, req)
|
||||
acquireCancel()
|
||||
var dialErr *openAIWSDialError
|
||||
if acquireErr != nil && s.isAgentIdentityAccount(ctx, account) && errors.As(acquireErr, &dialErr) && isAgentIdentityTaskInvalidWSDialError(dialErr) && !agentTaskRecoveryTried {
|
||||
agentTaskRecoveryTried = true
|
||||
if recoveryErr := s.recoverAgentIdentityTask(ctx, account, account.GetCredential("task_id")); recoveryErr != nil {
|
||||
return nil, fmt.Errorf("agent identity task recovery failed: %w", recoveryErr)
|
||||
}
|
||||
return acquireTurnLease(turn, preferred, forcePreferredConn)
|
||||
}
|
||||
if acquireErr != nil {
|
||||
dialStatus, dialClass, dialCloseStatus, dialCloseReason, dialRespServer, dialRespVia, dialRespCFRay, dialRespReqID := summarizeOpenAIWSDialError(acquireErr)
|
||||
logOpenAIWSModeInfo(
|
||||
|
||||
@@ -15,6 +15,16 @@ import (
|
||||
"github.com/tidwall/sjson"
|
||||
)
|
||||
|
||||
func validateOpenAIWSBearerToken(account *Account, token string) error {
|
||||
if account == nil {
|
||||
return errors.New("account is nil")
|
||||
}
|
||||
if strings.TrimSpace(token) == "" && !account.IsOpenAIAgentIdentity() {
|
||||
return errors.New("token is empty")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) buildOpenAIResponsesWSURL(account *Account) (string, error) {
|
||||
if account == nil {
|
||||
return "", errors.New("account is nil")
|
||||
@@ -67,7 +77,9 @@ func (s *OpenAIGatewayService) buildOpenAIWSHeaders(
|
||||
promptCacheKey string,
|
||||
) (http.Header, openAIWSSessionHeaderResolution, error) {
|
||||
headers := make(http.Header)
|
||||
headers.Set("authorization", "Bearer "+token)
|
||||
if account == nil || !account.IsOpenAIAgentIdentity() {
|
||||
headers.Set("authorization", "Bearer "+token)
|
||||
}
|
||||
|
||||
sessionResolution := resolveOpenAIWSSessionHeaders(c, promptCacheKey)
|
||||
if c != nil && c.Request != nil {
|
||||
|
||||
@@ -30,6 +30,7 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2(
|
||||
startTime time.Time,
|
||||
attempt int,
|
||||
lastFailureReason string,
|
||||
agentTaskRecoveryTried *bool,
|
||||
) (*OpenAIForwardResult, error) {
|
||||
if s == nil || account == nil {
|
||||
return nil, wrapOpenAIWSFallback("invalid_state", errors.New("service or account is nil"))
|
||||
@@ -174,9 +175,12 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2(
|
||||
defer acquireCancel()
|
||||
|
||||
lease, err := s.getOpenAIWSConnPool().Acquire(acquireCtx, openAIWSAcquireRequest{
|
||||
Account: account,
|
||||
WSURL: wsURL,
|
||||
Headers: wsHeaders,
|
||||
Account: account,
|
||||
WSURL: wsURL,
|
||||
Headers: wsHeaders,
|
||||
HeadersFactory: func(factoryCtx context.Context, headers http.Header) (http.Header, error) {
|
||||
return s.refreshOpenAIAgentIdentityHeaders(factoryCtx, account, headers)
|
||||
},
|
||||
PreferredConnID: preferredConnID,
|
||||
ForceNewConn: forceNewConn,
|
||||
ProxyURL: func() string {
|
||||
@@ -187,6 +191,14 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2(
|
||||
}(),
|
||||
})
|
||||
if err != nil {
|
||||
var agentDialErr *openAIWSDialError
|
||||
if s.isAgentIdentityAccount(ctx, account) && errors.As(err, &agentDialErr) && isAgentIdentityTaskInvalidWSDialError(agentDialErr) && agentTaskRecoveryTried != nil && !*agentTaskRecoveryTried {
|
||||
*agentTaskRecoveryTried = true
|
||||
if recoveryErr := s.recoverAgentIdentityTask(ctx, account, account.GetCredential("task_id")); recoveryErr != nil {
|
||||
return nil, fmt.Errorf("agent identity task recovery failed: %w", recoveryErr)
|
||||
}
|
||||
return nil, &agentIdentityTaskRecoveredError{}
|
||||
}
|
||||
dialStatus, dialClass, dialCloseStatus, dialCloseReason, dialRespServer, dialRespVia, dialRespCFRay, dialRespReqID := summarizeOpenAIWSDialError(err)
|
||||
logOpenAIWSModeInfo(
|
||||
"acquire_fail account_id=%d account_type=%s transport=%s reason=%s dial_status=%d dial_class=%s dial_close_status=%s dial_close_reason=%s dial_resp_server=%s dial_resp_via=%s dial_resp_cf_ray=%s dial_resp_x_request_id=%s cause=%s preferred_conn_id=%s force_new_conn=%v ws_host=%s ws_path=%s proxy_enabled=%v",
|
||||
|
||||
@@ -12,7 +12,6 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
@@ -193,7 +192,7 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn(
|
||||
releaseUpstreamCtx()
|
||||
return nil, fmt.Errorf("apply grok prompt cache identity: %w", err)
|
||||
}
|
||||
upstreamReq, err = buildGrokResponsesRequest(upstreamCtx, c, account, body, token, grokCacheIdentity)
|
||||
upstreamReq, err = buildGrokResponsesRequest(upstreamCtx, c, account, body, token, grokCacheIdentity, s.cfg)
|
||||
} else {
|
||||
upstreamReq, err = s.buildUpstreamRequestOpenAIPassthrough(upstreamCtx, c, account, body, token)
|
||||
}
|
||||
@@ -236,7 +235,7 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn(
|
||||
return nil, fmt.Errorf("upstream http bridge error: status=%d message=%s", resp.StatusCode, upstreamMsg)
|
||||
}
|
||||
if account.Platform == PlatformGrok {
|
||||
s.updateGrokUsageSnapshot(ctx, account, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode))
|
||||
s.updateGrokUsageFromResponse(ctx, account, resp.Header, resp.StatusCode)
|
||||
}
|
||||
|
||||
responseID := ""
|
||||
|
||||
@@ -39,6 +39,7 @@ var (
|
||||
type openAIWSDialError struct {
|
||||
StatusCode int
|
||||
ResponseHeaders http.Header
|
||||
ResponseBody []byte
|
||||
Err error
|
||||
}
|
||||
|
||||
@@ -60,9 +61,13 @@ func (e *openAIWSDialError) Unwrap() error {
|
||||
}
|
||||
|
||||
type openAIWSAcquireRequest struct {
|
||||
Account *Account
|
||||
WSURL string
|
||||
Headers http.Header
|
||||
Account *Account
|
||||
WSURL string
|
||||
Headers http.Header
|
||||
// HeadersFactory is evaluated inside dialConn. It exists so credentials
|
||||
// whose authorization is per-dial (Agent Identity) are never cached in
|
||||
// lastAcquire or delayed prewarm state.
|
||||
HeadersFactory func(context.Context, http.Header) (http.Header, error)
|
||||
ProxyURL string
|
||||
PreferredConnID string
|
||||
// ForceNewConn: 强制本次获取新连接(避免复用导致连接内续链状态互相污染)。
|
||||
@@ -544,6 +549,7 @@ type openAIWSAccountPool struct {
|
||||
pinnedConns map[string]int
|
||||
changedCh chan struct{}
|
||||
creating int
|
||||
generation uint64
|
||||
lastCleanupAt time.Time
|
||||
lastAcquire *openAIWSAcquireRequest
|
||||
prewarmActive bool
|
||||
@@ -1391,6 +1397,7 @@ func (p *openAIWSConnPool) ensureTargetIdleAsync(accountID int64) {
|
||||
}
|
||||
|
||||
var req openAIWSAcquireRequest
|
||||
generation := uint64(0)
|
||||
need := 0
|
||||
ap, ok := p.getAccountPool(accountID)
|
||||
if !ok || ap == nil {
|
||||
@@ -1425,6 +1432,7 @@ func (p *openAIWSConnPool) ensureTargetIdleAsync(accountID int64) {
|
||||
return
|
||||
}
|
||||
req = cloneOpenAIWSAcquireRequest(*ap.lastAcquire)
|
||||
generation = ap.generation
|
||||
ap.prewarmActive = true
|
||||
if cooldown := p.prewarmCooldown(); cooldown > 0 {
|
||||
ap.prewarmUntil = now.Add(cooldown)
|
||||
@@ -1432,7 +1440,7 @@ func (p *openAIWSConnPool) ensureTargetIdleAsync(accountID int64) {
|
||||
ap.creating += need
|
||||
p.metrics.scaleUpTotal.Add(int64(need))
|
||||
|
||||
go p.prewarmConns(accountID, req, need)
|
||||
go p.prewarmConns(accountID, req, need, generation)
|
||||
}
|
||||
|
||||
func (p *openAIWSConnPool) targetConnCountLocked(ap *openAIWSAccountPool, maxConns int) int {
|
||||
@@ -1475,7 +1483,11 @@ func (p *openAIWSConnPool) targetConnCountLocked(ap *openAIWSAccountPool, maxCon
|
||||
return target
|
||||
}
|
||||
|
||||
func (p *openAIWSConnPool) prewarmConns(accountID int64, req openAIWSAcquireRequest, total int) {
|
||||
func (p *openAIWSConnPool) prewarmConns(accountID int64, req openAIWSAcquireRequest, total int, generations ...uint64) {
|
||||
generation := uint64(0)
|
||||
if len(generations) > 0 {
|
||||
generation = generations[0]
|
||||
}
|
||||
defer func() {
|
||||
if ap, ok := p.getAccountPool(accountID); ok && ap != nil {
|
||||
ap.mu.Lock()
|
||||
@@ -1507,6 +1519,11 @@ func (p *openAIWSConnPool) prewarmConns(accountID int64, req openAIWSAcquireRequ
|
||||
ap.mu.Unlock()
|
||||
continue
|
||||
}
|
||||
if ap.generation != generation || ap.lastAcquire == nil {
|
||||
ap.mu.Unlock()
|
||||
conn.close()
|
||||
continue
|
||||
}
|
||||
if len(ap.conns) >= p.effectiveMaxConnsByAccount(req.Account) {
|
||||
ap.signalChangedLocked()
|
||||
ap.mu.Unlock()
|
||||
@@ -1521,6 +1538,35 @@ func (p *openAIWSConnPool) prewarmConns(accountID int64, req openAIWSAcquireRequ
|
||||
}
|
||||
}
|
||||
|
||||
// ClearAccount closes all pooled connections and discards delayed prewarm
|
||||
// state for one account. The generation guard prevents an in-flight prewarm
|
||||
// started before credential recovery from re-entering the pool afterwards.
|
||||
func (p *openAIWSConnPool) ClearAccount(accountID int64) {
|
||||
if p == nil || accountID <= 0 {
|
||||
return
|
||||
}
|
||||
ap, ok := p.getAccountPool(accountID)
|
||||
if !ok || ap == nil {
|
||||
return
|
||||
}
|
||||
ap.mu.Lock()
|
||||
ap.generation++
|
||||
conns := make([]*openAIWSConn, 0, len(ap.conns))
|
||||
for id, conn := range ap.conns {
|
||||
delete(ap.conns, id)
|
||||
delete(ap.pinnedConns, id)
|
||||
if conn != nil {
|
||||
conns = append(conns, conn)
|
||||
}
|
||||
}
|
||||
ap.lastAcquire = nil
|
||||
ap.prewarmUntil = time.Time{}
|
||||
ap.prewarmFails = 0
|
||||
ap.prewarmFailAt = time.Time{}
|
||||
ap.mu.Unlock()
|
||||
closeOpenAIWSConns(conns)
|
||||
}
|
||||
|
||||
func (p *openAIWSConnPool) evictConn(accountID int64, connID string) {
|
||||
if p == nil || accountID <= 0 || stringsTrim(connID) == "" {
|
||||
return
|
||||
@@ -1599,11 +1645,25 @@ func (p *openAIWSConnPool) dialConn(ctx context.Context, req openAIWSAcquireRequ
|
||||
if p == nil || p.clientDialer == nil {
|
||||
return nil, errors.New("openai ws client dialer is nil")
|
||||
}
|
||||
conn, status, handshakeHeaders, err := p.clientDialer.Dial(ctx, req.WSURL, req.Headers, req.ProxyURL)
|
||||
headers := cloneHeader(req.Headers)
|
||||
var err error
|
||||
if req.HeadersFactory != nil {
|
||||
headers, err = req.HeadersFactory(ctx, headers)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
conn, status, handshakeHeaders, err := p.clientDialer.Dial(ctx, req.WSURL, headers, req.ProxyURL)
|
||||
if err != nil {
|
||||
var handshakeErr *openAIWSHandshakeError
|
||||
var responseBody []byte
|
||||
if errors.As(err, &handshakeErr) && handshakeErr != nil {
|
||||
responseBody = append([]byte(nil), handshakeErr.Body...)
|
||||
}
|
||||
return nil, &openAIWSDialError{
|
||||
StatusCode: status,
|
||||
ResponseHeaders: cloneHeader(handshakeHeaders),
|
||||
ResponseBody: responseBody,
|
||||
Err: err,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -243,8 +243,8 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough(
|
||||
if account == nil {
|
||||
return errors.New("account is nil")
|
||||
}
|
||||
if strings.TrimSpace(token) == "" {
|
||||
return errors.New("token is empty")
|
||||
if err := validateOpenAIWSBearerToken(account, token); err != nil {
|
||||
return err
|
||||
}
|
||||
requestModel := strings.TrimSpace(gjson.GetBytes(firstClientMessage, "model").String())
|
||||
requestPreviousResponseID := strings.TrimSpace(gjson.GetBytes(firstClientMessage, "previous_response_id").String())
|
||||
@@ -359,10 +359,34 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough(
|
||||
return errors.New("openai ws passthrough dialer is nil")
|
||||
}
|
||||
|
||||
dialCtx, cancelDial := context.WithTimeout(ctx, s.openAIWSDialTimeout())
|
||||
defer cancelDial()
|
||||
upstreamConn, statusCode, handshakeHeaders, err := dialer.Dial(dialCtx, wsURL, headers, proxyURL)
|
||||
if err != nil {
|
||||
agentTaskRecoveryTried := false
|
||||
var upstreamConn openAIWSClientConn
|
||||
statusCode := 0
|
||||
var handshakeHeaders http.Header
|
||||
for {
|
||||
headers, err = s.refreshOpenAIAgentIdentityHeaders(ctx, account, headers)
|
||||
if err != nil {
|
||||
return fmt.Errorf("refresh ws authentication headers: %w", err)
|
||||
}
|
||||
dialCtx, cancelDial := context.WithTimeout(ctx, s.openAIWSDialTimeout())
|
||||
upstreamConn, statusCode, handshakeHeaders, err = dialer.Dial(dialCtx, wsURL, headers, proxyURL)
|
||||
cancelDial()
|
||||
if err == nil {
|
||||
break
|
||||
}
|
||||
var handshakeErr *openAIWSHandshakeError
|
||||
responseBody := []byte(nil)
|
||||
if errors.As(err, &handshakeErr) && handshakeErr != nil {
|
||||
responseBody = handshakeErr.Body
|
||||
}
|
||||
dialErr := &openAIWSDialError{StatusCode: statusCode, ResponseBody: responseBody, Err: err}
|
||||
if s.isAgentIdentityAccount(ctx, account) && isAgentIdentityTaskInvalidWSDialError(dialErr) && !agentTaskRecoveryTried {
|
||||
agentTaskRecoveryTried = true
|
||||
if recoveryErr := s.recoverAgentIdentityTask(ctx, account, account.GetCredential("task_id")); recoveryErr != nil {
|
||||
return fmt.Errorf("agent identity task recovery failed: %w", recoveryErr)
|
||||
}
|
||||
continue
|
||||
}
|
||||
logOpenAIWSV2Passthrough(
|
||||
"relay_dial_failed account_id=%d status_code=%d err=%s",
|
||||
account.ID,
|
||||
@@ -678,9 +702,15 @@ func (s *OpenAIGatewayService) mapOpenAIWSPassthroughDialError(
|
||||
wrappedErr := err
|
||||
var dialErr *openAIWSDialError
|
||||
if !errors.As(err, &dialErr) {
|
||||
var handshakeErr *openAIWSHandshakeError
|
||||
var responseBody []byte
|
||||
if errors.As(err, &handshakeErr) && handshakeErr != nil {
|
||||
responseBody = append([]byte(nil), handshakeErr.Body...)
|
||||
}
|
||||
wrappedErr = &openAIWSDialError{
|
||||
StatusCode: statusCode,
|
||||
ResponseHeaders: cloneHeader(handshakeHeaders),
|
||||
ResponseBody: responseBody,
|
||||
Err: err,
|
||||
}
|
||||
}
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user