From f479a5d10cb4d61379a1174b809948b5e911065d Mon Sep 17 00:00:00 2001 From: cat Date: Tue, 14 Jul 2026 18:09:59 +0800 Subject: [PATCH] =?UTF-8?q?fix(openai):=20=E6=94=B6=E7=B4=A7=20Agent=20Ide?= =?UTF-8?q?ntity=20task=20=E7=94=9F=E5=91=BD=E5=91=A8=E6=9C=9F?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/cmd/server/wire_gen.go | 22 +++---- .../internal/service/account_test_service.go | 11 ++-- .../internal/service/account_usage_service.go | 3 +- .../internal/service/openai_agent_identity.go | 20 +++--- .../openai_agent_identity_compat_test.go | 12 +++- .../service/openai_gateway_service.go | 6 ++ .../internal/service/openai_quota_service.go | 21 +++---- .../service/openai_quota_spark_window_test.go | 20 ++++++ backend/internal/service/wire.go | 63 ++++++++++++++++++- 9 files changed, 137 insertions(+), 41 deletions(-) diff --git a/backend/cmd/server/wire_gen.go b/backend/cmd/server/wire_gen.go index e148c5c363..088046de96 100644 --- a/backend/cmd/server/wire_gen.go +++ b/backend/cmd/server/wire_gen.go @@ -96,9 +96,6 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { usageLogRepository := repository.NewUsageLogRepository(client, db) usageService := service.NewUsageService(usageLogRepository, userRepository, client, apiKeyAuthCacheInvalidator) opsRepository := repository.NewOpsRepository(db) - batchImageRepository := repository.NewBatchImageRepository(db) - batchImageQueue := repository.NewBatchImageQueue(redisClient, configConfig) - batchImageDownloadLimiter := repository.NewBatchImageDownloadLimiter(redisClient, configConfig) usageBillingRepository := repository.NewUsageBillingRepository(client, db) gatewayCache := repository.NewGatewayCache(redisClient) schedulerOutboxRepository := repository.NewSchedulerOutboxRepository(db) @@ -137,11 +134,6 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { channelRepository := repository.NewChannelRepository(db) channelService := service.NewChannelService(channelRepository, groupRepository, apiKeyAuthCacheInvalidator, pricingService) modelPricingResolver := service.NewModelPricingResolver(channelService, billingService) - batchImageModelPricingResolver := service.ProvideBatchImageModelPricingResolver(modelPricingResolver) - batchImagePublicService := service.NewBatchImagePublicService(batchImageRepository, accountRepository, groupRepository, userGroupRateRepository, batchImageQueue, batchImageModelPricingResolver, usageBillingRepository, apiKeyAuthCacheInvalidator, configConfig) - batchImageDownloadService := service.NewBatchImageDownloadService(batchImageRepository, accountRepository, batchImageDownloadLimiter, configConfig) - batchImageCleanupService := service.ProvideBatchImageCleanupService(batchImageRepository, accountRepository, configConfig) - batchImageWorkerRuntime := service.ProvideBatchImageWorkerRuntime(batchImageRepository, accountRepository, batchImageQueue, usageBillingRepository, usageLogRepository, batchImageModelPricingResolver, apiKeyAuthCacheInvalidator, configConfig) notificationEmailService := service.NewNotificationEmailService(settingRepository, emailService) balanceNotifyService := service.ProvideBalanceNotifyService(emailService, settingRepository, accountRepository, notificationEmailService) gatewayService := service.NewGatewayService(accountRepository, groupRepository, usageLogRepository, usageBillingRepository, userRepository, userSubscriptionRepository, userGroupRateRepository, gatewayCache, configConfig, schedulerSnapshotService, concurrencyService, billingService, rateLimitService, billingCacheService, identityService, httpUpstream, deferredService, claudeTokenProvider, sessionLimitCache, rpmCache, digestSessionStore, settingService, tlsFingerprintProfileService, channelService, modelPricingResolver, balanceNotifyService, serviceUserPlatformQuotaRepository) @@ -190,10 +182,10 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { claudeUsageFetcher := repository.NewClaudeUsageFetcher(httpUpstream) antigravityQuotaFetcher := service.NewAntigravityQuotaFetcher(proxyRepository) grokQuotaFetcher := service.NewGrokQuotaFetcher() - openAIQuotaService := service.ProvideOpenAIQuotaService(accountRepository, proxyRepository, openAITokenProvider, privacyClientFactory) + openAIQuotaService := service.ProvideOpenAIQuotaService(accountRepository, proxyRepository, openAITokenProvider, privacyClientFactory, openAIGatewayService) usageCache := service.NewUsageCache() - accountUsageService := service.NewAccountUsageService(accountRepository, usageLogRepository, claudeUsageFetcher, geminiQuotaService, antigravityQuotaFetcher, grokQuotaFetcher, openAIQuotaService, usageCache, identityCache, tlsFingerprintProfileService) - accountTestService := service.NewAccountTestService(accountRepository, geminiTokenProvider, claudeTokenProvider, grokTokenProvider, antigravityGatewayService, httpUpstream, configConfig, tlsFingerprintProfileService) + accountUsageService := service.ProvideAccountUsageService(accountRepository, usageLogRepository, claudeUsageFetcher, geminiQuotaService, antigravityQuotaFetcher, grokQuotaFetcher, openAIQuotaService, usageCache, identityCache, tlsFingerprintProfileService, openAIGatewayService) + accountTestService := service.ProvideAccountTestService(accountRepository, geminiTokenProvider, claudeTokenProvider, grokTokenProvider, antigravityGatewayService, httpUpstream, configConfig, tlsFingerprintProfileService, openAIGatewayService) crsSyncService := service.NewCRSSyncService(accountRepository, proxyRepository, oAuthService, openAIOAuthService, geminiOAuthService, configConfig) accountHandler := admin.NewAccountHandler(adminService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, rateLimitService, accountUsageService, accountTestService, concurrencyService, crsSyncService, sessionLimitCache, rpmCache, compositeTokenCacheInvalidator) adminAnnouncementHandler := admin.NewAnnouncementHandler(announcementService) @@ -267,6 +259,13 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { handlerPaymentHandler := handler.NewPaymentHandler(paymentService, paymentConfigService, channelService) paymentWebhookHandler := handler.NewPaymentWebhookHandler(paymentService, registry) availableChannelHandler := handler.NewAvailableChannelHandler(channelService, apiKeyService, settingService) + batchImageRepository := repository.NewBatchImageRepository(db) + batchImageQueue := repository.NewBatchImageQueue(redisClient, configConfig) + batchImageModelPricingResolver := service.ProvideBatchImageModelPricingResolver(modelPricingResolver) + batchImagePublicService := service.NewBatchImagePublicService(batchImageRepository, accountRepository, groupRepository, userGroupRateRepository, batchImageQueue, batchImageModelPricingResolver, usageBillingRepository, apiKeyAuthCacheInvalidator, configConfig) + batchImageDownloadLimiter := repository.NewBatchImageDownloadLimiter(redisClient, configConfig) + batchImageDownloadService := service.NewBatchImageDownloadService(batchImageRepository, accountRepository, batchImageDownloadLimiter, configConfig) + batchImageCleanupService := service.ProvideBatchImageCleanupService(batchImageRepository, accountRepository, configConfig) batchImageHandler := handler.NewBatchImageHandler(batchImagePublicService, batchImageDownloadService, batchImageCleanupService) idempotencyCoordinator := service.ProvideIdempotencyCoordinator(idempotencyRepository, configConfig) idempotencyCleanupService := service.ProvideIdempotencyCleanupService(idempotencyRepository, configConfig) @@ -285,6 +284,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { accountExpiryService := service.ProvideAccountExpiryService(accountRepository) proxyExpiryService := service.ProvideProxyExpiryService(proxyRepository) subscriptionExpiryService := service.ProvideSubscriptionExpiryService(userSubscriptionRepository, settingRepository, notificationEmailService, leaderLockCache, db) + batchImageWorkerRuntime := service.ProvideBatchImageWorkerRuntime(batchImageRepository, accountRepository, batchImageQueue, usageBillingRepository, usageLogRepository, batchImageModelPricingResolver, apiKeyAuthCacheInvalidator, configConfig) scheduledTestRunnerService := service.ProvideScheduledTestRunnerService(scheduledTestPlanRepository, scheduledTestService, accountTestService, rateLimitService, configConfig) paymentOrderExpiryService := service.ProvidePaymentOrderExpiryService(paymentService, leaderLockCache, db) channelMonitorRunner := service.ProvideChannelMonitorRunner(channelMonitorService, settingService) diff --git a/backend/internal/service/account_test_service.go b/backend/internal/service/account_test_service.go index 68a00a5191..bb57a93cd7 100644 --- a/backend/internal/service/account_test_service.go +++ b/backend/internal/service/account_test_service.go @@ -74,6 +74,7 @@ type AccountTestService struct { cfg *config.Config tlsFPProfileService *TLSFingerprintProfileService agentIdentityTaskMu sync.Mutex + agentIdentityWS agentIdentityWSConnectionInvalidator } // NewAccountTestService creates a new AccountTestService @@ -605,7 +606,7 @@ func (s *AccountTestService) testOpenAIAccountConnection(c *gin.Context, account // Set common headers req.Header.Set("Content-Type", "application/json") if credentialAccount.IsOpenAIAgentIdentity() { - authHeaders, authErr := buildAgentIdentityAuthenticationHeaders(ctx, s.accountRepo, nil, &s.agentIdentityTaskMu, credentialAccount) + authHeaders, authErr := buildAgentIdentityAuthenticationHeaders(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, credentialAccount) if authErr != nil { return s.sendErrorAndEnd(c, "Failed to build Agent Identity authentication") } @@ -661,7 +662,7 @@ func (s *AccountTestService) testOpenAIAccountConnection(c *gin.Context, account body = redactAgentIdentitySensitiveBodyForAccount(ctx, s.accountRepo, credentialAccount, body) if !agentIdentityTaskRecoveryWasTried(ctx) && credentialAccount.IsOpenAIAgentIdentity() && isAgentIdentityTaskInvalidHTTPResponse(resp.StatusCode, body) { expectedTaskID := credentialAccount.GetCredential("task_id") - if err := ensureAgentIdentityTaskForAccount(ctx, s.accountRepo, nil, &s.agentIdentityTaskMu, credentialAccount, expectedTaskID); err != nil { + if err := ensureAgentIdentityTaskForAccount(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, credentialAccount, expectedTaskID); err != nil { return s.sendErrorAndEnd(c, fmt.Sprintf("Agent Identity task recovery failed: %s", err.Error())) } c.Request = c.Request.WithContext(markAgentIdentityTaskRecoveryTried(ctx)) @@ -895,7 +896,7 @@ func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account req.Header.Set("Content-Type", "application/json") req.Header.Set("Accept", "application/json") if credentialAccount.IsOpenAIAgentIdentity() { - authHeaders, authErr := buildAgentIdentityAuthenticationHeaders(ctx, s.accountRepo, nil, &s.agentIdentityTaskMu, credentialAccount) + authHeaders, authErr := buildAgentIdentityAuthenticationHeaders(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, credentialAccount) if authErr != nil { return s.sendErrorAndEnd(c, "Failed to build Agent Identity authentication") } @@ -943,7 +944,7 @@ func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account body = redactAgentIdentitySensitiveBodyForAccount(ctx, s.accountRepo, credentialAccount, body) if !agentIdentityTaskRecoveryWasTried(ctx) && credentialAccount.IsOpenAIAgentIdentity() && isAgentIdentityTaskInvalidHTTPResponse(resp.StatusCode, body) { expectedTaskID := credentialAccount.GetCredential("task_id") - if err := ensureAgentIdentityTaskForAccount(ctx, s.accountRepo, nil, &s.agentIdentityTaskMu, credentialAccount, expectedTaskID); err != nil { + if err := ensureAgentIdentityTaskForAccount(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, credentialAccount, expectedTaskID); err != nil { return s.sendErrorAndEnd(c, fmt.Sprintf("Agent Identity task recovery failed: %s", err.Error())) } c.Request = c.Request.WithContext(markAgentIdentityTaskRecoveryTried(ctx)) @@ -1785,7 +1786,7 @@ func (s *AccountTestService) testOpenAIImageOAuth(c *gin.Context, ctx context.Co req = req.WithContext(WithHTTPUpstreamProfile(req.Context(), HTTPUpstreamProfileOpenAI)) req.Host = "chatgpt.com" if credentialAccount.IsOpenAIAgentIdentity() { - authHeaders, authErr := buildAgentIdentityAuthenticationHeaders(ctx, s.accountRepo, nil, &s.agentIdentityTaskMu, credentialAccount) + authHeaders, authErr := buildAgentIdentityAuthenticationHeaders(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, credentialAccount) if authErr != nil { return s.sendErrorAndEnd(c, "Failed to build Agent Identity authentication") } diff --git a/backend/internal/service/account_usage_service.go b/backend/internal/service/account_usage_service.go index 4dc09c4109..622323c617 100644 --- a/backend/internal/service/account_usage_service.go +++ b/backend/internal/service/account_usage_service.go @@ -292,6 +292,7 @@ type AccountUsageService struct { identityCache IdentityCache tlsFPProfileService *TLSFingerprintProfileService agentIdentityTaskMu sync.Mutex + agentIdentityWS agentIdentityWSConnectionInvalidator } // NewAccountUsageService 创建AccountUsageService实例 @@ -706,7 +707,7 @@ func (s *AccountUsageService) probeOpenAICodexSnapshot(ctx context.Context, acco req.Host = "chatgpt.com" req.Header.Set("Content-Type", "application/json") if account.IsOpenAIAgentIdentity() { - authHeaders, authErr := buildAgentIdentityAuthenticationHeaders(ctx, s.accountRepo, nil, &s.agentIdentityTaskMu, account) + authHeaders, authErr := buildAgentIdentityAuthenticationHeaders(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, account) if authErr != nil { return nil, fmt.Errorf("build Agent Identity authentication: %w", authErr) } diff --git a/backend/internal/service/openai_agent_identity.go b/backend/internal/service/openai_agent_identity.go index 0492a2b82e..4bf447ad64 100644 --- a/backend/internal/service/openai_agent_identity.go +++ b/backend/internal/service/openai_agent_identity.go @@ -31,6 +31,10 @@ var openAIAgentIdentityAuthAPIBaseURL = agentIdentityAuthAPIBaseURL var agentIdentityTaskLocks sync.Map // map[int64]*sync.Mutex +type agentIdentityWSConnectionInvalidator interface { + InvalidateAgentIdentityWSConnections(accountID int64) +} + type agentIdentityKey struct { runtimeID string privateKey ed25519.PrivateKey @@ -231,7 +235,7 @@ func registerAgentIdentityTask(ctx context.Context, account *Account) (string, e return decryptAgentTaskID(key, encrypted) } -func ensureAgentIdentityTaskForAccount(ctx context.Context, repo AccountRepository, pool *openAIWSConnPool, taskMu *sync.Mutex, account *Account, expectedTaskID string) error { +func ensureAgentIdentityTaskForAccount(ctx context.Context, repo AccountRepository, wsInvalidator agentIdentityWSConnectionInvalidator, taskMu *sync.Mutex, account *Account, expectedTaskID string) error { if account == nil || !account.IsOpenAIAgentIdentity() { return nil } @@ -299,8 +303,8 @@ func ensureAgentIdentityTaskForAccount(ctx context.Context, repo AccountReposito if !account.IsShadow() && account != credAccount { account.Credentials = shallowCopyMap(credAccount.Credentials) } - if pool != nil { - pool.ClearAccount(credAccount.ID) + if wsInvalidator != nil { + wsInvalidator.InvalidateAgentIdentityWSConnections(credAccount.ID) } return nil } @@ -309,7 +313,7 @@ func (s *OpenAIGatewayService) ensureAgentIdentityTask(ctx context.Context, acco if s == nil { return errors.New("openai gateway service is nil") } - return ensureAgentIdentityTaskForAccount(ctx, s.accountRepo, s.openaiWSPool, &s.agentIdentityTaskMu, account, expectedTaskID) + return ensureAgentIdentityTaskForAccount(ctx, s.accountRepo, s, &s.agentIdentityTaskMu, account, expectedTaskID) } func isAgentIdentityTaskInvalidHTTPResponse(statusCode int, body []byte) bool { @@ -374,7 +378,7 @@ func (s *OpenAIGatewayService) buildOpenAIAuthenticationHeaders(ctx context.Cont } headers := make(http.Header) if credAccount != nil && credAccount.IsOpenAIAgentIdentity() { - agentHeaders, err := buildAgentIdentityAuthenticationHeaders(ctx, s.accountRepo, s.openaiWSPool, &s.agentIdentityTaskMu, credAccount) + agentHeaders, err := buildAgentIdentityAuthenticationHeaders(ctx, s.accountRepo, s, &s.agentIdentityTaskMu, credAccount) if err != nil { return nil, err } @@ -384,11 +388,11 @@ func (s *OpenAIGatewayService) buildOpenAIAuthenticationHeaders(ctx context.Cont return headers, nil } -func buildAgentIdentityAuthenticationHeaders(ctx context.Context, repo AccountRepository, pool *openAIWSConnPool, taskMu *sync.Mutex, account *Account) (http.Header, error) { +func buildAgentIdentityAuthenticationHeaders(ctx context.Context, repo AccountRepository, wsInvalidator agentIdentityWSConnectionInvalidator, taskMu *sync.Mutex, account *Account) (http.Header, error) { if account == nil || !account.IsOpenAIAgentIdentity() { return nil, errors.New("agent identity account is required") } - if err := ensureAgentIdentityTaskForAccount(ctx, repo, pool, taskMu, account, ""); err != nil { + if err := ensureAgentIdentityTaskForAccount(ctx, repo, wsInvalidator, taskMu, account, ""); err != nil { return nil, err } key, err := agentIdentityKeyFromAccount(account) @@ -423,7 +427,7 @@ func (s *OpenAIGatewayService) refreshOpenAIAgentIdentityHeaders(ctx context.Con if refreshed == nil { refreshed = make(http.Header) } - authHeaders, err := buildAgentIdentityAuthenticationHeaders(ctx, s.accountRepo, s.openaiWSPool, &s.agentIdentityTaskMu, credAccount) + authHeaders, err := buildAgentIdentityAuthenticationHeaders(ctx, s.accountRepo, s, &s.agentIdentityTaskMu, credAccount) if err != nil { return nil, err } diff --git a/backend/internal/service/openai_agent_identity_compat_test.go b/backend/internal/service/openai_agent_identity_compat_test.go index ff58276310..eb90dfb56e 100644 --- a/backend/internal/service/openai_agent_identity_compat_test.go +++ b/backend/internal/service/openai_agent_identity_compat_test.go @@ -89,7 +89,8 @@ func TestAccountTestServiceOpenAICompactAgentIdentityRecoversInvalidTaskOnce(t * {StatusCode: http.StatusUnauthorized, Header: http.Header{"Content-Type": []string{"application/json"}}, Body: io.NopCloser(strings.NewReader(`{"error":{"code":"invalid_task_id"}}`))}, {StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"application/json"}}, Body: io.NopCloser(strings.NewReader(`{"id":"compact-agent","status":"completed"}`))}, }} - svc := &AccountTestService{accountRepo: repo, httpUpstream: upstream} + invalidator := &agentIdentityWSInvalidationRecorder{} + svc := &AccountTestService{accountRepo: repo, httpUpstream: upstream, agentIdentityWS: invalidator} rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/22/test", bytes.NewReader(nil)) @@ -99,6 +100,7 @@ func TestAccountTestServiceOpenAICompactAgentIdentityRecoversInvalidTaskOnce(t * require.Len(t, upstream.requests, 2) require.Equal(t, "task-compact-new", account.GetCredential("task_id")) require.Equal(t, 0, repo.setErrorCalls) + require.Equal(t, []int64{account.ID}, invalidator.accountIDs) } func TestOpenAIAgentIdentityPassthroughKeepsSessionAndPromptCacheHeaders(t *testing.T) { @@ -431,6 +433,14 @@ type agentIdentityForwardRepo struct { account *Account } +type agentIdentityWSInvalidationRecorder struct { + accountIDs []int64 +} + +func (r *agentIdentityWSInvalidationRecorder) InvalidateAgentIdentityWSConnections(accountID int64) { + r.accountIDs = append(r.accountIDs, accountID) +} + type accountTestAgentIdentityRepo struct { AccountRepository account *Account diff --git a/backend/internal/service/openai_gateway_service.go b/backend/internal/service/openai_gateway_service.go index a17a6f6b05..24f4786e09 100644 --- a/backend/internal/service/openai_gateway_service.go +++ b/backend/internal/service/openai_gateway_service.go @@ -554,6 +554,12 @@ func (s *OpenAIGatewayService) CloseOpenAIWSPool() { } } +func (s *OpenAIGatewayService) InvalidateAgentIdentityWSConnections(accountID int64) { + if pool := s.getOpenAIWSConnPool(); pool != nil { + pool.ClearAccount(accountID) + } +} + func (s *OpenAIGatewayService) logOpenAIWSModeBootstrap() { if s == nil || s.cfg == nil { return diff --git a/backend/internal/service/openai_quota_service.go b/backend/internal/service/openai_quota_service.go index c736768b72..735a580570 100644 --- a/backend/internal/service/openai_quota_service.go +++ b/backend/internal/service/openai_quota_service.go @@ -24,6 +24,8 @@ import ( // errors.Is still matches it by identity since ResetCredit returns this var. var ErrSparkShadowResetNotSupported = infraerrors.New(http.StatusConflict, "SPARK_SHADOW_RESET_NOT_SUPPORTED", "spark shadow account does not support credit reset; reset the parent account") +var ErrAgentIdentityResetNotSupported = infraerrors.New(http.StatusConflict, "AGENT_IDENTITY_RESET_NOT_SUPPORTED", "agent identity does not support rate-limit reset credit consumption") + // Endpoints used by the OpenAI/ChatGPT/Codex quota query and reset feature. const ( chatGPTUsageURL = "https://chatgpt.com/backend-api/wham/usage" @@ -120,6 +122,7 @@ type OpenAIQuotaService struct { tokenProvider *OpenAITokenProvider privacyClientFactory PrivacyClientFactory agentIdentityTaskMu sync.Mutex + agentIdentityWS agentIdentityWSConnectionInvalidator } // NewOpenAIQuotaService constructs a quota service. token provider is required — @@ -241,6 +244,9 @@ func (s *OpenAIQuotaService) ResetCredit(ctx context.Context, accountID int64) ( if acc.IsShadow() { return nil, ErrSparkShadowResetNotSupported } + if acc.IsOpenAIAgentIdentity() { + return nil, ErrAgentIdentityResetNotSupported + } } accessToken, chatGPTAccountID, proxyURL, fedRAMP, err := s.prepareUpstreamCall(ctx, accountID) @@ -260,10 +266,8 @@ func (s *OpenAIQuotaService) ResetCredit(ctx context.Context, accountID int64) ( callCtx, cancel := context.WithTimeout(ctx, openaiQuotaUpstreamTimeout) defer cancel() - agentIdentity := s.isAgentIdentityAccount(ctx, accountID) - var payload OpenAIQuotaResetResult - for recovered := false; ; { + for { headers, headerErr := s.buildCodexQuotaHeaders(callCtx, accountID, accessToken, chatGPTAccountID, fedRAMP) if headerErr != nil { return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_QUOTA_AUTH_FAILED", "failed to build upstream authentication: %v", headerErr) @@ -279,13 +283,6 @@ func (s *OpenAIQuotaService) ResetCredit(ctx context.Context, accountID int64) ( return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_QUOTA_RESET_REQUEST_FAILED", "upstream request failed: %v", err) } if !resp.IsSuccessState() { - if agentIdentity && !recovered && isAgentIdentityTaskInvalidHTTPResponse(resp.StatusCode, []byte(resp.String())) { - recovered = true - if err := s.recoverAgentIdentityTask(ctx, accountID); err != nil { - return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_QUOTA_AUTH_FAILED", "agent identity task recovery failed: %v", err) - } - continue - } status := resp.StatusCode body := truncate(s.redactQuotaErrorBody(callCtx, accountID, resp.String()), 240) slog.Warn("openai_quota_reset_failed", "account_id", accountID, "status", status, "body", body) @@ -394,7 +391,7 @@ func (s *OpenAIQuotaService) recoverAgentIdentityTask(ctx context.Context, accou if !account.IsOpenAIAgentIdentity() { return nil } - return ensureAgentIdentityTaskForAccount(ctx, s.accountRepo, nil, &s.agentIdentityTaskMu, account, account.GetCredential("task_id")) + return ensureAgentIdentityTaskForAccount(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, account, account.GetCredential("task_id")) } func (s *OpenAIQuotaService) isAgentIdentityAccount(ctx context.Context, accountID int64) bool { @@ -436,7 +433,7 @@ func (s *OpenAIQuotaService) buildCodexQuotaHeaders(ctx context.Context, account if !account.IsOpenAIAgentIdentity() { return headers, nil } - if err := ensureAgentIdentityTaskForAccount(ctx, s.accountRepo, nil, &s.agentIdentityTaskMu, account, ""); err != nil { + if err := ensureAgentIdentityTaskForAccount(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, account, ""); err != nil { return nil, err } key, err := agentIdentityKeyFromAccount(account) diff --git a/backend/internal/service/openai_quota_spark_window_test.go b/backend/internal/service/openai_quota_spark_window_test.go index 99bfe15ca1..4e5fbc4a84 100644 --- a/backend/internal/service/openai_quota_spark_window_test.go +++ b/backend/internal/service/openai_quota_spark_window_test.go @@ -176,6 +176,23 @@ func TestResetCreditShadowRejected(t *testing.T) { "shadow ResetCredit 应映射为 409 Conflict 而非 500") } +func TestResetCreditAgentIdentityRejectedBeforeUpstream(t *testing.T) { + account := &Account{ + ID: 201, + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Credentials: map[string]any{ + "auth_mode": OpenAIAuthModeAgentIdentity, + }, + } + repo := &stubQuotaAccountRepo{accounts: map[int64]*Account{account.ID: account}} + svc := &OpenAIQuotaService{accountRepo: repo} + + _, err := svc.ResetCredit(context.Background(), account.ID) + require.ErrorIs(t, err, ErrAgentIdentityResetNotSupported) + require.Equal(t, http.StatusConflict, infraerrors.Code(err)) +} + // ── Part B: prepareUpstreamCall 影子 resolve ────────────────────────────── // TestPrepareUpstreamCallShadowResolve 验证影子账号（200）QueryUsage 时: @@ -306,13 +323,16 @@ func TestQueryUsageAgentIdentityRecoversInvalidTaskOnce(t *testing.T) { openAIAgentIdentityAuthAPIBaseURL = srv.URL t.Cleanup(func() { openAIAgentIdentityAuthAPIBaseURL = oldBase }) + invalidator := &agentIdentityWSInvalidationRecorder{} svc := NewOpenAIQuotaService(repo, nil, nil, newQuotaRedirectingFactory(srv)) + svc.agentIdentityWS = invalidator usage, err := svc.QueryUsage(context.Background(), account.ID) require.NoError(t, err) require.NotNil(t, usage) require.Equal(t, 2, usageCalls) require.Equal(t, 1, registerCalls) require.Equal(t, "task-quota-new", account.GetCredential("task_id")) + require.Equal(t, []int64{account.ID}, invalidator.accountIDs) } func TestParseOpenAIRateLimitResetCreditDetails_CompatibleContainers(t *testing.T) { diff --git a/backend/internal/service/wire.go b/backend/internal/service/wire.go index 7258ff05a3..8c20af661b 100644 --- a/backend/internal/service/wire.go +++ b/backend/internal/service/wire.go @@ -131,8 +131,65 @@ func ProvideOpenAIQuotaService( proxyRepo ProxyRepository, tokenProvider *OpenAITokenProvider, privacyClientFactory PrivacyClientFactory, + openAIGatewayService *OpenAIGatewayService, ) *OpenAIQuotaService { - return NewOpenAIQuotaService(accountRepo, proxyRepo, tokenProvider, privacyClientFactory) + service := NewOpenAIQuotaService(accountRepo, proxyRepo, tokenProvider, privacyClientFactory) + service.agentIdentityWS = openAIGatewayService + return service +} + +func ProvideAccountUsageService( + accountRepo AccountRepository, + usageLogRepo UsageLogRepository, + usageFetcher ClaudeUsageFetcher, + geminiQuotaService *GeminiQuotaService, + antigravityQuotaFetcher *AntigravityQuotaFetcher, + grokQuotaFetcher *GrokQuotaFetcher, + openAIQuotaService *OpenAIQuotaService, + cache *UsageCache, + identityCache IdentityCache, + tlsFPProfileService *TLSFingerprintProfileService, + openAIGatewayService *OpenAIGatewayService, +) *AccountUsageService { + service := NewAccountUsageService( + accountRepo, + usageLogRepo, + usageFetcher, + geminiQuotaService, + antigravityQuotaFetcher, + grokQuotaFetcher, + openAIQuotaService, + cache, + identityCache, + tlsFPProfileService, + ) + service.agentIdentityWS = openAIGatewayService + return service +} + +func ProvideAccountTestService( + accountRepo AccountRepository, + geminiTokenProvider *GeminiTokenProvider, + claudeTokenProvider *ClaudeTokenProvider, + grokTokenProvider *GrokTokenProvider, + antigravityGatewayService *AntigravityGatewayService, + httpUpstream HTTPUpstream, + cfg *config.Config, + tlsFPProfileService *TLSFingerprintProfileService, + openAIGatewayService *OpenAIGatewayService, +) *AccountTestService { + service := NewAccountTestService( + accountRepo, + geminiTokenProvider, + claudeTokenProvider, + grokTokenProvider, + antigravityGatewayService, + httpUpstream, + cfg, + tlsFPProfileService, + ) + service.agentIdentityWS = openAIGatewayService + return service } func ProvideGrokQuotaService( @@ -601,8 +658,8 @@ var ProviderSet = wire.NewSet( ProvideClaudeTokenProvider, NewAntigravityGatewayService, ProvideRateLimitService, - NewAccountUsageService, - NewAccountTestService, + ProvideAccountUsageService, + ProvideAccountTestService, ProvideSettingService, NewDataManagementService, ProvideBackupService,