Skip to content

Commit f6fd7c8

Browse files
author
QTom
committed
feat(antigravity): 从 LoadCodeAssist 复用 TierInfo 提取 plan_type
复用已有 GetTier() 返回的 tier ID（free-tier / g1-pro-tier / g1-ultra-tier），通过 TierIDToPlanType 映射为 Free / Pro / Ultra， 在 loadProjectIDWithRetry 中顺带提取并写入 credentials.plan_type； 前端增加 Abnormal 异常套餐红色标记。 Made-with: Cursor
1 parent c2965c0 commit f6fd7c8

8 files changed

Lines changed: 216 additions & 28 deletions

File tree

‎backend/internal/pkg/antigravity/client.go‎

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -78,7 +78,9 @@ type UserInfo struct {
7878
// LoadCodeAssistRequest loadCodeAssist 请求
7979
type LoadCodeAssistRequest struct {
8080
Metadata struct {
81-
IDEType string `json:"ideType"`
81+
IDEType string `json:"ideType"`
82+
IDEVersion string `json:"ideVersion"`
83+
IDEName string `json:"ideName"`
8284
} `json:"metadata"`
8385
}
8486

@@ -223,6 +225,23 @@ func (r *LoadCodeAssistResponse) GetAvailableCredits() []AvailableCredit {
223225
return r.PaidTier.AvailableCredits
224226
}
225227

228+
// TierIDToPlanType 将 tier ID 映射为用户可见的套餐名。
229+
func TierIDToPlanType(tierID string) string {
230+
switch strings.ToLower(strings.TrimSpace(tierID)) {
231+
case "free-tier":
232+
return "Free"
233+
case "g1-pro-tier":
234+
return "Pro"
235+
case "g1-ultra-tier":
236+
return "Ultra"
237+
default:
238+
if tierID == "" {
239+
return "Free"
240+
}
241+
return tierID
242+
}
243+
}
244+
226245
// Client Antigravity API 客户端
227246
type Client struct {
228247
httpClient *http.Client
@@ -421,6 +440,8 @@ func (c *Client) GetUserInfo(ctx context.Context, accessToken string) (*UserInfo
421440
func (c *Client) LoadCodeAssist(ctx context.Context, accessToken string) (*LoadCodeAssistResponse, map[string]any, error) {
422441
reqBody := LoadCodeAssistRequest{}
423442
reqBody.Metadata.IDEType = "ANTIGRAVITY"
443+
reqBody.Metadata.IDEVersion = "1.20.6"
444+
reqBody.Metadata.IDEName = "antigravity"
424445

425446
bodyBytes, err := json.Marshal(reqBody)
426447
if err != nil {

‎backend/internal/pkg/antigravity/client_test.go‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -250,6 +250,27 @@ func TestGetTier_两者都为nil(t *testing.T) {
250250
}
251251
}
252252

253+
func TestTierIDToPlanType(t *testing.T) {
254+
tests := []struct {
255+
tierID string
256+
want string
257+
}{
258+
{"free-tier", "Free"},
259+
{"g1-pro-tier", "Pro"},
260+
{"g1-ultra-tier", "Ultra"},
261+
{"FREE-TIER", "Free"},
262+
{"", "Free"},
263+
{"unknown-tier", "unknown-tier"},
264+
}
265+
for _, tt := range tests {
266+
t.Run(tt.tierID, func(t *testing.T) {
267+
if got := TierIDToPlanType(tt.tierID); got != tt.want {
268+
t.Errorf("TierIDToPlanType(%q) = %q, want %q", tt.tierID, got, tt.want)
269+
}
270+
})
271+
}
272+
}
273+
253274
// ---------------------------------------------------------------------------
254275
// NewClient
255276
// ---------------------------------------------------------------------------
@@ -800,6 +821,12 @@ type redirectRoundTripper struct {
800821
transport http.RoundTripper
801822
}
802823

824+
type roundTripperFunc func(*http.Request) (*http.Response, error)
825+
826+
func (f roundTripperFunc) RoundTrip(req *http.Request) (*http.Response, error) {
827+
return f(req)
828+
}
829+
803830
func (rt *redirectRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
804831
originalURL := req.URL.String()
805832
for prefix, target := range rt.redirects {
@@ -1271,6 +1298,12 @@ func TestClient_LoadCodeAssist_Success_RealCall(t *testing.T) {
12711298
if reqBody.Metadata.IDEType != "ANTIGRAVITY" {
12721299
t.Errorf("IDEType 不匹配: got %s, want ANTIGRAVITY", reqBody.Metadata.IDEType)
12731300
}
1301+
if strings.TrimSpace(reqBody.Metadata.IDEVersion) == "" {
1302+
t.Errorf("IDEVersion 不应为空")
1303+
}
1304+
if reqBody.Metadata.IDEName != "antigravity" {
1305+
t.Errorf("IDEName 不匹配: got %s, want antigravity", reqBody.Metadata.IDEName)
1306+
}
12741307

12751308
w.Header().Set("Content-Type", "application/json")
12761309
w.WriteHeader(http.StatusOK)

‎backend/internal/service/antigravity_oauth_service.go‎

Lines changed: 64 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -89,7 +89,8 @@ type AntigravityTokenInfo struct {
8989
TokenType string `json:"token_type"`
9090
Email string `json:"email,omitempty"`
9191
ProjectID string `json:"project_id,omitempty"`
92-
ProjectIDMissing bool `json:"-"` // LoadCodeAssist 未返回 project_id
92+
ProjectIDMissing bool `json:"-"`
93+
PlanType string `json:"-"`
9394
}
9495

9596
// ExchangeCode 用 authorization code 交换 token
@@ -145,13 +146,17 @@ func (s *AntigravityOAuthService) ExchangeCode(ctx context.Context, input *Antig
145146
result.Email = userInfo.Email
146147
}
147148

148-
// 获取 project_id（部分账户类型可能没有），失败时重试
149-
projectID, loadErr := s.loadProjectIDWithRetry(ctx, tokenResp.AccessToken, proxyURL, 3)
149+
// 获取 project_id + plan_type（部分账户类型可能没有），失败时重试
150+
loadResult, loadErr := s.loadProjectIDWithRetry(ctx, tokenResp.AccessToken, proxyURL, 3)
150151
if loadErr != nil {
151152
fmt.Printf("[AntigravityOAuth] 警告: 获取 project_id 失败（重试后）: %v\n", loadErr)
152153
result.ProjectIDMissing = true
153-
} else {
154-
result.ProjectID = projectID
154+
}
155+
if loadResult != nil {
156+
result.ProjectID = loadResult.ProjectID
157+
if loadResult.Subscription != nil {
158+
result.PlanType = loadResult.Subscription.PlanType
159+
}
155160
}
156161

157162
return result, nil
@@ -230,13 +235,17 @@ func (s *AntigravityOAuthService) ValidateRefreshToken(ctx context.Context, refr
230235
tokenInfo.Email = userInfo.Email
231236
}
232237

233-
// 获取 project_id（容错，失败不阻塞）
234-
projectID, loadErr := s.loadProjectIDWithRetry(ctx, tokenInfo.AccessToken, proxyURL, 3)
238+
// 获取 project_id + plan_type（容错，失败不阻塞）
239+
loadResult, loadErr := s.loadProjectIDWithRetry(ctx, tokenInfo.AccessToken, proxyURL, 3)
235240
if loadErr != nil {
236241
fmt.Printf("[AntigravityOAuth] 警告: 获取 project_id 失败（重试后）: %v\n", loadErr)
237242
tokenInfo.ProjectIDMissing = true
238-
} else {
239-
tokenInfo.ProjectID = projectID
243+
}
244+
if loadResult != nil {
245+
tokenInfo.ProjectID = loadResult.ProjectID
246+
if loadResult.Subscription != nil {
247+
tokenInfo.PlanType = loadResult.Subscription.PlanType
248+
}
240249
}
241250

242251
return tokenInfo, nil
@@ -288,33 +297,42 @@ func (s *AntigravityOAuthService) RefreshAccountToken(ctx context.Context, accou
288297
tokenInfo.Email = existingEmail
289298
}
290299

291-
// 每次刷新都调用 LoadCodeAssist 获取 project_id，失败时重试
300+
// 每次刷新都调用 LoadCodeAssist 获取 project_id + plan_type，失败时重试
292301
existingProjectID := strings.TrimSpace(account.GetCredential("project_id"))
293-
projectID, loadErr := s.loadProjectIDWithRetry(ctx, tokenInfo.AccessToken, proxyURL, 3)
302+
loadResult, loadErr := s.loadProjectIDWithRetry(ctx, tokenInfo.AccessToken, proxyURL, 3)
294303

295304
if loadErr != nil {
296-
// LoadCodeAssist 失败，保留原有 project_id
297305
tokenInfo.ProjectID = existingProjectID
298-
// 只有从未获取过 project_id 且本次也获取失败时，才标记为真正缺失
299-
// 如果之前有 project_id，本次只是临时故障，不应标记为错误
300306
if existingProjectID == "" {
301307
tokenInfo.ProjectIDMissing = true
302308
}
303-
} else {
304-
tokenInfo.ProjectID = projectID
309+
}
310+
if loadResult != nil {
311+
if loadResult.ProjectID != "" {
312+
tokenInfo.ProjectID = loadResult.ProjectID
313+
}
314+
if loadResult.Subscription != nil {
315+
tokenInfo.PlanType = loadResult.Subscription.PlanType
316+
}
305317
}
306318

307319
return tokenInfo, nil
308320
}
309321

310-
// loadProjectIDWithRetry 带重试机制获取 project_id
311-
// 返回 project_id 和错误，失败时会重试指定次数
312-
func (s *AntigravityOAuthService) loadProjectIDWithRetry(ctx context.Context, accessToken, proxyURL string, maxRetries int) (string, error) {
322+
// loadCodeAssistResult 封装 loadProjectIDWithRetry 的返回结果，
323+
// 同时携带从 LoadCodeAssist 响应中提取的 plan_type 信息。
324+
type loadCodeAssistResult struct {
325+
ProjectID string
326+
Subscription *AntigravitySubscriptionResult
327+
}
328+
329+
// loadProjectIDWithRetry 带重试机制获取 project_id，同时从响应中提取 plan_type。
330+
func (s *AntigravityOAuthService) loadProjectIDWithRetry(ctx context.Context, accessToken, proxyURL string, maxRetries int) (*loadCodeAssistResult, error) {
313331
var lastErr error
332+
var lastSubscription *AntigravitySubscriptionResult
314333

315334
for attempt := 0; attempt <= maxRetries; attempt++ {
316335
if attempt > 0 {
317-
// 指数退避：1s, 2s, 4s
318336
backoff := time.Duration(1<<uint(attempt-1)) * time.Second
319337
if backoff > 8*time.Second {
320338
backoff = 8 * time.Second
@@ -324,24 +342,34 @@ func (s *AntigravityOAuthService) loadProjectIDWithRetry(ctx context.Context, ac
324342

325343
client, err := antigravity.NewClient(proxyURL)
326344
if err != nil {
327-
return "", fmt.Errorf("create antigravity client failed: %w", err)
345+
return nil, fmt.Errorf("create antigravity client failed: %w", err)
328346
}
329347
loadResp, loadRaw, err := client.LoadCodeAssist(ctx, accessToken)
330348

349+
if loadResp != nil {
350+
sub := NormalizeAntigravitySubscription(loadResp)
351+
lastSubscription = &sub
352+
}
353+
331354
if err == nil && loadResp != nil && loadResp.CloudAICompanionProject != "" {
332-
return loadResp.CloudAICompanionProject, nil
355+
return &loadCodeAssistResult{
356+
ProjectID: loadResp.CloudAICompanionProject,
357+
Subscription: lastSubscription,
358+
}, nil
333359
}
334360

335361
if err == nil {
336362
if projectID, onboardErr := tryOnboardProjectID(ctx, client, accessToken, loadRaw); onboardErr == nil && projectID != "" {
337-
return projectID, nil
363+
return &loadCodeAssistResult{
364+
ProjectID: projectID,
365+
Subscription: lastSubscription,
366+
}, nil
338367
} else if onboardErr != nil {
339368
lastErr = onboardErr
340369
continue
341370
}
342371
}
343372

344-
// 记录错误
345373
if err != nil {
346374
lastErr = err
347375
} else if loadResp == nil {
@@ -351,7 +379,10 @@ func (s *AntigravityOAuthService) loadProjectIDWithRetry(ctx context.Context, ac
351379
}
352380
}
353381

354-
return "", fmt.Errorf("获取 project_id 失败 (重试 %d 次后): %w", maxRetries, lastErr)
382+
if lastSubscription != nil {
383+
return &loadCodeAssistResult{Subscription: lastSubscription}, fmt.Errorf("获取 project_id 失败 (重试 %d 次后): %w", maxRetries, lastErr)
384+
}
385+
return nil, fmt.Errorf("获取 project_id 失败 (重试 %d 次后): %w", maxRetries, lastErr)
355386
}
356387

357388
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
410441
proxyURL = proxy.URL()
411442
}
412443
}
413-
return s.loadProjectIDWithRetry(ctx, accessToken, proxyURL, 3)
444+
result, err := s.loadProjectIDWithRetry(ctx, accessToken, proxyURL, 3)
445+
if result != nil {
446+
return result.ProjectID, err
447+
}
448+
return "", err
414449
}
415450

416451
// BuildAccountCredentials 构建账户凭证
@@ -431,6 +466,9 @@ func (s *AntigravityOAuthService) BuildAccountCredentials(tokenInfo *Antigravity
431466
if tokenInfo.ProjectID != "" {
432467
creds["project_id"] = tokenInfo.ProjectID
433468
}
469+
if tokenInfo.PlanType != "" {
470+
creds["plan_type"] = tokenInfo.PlanType
471+
}
434472
return creds
435473
}
436474

‎backend/internal/service/antigravity_privacy_service_test.go‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,3 +16,26 @@ func TestApplyAntigravityPrivacyMode_SetsInMemoryExtra(t *testing.T) {
1616
t.Fatalf("expected privacy_mode %q, got %v", AntigravityPrivacySet, got)
1717
}
1818
}
19+
20+
func TestApplyAntigravityPrivacyMode_PreservedBySubscriptionResult(t *testing.T) {
21+
account := &Account{
22+
Credentials: map[string]any{
23+
"access_token": "token",
24+
},
25+
Extra: map[string]any{
26+
"existing": "value",
27+
},
28+
}
29+
applyAntigravityPrivacyMode(account, AntigravityPrivacySet)
30+
31+
_, extra := applyAntigravitySubscriptionResult(account, AntigravitySubscriptionResult{
32+
PlanType: "Pro",
33+
})
34+
35+
if got := extra["privacy_mode"]; got != AntigravityPrivacySet {
36+
t.Fatalf("expected subscription writeback to keep privacy_mode %q, got %v", AntigravityPrivacySet, got)
37+
}
38+
if got := extra["existing"]; got != "value" {
39+
t.Fatalf("expected existing extra fields to be preserved, got %v", got)
40+
}
41+
}

‎backend/internal/service/antigravity_subscription_service.go‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
package service
2+
3+
import (
4+
"strings"
5+
6+
"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
7+
)
8+
9+
const antigravitySubscriptionAbnormal = "abnormal"
10+
11+
// AntigravitySubscriptionResult 表示订阅检测后的规范化结果。
12+
type AntigravitySubscriptionResult struct {
13+
PlanType string
14+
SubscriptionStatus string
15+
SubscriptionError string
16+
}
17+
18+
// NormalizeAntigravitySubscription 从 LoadCodeAssistResponse 提取 plan_type + 异常状态。
19+
// 使用 GetTier()（返回 tier ID）+ TierIDToPlanType 映射。
20+
func NormalizeAntigravitySubscription(resp *antigravity.LoadCodeAssistResponse) AntigravitySubscriptionResult {
21+
if resp == nil {
22+
return AntigravitySubscriptionResult{PlanType: "Free"}
23+
}
24+
if len(resp.IneligibleTiers) > 0 {
25+
result := AntigravitySubscriptionResult{
26+
PlanType: "Abnormal",
27+
SubscriptionStatus: antigravitySubscriptionAbnormal,
28+
}
29+
if resp.IneligibleTiers[0] != nil {
30+
result.SubscriptionError = strings.TrimSpace(resp.IneligibleTiers[0].ReasonMessage)
31+
}
32+
return result
33+
}
34+
tierID := resp.GetTier()
35+
return AntigravitySubscriptionResult{
36+
PlanType: antigravity.TierIDToPlanType(tierID),
37+
}
38+
}
39+
40+
func applyAntigravitySubscriptionResult(account *Account, result AntigravitySubscriptionResult) (map[string]any, map[string]any) {
41+
credentials := make(map[string]any)
42+
for k, v := range account.Credentials {
43+
credentials[k] = v
44+
}
45+
credentials["plan_type"] = result.PlanType
46+
47+
extra := make(map[string]any)
48+
for k, v := range account.Extra {
49+
extra[k] = v
50+
}
51+
if result.SubscriptionStatus != "" {
52+
extra["subscription_status"] = result.SubscriptionStatus
53+
} else {
54+
delete(extra, "subscription_status")
55+
}
56+
if result.SubscriptionError != "" {
57+
extra["subscription_error"] = result.SubscriptionError
58+
} else {
59+
delete(extra, "subscription_error")
60+
}
61+
return credentials, extra
62+
}

0 commit comments

Comments
 (0)