diff --git a/backend/internal/handler/admin/account_codex_agent_identity_import_test.go b/backend/internal/handler/admin/account_codex_agent_identity_import_test.go index db55182b6b..85d1a21f4a 100644 --- a/backend/internal/handler/admin/account_codex_agent_identity_import_test.go +++ b/backend/internal/handler/admin/account_codex_agent_identity_import_test.go @@ -1,6 +1,7 @@ package admin import ( + "context" "crypto/ed25519" "crypto/rand" "crypto/x509" @@ -45,3 +46,38 @@ func TestNormalizeCodexImportEntryAcceptsAgentIdentityAuthJSON(t *testing.T) { require.NotContains(t, item.Credentials, "refresh_token") require.NotEmpty(t, item.WarningTexts) } + +func TestImportCodexSessionsCreatesAgentIdentityWithoutOAuthExpiry(t *testing.T) { + _, privateKey, err := ed25519.GenerateKey(rand.Reader) + require.NoError(t, err) + der, err := x509.MarshalPKCS8PrivateKey(privateKey) + require.NoError(t, err) + + svc := newCodexImportMemoryAdminService(nil) + handler := NewAccountHandler(svc, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil) + result, err := handler.importCodexSessions(context.Background(), CodexSessionImportRequest{ + SkipDefaultGroupBind: boolPtr(true), + }, []codexImportEntry{{ + Index: 1, + Value: map[string]any{ + "auth_mode": "agentIdentity", + "agent_identity": map[string]any{ + "agent_runtime_id": "runtime-import", + "agent_private_key": base64.StdEncoding.EncodeToString(der), + "task_id": "task-import", + "account_id": "account-import", + "chatgpt_user_id": "user-import", + }, + }, + }}) + require.NoError(t, err) + require.Equal(t, 1, result.Created) + require.Zero(t, result.Failed) + require.Len(t, svc.createdAccounts, 1) + created := svc.createdAccounts[0] + require.Nil(t, created.ExpiresAt) + require.Nil(t, created.AutoPauseOnExpired) + require.Equal(t, service.OpenAIAuthModeAgentIdentity, created.Credentials["auth_mode"]) + require.NotContains(t, created.Credentials, "access_token") + require.NotContains(t, created.Credentials, "refresh_token") +} diff --git a/backend/internal/handler/admin/account_codex_import.go b/backend/internal/handler/admin/account_codex_import.go index 57b97f8476..aa471494f1 100644 --- a/backend/internal/handler/admin/account_codex_import.go +++ b/backend/internal/handler/admin/account_codex_import.go @@ -774,6 +774,12 @@ func resolveCodexImportExpiry(req CodexSessionImportRequest, item *codexImportAc if item == nil { return nil, nil, nil, nil, errors.New("导入项为空") } + // Agent Identity has no OAuth access-token lifetime. Its runtime/task + // lifecycle is handled by the upstream task recovery path, so it must not + // be rejected or auto-paused by the OAuth import expiry policy. + if item.IsAgentIdentity { + return nil, nil, nil, nil, nil + } var requestExpiresAt *time.Time if req.ExpiresAt != nil && *req.ExpiresAt > 0 { diff --git a/backend/internal/service/openai_codex_models_service.go b/backend/internal/service/openai_codex_models_service.go index 1d9356441c..1ad532e9e8 100644 --- a/backend/internal/service/openai_codex_models_service.go +++ b/backend/internal/service/openai_codex_models_service.go @@ -42,8 +42,10 @@ type CodexModelsManifest struct { } type codexModelsManifestUpstreamError struct { - err error - retryable bool + err error + retryable bool + statusCode int + body []byte } func (e *codexModelsManifestUpstreamError) Error() string { return e.err.Error() } @@ -135,6 +137,7 @@ type codexModelsManifestRequest struct { proxyURL string accountID int64 credentialAccountID int64 + credentialAccount *Account accountConcurrency int useAPIKeyUpstream bool } @@ -308,15 +311,42 @@ func (s *OpenAIGatewayService) FetchCodexModelsManifest(ctx context.Context, acc proxyURL: proxyURL, accountID: account.ID, credentialAccountID: credAccount.ID, + credentialAccount: credAccount, accountConcurrency: account.Concurrency, useAPIKeyUpstream: useAPIKeyUpstream, } if useAPIKeyUpstream { return s.fetchCachedAPIKeyCodexModelsManifest(ctx, request, ifNoneMatch) } + manifest, fetchErr := s.fetchCodexModelsManifestUpstream(ctx, request, ifNoneMatch) + if !credAccount.IsOpenAIAgentIdentity() || !isAgentIdentityTaskInvalidCodexModelsError(fetchErr) { + return manifest, fetchErr + } + expectedTaskID := strings.TrimSpace(credAccount.GetCredential("task_id")) + if recoverErr := s.recoverAgentIdentityTask(ctx, credAccount, expectedTaskID); recoverErr != nil { + return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_CODEX_MODELS_AUTH_FAILED", "agent identity task recovery failed: %v", recoverErr) + } + authHeaders, authErr := s.buildOpenAIAuthenticationHeaders(ctx, credAccount, "") + if authErr != nil { + return nil, infraerrors.Newf(http.StatusBadGateway, "OPENAI_CODEX_MODELS_AUTH_FAILED", "build Codex models authentication after task recovery: %v", authErr) + } + request.headers.Del("Authorization") + request.headers.Del("ChatGPT-Account-ID") + for key, values := range authHeaders { + for _, value := range values { + request.headers.Add(key, value) + } + } + setOpenAIChatGPTAccountHeaders(request.headers, credAccount) return s.fetchCodexModelsManifestUpstream(ctx, request, ifNoneMatch) } +func isAgentIdentityTaskInvalidCodexModelsError(err error) bool { + var upstreamErr *codexModelsManifestUpstreamError + return errors.As(err, &upstreamErr) && + isAgentIdentityTaskInvalidHTTPResponse(upstreamErr.statusCode, upstreamErr.body) +} + func (s *OpenAIGatewayService) fetchCachedAPIKeyCodexModelsManifest(ctx context.Context, request codexModelsManifestRequest, ifNoneMatch string) (*CodexModelsManifest, error) { if err := ctx.Err(); err != nil { return nil, err @@ -410,12 +440,15 @@ 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 } return nil, &codexModelsManifestUpstreamError{ - err: infraerrors.Newf(http.StatusBadGateway, "OPENAI_CODEX_MODELS_UPSTREAM_FAILED", "codex models manifest upstream error %d: %s", resp.StatusCode, message), + err: infraerrors.Newf(http.StatusBadGateway, "OPENAI_CODEX_MODELS_UPSTREAM_FAILED", "codex models manifest upstream error %d: %s", resp.StatusCode, message), + statusCode: resp.StatusCode, + body: body, retryable: resp.StatusCode == http.StatusTooManyRequests || (resp.StatusCode >= http.StatusInternalServerError && resp.StatusCode < 600), } diff --git a/backend/internal/service/openai_codex_models_service_test.go b/backend/internal/service/openai_codex_models_service_test.go index 3f4ec96a2f..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" @@ -17,6 +18,7 @@ import ( "github.com/Wei-Shaw/sub2api/internal/config" infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" "github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint" + "github.com/stretchr/testify/require" "golang.org/x/net/http2" ) @@ -242,6 +244,93 @@ func TestFetchCodexModelsManifestAgentIdentityUsesAssertionWithoutOAuthToken(t * } } +func TestFetchCodexModelsManifestAgentIdentityRecoversInvalidTaskOnce(t *testing.T) { + key, privateKey := newTestAgentIdentityKey(t) + account := &Account{ + ID: 4, + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Credentials: map[string]any{ + "auth_mode": OpenAIAuthModeAgentIdentity, + "agent_runtime_id": key.runtimeID, + "agent_private_key": privateKey, + "task_id": "task-models-old", + "chatgpt_account_id": "acc-agent-recovery", + }, + } + repo := &stubQuotaAccountRepo{accounts: map[int64]*Account{account.ID: account}} + modelsCalls := 0 + registerCalls := 0 + var assertions []string + server := 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-models-new"}`)) + return + } + modelsCalls++ + assertions = append(assertions, r.Header.Get("Authorization")) + if modelsCalls == 1 { + w.WriteHeader(http.StatusUnauthorized) + _, _ = w.Write([]byte(`{"error":{"code":"invalid_task_id"}}`)) + return + } + _, _ = w.Write([]byte(`{"models":[]}`)) + })) + defer server.Close() + + originalModelsURL := chatgptCodexModelsURL + chatgptCodexModelsURL = server.URL + t.Cleanup(func() { chatgptCodexModelsURL = originalModelsURL }) + originalAuthBase := openAIAgentIdentityAuthAPIBaseURL + openAIAgentIdentityAuthAPIBaseURL = server.URL + t.Cleanup(func() { openAIAgentIdentityAuthAPIBaseURL = originalAuthBase }) + + s := &OpenAIGatewayService{accountRepo: repo} + manifest, err := s.FetchCodexModelsManifest(context.Background(), account, "0.137.0", "") + require.NoError(t, err) + require.Equal(t, `{"models":[]}`, string(manifest.Body)) + require.Equal(t, 2, modelsCalls) + require.Equal(t, 1, registerCalls) + require.Len(t, assertions, 2) + require.Equal(t, "task-models-old", decodeAgentAssertionTask(t, assertions[0])) + 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 a94d5fe3f8..c4bad643e7 100644 --- a/backend/internal/service/openai_quota_service.go +++ b/backend/internal/service/openai_quota_service.go @@ -22,8 +22,6 @@ 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" @@ -160,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) } @@ -175,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 @@ -209,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 @@ -260,9 +258,6 @@ 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) @@ -282,26 +277,38 @@ 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 - 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() { - 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, 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) + } + 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, expectedTaskID); 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", @@ -387,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") } @@ -404,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 { @@ -424,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 c521b149e9..ad871c3f2e 100644 --- a/backend/internal/service/openai_quota_spark_window_test.go +++ b/backend/internal/service/openai_quota_spark_window_test.go @@ -176,21 +176,120 @@ func TestResetCreditShadowRejected(t *testing.T) { "shadow ResetCredit 应映射为 409 Conflict 而非 500") } -func TestResetCreditAgentIdentityRejectedBeforeUpstream(t *testing.T) { +func TestResetCreditAgentIdentityUsesAssertionAndRecoversInvalidTaskOnce(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: 201, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Credentials: map[string]any{ - "auth_mode": OpenAIAuthModeAgentIdentity, + "auth_mode": OpenAIAuthModeAgentIdentity, + "agent_runtime_id": "runtime-reset-recovery", + "agent_private_key": base64.StdEncoding.EncodeToString(der), + "task_id": "task-reset-old", + "chatgpt_account_id": "account-reset-recovery", }, } repo := &stubQuotaAccountRepo{accounts: map[int64]*Account{account.ID: account}} - svc := &OpenAIQuotaService{accountRepo: repo} + 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-new"}`)) + return + } + resetCalls++ + assertions = append(assertions, r.Header.Get("authorization")) + require.Equal(t, "account-reset-recovery", r.Header.Get("chatgpt-account-id")) + if resetCalls == 1 { + w.WriteHeader(http.StatusUnauthorized) + _, _ = w.Write([]byte(`{"error":{"code":"invalid_task_id"}}`)) + return + } + _, _ = w.Write([]byte(`{"code":"ok","windows_reset":2}`)) + })) + defer srv.Close() + oldBase := openAIAgentIdentityAuthAPIBaseURL + openAIAgentIdentityAuthAPIBaseURL = srv.URL + t.Cleanup(func() { openAIAgentIdentityAuthAPIBaseURL = oldBase }) - _, err := svc.ResetCredit(context.Background(), account.ID) - require.ErrorIs(t, err, ErrAgentIdentityResetNotSupported) - require.Equal(t, http.StatusConflict, infraerrors.Code(err)) + invalidator := &agentIdentityWSInvalidationRecorder{} + svc := NewOpenAIQuotaService(repo, nil, nil, newQuotaRedirectingFactory(srv)) + svc.agentIdentityWS = invalidator + + result, err := svc.ResetCredit(context.Background(), account.ID) + require.NoError(t, err) + require.Equal(t, "ok", result.Code) + require.Equal(t, 2, result.WindowsReset) + require.Equal(t, 2, resetCalls) + require.Equal(t, 1, registerCalls) + require.Len(t, assertions, 2) + require.True(t, strings.HasPrefix(assertions[0], "AgentAssertion ")) + require.True(t, strings.HasPrefix(assertions[1], "AgentAssertion ")) + require.NotEqual(t, assertions[0], assertions[1]) + require.Equal(t, "task-reset-new", account.GetCredential("task_id")) + 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 ────────────────────────────── diff --git a/frontend/src/components/account/CreateAccountModal.vue b/frontend/src/components/account/CreateAccountModal.vue index 206c1f1494..b3dedacc7f 100644 --- a/frontend/src/components/account/CreateAccountModal.vue +++ b/frontend/src/components/account/CreateAccountModal.vue @@ -3117,6 +3117,7 @@ :show-session-token-option="false" :show-access-token-option="false" :show-codex-session-import-option="form.platform === 'openai'" + :show-agent-identity-option="form.platform === 'openai'" :show-codex-pat-option="form.platform === 'openai'" :show-sso-option="form.platform === 'grok'" :show-manual-option="true" @@ -5445,6 +5446,30 @@ const formatCodexImportMessages = (messages?: CodexSessionImportMessage[]) => { .join('\n') } +const isAgentIdentityImportContent = (content: string) => { + const isAgentIdentityValue = (value: unknown): boolean => { + if (Array.isArray(value)) return value.length > 0 && value.every(isAgentIdentityValue) + if (!value || typeof value !== 'object') return false + const record = value as Record + 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 { + return isAgentIdentityValue(JSON.parse(content)) + } catch { + const lines = content.split('\n').map((line) => line.trim()).filter(Boolean) + if (lines.length === 0) return false + try { + return lines.every((line) => isAgentIdentityValue(JSON.parse(line))) + } catch { + return false + } + } +} + const handleOpenAIImportCodexSession = async (content: string) => { const oauthClient = openaiOAuth const trimmed = content.trim() @@ -5452,6 +5477,10 @@ const handleOpenAIImportCodexSession = async (content: string) => { oauthClient.error.value = t('admin.accounts.oauth.openai.codexSessionEmpty') return } + if (oauthFlowRef.value?.inputMethod === 'agent_identity' && !isAgentIdentityImportContent(trimmed)) { + oauthClient.error.value = t('admin.accounts.oauth.openai.agentIdentityInvalid') + return + } const credentialExtras = buildOpenAICodexImportCredentialExtras() if (credentialExtras === null) { diff --git a/frontend/src/components/account/OAuthAuthorizationFlow.vue b/frontend/src/components/account/OAuthAuthorizationFlow.vue index 395b4ae923..f45aeeb5d7 100644 --- a/frontend/src/components/account/OAuthAuthorizationFlow.vue +++ b/frontend/src/components/account/OAuthAuthorizationFlow.vue @@ -103,6 +103,17 @@ t('admin.accounts.oauth.openai.codexSessionAuth') }} +