diff --git a/README.md b/README.md index 53a33780a..427e68e71 100644 --- a/README.md +++ b/README.md @@ -24,21 +24,24 @@ > Check out [DEEIX-AI / DEEIX-Chat](https://github.com/DEEIX-AI/DEEIX-Chat), a lightweight, integrated AI platform for model routing, chat, files, tools, billing, identity, and operations. > [!NOTE] -> This project is intended for technical research and self-hosted use. Make sure your accounts, network setup, and upstream usage comply with the applicable terms and local requirements. +> This project is for technical research and learning purposes only. Please comply with Grok's official terms of use and local laws when using it; otherwise, you will be solely responsible for all consequences! -## ❤️ Sponsors +## Sponsors +> [Want to sponsor this project?](mailto:chenyme03@gmail.com) - + - +
DEEIX AI / DEEIX ChatDEEIX-Chat is an open-source, deployable AI platform for individuals, teams, and organizations that need stable, long-term access to multiple models. It brings multimodal chat, model routing, files and RAG, MCP tools, usage billing, authentication, audit logs, and operational controls into one product.DEEIX-Chat is an open-source, self-hostable AI Chat platform for individuals, teams, and enterprises that need stable, long-term, unified access to multiple models. It brings models, conversations, files, tool calling, and administration together in one deployable and extensible system. Click here to start deploying.
RightCodeRight Code provides stable access services for Claude Code, Codex, Gemini, and other models. It supports invoicing and one-to-one assistance for businesses and teams. Part of the model capacity used to develop this project is provided by Right Code. Thanks to the Right Code team for supporting the project. Register here to get started.Right Code is an enterprise-grade AI Agent distribution platform that primarily provides stable access services for Claude Code, Codex, Gemini, and other models. It supports invoicing and dedicated one-to-one assistance for enterprises and teams. Thanks to Right Code for providing token support. Click here to register and get started.
+
+ Grok2API is a Go-based Grok API gateway with a built-in React admin console. It organizes Grok Build OAuth, Grok Web SSO, and Grok Console SSO credentials into independent account pools, exposes OpenAI- and Anthropic-style APIs, and provides one place to manage model routes, client keys, quotas, media, audits, and egress proxies. ## Highlights diff --git a/README.zh-CN.md b/README.zh-CN.md index 48486776b..2c03d1ca3 100644 --- a/README.zh-CN.md +++ b/README.zh-CN.md @@ -24,9 +24,11 @@ > 推荐个人新项目 [DEEIX-AI / DEEIX-Chat](https://github.com/DEEIX-AI/DEEIX-Chat)：面向多模型路由、对话、文件、工具、计费与运维的一体化轻量 AI 平台。 > [!NOTE] -> 本项目用于技术研究与个人部署。使用前请确认账号、网络和上游服务符合相应条款及所在地要求。 +> 本项目仅供技术研究与学习交流。使用时请务必遵循 Grok 官方的使用条款及当地法律法规，否则一切后果自负！ -## ❤️ 赞助商 +## 赞助商 + +> [希望赞助这个项目？](mailto:chenyme03@gmail.com) @@ -35,10 +37,12 @@ - +
RightCodeRight Code 是一个企业级 AI Agent 分发平台，主要提供稳定的 Claude Code、Codex、Gemini 等模型的中转服务。充值即可开票，企业、团队用户一对一对接。感谢 RightCode 提供的 Tokens 支持，点击 此处 注册并开始使用！Right Code 是一个企业级 AI Agent 分发平台，主要提供稳定的 Claude Code、Codex、Gemini 等模型的中转服务。充值即可开票，企业、团队用户一对一对接。感谢 Right Code 提供的 Tokens 支持，点击 此处 注册并开始使用！
+
+ Grok2API 是一个以 Go 为核心、内置 React 管理端的 Grok API 网关。它将 Grok Build OAuth、Grok Web SSO 与 Grok Console SSO 组织成相互独立的账号池，对外提供 OpenAI 与 Anthropic 风格接口，并统一管理模型路由、客户端密钥、额度、媒体、审计和出口代理。 ## 功能概览 @@ -377,5 +381,6 @@ make swagger ## 进一步阅读 +- [English README](./README.md) - [后端说明](./backend/README.md) - [前端说明](./frontend/README.md) diff --git a/backend/internal/application/account/provider_links.go b/backend/internal/application/account/provider_links.go index 2118da27e..5108bc04d 100644 --- a/backend/internal/application/account/provider_links.go +++ b/backend/internal/application/account/provider_links.go @@ -2,10 +2,12 @@ package account import ( "context" + "errors" "fmt" "strings" accountdomain "github.com/chenyme/grok2api/backend/internal/domain/account" + "github.com/chenyme/grok2api/backend/internal/infra/provider" ) type providerLinkRepository interface { @@ -14,7 +16,7 @@ type providerLinkRepository interface { } // SyncAccountIdentity 尽力补充 Web/Console 的稳定上游身份，并据此建立高可信弱关联。 -// 该操作不修改账号凭据、健康、额度或路由状态。 +// 只有明确的 401 会将当前 Provider 账号移出号池；其他同步失败不影响健康状态。 func (s *Service) SyncAccountIdentity(ctx context.Context, id uint64) error { _, err, _ := s.identitySyncs.Do(fmt.Sprintf("%d", id), func() (any, error) { return nil, s.syncAccountIdentity(ctx, id) @@ -48,6 +50,10 @@ func (s *Service) syncAccountIdentity(ctx context.Context, id uint64) error { } identity, err := adapter.SyncAccountIdentity(ctx, value) if err != nil { + if errors.Is(err, provider.ErrUnauthorized) { + markErr := s.markSSOCredentialRejected(ctx, value, fmt.Sprintf("%s SSO credential rejected", value.Provider)) + return errors.Join(err, markErr) + } return err } if len(identity.Email) > 255 || len(identity.UserID) > 255 || len(identity.TeamID) > 255 { @@ -69,8 +75,10 @@ func (s *Service) reconcileProviderLinksBestEffort(ctx context.Context, id uint6 } } -func (s *Service) syncAccountIdentityBestEffort(ctx context.Context, id uint64) { +func (s *Service) syncAccountIdentityBestEffort(ctx context.Context, id uint64) error { if err := s.SyncAccountIdentity(ctx, id); err != nil { s.logger.Warn("account_identity_sync_failed", "account_id", id, "error", err) + return err } + return nil } diff --git a/backend/internal/application/account/provider_links_test.go b/backend/internal/application/account/provider_links_test.go index ae94b27b7..9f3f71539 100644 --- a/backend/internal/application/account/provider_links_test.go +++ b/backend/internal/application/account/provider_links_test.go @@ -63,7 +63,7 @@ func TestSyncAccountIdentityLinksUniqueBuildWithoutSharingState(t *testing.T) { } } -func TestSyncAccountIdentityFailureDoesNotInvalidateAccount(t *testing.T) { +func TestSyncAccountIdentityUnauthorizedInvalidatesCurrentProviderAccount(t *testing.T) { t.Parallel() ctx := context.Background() service, repo, adapter := newWebAccountSettingsTestService(t) @@ -82,8 +82,8 @@ func TestSyncAccountIdentityFailureDoesNotInvalidateAccount(t *testing.T) { if err != nil { t.Fatal(err) } - if web.AuthStatus != accountdomain.AuthStatusActive || !web.Enabled || web.FailureCount != 0 { - t.Fatalf("identity failure changed account state: %#v", web) + if web.AuthStatus != accountdomain.AuthStatusReauthRequired || !web.Enabled || web.FailureCount != 0 { + t.Fatalf("identity unauthorized state = %#v", web) } } diff --git a/backend/internal/application/account/quota_refresh_test.go b/backend/internal/application/account/quota_refresh_test.go index c7e4119de..efd5d6824 100644 --- a/backend/internal/application/account/quota_refresh_test.go +++ b/backend/internal/application/account/quota_refresh_test.go @@ -2,6 +2,7 @@ package account import ( "context" + "errors" "path/filepath" "sync/atomic" "testing" @@ -132,6 +133,39 @@ func TestRefreshQuotaFetchesWebIdentityOnlyUntilDataExists(t *testing.T) { } } +func TestRefreshQuotaUnauthorizedMarksWebAccountInvalid(t *testing.T) { + ctx := context.Background() + database, err := relational.OpenSQLite(ctx, filepath.Join(t.TempDir(), "quota-unauthorized.db")) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = database.Close() }) + if err := database.InitializeSchema(ctx); err != nil { + t.Fatal(err) + } + accounts := relational.NewAccountRepository(database) + credential, _, err := accounts.UpsertByIdentity(ctx, accountdomain.Credential{ + Provider: accountdomain.ProviderWeb, AuthType: accountdomain.AuthTypeSSO, + Name: "web-unauthorized", SourceKey: "web-unauthorized", EncryptedAccessToken: "encrypted", + Enabled: true, AuthStatus: accountdomain.AuthStatusActive, + }) + if err != nil { + t.Fatal(err) + } + adapter := "aCountingAdapter{fullErr: provider.ErrUnauthorized} + service := NewService(accounts, nil, nil, nil, provider.NewRegistry(adapter), nil, nil) + if _, err := service.RefreshQuota(ctx, credential.ID); !errors.Is(err, provider.ErrUnauthorized) { + t.Fatalf("err = %v", err) + } + stored, err := accounts.Get(ctx, credential.ID) + if err != nil { + t.Fatal(err) + } + if stored.AuthStatus != accountdomain.AuthStatusReauthRequired || !stored.Enabled { + t.Fatalf("account state = %#v", stored) + } +} + type deniedQuotaRefreshLock struct{} func (deniedQuotaRefreshLock) Acquire(context.Context, string, time.Duration) (func(), bool, error) { @@ -142,6 +176,7 @@ type quotaCountingAdapter struct { modeCalls atomic.Int64 fullCalls atomic.Int64 identityCalls atomic.Int64 + fullErr error } func (a *quotaCountingAdapter) Provider() accountdomain.Provider { return accountdomain.ProviderWeb } @@ -155,7 +190,7 @@ func (a *quotaCountingAdapter) Definition() provider.Definition { func (a *quotaCountingAdapter) SyncQuota(context.Context, accountdomain.Credential) (provider.QuotaSnapshot, error) { a.fullCalls.Add(1) - return provider.QuotaSnapshot{}, nil + return provider.QuotaSnapshot{}, a.fullErr } func (a *quotaCountingAdapter) SyncAccountIdentity(context.Context, accountdomain.Credential) (provider.AccountIdentity, error) { diff --git a/backend/internal/application/account/service.go b/backend/internal/application/account/service.go index d6d3ff597..c22b8f86e 100644 --- a/backend/internal/application/account/service.go +++ b/backend/internal/application/account/service.go @@ -45,6 +45,7 @@ const ( credentialRefreshSafetyPoll time.Duration = time.Minute credentialRefreshTimeout time.Duration = 30 * time.Second credentialRefreshStateTTL time.Duration = 5 * time.Second + credentialStateWriteTimeout time.Duration = 5 * time.Second credentialRefreshBatchSize = 100 managedTaskWorkerCeiling = 50 webQuotaRefreshQueueSize = 4096 @@ -1254,7 +1255,7 @@ func (s *Service) convertWebAccountToBuild(ctx context.Context, id uint64, strat seed, err := converter.ConvertToBuild(ctx, value) if err != nil { if errors.Is(err, provider.ErrUnauthorized) { - _ = s.MarkReauthRequired(context.WithoutCancel(ctx), id, "Grok Web SSO credential rejected") + err = errors.Join(err, s.markSSOCredentialRejected(ctx, value, "Grok Web SSO credential rejected")) } return 0, false, false, err } @@ -1433,6 +1434,21 @@ func (s *Service) MarkReauthRequired(ctx context.Context, id uint64, reason stri return nil } +// markSSOCredentialRejected 在上游明确返回 401 后可靠持久化失效状态。 +// 状态写入不继承客户端取消，避免已经确认失效的账号因请求断开继续留在号池。 +func (s *Service) markSSOCredentialRejected(ctx context.Context, value accountdomain.Credential, reason string) error { + if value.AuthType != accountdomain.AuthTypeSSO { + return nil + } + writeCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), credentialStateWriteTimeout) + defer cancel() + if err := s.MarkReauthRequired(writeCtx, value.ID, reason); err != nil { + s.logger.Error("account_reauth_required_write_failed", "account_id", value.ID, "provider", value.Provider, "error", err) + return err + } + return nil +} + // EnsureCredential 在即将过期时刷新 token，同一账号并发请求只执行一次刷新。 func (s *Service) EnsureCredential(ctx context.Context, value accountdomain.Credential, force bool) (accountdomain.Credential, error) { return s.ensureCredential(ctx, value, force, false, false) @@ -1878,7 +1894,7 @@ func (s *Service) refreshQuota(ctx context.Context, id uint64) ([]accountdomain. snapshot, err := adapter.SyncQuota(ctx, value) if err != nil { if errors.Is(err, provider.ErrUnauthorized) { - _ = s.MarkReauthRequired(ctx, id, fmt.Sprintf("%s SSO credential rejected", value.Provider)) + err = errors.Join(err, s.markSSOCredentialRejected(ctx, value, fmt.Sprintf("%s SSO credential rejected", value.Provider))) } return nil, err } @@ -1904,7 +1920,9 @@ func (s *Service) refreshQuota(ctx context.Context, id uint64) ([]accountdomain. // 并沿用调用方取消语义，不能反向影响额度同步结果。 if (value.Provider == accountdomain.ProviderWeb || value.Provider == accountdomain.ProviderConsole) && ctx.Err() == nil { if strings.TrimSpace(value.UserID) == "" && strings.TrimSpace(value.Email) == "" { - s.syncAccountIdentityBestEffort(ctx, id) + if identityErr := s.syncAccountIdentityBestEffort(ctx, id); errors.Is(identityErr, provider.ErrUnauthorized) { + return snapshot.Windows, identityErr + } } else { // 已有 Session 身份时只做本地增量关联，不再访问上游。 s.reconcileProviderLinksBestEffort(ctx, id) @@ -1984,7 +2002,7 @@ func (s *Service) refreshQuotaMode(ctx context.Context, id uint64, mode string) window, err := adapter.SyncQuotaMode(ctx, value, mode) if err != nil { if errors.Is(err, provider.ErrUnauthorized) { - _ = s.MarkReauthRequired(ctx, id, fmt.Sprintf("%s SSO credential rejected", value.Provider)) + err = errors.Join(err, s.markSSOCredentialRejected(ctx, value, fmt.Sprintf("%s SSO credential rejected", value.Provider))) } return accountdomain.QuotaWindow{}, err } diff --git a/backend/internal/application/account/web_account_settings.go b/backend/internal/application/account/web_account_settings.go index 9cd1917e2..117e1c429 100644 --- a/backend/internal/application/account/web_account_settings.go +++ b/backend/internal/application/account/web_account_settings.go @@ -192,7 +192,7 @@ func (s *Service) runWebAccountSetting(ctx context.Context, credential accountdo return nil } if errors.Is(err, provider.ErrUnauthorized) { - _ = s.MarkReauthRequired(context.WithoutCancel(ctx), credential.ID, "Grok Web SSO credential rejected") + err = errors.Join(err, s.markSSOCredentialRejected(ctx, credential, "Grok Web SSO credential rejected")) } return fmt.Errorf("%s: %w", operation, err) } diff --git a/backend/internal/application/accountsync/service.go b/backend/internal/application/accountsync/service.go index 0afe07d31..ee82fee83 100644 --- a/backend/internal/application/accountsync/service.go +++ b/backend/internal/application/accountsync/service.go @@ -211,10 +211,14 @@ func (s *Service) syncAccount(ctx context.Context, accountID uint64) error { if view.Credential.Provider == accountdomain.ProviderWeb || view.Credential.Provider == accountdomain.ProviderConsole { if identity, ok := s.accounts.(identitySynchronizer); ok { operationCtx, cancel := context.WithTimeout(ctx, operationTimeout) - if err := identity.SyncAccountIdentity(operationCtx, accountID); err != nil { - s.logger.Warn("account_initial_identity_sync_failed", "account_id", accountID, "error", err) + identityErr := identity.SyncAccountIdentity(operationCtx, accountID) + if identityErr != nil { + s.logger.Warn("account_initial_identity_sync_failed", "account_id", accountID, "error", identityErr) } cancel() + if errors.Is(identityErr, provider.ErrUnauthorized) { + return fmt.Errorf("同步账号身份: %w", identityErr) + } } } if definition.Quota == provider.QuotaRemoteWindow || definition.Quota == provider.QuotaLocalWindow { diff --git a/backend/internal/application/accountsync/service_test.go b/backend/internal/application/accountsync/service_test.go index 8076ff01f..1b9044e02 100644 --- a/backend/internal/application/accountsync/service_test.go +++ b/backend/internal/application/accountsync/service_test.go @@ -230,6 +230,24 @@ func TestSyncAccountIgnoresBestEffortSSOIdentityFailure(t *testing.T) { } } +func TestSyncAccountStopsAfterSSOIdentityUnauthorized(t *testing.T) { + for _, providerValue := range []accountdomain.Provider{accountdomain.ProviderWeb, accountdomain.ProviderConsole} { + t.Run(string(providerValue), func(t *testing.T) { + reader := &identityAccountReaderStub{accountReaderStub: accountReaderStub{provider: providerValue}, err: provider.ErrUnauthorized} + quota := "aStub{} + models := &modelStub{hasSnapshot: true} + service := NewService(slog.Default(), reader, &billingStub{}, quota, models) + + if err := service.syncAccount(context.Background(), 10); !errors.Is(err, provider.ErrUnauthorized) { + t.Fatalf("err = %v", err) + } + if reader.calls != 1 || quota.syncs != 0 { + t.Fatalf("identity calls=%d quota syncs=%d", reader.calls, quota.syncs) + } + }) + } +} + func TestSyncAccountUsesDeclaredQuotaPolicyInsteadOfProviderName(t *testing.T) { billing := &billingStub{} quota := "aStub{} diff --git a/backend/internal/application/gateway/image.go b/backend/internal/application/gateway/image.go index a68778fa8..cba3b58bc 100644 --- a/backend/internal/application/gateway/image.go +++ b/backend/internal/application/gateway/image.go @@ -191,6 +191,14 @@ func (s *Service) executeImage( response, err = execute(ctx, route.Provider, credential, route.UpstreamModel) if err != nil { s.logger.Error("image_upstream_failed", "event_id", eventID, "request_id", requestID, "model", externalModel, "provider", route.Provider, "account_id", credential.ID, "error", err) + if isSSOCredentialRejected(err, credential) { + s.markSSOCredentialRejected(ctx, credential, fmt.Sprintf("%s SSO credential rejected", credential.Provider)) + failedCredential := credential + lastCredentialFailure = &failedCredential + lastCredentialError = provider.ErrUnauthorized + lease.Release() + continue + } if !provider.IsMediaPostProcessingError(err) { s.selector.MarkFailure(ctx, credential, 0, 0) } @@ -202,6 +210,16 @@ func (s *Service) executeImage( writeFailureAudit(http.StatusBadGateway, errorCode, &credential) return nil, err } + if response.StatusCode == http.StatusUnauthorized && credential.AuthType == accountdomain.AuthTypeSSO { + _, _ = readRetryableBody(response.Body) + s.markSSOCredentialRejected(ctx, credential, fmt.Sprintf("%s SSO credential rejected", credential.Provider)) + failedCredential := credential + lastCredentialFailure = &failedCredential + lastCredentialError = provider.ErrUnauthorized + response = nil + lease.Release() + continue + } if s.providers.RetryForbiddenAsEgress(credential.Provider) && response.StatusCode == http.StatusForbidden && attempt == 0 && attempt+1 < attempts { _, _ = readRetryableBody(response.Body) lease.Release() @@ -230,10 +248,6 @@ func (s *Service) executeImage( } return nil, fmt.Errorf("%w: %w", ErrNoAvailableAccount, lastCredentialError) } - if response.StatusCode == http.StatusUnauthorized && credential.AuthType == accountdomain.AuthTypeSSO { - _ = s.accounts.MarkReauthRequired(ctx, credential.ID, fmt.Sprintf("%s SSO credential rejected", credential.Provider)) - s.selector.MarkFailure(ctx, credential, http.StatusUnauthorized, 0) - } effectiveQuotaMode := lease.QuotaMode accountID := credential.ID var once sync.Once diff --git a/backend/internal/application/gateway/service.go b/backend/internal/application/gateway/service.go index 2f3849117..10a0ae659 100644 --- a/backend/internal/application/gateway/service.go +++ b/backend/internal/application/gateway/service.go @@ -551,6 +551,11 @@ attemptLoop: lastFailure = &UpstreamFailure{HTTPStatus: 499, Code: "request_canceled", PublicMessage: "请求已取消", AccountID: credential.ID, AccountName: credential.Name, Cause: firstError(ctx.Err(), err)} break } + if isSSOCredentialRejected(err, credential) { + s.markSSOCredentialRejected(ctx, credential, fmt.Sprintf("%s SSO credential rejected", credential.Provider)) + lastFailure = newHTTPUpstreamFailure(http.StatusUnauthorized, nil, credential.ID, credential.Name) + continue + } lastFailure = newTransportUpstreamFailure(err, credential.ID, credential.Name) failureFingerprints[lastFailure.Fingerprint]++ if failureFingerprints[lastFailure.Fingerprint] >= 2 { @@ -565,8 +570,7 @@ attemptLoop: if response.StatusCode == http.StatusUnauthorized { response.Body.Close() if credential.AuthType == accountdomain.AuthTypeSSO { - _ = s.accounts.MarkReauthRequired(ctx, credential.ID, fmt.Sprintf("%s SSO credential rejected", credential.Provider)) - s.selector.MarkFailure(ctx, credential, http.StatusUnauthorized, 0) + s.markSSOCredentialRejected(ctx, credential, fmt.Sprintf("%s SSO credential rejected", credential.Provider)) lease.Release() lastErr = fmt.Errorf("%s SSO 凭据已失效", credential.Provider) lastFailure = newHTTPUpstreamFailure(http.StatusUnauthorized, nil, credential.ID, credential.Name) @@ -830,6 +834,30 @@ attemptLoop: return nil, fmt.Errorf("%w: %w", ErrNoAvailableAccount, lastErr) } +func isSSOCredentialRejected(err error, credential accountdomain.Credential) bool { + if credential.AuthType != accountdomain.AuthTypeSSO || err == nil { + return false + } + if errors.Is(err, provider.ErrUnauthorized) { + return true + } + status, ok := provider.ErrorHTTPStatus(err) + return ok && status == http.StatusUnauthorized +} + +func (s *Service) markSSOCredentialRejected(ctx context.Context, credential accountdomain.Credential, reason string) { + if credential.AuthType != accountdomain.AuthTypeSSO { + return + } + writeCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), finalizationTimeout) + defer cancel() + if err := s.accounts.MarkReauthRequired(writeCtx, credential.ID, reason); err != nil { + s.logger.Error("account_reauth_required_write_failed", "account_id", credential.ID, "provider", credential.Provider, "error", err) + } + // 即使持久化失败也立即丢弃本进程的一秒候选快照，避免当前失效账号被下一请求再次选中。 + s.selector.MarkQuotaStateChanged(credential.Provider) +} + func (s *Service) queueAccountModelSync(accountID uint64) { syncer, ok := s.models.(accountModelSyncer) if !ok || accountID == 0 { @@ -954,11 +982,19 @@ func (s *Service) forwardOwnedResponse(ctx context.Context, input ResourceInput, } response, err := adapter.ForwardResponse(ctx, provider.ResponseResourceRequest{Credential: credential, Method: method, Path: path}) if err != nil { + if isSSOCredentialRejected(err, credential) { + s.markSSOCredentialRejected(ctx, credential, fmt.Sprintf("%s SSO credential rejected", credential.Provider)) + } lease.Release() return nil, err } if response.StatusCode == http.StatusUnauthorized { response.Body.Close() + if credential.AuthType == accountdomain.AuthTypeSSO { + s.markSSOCredentialRejected(ctx, credential, fmt.Sprintf("%s SSO credential rejected", credential.Provider)) + lease.Release() + return nil, ErrResponseAccountUnavailable + } if s.markPermanentlyUnrefreshableCredentialRejected(ctx, credential) { lease.Release() return nil, fmt.Errorf("%w: %w", ErrResponseAccountUnavailable, accountapp.ErrCredentialRefreshPermanent) diff --git a/backend/internal/application/gateway/service_test.go b/backend/internal/application/gateway/service_test.go index e8f72b966..a6375342f 100644 --- a/backend/internal/application/gateway/service_test.go +++ b/backend/internal/application/gateway/service_test.go @@ -212,6 +212,93 @@ func TestGatewayFailsOverBeforeReturningBody(t *testing.T) { } } +func TestGatewaySSOUnauthorizedMarksInvalidAndSwitchesAccount(t *testing.T) { + for _, providerValue := range []account.Provider{account.ProviderWeb, account.ProviderConsole} { + providerValue := providerValue + t.Run(string(providerValue), func(t *testing.T) { + ctx := context.Background() + database, err := relational.OpenSQLite(ctx, filepath.Join(t.TempDir(), "sso-401.db")) + if err != nil { + t.Fatal(err) + } + defer database.Close() + if err := database.InitializeSchema(ctx); err != nil { + t.Fatal(err) + } + accountRepo := relational.NewAccountRepository(database) + modelRepo := relational.NewModelRepository(database) + auditRepo := relational.NewAuditRepository(database) + responseRepo := relational.NewResponseRepository(database) + keyRepo := relational.NewClientKeyRepository(database) + credentials := make([]account.Credential, 0, 2) + for index, name := range []string{"rejected", "healthy"} { + credential, _, createErr := accountRepo.UpsertByIdentity(ctx, account.Credential{ + Provider: providerValue, AuthType: account.AuthTypeSSO, Name: name, SourceKey: string(providerValue) + "-" + name, + EncryptedAccessToken: "encrypted-" + name, Enabled: true, AuthStatus: account.AuthStatusActive, + Priority: 200 - index*100, MaxConcurrent: 1, + }) + if createErr != nil { + t.Fatal(createErr) + } + credentials = append(credentials, credential) + } + modelName := "grok-sso-401" + if err := modelRepo.UpsertDiscovered(ctx, providerValue, []string{modelName}); err != nil { + t.Fatal(err) + } + for _, credential := range credentials { + if err := modelRepo.ReplaceAccountCapabilities(ctx, credential.ID, []string{modelName}, time.Now().UTC()); err != nil { + t.Fatal(err) + } + } + key, err := keyRepo.Create(ctx, clientkey.Key{ + Name: "sso-401-key", Prefix: "sso-401", SecretHash: strings.Repeat("e", 64), EncryptedSecret: "encrypted-key", + Enabled: true, RPMLimit: 60, MaxConcurrent: 4, + }) + if err != nil { + t.Fatal(err) + } + adapter := &ssoUnauthorizedAdapter{providerValue: providerValue, rejectedID: credentials[0].ID} + registry := provider.NewRegistry(adapter) + sticky := memory.NewStickyStore() + accountService := accountapp.NewService(accountRepo, auditRepo, memory.NewDeviceSessionStore(), sticky, registry, testCipher(t), nil) + selector := NewSelector(accountRepo, memory.NewConcurrencyLimiter(), sticky, registry, time.Hour, time.Second, time.Minute) + service := NewService(modelRepo, auditRepo, accountService, clientkeyapp.NewService(nil, nil, nil, 60, 4, nil), registry, selector, responseRepo, 2) + + result, err := service.CreateResponse(ctx, Input{ + RequestID: "req-sso-401", ClientKey: key, PublicModel: modelName, + Body: []byte(`{"model":"grok-sso-401","input":"hello"}`), + }) + if err != nil { + t.Fatal(err) + } + body, err := io.ReadAll(result.Body) + if err != nil { + t.Fatal(err) + } + result.Finalize(Usage{}, "", "") + _ = result.Body.Close() + if string(body) != "ok" { + t.Fatalf("body = %q", body) + } + if attempts := adapter.Attempts(); len(attempts) != 2 || attempts[0] != credentials[0].ID || attempts[1] != credentials[1].ID { + t.Fatalf("attempts = %#v", attempts) + } + rejected, err := accountRepo.Get(ctx, credentials[0].ID) + if err != nil { + t.Fatal(err) + } + if rejected.AuthStatus != account.AuthStatusReauthRequired || !rejected.Enabled { + t.Fatalf("rejected account = %#v", rejected) + } + healthy, err := accountRepo.Get(ctx, credentials[1].ID) + if err != nil || healthy.AuthStatus != account.AuthStatusActive { + t.Fatalf("healthy account = %#v, err = %v", healthy, err) + } + }) + } +} + func TestGatewayTeamModelRateLimitOnlySkipsMatchingTeam(t *testing.T) { ctx := context.Background() database, err := relational.OpenSQLite(ctx, filepath.Join(t.TempDir(), "team-model-rate-limit.db")) @@ -982,6 +1069,77 @@ func TestImageStreamPropagatesWithoutTouchingChatQuota(t *testing.T) { } } +func TestWebImageUnauthorizedMarksInvalidAndSwitchesAccount(t *testing.T) { + ctx := context.Background() + database, err := relational.OpenSQLite(ctx, filepath.Join(t.TempDir(), "web-image-401.db")) + if err != nil { + t.Fatal(err) + } + defer database.Close() + if err := database.InitializeSchema(ctx); err != nil { + t.Fatal(err) + } + accountRepo := relational.NewAccountRepository(database) + modelRepo := relational.NewModelRepository(database) + auditRepo := relational.NewAuditRepository(database) + responseRepo := relational.NewResponseRepository(database) + keyRepo := relational.NewClientKeyRepository(database) + credentials := make([]account.Credential, 0, 2) + for index, name := range []string{"rejected-image", "healthy-image"} { + credential, _, createErr := accountRepo.UpsertByIdentity(ctx, account.Credential{ + Provider: account.ProviderWeb, AuthType: account.AuthTypeSSO, WebTier: account.WebTierSuper, + Name: name, SourceKey: name, EncryptedAccessToken: "encrypted-" + name, + Enabled: true, AuthStatus: account.AuthStatusActive, Priority: 200 - index*100, MaxConcurrent: 1, + }) + if createErr != nil { + t.Fatal(createErr) + } + credentials = append(credentials, credential) + } + now := time.Now().UTC() + if err := modelRepo.UpsertRoutes(ctx, []modeldomain.Route{{ + PublicID: "grok-image-401", Provider: account.ProviderWeb, UpstreamModel: "grok-image-401", + Capability: modeldomain.CapabilityImage, Enabled: true, + }}); err != nil { + t.Fatal(err) + } + for _, credential := range credentials { + if err := modelRepo.ReplaceAccountCapabilities(ctx, credential.ID, []string{"grok-image-401"}, now); err != nil { + t.Fatal(err) + } + } + key, err := keyRepo.Create(ctx, clientkey.Key{ + Name: "image-401-key", Prefix: "image-401", SecretHash: strings.Repeat("f", 64), EncryptedSecret: "encrypted-key", + Enabled: true, RPMLimit: 60, MaxConcurrent: 4, + }) + if err != nil { + t.Fatal(err) + } + adapter := &webImageStreamAdapter{synced: make(chan string, 1), unauthorizedID: credentials[0].ID} + registry := provider.NewRegistry(adapter) + sticky := memory.NewStickyStore() + accountService := accountapp.NewService(accountRepo, auditRepo, memory.NewDeviceSessionStore(), sticky, registry, testCipher(t), nil) + selector := NewSelector(accountRepo, memory.NewConcurrencyLimiter(), sticky, registry, time.Hour, time.Second, time.Minute) + service := NewService(modelRepo, auditRepo, accountService, clientkeyapp.NewService(keyRepo, nil, nil, 60, 4, nil), registry, selector, responseRepo, 2) + + result, err := service.GenerateImage(ctx, ImageGenerationInput{ + RequestID: "req-image-401", ClientKey: key, PublicModel: "grok-image-401", Prompt: "test", Count: 1, + }) + if err != nil { + t.Fatal(err) + } + _, _ = io.ReadAll(result.Body) + result.Finalize(Usage{}, "", "") + _ = result.Body.Close() + if attempts := adapter.Attempts(); len(attempts) != 2 || attempts[0] != credentials[0].ID || attempts[1] != credentials[1].ID { + t.Fatalf("attempts = %#v", attempts) + } + rejected, err := accountRepo.Get(ctx, credentials[0].ID) + if err != nil || rejected.AuthStatus != account.AuthStatusReauthRequired || !rejected.Enabled { + t.Fatalf("rejected account = %#v, err = %v", rejected, err) + } +} + func TestSuccessfulWebChatRefreshesCurrentModeQuota(t *testing.T) { ctx := context.Background() database, err := relational.OpenSQLite(ctx, filepath.Join(t.TempDir(), "chat-quota-refresh.db")) @@ -1101,6 +1259,35 @@ type failoverAdapter struct { resourceStatus int } +type ssoUnauthorizedAdapter struct { + mu sync.Mutex + providerValue account.Provider + rejectedID uint64 + attempts []uint64 +} + +func (a *ssoUnauthorizedAdapter) Provider() account.Provider { return a.providerValue } +func (a *ssoUnauthorizedAdapter) Definition() provider.Definition { + return testConversationDefinition(a.providerValue) +} +func (a *ssoUnauthorizedAdapter) ForwardResponse(_ context.Context, request provider.ResponseResourceRequest) (*provider.Response, error) { + a.mu.Lock() + a.attempts = append(a.attempts, request.Credential.ID) + a.mu.Unlock() + if request.Credential.ID == a.rejectedID { + return &provider.Response{ + StatusCode: http.StatusUnauthorized, Status: "401 Unauthorized", Header: make(http.Header), + Body: io.NopCloser(strings.NewReader(`{"error":{"code":"unauthorized","message":"credential rejected"}}`)), + }, nil + } + return &provider.Response{StatusCode: http.StatusOK, Status: "200 OK", Header: make(http.Header), Body: io.NopCloser(strings.NewReader("ok"))}, nil +} +func (a *ssoUnauthorizedAdapter) Attempts() []uint64 { + a.mu.Lock() + defer a.mu.Unlock() + return append([]uint64(nil), a.attempts...) +} + type statelessConsoleAdapter struct{} type teamModelRateLimitConsoleAttempt struct { @@ -1226,13 +1413,14 @@ func (a *systemicForbiddenAdapter) Attempts() []uint64 { type webRateLimitAdapter struct{} type webImageStreamAdapter struct { - mu sync.Mutex - streaming bool - partialImages int - editRequest provider.ImageEditRequest - synced chan string - failureEgress *infraegress.Manager - attempts []uint64 + mu sync.Mutex + streaming bool + partialImages int + editRequest provider.ImageEditRequest + synced chan string + failureEgress *infraegress.Manager + attempts []uint64 + unauthorizedID uint64 } type webChatQuotaAdapter struct { @@ -1295,7 +1483,14 @@ func (a *webImageStreamAdapter) GenerateImage(ctx context.Context, request provi a.partialImages = request.PartialImages failureEgress := a.failureEgress a.attempts = append(a.attempts, request.Credential.ID) + unauthorizedID := a.unauthorizedID a.mu.Unlock() + if request.Credential.ID == unauthorizedID { + return &provider.Response{ + StatusCode: http.StatusUnauthorized, Status: "401 Unauthorized", Header: make(http.Header), + Body: io.NopCloser(strings.NewReader(`{"error":{"code":"unauthorized"}}`)), + }, nil + } if failureEgress != nil { lease, err := failureEgress.Acquire(ctx, egressdomain.ScopeWeb, "image-failure") if err != nil { diff --git a/backend/internal/application/gateway/video.go b/backend/internal/application/gateway/video.go index aecb98610..7d66cd667 100644 --- a/backend/internal/application/gateway/video.go +++ b/backend/internal/application/gateway/video.go @@ -328,12 +328,14 @@ func (s *Service) runVideoJob(parent context.Context, job media.Job, route model failureHandled := false if errors.Is(err, provider.ErrUnauthorized) { if lease.Credential.AuthType == account.AuthTypeSSO { - _ = s.accounts.MarkReauthRequired(failureCtx, lease.Credential.ID, fmt.Sprintf("%s SSO credential rejected", lease.Credential.Provider)) + s.markSSOCredentialRejected(failureCtx, lease.Credential, fmt.Sprintf("%s SSO credential rejected", lease.Credential.Provider)) } - s.selector.MarkFailure(failureCtx, lease.Credential, http.StatusUnauthorized, 0) failureHandled = true } else if status, ok := provider.ErrorHTTPStatus(err); ok { switch { + case status == http.StatusUnauthorized && lease.Credential.AuthType == account.AuthTypeSSO: + s.markSSOCredentialRejected(failureCtx, lease.Credential, fmt.Sprintf("%s SSO credential rejected", lease.Credential.Provider)) + failureHandled = true case status == http.StatusForbidden && s.providers.RetryForbiddenAsEgress(lease.Credential.Provider): // Web Provider 已对 anti-bot 403 降低出口健康并重建浏览器会话； // 视频请求已提交，不能换号重试，也不能误伤账号池。