package auth import ( "context" "fmt" "net/http" "testing" "time" internalconfig "github.com/router-for-me/CLIProxyAPI/v6/internal/config" "github.com/router-for-me/CLIProxyAPI/v6/internal/registry" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor" ) type antigravityCreditsFallbackExecutor struct { streamCreditsRequested []bool } func (e *antigravityCreditsFallbackExecutor) Identifier() string { return "antigravity" } func (e *antigravityCreditsFallbackExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusNotImplemented, Message: "Execute not implemented"} } func (e *antigravityCreditsFallbackExecutor) ExecuteStream(ctx context.Context, _ *Auth, req cliproxyexecutor.Request, _ cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { creditsRequested := AntigravityCreditsRequested(ctx) e.streamCreditsRequested = append(e.streamCreditsRequested, creditsRequested) ch := make(chan cliproxyexecutor.StreamChunk, 1) if !creditsRequested { ch <- cliproxyexecutor.StreamChunk{Err: &Error{HTTPStatus: http.StatusTooManyRequests, Message: "quota exhausted"}} close(ch) return &cliproxyexecutor.StreamResult{Headers: http.Header{"X-Initial": {req.Model}}, Chunks: ch}, nil } ch <- cliproxyexecutor.StreamChunk{Payload: []byte("credits fallback")} close(ch) return &cliproxyexecutor.StreamResult{Headers: http.Header{"X-Credits": {req.Model}}, Chunks: ch}, nil } func (e *antigravityCreditsFallbackExecutor) Refresh(_ context.Context, auth *Auth) (*Auth, error) { return auth, nil } func (e *antigravityCreditsFallbackExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusNotImplemented, Message: "CountTokens not implemented"} } func (e *antigravityCreditsFallbackExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) { return nil, &Error{HTTPStatus: http.StatusNotImplemented, Message: "HttpRequest not implemented"} } func TestManagerExecuteStream_AntigravityCreditsFallbackAfterBootstrap429(t *testing.T) { const model = "claude-opus-4-6-thinking" executor := &antigravityCreditsFallbackExecutor{} manager := NewManager(nil, nil, nil) manager.SetConfig(&internalconfig.Config{ QuotaExceeded: internalconfig.QuotaExceeded{AntigravityCredits: true}, }) manager.RegisterExecutor(executor) registry.GetGlobalRegistry().RegisterClient("ag-credits", "antigravity", []*registry.ModelInfo{{ID: model}}) t.Cleanup(func() { registry.GetGlobalRegistry().UnregisterClient("ag-credits") }) if _, errRegister := manager.Register(context.Background(), &Auth{ID: "ag-credits", Provider: "antigravity"}); errRegister != nil { t.Fatalf("register auth: %v", errRegister) } streamResult, errExecute := manager.ExecuteStream(context.Background(), []string{"antigravity"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{}) if errExecute != nil { t.Fatalf("execute stream: %v", errExecute) } var payload []byte for chunk := range streamResult.Chunks { if chunk.Err != nil { t.Fatalf("unexpected stream error: %v", chunk.Err) } payload = append(payload, chunk.Payload...) } if string(payload) != "credits fallback" { t.Fatalf("payload = %q, want %q", string(payload), "credits fallback") } if got := streamResult.Headers.Get("X-Credits"); got != model { t.Fatalf("X-Credits header = %q, want routed model", got) } if len(executor.streamCreditsRequested) != 2 { t.Fatalf("stream calls = %d, want 2", len(executor.streamCreditsRequested)) } if executor.streamCreditsRequested[0] || !executor.streamCreditsRequested[1] { t.Fatalf("credits flags = %v, want [false true]", executor.streamCreditsRequested) } } func TestStatusCodeFromError_UnwrapsStreamBootstrap429(t *testing.T) { bootstrapErr := newStreamBootstrapError(&Error{HTTPStatus: http.StatusTooManyRequests, Message: "quota exhausted"}, nil) wrappedErr := fmt.Errorf("conductor stream failed: %w", bootstrapErr) if status := statusCodeFromError(wrappedErr); status != http.StatusTooManyRequests { t.Fatalf("statusCodeFromError() = %d, want %d", status, http.StatusTooManyRequests) } } func TestIsAuthBlockedForModel_ClaudeWithCreditsStillBlockedDuringCooldown(t *testing.T) { auth := &Auth{ ID: "ag-1", Provider: "antigravity", ModelStates: map[string]*ModelState{ "claude-sonnet-4-6": { Unavailable: true, NextRetryAfter: time.Now().Add(10 * time.Minute), Quota: QuotaState{ Exceeded: true, NextRecoverAt: time.Now().Add(10 * time.Minute), }, }, }, } SetAntigravityCreditsHint(auth.ID, AntigravityCreditsHint{ Known: true, Available: true, UpdatedAt: time.Now(), }) blocked, reason, _ := isAuthBlockedForModel(auth, "claude-sonnet-4-6", time.Now()) if !blocked || reason != blockReasonCooldown { t.Fatalf("expected auth to be blocked during cooldown even with credits, got blocked=%v reason=%v", blocked, reason) } } func TestIsAuthBlockedForModel_KeepsGeminiBlockedWithoutCreditsBypass(t *testing.T) { auth := &Auth{ ID: "ag-2", Provider: "antigravity", ModelStates: map[string]*ModelState{ "gemini-3-flash": { Unavailable: true, NextRetryAfter: time.Now().Add(10 * time.Minute), Quota: QuotaState{ Exceeded: true, NextRecoverAt: time.Now().Add(10 * time.Minute), }, }, }, } SetAntigravityCreditsHint(auth.ID, AntigravityCreditsHint{ Known: true, Available: true, UpdatedAt: time.Now(), }) blocked, reason, _ := isAuthBlockedForModel(auth, "gemini-3-flash", time.Now()) if !blocked || reason != blockReasonCooldown { t.Fatalf("expected gemini auth to remain blocked, got blocked=%v reason=%v", blocked, reason) } }