mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-07 15:57:53 +08:00
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:
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user