feat: add AI Credits balance handling and update model status indicators

2026-04-26 09:24:48 +08:00 · 2026-03-15 23:50:28 +08:00
parent 17e4033340
commit ced90e1d84
12 changed files with 465 additions and 77 deletions
--- a/backend/internal/service/account_usage_service.go
+++ b/backend/internal/service/account_usage_service.go
@@ -166,6 +166,13 @@ type AntigravityModelDetail struct {
 	SupportedMimeTypes map[string]bool `json:"supported_mime_types,omitempty"`
 }

+// AICredit 表示 Antigravity 账号的 AI Credits 余额信息。
+type AICredit struct {
+	CreditType     string  `json:"credit_type,omitempty"`
+	Amount         float64 `json:"amount,omitempty"`
+	MinimumBalance float64 `json:"minimum_balance,omitempty"`
+}
+
 // UsageInfo 账号使用量信息
 type UsageInfo struct {
 	UpdatedAt          *time.Time     `json:"updated_at,omitempty"`           // 更新时间
@@ -189,6 +196,9 @@ type UsageInfo struct {
 	// Antigravity 模型详细能力信息（与 antigravity_quota 同 key）
 	AntigravityQuotaDetails map[string]*AntigravityModelDetail `json:"antigravity_quota_details,omitempty"`

+	// Antigravity AI Credits 余额
+	AICredits []AICredit `json:"ai_credits,omitempty"`
+
 	// Antigravity 废弃模型转发规则 (old_model_id -> new_model_id)
 	ModelForwardingRules map[string]string `json:"model_forwarding_rules,omitempty"`

--- a/backend/internal/service/antigravity_quota_fetcher.go
+++ b/backend/internal/service/antigravity_quota_fetcher.go
@@ -78,11 +78,11 @@ func (f *AntigravityQuotaFetcher) FetchQuota(ctx context.Context, account *Accou
 		return nil, err
 	}

-	// 调用 LoadCodeAssist 获取订阅等级（非关键路径，失败不影响主流程）
-	tierRaw, tierNormalized := f.fetchSubscriptionTier(ctx, client, accessToken)
+	// 调用 LoadCodeAssist 获取订阅等级和 AI Credits 余额（非关键路径，失败不影响主流程）
+	tierRaw, tierNormalized, loadResp := f.fetchSubscriptionTier(ctx, client, accessToken)

 	// 转换为 UsageInfo
-	usageInfo := f.buildUsageInfo(modelsResp, tierRaw, tierNormalized)
+	usageInfo := f.buildUsageInfo(modelsResp, tierRaw, tierNormalized, loadResp)

 	return &QuotaResult{
 		UsageInfo: usageInfo,
@@ -90,20 +90,21 @@ func (f *AntigravityQuotaFetcher) FetchQuota(ctx context.Context, account *Accou
 	}, nil
 }

-// fetchSubscriptionTier 获取账号订阅等级，失败返回空字符串
-func (f *AntigravityQuotaFetcher) fetchSubscriptionTier(ctx context.Context, client *antigravity.Client, accessToken string) (raw, normalized string) {
+// fetchSubscriptionTier 获取账号订阅等级，失败返回空字符串。
+// 同时返回 LoadCodeAssistResponse，以便提取 AI Credits 余额。
+func (f *AntigravityQuotaFetcher) fetchSubscriptionTier(ctx context.Context, client *antigravity.Client, accessToken string) (raw, normalized string, loadResp *antigravity.LoadCodeAssistResponse) {
 	loadResp, _, err := client.LoadCodeAssist(ctx, accessToken)
 	if err != nil {
 		slog.Warn("failed to fetch subscription tier", "error", err)
-		return "", ""
+		return "", "", nil
 	}
 	if loadResp == nil {
-		return "", ""
+		return "", "", nil
 	}

 	raw = loadResp.GetTier() // 已有方法：paidTier > currentTier
 	normalized = normalizeTier(raw)
-	return raw, normalized
+	return raw, normalized, loadResp
 }

 // normalizeTier 将原始 tier 字符串归一化为 FREE/PRO/ULTRA/UNKNOWN
@@ -124,8 +125,8 @@ func normalizeTier(raw string) string {
 	}
 }

-// buildUsageInfo 将 API 响应转换为 UsageInfo
-func (f *AntigravityQuotaFetcher) buildUsageInfo(modelsResp *antigravity.FetchAvailableModelsResponse, tierRaw, tierNormalized string) *UsageInfo {
+// buildUsageInfo 将 API 响应转换为 UsageInfo。
+func (f *AntigravityQuotaFetcher) buildUsageInfo(modelsResp *antigravity.FetchAvailableModelsResponse, tierRaw, tierNormalized string, loadResp *antigravity.LoadCodeAssistResponse) *UsageInfo {
 	now := time.Now()
 	info := &UsageInfo{
 		UpdatedAt:               &now,
@@ -190,6 +191,16 @@ func (f *AntigravityQuotaFetcher) buildUsageInfo(modelsResp *antigravity.FetchAv
 		}
 	}

+	if loadResp != nil {
+		for _, credit := range loadResp.GetAvailableCredits() {
+			info.AICredits = append(info.AICredits, AICredit{
+				CreditType:     credit.CreditType,
+				Amount:         credit.GetAmount(),
+				MinimumBalance: credit.GetMinimumAmount(),
+			})
+		}
+	}
+
 	return info
 }

--- a/backend/internal/service/antigravity_quota_fetcher_test.go
+++ b/backend/internal/service/antigravity_quota_fetcher_test.go
@@ -81,7 +81,7 @@ func TestBuildUsageInfo_BasicModels(t *testing.T) {
 		},
 	}

-	info := fetcher.buildUsageInfo(modelsResp, "g1-pro-tier", "PRO")
+	info := fetcher.buildUsageInfo(modelsResp, "g1-pro-tier", "PRO", nil)

 	// 基本字段
 	require.NotNil(t, info.UpdatedAt, "UpdatedAt should be set")
@@ -141,7 +141,7 @@ func TestBuildUsageInfo_DeprecatedModels(t *testing.T) {
 		},
 	}

-	info := fetcher.buildUsageInfo(modelsResp, "", "")
+	info := fetcher.buildUsageInfo(modelsResp, "", "", nil)

 	require.Len(t, info.ModelForwardingRules, 2)
 	require.Equal(t, "claude-sonnet-4-20250514", info.ModelForwardingRules["claude-3-sonnet-20240229"])
@@ -159,7 +159,7 @@ func TestBuildUsageInfo_NoDeprecatedModels(t *testing.T) {
 		},
 	}

