From 2147da682bbbb2820620dacc89e8c9e5d787541a Mon Sep 17 00:00:00 2001 From: cat Date: Wed, 15 Jul 2026 18:30:16 +0800 Subject: [PATCH] =?UTF-8?q?fix(openai):=20=E4=BF=AE=E5=A4=8D=20Agent=20Ide?= =?UTF-8?q?ntity=20=E5=AE=A1=E6=9F=A5=E9=97=AE=E9=A2=98?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../service/openai_codex_models_service.go | 3 + .../openai_codex_models_service_test.go | 34 ++++++++++++ .../internal/service/openai_quota_service.go | 34 ++++++------ .../service/openai_quota_spark_window_test.go | 55 +++++++++++++++++++ .../components/account/CreateAccountModal.vue | 5 +- .../__tests__/CreateAccountModal.spec.ts | 16 ++++++ 6 files changed, 129 insertions(+), 18 deletions(-) diff --git a/backend/internal/service/openai_codex_models_service.go b/backend/internal/service/openai_codex_models_service.go index 04f29c80a9..1ad532e9e8 100644 --- a/backend/internal/service/openai_codex_models_service.go +++ b/backend/internal/service/openai_codex_models_service.go @@ -137,6 +137,7 @@ type codexModelsManifestRequest struct { proxyURL string accountID int64 credentialAccountID int64 + credentialAccount *Account accountConcurrency int useAPIKeyUpstream bool } @@ -310,6 +311,7 @@ func (s *OpenAIGatewayService) FetchCodexModelsManifest(ctx context.Context, acc proxyURL: proxyURL, accountID: account.ID, credentialAccountID: credAccount.ID, + credentialAccount: credAccount, accountConcurrency: account.Concurrency, useAPIKeyUpstream: useAPIKeyUpstream, } @@ -438,6 +440,7 @@ func (s *OpenAIGatewayService) fetchCodexModelsManifestUpstream(ctx context.Cont } if resp.StatusCode < 200 || resp.StatusCode >= 300 { body, _ := io.ReadAll(io.LimitReader(resp.Body, 2048)) + body = s.redactAgentIdentitySensitiveBody(reqCtx, request.credentialAccount, body) message := strings.TrimSpace(string(body)) if message == "" { message = resp.Status diff --git a/backend/internal/service/openai_codex_models_service_test.go b/backend/internal/service/openai_codex_models_service_test.go index 66ed2c2b5b..8554c07258 100644 --- a/backend/internal/service/openai_codex_models_service_test.go +++ b/backend/internal/service/openai_codex_models_service_test.go @@ -3,6 +3,7 @@ package service import ( "context" "errors" + "fmt" "io" "net" "net/http" @@ -297,6 +298,39 @@ func TestFetchCodexModelsManifestAgentIdentityRecoversInvalidTaskOnce(t *testing require.Equal(t, "task-models-new", decodeAgentAssertionTask(t, assertions[1])) } +func TestFetchCodexModelsManifestAgentIdentityRedactsUpstreamErrors(t *testing.T) { + key, privateKey := newTestAgentIdentityKey(t) + account := &Account{ + ID: 5, + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Credentials: map[string]any{ + "auth_mode": OpenAIAuthModeAgentIdentity, + "agent_runtime_id": key.runtimeID, + "agent_private_key": privateKey, + "task_id": key.taskID, + "chatgpt_account_id": "acc-agent-redaction", + }, + } + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusBadRequest) + _, _ = fmt.Fprintf(w, `{"error":"%s %s %s AgentAssertion leaked"}`, key.runtimeID, key.taskID, privateKey) + })) + defer server.Close() + original := chatgptCodexModelsURL + chatgptCodexModelsURL = server.URL + t.Cleanup(func() { chatgptCodexModelsURL = original }) + + s := &OpenAIGatewayService{} + _, err := s.FetchCodexModelsManifest(context.Background(), account, "0.137.0", "") + require.Error(t, err) + require.NotContains(t, err.Error(), key.runtimeID) + require.NotContains(t, err.Error(), key.taskID) + require.NotContains(t, err.Error(), privateKey) + require.NotContains(t, err.Error(), "AgentAssertion leaked") + require.Contains(t, err.Error(), "[redacted]") +} + func TestFetchCodexModelsManifestDefaultClientVersion(t *testing.T) { var gotClientVersion string server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { diff --git a/backend/internal/service/openai_quota_service.go b/backend/internal/service/openai_quota_service.go index ffda0f4f87..c4bad643e7 100644 --- a/backend/internal/service/openai_quota_service.go +++ b/backend/internal/service/openai_quota_service.go @@ -158,7 +158,7 @@ func (s *OpenAIQuotaService) QueryUsage(ctx context.Context, accountID int64) (* var payload OpenAIQuotaUsage for recovered := false; ; { - quotaHeaders, headerErr := s.buildCodexQuotaHeaders(callCtx, accountID, accessToken, chatGPTAccountID, fedRAMP) + quotaHeaders, expectedTaskID, 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) } @@ -173,7 +173,7 @@ func (s *OpenAIQuotaService) QueryUsage(ctx context.Context, accountID int64) (* if !resp.IsSuccessState() { if agentIdentity && !recovered && isAgentIdentityTaskInvalidHTTPResponse(resp.StatusCode, []byte(resp.String())) { recovered = true - if err := s.recoverAgentIdentityTask(ctx, accountID); err != nil { + if err := s.recoverAgentIdentityTask(ctx, accountID, expectedTaskID); err != nil { return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_QUOTA_AUTH_FAILED", "agent identity task recovery failed: %v", err) } continue @@ -207,7 +207,7 @@ func (s *OpenAIQuotaService) QueryUsage(ctx context.Context, accountID int64) (* } func (s *OpenAIQuotaService) queryResetCreditDetails(ctx context.Context, client *req.Client, accessToken, chatGPTAccountID string, fedRAMP bool, accountID int64) *openAIRateLimitResetCreditDetails { - quotaHeaders, headerErr := s.buildCodexQuotaHeaders(ctx, accountID, accessToken, chatGPTAccountID, fedRAMP) + quotaHeaders, _, headerErr := s.buildCodexQuotaHeaders(ctx, accountID, accessToken, chatGPTAccountID, fedRAMP) if headerErr != nil { slog.Warn("openai_quota_reset_credit_details_auth_failed", "account_id", accountID, "error", headerErr) return nil @@ -281,7 +281,7 @@ func (s *OpenAIQuotaService) ResetCredit(ctx context.Context, accountID int64) ( var payload OpenAIQuotaResetResult for recovered := false; ; { - headers, headerErr := s.buildCodexQuotaHeaders(callCtx, accountID, accessToken, chatGPTAccountID, fedRAMP) + headers, expectedTaskID, 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) } @@ -298,7 +298,7 @@ func (s *OpenAIQuotaService) ResetCredit(ctx context.Context, accountID int64) ( if !resp.IsSuccessState() { if agentIdentity && !recovered && isAgentIdentityTaskInvalidHTTPResponse(resp.StatusCode, []byte(resp.String())) { recovered = true - if err := s.recoverAgentIdentityTask(ctx, accountID); err != nil { + if err := s.recoverAgentIdentityTask(ctx, accountID, expectedTaskID); err != nil { return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_QUOTA_AUTH_FAILED", "agent identity task recovery failed: %v", err) } continue @@ -394,7 +394,7 @@ func (s *OpenAIQuotaService) prepareUpstreamCall(ctx context.Context, accountID return accessToken, chatGPTAccountID, proxyURL, fedRAMP, nil } -func (s *OpenAIQuotaService) recoverAgentIdentityTask(ctx context.Context, accountID int64) error { +func (s *OpenAIQuotaService) recoverAgentIdentityTask(ctx context.Context, accountID int64, expectedTaskID string) error { if s == nil || s.accountRepo == nil { return fmt.Errorf("account repository is unavailable") } @@ -411,7 +411,7 @@ func (s *OpenAIQuotaService) recoverAgentIdentityTask(ctx context.Context, accou if !account.IsOpenAIAgentIdentity() { return nil } - return ensureAgentIdentityTaskForAccount(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, account, account.GetCredential("task_id")) + return ensureAgentIdentityTaskForAccount(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, account, expectedTaskID) } func (s *OpenAIQuotaService) isAgentIdentityAccount(ctx context.Context, accountID int64) bool { @@ -431,41 +431,41 @@ func (s *OpenAIQuotaService) isAgentIdentityAccount(ctx context.Context, account return account.IsOpenAIAgentIdentity() } -func (s *OpenAIQuotaService) buildCodexQuotaHeaders(ctx context.Context, accountID int64, accessToken, chatGPTAccountID string, fedRAMP bool) (map[string]string, error) { +func (s *OpenAIQuotaService) buildCodexQuotaHeaders(ctx context.Context, accountID int64, accessToken, chatGPTAccountID string, fedRAMP bool) (map[string]string, string, error) { headers := buildCodexCommonHeaders(accessToken, chatGPTAccountID, fedRAMP) if s == nil || s.accountRepo == nil { - return headers, nil + return headers, "", nil } account, err := s.accountRepo.GetByID(ctx, accountID) if err != nil || account == nil { if strings.TrimSpace(accessToken) == "" { - return nil, fmt.Errorf("agent identity account credentials are unavailable") + return nil, "", fmt.Errorf("agent identity account credentials are unavailable") } - return headers, nil + return headers, "", nil } if account.IsShadow() { if resolved, resolveErr := resolveCredentialAccount(ctx, s.accountRepo, account); resolveErr == nil && resolved != nil { account = resolved } else if strings.TrimSpace(accessToken) == "" { - return nil, fmt.Errorf("agent identity shadow credentials are unavailable") + return nil, "", fmt.Errorf("agent identity shadow credentials are unavailable") } } if !account.IsOpenAIAgentIdentity() { - return headers, nil + return headers, "", nil } if err := ensureAgentIdentityTaskForAccount(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, account, ""); err != nil { - return nil, err + return nil, "", err } key, err := agentIdentityKeyFromAccount(account) if err != nil { - return nil, err + return nil, "", err } assertion, err := buildAgentAssertion(key, time.Now()) if err != nil { - return nil, err + return nil, "", err } headers["authorization"] = assertion - return headers, nil + return headers, key.taskID, nil } func (s *OpenAIQuotaService) redactQuotaErrorBody(ctx context.Context, accountID int64, body string) string { diff --git a/backend/internal/service/openai_quota_spark_window_test.go b/backend/internal/service/openai_quota_spark_window_test.go index acce031951..ad871c3f2e 100644 --- a/backend/internal/service/openai_quota_spark_window_test.go +++ b/backend/internal/service/openai_quota_spark_window_test.go @@ -237,6 +237,61 @@ func TestResetCreditAgentIdentityUsesAssertionAndRecoversInvalidTaskOnce(t *test require.Equal(t, []int64{account.ID}, invalidator.accountIDs) } +func TestResetCreditAgentIdentityReusesConcurrentlyRecoveredTask(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: 202, + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Credentials: map[string]any{ + "auth_mode": OpenAIAuthModeAgentIdentity, + "agent_runtime_id": "runtime-reset-concurrent", + "agent_private_key": base64.StdEncoding.EncodeToString(der), + "task_id": "task-reset-old", + "chatgpt_account_id": "account-reset-concurrent", + }, + } + repo := &stubQuotaAccountRepo{accounts: map[int64]*Account{account.ID: account}} + resetCalls := 0 + registerCalls := 0 + var assertions []string + 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-reset-unexpected"}`)) + return + } + resetCalls++ + assertions = append(assertions, r.Header.Get("authorization")) + if resetCalls == 1 { + credentials := shallowCopyMap(account.Credentials) + credentials["task_id"] = "task-reset-concurrent" + account.Credentials = credentials + w.WriteHeader(http.StatusUnauthorized) + _, _ = w.Write([]byte(`{"error":{"code":"invalid_task_id"}}`)) + return + } + _, _ = w.Write([]byte(`{"code":"ok","windows_reset":1}`)) + })) + defer srv.Close() + oldBase := openAIAgentIdentityAuthAPIBaseURL + openAIAgentIdentityAuthAPIBaseURL = srv.URL + t.Cleanup(func() { openAIAgentIdentityAuthAPIBaseURL = oldBase }) + + svc := NewOpenAIQuotaService(repo, nil, nil, newQuotaRedirectingFactory(srv)) + result, err := svc.ResetCredit(context.Background(), account.ID) + require.NoError(t, err) + require.Equal(t, "ok", result.Code) + require.Equal(t, 2, resetCalls) + require.Zero(t, registerCalls) + require.Equal(t, "task-reset-old", decodeAgentAssertionTask(t, assertions[0])) + require.Equal(t, "task-reset-concurrent", decodeAgentAssertionTask(t, assertions[1])) +} + // ── Part B: prepareUpstreamCall 影子 resolve ────────────────────────────── // TestPrepareUpstreamCallShadowResolve 验证影子账号（200）QueryUsage 时: diff --git a/frontend/src/components/account/CreateAccountModal.vue b/frontend/src/components/account/CreateAccountModal.vue index e5b29be95c..7216d5b4cb 100644 --- a/frontend/src/components/account/CreateAccountModal.vue +++ b/frontend/src/components/account/CreateAccountModal.vue @@ -5515,7 +5515,10 @@ const isAgentIdentityImportContent = (content: string) => { if (Array.isArray(value)) return value.length > 0 && value.every(isAgentIdentityValue) if (!value || typeof value !== 'object') return false const record = value as Record - return record.auth_mode === 'agentIdentity' && !!record.agent_identity && typeof record.agent_identity === 'object' + const authMode = record.auth_mode ?? record.authMode + const agentIdentity = record.agent_identity ?? record.agentIdentity + return (typeof authMode === 'string' && authMode.toLowerCase() === 'agentidentity') + || (!!agentIdentity && typeof agentIdentity === 'object') } try { diff --git a/frontend/src/components/account/__tests__/CreateAccountModal.spec.ts b/frontend/src/components/account/__tests__/CreateAccountModal.spec.ts index dcd34c0902..7bfd714011 100644 --- a/frontend/src/components/account/__tests__/CreateAccountModal.spec.ts +++ b/frontend/src/components/account/__tests__/CreateAccountModal.spec.ts @@ -171,6 +171,22 @@ describe('CreateAccountModal OpenAI long-context billing', () => { expect(flow.props('titleOverride')).toBe('Agent Identity') }) + it.each([ + ['camelCase', { authMode: 'agentIdentity', agentIdentity: { agentRuntimeId: 'runtime' } }], + ['nested identity without auth_mode', { agent_identity: { agent_runtime_id: 'runtime' } }], + ])('accepts backend-compatible %s Agent Identity imports', async (_name, content) => { + const wrapper = mountModal() + await selectButtonByText(wrapper, 'OpenAI') + await wrapper.get('[data-testid="openai-account-type-agent-identity"]').trigger('click') + await wrapper.get('form#create-account-form input[type="text"]').setValue('Agent Identity') + await wrapper.get('form#create-account-form').trigger('submit.prevent') + + wrapper.getComponent(OAuthAuthorizationFlowStub).vm.$emit('import-codex-session', JSON.stringify(content)) + await flushPromises() + + expect(importCodexSessionMock).toHaveBeenCalledTimes(1) + }) + it('shows Codex PAT as a separate OpenAI account type', async () => { const wrapper = mountModal() await selectButtonByText(wrapper, 'OpenAI')