diff --git a/cmd/fetch_antigravity_models/main.go b/cmd/fetch_antigravity_models/main.go new file mode 100644 index 00000000..0cf45d3b --- /dev/null +++ b/cmd/fetch_antigravity_models/main.go @@ -0,0 +1,275 @@ +// Command fetch_antigravity_models connects to the Antigravity API using the +// stored auth credentials and saves the dynamically fetched model list to a +// JSON file for inspection or offline use. +// +// Usage: +// +// go run ./cmd/fetch_antigravity_models [flags] +// +// Flags: +// +// --auths-dir Directory containing auth JSON files (default: "auths") +// --output Output JSON file path (default: "antigravity_models.json") +// --pretty Pretty-print the output JSON (default: true) +package main + +import ( + "context" + "encoding/json" + "flag" + "fmt" + "io" + "net/http" + "os" + "path/filepath" + "strings" + "time" + + "github.com/router-for-me/CLIProxyAPI/v6/internal/logging" + sdkauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/auth" + coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth" + "github.com/router-for-me/CLIProxyAPI/v6/sdk/proxyutil" + log "github.com/sirupsen/logrus" + "github.com/tidwall/gjson" +) + +const ( + antigravityBaseURLDaily = "https://daily-cloudcode-pa.googleapis.com" + antigravitySandboxBaseURLDaily = "https://daily-cloudcode-pa.sandbox.googleapis.com" + antigravityBaseURLProd = "https://cloudcode-pa.googleapis.com" + antigravityModelsPath = "/v1internal:fetchAvailableModels" +) + +func init() { + logging.SetupBaseLogger() + log.SetLevel(log.InfoLevel) +} + +// modelOutput wraps the fetched model list with fetch metadata. +type modelOutput struct { + Models []modelEntry `json:"models"` +} + +// modelEntry contains only the fields we want to keep for static model definitions. +type modelEntry struct { + ID string `json:"id"` + Object string `json:"object"` + OwnedBy string `json:"owned_by"` + Type string `json:"type"` + DisplayName string `json:"display_name"` + Name string `json:"name"` + Description string `json:"description"` + ContextLength int `json:"context_length,omitempty"` + MaxCompletionTokens int `json:"max_completion_tokens,omitempty"` +} + +func main() { + var authsDir string + var outputPath string + var pretty bool + + flag.StringVar(&authsDir, "auths-dir", "auths", "Directory containing auth JSON files") + flag.StringVar(&outputPath, "output", "antigravity_models.json", "Output JSON file path") + flag.BoolVar(&pretty, "pretty", true, "Pretty-print the output JSON") + flag.Parse() + + // Resolve relative paths against the working directory. + wd, err := os.Getwd() + if err != nil { + fmt.Fprintf(os.Stderr, "error: cannot get working directory: %v\n", err) + os.Exit(1) + } + if !filepath.IsAbs(authsDir) { + authsDir = filepath.Join(wd, authsDir) + } + if !filepath.IsAbs(outputPath) { + outputPath = filepath.Join(wd, outputPath) + } + + fmt.Printf("Scanning auth files in: %s\n", authsDir) + + // Load all auth records from the directory. + fileStore := sdkauth.NewFileTokenStore() + fileStore.SetBaseDir(authsDir) + + ctx := context.Background() + auths, err := fileStore.List(ctx) + if err != nil { + fmt.Fprintf(os.Stderr, "error: failed to list auth files: %v\n", err) + os.Exit(1) + } + if len(auths) == 0 { + fmt.Fprintf(os.Stderr, "error: no auth files found in %s\n", authsDir) + os.Exit(1) + } + + // Find the first enabled antigravity auth. + var chosen *coreauth.Auth + for _, a := range auths { + if a == nil || a.Disabled { + continue + } + if strings.EqualFold(strings.TrimSpace(a.Provider), "antigravity") { + chosen = a + break + } + } + if chosen == nil { + fmt.Fprintf(os.Stderr, "error: no enabled antigravity auth found in %s\n", authsDir) + os.Exit(1) + } + + fmt.Printf("Using auth: id=%s label=%s\n", chosen.ID, chosen.Label) + + // Fetch models from the upstream Antigravity API. + fmt.Println("Fetching Antigravity model list from upstream...") + + fetchCtx, cancel := context.WithTimeout(ctx, 30*time.Second) + defer cancel() + + models := fetchModels(fetchCtx, chosen) + if len(models) == 0 { + fmt.Fprintln(os.Stderr, "warning: no models returned (API may be unavailable or token expired)") + } else { + fmt.Printf("Fetched %d models.\n", len(models)) + } + + // Build the output payload. + out := modelOutput{ + Models: models, + } + + // Marshal to JSON. + var raw []byte + if pretty { + raw, err = json.MarshalIndent(out, "", " ") + } else { + raw, err = json.Marshal(out) + } + if err != nil { + fmt.Fprintf(os.Stderr, "error: failed to marshal JSON: %v\n", err) + os.Exit(1) + } + + if err = os.WriteFile(outputPath, raw, 0o644); err != nil { + fmt.Fprintf(os.Stderr, "error: failed to write output file %s: %v\n", outputPath, err) + os.Exit(1) + } + + fmt.Printf("Model list saved to: %s\n", outputPath) +} + +func fetchModels(ctx context.Context, auth *coreauth.Auth) []modelEntry { + accessToken := metaStringValue(auth.Metadata, "access_token") + if accessToken == "" { + fmt.Fprintln(os.Stderr, "error: no access token found in auth") + return nil + } + + baseURLs := []string{antigravityBaseURLProd, antigravityBaseURLDaily, antigravitySandboxBaseURLDaily} + + for _, baseURL := range baseURLs { + modelsURL := baseURL + antigravityModelsPath + + var payload []byte + if auth != nil && auth.Metadata != nil { + if pid, ok := auth.Metadata["project_id"].(string); ok && strings.TrimSpace(pid) != "" { + payload = []byte(fmt.Sprintf(`{"project": "%s"}`, strings.TrimSpace(pid))) + } + } + if len(payload) == 0 { + payload = []byte(`{}`) + } + + httpReq, errReq := http.NewRequestWithContext(ctx, http.MethodPost, modelsURL, strings.NewReader(string(payload))) + if errReq != nil { + continue + } + httpReq.Close = true + httpReq.Header.Set("Content-Type", "application/json") + httpReq.Header.Set("Authorization", "Bearer "+accessToken) + httpReq.Header.Set("User-Agent", "antigravity/1.19.6 darwin/arm64") + + httpClient := &http.Client{Timeout: 30 * time.Second} + if transport, _, errProxy := proxyutil.BuildHTTPTransport(auth.ProxyURL); errProxy == nil && transport != nil { + httpClient.Transport = transport + } + httpResp, errDo := httpClient.Do(httpReq) + if errDo != nil { + continue + } + + bodyBytes, errRead := io.ReadAll(httpResp.Body) + httpResp.Body.Close() + if errRead != nil { + continue + } + + if httpResp.StatusCode < http.StatusOK || httpResp.StatusCode >= http.StatusMultipleChoices { + continue + } + + result := gjson.GetBytes(bodyBytes, "models") + if !result.Exists() { + continue + } + + var models []modelEntry + + for originalName, modelData := range result.Map() { + modelID := strings.TrimSpace(originalName) + if modelID == "" { + continue + } + // Skip internal/experimental models + switch modelID { + case "chat_20706", "chat_23310", "tab_flash_lite_preview", "tab_jump_flash_lite_preview", "gemini-2.5-flash-thinking", "gemini-2.5-pro": + continue + } + + displayName := modelData.Get("displayName").String() + if displayName == "" { + displayName = modelID + } + + entry := modelEntry{ + ID: modelID, + Object: "model", + OwnedBy: "antigravity", + Type: "antigravity", + DisplayName: displayName, + Name: modelID, + Description: displayName, + } + + if maxTok := modelData.Get("maxTokens").Int(); maxTok > 0 { + entry.ContextLength = int(maxTok) + } + if maxOut := modelData.Get("maxOutputTokens").Int(); maxOut > 0 { + entry.MaxCompletionTokens = int(maxOut) + } + + models = append(models, entry) + } + + return models + } + + return nil +} + +func metaStringValue(m map[string]interface{}, key string) string { + if m == nil { + return "" + } + v, ok := m[key] + if !ok { + return "" + } + switch val := v.(type) { + case string: + return val + default: + return "" + } +}