-	info := fetcher.buildUsageInfo(modelsResp, "", "")
+	info := fetcher.buildUsageInfo(modelsResp, "", "", nil)

 	require.Nil(t, info.ModelForwardingRules, "ModelForwardingRules should be nil when no deprecated models")
 }
@@ -171,7 +171,7 @@ func TestBuildUsageInfo_EmptyModels(t *testing.T) {
 		Models: map[string]antigravity.ModelInfo{},
 	}

-	info := fetcher.buildUsageInfo(modelsResp, "", "")
+	info := fetcher.buildUsageInfo(modelsResp, "", "", nil)

 	require.NotNil(t, info)
 	require.NotNil(t, info.AntigravityQuota)
@@ -193,7 +193,7 @@ func TestBuildUsageInfo_ModelWithNilQuotaInfo(t *testing.T) {
 		},
 	}

-	info := fetcher.buildUsageInfo(modelsResp, "", "")
+	info := fetcher.buildUsageInfo(modelsResp, "", "", nil)

 	require.NotNil(t, info)
 	require.Empty(t, info.AntigravityQuota, "models with nil QuotaInfo should be skipped")
@@ -222,7 +222,7 @@ func TestBuildUsageInfo_FiveHourPriorityOrder(t *testing.T) {
 		},
 	}

-	info := fetcher.buildUsageInfo(modelsResp, "", "")
+	info := fetcher.buildUsageInfo(modelsResp, "", "", nil)

 	require.NotNil(t, info.FiveHour, "FiveHour should be set when a priority model exists")
 	// claude-sonnet-4-20250514 is first in priority list, so it should be used
