oadmin/sub2api
1
0
Fork 0
Code Issues Pull Requests 1 Actions 1 Packages Projects Releases Wiki Activity

feat(gemini): 添加 Google One 存储空间推断 Tier 功能

Browse Source 
## 功能概述
通过 Google Drive API 获取存储空间配额来推断 Google One 订阅等级，并优化统一的配额显示系统。

## 后端改动
- 新增 Drive API 客户端 (drive_client.go)
  - 支持代理和指数退避重试
  - 处理 403/429 错误
- 添加 Tier 推断逻辑 (inferGoogleOneTier)
  - 支持 6 种 tier 类型：AI_PREMIUM, GOOGLE_ONE_STANDARD, GOOGLE_ONE_BASIC, FREE, GOOGLE_ONE_UNKNOWN, GOOGLE_ONE_UNLIMITED
- 集成到 OAuth 流程
  - ExchangeCode: 授权时自动获取 tier
  - RefreshAccountToken: Token 刷新时更新 tier (24小时缓存)
- 新增管理 API 端点
  - POST /api/v1/admin/accounts/:id/refresh-tier (单个账号刷新)
  - POST /api/v1/admin/accounts/batch-refresh-tier (批量刷新)

## 前端改动
- 更新 AccountQuotaInfo.vue
  - 添加 Google One tier 标签映射
  - 添加 tier 颜色样式 (紫色/蓝色/绿色/灰色/琥珀色)
- 更新 AccountUsageCell.vue
  - 添加 Google One tier 显示逻辑
  - 根据 oauth_type 区分显示方式
- 添加国际化翻译 (en.ts, zh.ts)
  - aiPremium, standard, basic, free, personal, unlimited

## Tier 推断规则
- >= 2TB: AI Premium
- >= 200GB: Google One Standard
- >= 100GB: Google One Basic
- >= 15GB: Free
- > 100TB: Unlimited (G Suite legacy)
- 其他/失败: Unknown (显示为 Personal)

## 优雅降级
- Drive API 失败时使用 GOOGLE_ONE_UNKNOWN
- 不阻断 OAuth 流程
- 24小时缓存避免频繁调用

## 测试
- ✅ 后端编译成功
- ✅ 前端构建成功
- ✅ 所有代码符合现有规范
This commit is contained in:
ianshaw
2025-12-31 21:45:24 -08:00
parent 4f13c8de0d
commit 7df914af06
11 changed files with 691 additions and 110 deletions
Download Patch File Download Diff File Expand all files Collapse all files

165
backend/internal/handler/admin/account_handler.go
View File

