From 10aa88aab95f19a48b66142ae2803408fd37a2be Mon Sep 17 00:00:00 2001 From: cat Date: Tue, 14 Jul 2026 17:19:15 +0800 Subject: [PATCH] =?UTF-8?q?fix(openai):=20=E8=A1=A5=E9=BD=90=E8=BA=AB?= =?UTF-8?q?=E4=BB=BD=E5=A4=B1=E6=95=88=E6=81=A2=E5=A4=8D=E8=B7=AF=E5=BE=84?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../internal/service/account_test_service.go | 31 +++- .../internal/service/openai_agent_identity.go | 36 ++++- .../openai_agent_identity_compat_test.go | 147 +++++++++++++++++- .../openai_gateway_chat_completions.go | 7 + .../service/openai_gateway_messages.go | 7 + .../service/openai_images_responses.go | 8 + .../internal/service/openai_quota_service.go | 132 +++++++++++----- .../service/openai_quota_spark_window_test.go | 58 +++++++ 8 files changed, 378 insertions(+), 48 deletions(-) diff --git a/backend/internal/service/account_test_service.go b/backend/internal/service/account_test_service.go index 5a0bc40f40..68a00a5191 100644 --- a/backend/internal/service/account_test_service.go +++ b/backend/internal/service/account_test_service.go @@ -590,8 +590,11 @@ func (s *AccountTestService) testOpenAIAccountConnection(c *gin.Context, account payload := createOpenAITestPayload(testModelID, isOAuth) payloadBytes, _ := json.Marshal(payload) - // Send test_start event - s.sendEvent(c, TestEvent{Type: "test_start", Model: testModelID}) + // Send test_start event once. A task-invalid Agent Identity response may + // restart this probe after registering a replacement task. + if !agentIdentityTaskRecoveryWasTried(ctx) { + s.sendEvent(c, TestEvent{Type: "test_start", Model: testModelID}) + } req, err := http.NewRequestWithContext(ctx, "POST", apiURL, bytes.NewReader(payloadBytes)) if err != nil { @@ -656,6 +659,14 @@ func (s *AccountTestService) testOpenAIAccountConnection(c *gin.Context, account if resp.StatusCode != http.StatusOK { body, _ := io.ReadAll(resp.Body) 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 { + return s.sendErrorAndEnd(c, fmt.Sprintf("Agent Identity task recovery failed: %s", err.Error())) + } + c.Request = c.Request.WithContext(markAgentIdentityTaskRecoveryTried(ctx)) + return s.testOpenAIAccountConnection(c, account, modelID, prompt, mode) + } if resp.StatusCode == http.StatusTooManyRequests { s.reconcileOpenAI429State(ctx, account, resp.Header, body) } @@ -718,7 +729,9 @@ func (s *AccountTestService) testGrokAccountConnection(c *gin.Context, account * return s.sendErrorAndEnd(c, "Failed to create Grok test payload") } - s.sendEvent(c, TestEvent{Type: "test_start", Model: testModelID}) + if !agentIdentityTaskRecoveryWasTried(ctx) { + s.sendEvent(c, TestEvent{Type: "test_start", Model: testModelID}) + } req, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, bytes.NewReader(payloadBytes)) if err != nil { @@ -869,7 +882,9 @@ func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account c.Writer.Flush() payloadBytes, _ := json.Marshal(createOpenAICompactProbePayload(testModelID)) - s.sendEvent(c, TestEvent{Type: "test_start", Model: testModelID}) + if !agentIdentityTaskRecoveryWasTried(ctx) { + s.sendEvent(c, TestEvent{Type: "test_start", Model: testModelID}) + } req, err := http.NewRequestWithContext(ctx, "POST", apiURL, bytes.NewReader(payloadBytes)) if err != nil { @@ -926,6 +941,14 @@ func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account body, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) 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 { + return s.sendErrorAndEnd(c, fmt.Sprintf("Agent Identity task recovery failed: %s", err.Error())) + } + c.Request = c.Request.WithContext(markAgentIdentityTaskRecoveryTried(ctx)) + return s.testOpenAICompactConnection(c, account, testModelID) + } if s.accountRepo != nil { updates := buildOpenAICompactProbeExtraUpdates(resp, body, nil, time.Now()) diff --git a/backend/internal/service/openai_agent_identity.go b/backend/internal/service/openai_agent_identity.go index a5ea9b7612..0492a2b82e 100644 --- a/backend/internal/service/openai_agent_identity.go +++ b/backend/internal/service/openai_agent_identity.go @@ -317,13 +317,26 @@ func isAgentIdentityTaskInvalidHTTPResponse(statusCode int, body []byte) bool { return false } lower := strings.ToLower(string(body)) + compact := strings.NewReplacer(" ", "", "\t", "", "\r", "", "\n", "").Replace(lower) for _, marker := range []string{ - "invalid task", - "task_id", - "task id", - "task_not_found", - "task_expired", - "unknown task", + `"code":"invalid_task_id"`, + `"code":"task_not_found"`, + `"code":"task_expired"`, + `"error":"invalid_task_id"`, + } { + if strings.Contains(compact, marker) { + return true + } + } + for _, marker := range []string{ + "invalid task_id", + "invalid task id", + "task_id is invalid", + "task id is invalid", + "task not found", + "task expired", + "unknown task_id", + "unknown task id", } { if strings.Contains(lower, marker) { return true @@ -332,6 +345,17 @@ func isAgentIdentityTaskInvalidHTTPResponse(statusCode int, body []byte) bool { return false } +type agentIdentityTaskRecoveryContextKey struct{} + +func markAgentIdentityTaskRecoveryTried(ctx context.Context) context.Context { + return context.WithValue(ctx, agentIdentityTaskRecoveryContextKey{}, true) +} + +func agentIdentityTaskRecoveryWasTried(ctx context.Context) bool { + tried, _ := ctx.Value(agentIdentityTaskRecoveryContextKey{}).(bool) + return tried +} + func isAgentIdentityTaskInvalidWSDialError(err *openAIWSDialError) bool { return err != nil && isAgentIdentityTaskInvalidHTTPResponse(err.StatusCode, err.ResponseBody) } diff --git a/backend/internal/service/openai_agent_identity_compat_test.go b/backend/internal/service/openai_agent_identity_compat_test.go index c034af2bfa..bbfa82481d 100644 --- a/backend/internal/service/openai_agent_identity_compat_test.go +++ b/backend/internal/service/openai_agent_identity_compat_test.go @@ -55,6 +55,52 @@ func TestAccountTestServiceOpenAICompactAgentIdentityUsesFreshAssertion(t *testi require.NotContains(t, upstream.lastReq.Header.Get("Authorization"), privateKey) } +func TestAccountTestServiceOpenAICompactAgentIdentityRecoversInvalidTaskOnce(t *testing.T) { + gin.SetMode(gin.TestMode) + key, privateKey := newTestAgentIdentityKey(t) + account := &Account{ + ID: 22, + Name: "agent-identity-recovery", + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Credentials: map[string]any{ + "auth_mode": OpenAIAuthModeAgentIdentity, + "agent_runtime_id": key.runtimeID, + "agent_private_key": privateKey, + "task_id": "task-compact-old", + "chatgpt_account_id": "account-agent-compact-recovery", + }, + } + repo := &accountTestAgentIdentityRepo{account: account} + registerCalls := 0 + registerServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + registerCalls++ + _, _ = io.WriteString(w, `{"task_id":"task-compact-new"}`) + })) + defer registerServer.Close() + oldBase := openAIAgentIdentityAuthAPIBaseURL + openAIAgentIdentityAuthAPIBaseURL = registerServer.URL + t.Cleanup(func() { openAIAgentIdentityAuthAPIBaseURL = oldBase }) + + upstream := &httpUpstreamRecorder{responses: []*http.Response{ + {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} + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/22/test", bytes.NewReader(nil)) + + require.NoError(t, svc.TestAccountConnection(c, account.ID, "gpt-5.4", "", AccountTestModeCompact)) + require.Equal(t, 1, registerCalls) + require.Len(t, upstream.requests, 2) + require.Equal(t, "task-compact-new", account.GetCredential("task_id")) + require.Equal(t, 0, repo.setErrorCalls) +} + func TestOpenAIAgentIdentityPassthroughKeepsSessionAndPromptCacheHeaders(t *testing.T) { gin.SetMode(gin.TestMode) key, privateKey := newTestAgentIdentityKey(t) @@ -225,6 +271,7 @@ func TestOpenAIAgentIdentityTaskInvalidRetriesExactlyOnce(t *testing.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(successBody))}, }} + require.True(t, isAgentIdentityTaskInvalidHTTPResponse(http.StatusUnauthorized, []byte(`{"error":{"code":"invalid_task_id"}}`))) svc := &OpenAIGatewayService{cfg: &config.Config{}, accountRepo: repo, httpUpstream: upstream} rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) @@ -256,7 +303,7 @@ func TestOpenAIAgentIdentityTaskInvalidRetriesExactlyOnce(t *testing.T) { account.Credentials["task_id"] = "task-old-passthrough" upstream.responses = []*http.Response{ {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(successBody))}, + {StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"text/event-stream"}}, Body: io.NopCloser(strings.NewReader("data: {\"type\":\"response.completed\",\"response\":{\"usage\":{\"input_tokens\":1,\"output_tokens\":1,\"total_tokens\":2}}}\n\ndata: [DONE]\n\n"))}, } rec3 := httptest.NewRecorder() c3, _ := gin.CreateTestContext(rec3) @@ -267,6 +314,80 @@ func TestOpenAIAgentIdentityTaskInvalidRetriesExactlyOnce(t *testing.T) { require.Len(t, upstream.requests, 6) } +func TestOpenAIAgentIdentityCompatRoutesRecoverInvalidTaskOnce(t *testing.T) { + gin.SetMode(gin.TestMode) + tests := []struct { + name string + path string + body []byte + call func(*OpenAIGatewayService, context.Context, *gin.Context, *Account, []byte) (*OpenAIForwardResult, error) + }{ + { + name: "chat completions", + path: "/v1/chat/completions", + body: []byte(`{"model":"gpt-5.4","stream":false,"messages":[{"role":"user","content":"hi"}]}`), + call: func(s *OpenAIGatewayService, ctx context.Context, c *gin.Context, account *Account, body []byte) (*OpenAIForwardResult, error) { + return s.ForwardAsChatCompletions(ctx, c, account, body, "", "gpt-5.4") + }, + }, + { + name: "anthropic messages", + path: "/v1/messages", + body: []byte(`{"model":"gpt-5.4","stream":false,"max_tokens":32,"messages":[{"role":"user","content":"hi"}]}`), + call: func(s *OpenAIGatewayService, ctx context.Context, c *gin.Context, account *Account, body []byte) (*OpenAIForwardResult, error) { + return s.ForwardAsAnthropic(ctx, c, account, body, "", "gpt-5.4") + }, + }, + } + + for index, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + key, privateKey := newTestAgentIdentityKey(t) + account := &Account{ + ID: int64(40 + index), + Name: "agent-identity-compat", + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Credentials: map[string]any{ + "auth_mode": OpenAIAuthModeAgentIdentity, + "agent_runtime_id": key.runtimeID, + "agent_private_key": privateKey, + "task_id": "task-compat-old", + "chatgpt_account_id": "account-compat-recovery", + }, + } + repo := &agentIdentityForwardRepo{account: account} + registerCalls := 0 + registerServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + registerCalls++ + _, _ = io.WriteString(w, `{"task_id":"task-compat-new"}`) + })) + defer registerServer.Close() + oldBase := openAIAgentIdentityAuthAPIBaseURL + openAIAgentIdentityAuthAPIBaseURL = registerServer.URL + t.Cleanup(func() { openAIAgentIdentityAuthAPIBaseURL = oldBase }) + + upstream := &httpUpstreamRecorder{responses: []*http.Response{ + {StatusCode: http.StatusUnauthorized, Header: http.Header{"Content-Type": []string{"application/json"}}, Body: io.NopCloser(strings.NewReader(`{"error":{"code":"invalid_task_id"}}`))}, + {StatusCode: http.StatusUnauthorized, Header: http.Header{"Content-Type": []string{"application/json"}}, Body: io.NopCloser(strings.NewReader(`{"error":{"code":"invalid_task_id"}}`))}, + }} + svc := &OpenAIGatewayService{cfg: &config.Config{}, accountRepo: repo, httpUpstream: upstream} + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, tt.path, bytes.NewReader(tt.body)) + + _, err := tt.call(svc, context.Background(), c, account, tt.body) + require.Error(t, err) + require.Equal(t, 1, registerCalls) + require.Len(t, upstream.requests, 2) + require.Equal(t, "task-compat-new", account.GetCredential("task_id")) + }) + } +} + func decodeAgentAssertionTask(t *testing.T, header string) string { t.Helper() encoded := strings.TrimPrefix(header, "AgentAssertion ") @@ -284,6 +405,30 @@ type agentIdentityForwardRepo struct { account *Account } +type accountTestAgentIdentityRepo struct { + AccountRepository + account *Account + setErrorCalls int +} + +func (r *accountTestAgentIdentityRepo) GetByID(_ context.Context, _ int64) (*Account, error) { + return r.account, nil +} + +func (r *accountTestAgentIdentityRepo) UpdateCredentials(_ context.Context, _ int64, credentials map[string]any) error { + r.account.Credentials = credentials + return nil +} + +func (r *accountTestAgentIdentityRepo) UpdateExtra(_ context.Context, _ int64, _ map[string]any) error { + return nil +} + +func (r *accountTestAgentIdentityRepo) SetError(_ context.Context, _ int64, _ string) error { + r.setErrorCalls++ + return nil +} + func (r *agentIdentityForwardRepo) GetByID(_ context.Context, _ int64) (*Account, error) { return r.account, nil } diff --git a/backend/internal/service/openai_gateway_chat_completions.go b/backend/internal/service/openai_gateway_chat_completions.go index 15c4e47635..5d80292c47 100644 --- a/backend/internal/service/openai_gateway_chat_completions.go +++ b/backend/internal/service/openai_gateway_chat_completions.go @@ -267,6 +267,13 @@ func (s *OpenAIGatewayService) ForwardAsChatCompletions( // 8. Handle error response with failover if resp.StatusCode >= 400 { respBody, upstreamMsg := s.readOpenAIUpstreamError(resp) + if !agentIdentityTaskRecoveryWasTried(ctx) && s.isAgentIdentityAccount(ctx, account) && isAgentIdentityTaskInvalidHTTPResponse(resp.StatusCode, respBody) { + expectedTaskID := account.GetCredential("task_id") + if err := s.recoverAgentIdentityTask(ctx, account, expectedTaskID); err != nil { + return nil, fmt.Errorf("agent identity task recovery failed: %w", err) + } + return s.ForwardAsChatCompletions(markAgentIdentityTaskRecoveryTried(ctx), c, account, body, promptCacheKey, defaultMappedModel) + } if account.Type == AccountTypeAPIKey && openai_compat.ResolveResponsesSupport(account.Extra) == openai_compat.ResponsesSupportUnknown && !isResponsesEndpointSupportedByStatus(resp.StatusCode) { diff --git a/backend/internal/service/openai_gateway_messages.go b/backend/internal/service/openai_gateway_messages.go index 219b5e4be4..4d8b14cc28 100644 --- a/backend/internal/service/openai_gateway_messages.go +++ b/backend/internal/service/openai_gateway_messages.go @@ -316,6 +316,13 @@ func (s *OpenAIGatewayService) ForwardAsAnthropic( // 8. Handle error response with failover if resp.StatusCode >= 400 { respBody, upstreamMsg := s.readOpenAIUpstreamError(resp) + if !agentIdentityTaskRecoveryWasTried(ctx) && s.isAgentIdentityAccount(ctx, account) && isAgentIdentityTaskInvalidHTTPResponse(resp.StatusCode, respBody) { + expectedTaskID := account.GetCredential("task_id") + if err := s.recoverAgentIdentityTask(ctx, account, expectedTaskID); err != nil { + return nil, fmt.Errorf("agent identity task recovery failed: %w", err) + } + return s.ForwardAsAnthropic(markAgentIdentityTaskRecoveryTried(ctx), c, account, body, promptCacheKey, defaultMappedModel) + } if account.Platform == PlatformGrok { s.updateGrokUsageSnapshot(ctx, account.ID, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) s.handleGrokAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody) diff --git a/backend/internal/service/openai_images_responses.go b/backend/internal/service/openai_images_responses.go index 226cbf8e88..04347d6f90 100644 --- a/backend/internal/service/openai_images_responses.go +++ b/backend/internal/service/openai_images_responses.go @@ -1560,6 +1560,14 @@ func (s *OpenAIGatewayService) forwardOpenAIImagesOAuth( if resp.StatusCode >= 400 { respBody := s.readUpstreamErrorBody(resp) _ = resp.Body.Close() + respBody = s.redactAgentIdentitySensitiveBody(upstreamCtx, account, respBody) + if !agentIdentityTaskRecoveryWasTried(ctx) && s.isAgentIdentityAccount(ctx, account) && isAgentIdentityTaskInvalidHTTPResponse(resp.StatusCode, respBody) { + expectedTaskID := account.GetCredential("task_id") + if err := s.recoverAgentIdentityTask(ctx, account, expectedTaskID); err != nil { + return nil, fmt.Errorf("agent identity task recovery failed: %w", err) + } + return s.forwardOpenAIImagesOAuth(markAgentIdentityTaskRecoveryTried(ctx), c, account, parsed, channelMappedModel) + } resp.Body = io.NopCloser(bytes.NewReader(respBody)) upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody)) upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) diff --git a/backend/internal/service/openai_quota_service.go b/backend/internal/service/openai_quota_service.go index cebab0a1d7..c736768b72 100644 --- a/backend/internal/service/openai_quota_service.go +++ b/backend/internal/service/openai_quota_service.go @@ -155,25 +155,36 @@ func (s *OpenAIQuotaService) QueryUsage(ctx context.Context, accountID int64) (* callCtx, cancel := context.WithTimeout(ctx, openaiQuotaUpstreamTimeout) defer cancel() + agentIdentity := s.isAgentIdentityAccount(ctx, accountID) - quotaHeaders, 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) - } var payload OpenAIQuotaUsage - resp, err := client.R(). - SetContext(callCtx). - SetHeaders(quotaHeaders). - SetSuccessResult(&payload). - Get(chatGPTUsageURL) - if err != nil { - return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_QUOTA_REQUEST_FAILED", "upstream request failed: %v", err) - } - if !resp.IsSuccessState() { - status := resp.StatusCode - body := truncate(s.redactQuotaErrorBody(ctx, accountID, resp.String()), 240) - slog.Warn("openai_quota_query_failed", "account_id", accountID, "status", status, "body", body) - return nil, infraerrors.Newf(mapUpstreamStatus(status), "OPENAI_QUOTA_UPSTREAM_ERROR", "upstream returned %d: %s", status, body) + for recovered := false; ; { + quotaHeaders, 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) + } + resp, err := client.R(). + SetContext(callCtx). + SetHeaders(quotaHeaders). + SetSuccessResult(&payload). + Get(chatGPTUsageURL) + if err != nil { + return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_QUOTA_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(ctx, accountID, resp.String()), 240) + slog.Warn("openai_quota_query_failed", "account_id", accountID, "status", status, "body", body) + return nil, infraerrors.Newf(mapUpstreamStatus(status), "OPENAI_QUOTA_UPSTREAM_ERROR", "upstream returned %d: %s", status, body) + } + break } payload.FetchedAt = time.Now().Unix() @@ -249,28 +260,38 @@ func (s *OpenAIQuotaService) ResetCredit(ctx context.Context, accountID int64) ( callCtx, cancel := context.WithTimeout(ctx, openaiQuotaUpstreamTimeout) defer cancel() - - 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) - } - headers["content-type"] = "application/json" + agentIdentity := s.isAgentIdentityAccount(ctx, accountID) var payload OpenAIQuotaResetResult - resp, err := client.R(). - SetContext(callCtx). - SetHeaders(headers). - SetBody(map[string]string{"redeem_request_id": redeemRequestID}). - SetSuccessResult(&payload). - Post(chatGPTRateLimitResetURL) - if err != nil { - return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_QUOTA_RESET_REQUEST_FAILED", "upstream request failed: %v", err) - } - if !resp.IsSuccessState() { - 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) - return nil, infraerrors.Newf(mapUpstreamStatus(status), "OPENAI_QUOTA_RESET_UPSTREAM_ERROR", "upstream returned %d: %s", status, body) + for recovered := false; ; { + 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) + } + headers["content-type"] = "application/json" + resp, err := client.R(). + SetContext(callCtx). + SetHeaders(headers). + SetBody(map[string]string{"redeem_request_id": redeemRequestID}). + SetSuccessResult(&payload). + Post(chatGPTRateLimitResetURL) + if err != nil { + 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) + return nil, infraerrors.Newf(mapUpstreamStatus(status), "OPENAI_QUOTA_RESET_UPSTREAM_ERROR", "upstream returned %d: %s", status, body) + } + break } slog.Info("openai_quota_reset_success", @@ -356,6 +377,43 @@ func (s *OpenAIQuotaService) prepareUpstreamCall(ctx context.Context, accountID return accessToken, chatGPTAccountID, proxyURL, fedRAMP, nil } +func (s *OpenAIQuotaService) recoverAgentIdentityTask(ctx context.Context, accountID int64) error { + if s == nil || s.accountRepo == nil { + return fmt.Errorf("account repository is unavailable") + } + account, err := s.accountRepo.GetByID(ctx, accountID) + if err != nil || account == nil { + return fmt.Errorf("account is unavailable") + } + if account.IsShadow() { + account, err = resolveCredentialAccount(ctx, s.accountRepo, account) + if err != nil || account == nil { + return fmt.Errorf("credential account is unavailable") + } + } + if !account.IsOpenAIAgentIdentity() { + return nil + } + return ensureAgentIdentityTaskForAccount(ctx, s.accountRepo, nil, &s.agentIdentityTaskMu, account, account.GetCredential("task_id")) +} + +func (s *OpenAIQuotaService) isAgentIdentityAccount(ctx context.Context, accountID int64) bool { + if s == nil || s.accountRepo == nil { + return false + } + account, err := s.accountRepo.GetByID(ctx, accountID) + if err != nil || account == nil { + return false + } + if account.IsShadow() { + account, err = resolveCredentialAccount(ctx, s.accountRepo, account) + if err != nil || account == nil { + return false + } + } + return account.IsOpenAIAgentIdentity() +} + func (s *OpenAIQuotaService) buildCodexQuotaHeaders(ctx context.Context, accountID int64, accessToken, chatGPTAccountID string, fedRAMP bool) (map[string]string, error) { headers := buildCodexCommonHeaders(accessToken, chatGPTAccountID, fedRAMP) if s == nil || s.accountRepo == nil { diff --git a/backend/internal/service/openai_quota_spark_window_test.go b/backend/internal/service/openai_quota_spark_window_test.go index d2b18af05b..99bfe15ca1 100644 --- a/backend/internal/service/openai_quota_spark_window_test.go +++ b/backend/internal/service/openai_quota_spark_window_test.go @@ -38,6 +38,15 @@ func (r *stubQuotaAccountRepo) GetByID(_ context.Context, id int64) (*Account, e return acc, nil } +func (r *stubQuotaAccountRepo) UpdateCredentials(_ context.Context, id int64, credentials map[string]any) error { + acc, ok := r.accounts[id] + if !ok { + return fmt.Errorf("account %d not found", id) + } + acc.Credentials = credentials + return nil +} + // stubQuotaTokenCache 实现 OpenAITokenCache，返回预设静态 token。 type stubQuotaTokenCache struct { tokens map[string]string @@ -257,6 +266,55 @@ func TestQueryUsageAgentIdentityUsesAssertionWithoutOAuthToken(t *testing.T) { require.Equal(t, "true", fedrampHeader) } +func TestQueryUsageAgentIdentityRecoversInvalidTaskOnce(t *testing.T) { + _, privateKey, err := ed25519.GenerateKey(rand.Reader) + require.NoError(t, err) + der, err := x509.MarshalPKCS8PrivateKey(privateKey) + require.NoError(t, err) + account := &Account{ + ID: 301, + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Credentials: map[string]any{ + "auth_mode": OpenAIAuthModeAgentIdentity, + "agent_runtime_id": "runtime-quota-recovery", + "agent_private_key": base64.StdEncoding.EncodeToString(der), + "task_id": "task-quota-old", + "chatgpt_account_id": "account-quota-recovery", + }, + } + repo := &stubQuotaAccountRepo{accounts: map[int64]*Account{account.ID: account}} + usageCalls := 0 + registerCalls := 0 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("content-type", "application/json") + if strings.Contains(r.URL.Path, "/task/register") { + registerCalls++ + _, _ = w.Write([]byte(`{"task_id":"task-quota-new"}`)) + return + } + usageCalls++ + if usageCalls == 1 { + w.WriteHeader(http.StatusUnauthorized) + _, _ = w.Write([]byte(`{"error":{"code":"invalid_task_id"}}`)) + return + } + _, _ = w.Write([]byte(`{"plan_type":"pro","rate_limit":{"allowed":true}}`)) + })) + defer srv.Close() + oldBase := openAIAgentIdentityAuthAPIBaseURL + openAIAgentIdentityAuthAPIBaseURL = srv.URL + t.Cleanup(func() { openAIAgentIdentityAuthAPIBaseURL = oldBase }) + + svc := NewOpenAIQuotaService(repo, nil, nil, newQuotaRedirectingFactory(srv)) + 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")) +} + func TestParseOpenAIRateLimitResetCreditDetails_CompatibleContainers(t *testing.T) { tests := []struct { name string