From a4a46a861875cfc4f8753e911f090de0d48204d2 Mon Sep 17 00:00:00 2001 From: Edric Li Date: Wed, 11 Feb 2026 12:52:56 +0800 Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20feat(antigravity):=20=E6=B7=BB?= =?UTF-8?q?=E5=8A=A0=20onboardUser=20=E6=94=AF=E6=8C=81=E5=B9=B6=E4=BF=AE?= =?UTF-8?q?=E5=A4=8D=20project=5Fid=20=E8=A1=A5=E9=BD=90=E9=80=BB=E8=BE=91?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 新增 OnboardUser API 客户端方法，支持账号 onboarding 获取 project_id - loadProjectIDWithRetry 增加 onboard 回退：LoadCodeAssist 未返回 project_id 时自动触发 onboarding - GetAccessToken 中 project_id 补齐改用轻量 FillProjectID 替代全量 RefreshAccountToken - 补齐逻辑增加 5 分钟冷却机制，防止频繁重试 - OnboardUser 轮询等待改为 context 感知，支持提前取消 - 提取 mergeCredentials 辅助方法消除重复代码 - 新增 extractProjectIDFromOnboardResponse 和 resolveDefaultTierID 单元测试 --- backend/internal/pkg/antigravity/client.go | 128 ++++++++++++++++++ .../internal/pkg/antigravity/client_test.go | 76 +++++++++++ .../service/antigravity_oauth_service.go | 70 +++++++++- .../service/antigravity_oauth_service_test.go | 82 +++++++++++ .../service/antigravity_token_provider.go | 51 ++++++- 5 files changed, 399 insertions(+), 8 deletions(-) create mode 100644 backend/internal/pkg/antigravity/client_test.go create mode 100644 backend/internal/service/antigravity_oauth_service_test.go diff --git a/backend/internal/pkg/antigravity/client.go b/backend/internal/pkg/antigravity/client.go index a6279b11..ac32fae5 100644 --- a/backend/internal/pkg/antigravity/client.go +++ b/backend/internal/pkg/antigravity/client.go @@ -115,6 +115,23 @@ type LoadCodeAssistResponse struct { IneligibleTiers []*IneligibleTier `json:"ineligibleTiers,omitempty"` } +// OnboardUserRequest onboardUser 请求 +type OnboardUserRequest struct { + TierID string `json:"tierId"` + Metadata struct { + IDEType string `json:"ideType"` + Platform string `json:"platform,omitempty"` + PluginType string `json:"pluginType,omitempty"` + } `json:"metadata"` +} + +// OnboardUserResponse onboardUser 响应 +type OnboardUserResponse struct { + Name string `json:"name,omitempty"` + Done bool `json:"done"` + Response map[string]any `json:"response,omitempty"` +} + // GetTier 获取账户类型 // 优先返回 paidTier（付费订阅级别），否则返回 currentTier func (r *LoadCodeAssistResponse) GetTier() string { @@ -361,6 +378,117 @@ func (c *Client) LoadCodeAssist(ctx context.Context, accessToken string) (*LoadC return nil, nil, lastErr } +// OnboardUser 触发账号 onboarding，并返回 project_id +// 说明： +// 1) 部分账号 loadCodeAssist 不会立即返回 cloudaicompanionProject； +// 2) 这时需要调用 onboardUser 完成初始化，之后才能拿到 project_id。 +func (c *Client) OnboardUser(ctx context.Context, accessToken, tierID string) (string, error) { + tierID = strings.TrimSpace(tierID) + if tierID == "" { + return "", fmt.Errorf("tier_id 为空") + } + + reqBody := OnboardUserRequest{TierID: tierID} + reqBody.Metadata.IDEType = "ANTIGRAVITY" + reqBody.Metadata.Platform = "PLATFORM_UNSPECIFIED" + reqBody.Metadata.PluginType = "GEMINI" + + bodyBytes, err := json.Marshal(reqBody) + if err != nil { + return "", fmt.Errorf("序列化请求失败: %w", err) + } + + availableURLs := BaseURLs + var lastErr error + + for urlIdx, baseURL := range availableURLs { + apiURL := baseURL + "/v1internal:onboardUser" + + for attempt := 1; attempt <= 5; attempt++ { + req, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, bytes.NewReader(bodyBytes)) + if err != nil { + lastErr = fmt.Errorf("创建请求失败: %w", err) + break + } + req.Header.Set("Authorization", "Bearer "+accessToken) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("User-Agent", UserAgent) + + resp, err := c.httpClient.Do(req) + if err != nil { + lastErr = fmt.Errorf("onboardUser 请求失败: %w", err) + if shouldFallbackToNextURL(err, 0) && urlIdx < len(availableURLs)-1 { + log.Printf("[antigravity] onboardUser URL fallback: %s -> %s", baseURL, availableURLs[urlIdx+1]) + break + } + return "", lastErr + } + + respBodyBytes, err := io.ReadAll(resp.Body) + _ = resp.Body.Close() + if err != nil { + return "", fmt.Errorf("读取响应失败: %w", err) + } + + if shouldFallbackToNextURL(nil, resp.StatusCode) && urlIdx < len(availableURLs)-1 { + log.Printf("[antigravity] onboardUser URL fallback (HTTP %d): %s -> %s", resp.StatusCode, baseURL, availableURLs[urlIdx+1]) + break + } + + if resp.StatusCode != http.StatusOK { + lastErr = fmt.Errorf("onboardUser 失败 (HTTP %d): %s", resp.StatusCode, string(respBodyBytes)) + return "", lastErr + } + + var onboardResp OnboardUserResponse + if err := json.Unmarshal(respBodyBytes, &onboardResp); err != nil { + lastErr = fmt.Errorf("onboardUser 响应解析失败: %w", err) + return "", lastErr + } + + if onboardResp.Done { + if projectID := extractProjectIDFromOnboardResponse(onboardResp.Response); projectID != "" { + DefaultURLAvailability.MarkSuccess(baseURL) + return projectID, nil + } + lastErr = fmt.Errorf("onboardUser 完成但未返回 project_id") + return "", lastErr + } + + // done=false 时等待后重试（与 CLIProxyAPI 行为一致） + select { + case <-time.After(2 * time.Second): + case <-ctx.Done(): + return "", ctx.Err() + } + } + } + + if lastErr != nil { + return "", lastErr + } + return "", fmt.Errorf("onboardUser 未返回 project_id") +} + +func extractProjectIDFromOnboardResponse(resp map[string]any) string { + if len(resp) == 0 { + return "" + } + + if v, ok := resp["cloudaicompanionProject"]; ok { + switch project := v.(type) { + case string: + return strings.TrimSpace(project) + case map[string]any: + if id, ok := project["id"].(string); ok { + return strings.TrimSpace(id) + } + } + } + + return "" +} + // ModelQuotaInfo 模型配额信息 type ModelQuotaInfo struct { RemainingFraction float64 `json:"remainingFraction"` diff --git a/backend/internal/pkg/antigravity/client_test.go b/backend/internal/pkg/antigravity/client_test.go new file mode 100644 index 00000000..ac30093d --- /dev/null +++ b/backend/internal/pkg/antigravity/client_test.go @@ -0,0 +1,76 @@ +package antigravity + +import ( + "testing" +) + +func TestExtractProjectIDFromOnboardResponse(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + resp map[string]any + want string + }{ + { + name: "nil response", + resp: nil, + want: "", + }, + { + name: "empty response", + resp: map[string]any{}, + want: "", + }, + { + name: "project as string", + resp: map[string]any{ + "cloudaicompanionProject": "my-project-123", + }, + want: "my-project-123", + }, + { + name: "project as string with spaces", + resp: map[string]any{ + "cloudaicompanionProject": " my-project-123 ", + }, + want: "my-project-123", + }, + { + name: "project as map with id", + resp: map[string]any{ + "cloudaicompanionProject": map[string]any{ + "id": "proj-from-map", + }, + }, + want: "proj-from-map", + }, + { + name: "project as map without id", + resp: map[string]any{ + "cloudaicompanionProject": map[string]any{ + "name": "some-name", + }, + }, + want: "", + }, + { + name: "missing cloudaicompanionProject key", + resp: map[string]any{ + "otherField": "value", + }, + want: "", + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + got := extractProjectIDFromOnboardResponse(tc.resp) + if got != tc.want { + t.Fatalf("extractProjectIDFromOnboardResponse() = %q, want %q", got, tc.want) + } + }) + } +} diff --git a/backend/internal/service/antigravity_oauth_service.go b/backend/internal/service/antigravity_oauth_service.go index fa8379ed..f4f0ef4c 100644 --- a/backend/internal/service/antigravity_oauth_service.go +++ b/backend/internal/service/antigravity_oauth_service.go @@ -273,12 +273,21 @@ func (s *AntigravityOAuthService) loadProjectIDWithRetry(ctx context.Context, ac } client := antigravity.NewClient(proxyURL) - loadResp, _, err := client.LoadCodeAssist(ctx, accessToken) + loadResp, loadRaw, err := client.LoadCodeAssist(ctx, accessToken) if err == nil && loadResp != nil && loadResp.CloudAICompanionProject != "" { return loadResp.CloudAICompanionProject, nil } + if err == nil { + if projectID, onboardErr := tryOnboardProjectID(ctx, client, accessToken, loadRaw); onboardErr == nil && projectID != "" { + return projectID, nil + } else if onboardErr != nil { + lastErr = onboardErr + continue + } + } + // 记录错误 if err != nil { lastErr = err @@ -292,6 +301,65 @@ func (s *AntigravityOAuthService) loadProjectIDWithRetry(ctx context.Context, ac return "", fmt.Errorf("获取 project_id 失败 (重试 %d 次后): %w", maxRetries, lastErr) } +func tryOnboardProjectID(ctx context.Context, client *antigravity.Client, accessToken string, loadRaw map[string]any) (string, error) { + tierID := resolveDefaultTierID(loadRaw) + if tierID == "" { + return "", fmt.Errorf("loadCodeAssist 未返回可用的默认 tier") + } + + projectID, err := client.OnboardUser(ctx, accessToken, tierID) + if err != nil { + return "", fmt.Errorf("onboardUser 失败 (tier=%s): %w", tierID, err) + } + return projectID, nil +} + +func resolveDefaultTierID(loadRaw map[string]any) string { + if len(loadRaw) == 0 { + return "" + } + + rawTiers, ok := loadRaw["allowedTiers"] + if !ok { + return "" + } + + tiers, ok := rawTiers.([]any) + if !ok { + return "" + } + + for _, rawTier := range tiers { + tier, ok := rawTier.(map[string]any) + if !ok { + continue + } + if isDefault, _ := tier["isDefault"].(bool); !isDefault { + continue + } + if id, ok := tier["id"].(string); ok { + id = strings.TrimSpace(id) + if id != "" { + return id + } + } + } + + return "" +} + +// FillProjectID 仅获取 project_id，不刷新 OAuth token +func (s *AntigravityOAuthService) FillProjectID(ctx context.Context, account *Account, accessToken string) (string, error) { + var proxyURL string + if account.ProxyID != nil { + proxy, err := s.proxyRepo.GetByID(ctx, *account.ProxyID) + if err == nil && proxy != nil { + proxyURL = proxy.URL() + } + } + return s.loadProjectIDWithRetry(ctx, accessToken, proxyURL, 3) +} + // BuildAccountCredentials 构建账户凭证 func (s *AntigravityOAuthService) BuildAccountCredentials(tokenInfo *AntigravityTokenInfo) map[string]any { creds := map[string]any{ diff --git a/backend/internal/service/antigravity_oauth_service_test.go b/backend/internal/service/antigravity_oauth_service_test.go new file mode 100644 index 00000000..1d2d8235 --- /dev/null +++ b/backend/internal/service/antigravity_oauth_service_test.go @@ -0,0 +1,82 @@ +package service + +import ( + "testing" +) + +func TestResolveDefaultTierID(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + loadRaw map[string]any + want string + }{ + { + name: "nil loadRaw", + loadRaw: nil, + want: "", + }, + { + name: "missing allowedTiers", + loadRaw: map[string]any{ + "paidTier": map[string]any{"id": "g1-pro-tier"}, + }, + want: "", + }, + { + name: "empty allowedTiers", + loadRaw: map[string]any{"allowedTiers": []any{}}, + want: "", + }, + { + name: "tier missing id field", + loadRaw: map[string]any{ + "allowedTiers": []any{ + map[string]any{"isDefault": true}, + }, + }, + want: "", + }, + { + name: "allowedTiers but no default", + loadRaw: map[string]any{ + "allowedTiers": []any{ + map[string]any{"id": "free-tier", "isDefault": false}, + map[string]any{"id": "standard-tier", "isDefault": false}, + }, + }, + want: "", + }, + { + name: "default tier found", + loadRaw: map[string]any{ + "allowedTiers": []any{ + map[string]any{"id": "free-tier", "isDefault": true}, + map[string]any{"id": "standard-tier", "isDefault": false}, + }, + }, + want: "free-tier", + }, + { + name: "default tier id with spaces", + loadRaw: map[string]any{ + "allowedTiers": []any{ + map[string]any{"id": " standard-tier ", "isDefault": true}, + }, + }, + want: "standard-tier", + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + got := resolveDefaultTierID(tc.loadRaw) + if got != tc.want { + t.Fatalf("resolveDefaultTierID() = %q, want %q", got, tc.want) + } + }) + } +} diff --git a/backend/internal/service/antigravity_token_provider.go b/backend/internal/service/antigravity_token_provider.go index 1eb740f9..068d6a08 100644 --- a/backend/internal/service/antigravity_token_provider.go +++ b/backend/internal/service/antigravity_token_provider.go @@ -7,12 +7,14 @@ import ( "log/slog" "strconv" "strings" + "sync" "time" ) const ( antigravityTokenRefreshSkew = 3 * time.Minute antigravityTokenCacheSkew = 5 * time.Minute + antigravityBackfillCooldown = 5 * time.Minute ) // AntigravityTokenCache Token 缓存接口（复用 GeminiTokenCache 接口定义） @@ -23,6 +25,7 @@ type AntigravityTokenProvider struct { accountRepo AccountRepository tokenCache AntigravityTokenCache antigravityOAuthService *AntigravityOAuthService + backfillCooldown sync.Map // key: int64 (account.ID) → value: time.Time } func NewAntigravityTokenProvider( @@ -93,13 +96,7 @@ func (p *AntigravityTokenProvider) GetAccessToken(ctx context.Context, account * if err != nil { return "", err } - newCredentials := p.antigravityOAuthService.BuildAccountCredentials(tokenInfo) - for k, v := range account.Credentials { - if _, exists := newCredentials[k]; !exists { - newCredentials[k] = v - } - } - account.Credentials = newCredentials + p.mergeCredentials(account, tokenInfo) if updateErr := p.accountRepo.Update(ctx, account); updateErr != nil { log.Printf("[AntigravityTokenProvider] Failed to update account credentials: %v", updateErr) } @@ -113,6 +110,21 @@ func (p *AntigravityTokenProvider) GetAccessToken(ctx context.Context, account * return "", errors.New("access_token not found in credentials") } + // 如果账号还没有 project_id，尝试在线补齐，避免请求 daily/sandbox 时出现 + // "Invalid project resource name projects/"。 + // 仅调用 loadProjectIDWithRetry，不刷新 OAuth token；带冷却机制防止频繁重试。 + if strings.TrimSpace(account.GetCredential("project_id")) == "" && p.antigravityOAuthService != nil { + if p.shouldAttemptBackfill(account.ID) { + p.markBackfillAttempted(account.ID) + if projectID, err := p.antigravityOAuthService.FillProjectID(ctx, account, accessToken); err == nil && projectID != "" { + account.Credentials["project_id"] = projectID + if updateErr := p.accountRepo.Update(ctx, account); updateErr != nil { + log.Printf("[AntigravityTokenProvider] project_id 补齐持久化失败: %v", updateErr) + } + } + } + } + // 3. 存入缓存（验证版本后再写入，避免异步刷新任务与请求线程的竞态条件） if p.tokenCache != nil { latestAccount, isStale := CheckTokenVersion(ctx, account, p.accountRepo) @@ -144,6 +156,31 @@ func (p *AntigravityTokenProvider) GetAccessToken(ctx context.Context, account * return accessToken, nil } +// mergeCredentials 将 tokenInfo 构建的凭证合并到 account 中，保留原有未覆盖的字段 +func (p *AntigravityTokenProvider) mergeCredentials(account *Account, tokenInfo *AntigravityTokenInfo) { + newCredentials := p.antigravityOAuthService.BuildAccountCredentials(tokenInfo) + for k, v := range account.Credentials { + if _, exists := newCredentials[k]; !exists { + newCredentials[k] = v + } + } + account.Credentials = newCredentials +} + +// shouldAttemptBackfill 检查是否应该尝试补齐 project_id（冷却期内不重复尝试） +func (p *AntigravityTokenProvider) shouldAttemptBackfill(accountID int64) bool { + if v, ok := p.backfillCooldown.Load(accountID); ok { + if lastAttempt, ok := v.(time.Time); ok { + return time.Since(lastAttempt) > antigravityBackfillCooldown + } + } + return true +} + +func (p *AntigravityTokenProvider) markBackfillAttempted(accountID int64) { + p.backfillCooldown.Store(accountID, time.Now()) +} + func AntigravityTokenCacheKey(account *Account) string { projectID := strings.TrimSpace(account.GetCredential("project_id")) if projectID != "" {