@@ -3,6 +3,7 @@ package admin
import (
"strconv"
"strings"
"time"
"github.com/Wei-Shaw/sub2api/internal/handler/dto"
"github.com/Wei-Shaw/sub2api/internal/pkg/claude"
@@ -989,3 +990,167 @@ func (h *AccountHandler) GetAvailableModels(c *gin.Context) {
response.Success(c, models)
}
// 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) {
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.Credentials == nil || account.Credentials["oauth_type"] != "google_one" {
response.BadRequest(c, "Account is not a google_one OAuth account")
return
}
accessToken, ok := account.Credentials["access_token"].(string)
if !ok || accessToken == "" {
response.BadRequest(c, "Missing access_token in credentials")
return
}
var proxyURL string
if account.ProxyID != nil && account.Proxy != nil {
proxyURL = account.Proxy.URL()
}
tierID, storageInfo, err := h.geminiOAuthService.FetchGoogleOneTier(c.Request.Context(), accessToken, proxyURL)
if account.Extra == nil {
account.Extra = make(map[string]any)
}
if storageInfo != nil {
account.Extra["drive_storage_limit"] = storageInfo.Limit
account.Extra["drive_storage_usage"] = storageInfo.Usage
account.Extra["drive_tier_updated_at"] = timezone.Now().Format(time.RFC3339)
}
account.Credentials["tier_id"] = tierID
_, updateErr := h.adminService.UpdateAccount(c.Request.Context(), accountID, &service.UpdateAccountInput{
Credentials: account.Credentials,
Extra: account.Extra,
})
if updateErr != nil {
response.ErrorFrom(c, updateErr)
return
}
response.Success(c, gin.H{
"tier_id": tierID,
"drive_storage_limit": account.Extra["drive_storage_limit"],
"drive_storage_usage": account.Extra["drive_storage_usage"],
"updated_at": account.Extra["drive_tier_updated_at"],
})
}
// BatchRefreshTierRequest represents batch tier refresh request
type BatchRefreshTierRequest struct {
AccountIDs []int64 `json:"account_ids"`
}
// BatchRefreshTier handles batch refreshing Google One tier
// POST /api/v1/admin/accounts/batch-refresh-tier
func (h *AccountHandler) BatchRefreshTier(c *gin.Context) {
var req BatchRefreshTierRequest
if err := c.ShouldBindJSON(&req); err != nil {
req = BatchRefreshTierRequest{}
}
ctx := c.Request.Context()
var accounts []service.Account
if len(req.AccountIDs) == 0 {
allAccounts, _, err := h.adminService.ListAccounts(ctx, 1, 10000, "gemini", "oauth", "", "")
if err != nil {
response.ErrorFrom(c, err)
return
}
for _, acc := range allAccounts {
if acc.Credentials != nil && acc.Credentials["oauth_type"] == "google_one" {
accounts = append(accounts, acc)
}
}
} else {
for _, id := range req.AccountIDs {
acc, err := h.adminService.GetAccount(ctx, id)
if err != nil {
continue
}
if acc.Credentials != nil && acc.Credentials["oauth_type"] == "google_one" {
accounts = append(accounts, *acc)
}
}
}
total := len(accounts)
success := 0
failed := 0
errors := []gin.H{}
for _, account := range accounts {
accessToken, ok := account.Credentials["access_token"].(string)
if !ok || accessToken == "" {
failed++
errors = append(errors, gin.H{
"account_id": account.ID,
"error": "missing access_token",
})
continue
}
var proxyURL string
if account.ProxyID != nil && account.Proxy != nil {
proxyURL = account.Proxy.URL()
}
tierID, storageInfo, err := h.geminiOAuthService.FetchGoogleOneTier(ctx, accessToken, proxyURL)
if err != nil {
failed++
errors = append(errors, gin.H{
"account_id": account.ID,
"error": err.Error(),
})
continue
}
if account.Extra == nil {
account.Extra = make(map[string]any)
}
if storageInfo != nil {
account.Extra["drive_storage_limit"] = storageInfo.Limit
account.Extra["drive_storage_usage"] = storageInfo.Usage
account.Extra["drive_tier_updated_at"] = timezone.Now().Format(time.RFC3339)
}
account.Credentials["tier_id"] = tierID
_, updateErr := h.adminService.UpdateAccount(ctx, account.ID, &service.UpdateAccountInput{
Credentials: account.Credentials,
Extra: account.Extra,
})
if updateErr != nil {
failed++
errors = append(errors, gin.H{
"account_id": account.ID,
"error": updateErr.Error(),
})
continue
}
success++
}
response.Success(c, gin.H{
"total": total,
"success": success,
"failed": failed,
"errors": errors,
})
}

8
backend/internal/handler/admin/gemini_oauth_handler.go
View File

