diff --git a/backend/internal/handler/admin/account_data.go b/backend/internal/handler/admin/account_data.go index 322ae590..12139b51 100644 --- a/backend/internal/handler/admin/account_data.go +++ b/backend/internal/handler/admin/account_data.go @@ -267,6 +267,9 @@ func (h *AccountHandler) importData(ctx context.Context, req DataImportRequest) } } + // 收集需要异步设置隐私的 Antigravity OAuth 账号 + var privacyAccounts []*service.Account + for i := range dataPayload.Accounts { item := dataPayload.Accounts[i] if err := validateDataAccount(item); err != nil { @@ -314,7 +317,8 @@ func (h *AccountHandler) importData(ctx context.Context, req DataImportRequest) SkipDefaultGroupBind: skipDefaultGroupBind, } - if _, err := h.adminService.CreateAccount(ctx, accountInput); err != nil { + created, err := h.adminService.CreateAccount(ctx, accountInput) + if err != nil { result.AccountFailed++ result.Errors = append(result.Errors, DataImportError{ Kind: "account", @@ -323,9 +327,30 @@ func (h *AccountHandler) importData(ctx context.Context, req DataImportRequest) }) continue } + // 收集 Antigravity OAuth 账号，稍后异步设置隐私 + if created.Platform == service.PlatformAntigravity && created.Type == service.AccountTypeOAuth { + privacyAccounts = append(privacyAccounts, created) + } result.AccountCreated++ } + // 异步设置 Antigravity 隐私，避免大量导入时阻塞请求 + if len(privacyAccounts) > 0 { + adminSvc := h.adminService + go func() { + defer func() { + if r := recover(); r != nil { + slog.Error("import_antigravity_privacy_panic", "recover", r) + } + }() + bgCtx := context.Background() + for _, acc := range privacyAccounts { + adminSvc.ForceAntigravityPrivacy(bgCtx, acc) + } + slog.Info("import_antigravity_privacy_done", "count", len(privacyAccounts)) + }() + } + return result, nil } diff --git a/backend/internal/handler/admin/account_handler.go b/backend/internal/handler/admin/account_handler.go index 9eaf0bfd..6711abae 100644 --- a/backend/internal/handler/admin/account_handler.go +++ b/backend/internal/handler/admin/account_handler.go @@ -9,6 +9,7 @@ import ( "errors" "fmt" "log" + "log/slog" "net/http" "strconv" "strings" @@ -536,6 +537,8 @@ func (h *AccountHandler) Create(c *gin.Context) { if execErr != nil { return nil, execErr } + // Antigravity OAuth: 新账号直接设置隐私 + h.adminService.ForceAntigravityPrivacy(ctx, account) return h.buildAccountResponseWithRuntime(ctx, account), nil }) if err != nil { @@ -883,6 +886,8 @@ func (h *AccountHandler) refreshSingleAccount(ctx context.Context, account *serv // OpenAI OAuth: 刷新成功后检查并设置 privacy_mode h.adminService.EnsureOpenAIPrivacy(ctx, updatedAccount) + // Antigravity OAuth: 刷新成功后检查并设置 privacy_mode + h.adminService.EnsureAntigravityPrivacy(ctx, updatedAccount) return updatedAccount, "", nil } @@ -1154,6 +1159,8 @@ func (h *AccountHandler) BatchCreate(c *gin.Context) { success := 0 failed := 0 results := make([]gin.H, 0, len(req.Accounts)) + // 收集需要异步设置隐私的 Antigravity OAuth 账号 + var privacyAccounts []*service.Account for _, item := range req.Accounts { if item.RateMultiplier != nil && *item.RateMultiplier < 0 { @@ -1196,6 +1203,10 @@ func (h *AccountHandler) BatchCreate(c *gin.Context) { }) continue } + // 收集 Antigravity OAuth 账号，稍后异步设置隐私 + if account.Platform == service.PlatformAntigravity && account.Type == service.AccountTypeOAuth { + privacyAccounts = append(privacyAccounts, account) + } success++ results = append(results, gin.H{ "name": item.Name, @@ -1204,6 +1215,22 @@ func (h *AccountHandler) BatchCreate(c *gin.Context) { }) } + // 异步设置 Antigravity 隐私，避免批量创建时阻塞请求 + if len(privacyAccounts) > 0 { + adminSvc := h.adminService + go func() { + defer func() { + if r := recover(); r != nil { + slog.Error("batch_create_antigravity_privacy_panic", "recover", r) + } + }() + bgCtx := context.Background() + for _, acc := range privacyAccounts { + adminSvc.ForceAntigravityPrivacy(bgCtx, acc) + } + }() + } + return gin.H{ "success": success, "failed": failed, @@ -1869,6 +1896,42 @@ func (h *AccountHandler) GetAvailableModels(c *gin.Context) { response.Success(c, models) } +// SetPrivacy handles setting privacy for a single Antigravity OAuth account +// POST /api/v1/admin/accounts/:id/set-privacy +func (h *AccountHandler) SetPrivacy(c *gin.Context) { + accountID, err := strconv.ParseInt(c.Param("id"), 10, 64) + if err != nil { + response.BadRequest(c, "Invalid account ID") + return + } + account, err := h.adminService.GetAccount(c.Request.Context(), accountID) + if err != nil { + response.NotFound(c, "Account not found") + return + } + if account.Platform != service.PlatformAntigravity || account.Type != service.AccountTypeOAuth { + response.BadRequest(c, "Only Antigravity OAuth accounts support privacy setting") + return + } + mode := h.adminService.ForceAntigravityPrivacy(c.Request.Context(), account) + if mode == "" { + response.BadRequest(c, "Cannot set privacy: missing access_token") + return + } + // 从 DB 重新读取以确保返回最新状态 + updated, err := h.adminService.GetAccount(c.Request.Context(), accountID) + if err != nil { + // 隐私已设置成功但读取失败，回退到内存更新 + if account.Extra == nil { + account.Extra = make(map[string]any) + } + account.Extra["privacy_mode"] = mode + response.Success(c, h.buildAccountResponseWithRuntime(c.Request.Context(), account)) + return + } + response.Success(c, h.buildAccountResponseWithRuntime(c.Request.Context(), updated)) +} + // RefreshTier handles refreshing Google One tier for a single account // POST /api/v1/admin/accounts/:id/refresh-tier func (h *AccountHandler) RefreshTier(c *gin.Context) { diff --git a/backend/internal/handler/admin/admin_service_stub_test.go b/backend/internal/handler/admin/admin_service_stub_test.go index 4ed0a623..745c5610 100644 --- a/backend/internal/handler/admin/admin_service_stub_test.go +++ b/backend/internal/handler/admin/admin_service_stub_test.go @@ -445,6 +445,14 @@ func (s *stubAdminService) EnsureOpenAIPrivacy(ctx context.Context, account *ser return "" } +func (s *stubAdminService) EnsureAntigravityPrivacy(ctx context.Context, account *service.Account) string { + return "" +} + +func (s *stubAdminService) ForceAntigravityPrivacy(ctx context.Context, account *service.Account) string { + return "" +} + func (s *stubAdminService) ReplaceUserGroup(ctx context.Context, userID, oldGroupID, newGroupID int64) (*service.ReplaceUserGroupResult, error) { return &service.ReplaceUserGroupResult{MigratedKeys: 0}, nil } diff --git a/backend/internal/pkg/antigravity/client.go b/backend/internal/pkg/antigravity/client.go index f24ff5a8..fdd7fea1 100644 --- a/backend/internal/pkg/antigravity/client.go +++ b/backend/internal/pkg/antigravity/client.go @@ -78,7 +78,9 @@ type UserInfo struct { // LoadCodeAssistRequest loadCodeAssist 请求 type LoadCodeAssistRequest struct { Metadata struct { - IDEType string `json:"ideType"` + IDEType string `json:"ideType"` + IDEVersion string `json:"ideVersion"` + IDEName string `json:"ideName"` } `json:"metadata"` } @@ -223,6 +225,23 @@ func (r *LoadCodeAssistResponse) GetAvailableCredits() []AvailableCredit { return r.PaidTier.AvailableCredits } +// TierIDToPlanType 将 tier ID 映射为用户可见的套餐名。 +func TierIDToPlanType(tierID string) string { + switch strings.ToLower(strings.TrimSpace(tierID)) { + case "free-tier": + return "Free" + case "g1-pro-tier": + return "Pro" + case "g1-ultra-tier": + return "Ultra" + default: + if tierID == "" { + return "Free" + } + return tierID + } +} + // Client Antigravity API 客户端 type Client struct { httpClient *http.Client @@ -421,6 +440,8 @@ func (c *Client) GetUserInfo(ctx context.Context, accessToken string) (*UserInfo func (c *Client) LoadCodeAssist(ctx context.Context, accessToken string) (*LoadCodeAssistResponse, map[string]any, error) { reqBody := LoadCodeAssistRequest{} reqBody.Metadata.IDEType = "ANTIGRAVITY" + reqBody.Metadata.IDEVersion = "1.20.6" + reqBody.Metadata.IDEName = "antigravity" bodyBytes, err := json.Marshal(reqBody) if err != nil { @@ -704,3 +725,139 @@ func (c *Client) FetchAvailableModels(ctx context.Context, accessToken, projectI return nil, nil, lastErr } + +// ── Privacy API ────────────────────────────────────────────────────── + +// privacyBaseURL 隐私设置 API 仅使用 daily 端点（与 Antigravity 客户端行为一致） +const privacyBaseURL = antigravityDailyBaseURL + +// SetUserSettingsRequest setUserSettings 请求体 +type SetUserSettingsRequest struct { + UserSettings map[string]any `json:"user_settings"` +} + +// FetchUserInfoRequest fetchUserInfo 请求体 +type FetchUserInfoRequest struct { + Project string `json:"project"` +} + +// FetchUserInfoResponse fetchUserInfo 响应体 +type FetchUserInfoResponse struct { + UserSettings map[string]any `json:"userSettings,omitempty"` + RegionCode string `json:"regionCode,omitempty"` +} + +// IsPrivate 判断隐私是否已设置：userSettings 为空或不含 telemetryEnabled 表示已设置 +func (r *FetchUserInfoResponse) IsPrivate() bool { + if r == nil || r.UserSettings == nil { + return true + } + _, hasTelemetry := r.UserSettings["telemetryEnabled"] + return !hasTelemetry +} + +// SetUserSettingsResponse setUserSettings 响应体 +type SetUserSettingsResponse struct { + UserSettings map[string]any `json:"userSettings,omitempty"` +} + +// IsSuccess 判断 setUserSettings 是否成功：返回 {"userSettings":{}} 且无 telemetryEnabled +func (r *SetUserSettingsResponse) IsSuccess() bool { + if r == nil { + return false + } + // userSettings 为 nil 或空 map 均视为成功 + if len(r.UserSettings) == 0 { + return true + } + // 如果包含 telemetryEnabled 字段，说明未成功清除 + _, hasTelemetry := r.UserSettings["telemetryEnabled"] + return !hasTelemetry +} + +// SetUserSettings 调用 setUserSettings API 设置用户隐私，返回解析后的响应 +func (c *Client) SetUserSettings(ctx context.Context, accessToken string) (*SetUserSettingsResponse, error) { + // 发送空 user_settings 以清除隐私设置 + payload := SetUserSettingsRequest{UserSettings: map[string]any{}} + bodyBytes, err := json.Marshal(payload) + if err != nil { + return nil, fmt.Errorf("序列化请求失败: %w", err) + } + + apiURL := privacyBaseURL + "/v1internal:setUserSettings" + req, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, bytes.NewReader(bodyBytes)) + if err != nil { + return nil, fmt.Errorf("创建请求失败: %w", err) + } + req.Header.Set("Authorization", "Bearer "+accessToken) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "*/*") + req.Header.Set("User-Agent", GetUserAgent()) + req.Header.Set("X-Goog-Api-Client", "gl-node/22.21.1") + req.Host = "daily-cloudcode-pa.googleapis.com" + + resp, err := c.httpClient.Do(req) + if err != nil { + return nil, fmt.Errorf("setUserSettings 请求失败: %w", err) + } + defer func() { _ = resp.Body.Close() }() + + respBody, err := io.ReadAll(resp.Body) + if err != nil { + return nil, fmt.Errorf("读取响应失败: %w", err) + } + + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("setUserSettings 失败 (HTTP %d): %s", resp.StatusCode, string(respBody)) + } + + var result SetUserSettingsResponse + if err := json.Unmarshal(respBody, &result); err != nil { + return nil, fmt.Errorf("响应解析失败: %w", err) + } + + return &result, nil +} + +// FetchUserInfo 调用 fetchUserInfo API 获取用户隐私设置状态 +func (c *Client) FetchUserInfo(ctx context.Context, accessToken, projectID string) (*FetchUserInfoResponse, error) { + reqBody := FetchUserInfoRequest{Project: projectID} + bodyBytes, err := json.Marshal(reqBody) + if err != nil { + return nil, fmt.Errorf("序列化请求失败: %w", err) + } + + apiURL := privacyBaseURL + "/v1internal:fetchUserInfo" + req, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, bytes.NewReader(bodyBytes)) + if err != nil { + return nil, fmt.Errorf("创建请求失败: %w", err) + } + req.Header.Set("Authorization", "Bearer "+accessToken) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "*/*") + req.Header.Set("User-Agent", GetUserAgent()) + req.Header.Set("X-Goog-Api-Client", "gl-node/22.21.1") + req.Host = "daily-cloudcode-pa.googleapis.com" + + resp, err := c.httpClient.Do(req) + if err != nil { + return nil, fmt.Errorf("fetchUserInfo 请求失败: %w", err) + } + defer func() { _ = resp.Body.Close() }() + + respBody, err := io.ReadAll(resp.Body) + if err != nil { + return nil, fmt.Errorf("读取响应失败: %w", err) + } + + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("fetchUserInfo 失败 (HTTP %d): %s", resp.StatusCode, string(respBody)) + } + + var result FetchUserInfoResponse + if err := json.Unmarshal(respBody, &result); err != nil { + return nil, fmt.Errorf("响应解析失败: %w", err) + } + + return &result, nil +} diff --git a/backend/internal/pkg/antigravity/client_test.go b/backend/internal/pkg/antigravity/client_test.go index 61a08c3d..b6c2e6a5 100644 --- a/backend/internal/pkg/antigravity/client_test.go +++ b/backend/internal/pkg/antigravity/client_test.go @@ -250,6 +250,27 @@ func TestGetTier_两者都为nil(t *testing.T) { } } +func TestTierIDToPlanType(t *testing.T) { + tests := []struct { + tierID string + want string + }{ + {"free-tier", "Free"}, + {"g1-pro-tier", "Pro"}, + {"g1-ultra-tier", "Ultra"}, + {"FREE-TIER", "Free"}, + {"", "Free"}, + {"unknown-tier", "unknown-tier"}, + } + for _, tt := range tests { + t.Run(tt.tierID, func(t *testing.T) { + if got := TierIDToPlanType(tt.tierID); got != tt.want { + t.Errorf("TierIDToPlanType(%q) = %q, want %q", tt.tierID, got, tt.want) + } + }) + } +} + // --------------------------------------------------------------------------- // NewClient // --------------------------------------------------------------------------- @@ -800,6 +821,12 @@ type redirectRoundTripper struct { transport http.RoundTripper } +type roundTripperFunc func(*http.Request) (*http.Response, error) + +func (f roundTripperFunc) RoundTrip(req *http.Request) (*http.Response, error) { + return f(req) +} + func (rt *redirectRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { originalURL := req.URL.String() for prefix, target := range rt.redirects { @@ -1271,6 +1298,12 @@ func TestClient_LoadCodeAssist_Success_RealCall(t *testing.T) { if reqBody.Metadata.IDEType != "ANTIGRAVITY" { t.Errorf("IDEType 不匹配: got %s, want ANTIGRAVITY", reqBody.Metadata.IDEType) } + if strings.TrimSpace(reqBody.Metadata.IDEVersion) == "" { + t.Errorf("IDEVersion 不应为空") + } + if reqBody.Metadata.IDEName != "antigravity" { + t.Errorf("IDEName 不匹配: got %s, want antigravity", reqBody.Metadata.IDEName) + } w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusOK) diff --git a/backend/internal/server/routes/admin.go b/backend/internal/server/routes/admin.go index c4ddeab3..6fd239bb 100644 --- a/backend/internal/server/routes/admin.go +++ b/backend/internal/server/routes/admin.go @@ -257,6 +257,7 @@ func registerAccountRoutes(admin *gin.RouterGroup, h *handler.Handlers) { accounts.POST("/:id/test", h.Admin.Account.Test) accounts.POST("/:id/recover-state", h.Admin.Account.RecoverState) accounts.POST("/:id/refresh", h.Admin.Account.Refresh) + accounts.POST("/:id/set-privacy", h.Admin.Account.SetPrivacy) accounts.POST("/:id/refresh-tier", h.Admin.Account.RefreshTier) accounts.GET("/:id/stats", h.Admin.Account.GetStats) accounts.POST("/:id/clear-error", h.Admin.Account.ClearError) diff --git a/backend/internal/service/admin_service.go b/backend/internal/service/admin_service.go index ed85ee34..10f71bbc 100644 --- a/backend/internal/service/admin_service.go +++ b/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 +} diff --git a/backend/internal/service/antigravity_oauth_service.go b/backend/internal/service/antigravity_oauth_service.go index 3d5ae524..a300d898 100644 --- a/backend/internal/service/antigravity_oauth_service.go +++ b/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< 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 } diff --git a/backend/internal/service/antigravity_privacy_service.go b/backend/internal/service/antigravity_privacy_service.go new file mode 100644 index 00000000..50fe07f6 --- /dev/null +++ b/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 +} diff --git a/backend/internal/service/antigravity_privacy_service_test.go b/backend/internal/service/antigravity_privacy_service_test.go new file mode 100644 index 00000000..893500a6 --- /dev/null +++ b/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) + } +} diff --git a/backend/internal/service/antigravity_subscription_service.go b/backend/internal/service/antigravity_subscription_service.go new file mode 100644 index 00000000..04559be8 --- /dev/null +++ b/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), + } +} diff --git a/backend/internal/service/token_refresh_service.go b/backend/internal/service/token_refresh_service.go index 24b7424f..ac14aa56 100644 --- a/backend/internal/service/token_refresh_service.go +++ b/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, + ) + } +} diff --git a/frontend/src/api/admin/accounts.ts b/frontend/src/api/admin/accounts.ts index ece5a30f..fd93fe7e 100644 --- a/frontend/src/api/admin/accounts.ts +++ b/frontend/src/api/admin/accounts.ts @@ -627,6 +627,16 @@ export async function batchRefresh(accountIds: number[]): Promise { + const { data } = await apiClient.post(`/admin/accounts/${id}/set-privacy`) + return data +} + export const accountsAPI = { list, listWithEtag, @@ -663,7 +673,8 @@ export const accountsAPI = { importData, getAntigravityDefaultModelMapping, batchClearError, - batchRefresh + batchRefresh, + setPrivacy } export default accountsAPI diff --git a/frontend/src/components/admin/account/AccountActionMenu.vue b/frontend/src/components/admin/account/AccountActionMenu.vue index f5bc5aa0..e682ddaf 100644 --- a/frontend/src/components/admin/account/AccountActionMenu.vue +++ b/frontend/src/components/admin/account/AccountActionMenu.vue @@ -32,6 +32,10 @@ {{ t('admin.accounts.refreshToken') }} +