| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910 |
- // Package antigravity provides a client for the Antigravity API.
- package antigravity
- import (
- "bytes"
- "context"
- "encoding/json"
- "errors"
- "fmt"
- "io"
- "log"
- "net"
- "net/http"
- "net/url"
- "strings"
- "time"
- "github.com/Wei-Shaw/sub2api/internal/pkg/proxyurl"
- "github.com/Wei-Shaw/sub2api/internal/pkg/proxyutil"
- "github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
- )
- // ForbiddenError 表示上游返回 403 Forbidden
- type ForbiddenError struct {
- StatusCode int
- Body string
- }
- func (e *ForbiddenError) Error() string {
- return fmt.Sprintf("fetchAvailableModels 失败 (HTTP %d): %s", e.StatusCode, e.Body)
- }
- // NewAPIRequestWithURL 使用指定的 base URL 创建 Antigravity API 请求(v1internal 端点)
- func NewAPIRequestWithURL(ctx context.Context, baseURL, action, accessToken string, body []byte) (*http.Request, error) {
- // 构建 URL,流式请求添加 ?alt=sse 参数
- apiURL := fmt.Sprintf("%s/v1internal:%s", baseURL, action)
- isStream := action == "streamGenerateContent"
- if isStream {
- apiURL += "?alt=sse"
- }
- req, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, bytes.NewReader(body))
- if err != nil {
- return nil, err
- }
- // 基础 Headers(与 Antigravity-Manager 保持一致,只设置这 3 个)
- req.Header.Set("Content-Type", "application/json")
- req.Header.Set("Authorization", "Bearer "+accessToken)
- req.Header.Set("User-Agent", GetUserAgentForContext(ctx))
- return req, nil
- }
- // NewAPIRequest 使用默认 URL 创建 Antigravity API 请求(v1internal 端点)
- // 向后兼容:仅使用默认 BaseURL
- func NewAPIRequest(ctx context.Context, action, accessToken string, body []byte) (*http.Request, error) {
- return NewAPIRequestWithURL(ctx, BaseURL, action, accessToken, body)
- }
- // 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"`
- IDEVersion string `json:"ideVersion"`
- IDEName string `json:"ideName"`
- } `json:"metadata"`
- }
- // TierInfo 账户类型信息
- type TierInfo struct {
- ID string `json:"id"` // free-tier, g1-pro-tier, g1-ultra-tier
- Name string `json:"name"` // 显示名称
- Description string `json:"description"` // 描述
- }
- // UnmarshalJSON supports both legacy string tiers and object tiers.
- func (t *TierInfo) UnmarshalJSON(data []byte) error {
- data = bytes.TrimSpace(data)
- if len(data) == 0 || string(data) == "null" {
- return nil
- }
- if data[0] == '"' {
- var id string
- if err := json.Unmarshal(data, &id); err != nil {
- return err
- }
- t.ID = id
- return nil
- }
- type alias TierInfo
- var decoded alias
- if err := json.Unmarshal(data, &decoded); err != nil {
- return err
- }
- *t = TierInfo(decoded)
- return nil
- }
- // IneligibleTier 不符合条件的层级信息
- type IneligibleTier struct {
- Tier *TierInfo `json:"tier,omitempty"`
- // ReasonCode 不符合条件的原因代码,如 INELIGIBLE_ACCOUNT
- ReasonCode string `json:"reasonCode,omitempty"`
- ReasonMessage string `json:"reasonMessage,omitempty"`
- }
- // LoadCodeAssistResponse loadCodeAssist 响应
- type LoadCodeAssistResponse struct {
- CloudAICompanionProject string `json:"cloudaicompanionProject"`
- CurrentTier *TierInfo `json:"currentTier,omitempty"`
- PaidTier *PaidTierInfo `json:"paidTier,omitempty"`
- IneligibleTiers []*IneligibleTier `json:"ineligibleTiers,omitempty"`
- }
- // PaidTierInfo 付费等级信息,包含 AI Credits 余额。
- type PaidTierInfo struct {
- ID string `json:"id"`
- Name string `json:"name"`
- Description string `json:"description"`
- AvailableCredits []AvailableCredit `json:"availableCredits,omitempty"`
- }
- // UnmarshalJSON 兼容 paidTier 既可能是字符串也可能是对象的情况。
- func (p *PaidTierInfo) UnmarshalJSON(data []byte) error {
- data = bytes.TrimSpace(data)
- if len(data) == 0 || string(data) == "null" {
- return nil
- }
- if data[0] == '"' {
- var id string
- if err := json.Unmarshal(data, &id); err != nil {
- return err
- }
- p.ID = id
- return nil
- }
- type alias PaidTierInfo
- var raw alias
- if err := json.Unmarshal(data, &raw); err != nil {
- return err
- }
- *p = PaidTierInfo(raw)
- return nil
- }
- // AvailableCredit 表示一条 AI Credits 余额记录。
- type AvailableCredit struct {
- CreditType string `json:"creditType,omitempty"`
- CreditAmount string `json:"creditAmount,omitempty"`
- MinimumCreditAmountForUsage string `json:"minimumCreditAmountForUsage,omitempty"`
- }
- // GetAmount 将 creditAmount 解析为浮点数。
- func (c *AvailableCredit) GetAmount() float64 {
- if c.CreditAmount == "" {
- return 0
- }
- var value float64
- _, _ = fmt.Sscanf(c.CreditAmount, "%f", &value)
- return value
- }
- // GetMinimumAmount 将 minimumCreditAmountForUsage 解析为浮点数。
- func (c *AvailableCredit) GetMinimumAmount() float64 {
- if c.MinimumCreditAmountForUsage == "" {
- return 0
- }
- var value float64
- _, _ = fmt.Sscanf(c.MinimumCreditAmountForUsage, "%f", &value)
- return value
- }
- // OnboardUserRequest onboardUser 请求
- type OnboardUserRequest struct {
- TierID string `json:"tierId"`
- Metadata struct {
- IDEType string `json:"ideType"`
- Platform string `json:"platform,omitempty"`
- PluginType string `json:"pluginType,omitempty"`
- } `json:"metadata"`
- }
- // OnboardUserResponse onboardUser 响应
- type OnboardUserResponse struct {
- Name string `json:"name,omitempty"`
- Done bool `json:"done"`
- Response map[string]any `json:"response,omitempty"`
- }
- // GetTier 获取账户类型
- // 优先返回 paidTier(付费订阅级别),否则返回 currentTier
- func (r *LoadCodeAssistResponse) GetTier() string {
- if r.PaidTier != nil && r.PaidTier.ID != "" {
- return r.PaidTier.ID
- }
- if r.CurrentTier != nil {
- return r.CurrentTier.ID
- }
- return ""
- }
- // GetAvailableCredits 返回 paid tier 中的 AI Credits 余额列表。
- func (r *LoadCodeAssistResponse) GetAvailableCredits() []AvailableCredit {
- if r.PaidTier == nil {
- return nil
- }
- return r.PaidTier.AvailableCredits
- }
- // TierIDToPlanType 将 tier ID 映射为用户可见的套餐名。
- func TierIDToPlanType(tierID string) string {
- switch strings.ToLower(strings.TrimSpace(tierID)) {
- case "free-tier":
- return "Free"
- case "g1-pro-tier":
- return "Pro"
- case "g1-ultra-tier":
- return "Ultra"
- default:
- if tierID == "" {
- return "Free"
- }
- return tierID
- }
- }
- // Client Antigravity API 客户端
- type Client struct {
- httpClient *http.Client
- }
- const (
- // proxyDialTimeout 代理 TCP 连接超时(含代理握手),代理不通时快速失败
- proxyDialTimeout = 5 * time.Second
- // proxyTLSHandshakeTimeout 代理 TLS 握手超时
- proxyTLSHandshakeTimeout = 5 * time.Second
- // clientTimeout 整体请求超时(含连接、发送、等待响应、读取 body)
- clientTimeout = 10 * time.Second
- )
- func NewClient(proxyURL string) (*Client, error) {
- client := &http.Client{
- Timeout: clientTimeout,
- }
- _, parsed, err := proxyurl.Parse(proxyURL)
- if err != nil {
- return nil, err
- }
- if parsed != nil {
- transport := &http.Transport{
- DialContext: (&net.Dialer{
- Timeout: proxyDialTimeout,
- }).DialContext,
- TLSHandshakeTimeout: proxyTLSHandshakeTimeout,
- }
- if err := proxyutil.ConfigureTransportProxy(transport, parsed); err != nil {
- return nil, fmt.Errorf("configure proxy: %w", err)
- }
- client.Transport = transport
- }
- return &Client{
- httpClient: client,
- }, nil
- }
- // IsConnectionError 判断是否为连接错误(网络超时、DNS 失败、连接拒绝)
- func IsConnectionError(err error) bool {
- if err == nil {
- return false
- }
- // 检查超时错误
- var netErr net.Error
- if errors.As(err, &netErr) && netErr.Timeout() {
- return true
- }
- // 检查连接错误(DNS 失败、连接拒绝)
- var opErr *net.OpError
- if errors.As(err, &opErr) {
- return true
- }
- // 检查 URL 错误
- var urlErr *url.Error
- return errors.As(err, &urlErr)
- }
- // shouldFallbackToNextURL 判断是否应切换到下一个 URL
- // 与 Antigravity-Manager 保持一致:连接错误、429、408、404、5xx 触发 URL 降级
- func shouldFallbackToNextURL(err error, statusCode int) bool {
- if IsConnectionError(err) {
- return true
- }
- return statusCode == http.StatusTooManyRequests ||
- statusCode == http.StatusRequestTimeout ||
- statusCode == http.StatusNotFound ||
- statusCode >= 500
- }
- // ExchangeCode 用 authorization code 交换 token
- func (c *Client) ExchangeCode(ctx context.Context, code, codeVerifier string) (*TokenResponse, error) {
- clientSecret, err := getClientSecret()
- if err != nil {
- return nil, err
- }
- 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 := servertiming.Do(c.httpClient, 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) {
- clientSecret, err := getClientSecret()
- if err != nil {
- return nil, err
- }
- 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 := servertiming.Do(c.httpClient, 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 := servertiming.Do(c.httpClient, 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 获取账户信息,返回解析后的结构体和原始 JSON
- // 支持 URL fallback:sandbox → daily → prod
- func (c *Client) LoadCodeAssist(ctx context.Context, accessToken string) (*LoadCodeAssistResponse, map[string]any, error) {
- reqBody := LoadCodeAssistRequest{}
- reqBody.Metadata.IDEType = "ANTIGRAVITY"
- reqBody.Metadata.IDEVersion = GetUserAgentVersionForContext(ctx)
- reqBody.Metadata.IDEName = "antigravity"
- bodyBytes, err := json.Marshal(reqBody)
- if err != nil {
- return nil, nil, fmt.Errorf("序列化请求失败: %w", err)
- }
- // 固定顺序:prod -> daily
- availableURLs := BaseURLs
- var lastErr error
- for urlIdx, baseURL := range availableURLs {
- apiURL := baseURL + "/v1internal:loadCodeAssist"
- req, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, strings.NewReader(string(bodyBytes)))
- if err != nil {
- lastErr = fmt.Errorf("创建请求失败: %w", err)
- continue
- }
- req.Header.Set("Authorization", "Bearer "+accessToken)
- req.Header.Set("Content-Type", "application/json")
- req.Header.Set("User-Agent", GetUserAgentForContext(ctx))
- resp, err := servertiming.Do(c.httpClient, req)
- if err != nil {
- lastErr = fmt.Errorf("loadCodeAssist 请求失败: %w", err)
- if shouldFallbackToNextURL(err, 0) && urlIdx < len(availableURLs)-1 {
- log.Printf("[antigravity] loadCodeAssist URL fallback: %s -> %s", baseURL, availableURLs[urlIdx+1])
- continue
- }
- return nil, nil, lastErr
- }
- respBodyBytes, err := io.ReadAll(resp.Body)
- _ = resp.Body.Close() // 立即关闭,避免循环内 defer 导致的资源泄漏
- if err != nil {
- return nil, nil, fmt.Errorf("读取响应失败: %w", err)
- }
- // 检查是否需要 URL 降级
- if shouldFallbackToNextURL(nil, resp.StatusCode) && urlIdx < len(availableURLs)-1 {
- log.Printf("[antigravity] loadCodeAssist URL fallback (HTTP %d): %s -> %s", resp.StatusCode, baseURL, availableURLs[urlIdx+1])
- continue
- }
- if resp.StatusCode != http.StatusOK {
- return nil, nil, fmt.Errorf("loadCodeAssist 失败 (HTTP %d): %s", resp.StatusCode, string(respBodyBytes))
- }
- var loadResp LoadCodeAssistResponse
- if err := json.Unmarshal(respBodyBytes, &loadResp); err != nil {
- return nil, nil, fmt.Errorf("响应解析失败: %w", err)
- }
- // 解析原始 JSON 为 map
- var rawResp map[string]any
- _ = json.Unmarshal(respBodyBytes, &rawResp)
- // 标记成功的 URL,下次优先使用
- DefaultURLAvailability.MarkSuccess(baseURL)
- return &loadResp, rawResp, nil
- }
- return nil, nil, lastErr
- }
- // OnboardUser 触发账号 onboarding,并返回 project_id
- // 说明:
- // 1) 部分账号 loadCodeAssist 不会立即返回 cloudaicompanionProject;
- // 2) 这时需要调用 onboardUser 完成初始化,之后才能拿到 project_id。
- func (c *Client) OnboardUser(ctx context.Context, accessToken, tierID string) (string, error) {
- tierID = strings.TrimSpace(tierID)
- if tierID == "" {
- return "", fmt.Errorf("tier_id 为空")
- }
- reqBody := OnboardUserRequest{TierID: tierID}
- reqBody.Metadata.IDEType = "ANTIGRAVITY"
- reqBody.Metadata.Platform = "PLATFORM_UNSPECIFIED"
- reqBody.Metadata.PluginType = "GEMINI"
- bodyBytes, err := json.Marshal(reqBody)
- if err != nil {
- return "", fmt.Errorf("序列化请求失败: %w", err)
- }
- availableURLs := BaseURLs
- var lastErr error
- for urlIdx, baseURL := range availableURLs {
- apiURL := baseURL + "/v1internal:onboardUser"
- for attempt := 1; attempt <= 5; attempt++ {
- req, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, bytes.NewReader(bodyBytes))
- if err != nil {
- lastErr = fmt.Errorf("创建请求失败: %w", err)
- break
- }
- req.Header.Set("Authorization", "Bearer "+accessToken)
- req.Header.Set("Content-Type", "application/json")
- req.Header.Set("User-Agent", GetUserAgentForContext(ctx))
- resp, err := servertiming.Do(c.httpClient, req)
- if err != nil {
- lastErr = fmt.Errorf("onboardUser 请求失败: %w", err)
- if shouldFallbackToNextURL(err, 0) && urlIdx < len(availableURLs)-1 {
- log.Printf("[antigravity] onboardUser URL fallback: %s -> %s", baseURL, availableURLs[urlIdx+1])
- break
- }
- return "", lastErr
- }
- respBodyBytes, err := io.ReadAll(resp.Body)
- _ = resp.Body.Close()
- if err != nil {
- return "", fmt.Errorf("读取响应失败: %w", err)
- }
- if shouldFallbackToNextURL(nil, resp.StatusCode) && urlIdx < len(availableURLs)-1 {
- log.Printf("[antigravity] onboardUser URL fallback (HTTP %d): %s -> %s", resp.StatusCode, baseURL, availableURLs[urlIdx+1])
- break
- }
- if resp.StatusCode != http.StatusOK {
- lastErr = fmt.Errorf("onboardUser 失败 (HTTP %d): %s", resp.StatusCode, string(respBodyBytes))
- return "", lastErr
- }
- var onboardResp OnboardUserResponse
- if err := json.Unmarshal(respBodyBytes, &onboardResp); err != nil {
- lastErr = fmt.Errorf("onboardUser 响应解析失败: %w", err)
- return "", lastErr
- }
- if onboardResp.Done {
- if projectID := extractProjectIDFromOnboardResponse(onboardResp.Response); projectID != "" {
- DefaultURLAvailability.MarkSuccess(baseURL)
- return projectID, nil
- }
- lastErr = fmt.Errorf("onboardUser 完成但未返回 project_id")
- return "", lastErr
- }
- // done=false 时等待后重试,避免在上游仍未完成时提前结束轮询。
- select {
- case <-time.After(2 * time.Second):
- case <-ctx.Done():
- return "", ctx.Err()
- }
- }
- }
- if lastErr != nil {
- return "", lastErr
- }
- return "", fmt.Errorf("onboardUser 未返回 project_id")
- }
- func extractProjectIDFromOnboardResponse(resp map[string]any) string {
- if len(resp) == 0 {
- return ""
- }
- if v, ok := resp["cloudaicompanionProject"]; ok {
- switch project := v.(type) {
- case string:
- return strings.TrimSpace(project)
- case map[string]any:
- if id, ok := project["id"].(string); ok {
- return strings.TrimSpace(id)
- }
- }
- }
- return ""
- }
- // ModelQuotaInfo 模型配额信息
- type ModelQuotaInfo struct {
- RemainingFraction float64 `json:"remainingFraction"`
- ResetTime string `json:"resetTime,omitempty"`
- }
- // ModelInfo 模型信息
- type ModelInfo struct {
- QuotaInfo *ModelQuotaInfo `json:"quotaInfo,omitempty"`
- DisplayName string `json:"displayName,omitempty"`
- SupportsImages *bool `json:"supportsImages,omitempty"`
- SupportsThinking *bool `json:"supportsThinking,omitempty"`
- ThinkingBudget *int `json:"thinkingBudget,omitempty"`
- Recommended *bool `json:"recommended,omitempty"`
- MaxTokens *int `json:"maxTokens,omitempty"`
- MaxOutputTokens *int `json:"maxOutputTokens,omitempty"`
- SupportedMimeTypes map[string]bool `json:"supportedMimeTypes,omitempty"`
- }
- // DeprecatedModelInfo 废弃模型转发信息
- type DeprecatedModelInfo struct {
- NewModelID string `json:"newModelId"`
- }
- // FetchAvailableModelsRequest fetchAvailableModels 请求
- type FetchAvailableModelsRequest struct {
- Project string `json:"project"`
- }
- // FetchAvailableModelsResponse fetchAvailableModels 响应
- type FetchAvailableModelsResponse struct {
- Models map[string]ModelInfo `json:"models"`
- DeprecatedModelIDs map[string]DeprecatedModelInfo `json:"deprecatedModelIds,omitempty"`
- }
- // FetchAvailableModels 获取可用模型和配额信息,返回解析后的结构体和原始 JSON
- // 支持 URL fallback:sandbox → daily → prod
- func (c *Client) FetchAvailableModels(ctx context.Context, accessToken, projectID string, bodyLimit int64) (*FetchAvailableModelsResponse, map[string]any, error) {
- if c == nil || c.httpClient == nil {
- return nil, nil, errors.New("antigravity client is not configured")
- }
- if bodyLimit <= 0 {
- return nil, nil, errors.New("fetchAvailableModels body limit must be positive")
- }
- reqBody := FetchAvailableModelsRequest{Project: projectID}
- bodyBytes, err := json.Marshal(reqBody)
- if err != nil {
- return nil, nil, fmt.Errorf("序列化请求失败: %w", err)
- }
- // 固定顺序:prod -> daily
- availableURLs := BaseURLs
- fetchClient := c.fetchAvailableModelsHTTPClient()
- var lastErr error
- for urlIdx, baseURL := range availableURLs {
- apiURL := baseURL + "/v1internal:fetchAvailableModels"
- req, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, strings.NewReader(string(bodyBytes)))
- if err != nil {
- lastErr = fmt.Errorf("创建请求失败: %w", err)
- continue
- }
- req.Header.Set("Authorization", "Bearer "+accessToken)
- req.Header.Set("Content-Type", "application/json")
- req.Header.Set("User-Agent", GetUserAgentForContext(ctx))
- resp, err := servertiming.Do(fetchClient, req)
- if err != nil {
- lastErr = fmt.Errorf("fetchAvailableModels 请求失败: %w", err)
- if shouldFallbackToNextURL(err, 0) && urlIdx < len(availableURLs)-1 {
- log.Printf("[antigravity] fetchAvailableModels URL fallback: %s -> %s", baseURL, availableURLs[urlIdx+1])
- continue
- }
- return nil, nil, lastErr
- }
- respBodyBytes, err := io.ReadAll(io.LimitReader(resp.Body, bodyLimit+1))
- _ = resp.Body.Close() // 立即关闭,避免循环内 defer 导致的资源泄漏
- if err != nil {
- return nil, nil, fmt.Errorf("读取响应失败: %w", err)
- }
- if int64(len(respBodyBytes)) > bodyLimit {
- return nil, nil, fmt.Errorf("响应超过 %d 字节", bodyLimit)
- }
- // 检查是否需要 URL 降级
- if shouldFallbackToNextURL(nil, resp.StatusCode) && urlIdx < len(availableURLs)-1 {
- log.Printf("[antigravity] fetchAvailableModels URL fallback (HTTP %d): %s -> %s", resp.StatusCode, baseURL, availableURLs[urlIdx+1])
- continue
- }
- if resp.StatusCode == http.StatusForbidden {
- return nil, nil, &ForbiddenError{
- StatusCode: resp.StatusCode,
- Body: string(respBodyBytes),
- }
- }
- if resp.StatusCode != http.StatusOK {
- return nil, nil, fmt.Errorf("fetchAvailableModels 失败 (HTTP %d): %s", resp.StatusCode, string(respBodyBytes))
- }
- var modelsResp FetchAvailableModelsResponse
- if err := json.Unmarshal(respBodyBytes, &modelsResp); err != nil {
- return nil, nil, fmt.Errorf("响应解析失败: %w", err)
- }
- // 解析原始 JSON 为 map
- var rawResp map[string]any
- _ = json.Unmarshal(respBodyBytes, &rawResp)
- // 标记成功的 URL,下次优先使用
- DefaultURLAvailability.MarkSuccess(baseURL)
- return &modelsResp, rawResp, nil
- }
- return nil, nil, lastErr
- }
- func (c *Client) fetchAvailableModelsHTTPClient() *http.Client {
- fetchClient := *c.httpClient
- fetchClient.CheckRedirect = checkFetchAvailableModelsRedirect
- return &fetchClient
- }
- func checkFetchAvailableModelsRedirect(req *http.Request, via []*http.Request) error {
- if len(via) >= 10 {
- return errors.New("stopped after 10 redirects")
- }
- if req == nil || req.URL == nil {
- return errors.New("redirect url is nil")
- }
- if !isAllowedFetchAvailableModelsRedirectHost(req.URL.Hostname()) {
- return fmt.Errorf("redirect to unsupported host: %s", req.URL.Hostname())
- }
- return nil
- }
- func isAllowedFetchAvailableModelsRedirectHost(host string) bool {
- host = strings.ToLower(strings.TrimSpace(host))
- if host == "" {
- return false
- }
- for _, baseURL := range BaseURLs {
- parsed, err := url.Parse(baseURL)
- if err != nil {
- continue
- }
- if strings.EqualFold(host, parsed.Hostname()) {
- return true
- }
- }
- return false
- }
- // ── Privacy API ──────────────────────────────────────────────────────
- // privacyBaseURL 隐私设置 API 仅使用 daily 端点(与 Antigravity 客户端行为一致)
- const privacyBaseURL = antigravityDailyBaseURL
- // SetUserSettingsRequest setUserSettings 请求体
- type SetUserSettingsRequest struct {
- UserSettings map[string]any `json:"user_settings"`
- }
- // FetchUserInfoRequest fetchUserInfo 请求体
- type FetchUserInfoRequest struct {
- Project string `json:"project"`
- }
- // FetchUserInfoResponse fetchUserInfo 响应体
- type FetchUserInfoResponse struct {
- UserSettings map[string]any `json:"userSettings,omitempty"`
- RegionCode string `json:"regionCode,omitempty"`
- }
- // IsPrivate 判断隐私是否已设置:userSettings 为空或不含 telemetryEnabled 表示已设置
- func (r *FetchUserInfoResponse) IsPrivate() bool {
- if r == nil || r.UserSettings == nil {
- return true
- }
- _, hasTelemetry := r.UserSettings["telemetryEnabled"]
- return !hasTelemetry
- }
- // SetUserSettingsResponse setUserSettings 响应体
- type SetUserSettingsResponse struct {
- UserSettings map[string]any `json:"userSettings,omitempty"`
- }
- // IsSuccess 判断 setUserSettings 是否成功:返回 {"userSettings":{}} 且无 telemetryEnabled
- func (r *SetUserSettingsResponse) IsSuccess() bool {
- if r == nil {
- return false
- }
- // userSettings 为 nil 或空 map 均视为成功
- if len(r.UserSettings) == 0 {
- return true
- }
- // 如果包含 telemetryEnabled 字段,说明未成功清除
- _, hasTelemetry := r.UserSettings["telemetryEnabled"]
- return !hasTelemetry
- }
- // SetUserSettings 调用 setUserSettings API 设置用户隐私,返回解析后的响应
- func (c *Client) SetUserSettings(ctx context.Context, accessToken string) (*SetUserSettingsResponse, error) {
- // 发送空 user_settings 以清除隐私设置
- payload := SetUserSettingsRequest{UserSettings: map[string]any{}}
- bodyBytes, err := json.Marshal(payload)
- if err != nil {
- return nil, fmt.Errorf("序列化请求失败: %w", err)
- }
- apiURL := privacyBaseURL + "/v1internal:setUserSettings"
- req, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, bytes.NewReader(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("Accept", "*/*")
- req.Header.Set("User-Agent", GetUserAgentForContext(ctx))
- req.Header.Set("X-Goog-Api-Client", "gl-node/22.21.1")
- req.Host = "daily-cloudcode-pa.googleapis.com"
- resp, err := servertiming.Do(c.httpClient, req)
- if err != nil {
- return nil, fmt.Errorf("setUserSettings 请求失败: %w", err)
- }
- defer func() { _ = resp.Body.Close() }()
- respBody, err := io.ReadAll(resp.Body)
- if err != nil {
- return nil, fmt.Errorf("读取响应失败: %w", err)
- }
- if resp.StatusCode != http.StatusOK {
- return nil, fmt.Errorf("setUserSettings 失败 (HTTP %d): %s", resp.StatusCode, string(respBody))
- }
- var result SetUserSettingsResponse
- if err := json.Unmarshal(respBody, &result); err != nil {
- return nil, fmt.Errorf("响应解析失败: %w", err)
- }
- return &result, nil
- }
- // FetchUserInfo 调用 fetchUserInfo API 获取用户隐私设置状态
- func (c *Client) FetchUserInfo(ctx context.Context, accessToken, projectID string) (*FetchUserInfoResponse, error) {
- reqBody := FetchUserInfoRequest{Project: projectID}
- bodyBytes, err := json.Marshal(reqBody)
- if err != nil {
- return nil, fmt.Errorf("序列化请求失败: %w", err)
- }
- apiURL := privacyBaseURL + "/v1internal:fetchUserInfo"
- req, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, bytes.NewReader(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("Accept", "*/*")
- req.Header.Set("User-Agent", GetUserAgentForContext(ctx))
- req.Header.Set("X-Goog-Api-Client", "gl-node/22.21.1")
- req.Host = "daily-cloudcode-pa.googleapis.com"
- resp, err := servertiming.Do(c.httpClient, req)
- if err != nil {
- return nil, fmt.Errorf("fetchUserInfo 请求失败: %w", err)
- }
- defer func() { _ = resp.Body.Close() }()
- respBody, err := io.ReadAll(resp.Body)
- if err != nil {
- return nil, fmt.Errorf("读取响应失败: %w", err)
- }
- if resp.StatusCode != http.StatusOK {
- return nil, fmt.Errorf("fetchUserInfo 失败 (HTTP %d): %s", resp.StatusCode, string(respBody))
- }
- var result FetchUserInfoResponse
- if err := json.Unmarshal(respBody, &result); err != nil {
- return nil, fmt.Errorf("响应解析失败: %w", err)
- }
- return &result, nil
- }
|