Merge pull request #1299 from DaydreamCoding/feat/antigravity-privacy-and-subscription

feat(antigravity): 自动隐私设置 + 订阅状态检测

backend/internal/service/admin_service.go

@@ -65,6 +65,10 @@ type AdminService interface {
 	SetAccountError(ctx context.Context, id int64, errorMsg string) error
 	// EnsureOpenAIPrivacy 检查 OpenAI OAuth 账号 privacy_mode，未设置则尝试关闭训练数据共享并持久化。
 	EnsureOpenAIPrivacy(ctx context.Context, account *Account) string
+	// EnsureAntigravityPrivacy 检查 Antigravity OAuth 账号 privacy_mode，未设置则调用 setUserSettings 并持久化。
+	EnsureAntigravityPrivacy(ctx context.Context, account *Account) string
+	// ForceAntigravityPrivacy 强制重新设置 Antigravity OAuth 账号隐私，无论当前状态。
+	ForceAntigravityPrivacy(ctx context.Context, account *Account) string
 	SetAccountSchedulable(ctx context.Context, id int64, schedulable bool) (*Account, error)
 	BulkUpdateAccounts(ctx context.Context, input *BulkUpdateAccountsInput) (*BulkUpdateAccountsResult, error)
 	CheckMixedChannelRisk(ctx context.Context, currentAccountID int64, currentAccountPlatform string, groupIDs []int64) error
@@ -2661,3 +2665,78 @@ func (s *adminServiceImpl) EnsureOpenAIPrivacy(ctx context.Context, account *Acc
 	_ = s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{"privacy_mode": mode})
 	return mode
 }
+
+// EnsureAntigravityPrivacy 检查 Antigravity OAuth 账号隐私状态。
+// 如果 Extra["privacy_mode"] 已存在（无论成功或失败），直接跳过。
+// 仅对从未设置过隐私的账号执行 setUserSettings + fetchUserInfo 流程。
+// 用户可通过前端 ForceAntigravityPrivacy（SetPrivacy 按钮）强制重新设置。
+func (s *adminServiceImpl) EnsureAntigravityPrivacy(ctx context.Context, account *Account) string {
+	if account.Platform != PlatformAntigravity || account.Type != AccountTypeOAuth {
+		return ""
+	}
+	// 已设置过则跳过（无论成功或失败），用户可通过 Force 手动重试
+	if account.Extra != nil {
+		if existing, ok := account.Extra["privacy_mode"].(string); ok && existing != "" {
+			return existing
+		}
+	}
+
+	token, _ := account.Credentials["access_token"].(string)
+	if token == "" {
+		return ""
+	}
+
+	projectID, _ := account.Credentials["project_id"].(string)
+
+	var proxyURL string
+	if account.ProxyID != nil {
+		if p, err := s.proxyRepo.GetByID(ctx, *account.ProxyID); err == nil && p != nil {
+			proxyURL = p.URL()
+		}
+	}
+
+	mode := setAntigravityPrivacy(ctx, token, projectID, proxyURL)
+	if mode == "" {
+		return ""
+	}
+
+	if err := s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{"privacy_mode": mode}); err != nil {
+		logger.LegacyPrintf("service.admin", "update_antigravity_privacy_mode_failed: account_id=%d err=%v", account.ID, err)
+		return mode
+	}
+	applyAntigravityPrivacyMode(account, mode)
+	return mode
+}
+
+// ForceAntigravityPrivacy 强制重新设置 Antigravity OAuth 账号隐私，无论当前状态。
+func (s *adminServiceImpl) ForceAntigravityPrivacy(ctx context.Context, account *Account) string {
+	if account.Platform != PlatformAntigravity || account.Type != AccountTypeOAuth {
+		return ""
+	}
+
+	token, _ := account.Credentials["access_token"].(string)
+	if token == "" {
+		return ""
+	}
+
+	projectID, _ := account.Credentials["project_id"].(string)
+
+	var proxyURL string
+	if account.ProxyID != nil {
+		if p, err := s.proxyRepo.GetByID(ctx, *account.ProxyID); err == nil && p != nil {
+			proxyURL = p.URL()
+		}
+	}
+
+	mode := setAntigravityPrivacy(ctx, token, projectID, proxyURL)
+	if mode == "" {
+		return ""
+	}
+
+	if err := s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{"privacy_mode": mode}); err != nil {
+		logger.LegacyPrintf("service.admin", "force_update_antigravity_privacy_mode_failed: account_id=%d err=%v", account.ID, err)
+		return mode
+	}
+	applyAntigravityPrivacyMode(account, mode)
+	return mode
+}

