Commit c471951a authored by weak-fox's avatar weak-fox Committed by 陈曦
Browse files

Retry OpenAI privacy opt-out after failed states

parent b22ab6ad
......@@ -2639,10 +2639,8 @@ func (s *adminServiceImpl) EnsureOpenAIPrivacy(ctx context.Context, account *Acc
if s.privacyClientFactory == nil {
return ""
}
if account.Extra != nil {
if _, ok := account.Extra["privacy_mode"]; ok {
return ""
}
if shouldSkipOpenAIPrivacyEnsure(account.Extra) {
return ""
}
token, _ := account.Credentials["access_token"].(string)
......
//go:build unit
package service
import (
"context"
"errors"
"testing"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/imroc/req/v3"
"github.com/stretchr/testify/require"
)
func TestAdminService_EnsureOpenAIPrivacy_RetriesNonSuccessModes(t *testing.T) {
t.Parallel()
for _, mode := range []string{PrivacyModeFailed, PrivacyModeCFBlocked} {
t.Run(mode, func(t *testing.T) {
t.Parallel()
privacyCalls := 0
svc := &adminServiceImpl{
accountRepo: &mockAccountRepoForGemini{},
privacyClientFactory: func(proxyURL string) (*req.Client, error) {
privacyCalls++
return nil, errors.New("factory failed")
},
}
account := &Account{
ID: 101,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Credentials: map[string]any{
"access_token": "token-1",
},
Extra: map[string]any{
"privacy_mode": mode,
},
}
got := svc.EnsureOpenAIPrivacy(context.Background(), account)
require.Equal(t, PrivacyModeFailed, got)
require.Equal(t, 1, privacyCalls)
})
}
}
func TestTokenRefreshService_ensureOpenAIPrivacy_RetriesNonSuccessModes(t *testing.T) {
t.Parallel()
cfg := &config.Config{
TokenRefresh: config.TokenRefreshConfig{
MaxRetries: 1,
RetryBackoffSeconds: 0,
},
}
for _, mode := range []string{PrivacyModeFailed, PrivacyModeCFBlocked} {
t.Run(mode, func(t *testing.T) {
t.Parallel()
service := NewTokenRefreshService(&tokenRefreshAccountRepo{}, nil, nil, nil, nil, nil, nil, cfg, nil)
privacyCalls := 0
service.SetPrivacyDeps(func(proxyURL string) (*req.Client, error) {
privacyCalls++
return nil, errors.New("factory failed")
}, nil)
account := &Account{
ID: 202,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Credentials: map[string]any{
"access_token": "token-2",
},
Extra: map[string]any{
"privacy_mode": mode,
},
}
service.ensureOpenAIPrivacy(context.Background(), account)
require.Equal(t, 1, privacyCalls)
})
}
}
......@@ -22,6 +22,19 @@ const (
PrivacyModeCFBlocked = "training_set_cf_blocked"
)
func shouldSkipOpenAIPrivacyEnsure(extra map[string]any) bool {
if extra == nil {
return false
}
raw, ok := extra["privacy_mode"]
if !ok {
return false
}
mode, _ := raw.(string)
mode = strings.TrimSpace(mode)
return mode != PrivacyModeFailed && mode != PrivacyModeCFBlocked
}
// disableOpenAITraining calls ChatGPT settings API to turn off "Improve the model for everyone".
// Returns privacy_mode value: "training_off" on success, "cf_blocked" / "failed" on failure.
func disableOpenAITraining(ctx context.Context, clientFactory PrivacyClientFactory, accessToken, proxyURL string) string {
......
......@@ -443,11 +443,8 @@ func (s *TokenRefreshService) ensureOpenAIPrivacy(ctx context.Context, account *
if s.privacyClientFactory == nil {
return
}
// 已设置过则跳过
if account.Extra != nil {
if _, ok := account.Extra["privacy_mode"]; ok {
return
}
if shouldSkipOpenAIPrivacyEnsure(account.Extra) {
return
}
token, _ := account.Credentials["access_token"].(string)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment