08ce6de4db
后端: - 新增 AntigravityQuotaRefresher 定时刷新配额 - Client 添加 FetchAvailableModels 方法获取模型配额 - 配额数据存入 account.extra.quota 字段 前端: - AccountUsageCell 支持显示 Antigravity 账户配额 - UsageProgressBar 新增 amber 颜色 - 显示 G3P/G3F/G3I/C4.5 四个配额进度条
278 lines
8.0 KiB
Go
278 lines
8.0 KiB
Go
package antigravity
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"net/url"
|
|
"strings"
|
|
"time"
|
|
)
|
|
|
|
// TokenResponse Google OAuth token 响应
|
|
type TokenResponse struct {
|
|
AccessToken string `json:"access_token"`
|
|
ExpiresIn int64 `json:"expires_in"`
|
|
TokenType string `json:"token_type"`
|
|
Scope string `json:"scope,omitempty"`
|
|
RefreshToken string `json:"refresh_token,omitempty"`
|
|
}
|
|
|
|
// UserInfo Google 用户信息
|
|
type UserInfo struct {
|
|
Email string `json:"email"`
|
|
Name string `json:"name,omitempty"`
|
|
GivenName string `json:"given_name,omitempty"`
|
|
FamilyName string `json:"family_name,omitempty"`
|
|
Picture string `json:"picture,omitempty"`
|
|
}
|
|
|
|
// LoadCodeAssistRequest loadCodeAssist 请求
|
|
type LoadCodeAssistRequest struct {
|
|
Metadata struct {
|
|
IDEType string `json:"ideType"`
|
|
} `json:"metadata"`
|
|
}
|
|
|
|
// LoadCodeAssistResponse loadCodeAssist 响应
|
|
type LoadCodeAssistResponse struct {
|
|
CloudAICompanionProject string `json:"cloudaicompanionProject"`
|
|
}
|
|
|
|
// Client Antigravity API 客户端
|
|
type Client struct {
|
|
httpClient *http.Client
|
|
proxyURL string
|
|
}
|
|
|
|
func NewClient(proxyURL string) *Client {
|
|
client := &http.Client{
|
|
Timeout: 30 * time.Second,
|
|
}
|
|
|
|
if strings.TrimSpace(proxyURL) != "" {
|
|
if proxyURLParsed, err := url.Parse(proxyURL); err == nil {
|
|
client.Transport = &http.Transport{
|
|
Proxy: http.ProxyURL(proxyURLParsed),
|
|
}
|
|
}
|
|
}
|
|
|
|
return &Client{
|
|
httpClient: client,
|
|
proxyURL: proxyURL,
|
|
}
|
|
}
|
|
|
|
// ExchangeCode 用 authorization code 交换 token
|
|
func (c *Client) ExchangeCode(ctx context.Context, code, codeVerifier string) (*TokenResponse, error) {
|
|
params := url.Values{}
|
|
params.Set("client_id", ClientID)
|
|
params.Set("client_secret", ClientSecret)
|
|
params.Set("code", code)
|
|
params.Set("redirect_uri", RedirectURI)
|
|
params.Set("grant_type", "authorization_code")
|
|
params.Set("code_verifier", codeVerifier)
|
|
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodPost, TokenURL, strings.NewReader(params.Encode()))
|
|
if err != nil {
|
|
return nil, fmt.Errorf("创建请求失败: %w", err)
|
|
}
|
|
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
|
|
|
|
resp, err := c.httpClient.Do(req)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("Token 交换请求失败: %w", err)
|
|
}
|
|
defer func() { _ = resp.Body.Close() }()
|
|
|
|
bodyBytes, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("读取响应失败: %w", err)
|
|
}
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
return nil, fmt.Errorf("Token 交换失败 (HTTP %d): %s", resp.StatusCode, string(bodyBytes))
|
|
}
|
|
|
|
var tokenResp TokenResponse
|
|
if err := json.Unmarshal(bodyBytes, &tokenResp); err != nil {
|
|
return nil, fmt.Errorf("Token 解析失败: %w", err)
|
|
}
|
|
|
|
return &tokenResp, nil
|
|
}
|
|
|
|
// RefreshToken 刷新 access_token
|
|
func (c *Client) RefreshToken(ctx context.Context, refreshToken string) (*TokenResponse, error) {
|
|
params := url.Values{}
|
|
params.Set("client_id", ClientID)
|
|
params.Set("client_secret", ClientSecret)
|
|
params.Set("refresh_token", refreshToken)
|
|
params.Set("grant_type", "refresh_token")
|
|
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodPost, TokenURL, strings.NewReader(params.Encode()))
|
|
if err != nil {
|
|
return nil, fmt.Errorf("创建请求失败: %w", err)
|
|
}
|
|
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
|
|
|
|
resp, err := c.httpClient.Do(req)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("Token 刷新请求失败: %w", err)
|
|
}
|
|
defer func() { _ = resp.Body.Close() }()
|
|
|
|
bodyBytes, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("读取响应失败: %w", err)
|
|
}
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
return nil, fmt.Errorf("Token 刷新失败 (HTTP %d): %s", resp.StatusCode, string(bodyBytes))
|
|
}
|
|
|
|
var tokenResp TokenResponse
|
|
if err := json.Unmarshal(bodyBytes, &tokenResp); err != nil {
|
|
return nil, fmt.Errorf("Token 解析失败: %w", err)
|
|
}
|
|
|
|
return &tokenResp, nil
|
|
}
|
|
|
|
// GetUserInfo 获取用户信息
|
|
func (c *Client) GetUserInfo(ctx context.Context, accessToken string) (*UserInfo, error) {
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodGet, UserInfoURL, nil)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("创建请求失败: %w", err)
|
|
}
|
|
req.Header.Set("Authorization", "Bearer "+accessToken)
|
|
|
|
resp, err := c.httpClient.Do(req)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("用户信息请求失败: %w", err)
|
|
}
|
|
defer func() { _ = resp.Body.Close() }()
|
|
|
|
bodyBytes, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("读取响应失败: %w", err)
|
|
}
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
return nil, fmt.Errorf("获取用户信息失败 (HTTP %d): %s", resp.StatusCode, string(bodyBytes))
|
|
}
|
|
|
|
var userInfo UserInfo
|
|
if err := json.Unmarshal(bodyBytes, &userInfo); err != nil {
|
|
return nil, fmt.Errorf("用户信息解析失败: %w", err)
|
|
}
|
|
|
|
return &userInfo, nil
|
|
}
|
|
|
|
// LoadCodeAssist 获取 project_id
|
|
func (c *Client) LoadCodeAssist(ctx context.Context, accessToken string) (*LoadCodeAssistResponse, error) {
|
|
reqBody := LoadCodeAssistRequest{}
|
|
reqBody.Metadata.IDEType = "ANTIGRAVITY"
|
|
|
|
bodyBytes, err := json.Marshal(reqBody)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("序列化请求失败: %w", err)
|
|
}
|
|
|
|
url := BaseURL + "/v1internal:loadCodeAssist"
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, strings.NewReader(string(bodyBytes)))
|
|
if err != nil {
|
|
return nil, fmt.Errorf("创建请求失败: %w", err)
|
|
}
|
|
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 {
|
|
return nil, fmt.Errorf("loadCodeAssist 请求失败: %w", err)
|
|
}
|
|
defer func() { _ = resp.Body.Close() }()
|
|
|
|
respBodyBytes, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("读取响应失败: %w", err)
|
|
}
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
return nil, fmt.Errorf("loadCodeAssist 失败 (HTTP %d): %s", resp.StatusCode, string(respBodyBytes))
|
|
}
|
|
|
|
var loadResp LoadCodeAssistResponse
|
|
if err := json.Unmarshal(respBodyBytes, &loadResp); err != nil {
|
|
return nil, fmt.Errorf("响应解析失败: %w", err)
|
|
}
|
|
|
|
return &loadResp, nil
|
|
}
|
|
|
|
// ModelQuotaInfo 模型配额信息
|
|
type ModelQuotaInfo struct {
|
|
RemainingFraction float64 `json:"remainingFraction"`
|
|
ResetTime string `json:"resetTime,omitempty"`
|
|
}
|
|
|
|
// ModelInfo 模型信息
|
|
type ModelInfo struct {
|
|
QuotaInfo *ModelQuotaInfo `json:"quotaInfo,omitempty"`
|
|
}
|
|
|
|
// FetchAvailableModelsRequest fetchAvailableModels 请求
|
|
type FetchAvailableModelsRequest struct {
|
|
Project string `json:"project"`
|
|
}
|
|
|
|
// FetchAvailableModelsResponse fetchAvailableModels 响应
|
|
type FetchAvailableModelsResponse struct {
|
|
Models map[string]ModelInfo `json:"models"`
|
|
}
|
|
|
|
// FetchAvailableModels 获取可用模型和配额信息
|
|
func (c *Client) FetchAvailableModels(ctx context.Context, accessToken, projectID string) (*FetchAvailableModelsResponse, error) {
|
|
reqBody := FetchAvailableModelsRequest{Project: projectID}
|
|
bodyBytes, err := json.Marshal(reqBody)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("序列化请求失败: %w", err)
|
|
}
|
|
|
|
apiURL := BaseURL + "/v1internal:fetchAvailableModels"
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, strings.NewReader(string(bodyBytes)))
|
|
if err != nil {
|
|
return nil, fmt.Errorf("创建请求失败: %w", err)
|
|
}
|
|
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 {
|
|
return nil, fmt.Errorf("fetchAvailableModels 请求失败: %w", err)
|
|
}
|
|
defer func() { _ = resp.Body.Close() }()
|
|
|
|
respBodyBytes, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("读取响应失败: %w", err)
|
|
}
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
return nil, fmt.Errorf("fetchAvailableModels 失败 (HTTP %d): %s", resp.StatusCode, string(respBodyBytes))
|
|
}
|
|
|
|
var modelsResp FetchAvailableModelsResponse
|
|
if err := json.Unmarshal(respBodyBytes, &modelsResp); err != nil {
|
|
return nil, fmt.Errorf("响应解析失败: %w", err)
|
|
}
|
|
|
|
return &modelsResp, nil
|
|
}
|