fix(grok): fail over OAuth credential errors safely

This commit is contained in:
Heatherm Huang
2026-07-14 14:55:30 +08:00
parent 41cec0db05
commit 343390057d
63 changed files with 5291 additions and 522 deletions
@@ -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"))
@@ -71,7 +71,13 @@ func (s *FailoverState) HandleFailoverError(
retryLimit int,
failoverErr *service.UpstreamFailoverError,
) FailoverAction {
if ctx != nil && ctx.Err() != nil {
return FailoverCanceled
}
s.LastFailoverErr = failoverErr
if failoverErr == nil || !failoverErr.ShouldRetryNextAccount() {
return FailoverExhausted
}
// 缓存计费判断
if needForceCacheBilling(s.hasBoundSession, failoverErr) {
+47 -2
View File
@@ -2,6 +2,7 @@ package handler
import (
"context"
"net/http"
"testing"
"time"
@@ -128,6 +129,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)
@@ -452,8 +497,8 @@ 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])
// 入口已取消时不得改变任何重试状态。
require.Zero(t, fs.SameAccountRetryCount[100])
})
t.Run("Antigravity延迟期间context取消", func(t *testing.T) {
@@ -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 {
@@ -328,6 +331,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 {
@@ -307,6 +310,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
+7 -12
View File
@@ -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())
+13 -1
View File
@@ -174,6 +174,9 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
routingStart := time.Now()
for {
if requestCtx.Err() != nil {
return
}
selection, scheduleDecision, err := h.gatewayService.SelectAccountWithSchedulerForCapability(
requestCtx,
apiKey.GroupID,
@@ -257,11 +260,20 @@ 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 requestCtx.Err() != nil {
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 {
@@ -135,6 +135,9 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) {
var lastFailoverErr *service.UpstreamFailoverError
for {
if c.Request.Context().Err() != nil {
return
}
reqLog.Debug("openai_chat_completions.account_selecting", zap.Int("excluded_account_count", len(failedAccountIDs)))
selection, scheduleDecision, err := h.gatewayService.SelectAccountWithSchedulerForCapability(
c.Request.Context(),
@@ -231,11 +234,20 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) {
} else {
var failoverErr *service.UpstreamFailoverError
if errors.As(err, &failoverErr) {
if c.Request.Context().Err() != nil {
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()
@@ -0,0 +1,732 @@
//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
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 &copy, 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) 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
}
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
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
cancelRequest := u.cancelRequest
u.mu.Unlock()
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 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" {
accounts[0].Credentials["expires_at"] = time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339)
}
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{}
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)
}
@@ -334,6 +334,9 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
var lastFailoverErr *service.UpstreamFailoverError
for {
if c.Request.Context().Err() != nil {
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(
@@ -443,11 +446,20 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
} else {
var failoverErr *service.UpstreamFailoverError
if errors.As(err, &failoverErr) {
if c.Request.Context().Err() != nil {
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()
@@ -835,6 +847,9 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) {
effectiveMappedModel := preferredMappedModel
for {
if c.Request.Context().Err() != nil {
return
}
currentRoutingModel := routingModel
if effectiveMappedModel != "" {
currentRoutingModel = effectiveMappedModel
@@ -935,11 +950,20 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) {
} else {
var failoverErr *service.UpstreamFailoverError
if errors.As(err, &failoverErr) {
if c.Request.Context().Err() != nil {
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()
@@ -1093,6 +1117,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 +1496,49 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
switchCount := 0
failedAccountIDs := make(map[int64]struct{})
var lastFailoverErr *service.UpstreamFailoverError
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) {
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 +1606,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 +1775,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) {
@@ -1946,6 +2009,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 +2060,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 +2302,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")
+117 -65
View File
@@ -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":
+103
View File
@@ -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
@@ -1414,6 +1472,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
@@ -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) {
+37 -5
View File
@@ -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)
}
+10 -23
View File
@@ -1268,16 +1268,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,27 +1283,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
return a.GetGrokBaseURL()
}
func isOfficialGrokAPIBaseURL(raw string) bool {
@@ -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())
}
@@ -724,7 +724,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 {
@@ -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",
},
@@ -133,8 +134,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 +168,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}}
+61 -1
View File
@@ -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,650 @@
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())
return actual.(*oauthRefreshLocalLock)
}
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
+5 -2
View File
@@ -304,7 +304,7 @@ 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
}
@@ -339,7 +339,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 == "" {
@@ -635,6 +637,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,9 @@ package service
import (
"context"
"crypto/subtle"
"encoding/base64"
"encoding/json"
"errors"
"net/http"
"strings"
"time"
@@ -314,10 +317,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
}
@@ -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 {
@@ -387,6 +389,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 +399,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 +411,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()
}
}
@@ -180,19 +180,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 +243,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 +270,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},
@@ -352,17 +341,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,16 +374,7 @@ 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)
repo := &grokQuotaAccountRepo{
mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
accountsByID: map[int64]*Account{45: account},
@@ -435,15 +406,7 @@ func TestGrokQuotaServiceProbeUsageStoresNoHeadersState(t *testing.T) {
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},
+164 -54
View File
@@ -3,19 +3,26 @@ 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")
errGrokOAuthRefreshLockTimeout = errors.New("grok oauth refresh lock wait timed out")
)
type GrokTokenCache = GeminiTokenCache
@@ -36,7 +43,7 @@ func NewGrokTokenProvider(
return &GrokTokenProvider{
accountRepo: accountRepo,
tokenCache: tokenCache,
refreshPolicy: AntigravityProviderRefreshPolicy(),
refreshPolicy: GrokProviderRefreshPolicy(),
}
}
@@ -60,32 +67,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 +125,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 +168,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 {
+110 -22
View File
@@ -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"
)
@@ -96,6 +97,20 @@ 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{})
return actual.(*sync.Mutex)
}
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 +122,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 +130,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 +147,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 +168,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
}
@@ -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)
}
@@ -109,7 +109,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 +162,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 +189,7 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions(
return nil, &UpstreamFailoverError{
StatusCode: resp.StatusCode,
ResponseBody: respBody,
ResponseHeaders: resp.Header.Clone(),
RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode),
}
}
@@ -54,7 +54,7 @@ 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
}
@@ -100,6 +100,7 @@ func (s *OpenAIGatewayService) forwardGrokResponses(
return nil, &UpstreamFailoverError{
StatusCode: resp.StatusCode,
ResponseBody: respBody,
ResponseHeaders: resp.Header.Clone(),
RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode),
}
}
@@ -610,6 +611,7 @@ func (s *OpenAIGatewayService) describeGrokComposerImage(
return "", OpenAIUsage{}, &UpstreamFailoverError{
StatusCode: resp.StatusCode,
ResponseBody: respBody,
ResponseHeaders: resp.Header.Clone(),
RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode),
}
}
@@ -750,7 +752,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) != "" {
@@ -253,7 +253,7 @@ 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)
}
@@ -294,6 +294,7 @@ func (s *OpenAIGatewayService) forwardGrokChatCompletionsViaResponses(
return nil, &UpstreamFailoverError{
StatusCode: resp.StatusCode,
ResponseBody: respBody,
ResponseHeaders: resp.Header.Clone(),
RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode),
}
}
@@ -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,
},
}
}
@@ -259,7 +259,7 @@ func TestBuildGrokResponsesRequestUsesAccountBaseURLAndBearerToken(t *testing.T)
req, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "access-token", "isolated-cache-id")
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")
@@ -284,6 +284,8 @@ func TestBuildGrokResponsesRequestAllowsPublicAPIKeyBaseURLByDefault(t *testing.
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) {
@@ -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)
@@ -1017,18 +1086,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,18 +1143,7 @@ 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")
repo := &grokQuotaAccountRepo{
mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
accountsByID: map[int64]*Account{56: account},
@@ -1145,18 +1192,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 +1251,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 +1311,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 +1374,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 +1424,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},
@@ -258,7 +258,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)
}
@@ -401,6 +401,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
+8 -7
View File
@@ -27,7 +27,7 @@ type OpsErrorLog struct {
CreatedAt time.Time `json:"created_at"`
// Standardized classification
// - phase: request|auth|routing|upstream|network|internal
// - phase: request|auth|account_auth|routing|upstream|network|internal
// - owner: client|provider|platform
// - source: client_request|upstream_http|gateway
Phase string `json:"phase"`
@@ -121,7 +121,7 @@ type OpsErrorLogFilter struct {
StatusCodes []int
StatusCodesOther bool
Phase string // Special: Phase=="upstream" bypasses status>=400 clause; do not set together with ErrorPhasesAny.
Phase string // Recovered provider rows bypass status>=400 only with the explicit opt-in below.
Owner string
Source string
Resolved *bool
@@ -150,14 +150,15 @@ type OpsErrorLogFilter struct {
// ExcludeCountTokens drops count_tokens probe errors (is_count_tokens=true).
ExcludeCountTokens bool
// IncludeRecoveredUpstream 显式豁免 status>=400 守卫(仅在 Phase=="upstream" 时生效):
// ops 专用上游错误列表需要看到 status<400 的 recovered upstream 行。
// 请求错误语义的端点不设此开关,phase=upstream 过滤照常生效且守卫保留。
// IncludeRecoveredUpstream explicitly exempts provider-health phases
// (upstream and account_auth) from the status>=400 guard. Ops provider
// health lists need status<400 recovered rows; request-error endpoints do
// not set this flag and retain client-error semantics.
IncludeRecoveredUpstream bool
// ErrorPhasesAny / ErrorTypesAny add plain ANY() filters WITHOUT touching the
// special-cased single `Phase` field (only Phase=="upstream" with
// IncludeRecoveredUpstream bypasses the status>=400 clause).
// special-cased single `Phase` field. With IncludeRecoveredUpstream, an ANY
// list containing only upstream/account_auth also bypasses status>=400.
// NOTE: these ANY filters do NOT bypass status>=400; records with error_phase='upstream'
// but status_code<400 (recovered upstream errors) remain excluded.
// Used to map user-facing coarse categories to backend conditions.
+28 -1
View File
@@ -222,6 +222,33 @@ func (s *OpsService) prepareErrorLogInput(ctx context.Context, entry *OpsInsertE
entry.ErrorType = "api_error"
}
// Credential acquisition is a gateway/account-auth stage, not an inference
// HTTP attempt. Enforce that ownership at the persistence boundary so an
// earlier inference attempt cannot leak its status or text into top-level
// auth fields even if a caller supplied stale single-value context.
for i := len(entry.UpstreamErrors) - 1; i >= 0; i-- {
last := entry.UpstreamErrors[i]
if last == nil {
continue
}
if last.Stage == string(GatewayFailureStageAccountAuth) {
entry.ErrorPhase = string(GatewayFailureStageAccountAuth)
entry.ErrorOwner = "provider"
entry.ErrorSource = "gateway"
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
}
}
break
}
// Sanitize + truncate error_body to avoid storing sensitive data.
if strings.TrimSpace(entry.ErrorBody) != "" {
sanitized, _ := sanitizeErrorBodyForStorage(entry.ErrorBody, opsMaxStoredErrorBodyBytes)
@@ -229,7 +256,7 @@ func (s *OpsService) prepareErrorLogInput(ctx context.Context, entry *OpsInsertE
}
// Sanitize upstream error context if provided by gateway services.
if entry.UpstreamStatusCode != nil && *entry.UpstreamStatusCode <= 0 {
if entry.UpstreamStatusCode != nil && *entry.UpstreamStatusCode <= 0 && entry.ErrorPhase != string(GatewayFailureStageAccountAuth) {
entry.UpstreamStatusCode = nil
}
if entry.UpstreamErrorMessage != nil {
@@ -97,6 +97,54 @@ func TestOpsServiceRecordErrorBatch_FallsBackToSingleInsert(t *testing.T) {
require.Equal(t, 2, singleCalls)
}
func TestOpsServiceRecordErrorPersistsExplicitAccountAuthStatusZero(t *testing.T) {
t.Parallel()
var captured *OpsInsertErrorLogInput
repo := &opsRepoMock{
InsertErrorLogFn: func(_ context.Context, input *OpsInsertErrorLogInput) (int64, error) {
captured = input
return 1, nil
},
}
svc := NewOpsService(repo, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
staleStatus := 403
staleMessage := "stale inference message"
staleDetail := "stale inference detail"
err := svc.RecordError(context.Background(), &OpsInsertErrorLogInput{
ErrorPhase: "upstream",
ErrorType: "upstream_error",
ErrorOwner: "provider",
ErrorSource: "upstream_http",
UpstreamStatusCode: &staleStatus,
UpstreamErrorMessage: &staleMessage,
UpstreamErrorDetail: &staleDetail,
UpstreamErrors: []*OpsUpstreamErrorEvent{
{Stage: string(GatewayFailureStageInference), UpstreamStatusCode: 403, Message: staleMessage, Detail: staleDetail},
{
Stage: string(GatewayFailureStageAccountAuth), Scope: string(GatewayFailureScopeAccount),
Reason: string(GrokCredentialReasonRevoked), Message: "Grok OAuth credentials require account action",
},
},
})
require.NoError(t, err)
require.NotNil(t, captured)
require.Equal(t, "account_auth", captured.ErrorPhase)
require.Equal(t, "provider", captured.ErrorOwner)
require.Equal(t, "gateway", captured.ErrorSource)
require.NotNil(t, captured.UpstreamStatusCode)
require.Zero(t, *captured.UpstreamStatusCode)
require.NotNil(t, captured.UpstreamErrorMessage)
require.Equal(t, "Grok OAuth credentials require account action", *captured.UpstreamErrorMessage)
require.Nil(t, captured.UpstreamErrorDetail)
require.Nil(t, captured.UpstreamErrors)
require.NotNil(t, captured.UpstreamErrorsJSON)
require.Contains(t, *captured.UpstreamErrorsJSON, `"upstream_status_code":403`)
require.Contains(t, *captured.UpstreamErrorsJSON, `"stage":"account_auth"`)
}
func strPtr(v string) *string {
return &v
}
@@ -188,6 +188,11 @@ type OpsUpstreamErrorEvent struct {
// Kind: http_error | request_error | retry_exhausted | failover
Kind string `json:"kind,omitempty"`
// Stage/Scope/Reason distinguish credential acquisition from inference
// without overloading upstream_status_code with a synthetic HTTP status.
Stage string `json:"stage,omitempty"`
Scope string `json:"scope,omitempty"`
Reason string `json:"reason,omitempty"`
Message string `json:"message,omitempty"`
Detail string `json:"detail,omitempty"`
@@ -204,6 +209,9 @@ func appendOpsUpstreamError(c *gin.Context, ev OpsUpstreamErrorEvent) {
ev.UpstreamRequestID = strings.TrimSpace(ev.UpstreamRequestID)
ev.UpstreamResponseBody = strings.TrimSpace(ev.UpstreamResponseBody)
ev.Kind = strings.TrimSpace(ev.Kind)
ev.Stage = strings.TrimSpace(ev.Stage)
ev.Scope = strings.TrimSpace(ev.Scope)
ev.Reason = strings.TrimSpace(ev.Reason)
ev.UpstreamURL = strings.TrimSpace(ev.UpstreamURL)
ev.Message = strings.TrimSpace(ev.Message)
ev.Detail = strings.TrimSpace(ev.Detail)
+2 -2
View File
@@ -44,7 +44,7 @@ func MapUserErrorCategory(phase, errType string) string {
return "auth"
case "routing":
return "service_unavailable"
case "upstream", "network":
case "account_auth", "upstream", "network":
return "upstream"
case "internal":
return "internal"
@@ -73,7 +73,7 @@ func CategoryToFilter(category string) (phases []string, errorTypes []string) {
case "service_unavailable":
return []string{"routing"}, nil
case "upstream":
return []string{"upstream", "network"}, nil
return []string{"account_auth", "upstream", "network"}, nil
case "internal":
return []string{"internal"}, nil
case "rate_limit":
@@ -17,6 +17,7 @@ func TestMapUserErrorCategory(t *testing.T) {
{"request", "subscription_error", "quota"},
{"request", "invalid_request_error", "invalid_request"},
{"routing", "api_error", "service_unavailable"},
{"account_auth", "upstream_error", "upstream"},
{"upstream", "upstream_error", "upstream"},
{"network", "api_error", "upstream"},
{"internal", "api_error", "internal"},
@@ -43,7 +44,7 @@ func TestCategoryToFilter(t *testing.T) {
t.Fatalf("service_unavailable => phases=%v types=%v", phases, types)
}
phases, types = CategoryToFilter("upstream")
if len(phases) != 2 || phases[0] != "upstream" || phases[1] != "network" || len(types) != 0 {
if len(phases) != 3 || phases[0] != "account_auth" || phases[1] != "upstream" || phases[2] != "network" || len(types) != 0 {
t.Fatalf("upstream => phases=%v types=%v", phases, types)
}
phases, types = CategoryToFilter("internal")
@@ -61,6 +61,14 @@ func AntigravityProviderRefreshPolicy() ProviderRefreshPolicy {
}
}
func GrokProviderRefreshPolicy() ProviderRefreshPolicy {
return ProviderRefreshPolicy{
OnRefreshError: ProviderRefreshErrorReturn,
OnLockHeld: ProviderLockHeldWaitForCache,
FailureTTL: 0,
}
}
// BackgroundSkipAction 定义后台刷新服务在“未实际刷新”场景的计数方式。
type BackgroundSkipAction int
@@ -26,6 +26,9 @@ type tokenRefreshAccountRepo struct {
lastExtraUpdates map[string]any
lastAccount *Account
updateErr error
setErrorErr error
setTempUnschedErr error
beforeConditionalState func()
}
func (r *tokenRefreshAccountRepo) Update(ctx context.Context, account *Account) error {
@@ -56,7 +59,7 @@ func (r *tokenRefreshAccountRepo) UpdateCredentials(ctx context.Context, id int6
func (r *tokenRefreshAccountRepo) SetError(ctx context.Context, id int64, errorMsg string) error {
r.setErrorCalls++
r.lastErrorMessage = errorMsg
return nil
return r.setErrorErr
}
func (r *tokenRefreshAccountRepo) ClearTempUnschedulable(ctx context.Context, id int64) error {
@@ -67,7 +70,66 @@ func (r *tokenRefreshAccountRepo) ClearTempUnschedulable(ctx context.Context, id
func (r *tokenRefreshAccountRepo) SetTempUnschedulable(ctx context.Context, id int64, until time.Time, reason string) error {
r.setTempUnschedCalls++
r.lastTempUnschedReason = reason
return nil
return r.setTempUnschedErr
}
func (r *tokenRefreshAccountRepo) SetGrokCredentialErrorIfMatch(
_ context.Context,
id int64,
snapshot GrokCredentialMutationSnapshot,
errorMsg string,
) (bool, error) {
if r.beforeConditionalState != nil {
hook := r.beforeConditionalState
r.beforeConditionalState = nil
hook()
}
account := r.accountsByID[id]
if !grokCredentialSnapshotMatchesAccount(account, snapshot) ||
(errorMsg == string(GrokCredentialReasonProxyInvalid) && account.Proxy != nil) {
return false, nil
}
r.setErrorCalls++
r.lastErrorMessage = errorMsg
if r.setErrorErr != nil {
return false, r.setErrorErr
}
account.Status = StatusError
account.Schedulable = false
account.ErrorMessage = errorMsg
return true, nil
}
func (r *tokenRefreshAccountRepo) SetGrokCredentialTempUnschedulableIfMatch(
_ context.Context,
id int64,
snapshot GrokCredentialMutationSnapshot,
until time.Time,
reason string,
) (bool, error) {
if r.beforeConditionalState != nil {
hook := r.beforeConditionalState
r.beforeConditionalState = nil
hook()
}
account := r.accountsByID[id]
if !grokCredentialSnapshotMatchesAccount(account, snapshot) {
return false, nil
}
r.setTempUnschedCalls++
r.lastTempUnschedReason = reason
if r.setTempUnschedErr != nil {
return false, r.setTempUnschedErr
}
value := until
account.TempUnschedulableUntil = &value
return true, nil
}
func grokCredentialSnapshotMatchesAccount(account *Account, snapshot GrokCredentialMutationSnapshot) bool {
return account != nil && account.IsGrokOAuth() && account.IsSchedulable() &&
grokCredentialMutationSnapshot(account).CredentialsJSON == snapshot.CredentialsJSON &&
grokCredentialProxyIDsEqual(account.ProxyID, snapshot.ProxyID)
}
func (r *tokenRefreshAccountRepo) UpdateExtra(ctx context.Context, id int64, updates map[string]any) error {
@@ -750,6 +812,7 @@ func TestPathA_Success(t *testing.T) {
ID: 100,
Platform: PlatformGemini,
Type: AccountTypeOAuth,
Status: StatusActive,
}
repo := &tokenRefreshAccountRepo{}
repo.accountsByID = map[int64]*Account{account.ID: account}
@@ -771,6 +834,7 @@ func TestPathA_LockHeld(t *testing.T) {
ID: 101,
Platform: PlatformGemini,
Type: AccountTypeOAuth,
Status: StatusActive,
}
repo := &tokenRefreshAccountRepo{}
invalidator := &tokenCacheInvalidatorStub{}
@@ -791,6 +855,7 @@ func TestPathA_AlreadyRefreshed(t *testing.T) {
ID: 102,
Platform: PlatformGemini,
Type: AccountTypeOAuth,
Status: StatusActive,
}
repo := &tokenRefreshAccountRepo{}
repo.accountsByID = map[int64]*Account{account.ID: account}
@@ -830,6 +895,7 @@ func TestPathA_NonRetryableError(t *testing.T) {
ID: 103,
Platform: PlatformGemini,
Type: AccountTypeOAuth,
Status: StatusActive,
}
repo := &tokenRefreshAccountRepo{}
repo.accountsByID = map[int64]*Account{account.ID: account}
@@ -855,6 +921,7 @@ func TestPathA_RetryableErrorExhausted(t *testing.T) {
ID: 104,
Platform: PlatformGemini,
Type: AccountTypeOAuth,
Status: StatusActive,
}
repo := &tokenRefreshAccountRepo{}
repo.accountsByID = map[int64]*Account{account.ID: account}
@@ -888,6 +955,7 @@ func TestPathA_DBUpdateFailed(t *testing.T) {
ID: 105,
Platform: PlatformGemini,
Type: AccountTypeOAuth,
Status: StatusActive,
}
repo := &tokenRefreshAccountRepo{updateErr: errors.New("db connection lost")}
repo.accountsByID = map[int64]*Account{account.ID: account}
@@ -898,7 +966,7 @@ func TestPathA_DBUpdateFailed(t *testing.T) {
err := service.refreshWithRetry(context.Background(), account, refresher, refresher, time.Hour)
require.Error(t, err)
require.Contains(t, err.Error(), "DB update failed")
require.ErrorIs(t, err, errOAuthRefreshCredentialPersist)
require.Equal(t, 1, repo.updateCalls) // DB 更新被尝试
require.Equal(t, 0, invalidator.calls) // DB 失败时不应触发缓存失效
}
+1 -1
View File
@@ -186,7 +186,7 @@ func ProvideGrokTokenProvider(
p := NewGrokTokenProvider(accountRepo, tokenCache)
executor := NewGrokTokenRefresher(grokOAuthService)
p.SetRefreshAPI(refreshAPI, executor)
p.SetRefreshPolicy(AntigravityProviderRefreshPolicy())
p.SetRefreshPolicy(GrokProviderRefreshPolicy())
p.SetTempUnschedCache(tempUnschedCache)
return p
}
@@ -277,6 +277,7 @@ const billingTypeOptions = ref<SelectOption[]>([
const errorPhaseOptions = computed<SelectOption[]>(() => [
{ value: null, label: t('admin.usage.allTypes') },
{ value: 'upstream', label: t('admin.ops.errorLog.typeUpstream') },
{ value: 'account_auth', label: t('admin.ops.errorLog.typeAccountAuth') },
{ value: 'request', label: t('admin.ops.errorLog.typeRequest') },
{ value: 'auth', label: t('admin.ops.errorLog.typeAuth') },
{ value: 'routing', label: t('admin.ops.errorLog.typeRouting') },
@@ -266,6 +266,7 @@ export default {
typeUpstream: 'Upstream',
typeRequest: 'Request',
typeAuth: 'Auth',
typeAccountAuth: 'Account Auth',
typeRouting: 'Routing',
typeInternal: 'Internal',
endpoint: 'Endpoint',
@@ -291,6 +292,7 @@ export default {
phase: {
request: 'Request',
auth: 'Auth',
account_auth: 'Account Auth',
routing: 'Routing',
upstream: 'Upstream',
network: 'Network',
@@ -266,6 +266,7 @@ export default {
typeUpstream: '上游',
typeRequest: '请求',
typeAuth: '认证',
typeAccountAuth: '账号认证',
typeRouting: '路由',
typeInternal: '内部',
endpoint: '端点',
@@ -291,6 +292,7 @@ export default {
phase: {
request: '请求',
auth: '认证',
account_auth: '账号认证',
routing: '路由',
upstream: '上游',
network: '网络',
+1
View File
@@ -9,6 +9,7 @@ export function mapErrorCategory(phase?: string | null, errType?: string | null)
return 'auth'
case 'routing':
return 'service_unavailable'
case 'account_auth':
case 'upstream':
case 'network':
return 'upstream'
@@ -72,6 +72,7 @@ const phaseSelectOptions = computed(() => {
{ value: '', label: t('common.all') },
{ value: 'request', label: t('admin.ops.errorDetails.phase.request') || 'request' },
{ value: 'auth', label: t('admin.ops.errorDetails.phase.auth') || 'auth' },
{ value: 'account_auth', label: t('admin.ops.errorDetails.phase.account_auth') || 'account_auth' },
{ value: 'routing', label: t('admin.ops.errorDetails.phase.routing') || 'routing' },
{ value: 'upstream', label: t('admin.ops.errorDetails.phase.upstream') || 'upstream' },
{ value: 'network', label: t('admin.ops.errorDetails.phase.network') || 'network' },
@@ -277,6 +277,9 @@ function getTypeBadge(log: OpsErrorLog): { label: string; className: string } {
if (phase === 'auth' && owner === 'client') {
return { label: t('admin.ops.errorLog.typeAuth'), className: 'bg-blue-100 text-blue-800 dark:bg-blue-900 dark:text-blue-200' }
}
if (phase === 'account_auth') {
return { label: t('admin.ops.errorLog.typeAccountAuth'), className: 'bg-orange-100 text-orange-800 dark:bg-orange-900 dark:text-orange-200' }
}
if (phase === 'routing' && owner === 'platform') {
return { label: t('admin.ops.errorLog.typeRouting'), className: 'bg-purple-100 text-purple-800 dark:bg-purple-900 dark:text-purple-200' }
}