client.go 28 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910
  1. // Package antigravity provides a client for the Antigravity API.
  2. package antigravity
  3. import (
  4. "bytes"
  5. "context"
  6. "encoding/json"
  7. "errors"
  8. "fmt"
  9. "io"
  10. "log"
  11. "net"
  12. "net/http"
  13. "net/url"
  14. "strings"
  15. "time"
  16. "github.com/Wei-Shaw/sub2api/internal/pkg/proxyurl"
  17. "github.com/Wei-Shaw/sub2api/internal/pkg/proxyutil"
  18. "github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
  19. )
  20. // ForbiddenError 表示上游返回 403 Forbidden
  21. type ForbiddenError struct {
  22. StatusCode int
  23. Body string
  24. }
  25. func (e *ForbiddenError) Error() string {
  26. return fmt.Sprintf("fetchAvailableModels 失败 (HTTP %d): %s", e.StatusCode, e.Body)
  27. }
  28. // NewAPIRequestWithURL 使用指定的 base URL 创建 Antigravity API 请求(v1internal 端点)
  29. func NewAPIRequestWithURL(ctx context.Context, baseURL, action, accessToken string, body []byte) (*http.Request, error) {
  30. // 构建 URL,流式请求添加 ?alt=sse 参数
  31. apiURL := fmt.Sprintf("%s/v1internal:%s", baseURL, action)
  32. isStream := action == "streamGenerateContent"
  33. if isStream {
  34. apiURL += "?alt=sse"
  35. }
  36. req, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, bytes.NewReader(body))
  37. if err != nil {
  38. return nil, err
  39. }
  40. // 基础 Headers(与 Antigravity-Manager 保持一致,只设置这 3 个)
  41. req.Header.Set("Content-Type", "application/json")
  42. req.Header.Set("Authorization", "Bearer "+accessToken)
  43. req.Header.Set("User-Agent", GetUserAgentForContext(ctx))
  44. return req, nil
  45. }
  46. // NewAPIRequest 使用默认 URL 创建 Antigravity API 请求(v1internal 端点)
  47. // 向后兼容:仅使用默认 BaseURL
  48. func NewAPIRequest(ctx context.Context, action, accessToken string, body []byte) (*http.Request, error) {
  49. return NewAPIRequestWithURL(ctx, BaseURL, action, accessToken, body)
  50. }
  51. // TokenResponse Google OAuth token 响应
  52. type TokenResponse struct {
  53. AccessToken string `json:"access_token"`
  54. ExpiresIn int64 `json:"expires_in"`
  55. TokenType string `json:"token_type"`
  56. Scope string `json:"scope,omitempty"`
  57. RefreshToken string `json:"refresh_token,omitempty"`
  58. }
  59. // UserInfo Google 用户信息
  60. type UserInfo struct {
  61. Email string `json:"email"`
  62. Name string `json:"name,omitempty"`
  63. GivenName string `json:"given_name,omitempty"`
  64. FamilyName string `json:"family_name,omitempty"`
  65. Picture string `json:"picture,omitempty"`
  66. }
  67. // LoadCodeAssistRequest loadCodeAssist 请求
  68. type LoadCodeAssistRequest struct {
  69. Metadata struct {
  70. IDEType string `json:"ideType"`
  71. IDEVersion string `json:"ideVersion"`
  72. IDEName string `json:"ideName"`
  73. } `json:"metadata"`
  74. }
  75. // TierInfo 账户类型信息
  76. type TierInfo struct {
  77. ID string `json:"id"` // free-tier, g1-pro-tier, g1-ultra-tier
  78. Name string `json:"name"` // 显示名称
  79. Description string `json:"description"` // 描述
  80. }
  81. // UnmarshalJSON supports both legacy string tiers and object tiers.
  82. func (t *TierInfo) UnmarshalJSON(data []byte) error {
  83. data = bytes.TrimSpace(data)
  84. if len(data) == 0 || string(data) == "null" {
  85. return nil
  86. }
  87. if data[0] == '"' {
  88. var id string
  89. if err := json.Unmarshal(data, &id); err != nil {
  90. return err
  91. }
  92. t.ID = id
  93. return nil
  94. }
  95. type alias TierInfo
  96. var decoded alias
  97. if err := json.Unmarshal(data, &decoded); err != nil {
  98. return err
  99. }
  100. *t = TierInfo(decoded)
  101. return nil
  102. }
  103. // IneligibleTier 不符合条件的层级信息
  104. type IneligibleTier struct {
  105. Tier *TierInfo `json:"tier,omitempty"`
  106. // ReasonCode 不符合条件的原因代码,如 INELIGIBLE_ACCOUNT
  107. ReasonCode string `json:"reasonCode,omitempty"`
  108. ReasonMessage string `json:"reasonMessage,omitempty"`
  109. }
  110. // LoadCodeAssistResponse loadCodeAssist 响应
  111. type LoadCodeAssistResponse struct {
  112. CloudAICompanionProject string `json:"cloudaicompanionProject"`
  113. CurrentTier *TierInfo `json:"currentTier,omitempty"`
  114. PaidTier *PaidTierInfo `json:"paidTier,omitempty"`
  115. IneligibleTiers []*IneligibleTier `json:"ineligibleTiers,omitempty"`
  116. }
  117. // PaidTierInfo 付费等级信息,包含 AI Credits 余额。
  118. type PaidTierInfo struct {
  119. ID string `json:"id"`
  120. Name string `json:"name"`
  121. Description string `json:"description"`
  122. AvailableCredits []AvailableCredit `json:"availableCredits,omitempty"`
  123. }
  124. // UnmarshalJSON 兼容 paidTier 既可能是字符串也可能是对象的情况。
  125. func (p *PaidTierInfo) UnmarshalJSON(data []byte) error {
  126. data = bytes.TrimSpace(data)
  127. if len(data) == 0 || string(data) == "null" {
  128. return nil
  129. }
  130. if data[0] == '"' {
  131. var id string
  132. if err := json.Unmarshal(data, &id); err != nil {
  133. return err
  134. }
  135. p.ID = id
  136. return nil
  137. }
  138. type alias PaidTierInfo
  139. var raw alias
  140. if err := json.Unmarshal(data, &raw); err != nil {
  141. return err
  142. }
  143. *p = PaidTierInfo(raw)
  144. return nil
  145. }
  146. // AvailableCredit 表示一条 AI Credits 余额记录。
  147. type AvailableCredit struct {
  148. CreditType string `json:"creditType,omitempty"`
  149. CreditAmount string `json:"creditAmount,omitempty"`
  150. MinimumCreditAmountForUsage string `json:"minimumCreditAmountForUsage,omitempty"`
  151. }
  152. // GetAmount 将 creditAmount 解析为浮点数。
  153. func (c *AvailableCredit) GetAmount() float64 {
  154. if c.CreditAmount == "" {
  155. return 0
  156. }
  157. var value float64
  158. _, _ = fmt.Sscanf(c.CreditAmount, "%f", &value)
  159. return value
  160. }
  161. // GetMinimumAmount 将 minimumCreditAmountForUsage 解析为浮点数。
  162. func (c *AvailableCredit) GetMinimumAmount() float64 {
  163. if c.MinimumCreditAmountForUsage == "" {
  164. return 0
  165. }
  166. var value float64
  167. _, _ = fmt.Sscanf(c.MinimumCreditAmountForUsage, "%f", &value)
  168. return value
  169. }
  170. // OnboardUserRequest onboardUser 请求
  171. type OnboardUserRequest struct {
  172. TierID string `json:"tierId"`
  173. Metadata struct {
  174. IDEType string `json:"ideType"`
  175. Platform string `json:"platform,omitempty"`
  176. PluginType string `json:"pluginType,omitempty"`
  177. } `json:"metadata"`
  178. }
  179. // OnboardUserResponse onboardUser 响应
  180. type OnboardUserResponse struct {
  181. Name string `json:"name,omitempty"`
  182. Done bool `json:"done"`
  183. Response map[string]any `json:"response,omitempty"`
  184. }
  185. // GetTier 获取账户类型
  186. // 优先返回 paidTier(付费订阅级别),否则返回 currentTier
  187. func (r *LoadCodeAssistResponse) GetTier() string {
  188. if r.PaidTier != nil && r.PaidTier.ID != "" {
  189. return r.PaidTier.ID
  190. }
  191. if r.CurrentTier != nil {
  192. return r.CurrentTier.ID
  193. }
  194. return ""
  195. }
  196. // GetAvailableCredits 返回 paid tier 中的 AI Credits 余额列表。
  197. func (r *LoadCodeAssistResponse) GetAvailableCredits() []AvailableCredit {
  198. if r.PaidTier == nil {
  199. return nil
  200. }
  201. return r.PaidTier.AvailableCredits
  202. }
  203. // TierIDToPlanType 将 tier ID 映射为用户可见的套餐名。
  204. func TierIDToPlanType(tierID string) string {
  205. switch strings.ToLower(strings.TrimSpace(tierID)) {
  206. case "free-tier":
  207. return "Free"
  208. case "g1-pro-tier":
  209. return "Pro"
  210. case "g1-ultra-tier":
  211. return "Ultra"
  212. default:
  213. if tierID == "" {
  214. return "Free"
  215. }
  216. return tierID
  217. }
  218. }
  219. // Client Antigravity API 客户端
  220. type Client struct {
  221. httpClient *http.Client
  222. }
  223. const (
  224. // proxyDialTimeout 代理 TCP 连接超时(含代理握手),代理不通时快速失败
  225. proxyDialTimeout = 5 * time.Second
  226. // proxyTLSHandshakeTimeout 代理 TLS 握手超时
  227. proxyTLSHandshakeTimeout = 5 * time.Second
  228. // clientTimeout 整体请求超时(含连接、发送、等待响应、读取 body)
  229. clientTimeout = 10 * time.Second
  230. )
  231. func NewClient(proxyURL string) (*Client, error) {
  232. client := &http.Client{
  233. Timeout: clientTimeout,
  234. }
  235. _, parsed, err := proxyurl.Parse(proxyURL)
  236. if err != nil {
  237. return nil, err
  238. }
  239. if parsed != nil {
  240. transport := &http.Transport{
  241. DialContext: (&net.Dialer{
  242. Timeout: proxyDialTimeout,
  243. }).DialContext,
  244. TLSHandshakeTimeout: proxyTLSHandshakeTimeout,
  245. }
  246. if err := proxyutil.ConfigureTransportProxy(transport, parsed); err != nil {
  247. return nil, fmt.Errorf("configure proxy: %w", err)
  248. }
  249. client.Transport = transport
  250. }
  251. return &Client{
  252. httpClient: client,
  253. }, nil
  254. }
  255. // IsConnectionError 判断是否为连接错误(网络超时、DNS 失败、连接拒绝)
  256. func IsConnectionError(err error) bool {
  257. if err == nil {
  258. return false
  259. }
  260. // 检查超时错误
  261. var netErr net.Error
  262. if errors.As(err, &netErr) && netErr.Timeout() {
  263. return true
  264. }
  265. // 检查连接错误(DNS 失败、连接拒绝)
  266. var opErr *net.OpError
  267. if errors.As(err, &opErr) {
  268. return true
  269. }
  270. // 检查 URL 错误
  271. var urlErr *url.Error
  272. return errors.As(err, &urlErr)
  273. }
  274. // shouldFallbackToNextURL 判断是否应切换到下一个 URL
  275. // 与 Antigravity-Manager 保持一致:连接错误、429、408、404、5xx 触发 URL 降级
  276. func shouldFallbackToNextURL(err error, statusCode int) bool {
  277. if IsConnectionError(err) {
  278. return true
  279. }
  280. return statusCode == http.StatusTooManyRequests ||
  281. statusCode == http.StatusRequestTimeout ||
  282. statusCode == http.StatusNotFound ||
  283. statusCode >= 500
  284. }
  285. // ExchangeCode 用 authorization code 交换 token
  286. func (c *Client) ExchangeCode(ctx context.Context, code, codeVerifier string) (*TokenResponse, error) {
  287. clientSecret, err := getClientSecret()
  288. if err != nil {
  289. return nil, err
  290. }
  291. params := url.Values{}
  292. params.Set("client_id", ClientID)
  293. params.Set("client_secret", clientSecret)
  294. params.Set("code", code)
  295. params.Set("redirect_uri", RedirectURI)
  296. params.Set("grant_type", "authorization_code")
  297. params.Set("code_verifier", codeVerifier)
  298. req, err := http.NewRequestWithContext(ctx, http.MethodPost, TokenURL, strings.NewReader(params.Encode()))
  299. if err != nil {
  300. return nil, fmt.Errorf("创建请求失败: %w", err)
  301. }
  302. req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
  303. resp, err := servertiming.Do(c.httpClient, req)
  304. if err != nil {
  305. return nil, fmt.Errorf("token 交换请求失败: %w", err)
  306. }
  307. defer func() { _ = resp.Body.Close() }()
  308. bodyBytes, err := io.ReadAll(resp.Body)
  309. if err != nil {
  310. return nil, fmt.Errorf("读取响应失败: %w", err)
  311. }
  312. if resp.StatusCode != http.StatusOK {
  313. return nil, fmt.Errorf("token 交换失败 (HTTP %d): %s", resp.StatusCode, string(bodyBytes))
  314. }
  315. var tokenResp TokenResponse
  316. if err := json.Unmarshal(bodyBytes, &tokenResp); err != nil {
  317. return nil, fmt.Errorf("token 解析失败: %w", err)
  318. }
  319. return &tokenResp, nil
  320. }
  321. // RefreshToken 刷新 access_token
  322. func (c *Client) RefreshToken(ctx context.Context, refreshToken string) (*TokenResponse, error) {
  323. clientSecret, err := getClientSecret()
  324. if err != nil {
  325. return nil, err
  326. }
  327. params := url.Values{}
  328. params.Set("client_id", ClientID)
  329. params.Set("client_secret", clientSecret)
  330. params.Set("refresh_token", refreshToken)
  331. params.Set("grant_type", "refresh_token")
  332. req, err := http.NewRequestWithContext(ctx, http.MethodPost, TokenURL, strings.NewReader(params.Encode()))
  333. if err != nil {
  334. return nil, fmt.Errorf("创建请求失败: %w", err)
  335. }
  336. req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
  337. resp, err := servertiming.Do(c.httpClient, req)
  338. if err != nil {
  339. return nil, fmt.Errorf("token 刷新请求失败: %w", err)
  340. }
  341. defer func() { _ = resp.Body.Close() }()
  342. bodyBytes, err := io.ReadAll(resp.Body)
  343. if err != nil {
  344. return nil, fmt.Errorf("读取响应失败: %w", err)
  345. }
  346. if resp.StatusCode != http.StatusOK {
  347. return nil, fmt.Errorf("token 刷新失败 (HTTP %d): %s", resp.StatusCode, string(bodyBytes))
  348. }
  349. var tokenResp TokenResponse
  350. if err := json.Unmarshal(bodyBytes, &tokenResp); err != nil {
  351. return nil, fmt.Errorf("token 解析失败: %w", err)
  352. }
  353. return &tokenResp, nil
  354. }
  355. // GetUserInfo 获取用户信息
  356. func (c *Client) GetUserInfo(ctx context.Context, accessToken string) (*UserInfo, error) {
  357. req, err := http.NewRequestWithContext(ctx, http.MethodGet, UserInfoURL, nil)
  358. if err != nil {
  359. return nil, fmt.Errorf("创建请求失败: %w", err)
  360. }
  361. req.Header.Set("Authorization", "Bearer "+accessToken)
  362. resp, err := servertiming.Do(c.httpClient, req)
  363. if err != nil {
  364. return nil, fmt.Errorf("用户信息请求失败: %w", err)
  365. }
  366. defer func() { _ = resp.Body.Close() }()
  367. bodyBytes, err := io.ReadAll(resp.Body)
  368. if err != nil {
  369. return nil, fmt.Errorf("读取响应失败: %w", err)
  370. }
  371. if resp.StatusCode != http.StatusOK {
  372. return nil, fmt.Errorf("获取用户信息失败 (HTTP %d): %s", resp.StatusCode, string(bodyBytes))
  373. }
  374. var userInfo UserInfo
  375. if err := json.Unmarshal(bodyBytes, &userInfo); err != nil {
  376. return nil, fmt.Errorf("用户信息解析失败: %w", err)
  377. }
  378. return &userInfo, nil
  379. }
  380. // LoadCodeAssist 获取账户信息,返回解析后的结构体和原始 JSON
  381. // 支持 URL fallback:sandbox → daily → prod
  382. func (c *Client) LoadCodeAssist(ctx context.Context, accessToken string) (*LoadCodeAssistResponse, map[string]any, error) {
  383. reqBody := LoadCodeAssistRequest{}
  384. reqBody.Metadata.IDEType = "ANTIGRAVITY"
  385. reqBody.Metadata.IDEVersion = GetUserAgentVersionForContext(ctx)
  386. reqBody.Metadata.IDEName = "antigravity"
  387. bodyBytes, err := json.Marshal(reqBody)
  388. if err != nil {
  389. return nil, nil, fmt.Errorf("序列化请求失败: %w", err)
  390. }
  391. // 固定顺序:prod -> daily
  392. availableURLs := BaseURLs
  393. var lastErr error
  394. for urlIdx, baseURL := range availableURLs {
  395. apiURL := baseURL + "/v1internal:loadCodeAssist"
  396. req, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, strings.NewReader(string(bodyBytes)))
  397. if err != nil {
  398. lastErr = fmt.Errorf("创建请求失败: %w", err)
  399. continue
  400. }
  401. req.Header.Set("Authorization", "Bearer "+accessToken)
  402. req.Header.Set("Content-Type", "application/json")
  403. req.Header.Set("User-Agent", GetUserAgentForContext(ctx))
  404. resp, err := servertiming.Do(c.httpClient, req)
  405. if err != nil {
  406. lastErr = fmt.Errorf("loadCodeAssist 请求失败: %w", err)
  407. if shouldFallbackToNextURL(err, 0) && urlIdx < len(availableURLs)-1 {
  408. log.Printf("[antigravity] loadCodeAssist URL fallback: %s -> %s", baseURL, availableURLs[urlIdx+1])
  409. continue
  410. }
  411. return nil, nil, lastErr
  412. }
  413. respBodyBytes, err := io.ReadAll(resp.Body)
  414. _ = resp.Body.Close() // 立即关闭,避免循环内 defer 导致的资源泄漏
  415. if err != nil {
  416. return nil, nil, fmt.Errorf("读取响应失败: %w", err)
  417. }
  418. // 检查是否需要 URL 降级
  419. if shouldFallbackToNextURL(nil, resp.StatusCode) && urlIdx < len(availableURLs)-1 {
  420. log.Printf("[antigravity] loadCodeAssist URL fallback (HTTP %d): %s -> %s", resp.StatusCode, baseURL, availableURLs[urlIdx+1])
  421. continue
  422. }
  423. if resp.StatusCode != http.StatusOK {
  424. return nil, nil, fmt.Errorf("loadCodeAssist 失败 (HTTP %d): %s", resp.StatusCode, string(respBodyBytes))
  425. }
  426. var loadResp LoadCodeAssistResponse
  427. if err := json.Unmarshal(respBodyBytes, &loadResp); err != nil {
  428. return nil, nil, fmt.Errorf("响应解析失败: %w", err)
  429. }
  430. // 解析原始 JSON 为 map
  431. var rawResp map[string]any
  432. _ = json.Unmarshal(respBodyBytes, &rawResp)
  433. // 标记成功的 URL,下次优先使用
  434. DefaultURLAvailability.MarkSuccess(baseURL)
  435. return &loadResp, rawResp, nil
  436. }
  437. return nil, nil, lastErr
  438. }
  439. // OnboardUser 触发账号 onboarding,并返回 project_id
  440. // 说明:
  441. // 1) 部分账号 loadCodeAssist 不会立即返回 cloudaicompanionProject;
  442. // 2) 这时需要调用 onboardUser 完成初始化,之后才能拿到 project_id。
  443. func (c *Client) OnboardUser(ctx context.Context, accessToken, tierID string) (string, error) {
  444. tierID = strings.TrimSpace(tierID)
  445. if tierID == "" {
  446. return "", fmt.Errorf("tier_id 为空")
  447. }
  448. reqBody := OnboardUserRequest{TierID: tierID}
  449. reqBody.Metadata.IDEType = "ANTIGRAVITY"
  450. reqBody.Metadata.Platform = "PLATFORM_UNSPECIFIED"
  451. reqBody.Metadata.PluginType = "GEMINI"
  452. bodyBytes, err := json.Marshal(reqBody)
  453. if err != nil {
  454. return "", fmt.Errorf("序列化请求失败: %w", err)
  455. }
  456. availableURLs := BaseURLs
  457. var lastErr error
  458. for urlIdx, baseURL := range availableURLs {
  459. apiURL := baseURL + "/v1internal:onboardUser"
  460. for attempt := 1; attempt <= 5; attempt++ {
  461. req, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, bytes.NewReader(bodyBytes))
  462. if err != nil {
  463. lastErr = fmt.Errorf("创建请求失败: %w", err)
  464. break
  465. }
  466. req.Header.Set("Authorization", "Bearer "+accessToken)
  467. req.Header.Set("Content-Type", "application/json")
  468. req.Header.Set("User-Agent", GetUserAgentForContext(ctx))
  469. resp, err := servertiming.Do(c.httpClient, req)
  470. if err != nil {
  471. lastErr = fmt.Errorf("onboardUser 请求失败: %w", err)
  472. if shouldFallbackToNextURL(err, 0) && urlIdx < len(availableURLs)-1 {
  473. log.Printf("[antigravity] onboardUser URL fallback: %s -> %s", baseURL, availableURLs[urlIdx+1])
  474. break
  475. }
  476. return "", lastErr
  477. }
  478. respBodyBytes, err := io.ReadAll(resp.Body)
  479. _ = resp.Body.Close()
  480. if err != nil {
  481. return "", fmt.Errorf("读取响应失败: %w", err)
  482. }
  483. if shouldFallbackToNextURL(nil, resp.StatusCode) && urlIdx < len(availableURLs)-1 {
  484. log.Printf("[antigravity] onboardUser URL fallback (HTTP %d): %s -> %s", resp.StatusCode, baseURL, availableURLs[urlIdx+1])
  485. break
  486. }
  487. if resp.StatusCode != http.StatusOK {
  488. lastErr = fmt.Errorf("onboardUser 失败 (HTTP %d): %s", resp.StatusCode, string(respBodyBytes))
  489. return "", lastErr
  490. }
  491. var onboardResp OnboardUserResponse
  492. if err := json.Unmarshal(respBodyBytes, &onboardResp); err != nil {
  493. lastErr = fmt.Errorf("onboardUser 响应解析失败: %w", err)
  494. return "", lastErr
  495. }
  496. if onboardResp.Done {
  497. if projectID := extractProjectIDFromOnboardResponse(onboardResp.Response); projectID != "" {
  498. DefaultURLAvailability.MarkSuccess(baseURL)
  499. return projectID, nil
  500. }
  501. lastErr = fmt.Errorf("onboardUser 完成但未返回 project_id")
  502. return "", lastErr
  503. }
  504. // done=false 时等待后重试,避免在上游仍未完成时提前结束轮询。
  505. select {
  506. case <-time.After(2 * time.Second):
  507. case <-ctx.Done():
  508. return "", ctx.Err()
  509. }
  510. }
  511. }
  512. if lastErr != nil {
  513. return "", lastErr
  514. }
  515. return "", fmt.Errorf("onboardUser 未返回 project_id")
  516. }
  517. func extractProjectIDFromOnboardResponse(resp map[string]any) string {
  518. if len(resp) == 0 {
  519. return ""
  520. }
  521. if v, ok := resp["cloudaicompanionProject"]; ok {
  522. switch project := v.(type) {
  523. case string:
  524. return strings.TrimSpace(project)
  525. case map[string]any:
  526. if id, ok := project["id"].(string); ok {
  527. return strings.TrimSpace(id)
  528. }
  529. }
  530. }
  531. return ""
  532. }
  533. // ModelQuotaInfo 模型配额信息
  534. type ModelQuotaInfo struct {
  535. RemainingFraction float64 `json:"remainingFraction"`
  536. ResetTime string `json:"resetTime,omitempty"`
  537. }
  538. // ModelInfo 模型信息
  539. type ModelInfo struct {
  540. QuotaInfo *ModelQuotaInfo `json:"quotaInfo,omitempty"`
  541. DisplayName string `json:"displayName,omitempty"`
  542. SupportsImages *bool `json:"supportsImages,omitempty"`
  543. SupportsThinking *bool `json:"supportsThinking,omitempty"`
  544. ThinkingBudget *int `json:"thinkingBudget,omitempty"`
  545. Recommended *bool `json:"recommended,omitempty"`
  546. MaxTokens *int `json:"maxTokens,omitempty"`
  547. MaxOutputTokens *int `json:"maxOutputTokens,omitempty"`
  548. SupportedMimeTypes map[string]bool `json:"supportedMimeTypes,omitempty"`
  549. }
  550. // DeprecatedModelInfo 废弃模型转发信息
  551. type DeprecatedModelInfo struct {
  552. NewModelID string `json:"newModelId"`
  553. }
  554. // FetchAvailableModelsRequest fetchAvailableModels 请求
  555. type FetchAvailableModelsRequest struct {
  556. Project string `json:"project"`
  557. }
  558. // FetchAvailableModelsResponse fetchAvailableModels 响应
  559. type FetchAvailableModelsResponse struct {
  560. Models map[string]ModelInfo `json:"models"`
  561. DeprecatedModelIDs map[string]DeprecatedModelInfo `json:"deprecatedModelIds,omitempty"`
  562. }
  563. // FetchAvailableModels 获取可用模型和配额信息,返回解析后的结构体和原始 JSON
  564. // 支持 URL fallback:sandbox → daily → prod
  565. func (c *Client) FetchAvailableModels(ctx context.Context, accessToken, projectID string, bodyLimit int64) (*FetchAvailableModelsResponse, map[string]any, error) {
  566. if c == nil || c.httpClient == nil {
  567. return nil, nil, errors.New("antigravity client is not configured")
  568. }
  569. if bodyLimit <= 0 {
  570. return nil, nil, errors.New("fetchAvailableModels body limit must be positive")
  571. }
  572. reqBody := FetchAvailableModelsRequest{Project: projectID}
  573. bodyBytes, err := json.Marshal(reqBody)
  574. if err != nil {
  575. return nil, nil, fmt.Errorf("序列化请求失败: %w", err)
  576. }
  577. // 固定顺序:prod -> daily
  578. availableURLs := BaseURLs
  579. fetchClient := c.fetchAvailableModelsHTTPClient()
  580. var lastErr error
  581. for urlIdx, baseURL := range availableURLs {
  582. apiURL := baseURL + "/v1internal:fetchAvailableModels"
  583. req, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, strings.NewReader(string(bodyBytes)))
  584. if err != nil {
  585. lastErr = fmt.Errorf("创建请求失败: %w", err)
  586. continue
  587. }
  588. req.Header.Set("Authorization", "Bearer "+accessToken)
  589. req.Header.Set("Content-Type", "application/json")
  590. req.Header.Set("User-Agent", GetUserAgentForContext(ctx))
  591. resp, err := servertiming.Do(fetchClient, req)
  592. if err != nil {
  593. lastErr = fmt.Errorf("fetchAvailableModels 请求失败: %w", err)
  594. if shouldFallbackToNextURL(err, 0) && urlIdx < len(availableURLs)-1 {
  595. log.Printf("[antigravity] fetchAvailableModels URL fallback: %s -> %s", baseURL, availableURLs[urlIdx+1])
  596. continue
  597. }
  598. return nil, nil, lastErr
  599. }
  600. respBodyBytes, err := io.ReadAll(io.LimitReader(resp.Body, bodyLimit+1))
  601. _ = resp.Body.Close() // 立即关闭,避免循环内 defer 导致的资源泄漏
  602. if err != nil {
  603. return nil, nil, fmt.Errorf("读取响应失败: %w", err)
  604. }
  605. if int64(len(respBodyBytes)) > bodyLimit {
  606. return nil, nil, fmt.Errorf("响应超过 %d 字节", bodyLimit)
  607. }
  608. // 检查是否需要 URL 降级
  609. if shouldFallbackToNextURL(nil, resp.StatusCode) && urlIdx < len(availableURLs)-1 {
  610. log.Printf("[antigravity] fetchAvailableModels URL fallback (HTTP %d): %s -> %s", resp.StatusCode, baseURL, availableURLs[urlIdx+1])
  611. continue
  612. }
  613. if resp.StatusCode == http.StatusForbidden {
  614. return nil, nil, &ForbiddenError{
  615. StatusCode: resp.StatusCode,
  616. Body: string(respBodyBytes),
  617. }
  618. }
  619. if resp.StatusCode != http.StatusOK {
  620. return nil, nil, fmt.Errorf("fetchAvailableModels 失败 (HTTP %d): %s", resp.StatusCode, string(respBodyBytes))
  621. }
  622. var modelsResp FetchAvailableModelsResponse
  623. if err := json.Unmarshal(respBodyBytes, &modelsResp); err != nil {
  624. return nil, nil, fmt.Errorf("响应解析失败: %w", err)
  625. }
  626. // 解析原始 JSON 为 map
  627. var rawResp map[string]any
  628. _ = json.Unmarshal(respBodyBytes, &rawResp)
  629. // 标记成功的 URL,下次优先使用
  630. DefaultURLAvailability.MarkSuccess(baseURL)
  631. return &modelsResp, rawResp, nil
  632. }
  633. return nil, nil, lastErr
  634. }
  635. func (c *Client) fetchAvailableModelsHTTPClient() *http.Client {
  636. fetchClient := *c.httpClient
  637. fetchClient.CheckRedirect = checkFetchAvailableModelsRedirect
  638. return &fetchClient
  639. }
  640. func checkFetchAvailableModelsRedirect(req *http.Request, via []*http.Request) error {
  641. if len(via) >= 10 {
  642. return errors.New("stopped after 10 redirects")
  643. }
  644. if req == nil || req.URL == nil {
  645. return errors.New("redirect url is nil")
  646. }
  647. if !isAllowedFetchAvailableModelsRedirectHost(req.URL.Hostname()) {
  648. return fmt.Errorf("redirect to unsupported host: %s", req.URL.Hostname())
  649. }
  650. return nil
  651. }
  652. func isAllowedFetchAvailableModelsRedirectHost(host string) bool {
  653. host = strings.ToLower(strings.TrimSpace(host))
  654. if host == "" {
  655. return false
  656. }
  657. for _, baseURL := range BaseURLs {
  658. parsed, err := url.Parse(baseURL)
  659. if err != nil {
  660. continue
  661. }
  662. if strings.EqualFold(host, parsed.Hostname()) {
  663. return true
  664. }
  665. }
  666. return false
  667. }
  668. // ── Privacy API ──────────────────────────────────────────────────────
  669. // privacyBaseURL 隐私设置 API 仅使用 daily 端点(与 Antigravity 客户端行为一致)
  670. const privacyBaseURL = antigravityDailyBaseURL
  671. // SetUserSettingsRequest setUserSettings 请求体
  672. type SetUserSettingsRequest struct {
  673. UserSettings map[string]any `json:"user_settings"`
  674. }
  675. // FetchUserInfoRequest fetchUserInfo 请求体
  676. type FetchUserInfoRequest struct {
  677. Project string `json:"project"`
  678. }
  679. // FetchUserInfoResponse fetchUserInfo 响应体
  680. type FetchUserInfoResponse struct {
  681. UserSettings map[string]any `json:"userSettings,omitempty"`
  682. RegionCode string `json:"regionCode,omitempty"`
  683. }
  684. // IsPrivate 判断隐私是否已设置:userSettings 为空或不含 telemetryEnabled 表示已设置
  685. func (r *FetchUserInfoResponse) IsPrivate() bool {
  686. if r == nil || r.UserSettings == nil {
  687. return true
  688. }
  689. _, hasTelemetry := r.UserSettings["telemetryEnabled"]
  690. return !hasTelemetry
  691. }
  692. // SetUserSettingsResponse setUserSettings 响应体
  693. type SetUserSettingsResponse struct {
  694. UserSettings map[string]any `json:"userSettings,omitempty"`
  695. }
  696. // IsSuccess 判断 setUserSettings 是否成功:返回 {"userSettings":{}} 且无 telemetryEnabled
  697. func (r *SetUserSettingsResponse) IsSuccess() bool {
  698. if r == nil {
  699. return false
  700. }
  701. // userSettings 为 nil 或空 map 均视为成功
  702. if len(r.UserSettings) == 0 {
  703. return true
  704. }
  705. // 如果包含 telemetryEnabled 字段,说明未成功清除
  706. _, hasTelemetry := r.UserSettings["telemetryEnabled"]
  707. return !hasTelemetry
  708. }
  709. // SetUserSettings 调用 setUserSettings API 设置用户隐私,返回解析后的响应
  710. func (c *Client) SetUserSettings(ctx context.Context, accessToken string) (*SetUserSettingsResponse, error) {
  711. // 发送空 user_settings 以清除隐私设置
  712. payload := SetUserSettingsRequest{UserSettings: map[string]any{}}
  713. bodyBytes, err := json.Marshal(payload)
  714. if err != nil {
  715. return nil, fmt.Errorf("序列化请求失败: %w", err)
  716. }
  717. apiURL := privacyBaseURL + "/v1internal:setUserSettings"
  718. req, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, bytes.NewReader(bodyBytes))
  719. if err != nil {
  720. return nil, fmt.Errorf("创建请求失败: %w", err)
  721. }
  722. req.Header.Set("Authorization", "Bearer "+accessToken)
  723. req.Header.Set("Content-Type", "application/json")
  724. req.Header.Set("Accept", "*/*")
  725. req.Header.Set("User-Agent", GetUserAgentForContext(ctx))
  726. req.Header.Set("X-Goog-Api-Client", "gl-node/22.21.1")
  727. req.Host = "daily-cloudcode-pa.googleapis.com"
  728. resp, err := servertiming.Do(c.httpClient, req)
  729. if err != nil {
  730. return nil, fmt.Errorf("setUserSettings 请求失败: %w", err)
  731. }
  732. defer func() { _ = resp.Body.Close() }()
  733. respBody, err := io.ReadAll(resp.Body)
  734. if err != nil {
  735. return nil, fmt.Errorf("读取响应失败: %w", err)
  736. }
  737. if resp.StatusCode != http.StatusOK {
  738. return nil, fmt.Errorf("setUserSettings 失败 (HTTP %d): %s", resp.StatusCode, string(respBody))
  739. }
  740. var result SetUserSettingsResponse
  741. if err := json.Unmarshal(respBody, &result); err != nil {
  742. return nil, fmt.Errorf("响应解析失败: %w", err)
  743. }
  744. return &result, nil
  745. }
  746. // FetchUserInfo 调用 fetchUserInfo API 获取用户隐私设置状态
  747. func (c *Client) FetchUserInfo(ctx context.Context, accessToken, projectID string) (*FetchUserInfoResponse, error) {
  748. reqBody := FetchUserInfoRequest{Project: projectID}
  749. bodyBytes, err := json.Marshal(reqBody)
  750. if err != nil {
  751. return nil, fmt.Errorf("序列化请求失败: %w", err)
  752. }
  753. apiURL := privacyBaseURL + "/v1internal:fetchUserInfo"
  754. req, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, bytes.NewReader(bodyBytes))
  755. if err != nil {
  756. return nil, fmt.Errorf("创建请求失败: %w", err)
  757. }
  758. req.Header.Set("Authorization", "Bearer "+accessToken)
  759. req.Header.Set("Content-Type", "application/json")
  760. req.Header.Set("Accept", "*/*")
  761. req.Header.Set("User-Agent", GetUserAgentForContext(ctx))
  762. req.Header.Set("X-Goog-Api-Client", "gl-node/22.21.1")
  763. req.Host = "daily-cloudcode-pa.googleapis.com"
  764. resp, err := servertiming.Do(c.httpClient, req)
  765. if err != nil {
  766. return nil, fmt.Errorf("fetchUserInfo 请求失败: %w", err)
  767. }
  768. defer func() { _ = resp.Body.Close() }()
  769. respBody, err := io.ReadAll(resp.Body)
  770. if err != nil {
  771. return nil, fmt.Errorf("读取响应失败: %w", err)
  772. }
  773. if resp.StatusCode != http.StatusOK {
  774. return nil, fmt.Errorf("fetchUserInfo 失败 (HTTP %d): %s", resp.StatusCode, string(respBody))
  775. }
  776. var result FetchUserInfoResponse
  777. if err := json.Unmarshal(respBody, &result); err != nil {
  778. return nil, fmt.Errorf("响应解析失败: %w", err)
  779. }
  780. return &result, nil
  781. }