mirror of
https://github.com/sipeed/picoclaw.git
synced 2026-06-12 18:08:54 +00:00
943385105f
io.ReadAll errors were silently discarded with `body, _ := io.ReadAll(...)`, which could cause empty or partial data to be used for JSON unmarshaling or error messages. This adds proper error checks for all instances.
804 lines
22 KiB
Go
804 lines
22 KiB
Go
package providers
|
|
|
|
import (
|
|
"bufio"
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"math/rand"
|
|
"net/http"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/sipeed/picoclaw/pkg/auth"
|
|
"github.com/sipeed/picoclaw/pkg/logger"
|
|
)
|
|
|
|
const (
|
|
antigravityBaseURL = "https://cloudcode-pa.googleapis.com"
|
|
antigravityDefaultModel = "gemini-3-flash"
|
|
antigravityUserAgent = "antigravity"
|
|
antigravityXGoogClient = "google-cloud-sdk vscode_cloudshelleditor/0.1"
|
|
antigravityVersion = "1.15.8"
|
|
)
|
|
|
|
// AntigravityProvider implements LLMProvider using Google's Cloud Code Assist (Antigravity) API.
|
|
// This provider authenticates via Google OAuth and provides access to models like Claude and Gemini
|
|
// through Google's infrastructure.
|
|
type AntigravityProvider struct {
|
|
tokenSource func() (string, string, error) // Returns (accessToken, projectID, error)
|
|
httpClient *http.Client
|
|
}
|
|
|
|
// NewAntigravityProvider creates a new Antigravity provider using stored auth credentials.
|
|
func NewAntigravityProvider() *AntigravityProvider {
|
|
return &AntigravityProvider{
|
|
tokenSource: createAntigravityTokenSource(),
|
|
httpClient: &http.Client{
|
|
Timeout: 120 * time.Second,
|
|
},
|
|
}
|
|
}
|
|
|
|
// Chat implements LLMProvider.Chat using the Cloud Code Assist v1internal API.
|
|
// The v1internal endpoint wraps the standard Gemini request in an envelope with
|
|
// project, model, request, requestType, userAgent, and requestId fields.
|
|
func (p *AntigravityProvider) Chat(
|
|
ctx context.Context,
|
|
messages []Message,
|
|
tools []ToolDefinition,
|
|
model string,
|
|
options map[string]any,
|
|
) (*LLMResponse, error) {
|
|
accessToken, projectID, err := p.tokenSource()
|
|
if err != nil {
|
|
return nil, fmt.Errorf("antigravity auth: %w", err)
|
|
}
|
|
|
|
if model == "" || model == "antigravity" || model == "google-antigravity" {
|
|
model = antigravityDefaultModel
|
|
}
|
|
// Strip provider prefixes if present
|
|
model = strings.TrimPrefix(model, "google-antigravity/")
|
|
model = strings.TrimPrefix(model, "antigravity/")
|
|
|
|
logger.DebugCF("provider.antigravity", "Starting chat", map[string]any{
|
|
"model": model,
|
|
"project": projectID,
|
|
"requestId": fmt.Sprintf("agent-%d-%s", time.Now().UnixMilli(), randomString(9)),
|
|
})
|
|
|
|
// Build the inner Gemini-format request
|
|
innerRequest := p.buildRequest(messages, tools, model, options)
|
|
|
|
// Wrap in v1internal envelope (matches pi-ai SDK format)
|
|
envelope := map[string]any{
|
|
"project": projectID,
|
|
"model": model,
|
|
"request": innerRequest,
|
|
"requestType": "agent",
|
|
"userAgent": antigravityUserAgent,
|
|
"requestId": fmt.Sprintf("agent-%d-%s", time.Now().UnixMilli(), randomString(9)),
|
|
}
|
|
|
|
bodyBytes, err := json.Marshal(envelope)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("marshaling request: %w", err)
|
|
}
|
|
|
|
// Build API URL — uses Cloud Code Assist v1internal streaming endpoint
|
|
apiURL := fmt.Sprintf("%s/v1internal:streamGenerateContent?alt=sse", antigravityBaseURL)
|
|
|
|
req, err := http.NewRequestWithContext(ctx, "POST", apiURL, bytes.NewReader(bodyBytes))
|
|
if err != nil {
|
|
return nil, fmt.Errorf("creating request: %w", err)
|
|
}
|
|
|
|
// Headers matching the pi-ai SDK antigravity format
|
|
clientMetadata, _ := json.Marshal(map[string]string{
|
|
"ideType": "IDE_UNSPECIFIED",
|
|
"platform": "PLATFORM_UNSPECIFIED",
|
|
"pluginType": "GEMINI",
|
|
})
|
|
req.Header.Set("Authorization", "Bearer "+accessToken)
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("Accept", "text/event-stream")
|
|
req.Header.Set("User-Agent", fmt.Sprintf("antigravity/%s linux/amd64", antigravityVersion))
|
|
req.Header.Set("X-Goog-Api-Client", antigravityXGoogClient)
|
|
req.Header.Set("Client-Metadata", string(clientMetadata))
|
|
|
|
resp, err := p.httpClient.Do(req)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("antigravity API call: %w", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
respBody, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("reading response: %w", err)
|
|
}
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
logger.ErrorCF("provider.antigravity", "API call failed", map[string]any{
|
|
"status_code": resp.StatusCode,
|
|
"response": string(respBody),
|
|
"model": model,
|
|
})
|
|
|
|
return nil, p.parseAntigravityError(resp.StatusCode, respBody)
|
|
}
|
|
|
|
// Response is always SSE from streamGenerateContent — each line is "data: {...}"
|
|
// with a "response" wrapper containing the standard Gemini response
|
|
llmResp, err := p.parseSSEResponse(string(respBody))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
// Check for empty response (some models might return valid success but empty text)
|
|
if llmResp.Content == "" && len(llmResp.ToolCalls) == 0 {
|
|
return nil, fmt.Errorf(
|
|
"antigravity: model returned an empty response (this model might be invalid or restricted)",
|
|
)
|
|
}
|
|
|
|
return llmResp, nil
|
|
}
|
|
|
|
// GetDefaultModel returns the default model identifier.
|
|
func (p *AntigravityProvider) GetDefaultModel() string {
|
|
return antigravityDefaultModel
|
|
}
|
|
|
|
// --- Request building ---
|
|
|
|
type antigravityRequest struct {
|
|
Contents []antigravityContent `json:"contents"`
|
|
Tools []antigravityTool `json:"tools,omitempty"`
|
|
SystemPrompt *antigravitySystemPrompt `json:"systemInstruction,omitempty"`
|
|
Config *antigravityGenConfig `json:"generationConfig,omitempty"`
|
|
}
|
|
|
|
type antigravityContent struct {
|
|
Role string `json:"role"`
|
|
Parts []antigravityPart `json:"parts"`
|
|
}
|
|
|
|
type antigravityPart struct {
|
|
Text string `json:"text,omitempty"`
|
|
ThoughtSignature string `json:"thoughtSignature,omitempty"`
|
|
ThoughtSignatureSnake string `json:"thought_signature,omitempty"`
|
|
FunctionCall *antigravityFunctionCall `json:"functionCall,omitempty"`
|
|
FunctionResponse *antigravityFunctionResponse `json:"functionResponse,omitempty"`
|
|
}
|
|
|
|
type antigravityFunctionCall struct {
|
|
Name string `json:"name"`
|
|
Args map[string]any `json:"args"`
|
|
}
|
|
|
|
type antigravityFunctionResponse struct {
|
|
Name string `json:"name"`
|
|
Response map[string]any `json:"response"`
|
|
}
|
|
|
|
type antigravityTool struct {
|
|
FunctionDeclarations []antigravityFuncDecl `json:"functionDeclarations"`
|
|
}
|
|
|
|
type antigravityFuncDecl struct {
|
|
Name string `json:"name"`
|
|
Description string `json:"description,omitempty"`
|
|
Parameters any `json:"parameters,omitempty"`
|
|
}
|
|
|
|
type antigravitySystemPrompt struct {
|
|
Parts []antigravityPart `json:"parts"`
|
|
}
|
|
|
|
type antigravityGenConfig struct {
|
|
MaxOutputTokens int `json:"maxOutputTokens,omitempty"`
|
|
Temperature float64 `json:"temperature,omitempty"`
|
|
}
|
|
|
|
func (p *AntigravityProvider) buildRequest(
|
|
messages []Message,
|
|
tools []ToolDefinition,
|
|
model string,
|
|
options map[string]any,
|
|
) antigravityRequest {
|
|
req := antigravityRequest{}
|
|
toolCallNames := make(map[string]string)
|
|
|
|
// Build contents from messages
|
|
for _, msg := range messages {
|
|
switch msg.Role {
|
|
case "system":
|
|
req.SystemPrompt = &antigravitySystemPrompt{
|
|
Parts: []antigravityPart{{Text: msg.Content}},
|
|
}
|
|
case "user":
|
|
if msg.ToolCallID != "" {
|
|
toolName := resolveToolResponseName(msg.ToolCallID, toolCallNames)
|
|
// Tool result
|
|
req.Contents = append(req.Contents, antigravityContent{
|
|
Role: "user",
|
|
Parts: []antigravityPart{{
|
|
FunctionResponse: &antigravityFunctionResponse{
|
|
Name: toolName,
|
|
Response: map[string]any{
|
|
"result": msg.Content,
|
|
},
|
|
},
|
|
}},
|
|
})
|
|
} else {
|
|
req.Contents = append(req.Contents, antigravityContent{
|
|
Role: "user",
|
|
Parts: []antigravityPart{{Text: msg.Content}},
|
|
})
|
|
}
|
|
case "assistant":
|
|
content := antigravityContent{
|
|
Role: "model",
|
|
}
|
|
if msg.Content != "" {
|
|
content.Parts = append(content.Parts, antigravityPart{Text: msg.Content})
|
|
}
|
|
for _, tc := range msg.ToolCalls {
|
|
toolName, toolArgs, thoughtSignature := normalizeStoredToolCall(tc)
|
|
if toolName == "" {
|
|
logger.WarnCF(
|
|
"provider.antigravity",
|
|
"Skipping tool call with empty name in history",
|
|
map[string]any{
|
|
"tool_call_id": tc.ID,
|
|
},
|
|
)
|
|
continue
|
|
}
|
|
if tc.ID != "" {
|
|
toolCallNames[tc.ID] = toolName
|
|
}
|
|
content.Parts = append(content.Parts, antigravityPart{
|
|
ThoughtSignature: thoughtSignature,
|
|
ThoughtSignatureSnake: thoughtSignature,
|
|
FunctionCall: &antigravityFunctionCall{
|
|
Name: toolName,
|
|
Args: toolArgs,
|
|
},
|
|
})
|
|
}
|
|
if len(content.Parts) > 0 {
|
|
req.Contents = append(req.Contents, content)
|
|
}
|
|
case "tool":
|
|
toolName := resolveToolResponseName(msg.ToolCallID, toolCallNames)
|
|
req.Contents = append(req.Contents, antigravityContent{
|
|
Role: "user",
|
|
Parts: []antigravityPart{{
|
|
FunctionResponse: &antigravityFunctionResponse{
|
|
Name: toolName,
|
|
Response: map[string]any{
|
|
"result": msg.Content,
|
|
},
|
|
},
|
|
}},
|
|
})
|
|
}
|
|
}
|
|
|
|
// Build tools (sanitize schemas for Gemini compatibility)
|
|
if len(tools) > 0 {
|
|
var funcDecls []antigravityFuncDecl
|
|
for _, t := range tools {
|
|
if t.Type != "function" {
|
|
continue
|
|
}
|
|
params := sanitizeSchemaForGemini(t.Function.Parameters)
|
|
funcDecls = append(funcDecls, antigravityFuncDecl{
|
|
Name: t.Function.Name,
|
|
Description: t.Function.Description,
|
|
Parameters: params,
|
|
})
|
|
}
|
|
if len(funcDecls) > 0 {
|
|
req.Tools = []antigravityTool{{FunctionDeclarations: funcDecls}}
|
|
}
|
|
}
|
|
|
|
// Generation config
|
|
config := &antigravityGenConfig{}
|
|
if val, ok := options["max_tokens"]; ok {
|
|
if maxTokens, ok := val.(int); ok && maxTokens > 0 {
|
|
config.MaxOutputTokens = maxTokens
|
|
} else if maxTokens, ok := val.(float64); ok && maxTokens > 0 {
|
|
config.MaxOutputTokens = int(maxTokens)
|
|
}
|
|
}
|
|
if temp, ok := options["temperature"].(float64); ok {
|
|
config.Temperature = temp
|
|
}
|
|
if config.MaxOutputTokens > 0 || config.Temperature > 0 {
|
|
req.Config = config
|
|
}
|
|
|
|
return req
|
|
}
|
|
|
|
func normalizeStoredToolCall(tc ToolCall) (string, map[string]any, string) {
|
|
name := tc.Name
|
|
args := tc.Arguments
|
|
thoughtSignature := ""
|
|
|
|
if name == "" && tc.Function != nil {
|
|
name = tc.Function.Name
|
|
thoughtSignature = tc.Function.ThoughtSignature
|
|
} else if tc.Function != nil {
|
|
thoughtSignature = tc.Function.ThoughtSignature
|
|
}
|
|
|
|
if args == nil {
|
|
args = map[string]any{}
|
|
}
|
|
|
|
if len(args) == 0 && tc.Function != nil && tc.Function.Arguments != "" {
|
|
var parsed map[string]any
|
|
if err := json.Unmarshal([]byte(tc.Function.Arguments), &parsed); err == nil && parsed != nil {
|
|
args = parsed
|
|
}
|
|
}
|
|
|
|
return name, args, thoughtSignature
|
|
}
|
|
|
|
func resolveToolResponseName(toolCallID string, toolCallNames map[string]string) string {
|
|
if toolCallID == "" {
|
|
return ""
|
|
}
|
|
|
|
if name, ok := toolCallNames[toolCallID]; ok && name != "" {
|
|
return name
|
|
}
|
|
|
|
return inferToolNameFromCallID(toolCallID)
|
|
}
|
|
|
|
func inferToolNameFromCallID(toolCallID string) string {
|
|
if !strings.HasPrefix(toolCallID, "call_") {
|
|
return toolCallID
|
|
}
|
|
|
|
rest := strings.TrimPrefix(toolCallID, "call_")
|
|
if idx := strings.LastIndex(rest, "_"); idx > 0 {
|
|
candidate := rest[:idx]
|
|
if candidate != "" {
|
|
return candidate
|
|
}
|
|
}
|
|
|
|
return toolCallID
|
|
}
|
|
|
|
// --- Response parsing ---
|
|
|
|
type antigravityJSONResponse struct {
|
|
Candidates []struct {
|
|
Content struct {
|
|
Parts []struct {
|
|
Text string `json:"text,omitempty"`
|
|
ThoughtSignature string `json:"thoughtSignature,omitempty"`
|
|
ThoughtSignatureSnake string `json:"thought_signature,omitempty"`
|
|
FunctionCall *antigravityFunctionCall `json:"functionCall,omitempty"`
|
|
} `json:"parts"`
|
|
Role string `json:"role"`
|
|
} `json:"content"`
|
|
FinishReason string `json:"finishReason"`
|
|
} `json:"candidates"`
|
|
UsageMetadata struct {
|
|
PromptTokenCount int `json:"promptTokenCount"`
|
|
CandidatesTokenCount int `json:"candidatesTokenCount"`
|
|
TotalTokenCount int `json:"totalTokenCount"`
|
|
} `json:"usageMetadata"`
|
|
}
|
|
|
|
func (p *AntigravityProvider) parseSSEResponse(body string) (*LLMResponse, error) {
|
|
var contentParts []string
|
|
var toolCalls []ToolCall
|
|
var usage *UsageInfo
|
|
var finishReason string
|
|
|
|
scanner := bufio.NewScanner(strings.NewReader(body))
|
|
for scanner.Scan() {
|
|
line := scanner.Text()
|
|
if !strings.HasPrefix(line, "data: ") {
|
|
continue
|
|
}
|
|
data := strings.TrimPrefix(line, "data: ")
|
|
if data == "[DONE]" {
|
|
break
|
|
}
|
|
|
|
// v1internal SSE wraps the Gemini response in a "response" field
|
|
var sseChunk struct {
|
|
Response antigravityJSONResponse `json:"response"`
|
|
}
|
|
if err := json.Unmarshal([]byte(data), &sseChunk); err != nil {
|
|
continue
|
|
}
|
|
resp := sseChunk.Response
|
|
|
|
for _, candidate := range resp.Candidates {
|
|
for _, part := range candidate.Content.Parts {
|
|
if part.Text != "" {
|
|
contentParts = append(contentParts, part.Text)
|
|
}
|
|
if part.FunctionCall != nil {
|
|
argumentsJSON, _ := json.Marshal(part.FunctionCall.Args)
|
|
toolCalls = append(toolCalls, ToolCall{
|
|
ID: fmt.Sprintf("call_%s_%d", part.FunctionCall.Name, time.Now().UnixNano()),
|
|
Name: part.FunctionCall.Name,
|
|
Arguments: part.FunctionCall.Args,
|
|
Function: &FunctionCall{
|
|
Name: part.FunctionCall.Name,
|
|
Arguments: string(argumentsJSON),
|
|
ThoughtSignature: extractPartThoughtSignature(
|
|
part.ThoughtSignature,
|
|
part.ThoughtSignatureSnake,
|
|
),
|
|
},
|
|
})
|
|
}
|
|
}
|
|
if candidate.FinishReason != "" {
|
|
finishReason = candidate.FinishReason
|
|
}
|
|
}
|
|
|
|
if resp.UsageMetadata.TotalTokenCount > 0 {
|
|
usage = &UsageInfo{
|
|
PromptTokens: resp.UsageMetadata.PromptTokenCount,
|
|
CompletionTokens: resp.UsageMetadata.CandidatesTokenCount,
|
|
TotalTokens: resp.UsageMetadata.TotalTokenCount,
|
|
}
|
|
}
|
|
}
|
|
|
|
mappedFinish := "stop"
|
|
if len(toolCalls) > 0 {
|
|
mappedFinish = "tool_calls"
|
|
}
|
|
if finishReason == "MAX_TOKENS" {
|
|
mappedFinish = "length"
|
|
}
|
|
|
|
return &LLMResponse{
|
|
Content: strings.Join(contentParts, ""),
|
|
ToolCalls: toolCalls,
|
|
FinishReason: mappedFinish,
|
|
Usage: usage,
|
|
}, nil
|
|
}
|
|
|
|
func extractPartThoughtSignature(thoughtSignature string, thoughtSignatureSnake string) string {
|
|
if thoughtSignature != "" {
|
|
return thoughtSignature
|
|
}
|
|
if thoughtSignatureSnake != "" {
|
|
return thoughtSignatureSnake
|
|
}
|
|
return ""
|
|
}
|
|
|
|
// --- Schema sanitization ---
|
|
|
|
// Google/Gemini doesn't support many JSON Schema keywords that other providers accept.
|
|
var geminiUnsupportedKeywords = map[string]bool{
|
|
"patternProperties": true,
|
|
"additionalProperties": true,
|
|
"$schema": true,
|
|
"$id": true,
|
|
"$ref": true,
|
|
"$defs": true,
|
|
"definitions": true,
|
|
"examples": true,
|
|
"minLength": true,
|
|
"maxLength": true,
|
|
"minimum": true,
|
|
"maximum": true,
|
|
"multipleOf": true,
|
|
"pattern": true,
|
|
"format": true,
|
|
"minItems": true,
|
|
"maxItems": true,
|
|
"uniqueItems": true,
|
|
"minProperties": true,
|
|
"maxProperties": true,
|
|
}
|
|
|
|
func sanitizeSchemaForGemini(schema map[string]any) map[string]any {
|
|
if schema == nil {
|
|
return nil
|
|
}
|
|
|
|
result := make(map[string]any)
|
|
for k, v := range schema {
|
|
if geminiUnsupportedKeywords[k] {
|
|
continue
|
|
}
|
|
// Recursively sanitize nested objects
|
|
switch val := v.(type) {
|
|
case map[string]any:
|
|
result[k] = sanitizeSchemaForGemini(val)
|
|
case []any:
|
|
sanitized := make([]any, len(val))
|
|
for i, item := range val {
|
|
if m, ok := item.(map[string]any); ok {
|
|
sanitized[i] = sanitizeSchemaForGemini(m)
|
|
} else {
|
|
sanitized[i] = item
|
|
}
|
|
}
|
|
result[k] = sanitized
|
|
default:
|
|
result[k] = v
|
|
}
|
|
}
|
|
|
|
// Ensure top-level has type: "object" if properties are present
|
|
if _, hasProps := result["properties"]; hasProps {
|
|
if _, hasType := result["type"]; !hasType {
|
|
result["type"] = "object"
|
|
}
|
|
}
|
|
|
|
return result
|
|
}
|
|
|
|
// --- Token source ---
|
|
|
|
func createAntigravityTokenSource() func() (string, string, error) {
|
|
return func() (string, string, error) {
|
|
cred, err := auth.GetCredential("google-antigravity")
|
|
if err != nil {
|
|
return "", "", fmt.Errorf("loading auth credentials: %w", err)
|
|
}
|
|
if cred == nil {
|
|
return "", "", fmt.Errorf(
|
|
"no credentials for google-antigravity. Run: picoclaw auth login --provider google-antigravity",
|
|
)
|
|
}
|
|
|
|
// Refresh if needed
|
|
if cred.NeedsRefresh() && cred.RefreshToken != "" {
|
|
oauthCfg := auth.GoogleAntigravityOAuthConfig()
|
|
refreshed, err := auth.RefreshAccessToken(cred, oauthCfg)
|
|
if err != nil {
|
|
return "", "", fmt.Errorf("refreshing token: %w", err)
|
|
}
|
|
refreshed.Email = cred.Email
|
|
if refreshed.ProjectID == "" {
|
|
refreshed.ProjectID = cred.ProjectID
|
|
}
|
|
if err := auth.SetCredential("google-antigravity", refreshed); err != nil {
|
|
return "", "", fmt.Errorf("saving refreshed token: %w", err)
|
|
}
|
|
cred = refreshed
|
|
}
|
|
|
|
if cred.IsExpired() {
|
|
return "", "", fmt.Errorf(
|
|
"antigravity credentials expired. Run: picoclaw auth login --provider google-antigravity",
|
|
)
|
|
}
|
|
|
|
projectID := cred.ProjectID
|
|
if projectID == "" {
|
|
// Try to fetch project ID from API
|
|
fetchedID, err := FetchAntigravityProjectID(cred.AccessToken)
|
|
if err != nil {
|
|
logger.WarnCF("provider.antigravity", "Could not fetch project ID, using fallback", map[string]any{
|
|
"error": err.Error(),
|
|
})
|
|
projectID = "rising-fact-p41fc" // Default fallback (same as OpenCode)
|
|
} else {
|
|
projectID = fetchedID
|
|
cred.ProjectID = projectID
|
|
_ = auth.SetCredential("google-antigravity", cred)
|
|
}
|
|
}
|
|
|
|
return cred.AccessToken, projectID, nil
|
|
}
|
|
}
|
|
|
|
// FetchAntigravityProjectID retrieves the Google Cloud project ID from the loadCodeAssist endpoint.
|
|
func FetchAntigravityProjectID(accessToken string) (string, error) {
|
|
reqBody, _ := json.Marshal(map[string]any{
|
|
"metadata": map[string]any{
|
|
"ideType": "IDE_UNSPECIFIED",
|
|
"platform": "PLATFORM_UNSPECIFIED",
|
|
"pluginType": "GEMINI",
|
|
},
|
|
})
|
|
|
|
req, err := http.NewRequest("POST", antigravityBaseURL+"/v1internal:loadCodeAssist", bytes.NewReader(reqBody))
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
req.Header.Set("Authorization", "Bearer "+accessToken)
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("User-Agent", antigravityUserAgent)
|
|
req.Header.Set("X-Goog-Api-Client", antigravityXGoogClient)
|
|
|
|
client := &http.Client{Timeout: 15 * time.Second}
|
|
resp, err := client.Do(req)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
body, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
return "", fmt.Errorf("reading loadCodeAssist response: %w", err)
|
|
}
|
|
if resp.StatusCode != http.StatusOK {
|
|
return "", fmt.Errorf("loadCodeAssist failed: %s", string(body))
|
|
}
|
|
|
|
var result struct {
|
|
CloudAICompanionProject string `json:"cloudaicompanionProject"`
|
|
}
|
|
if err := json.Unmarshal(body, &result); err != nil {
|
|
return "", err
|
|
}
|
|
|
|
if result.CloudAICompanionProject == "" {
|
|
return "", fmt.Errorf("no project ID in loadCodeAssist response")
|
|
}
|
|
|
|
return result.CloudAICompanionProject, nil
|
|
}
|
|
|
|
// FetchAntigravityModels fetches available models from the Cloud Code Assist API.
|
|
func FetchAntigravityModels(accessToken, projectID string) ([]AntigravityModelInfo, error) {
|
|
reqBody, _ := json.Marshal(map[string]any{
|
|
"project": projectID,
|
|
})
|
|
|
|
req, err := http.NewRequest("POST", antigravityBaseURL+"/v1internal:fetchAvailableModels", bytes.NewReader(reqBody))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
req.Header.Set("Authorization", "Bearer "+accessToken)
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("User-Agent", antigravityUserAgent)
|
|
req.Header.Set("X-Goog-Api-Client", antigravityXGoogClient)
|
|
|
|
client := &http.Client{Timeout: 15 * time.Second}
|
|
resp, err := client.Do(req)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
body, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("reading fetchAvailableModels response: %w", err)
|
|
}
|
|
if resp.StatusCode != http.StatusOK {
|
|
return nil, fmt.Errorf(
|
|
"fetchAvailableModels failed (HTTP %d): %s",
|
|
resp.StatusCode,
|
|
truncateString(string(body), 200),
|
|
)
|
|
}
|
|
|
|
var result struct {
|
|
Models map[string]struct {
|
|
DisplayName string `json:"displayName"`
|
|
QuotaInfo struct {
|
|
RemainingFraction any `json:"remainingFraction"`
|
|
ResetTime string `json:"resetTime"`
|
|
IsExhausted bool `json:"isExhausted"`
|
|
} `json:"quotaInfo"`
|
|
} `json:"models"`
|
|
}
|
|
if err := json.Unmarshal(body, &result); err != nil {
|
|
return nil, fmt.Errorf("parsing models response: %w", err)
|
|
}
|
|
|
|
var models []AntigravityModelInfo
|
|
for id, info := range result.Models {
|
|
models = append(models, AntigravityModelInfo{
|
|
ID: id,
|
|
DisplayName: info.DisplayName,
|
|
IsExhausted: info.QuotaInfo.IsExhausted,
|
|
})
|
|
}
|
|
|
|
// Ensure gemini-3-flash-preview and gemini-3-flash are in the list if they aren't already
|
|
hasFlashPreview := false
|
|
hasFlash := false
|
|
for _, m := range models {
|
|
if m.ID == "gemini-3-flash-preview" {
|
|
hasFlashPreview = true
|
|
}
|
|
if m.ID == "gemini-3-flash" {
|
|
hasFlash = true
|
|
}
|
|
}
|
|
if !hasFlashPreview {
|
|
models = append(models, AntigravityModelInfo{
|
|
ID: "gemini-3-flash-preview",
|
|
DisplayName: "Gemini 3 Flash (Preview)",
|
|
})
|
|
}
|
|
if !hasFlash {
|
|
models = append(models, AntigravityModelInfo{
|
|
ID: "gemini-3-flash",
|
|
DisplayName: "Gemini 3 Flash",
|
|
})
|
|
}
|
|
|
|
return models, nil
|
|
}
|
|
|
|
type AntigravityModelInfo struct {
|
|
ID string `json:"id"`
|
|
DisplayName string `json:"display_name"`
|
|
IsExhausted bool `json:"is_exhausted"`
|
|
}
|
|
|
|
// --- Helpers ---
|
|
|
|
func truncateString(s string, maxLen int) string {
|
|
if len(s) <= maxLen {
|
|
return s
|
|
}
|
|
return s[:maxLen] + "..."
|
|
}
|
|
|
|
func randomString(n int) string {
|
|
const letters = "abcdefghijklmnopqrstuvwxyz0123456789"
|
|
b := make([]byte, n)
|
|
for i := range b {
|
|
b[i] = letters[rand.Intn(len(letters))]
|
|
}
|
|
return string(b)
|
|
}
|
|
|
|
func (p *AntigravityProvider) parseAntigravityError(statusCode int, body []byte) error {
|
|
var errResp struct {
|
|
Error struct {
|
|
Code int `json:"code"`
|
|
Message string `json:"message"`
|
|
Status string `json:"status"`
|
|
Details []map[string]any `json:"details"`
|
|
} `json:"error"`
|
|
}
|
|
|
|
if err := json.Unmarshal(body, &errResp); err != nil {
|
|
return fmt.Errorf("antigravity API error (HTTP %d): %s", statusCode, truncateString(string(body), 500))
|
|
}
|
|
|
|
msg := errResp.Error.Message
|
|
if statusCode == 429 {
|
|
// Try to extract quota reset info
|
|
for _, detail := range errResp.Error.Details {
|
|
if typeVal, ok := detail["@type"].(string); ok && strings.HasSuffix(typeVal, "ErrorInfo") {
|
|
if metadata, ok := detail["metadata"].(map[string]any); ok {
|
|
if delay, ok := metadata["quotaResetDelay"].(string); ok {
|
|
return fmt.Errorf("antigravity rate limit exceeded: %s (reset in %s)", msg, delay)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
return fmt.Errorf("antigravity rate limit exceeded: %s", msg)
|
|
}
|
|
|
|
return fmt.Errorf("antigravity API error (%s): %s", errResp.Error.Status, msg)
|
|
}
|