a185ad1144
主要改动： - 后端：重构 Gemini 配额服务，支持多层级配额策略（GCP Standard/Free, Google One, AI Studio, Code Assist） - 后端：优化 OAuth 服务，增强 tier_id 识别和存储逻辑 - 后端：改进用量统计服务，支持不同平台的配额查询 - 后端：优化限流服务，增加临时解除调度状态管理 - 前端：统一四种授权方式的用量显示格式和徽标样式 - 前端：增强账户配额信息展示，支持多种配额类型 - 前端：改进创建和重新授权模态框的用户体验 - 国际化：完善中英文配额相关文案 - 移除 CHANGELOG.md 文件 测试：所有单元测试通过
118 lines
3.4 KiB
Go
118 lines
3.4 KiB
Go
package repository
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"net/url"
|
|
"time"
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/config"
|
|
"github.com/Wei-Shaw/sub2api/internal/pkg/geminicli"
|
|
"github.com/Wei-Shaw/sub2api/internal/service"
|
|
|
|
"github.com/imroc/req/v3"
|
|
)
|
|
|
|
type geminiOAuthClient struct {
|
|
tokenURL string
|
|
cfg *config.Config
|
|
}
|
|
|
|
func NewGeminiOAuthClient(cfg *config.Config) service.GeminiOAuthClient {
|
|
return &geminiOAuthClient{
|
|
tokenURL: geminicli.TokenURL,
|
|
cfg: cfg,
|
|
}
|
|
}
|
|
|
|
func (c *geminiOAuthClient) ExchangeCode(ctx context.Context, oauthType, code, codeVerifier, redirectURI, proxyURL string) (*geminicli.TokenResponse, error) {
|
|
client := createGeminiReqClient(proxyURL)
|
|
|
|
// Use different OAuth clients based on oauthType:
|
|
// - code_assist: always use built-in Gemini CLI OAuth client (public)
|
|
// - google_one: uses configured OAuth client when provided; otherwise falls back to built-in client
|
|
// - ai_studio: requires a user-provided OAuth client
|
|
oauthCfgInput := geminicli.OAuthConfig{
|
|
ClientID: c.cfg.Gemini.OAuth.ClientID,
|
|
ClientSecret: c.cfg.Gemini.OAuth.ClientSecret,
|
|
Scopes: c.cfg.Gemini.OAuth.Scopes,
|
|
}
|
|
if oauthType == "code_assist" {
|
|
oauthCfgInput.ClientID = ""
|
|
oauthCfgInput.ClientSecret = ""
|
|
}
|
|
|
|
oauthCfg, err := geminicli.EffectiveOAuthConfig(oauthCfgInput, oauthType)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
formData := url.Values{}
|
|
formData.Set("grant_type", "authorization_code")
|
|
formData.Set("client_id", oauthCfg.ClientID)
|
|
formData.Set("client_secret", oauthCfg.ClientSecret)
|
|
formData.Set("code", code)
|
|
formData.Set("code_verifier", codeVerifier)
|
|
formData.Set("redirect_uri", redirectURI)
|
|
|
|
var tokenResp geminicli.TokenResponse
|
|
resp, err := client.R().
|
|
SetContext(ctx).
|
|
SetFormDataFromValues(formData).
|
|
SetSuccessResult(&tokenResp).
|
|
Post(c.tokenURL)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("request failed: %w", err)
|
|
}
|
|
if !resp.IsSuccessState() {
|
|
return nil, fmt.Errorf("token exchange failed: status %d, body: %s", resp.StatusCode, geminicli.SanitizeBodyForLogs(resp.String()))
|
|
}
|
|
return &tokenResp, nil
|
|
}
|
|
|
|
func (c *geminiOAuthClient) RefreshToken(ctx context.Context, oauthType, refreshToken, proxyURL string) (*geminicli.TokenResponse, error) {
|
|
client := createGeminiReqClient(proxyURL)
|
|
|
|
oauthCfgInput := geminicli.OAuthConfig{
|
|
ClientID: c.cfg.Gemini.OAuth.ClientID,
|
|
ClientSecret: c.cfg.Gemini.OAuth.ClientSecret,
|
|
Scopes: c.cfg.Gemini.OAuth.Scopes,
|
|
}
|
|
if oauthType == "code_assist" {
|
|
oauthCfgInput.ClientID = ""
|
|
oauthCfgInput.ClientSecret = ""
|
|
}
|
|
|
|
oauthCfg, err := geminicli.EffectiveOAuthConfig(oauthCfgInput, oauthType)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
formData := url.Values{}
|
|
formData.Set("grant_type", "refresh_token")
|
|
formData.Set("refresh_token", refreshToken)
|
|
formData.Set("client_id", oauthCfg.ClientID)
|
|
formData.Set("client_secret", oauthCfg.ClientSecret)
|
|
|
|
var tokenResp geminicli.TokenResponse
|
|
resp, err := client.R().
|
|
SetContext(ctx).
|
|
SetFormDataFromValues(formData).
|
|
SetSuccessResult(&tokenResp).
|
|
Post(c.tokenURL)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("request failed: %w", err)
|
|
}
|
|
if !resp.IsSuccessState() {
|
|
return nil, fmt.Errorf("token refresh failed: status %d, body: %s", resp.StatusCode, geminicli.SanitizeBodyForLogs(resp.String()))
|
|
}
|
|
return &tokenResp, nil
|
|
}
|
|
|
|
func createGeminiReqClient(proxyURL string) *req.Client {
|
|
return getSharedReqClient(reqClientOptions{
|
|
ProxyURL: proxyURL,
|
|
Timeout: 60 * time.Second,
|
|
})
|
|
}
|