@@ -251,7 +251,7 @@ func TestBuildUsageInfo_FiveHourFallbackToClaude4(t *testing.T) {
 		},
 	}

-	info := fetcher.buildUsageInfo(modelsResp, "", "")
+	info := fetcher.buildUsageInfo(modelsResp, "", "", nil)

 	require.NotNil(t, info.FiveHour)
 	expectedUtilization := (1.0 - 0.60) * 100 // 40
@@ -277,7 +277,7 @@ func TestBuildUsageInfo_FiveHourFallbackToGemini(t *testing.T) {
 		},
 	}

-	info := fetcher.buildUsageInfo(modelsResp, "", "")
+	info := fetcher.buildUsageInfo(modelsResp, "", "", nil)

 	require.NotNil(t, info.FiveHour)
 	expectedUtilization := (1.0 - 0.30) * 100 // 70
@@ -298,7 +298,7 @@ func TestBuildUsageInfo_FiveHourNoPriorityModel(t *testing.T) {
 		},
 	}

-	info := fetcher.buildUsageInfo(modelsResp, "", "")
+	info := fetcher.buildUsageInfo(modelsResp, "", "", nil)

 	require.Nil(t, info.FiveHour, "FiveHour should be nil when no priority model exists")
 }
@@ -317,7 +317,7 @@ func TestBuildUsageInfo_FiveHourWithEmptyResetTime(t *testing.T) {
 		},
 	}

-	info := fetcher.buildUsageInfo(modelsResp, "", "")
+	info := fetcher.buildUsageInfo(modelsResp, "", "", nil)

 	require.NotNil(t, info.FiveHour)
 	require.Nil(t, info.FiveHour.ResetsAt, "ResetsAt should be nil when ResetTime is empty")
@@ -338,7 +338,7 @@ func TestBuildUsageInfo_FullUtilization(t *testing.T) {
 		},
 	}

-	info := fetcher.buildUsageInfo(modelsResp, "", "")
+	info := fetcher.buildUsageInfo(modelsResp, "", "", nil)

 	quota := info.AntigravityQuota["claude-sonnet-4-20250514"]
 	require.NotNil(t, quota)
@@ -358,13 +358,38 @@ func TestBuildUsageInfo_ZeroUtilization(t *testing.T) {
 		},
 	}

-	info := fetcher.buildUsageInfo(modelsResp, "", "")
-
+	info := fetcher.buildUsageInfo(modelsResp, "", "", nil)
 	quota := info.AntigravityQuota["claude-sonnet-4-20250514"]
 	require.NotNil(t, quota)
 	require.Equal(t, 0, quota.Utilization)
 }

+func TestBuildUsageInfo_AICredits(t *testing.T) {
+	fetcher := &AntigravityQuotaFetcher{}
+	modelsResp := &antigravity.FetchAvailableModelsResponse{
+		Models: map[string]antigravity.ModelInfo{},
+	}
+	loadResp := &antigravity.LoadCodeAssistResponse{
+		PaidTier: &antigravity.PaidTierInfo{
+			ID: "g1-pro-tier",
+			AvailableCredits: []antigravity.AvailableCredit{
+				{
+					CreditType:                  "GOOGLE_ONE_AI",
+					CreditAmount:                "25",
+					MinimumCreditAmountForUsage: "5",
+				},
+			},
+		},
+	}
+
+	info := fetcher.buildUsageInfo(modelsResp, "g1-pro-tier", "PRO", loadResp)
+
+	require.Len(t, info.AICredits, 1)
+	require.Equal(t, "GOOGLE_ONE_AI", info.AICredits[0].CreditType)
+	require.Equal(t, 25.0, info.AICredits[0].Amount)
+	require.Equal(t, 5.0, info.AICredits[0].MinimumBalance)
+}
+
 func TestFetchQuota_ForbiddenReturnsIsForbidden(t *testing.T) {
 	// 模拟 FetchQuota 遇到 403 时的行为：
 	// FetchAvailableModels 返回 ForbiddenError → FetchQuota 应返回 is_forbidden=true