@@ -46,8 +46,8 @@ func (h *GeminiOAuthHandler) GenerateAuthURL(c *gin.Context) {
if oauthType == "" {
oauthType = "code_assist"
}
if oauthType != "code_assist" && oauthType != "ai_studio" {
response.BadRequest(c, "Invalid oauth_type: must be 'code_assist' or 'ai_studio'")
if oauthType != "code_assist" && oauthType != "google_one" && oauthType != "ai_studio" {
response.BadRequest(c, "Invalid oauth_type: must be 'code_assist', 'google_one', or 'ai_studio'")
return
}
@@ -92,8 +92,8 @@ func (h *GeminiOAuthHandler) ExchangeCode(c *gin.Context) {
if oauthType == "" {
oauthType = "code_assist"
}
if oauthType != "code_assist" && oauthType != "ai_studio" {
response.BadRequest(c, "Invalid oauth_type: must be 'code_assist' or 'ai_studio'")
if oauthType != "code_assist" && oauthType != "google_one" && oauthType != "ai_studio" {
response.BadRequest(c, "Invalid oauth_type: must be 'code_assist', 'google_one', or 'ai_studio'")
return
}

113
backend/internal/pkg/geminicli/drive_client.go Normal file
View File

@@ -0,0 +1,113 @@
package geminicli
import (
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/httpclient"
)
// DriveStorageInfo represents Google Drive storage quota information
type DriveStorageInfo struct {
Limit int64 `json:"limit"` // Storage limit in bytes
Usage int64 `json:"usage"` // Current usage in bytes
}
// DriveClient interface for Google Drive API operations
type DriveClient interface {
GetStorageQuota(ctx context.Context, accessToken, proxyURL string) (*DriveStorageInfo, error)
}
type driveClient struct {
httpClient *http.Client
}
// NewDriveClient creates a new Drive API client
func NewDriveClient() DriveClient {
return &driveClient{
httpClient: &http.Client{
Timeout: 10 * time.Second,
},
}
}
// GetStorageQuota fetches storage quota from Google Drive API
func (c *driveClient) GetStorageQuota(ctx context.Context, accessToken, proxyURL string) (*DriveStorageInfo, error) {
const driveAPIURL = "https://www.googleapis.com/drive/v3/about?fields=storageQuota"
req, err := http.NewRequestWithContext(ctx, "GET", driveAPIURL, nil)
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Authorization", "Bearer "+accessToken)
// Get HTTP client with proxy support
client, err := httpclient.GetClient(httpclient.Options{
ProxyURL: proxyURL,
Timeout: 10 * time.Second,
})
if err != nil {
return nil, fmt.Errorf("failed to create HTTP client: %w", err)
}
// Retry logic with exponential backoff for rate limits
var resp *http.Response
maxRetries := 3
for attempt := 0; attempt < maxRetries; attempt++ {
resp, err = client.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to execute request: %w", err)
}
// Success
if resp.StatusCode == http.StatusOK {
break
}
// Rate limit - retry with exponential backoff
if resp.StatusCode == http.StatusTooManyRequests && attempt < maxRetries-1 {
resp.Body.Close()
backoff := time.Duration(1<<uint(attempt)) * time.Second // 1s, 2s, 4s
time.Sleep(backoff)
continue
}
// Other errors - return immediately
body, _ := io.ReadAll(resp.Body)
resp.Body.Close()
return nil, fmt.Errorf("Drive API error (status %d): %s", resp.StatusCode, string(body))
}
defer resp.Body.Close()
// Parse response
var result struct {
StorageQuota struct {
Limit string `json:"limit"` // Can be string or number
Usage string `json:"usage"`
} `json:"storageQuota"`
}
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
return nil, fmt.Errorf("failed to decode response: %w", err)
}
// Parse limit and usage (handle both string and number formats)
var limit, usage int64
if result.StorageQuota.Limit != "" {
fmt.Sscanf(result.StorageQuota.Limit, "%d", &limit)
}
if result.StorageQuota.Usage != "" {
fmt.Sscanf(result.StorageQuota.Usage, "%d", &usage)
}
return &DriveStorageInfo{
Limit: limit,
Usage: usage,
}, nil
}

2
backend/internal/server/routes/admin.go
View File

@@ -110,6 +110,7 @@ func registerAccountRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
accounts.DELETE("/:id", h.Admin.Account.Delete)
accounts.POST("/:id/test", h.Admin.Account.Test)
accounts.POST("/:id/refresh", h.Admin.Account.Refresh)
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)
accounts.GET("/:id/usage", h.Admin.Account.GetUsage)
@@ -119,6 +120,7 @@ func registerAccountRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
accounts.GET("/:id/models", h.Admin.Account.GetAvailableModels)
accounts.POST("/batch", h.Admin.Account.BatchCreate)
accounts.POST("/batch-update-credentials", h.Admin.Account.BatchUpdateCredentials)
accounts.POST("/batch-refresh-tier", h.Admin.Account.BatchRefreshTier)
accounts.POST("/bulk-update", h.Admin.Account.BulkUpdate)
// Claude OAuth routes

159
backend/internal/service/gemini_oauth_service.go
View File

@@ -17,6 +17,15 @@ import (
"github.com/Wei-Shaw/sub2api/internal/pkg/httpclient"
)
const (
TierAIPremium = "AI_PREMIUM"
TierGoogleOneStandard = "GOOGLE_ONE_STANDARD"
TierGoogleOneBasic = "GOOGLE_ONE_BASIC"
TierFree = "FREE"
TierGoogleOneUnknown = "GOOGLE_ONE_UNKNOWN"
TierGoogleOneUnlimited = "GOOGLE_ONE_UNLIMITED"
)
type GeminiOAuthService struct {
sessionStore *geminicli.SessionStore
proxyRepo ProxyRepository
@@ -89,13 +98,14 @@ func (s *GeminiOAuthService) GenerateAuthURL(ctx context.Context, proxyID *int64
// OAuth client selection:
// - code_assist: always use built-in Gemini CLI OAuth client (public), regardless of configured client_id/secret.
// - google_one: same as code_assist, uses built-in client for personal Google accounts.
// - ai_studio: requires a user-provided OAuth client.
oauthCfg := geminicli.OAuthConfig{
ClientID: s.cfg.Gemini.OAuth.ClientID,
ClientSecret: s.cfg.Gemini.OAuth.ClientSecret,
Scopes: s.cfg.Gemini.OAuth.Scopes,
}
if oauthType == "code_assist" {
if oauthType == "code_assist" || oauthType == "google_one" {
oauthCfg.ClientID = ""
oauthCfg.ClientSecret = ""
}
@@ -156,15 +166,16 @@ type GeminiExchangeCodeInput struct {
}
type GeminiTokenInfo struct {
AccessToken string `json:"access_token"`
RefreshToken string `json:"refresh_token"`
ExpiresIn int64 `json:"expires_in"`
ExpiresAt int64 `json:"expires_at"`
TokenType string `json:"token_type"`
Scope string `json:"scope,omitempty"`
ProjectID string `json:"project_id,omitempty"`
OAuthType string `json:"oauth_type,omitempty"` // "code_assist" 或 "ai_studio"
TierID string `json:"tier_id,omitempty"` // Gemini Code Assist tier: LEGACY/PRO/ULTRA
AccessToken string `json:"access_token"`
RefreshToken string `json:"refresh_token"`
ExpiresIn int64 `json:"expires_in"`
ExpiresAt int64 `json:"expires_at"`
TokenType string `json:"token_type"`
Scope string `json:"scope,omitempty"`
ProjectID string `json:"project_id,omitempty"`
OAuthType string `json:"oauth_type,omitempty"` // "code_assist" 或 "ai_studio"
TierID string `json:"tier_id,omitempty"` // Gemini Code Assist tier: LEGACY/PRO/ULTRA
Extra map[string]any `json:"extra,omitempty"` // Drive metadata
}
// validateTierID validates tier_id format and length
@@ -205,6 +216,60 @@ func extractTierIDFromAllowedTiers(allowedTiers []geminicli.AllowedTier) string
return tierID
}
// inferGoogleOneTier infers Google One tier from Drive storage limit
func inferGoogleOneTier(storageBytes int64) string {
if storageBytes <= 0 {
return TierGoogleOneUnknown
}
// Unlimited storage (G Suite legacy)
if storageBytes > 100*1024*1024*1024*1024 { // > 100TB
return TierGoogleOneUnlimited
}
// AI Premium (2TB+)
if storageBytes >= 2*1024*1024*1024*1024 { // >= 2TB
return TierAIPremium
}
// Google One Standard (200GB)
if storageBytes >= 200*1024*1024*1024 { // >= 200GB
return TierGoogleOneStandard
}
// Google One Basic (100GB)
if storageBytes >= 100*1024*1024*1024 { // >= 100GB
return TierGoogleOneBasic
}
// Free (15GB)
if storageBytes >= 15*1024*1024*1024 { // >= 15GB
return TierFree
}
return TierGoogleOneUnknown
}
// fetchGoogleOneTier fetches Google One tier from Drive API
func (s *GeminiOAuthService) FetchGoogleOneTier(ctx context.Context, accessToken, proxyURL string) (string, *geminicli.DriveStorageInfo, error) {
driveClient := geminicli.NewDriveClient()
storageInfo, err := driveClient.GetStorageQuota(ctx, accessToken, proxyURL)
if err != nil {
// Check if it's a 403 (scope not granted)
if strings.Contains(err.Error(), "status 403") {
fmt.Printf("[GeminiOAuth] Drive API scope not available: %v\n", err)
return TierGoogleOneUnknown, nil, err
}
// Other errors
fmt.Printf("[GeminiOAuth] Failed to fetch Drive storage: %v\n", err)
return TierGoogleOneUnknown, nil, err
}
tierID := inferGoogleOneTier(storageInfo.Limit)
return tierID, storageInfo, nil
}
func (s *GeminiOAuthService) ExchangeCode(ctx context.Context, input *GeminiExchangeCodeInput) (*GeminiTokenInfo, error) {
session, ok := s.sessionStore.Get(input.SessionID)
if !ok {
@@ -272,7 +337,8 @@ func (s *GeminiOAuthService) ExchangeCode(ctx context.Context, input *GeminiExch
projectID := sessionProjectID
var tierID string
// 对于 code_assist 模式，project_id 是必需的
// 对于 code_assist 模式，project_id 是必需的，需要调用 Code Assist API
// 对于 google_one 模式，使用个人 Google 账号，不需要 project_id，配额由 Google 网关自动识别
// 对于 ai_studio 模式，project_id 是可选的（不影响使用 AI Studio API）
if oauthType == "code_assist" {
if projectID == "" {
@@ -298,7 +364,37 @@ func (s *GeminiOAuthService) ExchangeCode(ctx context.Context, input *GeminiExch
if tierID == "" {
tierID = "LEGACY"
}
} else if oauthType == "google_one" {
// Attempt to fetch Drive storage tier
tierID, storageInfo, err := s.FetchGoogleOneTier(ctx, tokenResp.AccessToken, proxyURL)
if err != nil {
// Log warning but don't block - use fallback
fmt.Printf("[GeminiOAuth] Warning: Failed to fetch Drive tier: %v\n", err)
tierID = TierGoogleOneUnknown
}
// Store Drive info in extra field for caching
if storageInfo != nil {
tokenInfo := &GeminiTokenInfo{
AccessToken: tokenResp.AccessToken,
RefreshToken: tokenResp.RefreshToken,
TokenType: tokenResp.TokenType,
ExpiresIn: tokenResp.ExpiresIn,
ExpiresAt: expiresAt,
Scope: tokenResp.Scope,
ProjectID: projectID,
TierID: tierID,
OAuthType: oauthType,
Extra: map[string]any{
"drive_storage_limit": storageInfo.Limit,
"drive_storage_usage": storageInfo.Usage,
"drive_tier_updated_at": time.Now().Format(time.RFC3339),
},
}
return tokenInfo, nil
}
}
// ai_studio 模式不设置 tierID，保持为空
return &GeminiTokenInfo{
AccessToken: tokenResp.AccessToken,
@@ -455,6 +551,41 @@ func (s *GeminiOAuthService) RefreshAccountToken(ctx context.Context, account *A
if strings.TrimSpace(tokenInfo.ProjectID) == "" {
return nil, fmt.Errorf("failed to auto-detect project_id: empty result")
}
} else if oauthType == "google_one" {
// Check if tier cache is stale (> 24 hours)
needsRefresh := true
if account.Extra != nil {
if updatedAtStr, ok := account.Extra["drive_tier_updated_at"].(string); ok {
if updatedAt, err := time.Parse(time.RFC3339, updatedAtStr); err == nil {
if time.Since(updatedAt) <= 24*time.Hour {
needsRefresh = false
// Use cached tier
if existingTierID != "" {
tokenInfo.TierID = existingTierID
}
}
}
}
}
if needsRefresh {
tierID, storageInfo, err := s.FetchGoogleOneTier(ctx, tokenInfo.AccessToken, proxyURL)
if err == nil && storageInfo != nil {
tokenInfo.TierID = tierID
tokenInfo.Extra = map[string]any{
"drive_storage_limit": storageInfo.Limit,
"drive_storage_usage": storageInfo.Usage,
"drive_tier_updated_at": time.Now().Format(time.RFC3339),
}
} else {
// Fallback to cached or unknown
if existingTierID != "" {
tokenInfo.TierID = existingTierID
} else {
tokenInfo.TierID = TierGoogleOneUnknown
}
}
}
}
return tokenInfo, nil
@@ -487,6 +618,12 @@ func (s *GeminiOAuthService) BuildAccountCredentials(tokenInfo *GeminiTokenInfo)
if tokenInfo.OAuthType != "" {
creds["oauth_type"] = tokenInfo.OAuthType
}
// Store extra metadata (Drive info) if present
if len(tokenInfo.Extra) > 0 {
for k, v := range tokenInfo.Extra {
creds[k] = v
}
}
return creds
}