mirror of
https://github.com/router-for-me/CLIProxyAPIPlus.git
synced 2026-03-09 15:25:17 +00:00
fix(antigravity): keep primary model list and backfill empty auths
This commit is contained in:
@@ -54,8 +54,58 @@ 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
|
||||
}
|
||||
clone := *model
|
||||
out = append(out, &clone)
|
||||
}
|
||||
if len(out) == 0 {
|
||||
return nil
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
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
|
||||
@@ -1007,7 +1057,7 @@ func FetchAntigravityModels(ctx context.Context, auth *cliproxyauth.Auth, cfg *c
|
||||
exec := &AntigravityExecutor{cfg: cfg}
|
||||
token, updatedAuth, errToken := exec.ensureAccessToken(ctx, auth)
|
||||
if errToken != nil || token == "" {
|
||||
return nil
|
||||
return fallbackAntigravityPrimaryModels()
|
||||
}
|
||||
if updatedAuth != nil {
|
||||
auth = updatedAuth
|
||||
@@ -1020,7 +1070,7 @@ func FetchAntigravityModels(ctx context.Context, auth *cliproxyauth.Auth, cfg *c
|
||||
modelsURL := baseURL + antigravityModelsPath
|
||||
httpReq, errReq := http.NewRequestWithContext(ctx, http.MethodPost, modelsURL, bytes.NewReader([]byte(`{}`)))
|
||||
if errReq != nil {
|
||||
return nil
|
||||
return fallbackAntigravityPrimaryModels()
|
||||
}
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
httpReq.Header.Set("Authorization", "Bearer "+token)
|
||||
@@ -1032,13 +1082,13 @@ func FetchAntigravityModels(ctx context.Context, auth *cliproxyauth.Auth, cfg *c
|
||||
httpResp, errDo := httpClient.Do(httpReq)
|
||||
if errDo != nil {
|
||||
if errors.Is(errDo, context.Canceled) || errors.Is(errDo, context.DeadlineExceeded) {
|
||||
return nil
|
||||
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 nil
|
||||
return fallbackAntigravityPrimaryModels()
|
||||
}
|
||||
|
||||
bodyBytes, errRead := io.ReadAll(httpResp.Body)
|
||||
@@ -1050,19 +1100,27 @@ func FetchAntigravityModels(ctx context.Context, auth *cliproxyauth.Auth, cfg *c
|
||||
log.Debugf("antigravity executor: models read error on base url %s, retrying with fallback base url: %s", baseURL, baseURLs[idx+1])
|
||||
continue
|
||||
}
|
||||
return nil
|
||||
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
|
||||
}
|
||||
return nil
|
||||
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() {
|
||||
return nil
|
||||
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()
|
||||
@@ -1107,9 +1165,18 @@ func FetchAntigravityModels(ctx context.Context, auth *cliproxyauth.Auth, cfg *c
|
||||
}
|
||||
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 nil
|
||||
return fallbackAntigravityPrimaryModels()
|
||||
}
|
||||
|
||||
func (e *AntigravityExecutor) ensureAccessToken(ctx context.Context, auth *cliproxyauth.Auth) (string, *cliproxyauth.Auth, error) {
|
||||
|
||||
@@ -0,0 +1,64 @@
|
||||
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"}}); !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"
|
||||
|
||||
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)
|
||||
}
|
||||
}
|
||||
@@ -925,6 +925,9 @@ 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
|
||||
}
|
||||
|
||||
@@ -1069,6 +1072,56 @@ 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
|
||||
|
||||
135
sdk/cliproxy/service_antigravity_backfill_test.go
Normal file
135
sdk/cliproxy/service_antigravity_backfill_test.go
Normal file
@@ -0,0 +1,135 @@
|
||||
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