fix(grok): let manual connection tests bypass the scheduling gate

Grok OAuth 'test connection' reused the production GrokTokenProvider
request path, whose eligibility gate rejects any account that is not
currently schedulable (manual switch off, rate limited, overloaded,
temporarily cooled down). Every such test failed before reaching xAI
with the generic 'oauth refresh account state changed', while
Codex/OpenAI account tests read credentials regardless of scheduling
state (#4598).

Add GetAccessTokenForManualTest and use it from the account test
service:

- skips the request-path scheduling eligibility gate so admins can
  probe excluded accounts; credential integrity checks (configured
  proxy present, refresh token present) still apply
- refreshes expired tokens through the shared OAuthRefreshAPI without
  the request-path marker, so the lock protocol and DB re-read keep
  working while scheduling state no longer blocks the refresh;
  credential rotation for non-active accounts stays blocked inside
  RefreshIfNeeded
- falls back to the still-valid current token when a refresh attempt
  fails or another worker holds the refresh lock, and otherwise
  surfaces the concrete refresh error instead of the generic message

Fixes #4598
This commit is contained in:
superman2003
2026-07-19 22:40:58 +08:00
parent d4b9797ff7
commit 4a97491be6
3 changed files with 218 additions and 1 deletions
@@ -711,7 +711,9 @@ func (s *AccountTestService) testGrokAccountConnection(c *gin.Context, account *
return s.sendErrorAndEnd(c, "Grok token provider not configured")
}
var err error
authToken, err = s.grokTokenProvider.GetAccessToken(ctx, account)
// 手动测试不走生产调度资格门:关闭调度、限流/过载/临时冷却中的账号
// 也应能被管理员探测(#4598),与 Codex/OpenAI 测试行为一致。
authToken, err = s.grokTokenProvider.GetAccessTokenForManualTest(ctx, account)
if err != nil {
return s.sendErrorAndEnd(c, fmt.Sprintf("Failed to get Grok access token: %s", err.Error()))
}
@@ -167,6 +167,77 @@ func (p *GrokTokenProvider) GetAccessToken(ctx context.Context, account *Account
return accessToken, nil
}
// GetAccessTokenForManualTest returns an access token for an admin-initiated
// "test connection" probe. Unlike GetAccessToken it does not apply the
// request-path scheduling eligibility gate (manual Schedulable switch,
// rate-limit / overload / temp-unschedulable cooldowns): a manual test exists
// precisely to check accounts in those states, matching how Codex/OpenAI
// account tests read credentials regardless of scheduling state (#4598).
//
// Credential integrity still applies: the configured-proxy-missing check, the
// shared refresh lock protocol, and the refresh API's own account re-read.
// Credential rotation for non-active (disabled/error) accounts remains
// blocked inside RefreshIfNeeded; their still-valid tokens are probed as-is.
func (p *GrokTokenProvider) GetAccessTokenForManualTest(ctx context.Context, account *Account) (string, error) {
if account == nil {
return "", errors.New("account is nil")
}
if account.Platform != PlatformGrok || account.Type != AccountTypeOAuth {
return "", errors.New("not a grok oauth account")
}
if account.ProxyID != nil && account.Proxy == nil {
return "", errGrokOAuthConfiguredProxyMiss
}
if strings.TrimSpace(account.GetGrokRefreshToken()) == "" {
return "", errGrokOAuthRefreshTokenMissing
}
accessToken := strings.TrimSpace(account.GetGrokAccessToken())
expiresAt := account.GetCredentialAsTime("expires_at")
tokenValid := accessToken != "" && expiresAt != nil && time.Now().Before(*expiresAt)
if accessToken != "" && expiresAt != nil && time.Until(*expiresAt) > grokTokenRefreshSkew {
return accessToken, nil
}
if p.refreshAPI == nil || p.executor == nil {
if tokenValid {
return accessToken, nil
}
return "", errGrokOAuthRefreshNotConfigured
}
// Deliberately not marked as a request-path refresh: the request path
// re-applies scheduling eligibility inside RefreshIfNeeded, which is
// exactly what a manual test must bypass.
refreshCtx, cancel := context.WithTimeout(ctx, grokRequestRefreshTimeout)
defer cancel()
result, err := p.refreshAPI.RefreshIfNeeded(refreshCtx, account, p.executor, grokTokenRefreshSkew)
if err != nil {
if tokenValid {
return accessToken, nil
}
return "", err
}
if result != nil && result.LockHeld {
if tokenValid {
return accessToken, nil
}
return "", errors.New("token refresh is already in progress on another worker; retry in a few seconds")
}
if result != nil && result.Account != nil {
account = result.Account
}
accessToken = strings.TrimSpace(account.GetGrokAccessToken())
if accessToken == "" {
return "", errGrokOAuthAccessTokenMissing
}
if latestExpiry := account.GetCredentialAsTime("expires_at"); latestExpiry != nil && !time.Now().Before(*latestExpiry) {
return "", errGrokOAuthAccessTokenExpired
}
return accessToken, nil
}
func (p *GrokTokenProvider) waitForRefreshedToken(ctx context.Context, account *Account, cacheKey string) (string, error) {
waitCtx, cancel := context.WithTimeout(ctx, grokRefreshLockWaitTimeout)
defer cancel()
@@ -287,6 +287,150 @@ func TestGrokTokenProviderRejectsStaleDBTokenWithoutExpiry(t *testing.T) {
require.Empty(t, token)
}
// TestGrokTokenProviderManualTestBypassesSchedulingGate reproduces #4598:
// admins must be able to run "test connection" against accounts that the
// scheduler currently excludes (manual switch off, rate limited, overloaded,
// temporarily cooled down). The production request path keeps rejecting them.
func TestGrokTokenProviderManualTestBypassesSchedulingGate(t *testing.T) {
future := time.Now().Add(time.Hour)
tests := []struct {
name string
mutate func(*Account)
}{
{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 }},
{name: "disabled by error", mutate: func(account *Account) { account.Status = StatusError }},
}
for index, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
account := expiredGrokOAuthAccountForCredentialTest(int64(120 + index))
account.Credentials["access_token"] = "still-valid-token"
account.Credentials["expires_at"] = time.Now().Add(2 * grokTokenRefreshSkew).UTC().Format(time.RFC3339)
tt.mutate(account)
provider := NewGrokTokenProvider(&tokenRefreshAccountRepo{}, &grokTokenCacheForProviderTest{})
// Production request path keeps excluding this account.
_, requestErr := provider.GetAccessToken(context.Background(), account)
require.ErrorIs(t, requestErr, errOAuthRefreshAccountStateChanged)
// Manual test path returns the valid credential for probing.
token, err := provider.GetAccessTokenForManualTest(context.Background(), account)
require.NoError(t, err)
require.Equal(t, "still-valid-token", token)
})
}
}
func TestGrokTokenProviderManualTestRefreshesExpiredTokenWhileUnschedulable(t *testing.T) {
t.Setenv(xai.EnvBaseURL, xai.DefaultCLIBaseURL)
expiredAt := time.Now().Add(-time.Minute).UTC().Format(time.RFC3339)
account := &Account{
ID: 130,
Platform: PlatformGrok,
Type: AccountTypeOAuth,
Status: StatusActive,
Schedulable: false,
Credentials: map[string]any{
"access_token": "expired-access-token",
"refresh_token": "refresh-token",
"expires_at": expiredAt,
"base_url": xai.DefaultCLIBaseURL,
"client_id": "client-id",
},
}
repo := &tokenRefreshAccountRepo{}
repo.accountsByID = map[int64]*Account{130: account}
cache := &grokTokenCacheForProviderTest{lockResult: true}
oauthSvc := NewGrokOAuthService(nil, &grokOAuthClientStub{
refreshResponse: &xai.TokenResponse{
AccessToken: "manual-test-refreshed-token",
TokenType: "Bearer",
ExpiresIn: 3600,
},
})
defer oauthSvc.Stop()
provider := NewGrokTokenProvider(repo, cache)
provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), NewGrokTokenRefresher(oauthSvc))
token, err := provider.GetAccessTokenForManualTest(context.Background(), account)
require.NoError(t, err)
require.Equal(t, "manual-test-refreshed-token", token)
require.Equal(t, 1, repo.updateCredentialsCalls)
}
func TestGrokTokenProviderManualTestFallsBackToValidTokenOnRefreshFailure(t *testing.T) {
account := &Account{
ID: 131,
Platform: PlatformGrok,
Type: AccountTypeOAuth,
Status: StatusActive,
Schedulable: false,
Credentials: map[string]any{
"access_token": "near-expiry-token",
"refresh_token": "refresh-token",
// Inside the refresh window but not expired yet.
"expires_at": time.Now().Add(10 * time.Minute).UTC().Format(time.RFC3339),
},
}
repo := &tokenRefreshAccountRepo{}
repo.accountsByID = map[int64]*Account{131: account}
cache := &grokTokenCacheForProviderTest{lockResult: true}
provider := NewGrokTokenProvider(repo, cache)
provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{
err: errors.New("upstream refresh unavailable"),
})
token, err := provider.GetAccessTokenForManualTest(context.Background(), account)
require.NoError(t, err)
require.Equal(t, "near-expiry-token", token)
}
func TestGrokTokenProviderManualTestReportsRefreshFailureWhenTokenExpired(t *testing.T) {
account := expiredGrokOAuthAccountForCredentialTest(132)
account.Schedulable = false
repo := &tokenRefreshAccountRepo{}
repo.accountsByID = map[int64]*Account{account.ID: account}
cache := &grokTokenCacheForProviderTest{lockResult: true}
provider := NewGrokTokenProvider(repo, cache)
provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{
err: errors.New("invalid_client: client credentials rejected"),
})
token, err := provider.GetAccessTokenForManualTest(context.Background(), account)
require.Error(t, err)
require.Empty(t, token)
require.Contains(t, err.Error(), "invalid_client")
}
func TestGrokTokenProviderManualTestLockHeldWithExpiredTokenReturnsSpecificError(t *testing.T) {
account := expiredGrokOAuthAccountForCredentialTest(133)
repo := &tokenRefreshAccountRepo{}
repo.accountsByID = map[int64]*Account{account.ID: account}
cache := &grokTokenCacheForProviderTest{lockResult: false}
provider := NewGrokTokenProvider(repo, cache)
provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &tokenRefresherStub{})
token, err := provider.GetAccessTokenForManualTest(context.Background(), account)
require.Error(t, err)
require.Empty(t, token)
require.Contains(t, err.Error(), "refresh is already in progress")
}
func TestGrokTokenProviderManualTestRequiresRefreshToken(t *testing.T) {
account := expiredGrokOAuthAccountForCredentialTest(134)
delete(account.Credentials, "refresh_token")
provider := NewGrokTokenProvider(&tokenRefreshAccountRepo{}, &grokTokenCacheForProviderTest{})
token, err := provider.GetAccessTokenForManualTest(context.Background(), account)
require.ErrorIs(t, err, errGrokOAuthRefreshTokenMissing)
require.Empty(t, token)
}
func TestGrokTokenProviderRejectsIneligibleSelectedAccountBeforeWarmCache(t *testing.T) {
future := time.Now().Add(time.Hour)
tests := []struct {