fix: harden Antigravity model list requests
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
This commit is contained in:
parent
ba676e43fd
commit
b9ecf25207
@ -254,6 +254,8 @@ const (
|
|||||||
proxyTLSHandshakeTimeout = 5 * time.Second
|
proxyTLSHandshakeTimeout = 5 * time.Second
|
||||||
// clientTimeout 整体请求超时(含连接、发送、等待响应、读取 body)
|
// clientTimeout 整体请求超时(含连接、发送、等待响应、读取 body)
|
||||||
clientTimeout = 10 * time.Second
|
clientTimeout = 10 * time.Second
|
||||||
|
// fetchAvailableModelsBodyLimit limits model-list responses to avoid unbounded memory use.
|
||||||
|
fetchAvailableModelsBodyLimit int64 = 8 << 20
|
||||||
)
|
)
|
||||||
|
|
||||||
func NewClient(proxyURL string) (*Client, error) {
|
func NewClient(proxyURL string) (*Client, error) {
|
||||||
@ -655,6 +657,10 @@ type FetchAvailableModelsResponse struct {
|
|||||||
// FetchAvailableModels 获取可用模型和配额信息,返回解析后的结构体和原始 JSON
|
// FetchAvailableModels 获取可用模型和配额信息,返回解析后的结构体和原始 JSON
|
||||||
// 支持 URL fallback:sandbox → daily → prod
|
// 支持 URL fallback:sandbox → daily → prod
|
||||||
func (c *Client) FetchAvailableModels(ctx context.Context, accessToken, projectID string) (*FetchAvailableModelsResponse, map[string]any, error) {
|
func (c *Client) FetchAvailableModels(ctx context.Context, accessToken, projectID string) (*FetchAvailableModelsResponse, map[string]any, error) {
|
||||||
|
if c == nil || c.httpClient == nil {
|
||||||
|
return nil, nil, errors.New("antigravity client is not configured")
|
||||||
|
}
|
||||||
|
|
||||||
reqBody := FetchAvailableModelsRequest{Project: projectID}
|
reqBody := FetchAvailableModelsRequest{Project: projectID}
|
||||||
bodyBytes, err := json.Marshal(reqBody)
|
bodyBytes, err := json.Marshal(reqBody)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@ -664,6 +670,7 @@ func (c *Client) FetchAvailableModels(ctx context.Context, accessToken, projectI
|
|||||||
// 固定顺序:prod -> daily
|
// 固定顺序:prod -> daily
|
||||||
availableURLs := BaseURLs
|
availableURLs := BaseURLs
|
||||||
|
|
||||||
|
fetchClient := c.fetchAvailableModelsHTTPClient()
|
||||||
var lastErr error
|
var lastErr error
|
||||||
for urlIdx, baseURL := range availableURLs {
|
for urlIdx, baseURL := range availableURLs {
|
||||||
apiURL := baseURL + "/v1internal:fetchAvailableModels"
|
apiURL := baseURL + "/v1internal:fetchAvailableModels"
|
||||||
@ -676,7 +683,7 @@ func (c *Client) FetchAvailableModels(ctx context.Context, accessToken, projectI
|
|||||||
req.Header.Set("Content-Type", "application/json")
|
req.Header.Set("Content-Type", "application/json")
|
||||||
req.Header.Set("User-Agent", GetUserAgentForContext(ctx))
|
req.Header.Set("User-Agent", GetUserAgentForContext(ctx))
|
||||||
|
|
||||||
resp, err := c.httpClient.Do(req)
|
resp, err := fetchClient.Do(req)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
lastErr = fmt.Errorf("fetchAvailableModels 请求失败: %w", err)
|
lastErr = fmt.Errorf("fetchAvailableModels 请求失败: %w", err)
|
||||||
if shouldFallbackToNextURL(err, 0) && urlIdx < len(availableURLs)-1 {
|
if shouldFallbackToNextURL(err, 0) && urlIdx < len(availableURLs)-1 {
|
||||||
@ -686,11 +693,14 @@ func (c *Client) FetchAvailableModels(ctx context.Context, accessToken, projectI
|
|||||||
return nil, nil, lastErr
|
return nil, nil, lastErr
|
||||||
}
|
}
|
||||||
|
|
||||||
respBodyBytes, err := io.ReadAll(resp.Body)
|
respBodyBytes, err := io.ReadAll(io.LimitReader(resp.Body, fetchAvailableModelsBodyLimit+1))
|
||||||
_ = resp.Body.Close() // 立即关闭,避免循环内 defer 导致的资源泄漏
|
_ = resp.Body.Close() // 立即关闭,避免循环内 defer 导致的资源泄漏
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, nil, fmt.Errorf("读取响应失败: %w", err)
|
return nil, nil, fmt.Errorf("读取响应失败: %w", err)
|
||||||
}
|
}
|
||||||
|
if int64(len(respBodyBytes)) > fetchAvailableModelsBodyLimit {
|
||||||
|
return nil, nil, fmt.Errorf("响应超过 %d 字节", fetchAvailableModelsBodyLimit)
|
||||||
|
}
|
||||||
|
|
||||||
// 检查是否需要 URL 降级
|
// 检查是否需要 URL 降级
|
||||||
if shouldFallbackToNextURL(nil, resp.StatusCode) && urlIdx < len(availableURLs)-1 {
|
if shouldFallbackToNextURL(nil, resp.StatusCode) && urlIdx < len(availableURLs)-1 {
|
||||||
@ -726,6 +736,42 @@ func (c *Client) FetchAvailableModels(ctx context.Context, accessToken, projectI
|
|||||||
return nil, nil, lastErr
|
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 ──────────────────────────────────────────────────────
|
// ── Privacy API ──────────────────────────────────────────────────────
|
||||||
|
|
||||||
// privacyBaseURL 隐私设置 API 仅使用 daily 端点(与 Antigravity 客户端行为一致)
|
// privacyBaseURL 隐私设置 API 仅使用 daily 端点(与 Antigravity 客户端行为一致)
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user