Merge pull request #2082 from router-for-me/antigravity
Refactor Antigravity model handling and improve logging
This commit is contained in:
275
cmd/fetch_antigravity_models/main.go
Normal file
275
cmd/fetch_antigravity_models/main.go
Normal file
@@ -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 <path> Directory containing auth JSON files (default: "auths")
|
||||
// --output <path> 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 ""
|
||||
}
|
||||
}
|
||||
@@ -3,32 +3,24 @@
|
||||
package registry
|
||||
|
||||
import (
|
||||
"sort"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// AntigravityModelConfig captures static antigravity model overrides, including
|
||||
// Thinking budget limits and provider max completion tokens.
|
||||
type AntigravityModelConfig struct {
|
||||
Thinking *ThinkingSupport `json:"thinking,omitempty"`
|
||||
MaxCompletionTokens int `json:"max_completion_tokens,omitempty"`
|
||||
}
|
||||
|
||||
// staticModelsJSON mirrors the top-level structure of models.json.
|
||||
type staticModelsJSON struct {
|
||||
Claude []*ModelInfo `json:"claude"`
|
||||
Gemini []*ModelInfo `json:"gemini"`
|
||||
Vertex []*ModelInfo `json:"vertex"`
|
||||
GeminiCLI []*ModelInfo `json:"gemini-cli"`
|
||||
AIStudio []*ModelInfo `json:"aistudio"`
|
||||
CodexFree []*ModelInfo `json:"codex-free"`
|
||||
CodexTeam []*ModelInfo `json:"codex-team"`
|
||||
CodexPlus []*ModelInfo `json:"codex-plus"`
|
||||
CodexPro []*ModelInfo `json:"codex-pro"`
|
||||
Qwen []*ModelInfo `json:"qwen"`
|
||||
IFlow []*ModelInfo `json:"iflow"`
|
||||
Kimi []*ModelInfo `json:"kimi"`
|
||||
Antigravity map[string]*AntigravityModelConfig `json:"antigravity"`
|
||||
Claude []*ModelInfo `json:"claude"`
|
||||
Gemini []*ModelInfo `json:"gemini"`
|
||||
Vertex []*ModelInfo `json:"vertex"`
|
||||
GeminiCLI []*ModelInfo `json:"gemini-cli"`
|
||||
AIStudio []*ModelInfo `json:"aistudio"`
|
||||
CodexFree []*ModelInfo `json:"codex-free"`
|
||||
CodexTeam []*ModelInfo `json:"codex-team"`
|
||||
CodexPlus []*ModelInfo `json:"codex-plus"`
|
||||
CodexPro []*ModelInfo `json:"codex-pro"`
|
||||
Qwen []*ModelInfo `json:"qwen"`
|
||||
IFlow []*ModelInfo `json:"iflow"`
|
||||
Kimi []*ModelInfo `json:"kimi"`
|
||||
Antigravity []*ModelInfo `json:"antigravity"`
|
||||
}
|
||||
|
||||
// GetClaudeModels returns the standard Claude model definitions.
|
||||
@@ -91,33 +83,9 @@ func GetKimiModels() []*ModelInfo {
|
||||
return cloneModelInfos(getModels().Kimi)
|
||||
}
|
||||
|
||||
// GetAntigravityModelConfig returns static configuration for antigravity models.
|
||||
// Keys use upstream model names returned by the Antigravity models endpoint.
|
||||
func GetAntigravityModelConfig() map[string]*AntigravityModelConfig {
|
||||
data := getModels()
|
||||
if len(data.Antigravity) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make(map[string]*AntigravityModelConfig, len(data.Antigravity))
|
||||
for k, v := range data.Antigravity {
|
||||
out[k] = cloneAntigravityModelConfig(v)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func cloneAntigravityModelConfig(cfg *AntigravityModelConfig) *AntigravityModelConfig {
|
||||
if cfg == nil {
|
||||
return nil
|
||||
}
|
||||
copyConfig := *cfg
|
||||
if cfg.Thinking != nil {
|
||||
copyThinking := *cfg.Thinking
|
||||
if len(cfg.Thinking.Levels) > 0 {
|
||||
copyThinking.Levels = append([]string(nil), cfg.Thinking.Levels...)
|
||||
}
|
||||
copyConfig.Thinking = ©Thinking
|
||||
}
|
||||
return ©Config
|
||||
// GetAntigravityModels returns the standard Antigravity model definitions.
|
||||
func GetAntigravityModels() []*ModelInfo {
|
||||
return cloneModelInfos(getModels().Antigravity)
|
||||
}
|
||||
|
||||
// cloneModelInfos returns a shallow copy of the slice with each element deep-cloned.
|
||||
@@ -145,7 +113,7 @@ func cloneModelInfos(models []*ModelInfo) []*ModelInfo {
|
||||
// - qwen
|
||||
// - iflow
|
||||
// - kimi
|
||||
// - antigravity (returns static overrides only)
|
||||
// - antigravity
|
||||
func GetStaticModelDefinitionsByChannel(channel string) []*ModelInfo {
|
||||
key := strings.ToLower(strings.TrimSpace(channel))
|
||||
switch key {
|
||||
@@ -168,28 +136,7 @@ func GetStaticModelDefinitionsByChannel(channel string) []*ModelInfo {
|
||||
case "kimi":
|
||||
return GetKimiModels()
|
||||
case "antigravity":
|
||||
cfg := GetAntigravityModelConfig()
|
||||
if len(cfg) == 0 {
|
||||
return nil
|
||||
}
|
||||
models := make([]*ModelInfo, 0, len(cfg))
|
||||
for modelID, entry := range cfg {
|
||||
if modelID == "" || entry == nil {
|
||||
continue
|
||||
}
|
||||
models = append(models, &ModelInfo{
|
||||
ID: modelID,
|
||||
Object: "model",
|
||||
OwnedBy: "antigravity",
|
||||
Type: "antigravity",
|
||||
Thinking: entry.Thinking,
|
||||
MaxCompletionTokens: entry.MaxCompletionTokens,
|
||||
})
|
||||
}
|
||||
sort.Slice(models, func(i, j int) bool {
|
||||
return strings.ToLower(models[i].ID) < strings.ToLower(models[j].ID)
|
||||
})
|
||||
return models
|
||||
return GetAntigravityModels()
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
@@ -213,6 +160,7 @@ func LookupStaticModelInfo(modelID string) *ModelInfo {
|
||||
data.Qwen,
|
||||
data.IFlow,
|
||||
data.Kimi,
|
||||
data.Antigravity,
|
||||
}
|
||||
for _, models := range allModels {
|
||||
for _, m := range models {
|
||||
@@ -222,14 +170,5 @@ func LookupStaticModelInfo(modelID string) *ModelInfo {
|
||||
}
|
||||
}
|
||||
|
||||
// Check Antigravity static config
|
||||
if cfg := cloneAntigravityModelConfig(data.Antigravity[modelID]); cfg != nil {
|
||||
return &ModelInfo{
|
||||
ID: modelID,
|
||||
Thinking: cfg.Thinking,
|
||||
MaxCompletionTokens: cfg.MaxCompletionTokens,
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -100,7 +100,7 @@ func tryRefreshModels(ctx context.Context) {
|
||||
log.Infof("models updated from %s", url)
|
||||
return
|
||||
}
|
||||
log.Warn("models refresh failed from all URLs, using current data")
|
||||
log.Warn("models refresh failed from all URLs, using local data")
|
||||
}
|
||||
|
||||
func loadModelsFromBytes(data []byte, source string) error {
|
||||
@@ -145,6 +145,7 @@ func validateModelsCatalog(data *staticModelsJSON) error {
|
||||
{name: "qwen", models: data.Qwen},
|
||||
{name: "iflow", models: data.IFlow},
|
||||
{name: "kimi", models: data.Kimi},
|
||||
{name: "antigravity", models: data.Antigravity},
|
||||
}
|
||||
|
||||
for _, section := range requiredSections {
|
||||
@@ -152,9 +153,6 @@ func validateModelsCatalog(data *staticModelsJSON) error {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if err := validateAntigravitySection(data.Antigravity); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -179,20 +177,3 @@ func validateModelSection(section string, models []*ModelInfo) error {
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func validateAntigravitySection(configs map[string]*AntigravityModelConfig) error {
|
||||
if len(configs) == 0 {
|
||||
return fmt.Errorf("antigravity section is empty")
|
||||
}
|
||||
|
||||
for modelID, cfg := range configs {
|
||||
trimmedID := strings.TrimSpace(modelID)
|
||||
if trimmedID == "" {
|
||||
return fmt.Errorf("antigravity contains empty model id")
|
||||
}
|
||||
if cfg == nil {
|
||||
return fmt.Errorf("antigravity[%q] is null", trimmedID)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -2481,40 +2481,83 @@
|
||||
}
|
||||
}
|
||||
],
|
||||
"antigravity": {
|
||||
"claude-opus-4-6-thinking": {
|
||||
"antigravity": [
|
||||
{
|
||||
"id": "claude-opus-4-6-thinking",
|
||||
"object": "model",
|
||||
"owned_by": "antigravity",
|
||||
"type": "antigravity",
|
||||
"display_name": "Claude Opus 4.6 (Thinking)",
|
||||
"name": "claude-opus-4-6-thinking",
|
||||
"description": "Claude Opus 4.6 (Thinking)",
|
||||
"context_length": 200000,
|
||||
"max_completion_tokens": 64000,
|
||||
"thinking": {
|
||||
"min": 1024,
|
||||
"max": 64000,
|
||||
"zero_allowed": true,
|
||||
"dynamic_allowed": true
|
||||
},
|
||||
"max_completion_tokens": 64000
|
||||
}
|
||||
},
|
||||
"claude-sonnet-4-6": {
|
||||
{
|
||||
"id": "claude-sonnet-4-6",
|
||||
"object": "model",
|
||||
"owned_by": "antigravity",
|
||||
"type": "antigravity",
|
||||
"display_name": "Claude Sonnet 4.6 (Thinking)",
|
||||
"name": "claude-sonnet-4-6",
|
||||
"description": "Claude Sonnet 4.6 (Thinking)",
|
||||
"context_length": 200000,
|
||||
"max_completion_tokens": 64000,
|
||||
"thinking": {
|
||||
"min": 1024,
|
||||
"max": 64000,
|
||||
"zero_allowed": true,
|
||||
"dynamic_allowed": true
|
||||
},
|
||||
"max_completion_tokens": 64000
|
||||
}
|
||||
},
|
||||
"gemini-2.5-flash": {
|
||||
{
|
||||
"id": "gemini-2.5-flash",
|
||||
"object": "model",
|
||||
"owned_by": "antigravity",
|
||||
"type": "antigravity",
|
||||
"display_name": "Gemini 2.5 Flash",
|
||||
"name": "gemini-2.5-flash",
|
||||
"description": "Gemini 2.5 Flash",
|
||||
"context_length": 1048576,
|
||||
"max_completion_tokens": 65535,
|
||||
"thinking": {
|
||||
"max": 24576,
|
||||
"zero_allowed": true,
|
||||
"dynamic_allowed": true
|
||||
}
|
||||
},
|
||||
"gemini-2.5-flash-lite": {
|
||||
{
|
||||
"id": "gemini-2.5-flash-lite",
|
||||
"object": "model",
|
||||
"owned_by": "antigravity",
|
||||
"type": "antigravity",
|
||||
"display_name": "Gemini 2.5 Flash Lite",
|
||||
"name": "gemini-2.5-flash-lite",
|
||||
"description": "Gemini 2.5 Flash Lite",
|
||||
"context_length": 1048576,
|
||||
"max_completion_tokens": 65535,
|
||||
"thinking": {
|
||||
"max": 24576,
|
||||
"zero_allowed": true,
|
||||
"dynamic_allowed": true
|
||||
}
|
||||
},
|
||||
"gemini-3-flash": {
|
||||
{
|
||||
"id": "gemini-3-flash",
|
||||
"object": "model",
|
||||
"owned_by": "antigravity",
|
||||
"type": "antigravity",
|
||||
"display_name": "Gemini 3 Flash",
|
||||
"name": "gemini-3-flash",
|
||||
"description": "Gemini 3 Flash",
|
||||
"context_length": 1048576,
|
||||
"max_completion_tokens": 65536,
|
||||
"thinking": {
|
||||
"min": 128,
|
||||
"max": 32768,
|
||||
@@ -2527,7 +2570,16 @@
|
||||
]
|
||||
}
|
||||
},
|
||||
"gemini-3-pro-high": {
|
||||
{
|
||||
"id": "gemini-3-pro-high",
|
||||
"object": "model",
|
||||
"owned_by": "antigravity",
|
||||
"type": "antigravity",
|
||||
"display_name": "Gemini 3 Pro (High)",
|
||||
"name": "gemini-3-pro-high",
|
||||
"description": "Gemini 3 Pro (High)",
|
||||
"context_length": 1048576,
|
||||
"max_completion_tokens": 65535,
|
||||
"thinking": {
|
||||
"min": 128,
|
||||
"max": 32768,
|
||||
@@ -2538,7 +2590,16 @@
|
||||
]
|
||||
}
|
||||
},
|
||||
"gemini-3-pro-low": {
|
||||
{
|
||||
"id": "gemini-3-pro-low",
|
||||
"object": "model",
|
||||
"owned_by": "antigravity",
|
||||
"type": "antigravity",
|
||||
"display_name": "Gemini 3 Pro (Low)",
|
||||
"name": "gemini-3-pro-low",
|
||||
"description": "Gemini 3 Pro (Low)",
|
||||
"context_length": 1048576,
|
||||
"max_completion_tokens": 65535,
|
||||
"thinking": {
|
||||
"min": 128,
|
||||
"max": 32768,
|
||||
@@ -2549,7 +2610,14 @@
|
||||
]
|
||||
}
|
||||
},
|
||||
"gemini-3.1-flash-image": {
|
||||
{
|
||||
"id": "gemini-3.1-flash-image",
|
||||
"object": "model",
|
||||
"owned_by": "antigravity",
|
||||
"type": "antigravity",
|
||||
"display_name": "Gemini 3.1 Flash Image",
|
||||
"name": "gemini-3.1-flash-image",
|
||||
"description": "Gemini 3.1 Flash Image",
|
||||
"thinking": {
|
||||
"min": 128,
|
||||
"max": 32768,
|
||||
@@ -2560,18 +2628,16 @@
|
||||
]
|
||||
}
|
||||
},
|
||||
"gemini-3.1-flash-lite-preview": {
|
||||
"thinking": {
|
||||
"min": 128,
|
||||
"max": 32768,
|
||||
"dynamic_allowed": true,
|
||||
"levels": [
|
||||
"minimal",
|
||||
"high"
|
||||
]
|
||||
}
|
||||
},
|
||||
"gemini-3.1-pro-high": {
|
||||
{
|
||||
"id": "gemini-3.1-pro-high",
|
||||
"object": "model",
|
||||
"owned_by": "antigravity",
|
||||
"type": "antigravity",
|
||||
"display_name": "Gemini 3.1 Pro (High)",
|
||||
"name": "gemini-3.1-pro-high",
|
||||
"description": "Gemini 3.1 Pro (High)",
|
||||
"context_length": 1048576,
|
||||
"max_completion_tokens": 65535,
|
||||
"thinking": {
|
||||
"min": 128,
|
||||
"max": 32768,
|
||||
@@ -2582,7 +2648,16 @@
|
||||
]
|
||||
}
|
||||
},
|
||||
"gemini-3.1-pro-low": {
|
||||
{
|
||||
"id": "gemini-3.1-pro-low",
|
||||
"object": "model",
|
||||
"owned_by": "antigravity",
|
||||
"type": "antigravity",
|
||||
"display_name": "Gemini 3.1 Pro (Low)",
|
||||
"name": "gemini-3.1-pro-low",
|
||||
"description": "Gemini 3.1 Pro (Low)",
|
||||
"context_length": 1048576,
|
||||
"max_completion_tokens": 65535,
|
||||
"thinking": {
|
||||
"min": 128,
|
||||
"max": 32768,
|
||||
@@ -2593,6 +2668,16 @@
|
||||
]
|
||||
}
|
||||
},
|
||||
"gpt-oss-120b-medium": {}
|
||||
}
|
||||
{
|
||||
"id": "gpt-oss-120b-medium",
|
||||
"object": "model",
|
||||
"owned_by": "antigravity",
|
||||
"type": "antigravity",
|
||||
"display_name": "GPT-OSS 120B (Medium)",
|
||||
"name": "gpt-oss-120b-medium",
|
||||
"description": "GPT-OSS 120B (Medium)",
|
||||
"context_length": 114000,
|
||||
"max_completion_tokens": 32768
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -24,7 +24,6 @@ import (
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/registry"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/thinking"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
|
||||
sdkAuth "github.com/router-for-me/CLIProxyAPI/v6/sdk/auth"
|
||||
@@ -43,7 +42,6 @@ const (
|
||||
antigravityCountTokensPath = "/v1internal:countTokens"
|
||||
antigravityStreamPath = "/v1internal:streamGenerateContent"
|
||||
antigravityGeneratePath = "/v1internal:generateContent"
|
||||
antigravityModelsPath = "/v1internal:fetchAvailableModels"
|
||||
antigravityClientID = "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com"
|
||||
antigravityClientSecret = "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"
|
||||
defaultAntigravityAgent = "antigravity/1.19.6 darwin/arm64"
|
||||
@@ -55,78 +53,8 @@ const (
|
||||
var (
|
||||
randSource = rand.New(rand.NewSource(time.Now().UnixNano()))
|
||||
randSourceMutex sync.Mutex
|
||||
// antigravityPrimaryModelsCache keeps the latest non-empty model list fetched
|
||||
// from any antigravity auth. Empty fetches never overwrite this cache.
|
||||
antigravityPrimaryModelsCache struct {
|
||||
mu sync.RWMutex
|
||||
models []*registry.ModelInfo
|
||||
}
|
||||
)
|
||||
|
||||
func cloneAntigravityModels(models []*registry.ModelInfo) []*registry.ModelInfo {
|
||||
if len(models) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]*registry.ModelInfo, 0, len(models))
|
||||
for _, model := range models {
|
||||
if model == nil || strings.TrimSpace(model.ID) == "" {
|
||||
continue
|
||||
}
|
||||
out = append(out, cloneAntigravityModelInfo(model))
|
||||
}
|
||||
if len(out) == 0 {
|
||||
return nil
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func cloneAntigravityModelInfo(model *registry.ModelInfo) *registry.ModelInfo {
|
||||
if model == nil {
|
||||
return nil
|
||||
}
|
||||
clone := *model
|
||||
if len(model.SupportedGenerationMethods) > 0 {
|
||||
clone.SupportedGenerationMethods = append([]string(nil), model.SupportedGenerationMethods...)
|
||||
}
|
||||
if len(model.SupportedParameters) > 0 {
|
||||
clone.SupportedParameters = append([]string(nil), model.SupportedParameters...)
|
||||
}
|
||||
if model.Thinking != nil {
|
||||
thinkingClone := *model.Thinking
|
||||
if len(model.Thinking.Levels) > 0 {
|
||||
thinkingClone.Levels = append([]string(nil), model.Thinking.Levels...)
|
||||
}
|
||||
clone.Thinking = &thinkingClone
|
||||
}
|
||||
return &clone
|
||||
}
|
||||
|
||||
func storeAntigravityPrimaryModels(models []*registry.ModelInfo) bool {
|
||||
cloned := cloneAntigravityModels(models)
|
||||
if len(cloned) == 0 {
|
||||
return false
|
||||
}
|
||||
antigravityPrimaryModelsCache.mu.Lock()
|
||||
antigravityPrimaryModelsCache.models = cloned
|
||||
antigravityPrimaryModelsCache.mu.Unlock()
|
||||
return true
|
||||
}
|
||||
|
||||
func loadAntigravityPrimaryModels() []*registry.ModelInfo {
|
||||
antigravityPrimaryModelsCache.mu.RLock()
|
||||
cloned := cloneAntigravityModels(antigravityPrimaryModelsCache.models)
|
||||
antigravityPrimaryModelsCache.mu.RUnlock()
|
||||
return cloned
|
||||
}
|
||||
|
||||
func fallbackAntigravityPrimaryModels() []*registry.ModelInfo {
|
||||
models := loadAntigravityPrimaryModels()
|
||||
if len(models) > 0 {
|
||||
log.Debugf("antigravity executor: using cached primary model list (%d models)", len(models))
|
||||
}
|
||||
return models
|
||||
}
|
||||
|
||||
// AntigravityExecutor proxies requests to the antigravity upstream.
|
||||
type AntigravityExecutor struct {
|
||||
cfg *config.Config
|
||||
@@ -1150,168 +1078,6 @@ func (e *AntigravityExecutor) CountTokens(ctx context.Context, auth *cliproxyaut
|
||||
}
|
||||
}
|
||||
|
||||
// FetchAntigravityModels retrieves available models using the supplied auth.
|
||||
func FetchAntigravityModels(ctx context.Context, auth *cliproxyauth.Auth, cfg *config.Config) []*registry.ModelInfo {
|
||||
exec := &AntigravityExecutor{cfg: cfg}
|
||||
token, updatedAuth, errToken := exec.ensureAccessToken(ctx, auth)
|
||||
if errToken != nil || token == "" {
|
||||
return fallbackAntigravityPrimaryModels()
|
||||
}
|
||||
if updatedAuth != nil {
|
||||
auth = updatedAuth
|
||||
}
|
||||
|
||||
baseURLs := antigravityBaseURLFallbackOrder(auth)
|
||||
httpClient := newAntigravityHTTPClient(ctx, cfg, auth, 0)
|
||||
|
||||
for idx, 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, bytes.NewReader(payload))
|
||||
if errReq != nil {
|
||||
return fallbackAntigravityPrimaryModels()
|
||||
}
|
||||
httpReq.Close = true
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
httpReq.Header.Set("Authorization", "Bearer "+token)
|
||||
httpReq.Header.Set("User-Agent", resolveUserAgent(auth))
|
||||
if host := resolveHost(baseURL); host != "" {
|
||||
httpReq.Host = host
|
||||
}
|
||||
|
||||
httpResp, errDo := httpClient.Do(httpReq)
|
||||
if errDo != nil {
|
||||
if errors.Is(errDo, context.Canceled) || errors.Is(errDo, context.DeadlineExceeded) {
|
||||
return fallbackAntigravityPrimaryModels()
|
||||
}
|
||||
if idx+1 < len(baseURLs) {
|
||||
log.Debugf("antigravity executor: models request error on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1])
|
||||
continue
|
||||
}
|
||||
return fallbackAntigravityPrimaryModels()
|
||||
}
|
||||
|
||||
bodyBytes, errRead := io.ReadAll(httpResp.Body)
|
||||
if errClose := httpResp.Body.Close(); errClose != nil {
|
||||
log.Errorf("antigravity executor: close response body error: %v", errClose)
|
||||
}
|
||||
if errRead != nil {
|
||||
if idx+1 < len(baseURLs) {
|
||||
log.Debugf("antigravity executor: models read error on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1])
|
||||
continue
|
||||
}
|
||||
return fallbackAntigravityPrimaryModels()
|
||||
}
|
||||
if httpResp.StatusCode < http.StatusOK || httpResp.StatusCode >= http.StatusMultipleChoices {
|
||||
if httpResp.StatusCode == http.StatusTooManyRequests && idx+1 < len(baseURLs) {
|
||||
log.Debugf("antigravity executor: models request rate limited on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1])
|
||||
continue
|
||||
}
|
||||
if idx+1 < len(baseURLs) {
|
||||
log.Debugf("antigravity executor: models request failed with status %d on base url %s, retrying with fallback base url: %s", httpResp.StatusCode, baseURL, baseURLs[idx+1])
|
||||
continue
|
||||
}
|
||||
return fallbackAntigravityPrimaryModels()
|
||||
}
|
||||
|
||||
result := gjson.GetBytes(bodyBytes, "models")
|
||||
if !result.Exists() {
|
||||
if idx+1 < len(baseURLs) {
|
||||
log.Debugf("antigravity executor: models field missing on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1])
|
||||
continue
|
||||
}
|
||||
return fallbackAntigravityPrimaryModels()
|
||||
}
|
||||
|
||||
now := time.Now().Unix()
|
||||
modelConfig := registry.GetAntigravityModelConfig()
|
||||
models := make([]*registry.ModelInfo, 0, len(result.Map()))
|
||||
for originalName, modelData := range result.Map() {
|
||||
modelID := strings.TrimSpace(originalName)
|
||||
if modelID == "" {
|
||||
continue
|
||||
}
|
||||
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
|
||||
}
|
||||
modelCfg := modelConfig[modelID]
|
||||
|
||||
// Extract displayName from upstream response, fallback to modelID
|
||||
displayName := modelData.Get("displayName").String()
|
||||
if displayName == "" {
|
||||
displayName = modelID
|
||||
}
|
||||
|
||||
modelInfo := ®istry.ModelInfo{
|
||||
ID: modelID,
|
||||
Name: modelID,
|
||||
Description: displayName,
|
||||
DisplayName: displayName,
|
||||
Version: modelID,
|
||||
Object: "model",
|
||||
Created: now,
|
||||
OwnedBy: antigravityAuthType,
|
||||
Type: antigravityAuthType,
|
||||
}
|
||||
|
||||
// Build input modalities from upstream capability flags.
|
||||
inputModalities := []string{"TEXT"}
|
||||
if modelData.Get("supportsImages").Bool() {
|
||||
inputModalities = append(inputModalities, "IMAGE")
|
||||
}
|
||||
if modelData.Get("supportsVideo").Bool() {
|
||||
inputModalities = append(inputModalities, "VIDEO")
|
||||
}
|
||||
modelInfo.SupportedInputModalities = inputModalities
|
||||
modelInfo.SupportedOutputModalities = []string{"TEXT"}
|
||||
|
||||
// Token limits from upstream.
|
||||
if maxTok := modelData.Get("maxTokens").Int(); maxTok > 0 {
|
||||
modelInfo.InputTokenLimit = int(maxTok)
|
||||
}
|
||||
if maxOut := modelData.Get("maxOutputTokens").Int(); maxOut > 0 {
|
||||
modelInfo.OutputTokenLimit = int(maxOut)
|
||||
}
|
||||
|
||||
// Supported generation methods (Gemini v1beta convention).
|
||||
modelInfo.SupportedGenerationMethods = []string{"generateContent", "countTokens"}
|
||||
|
||||
// Look up Thinking support from static config using upstream model name.
|
||||
if modelCfg != nil {
|
||||
if modelCfg.Thinking != nil {
|
||||
modelInfo.Thinking = modelCfg.Thinking
|
||||
}
|
||||
if modelCfg.MaxCompletionTokens > 0 {
|
||||
modelInfo.MaxCompletionTokens = modelCfg.MaxCompletionTokens
|
||||
}
|
||||
}
|
||||
models = append(models, modelInfo)
|
||||
}
|
||||
if len(models) == 0 {
|
||||
if idx+1 < len(baseURLs) {
|
||||
log.Debugf("antigravity executor: empty models list on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1])
|
||||
continue
|
||||
}
|
||||
log.Debug("antigravity executor: fetched empty model list; retaining cached primary model list")
|
||||
return fallbackAntigravityPrimaryModels()
|
||||
}
|
||||
storeAntigravityPrimaryModels(models)
|
||||
return models
|
||||
}
|
||||
return fallbackAntigravityPrimaryModels()
|
||||
}
|
||||
|
||||
func (e *AntigravityExecutor) ensureAccessToken(ctx context.Context, auth *cliproxyauth.Auth) (string, *cliproxyauth.Auth, error) {
|
||||
if auth == nil {
|
||||
return "", nil, statusErr{code: http.StatusUnauthorized, msg: "missing auth"}
|
||||
|
||||
@@ -1,90 +0,0 @@
|
||||
package executor
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/registry"
|
||||
)
|
||||
|
||||
func resetAntigravityPrimaryModelsCacheForTest() {
|
||||
antigravityPrimaryModelsCache.mu.Lock()
|
||||
antigravityPrimaryModelsCache.models = nil
|
||||
antigravityPrimaryModelsCache.mu.Unlock()
|
||||
}
|
||||
|
||||
func TestStoreAntigravityPrimaryModels_EmptyDoesNotOverwrite(t *testing.T) {
|
||||
resetAntigravityPrimaryModelsCacheForTest()
|
||||
t.Cleanup(resetAntigravityPrimaryModelsCacheForTest)
|
||||
|
||||
seed := []*registry.ModelInfo{
|
||||
{ID: "claude-sonnet-4-5"},
|
||||
{ID: "gemini-2.5-pro"},
|
||||
}
|
||||
if updated := storeAntigravityPrimaryModels(seed); !updated {
|
||||
t.Fatal("expected non-empty model list to update primary cache")
|
||||
}
|
||||
|
||||
if updated := storeAntigravityPrimaryModels(nil); updated {
|
||||
t.Fatal("expected nil model list not to overwrite primary cache")
|
||||
}
|
||||
if updated := storeAntigravityPrimaryModels([]*registry.ModelInfo{}); updated {
|
||||
t.Fatal("expected empty model list not to overwrite primary cache")
|
||||
}
|
||||
|
||||
got := loadAntigravityPrimaryModels()
|
||||
if len(got) != 2 {
|
||||
t.Fatalf("expected cached model count 2, got %d", len(got))
|
||||
}
|
||||
if got[0].ID != "claude-sonnet-4-5" || got[1].ID != "gemini-2.5-pro" {
|
||||
t.Fatalf("unexpected cached model ids: %q, %q", got[0].ID, got[1].ID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadAntigravityPrimaryModels_ReturnsClone(t *testing.T) {
|
||||
resetAntigravityPrimaryModelsCacheForTest()
|
||||
t.Cleanup(resetAntigravityPrimaryModelsCacheForTest)
|
||||
|
||||
if updated := storeAntigravityPrimaryModels([]*registry.ModelInfo{{
|
||||
ID: "gpt-5",
|
||||
DisplayName: "GPT-5",
|
||||
SupportedGenerationMethods: []string{"generateContent"},
|
||||
SupportedParameters: []string{"temperature"},
|
||||
Thinking: ®istry.ThinkingSupport{
|
||||
Levels: []string{"high"},
|
||||
},
|
||||
}}); !updated {
|
||||
t.Fatal("expected model cache update")
|
||||
}
|
||||
|
||||
got := loadAntigravityPrimaryModels()
|
||||
if len(got) != 1 {
|
||||
t.Fatalf("expected one cached model, got %d", len(got))
|
||||
}
|
||||
got[0].ID = "mutated-id"
|
||||
if len(got[0].SupportedGenerationMethods) > 0 {
|
||||
got[0].SupportedGenerationMethods[0] = "mutated-method"
|
||||
}
|
||||
if len(got[0].SupportedParameters) > 0 {
|
||||
got[0].SupportedParameters[0] = "mutated-parameter"
|
||||
}
|
||||
if got[0].Thinking != nil && len(got[0].Thinking.Levels) > 0 {
|
||||
got[0].Thinking.Levels[0] = "mutated-level"
|
||||
}
|
||||
|
||||
again := loadAntigravityPrimaryModels()
|
||||
if len(again) != 1 {
|
||||
t.Fatalf("expected one cached model after mutation, got %d", len(again))
|
||||
}
|
||||
if again[0].ID != "gpt-5" {
|
||||
t.Fatalf("expected cached model id to remain %q, got %q", "gpt-5", again[0].ID)
|
||||
}
|
||||
if len(again[0].SupportedGenerationMethods) == 0 || again[0].SupportedGenerationMethods[0] != "generateContent" {
|
||||
t.Fatalf("expected cached generation methods to be unmutated, got %v", again[0].SupportedGenerationMethods)
|
||||
}
|
||||
if len(again[0].SupportedParameters) == 0 || again[0].SupportedParameters[0] != "temperature" {
|
||||
t.Fatalf("expected cached supported parameters to be unmutated, got %v", again[0].SupportedParameters)
|
||||
}
|
||||
if again[0].Thinking == nil || len(again[0].Thinking.Levels) == 0 || again[0].Thinking.Levels[0] != "high" {
|
||||
t.Fatalf("expected cached model thinking levels to be unmutated, got %v", again[0].Thinking)
|
||||
}
|
||||
}
|
||||
@@ -282,8 +282,6 @@ func (s *Service) applyCoreAuthAddOrUpdate(ctx context.Context, auth *coreauth.A
|
||||
// IMPORTANT: Update coreManager FIRST, before model registration.
|
||||
// This ensures that configuration changes (proxy_url, prefix, etc.) take effect
|
||||
// immediately for API calls, rather than waiting for model registration to complete.
|
||||
// Model registration may involve network calls (e.g., FetchAntigravityModels) that
|
||||
// could timeout if the new proxy_url is unreachable.
|
||||
op := "register"
|
||||
var err error
|
||||
if existing, ok := s.coreManager.GetByID(auth.ID); ok {
|
||||
@@ -813,9 +811,7 @@ func (s *Service) registerModelsForAuth(a *coreauth.Auth) {
|
||||
models = registry.GetAIStudioModels()
|
||||
models = applyExcludedModels(models, excluded)
|
||||
case "antigravity":
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
|
||||
models = executor.FetchAntigravityModels(ctx, a, s.cfg)
|
||||
cancel()
|
||||
models = registry.GetAntigravityModels()
|
||||
models = applyExcludedModels(models, excluded)
|
||||
case "claude":
|
||||
models = registry.GetClaudeModels()
|
||||
@@ -952,9 +948,6 @@ func (s *Service) registerModelsForAuth(a *coreauth.Auth) {
|
||||
key = strings.ToLower(strings.TrimSpace(a.Provider))
|
||||
}
|
||||
GlobalModelRegistry().RegisterClient(a.ID, key, applyModelPrefixes(models, a.Prefix, s.cfg != nil && s.cfg.ForceModelPrefix))
|
||||
if provider == "antigravity" {
|
||||
s.backfillAntigravityModels(a, models)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
@@ -1099,56 +1092,6 @@ func (s *Service) oauthExcludedModels(provider, authKind string) []string {
|
||||
return cfg.OAuthExcludedModels[providerKey]
|
||||
}
|
||||
|
||||
func (s *Service) backfillAntigravityModels(source *coreauth.Auth, primaryModels []*ModelInfo) {
|
||||
if s == nil || s.coreManager == nil || len(primaryModels) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
sourceID := ""
|
||||
if source != nil {
|
||||
sourceID = strings.TrimSpace(source.ID)
|
||||
}
|
||||
|
||||
reg := registry.GetGlobalRegistry()
|
||||
for _, candidate := range s.coreManager.List() {
|
||||
if candidate == nil || candidate.Disabled {
|
||||
continue
|
||||
}
|
||||
candidateID := strings.TrimSpace(candidate.ID)
|
||||
if candidateID == "" || candidateID == sourceID {
|
||||
continue
|
||||
}
|
||||
if !strings.EqualFold(strings.TrimSpace(candidate.Provider), "antigravity") {
|
||||
continue
|
||||
}
|
||||
if len(reg.GetModelsForClient(candidateID)) > 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
authKind := strings.ToLower(strings.TrimSpace(candidate.Attributes["auth_kind"]))
|
||||
if authKind == "" {
|
||||
if kind, _ := candidate.AccountInfo(); strings.EqualFold(kind, "api_key") {
|
||||
authKind = "apikey"
|
||||
}
|
||||
}
|
||||
excluded := s.oauthExcludedModels("antigravity", authKind)
|
||||
if candidate.Attributes != nil {
|
||||
if val, ok := candidate.Attributes["excluded_models"]; ok && strings.TrimSpace(val) != "" {
|
||||
excluded = strings.Split(val, ",")
|
||||
}
|
||||
}
|
||||
|
||||
models := applyExcludedModels(primaryModels, excluded)
|
||||
models = applyOAuthModelAlias(s.cfg, "antigravity", authKind, models)
|
||||
if len(models) == 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
reg.RegisterClient(candidateID, "antigravity", applyModelPrefixes(models, candidate.Prefix, s.cfg != nil && s.cfg.ForceModelPrefix))
|
||||
log.Debugf("antigravity models backfilled for auth %s using primary model list", candidateID)
|
||||
}
|
||||
}
|
||||
|
||||
func applyExcludedModels(models []*ModelInfo, excluded []string) []*ModelInfo {
|
||||
if len(models) == 0 || len(excluded) == 0 {
|
||||
return models
|
||||
|
||||
@@ -1,135 +0,0 @@
|
||||
package cliproxy
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/registry"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/sdk/config"
|
||||
)
|
||||
|
||||
func TestBackfillAntigravityModels_RegistersMissingAuth(t *testing.T) {
|
||||
source := &coreauth.Auth{
|
||||
ID: "ag-backfill-source",
|
||||
Provider: "antigravity",
|
||||
Status: coreauth.StatusActive,
|
||||
Attributes: map[string]string{
|
||||
"auth_kind": "oauth",
|
||||
},
|
||||
}
|
||||
target := &coreauth.Auth{
|
||||
ID: "ag-backfill-target",
|
||||
Provider: "antigravity",
|
||||
Status: coreauth.StatusActive,
|
||||
Attributes: map[string]string{
|
||||
"auth_kind": "oauth",
|
||||
},
|
||||
}
|
||||
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
if _, err := manager.Register(context.Background(), source); err != nil {
|
||||
t.Fatalf("register source auth: %v", err)
|
||||
}
|
||||
if _, err := manager.Register(context.Background(), target); err != nil {
|
||||
t.Fatalf("register target auth: %v", err)
|
||||
}
|
||||
|
||||
service := &Service{
|
||||
cfg: &config.Config{},
|
||||
coreManager: manager,
|
||||
}
|
||||
|
||||
reg := registry.GetGlobalRegistry()
|
||||
reg.UnregisterClient(source.ID)
|
||||
reg.UnregisterClient(target.ID)
|
||||
t.Cleanup(func() {
|
||||
reg.UnregisterClient(source.ID)
|
||||
reg.UnregisterClient(target.ID)
|
||||
})
|
||||
|
||||
primary := []*ModelInfo{
|
||||
{ID: "claude-sonnet-4-5"},
|
||||
{ID: "gemini-2.5-pro"},
|
||||
}
|
||||
reg.RegisterClient(source.ID, "antigravity", primary)
|
||||
|
||||
service.backfillAntigravityModels(source, primary)
|
||||
|
||||
got := reg.GetModelsForClient(target.ID)
|
||||
if len(got) != 2 {
|
||||
t.Fatalf("expected target auth to be backfilled with 2 models, got %d", len(got))
|
||||
}
|
||||
|
||||
ids := make(map[string]struct{}, len(got))
|
||||
for _, model := range got {
|
||||
if model == nil {
|
||||
continue
|
||||
}
|
||||
ids[strings.ToLower(strings.TrimSpace(model.ID))] = struct{}{}
|
||||
}
|
||||
if _, ok := ids["claude-sonnet-4-5"]; !ok {
|
||||
t.Fatal("expected backfilled model claude-sonnet-4-5")
|
||||
}
|
||||
if _, ok := ids["gemini-2.5-pro"]; !ok {
|
||||
t.Fatal("expected backfilled model gemini-2.5-pro")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBackfillAntigravityModels_RespectsExcludedModels(t *testing.T) {
|
||||
source := &coreauth.Auth{
|
||||
ID: "ag-backfill-source-excluded",
|
||||
Provider: "antigravity",
|
||||
Status: coreauth.StatusActive,
|
||||
Attributes: map[string]string{
|
||||
"auth_kind": "oauth",
|
||||
},
|
||||
}
|
||||
target := &coreauth.Auth{
|
||||
ID: "ag-backfill-target-excluded",
|
||||
Provider: "antigravity",
|
||||
Status: coreauth.StatusActive,
|
||||
Attributes: map[string]string{
|
||||
"auth_kind": "oauth",
|
||||
"excluded_models": "gemini-2.5-pro",
|
||||
},
|
||||
}
|
||||
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
if _, err := manager.Register(context.Background(), source); err != nil {
|
||||
t.Fatalf("register source auth: %v", err)
|
||||
}
|
||||
if _, err := manager.Register(context.Background(), target); err != nil {
|
||||
t.Fatalf("register target auth: %v", err)
|
||||
}
|
||||
|
||||
service := &Service{
|
||||
cfg: &config.Config{},
|
||||
coreManager: manager,
|
||||
}
|
||||
|
||||
reg := registry.GetGlobalRegistry()
|
||||
reg.UnregisterClient(source.ID)
|
||||
reg.UnregisterClient(target.ID)
|
||||
t.Cleanup(func() {
|
||||
reg.UnregisterClient(source.ID)
|
||||
reg.UnregisterClient(target.ID)
|
||||
})
|
||||
|
||||
primary := []*ModelInfo{
|
||||
{ID: "claude-sonnet-4-5"},
|
||||
{ID: "gemini-2.5-pro"},
|
||||
}
|
||||
reg.RegisterClient(source.ID, "antigravity", primary)
|
||||
|
||||
service.backfillAntigravityModels(source, primary)
|
||||
|
||||
got := reg.GetModelsForClient(target.ID)
|
||||
if len(got) != 1 {
|
||||
t.Fatalf("expected 1 model after exclusion, got %d", len(got))
|
||||
}
|
||||
if got[0] == nil || !strings.EqualFold(strings.TrimSpace(got[0].ID), "claude-sonnet-4-5") {
|
||||
t.Fatalf("expected remaining model %q, got %+v", "claude-sonnet-4-5", got[0])
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user