backend/internal/service/antigravity_oauth_service.go

@@ -89,7 +89,8 @@ type AntigravityTokenInfo struct {
 	TokenType        string `json:"token_type"`
 	Email            string `json:"email,omitempty"`
 	ProjectID        string `json:"project_id,omitempty"`
-	ProjectIDMissing bool   `json:"-"` // LoadCodeAssist 未返回 project_id
+	ProjectIDMissing bool   `json:"-"`
+	PlanType         string `json:"-"`
 }
 
 // ExchangeCode 用 authorization code 交换 token
@@ -145,13 +146,17 @@ func (s *AntigravityOAuthService) ExchangeCode(ctx context.Context, input *Antig
 		result.Email = userInfo.Email
 	}
 
-	// 获取 project_id（部分账户类型可能没有），失败时重试
-	projectID, loadErr := s.loadProjectIDWithRetry(ctx, tokenResp.AccessToken, proxyURL, 3)
+	// 获取 project_id + plan_type（部分账户类型可能没有），失败时重试
+	loadResult, loadErr := s.loadProjectIDWithRetry(ctx, tokenResp.AccessToken, proxyURL, 3)
 	if loadErr != nil {
 		fmt.Printf("[AntigravityOAuth] 警告: 获取 project_id 失败（重试后）: %v\n", loadErr)
 		result.ProjectIDMissing = true
-	} else {
-		result.ProjectID = projectID
+	}
+	if loadResult != nil {
+		result.ProjectID = loadResult.ProjectID
+		if loadResult.Subscription != nil {
+			result.PlanType = loadResult.Subscription.PlanType
+		}
 	}
 
 	return result, nil
@@ -230,13 +235,17 @@ func (s *AntigravityOAuthService) ValidateRefreshToken(ctx context.Context, refr
 		tokenInfo.Email = userInfo.Email
 	}
 
-	// 获取 project_id（容错，失败不阻塞）
-	projectID, loadErr := s.loadProjectIDWithRetry(ctx, tokenInfo.AccessToken, proxyURL, 3)
+	// 获取 project_id + plan_type（容错，失败不阻塞）
+	loadResult, loadErr := s.loadProjectIDWithRetry(ctx, tokenInfo.AccessToken, proxyURL, 3)
 	if loadErr != nil {
 		fmt.Printf("[AntigravityOAuth] 警告: 获取 project_id 失败（重试后）: %v\n", loadErr)
 		tokenInfo.ProjectIDMissing = true
-	} else {
-		tokenInfo.ProjectID = projectID
+	}
+	if loadResult != nil {
+		tokenInfo.ProjectID = loadResult.ProjectID
+		if loadResult.Subscription != nil {
+			tokenInfo.PlanType = loadResult.Subscription.PlanType
+		}
 	}
 
 	return tokenInfo, nil
@@ -288,33 +297,42 @@ func (s *AntigravityOAuthService) RefreshAccountToken(ctx context.Context, accou
 		tokenInfo.Email = existingEmail
 	}
 
-	// 每次刷新都调用 LoadCodeAssist 获取 project_id，失败时重试
+	// 每次刷新都调用 LoadCodeAssist 获取 project_id + plan_type，失败时重试
 	existingProjectID := strings.TrimSpace(account.GetCredential("project_id"))
-	projectID, loadErr := s.loadProjectIDWithRetry(ctx, tokenInfo.AccessToken, proxyURL, 3)
+	loadResult, loadErr := s.loadProjectIDWithRetry(ctx, tokenInfo.AccessToken, proxyURL, 3)
 
 	if loadErr != nil {
-		// LoadCodeAssist 失败，保留原有 project_id
 		tokenInfo.ProjectID = existingProjectID
-		// 只有从未获取过 project_id 且本次也获取失败时，才标记为真正缺失
-		// 如果之前有 project_id，本次只是临时故障，不应标记为错误
 		if existingProjectID == "" {
 			tokenInfo.ProjectIDMissing = true
 		}
-	} else {
-		tokenInfo.ProjectID = projectID
+	}
+	if loadResult != nil {
+		if loadResult.ProjectID != "" {
+			tokenInfo.ProjectID = loadResult.ProjectID
+		}
+		if loadResult.Subscription != nil {
+			tokenInfo.PlanType = loadResult.Subscription.PlanType
+		}
 	}
 
 	return tokenInfo, nil
 }
 
-// loadProjectIDWithRetry 带重试机制获取 project_id
-// 返回 project_id 和错误，失败时会重试指定次数
-func (s *AntigravityOAuthService) loadProjectIDWithRetry(ctx context.Context, accessToken, proxyURL string, maxRetries int) (string, error) {
+// loadCodeAssistResult 封装 loadProjectIDWithRetry 的返回结果，
+// 同时携带从 LoadCodeAssist 响应中提取的 plan_type 信息。
+type loadCodeAssistResult struct {
+	ProjectID    string
+	Subscription *AntigravitySubscriptionResult
+}
+
+// loadProjectIDWithRetry 带重试机制获取 project_id，同时从响应中提取 plan_type。
+func (s *AntigravityOAuthService) loadProjectIDWithRetry(ctx context.Context, accessToken, proxyURL string, maxRetries int) (*loadCodeAssistResult, error) {
 	var lastErr error
+	var lastSubscription *AntigravitySubscriptionResult
 
 	for attempt := 0; attempt <= maxRetries; attempt++ {
 		if attempt > 0 {
-			// 指数退避：1s, 2s, 4s
 			backoff := time.Duration(1<<uint(attempt-1)) * time.Second
 			if backoff > 8*time.Second {
 				backoff = 8 * time.Second
@@ -324,24 +342,34 @@ func (s *AntigravityOAuthService) loadProjectIDWithRetry(ctx context.Context, ac
 
 		client, err := antigravity.NewClient(proxyURL)
 		if err != nil {
-			return "", fmt.Errorf("create antigravity client failed: %w", err)
+			return nil, fmt.Errorf("create antigravity client failed: %w", err)
 		}
 		loadResp, loadRaw, err := client.LoadCodeAssist(ctx, accessToken)
 
+		if loadResp != nil {
+			sub := NormalizeAntigravitySubscription(loadResp)
+			lastSubscription = &sub
+		}
+
 		if err == nil && loadResp != nil && loadResp.CloudAICompanionProject != "" {
-			return loadResp.CloudAICompanionProject, nil
+			return &loadCodeAssistResult{
+				ProjectID:    loadResp.CloudAICompanionProject,
+				Subscription: lastSubscription,
+			}, nil
 		}
 
 		if err == nil {
 			if projectID, onboardErr := tryOnboardProjectID(ctx, client, accessToken, loadRaw); onboardErr == nil && projectID != "" {
-				return projectID, nil
+				return &loadCodeAssistResult{
+					ProjectID:    projectID,
+					Subscription: lastSubscription,
+				}, nil
 			} else if onboardErr != nil {
 				lastErr = onboardErr
 				continue
 			}
 		}
 
-		// 记录错误
 		if err != nil {
 			lastErr = err
 		} else if loadResp == nil {
@@ -351,7 +379,10 @@ func (s *AntigravityOAuthService) loadProjectIDWithRetry(ctx context.Context, ac
 		}
 	}
 
-	return "", fmt.Errorf("获取 project_id 失败 (重试 %d 次后): %w", maxRetries, lastErr)
+	if lastSubscription != nil {
+		return &loadCodeAssistResult{Subscription: lastSubscription}, fmt.Errorf("获取 project_id 失败 (重试 %d 次后): %w", maxRetries, lastErr)
+	}
+	return nil, fmt.Errorf("获取 project_id 失败 (重试 %d 次后): %w", maxRetries, lastErr)
 }
 
 func tryOnboardProjectID(ctx context.Context, client *antigravity.Client, accessToken string, loadRaw map[string]any) (string, error) {
@@ -410,7 +441,11 @@ func (s *AntigravityOAuthService) FillProjectID(ctx context.Context, account *Ac
 			proxyURL = proxy.URL()
 		}
 	}
-	return s.loadProjectIDWithRetry(ctx, accessToken, proxyURL, 3)
+	result, err := s.loadProjectIDWithRetry(ctx, accessToken, proxyURL, 3)
+	if result != nil {
+		return result.ProjectID, err
+	}
+	return "", err
 }
 
 // BuildAccountCredentials 构建账户凭证
@@ -431,6 +466,9 @@ func (s *AntigravityOAuthService) BuildAccountCredentials(tokenInfo *Antigravity
 	if tokenInfo.ProjectID != "" {
 		creds["project_id"] = tokenInfo.ProjectID
 	}
+	if tokenInfo.PlanType != "" {
+		creds["plan_type"] = tokenInfo.PlanType
+	}
 	return creds
 }
 

backend/internal/service/antigravity_privacy_service.go

@@ -0,0 +1,81 @@
+package service
+
+import (
+	"context"
+	"log/slog"
+	"strings"
+	"time"
+
+	"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
+)
+
+const (
+	AntigravityPrivacySet    = "privacy_set"
+	AntigravityPrivacyFailed = "privacy_set_failed"
+)
+
+// setAntigravityPrivacy 调用 Antigravity API 设置隐私并验证结果。
+// 流程：
+//  1. setUserSettings 清空设置 → 检查返回值 {"userSettings":{}}
+//  2. fetchUserInfo 二次验证隐私是否已生效（需要 project_id）
+//
+// 返回 privacy_mode 值："privacy_set" 成功，"privacy_set_failed" 失败，空串表示无法执行。
+func setAntigravityPrivacy(ctx context.Context, accessToken, projectID, proxyURL string) string {
+	if accessToken == "" {
+		return ""
+	}
+
+	ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
+	defer cancel()
+
+	client, err := antigravity.NewClient(proxyURL)
+	if err != nil {
+		slog.Warn("antigravity_privacy_client_error", "error", err.Error())
+		return AntigravityPrivacyFailed
+	}
+
+	// 第 1 步：调用 setUserSettings，检查返回值
+	setResp, err := client.SetUserSettings(ctx, accessToken)
+	if err != nil {
+		slog.Warn("antigravity_privacy_set_failed", "error", err.Error())
+		return AntigravityPrivacyFailed
+	}
+	if !setResp.IsSuccess() {
+		slog.Warn("antigravity_privacy_set_response_not_empty",
+			"user_settings", setResp.UserSettings,
+		)
+		return AntigravityPrivacyFailed
+	}
+
+	// 第 2 步：调用 fetchUserInfo 二次验证隐私是否已生效
+	if strings.TrimSpace(projectID) == "" {
+		slog.Warn("antigravity_privacy_missing_project_id")
+		return AntigravityPrivacyFailed
+	}
+	userInfo, err := client.FetchUserInfo(ctx, accessToken, projectID)
+	if err != nil {
+		slog.Warn("antigravity_privacy_verify_failed", "error", err.Error())
+		return AntigravityPrivacyFailed
+	}
+	if !userInfo.IsPrivate() {
+		slog.Warn("antigravity_privacy_verify_not_private",
+			"user_settings", userInfo.UserSettings,
+		)
+		return AntigravityPrivacyFailed
+	}
+
+	slog.Info("antigravity_privacy_set_success")
+	return AntigravityPrivacySet
+}
+
+func applyAntigravityPrivacyMode(account *Account, mode string) {
+	if account == nil || strings.TrimSpace(mode) == "" {
+		return
+	}
+	extra := make(map[string]any, len(account.Extra)+1)
+	for k, v := range account.Extra {
+		extra[k] = v
+	}
+	extra["privacy_mode"] = mode
+	account.Extra = extra
+}

backend/internal/service/antigravity_privacy_service_test.go

@@ -0,0 +1,67 @@
+//go:build unit
+
+package service
+
+import (
+	"testing"
+)
+
+func applyAntigravitySubscriptionResult(account *Account, result AntigravitySubscriptionResult) (map[string]any, map[string]any) {
+	credentials := make(map[string]any)
+	for k, v := range account.Credentials {
+		credentials[k] = v
+	}
+	credentials["plan_type"] = result.PlanType
+
+	extra := make(map[string]any)
+	for k, v := range account.Extra {
+		extra[k] = v
+	}
+	if result.SubscriptionStatus != "" {
+		extra["subscription_status"] = result.SubscriptionStatus
+	} else {
+		delete(extra, "subscription_status")
+	}
+	if result.SubscriptionError != "" {
+		extra["subscription_error"] = result.SubscriptionError
+	} else {
+		delete(extra, "subscription_error")
+	}
+	return credentials, extra
+}
+
+func TestApplyAntigravityPrivacyMode_SetsInMemoryExtra(t *testing.T) {
+	account := &Account{}
+
+	applyAntigravityPrivacyMode(account, AntigravityPrivacySet)
+
+	if account.Extra == nil {
+		t.Fatal("expected account.Extra to be initialized")
+	}
+	if got := account.Extra["privacy_mode"]; got != AntigravityPrivacySet {
+		t.Fatalf("expected privacy_mode %q, got %v", AntigravityPrivacySet, got)
+	}
+}
+
+func TestApplyAntigravityPrivacyMode_PreservedBySubscriptionResult(t *testing.T) {
+	account := &Account{
+		Credentials: map[string]any{
+			"access_token": "token",
+		},
+		Extra: map[string]any{
+			"existing": "value",
+		},
+	}
+	applyAntigravityPrivacyMode(account, AntigravityPrivacySet)
+
+	_, extra := applyAntigravitySubscriptionResult(account, AntigravitySubscriptionResult{
+		PlanType: "Pro",
+	})
+
+	if got := extra["privacy_mode"]; got != AntigravityPrivacySet {
+		t.Fatalf("expected subscription writeback to keep privacy_mode %q, got %v", AntigravityPrivacySet, got)
+	}
+	if got := extra["existing"]; got != "value" {
+		t.Fatalf("expected existing extra fields to be preserved, got %v", got)
+	}
+}

backend/internal/service/antigravity_subscription_service.go

@@ -0,0 +1,38 @@
+package service
+
+import (
+	"strings"
+
+	"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
+)
+
+const antigravitySubscriptionAbnormal = "abnormal"
+
+// AntigravitySubscriptionResult 表示订阅检测后的规范化结果。
+type AntigravitySubscriptionResult struct {
+	PlanType           string
+	SubscriptionStatus string
+	SubscriptionError  string
+}
+
+// NormalizeAntigravitySubscription 从 LoadCodeAssistResponse 提取 plan_type + 异常状态。
+// 使用 GetTier()（返回 tier ID）+ TierIDToPlanType 映射。
+func NormalizeAntigravitySubscription(resp *antigravity.LoadCodeAssistResponse) AntigravitySubscriptionResult {
+	if resp == nil {
+		return AntigravitySubscriptionResult{PlanType: "Free"}
+	}
+	if len(resp.IneligibleTiers) > 0 {
+		result := AntigravitySubscriptionResult{
+			PlanType:           "Abnormal",
+			SubscriptionStatus: antigravitySubscriptionAbnormal,
+		}
+		if resp.IneligibleTiers[0] != nil {
+			result.SubscriptionError = strings.TrimSpace(resp.IneligibleTiers[0].ReasonMessage)
+		}
+		return result
+	}
+	tierID := resp.GetTier()
+	return AntigravitySubscriptionResult{
+		PlanType: antigravity.TierIDToPlanType(tierID),
+	}
+}

backend/internal/service/token_refresh_service.go

@@ -128,7 +128,7 @@ func (s *TokenRefreshService) Start() {
 	)
 }
 
-// Stop 停止刷新服务
+// Stop 停止刷新服务（可安全多次调用）
 func (s *TokenRefreshService) Stop() {
 	close(s.stopCh)
 	s.wg.Wait()
@@ -404,6 +404,8 @@ func (s *TokenRefreshService) postRefreshActions(ctx context.Context, account *A
 	}
 	// OpenAI OAuth: 刷新成功后，检查是否已设置 privacy_mode，未设置则尝试关闭训练数据共享
 	s.ensureOpenAIPrivacy(ctx, account)
+	// Antigravity OAuth: 刷新成功后，检查是否已设置 privacy_mode，未设置则调用 setUserSettings
+	s.ensureAntigravityPrivacy(ctx, account)
 }
 
 // errRefreshSkipped 表示刷新被跳过（锁竞争或已被其他路径刷新），不计入 failed 或 refreshed
@@ -477,3 +479,50 @@ func (s *TokenRefreshService) ensureOpenAIPrivacy(ctx context.Context, account *
 		)
 	}
 }
+
+// ensureAntigravityPrivacy 后台刷新中检查 Antigravity OAuth 账号隐私状态。
+// 仅做 Extra["privacy_mode"] 存在性检查，不发起 HTTP 请求，避免每轮循环产生额外网络开销。
+// 用户可通过前端 SetPrivacy 按钮强制重新设置。
+func (s *TokenRefreshService) ensureAntigravityPrivacy(ctx context.Context, account *Account) {
+	if account.Platform != PlatformAntigravity || account.Type != AccountTypeOAuth {
+		return
+	}
+	// 已设置过（无论成功或失败）则跳过，不发 HTTP
+	if account.Extra != nil {
+		if _, ok := account.Extra["privacy_mode"]; ok {
+			return
+		}
+	}
+
+	token, _ := account.Credentials["access_token"].(string)
+	if token == "" {
+		return
+	}
+
+	projectID, _ := account.Credentials["project_id"].(string)
+
+	var proxyURL string
+	if account.ProxyID != nil && s.proxyRepo != nil {
+		if p, err := s.proxyRepo.GetByID(ctx, *account.ProxyID); err == nil && p != nil {
+			proxyURL = p.URL()
+		}
+	}
+
+	mode := setAntigravityPrivacy(ctx, token, projectID, proxyURL)
+	if mode == "" {
+		return
+	}
+
+	if err := s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{"privacy_mode": mode}); err != nil {
+		slog.Warn("token_refresh.update_antigravity_privacy_mode_failed",
+			"account_id", account.ID,
+			"error", err,
+		)
+	} else {
+		applyAntigravityPrivacyMode(account, mode)
+		slog.Info("token_refresh.antigravity_privacy_mode_set",
+			"account_id", account.ID,
+			"privacy_mode", mode,
+		)
+	}
+}