sususu98
239 lines
9.1 KiB
Go
239 lines
9.1 KiB
Go
package auth
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"net/http"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
|
"github.com/router-for-me/CLIProxyAPI/v7/internal/registry"
|
|
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
|
|
log "github.com/sirupsen/logrus"
|
|
)
|
|
|
|
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"}
|
|
}
|
|
|
|
type codexOnlyFailureExecutor struct{}
|
|
|
|
func (codexOnlyFailureExecutor) Identifier() string { return "codex" }
|
|
|
|
func (codexOnlyFailureExecutor) Execute(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
|
return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusTooManyRequests, Message: "codex quota exhausted"}
|
|
}
|
|
|
|
func (codexOnlyFailureExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) {
|
|
return nil, &Error{HTTPStatus: http.StatusTooManyRequests, Message: "codex quota exhausted"}
|
|
}
|
|
|
|
func (codexOnlyFailureExecutor) Refresh(_ context.Context, auth *Auth) (*Auth, error) {
|
|
return auth, nil
|
|
}
|
|
|
|
func (codexOnlyFailureExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
|
return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusTooManyRequests, Message: "codex quota exhausted"}
|
|
}
|
|
|
|
func (codexOnlyFailureExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) {
|
|
return nil, &Error{HTTPStatus: http.StatusTooManyRequests, Message: "codex quota exhausted"}
|
|
}
|
|
|
|
type captureLogHook struct {
|
|
messages []string
|
|
}
|
|
|
|
func (h *captureLogHook) Levels() []log.Level {
|
|
return log.AllLevels
|
|
}
|
|
|
|
func (h *captureLogHook) Fire(entry *log.Entry) error {
|
|
h.messages = append(h.messages, entry.Message)
|
|
return nil
|
|
}
|
|
|
|
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 TestManagerExecuteStream_CodexOnlyDoesNotEnterAntigravityCreditsFallback(t *testing.T) {
|
|
const model = "gpt-5.5"
|
|
logger := log.StandardLogger()
|
|
oldLevel := logger.GetLevel()
|
|
oldHooks := logger.ReplaceHooks(make(log.LevelHooks))
|
|
hook := &captureLogHook{}
|
|
logger.SetLevel(log.DebugLevel)
|
|
logger.AddHook(hook)
|
|
t.Cleanup(func() {
|
|
logger.SetLevel(oldLevel)
|
|
logger.ReplaceHooks(oldHooks)
|
|
})
|
|
|
|
manager := NewManager(nil, nil, nil)
|
|
manager.SetConfig(&internalconfig.Config{
|
|
QuotaExceeded: internalconfig.QuotaExceeded{AntigravityCredits: true},
|
|
})
|
|
manager.RegisterExecutor(codexOnlyFailureExecutor{})
|
|
manager.RegisterExecutor(&antigravityCreditsFallbackExecutor{})
|
|
reg := registry.GetGlobalRegistry()
|
|
reg.RegisterClient("codex-only", "codex", []*registry.ModelInfo{{ID: model}})
|
|
reg.RegisterClient("ag-unrelated", "antigravity", []*registry.ModelInfo{{ID: "gemini-3-flash"}})
|
|
t.Cleanup(func() {
|
|
reg.UnregisterClient("codex-only")
|
|
reg.UnregisterClient("ag-unrelated")
|
|
})
|
|
if _, errRegister := manager.Register(context.Background(), &Auth{ID: "codex-only", Provider: "codex"}); errRegister != nil {
|
|
t.Fatalf("register codex auth: %v", errRegister)
|
|
}
|
|
if _, errRegister := manager.Register(context.Background(), &Auth{ID: "ag-unrelated", Provider: "antigravity"}); errRegister != nil {
|
|
t.Fatalf("register antigravity auth: %v", errRegister)
|
|
}
|
|
|
|
_, errExecute := manager.ExecuteStream(context.Background(), []string{"codex"}, cliproxyexecutor.Request{Model: model}, cliproxyexecutor.Options{})
|
|
if errExecute == nil {
|
|
t.Fatal("expected codex execution failure")
|
|
}
|
|
|
|
for _, message := range hook.messages {
|
|
if strings.Contains(message, "shouldAttemptAntigravityCreditsFallback") {
|
|
t.Fatalf("codex-only request entered antigravity credits fallback gate; messages=%v", hook.messages)
|
|
}
|
|
}
|
|
}
|
|
|
|
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)
|
|
}
|
|
}
|