Merge pull request #2426 from xixiwenxuanhe/feature/antigravity-credits
feat(antigravity): add AI credits quota fallback
This commit is contained in:
@@ -91,6 +91,7 @@ max-retry-interval: 30
|
|||||||
quota-exceeded:
|
quota-exceeded:
|
||||||
switch-project: true # Whether to automatically switch to another project when a quota is exceeded
|
switch-project: true # Whether to automatically switch to another project when a quota is exceeded
|
||||||
switch-preview-model: true # Whether to automatically switch to a preview model when a quota is exceeded
|
switch-preview-model: true # Whether to automatically switch to a preview model when a quota is exceeded
|
||||||
|
antigravity-credits: true # Whether to retry Antigravity quota_exhausted 429s once with enabledCreditTypes=["GOOGLE_ONE_AI"]
|
||||||
|
|
||||||
# Routing strategy for selecting credentials when multiple match.
|
# Routing strategy for selecting credentials when multiple match.
|
||||||
routing:
|
routing:
|
||||||
|
|||||||
@@ -194,6 +194,10 @@ type QuotaExceeded struct {
|
|||||||
|
|
||||||
// SwitchPreviewModel indicates whether to automatically switch to a preview model when a quota is exceeded.
|
// SwitchPreviewModel indicates whether to automatically switch to a preview model when a quota is exceeded.
|
||||||
SwitchPreviewModel bool `yaml:"switch-preview-model" json:"switch-preview-model"`
|
SwitchPreviewModel bool `yaml:"switch-preview-model" json:"switch-preview-model"`
|
||||||
|
|
||||||
|
// AntigravityCredits indicates whether to retry Antigravity quota_exhausted 429s once
|
||||||
|
// on the same credential with enabledCreditTypes=["GOOGLE_ONE_AI"].
|
||||||
|
AntigravityCredits bool `yaml:"antigravity-credits" json:"antigravity-credits"`
|
||||||
}
|
}
|
||||||
|
|
||||||
// RoutingConfig configures how credentials are selected for requests.
|
// RoutingConfig configures how credentials are selected for requests.
|
||||||
|
|||||||
@@ -47,12 +47,41 @@ const (
|
|||||||
defaultAntigravityAgent = "antigravity/1.19.6 darwin/arm64"
|
defaultAntigravityAgent = "antigravity/1.19.6 darwin/arm64"
|
||||||
antigravityAuthType = "antigravity"
|
antigravityAuthType = "antigravity"
|
||||||
refreshSkew = 3000 * time.Second
|
refreshSkew = 3000 * time.Second
|
||||||
|
antigravityCreditsRetryTTL = 5 * time.Hour
|
||||||
// systemInstruction = "You are Antigravity, a powerful agentic AI coding assistant designed by the Google Deepmind team working on Advanced Agentic Coding.You are pair programming with a USER to solve their coding task. The task may require creating a new codebase, modifying or debugging an existing codebase, or simply answering a question.**Absolute paths only****Proactiveness**"
|
// systemInstruction = "You are Antigravity, a powerful agentic AI coding assistant designed by the Google Deepmind team working on Advanced Agentic Coding.You are pair programming with a USER to solve their coding task. The task may require creating a new codebase, modifying or debugging an existing codebase, or simply answering a question.**Absolute paths only****Proactiveness**"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
type antigravity429Category string
|
||||||
|
|
||||||
|
const (
|
||||||
|
antigravity429Unknown antigravity429Category = "unknown"
|
||||||
|
antigravity429RateLimited antigravity429Category = "rate_limited"
|
||||||
|
antigravity429QuotaExhausted antigravity429Category = "quota_exhausted"
|
||||||
|
)
|
||||||
|
|
||||||
var (
|
var (
|
||||||
randSource = rand.New(rand.NewSource(time.Now().UnixNano()))
|
randSource = rand.New(rand.NewSource(time.Now().UnixNano()))
|
||||||
randSourceMutex sync.Mutex
|
randSourceMutex sync.Mutex
|
||||||
|
antigravityCreditsExhaustedByAuth sync.Map
|
||||||
|
antigravityPreferCreditsByModel sync.Map
|
||||||
|
antigravityQuotaExhaustedKeywords = []string{
|
||||||
|
"quota_exhausted",
|
||||||
|
"quota exhausted",
|
||||||
|
}
|
||||||
|
antigravityCreditsExhaustedKeywords = []string{
|
||||||
|
"google_one_ai",
|
||||||
|
"insufficient credit",
|
||||||
|
"insufficient credits",
|
||||||
|
"not enough credit",
|
||||||
|
"not enough credits",
|
||||||
|
"credit exhausted",
|
||||||
|
"credits exhausted",
|
||||||
|
"credit balance",
|
||||||
|
"minimumcreditamountforusage",
|
||||||
|
"minimum credit amount for usage",
|
||||||
|
"minimum credit",
|
||||||
|
"resource has been exhausted",
|
||||||
|
}
|
||||||
)
|
)
|
||||||
|
|
||||||
// AntigravityExecutor proxies requests to the antigravity upstream.
|
// AntigravityExecutor proxies requests to the antigravity upstream.
|
||||||
@@ -183,6 +212,231 @@ func (e *AntigravityExecutor) HttpRequest(ctx context.Context, auth *cliproxyaut
|
|||||||
return httpClient.Do(httpReq)
|
return httpClient.Do(httpReq)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func injectEnabledCreditTypes(payload []byte) []byte {
|
||||||
|
if len(payload) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if !gjson.ValidBytes(payload) {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
updated, err := sjson.SetRawBytes(payload, "enabledCreditTypes", []byte(`["GOOGLE_ONE_AI"]`))
|
||||||
|
if err != nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
return updated
|
||||||
|
}
|
||||||
|
|
||||||
|
func classifyAntigravity429(body []byte) antigravity429Category {
|
||||||
|
if len(body) == 0 {
|
||||||
|
return antigravity429Unknown
|
||||||
|
}
|
||||||
|
lowerBody := strings.ToLower(string(body))
|
||||||
|
for _, keyword := range antigravityQuotaExhaustedKeywords {
|
||||||
|
if strings.Contains(lowerBody, keyword) {
|
||||||
|
return antigravity429QuotaExhausted
|
||||||
|
}
|
||||||
|
}
|
||||||
|
status := strings.TrimSpace(gjson.GetBytes(body, "error.status").String())
|
||||||
|
if !strings.EqualFold(status, "RESOURCE_EXHAUSTED") {
|
||||||
|
return antigravity429Unknown
|
||||||
|
}
|
||||||
|
details := gjson.GetBytes(body, "error.details")
|
||||||
|
if !details.Exists() || !details.IsArray() {
|
||||||
|
return antigravity429Unknown
|
||||||
|
}
|
||||||
|
for _, detail := range details.Array() {
|
||||||
|
if detail.Get("@type").String() != "type.googleapis.com/google.rpc.ErrorInfo" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
reason := strings.TrimSpace(detail.Get("reason").String())
|
||||||
|
if strings.EqualFold(reason, "QUOTA_EXHAUSTED") {
|
||||||
|
return antigravity429QuotaExhausted
|
||||||
|
}
|
||||||
|
if strings.EqualFold(reason, "RATE_LIMIT_EXCEEDED") {
|
||||||
|
return antigravity429RateLimited
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return antigravity429Unknown
|
||||||
|
}
|
||||||
|
|
||||||
|
func antigravityCreditsRetryEnabled(cfg *config.Config) bool {
|
||||||
|
return cfg != nil && cfg.QuotaExceeded.AntigravityCredits
|
||||||
|
}
|
||||||
|
|
||||||
|
func antigravityCreditsExhausted(auth *cliproxyauth.Auth, now time.Time) bool {
|
||||||
|
if auth == nil || strings.TrimSpace(auth.ID) == "" {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
value, ok := antigravityCreditsExhaustedByAuth.Load(auth.ID)
|
||||||
|
if !ok {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
until, ok := value.(time.Time)
|
||||||
|
if !ok || until.IsZero() {
|
||||||
|
antigravityCreditsExhaustedByAuth.Delete(auth.ID)
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
if !until.After(now) {
|
||||||
|
antigravityCreditsExhaustedByAuth.Delete(auth.ID)
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
|
||||||
|
func markAntigravityCreditsExhausted(auth *cliproxyauth.Auth, now time.Time) {
|
||||||
|
if auth == nil || strings.TrimSpace(auth.ID) == "" {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
antigravityCreditsExhaustedByAuth.Store(auth.ID, now.Add(antigravityCreditsRetryTTL))
|
||||||
|
}
|
||||||
|
|
||||||
|
func clearAntigravityCreditsExhausted(auth *cliproxyauth.Auth) {
|
||||||
|
if auth == nil || strings.TrimSpace(auth.ID) == "" {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
antigravityCreditsExhaustedByAuth.Delete(auth.ID)
|
||||||
|
}
|
||||||
|
|
||||||
|
func antigravityPreferCreditsKey(auth *cliproxyauth.Auth, modelName string) string {
|
||||||
|
if auth == nil {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
authID := strings.TrimSpace(auth.ID)
|
||||||
|
modelName = strings.TrimSpace(modelName)
|
||||||
|
if authID == "" || modelName == "" {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
return authID + "|" + modelName
|
||||||
|
}
|
||||||
|
|
||||||
|
func antigravityShouldPreferCredits(auth *cliproxyauth.Auth, modelName string, now time.Time) bool {
|
||||||
|
key := antigravityPreferCreditsKey(auth, modelName)
|
||||||
|
if key == "" {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
value, ok := antigravityPreferCreditsByModel.Load(key)
|
||||||
|
if !ok {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
until, ok := value.(time.Time)
|
||||||
|
if !ok || until.IsZero() {
|
||||||
|
antigravityPreferCreditsByModel.Delete(key)
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
if !until.After(now) {
|
||||||
|
antigravityPreferCreditsByModel.Delete(key)
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
|
||||||
|
func markAntigravityPreferCredits(auth *cliproxyauth.Auth, modelName string, now time.Time, retryAfter *time.Duration) {
|
||||||
|
key := antigravityPreferCreditsKey(auth, modelName)
|
||||||
|
if key == "" {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
until := now.Add(antigravityCreditsRetryTTL)
|
||||||
|
if retryAfter != nil && *retryAfter > 0 {
|
||||||
|
until = now.Add(*retryAfter)
|
||||||
|
}
|
||||||
|
antigravityPreferCreditsByModel.Store(key, until)
|
||||||
|
}
|
||||||
|
|
||||||
|
func clearAntigravityPreferCredits(auth *cliproxyauth.Auth, modelName string) {
|
||||||
|
key := antigravityPreferCreditsKey(auth, modelName)
|
||||||
|
if key == "" {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
antigravityPreferCreditsByModel.Delete(key)
|
||||||
|
}
|
||||||
|
|
||||||
|
func shouldMarkAntigravityCreditsExhausted(statusCode int, body []byte, reqErr error) bool {
|
||||||
|
if reqErr != nil || statusCode == 0 {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
if statusCode >= http.StatusInternalServerError || statusCode == http.StatusRequestTimeout {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
lowerBody := strings.ToLower(string(body))
|
||||||
|
for _, keyword := range antigravityCreditsExhaustedKeywords {
|
||||||
|
if strings.Contains(lowerBody, keyword) {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
|
func newAntigravityStatusErr(statusCode int, body []byte) statusErr {
|
||||||
|
err := statusErr{code: statusCode, msg: string(body)}
|
||||||
|
if statusCode == http.StatusTooManyRequests {
|
||||||
|
if retryAfter, parseErr := parseRetryDelay(body); parseErr == nil && retryAfter != nil {
|
||||||
|
err.retryAfter = retryAfter
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
func (e *AntigravityExecutor) attemptCreditsFallback(
|
||||||
|
ctx context.Context,
|
||||||
|
auth *cliproxyauth.Auth,
|
||||||
|
httpClient *http.Client,
|
||||||
|
token string,
|
||||||
|
modelName string,
|
||||||
|
payload []byte,
|
||||||
|
stream bool,
|
||||||
|
alt string,
|
||||||
|
baseURL string,
|
||||||
|
originalBody []byte,
|
||||||
|
) (*http.Response, bool) {
|
||||||
|
if !antigravityCreditsRetryEnabled(e.cfg) {
|
||||||
|
return nil, false
|
||||||
|
}
|
||||||
|
if classifyAntigravity429(originalBody) != antigravity429QuotaExhausted {
|
||||||
|
return nil, false
|
||||||
|
}
|
||||||
|
now := time.Now()
|
||||||
|
if antigravityCreditsExhausted(auth, now) {
|
||||||
|
return nil, false
|
||||||
|
}
|
||||||
|
creditsPayload := injectEnabledCreditTypes(payload)
|
||||||
|
if len(creditsPayload) == 0 {
|
||||||
|
return nil, false
|
||||||
|
}
|
||||||
|
|
||||||
|
httpReq, errReq := e.buildRequest(ctx, auth, token, modelName, creditsPayload, stream, alt, baseURL)
|
||||||
|
if errReq != nil {
|
||||||
|
recordAPIResponseError(ctx, e.cfg, errReq)
|
||||||
|
return nil, true
|
||||||
|
}
|
||||||
|
httpResp, errDo := httpClient.Do(httpReq)
|
||||||
|
if errDo != nil {
|
||||||
|
recordAPIResponseError(ctx, e.cfg, errDo)
|
||||||
|
return nil, true
|
||||||
|
}
|
||||||
|
if httpResp.StatusCode >= http.StatusOK && httpResp.StatusCode < http.StatusMultipleChoices {
|
||||||
|
retryAfter, _ := parseRetryDelay(originalBody)
|
||||||
|
markAntigravityPreferCredits(auth, modelName, now, retryAfter)
|
||||||
|
clearAntigravityCreditsExhausted(auth)
|
||||||
|
return httpResp, true
|
||||||
|
}
|
||||||
|
|
||||||
|
recordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone())
|
||||||
|
bodyBytes, errRead := io.ReadAll(httpResp.Body)
|
||||||
|
if errClose := httpResp.Body.Close(); errClose != nil {
|
||||||
|
log.Errorf("antigravity executor: close credits fallback response body error: %v", errClose)
|
||||||
|
}
|
||||||
|
if errRead != nil {
|
||||||
|
recordAPIResponseError(ctx, e.cfg, errRead)
|
||||||
|
return nil, true
|
||||||
|
}
|
||||||
|
appendAPIResponseChunk(ctx, e.cfg, bodyBytes)
|
||||||
|
if shouldMarkAntigravityCreditsExhausted(httpResp.StatusCode, bodyBytes, nil) {
|
||||||
|
clearAntigravityPreferCredits(auth, modelName)
|
||||||
|
markAntigravityCreditsExhausted(auth, now)
|
||||||
|
}
|
||||||
|
return nil, true
|
||||||
|
}
|
||||||
|
|
||||||
// Execute performs a non-streaming request to the Antigravity API.
|
// Execute performs a non-streaming request to the Antigravity API.
|
||||||
func (e *AntigravityExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) {
|
func (e *AntigravityExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) {
|
||||||
if opts.Alt == "responses/compact" {
|
if opts.Alt == "responses/compact" {
|
||||||
@@ -237,7 +491,15 @@ attemptLoop:
|
|||||||
var lastErr error
|
var lastErr error
|
||||||
|
|
||||||
for idx, baseURL := range baseURLs {
|
for idx, baseURL := range baseURLs {
|
||||||
httpReq, errReq := e.buildRequest(ctx, auth, token, baseModel, translated, false, opts.Alt, baseURL)
|
requestPayload := translated
|
||||||
|
usedCreditsDirect := false
|
||||||
|
if antigravityCreditsRetryEnabled(e.cfg) && antigravityShouldPreferCredits(auth, baseModel, time.Now()) {
|
||||||
|
if creditsPayload := injectEnabledCreditTypes(translated); len(creditsPayload) > 0 {
|
||||||
|
requestPayload = creditsPayload
|
||||||
|
usedCreditsDirect = true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
httpReq, errReq := e.buildRequest(ctx, auth, token, baseModel, requestPayload, false, opts.Alt, baseURL)
|
||||||
if errReq != nil {
|
if errReq != nil {
|
||||||
err = errReq
|
err = errReq
|
||||||
return resp, err
|
return resp, err
|
||||||
@@ -272,6 +534,36 @@ attemptLoop:
|
|||||||
}
|
}
|
||||||
appendAPIResponseChunk(ctx, e.cfg, bodyBytes)
|
appendAPIResponseChunk(ctx, e.cfg, bodyBytes)
|
||||||
|
|
||||||
|
if httpResp.StatusCode == http.StatusTooManyRequests {
|
||||||
|
if usedCreditsDirect {
|
||||||
|
if shouldMarkAntigravityCreditsExhausted(httpResp.StatusCode, bodyBytes, nil) {
|
||||||
|
clearAntigravityPreferCredits(auth, baseModel)
|
||||||
|
markAntigravityCreditsExhausted(auth, time.Now())
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
creditsResp, _ := e.attemptCreditsFallback(ctx, auth, httpClient, token, baseModel, translated, false, opts.Alt, baseURL, bodyBytes)
|
||||||
|
if creditsResp != nil {
|
||||||
|
recordAPIResponseMetadata(ctx, e.cfg, creditsResp.StatusCode, creditsResp.Header.Clone())
|
||||||
|
creditsBody, errCreditsRead := io.ReadAll(creditsResp.Body)
|
||||||
|
if errClose := creditsResp.Body.Close(); errClose != nil {
|
||||||
|
log.Errorf("antigravity executor: close credits success response body error: %v", errClose)
|
||||||
|
}
|
||||||
|
if errCreditsRead != nil {
|
||||||
|
recordAPIResponseError(ctx, e.cfg, errCreditsRead)
|
||||||
|
err = errCreditsRead
|
||||||
|
return resp, err
|
||||||
|
}
|
||||||
|
appendAPIResponseChunk(ctx, e.cfg, creditsBody)
|
||||||
|
reporter.publish(ctx, parseAntigravityUsage(creditsBody))
|
||||||
|
var param any
|
||||||
|
converted := sdktranslator.TranslateNonStream(ctx, to, from, req.Model, opts.OriginalRequest, translated, creditsBody, ¶m)
|
||||||
|
resp = cliproxyexecutor.Response{Payload: converted, Headers: creditsResp.Header.Clone()}
|
||||||
|
reporter.ensurePublished(ctx)
|
||||||
|
return resp, nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
if httpResp.StatusCode < http.StatusOK || httpResp.StatusCode >= http.StatusMultipleChoices {
|
if httpResp.StatusCode < http.StatusOK || httpResp.StatusCode >= http.StatusMultipleChoices {
|
||||||
log.Debugf("antigravity executor: upstream error status: %d, body: %s", httpResp.StatusCode, summarizeErrorBody(httpResp.Header.Get("Content-Type"), bodyBytes))
|
log.Debugf("antigravity executor: upstream error status: %d, body: %s", httpResp.StatusCode, summarizeErrorBody(httpResp.Header.Get("Content-Type"), bodyBytes))
|
||||||
lastStatus = httpResp.StatusCode
|
lastStatus = httpResp.StatusCode
|
||||||
@@ -295,13 +587,7 @@ attemptLoop:
|
|||||||
continue attemptLoop
|
continue attemptLoop
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
sErr := statusErr{code: httpResp.StatusCode, msg: string(bodyBytes)}
|
err = newAntigravityStatusErr(httpResp.StatusCode, bodyBytes)
|
||||||
if httpResp.StatusCode == http.StatusTooManyRequests {
|
|
||||||
if retryAfter, parseErr := parseRetryDelay(bodyBytes); parseErr == nil && retryAfter != nil {
|
|
||||||
sErr.retryAfter = retryAfter
|
|
||||||
}
|
|
||||||
}
|
|
||||||
err = sErr
|
|
||||||
return resp, err
|
return resp, err
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -315,13 +601,7 @@ attemptLoop:
|
|||||||
|
|
||||||
switch {
|
switch {
|
||||||
case lastStatus != 0:
|
case lastStatus != 0:
|
||||||
sErr := statusErr{code: lastStatus, msg: string(lastBody)}
|
err = newAntigravityStatusErr(lastStatus, lastBody)
|
||||||
if lastStatus == http.StatusTooManyRequests {
|
|
||||||
if retryAfter, parseErr := parseRetryDelay(lastBody); parseErr == nil && retryAfter != nil {
|
|
||||||
sErr.retryAfter = retryAfter
|
|
||||||
}
|
|
||||||
}
|
|
||||||
err = sErr
|
|
||||||
case lastErr != nil:
|
case lastErr != nil:
|
||||||
err = lastErr
|
err = lastErr
|
||||||
default:
|
default:
|
||||||
@@ -379,7 +659,15 @@ attemptLoop:
|
|||||||
var lastErr error
|
var lastErr error
|
||||||
|
|
||||||
for idx, baseURL := range baseURLs {
|
for idx, baseURL := range baseURLs {
|
||||||
httpReq, errReq := e.buildRequest(ctx, auth, token, baseModel, translated, true, opts.Alt, baseURL)
|
requestPayload := translated
|
||||||
|
usedCreditsDirect := false
|
||||||
|
if antigravityCreditsRetryEnabled(e.cfg) && antigravityShouldPreferCredits(auth, baseModel, time.Now()) {
|
||||||
|
if creditsPayload := injectEnabledCreditTypes(translated); len(creditsPayload) > 0 {
|
||||||
|
requestPayload = creditsPayload
|
||||||
|
usedCreditsDirect = true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
httpReq, errReq := e.buildRequest(ctx, auth, token, baseModel, requestPayload, true, opts.Alt, baseURL)
|
||||||
if errReq != nil {
|
if errReq != nil {
|
||||||
err = errReq
|
err = errReq
|
||||||
return resp, err
|
return resp, err
|
||||||
@@ -428,6 +716,23 @@ attemptLoop:
|
|||||||
return resp, err
|
return resp, err
|
||||||
}
|
}
|
||||||
appendAPIResponseChunk(ctx, e.cfg, bodyBytes)
|
appendAPIResponseChunk(ctx, e.cfg, bodyBytes)
|
||||||
|
if httpResp.StatusCode == http.StatusTooManyRequests {
|
||||||
|
if usedCreditsDirect {
|
||||||
|
if shouldMarkAntigravityCreditsExhausted(httpResp.StatusCode, bodyBytes, nil) {
|
||||||
|
clearAntigravityPreferCredits(auth, baseModel)
|
||||||
|
markAntigravityCreditsExhausted(auth, time.Now())
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
creditsResp, _ := e.attemptCreditsFallback(ctx, auth, httpClient, token, baseModel, translated, true, opts.Alt, baseURL, bodyBytes)
|
||||||
|
if creditsResp != nil {
|
||||||
|
httpResp = creditsResp
|
||||||
|
recordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if httpResp.StatusCode >= http.StatusOK && httpResp.StatusCode < http.StatusMultipleChoices {
|
||||||
|
goto streamSuccessClaudeNonStream
|
||||||
|
}
|
||||||
lastStatus = httpResp.StatusCode
|
lastStatus = httpResp.StatusCode
|
||||||
lastBody = append([]byte(nil), bodyBytes...)
|
lastBody = append([]byte(nil), bodyBytes...)
|
||||||
lastErr = nil
|
lastErr = nil
|
||||||
@@ -449,16 +754,11 @@ attemptLoop:
|
|||||||
continue attemptLoop
|
continue attemptLoop
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
sErr := statusErr{code: httpResp.StatusCode, msg: string(bodyBytes)}
|
err = newAntigravityStatusErr(httpResp.StatusCode, bodyBytes)
|
||||||
if httpResp.StatusCode == http.StatusTooManyRequests {
|
|
||||||
if retryAfter, parseErr := parseRetryDelay(bodyBytes); parseErr == nil && retryAfter != nil {
|
|
||||||
sErr.retryAfter = retryAfter
|
|
||||||
}
|
|
||||||
}
|
|
||||||
err = sErr
|
|
||||||
return resp, err
|
return resp, err
|
||||||
}
|
}
|
||||||
|
|
||||||
|
streamSuccessClaudeNonStream:
|
||||||
out := make(chan cliproxyexecutor.StreamChunk)
|
out := make(chan cliproxyexecutor.StreamChunk)
|
||||||
go func(resp *http.Response) {
|
go func(resp *http.Response) {
|
||||||
defer close(out)
|
defer close(out)
|
||||||
@@ -520,13 +820,7 @@ attemptLoop:
|
|||||||
|
|
||||||
switch {
|
switch {
|
||||||
case lastStatus != 0:
|
case lastStatus != 0:
|
||||||
sErr := statusErr{code: lastStatus, msg: string(lastBody)}
|
err = newAntigravityStatusErr(lastStatus, lastBody)
|
||||||
if lastStatus == http.StatusTooManyRequests {
|
|
||||||
if retryAfter, parseErr := parseRetryDelay(lastBody); parseErr == nil && retryAfter != nil {
|
|
||||||
sErr.retryAfter = retryAfter
|
|
||||||
}
|
|
||||||
}
|
|
||||||
err = sErr
|
|
||||||
case lastErr != nil:
|
case lastErr != nil:
|
||||||
err = lastErr
|
err = lastErr
|
||||||
default:
|
default:
|
||||||
@@ -782,7 +1076,15 @@ attemptLoop:
|
|||||||
var lastErr error
|
var lastErr error
|
||||||
|
|
||||||
for idx, baseURL := range baseURLs {
|
for idx, baseURL := range baseURLs {
|
||||||
httpReq, errReq := e.buildRequest(ctx, auth, token, baseModel, translated, true, opts.Alt, baseURL)
|
requestPayload := translated
|
||||||
|
usedCreditsDirect := false
|
||||||
|
if antigravityCreditsRetryEnabled(e.cfg) && antigravityShouldPreferCredits(auth, baseModel, time.Now()) {
|
||||||
|
if creditsPayload := injectEnabledCreditTypes(translated); len(creditsPayload) > 0 {
|
||||||
|
requestPayload = creditsPayload
|
||||||
|
usedCreditsDirect = true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
httpReq, errReq := e.buildRequest(ctx, auth, token, baseModel, requestPayload, true, opts.Alt, baseURL)
|
||||||
if errReq != nil {
|
if errReq != nil {
|
||||||
err = errReq
|
err = errReq
|
||||||
return nil, err
|
return nil, err
|
||||||
@@ -830,6 +1132,23 @@ attemptLoop:
|
|||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
appendAPIResponseChunk(ctx, e.cfg, bodyBytes)
|
appendAPIResponseChunk(ctx, e.cfg, bodyBytes)
|
||||||
|
if httpResp.StatusCode == http.StatusTooManyRequests {
|
||||||
|
if usedCreditsDirect {
|
||||||
|
if shouldMarkAntigravityCreditsExhausted(httpResp.StatusCode, bodyBytes, nil) {
|
||||||
|
clearAntigravityPreferCredits(auth, baseModel)
|
||||||
|
markAntigravityCreditsExhausted(auth, time.Now())
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
creditsResp, _ := e.attemptCreditsFallback(ctx, auth, httpClient, token, baseModel, translated, true, opts.Alt, baseURL, bodyBytes)
|
||||||
|
if creditsResp != nil {
|
||||||
|
httpResp = creditsResp
|
||||||
|
recordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if httpResp.StatusCode >= http.StatusOK && httpResp.StatusCode < http.StatusMultipleChoices {
|
||||||
|
goto streamSuccessExecuteStream
|
||||||
|
}
|
||||||
lastStatus = httpResp.StatusCode
|
lastStatus = httpResp.StatusCode
|
||||||
lastBody = append([]byte(nil), bodyBytes...)
|
lastBody = append([]byte(nil), bodyBytes...)
|
||||||
lastErr = nil
|
lastErr = nil
|
||||||
@@ -851,16 +1170,11 @@ attemptLoop:
|
|||||||
continue attemptLoop
|
continue attemptLoop
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
sErr := statusErr{code: httpResp.StatusCode, msg: string(bodyBytes)}
|
err = newAntigravityStatusErr(httpResp.StatusCode, bodyBytes)
|
||||||
if httpResp.StatusCode == http.StatusTooManyRequests {
|
|
||||||
if retryAfter, parseErr := parseRetryDelay(bodyBytes); parseErr == nil && retryAfter != nil {
|
|
||||||
sErr.retryAfter = retryAfter
|
|
||||||
}
|
|
||||||
}
|
|
||||||
err = sErr
|
|
||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
|
|
||||||
|
streamSuccessExecuteStream:
|
||||||
out := make(chan cliproxyexecutor.StreamChunk)
|
out := make(chan cliproxyexecutor.StreamChunk)
|
||||||
go func(resp *http.Response) {
|
go func(resp *http.Response) {
|
||||||
defer close(out)
|
defer close(out)
|
||||||
@@ -911,13 +1225,7 @@ attemptLoop:
|
|||||||
|
|
||||||
switch {
|
switch {
|
||||||
case lastStatus != 0:
|
case lastStatus != 0:
|
||||||
sErr := statusErr{code: lastStatus, msg: string(lastBody)}
|
err = newAntigravityStatusErr(lastStatus, lastBody)
|
||||||
if lastStatus == http.StatusTooManyRequests {
|
|
||||||
if retryAfter, parseErr := parseRetryDelay(lastBody); parseErr == nil && retryAfter != nil {
|
|
||||||
sErr.retryAfter = retryAfter
|
|
||||||
}
|
|
||||||
}
|
|
||||||
err = sErr
|
|
||||||
case lastErr != nil:
|
case lastErr != nil:
|
||||||
err = lastErr
|
err = lastErr
|
||||||
default:
|
default:
|
||||||
@@ -1479,7 +1787,7 @@ func antigravityWait(ctx context.Context, wait time.Duration) error {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func antigravityBaseURLFallbackOrder(auth *cliproxyauth.Auth) []string {
|
var antigravityBaseURLFallbackOrder = func(auth *cliproxyauth.Auth) []string {
|
||||||
if base := resolveCustomAntigravityBaseURL(auth); base != "" {
|
if base := resolveCustomAntigravityBaseURL(auth); base != "" {
|
||||||
return []string{base}
|
return []string{base}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -0,0 +1,423 @@
|
|||||||
|
package executor
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"io"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"strings"
|
||||||
|
"sync"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
||||||
|
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
|
||||||
|
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor"
|
||||||
|
sdktranslator "github.com/router-for-me/CLIProxyAPI/v6/sdk/translator"
|
||||||
|
)
|
||||||
|
|
||||||
|
func resetAntigravityCreditsRetryState() {
|
||||||
|
antigravityCreditsExhaustedByAuth = sync.Map{}
|
||||||
|
antigravityPreferCreditsByModel = sync.Map{}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestClassifyAntigravity429(t *testing.T) {
|
||||||
|
t.Run("quota exhausted", func(t *testing.T) {
|
||||||
|
body := []byte(`{"error":{"status":"RESOURCE_EXHAUSTED","message":"QUOTA_EXHAUSTED"}}`)
|
||||||
|
if got := classifyAntigravity429(body); got != antigravity429QuotaExhausted {
|
||||||
|
t.Fatalf("classifyAntigravity429() = %q, want %q", got, antigravity429QuotaExhausted)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("structured rate limit", func(t *testing.T) {
|
||||||
|
body := []byte(`{
|
||||||
|
"error": {
|
||||||
|
"status": "RESOURCE_EXHAUSTED",
|
||||||
|
"details": [
|
||||||
|
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "reason": "RATE_LIMIT_EXCEEDED"},
|
||||||
|
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.5s"}
|
||||||
|
]
|
||||||
|
}
|
||||||
|
}`)
|
||||||
|
if got := classifyAntigravity429(body); got != antigravity429RateLimited {
|
||||||
|
t.Fatalf("classifyAntigravity429() = %q, want %q", got, antigravity429RateLimited)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("structured quota exhausted", func(t *testing.T) {
|
||||||
|
body := []byte(`{
|
||||||
|
"error": {
|
||||||
|
"status": "RESOURCE_EXHAUSTED",
|
||||||
|
"details": [
|
||||||
|
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "reason": "QUOTA_EXHAUSTED"}
|
||||||
|
]
|
||||||
|
}
|
||||||
|
}`)
|
||||||
|
if got := classifyAntigravity429(body); got != antigravity429QuotaExhausted {
|
||||||
|
t.Fatalf("classifyAntigravity429() = %q, want %q", got, antigravity429QuotaExhausted)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("unknown", func(t *testing.T) {
|
||||||
|
body := []byte(`{"error":{"message":"too many requests"}}`)
|
||||||
|
if got := classifyAntigravity429(body); got != antigravity429Unknown {
|
||||||
|
t.Fatalf("classifyAntigravity429() = %q, want %q", got, antigravity429Unknown)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestInjectEnabledCreditTypes(t *testing.T) {
|
||||||
|
body := []byte(`{"model":"gemini-2.5-flash","request":{}}`)
|
||||||
|
got := injectEnabledCreditTypes(body)
|
||||||
|
if got == nil {
|
||||||
|
t.Fatal("injectEnabledCreditTypes() returned nil")
|
||||||
|
}
|
||||||
|
if !strings.Contains(string(got), `"enabledCreditTypes":["GOOGLE_ONE_AI"]`) {
|
||||||
|
t.Fatalf("injectEnabledCreditTypes() = %s, want enabledCreditTypes", string(got))
|
||||||
|
}
|
||||||
|
|
||||||
|
if got := injectEnabledCreditTypes([]byte(`not json`)); got != nil {
|
||||||
|
t.Fatalf("injectEnabledCreditTypes() for invalid json = %s, want nil", string(got))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestShouldMarkAntigravityCreditsExhausted(t *testing.T) {
|
||||||
|
for _, body := range [][]byte{
|
||||||
|
[]byte(`{"error":{"message":"Insufficient GOOGLE_ONE_AI credits"}}`),
|
||||||
|
[]byte(`{"error":{"message":"minimumCreditAmountForUsage requirement not met"}}`),
|
||||||
|
[]byte(`{"error":{"message":"Resource has been exhausted"}}`),
|
||||||
|
} {
|
||||||
|
if !shouldMarkAntigravityCreditsExhausted(http.StatusForbidden, body, nil) {
|
||||||
|
t.Fatalf("shouldMarkAntigravityCreditsExhausted(%s) = false, want true", string(body))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if shouldMarkAntigravityCreditsExhausted(http.StatusServiceUnavailable, []byte(`{"error":{"message":"credits exhausted"}}`), nil) {
|
||||||
|
t.Fatal("shouldMarkAntigravityCreditsExhausted() = true for 5xx, want false")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAntigravityExecute_RetriesQuotaExhaustedWithCredits(t *testing.T) {
|
||||||
|
resetAntigravityCreditsRetryState()
|
||||||
|
t.Cleanup(resetAntigravityCreditsRetryState)
|
||||||
|
|
||||||
|
var (
|
||||||
|
mu sync.Mutex
|
||||||
|
requestBodies []string
|
||||||
|
)
|
||||||
|
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
body, _ := io.ReadAll(r.Body)
|
||||||
|
_ = r.Body.Close()
|
||||||
|
|
||||||
|
mu.Lock()
|
||||||
|
requestBodies = append(requestBodies, string(body))
|
||||||
|
reqNum := len(requestBodies)
|
||||||
|
mu.Unlock()
|
||||||
|
|
||||||
|
if reqNum == 1 {
|
||||||
|
w.WriteHeader(http.StatusTooManyRequests)
|
||||||
|
_, _ = w.Write([]byte(`{"error":{"status":"RESOURCE_EXHAUSTED","message":"QUOTA_EXHAUSTED"}}`))
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
if !strings.Contains(string(body), `"enabledCreditTypes":["GOOGLE_ONE_AI"]`) {
|
||||||
|
t.Fatalf("second request body missing enabledCreditTypes: %s", string(body))
|
||||||
|
}
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
_, _ = w.Write([]byte(`{"response":{"candidates":[{"content":{"role":"model","parts":[{"text":"ok"}]}}],"usageMetadata":{"promptTokenCount":1,"candidatesTokenCount":1,"totalTokenCount":2}}}`))
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
exec := NewAntigravityExecutor(&config.Config{
|
||||||
|
QuotaExceeded: config.QuotaExceeded{AntigravityCredits: true},
|
||||||
|
})
|
||||||
|
auth := &cliproxyauth.Auth{
|
||||||
|
ID: "auth-credits-ok",
|
||||||
|
Attributes: map[string]string{
|
||||||
|
"base_url": server.URL,
|
||||||
|
},
|
||||||
|
Metadata: map[string]any{
|
||||||
|
"access_token": "token",
|
||||||
|
"project_id": "project-1",
|
||||||
|
"expired": time.Now().Add(1 * time.Hour).Format(time.RFC3339),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
resp, err := exec.Execute(context.Background(), auth, cliproxyexecutor.Request{
|
||||||
|
Model: "gemini-2.5-flash",
|
||||||
|
Payload: []byte(`{"request":{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}}`),
|
||||||
|
}, cliproxyexecutor.Options{
|
||||||
|
SourceFormat: sdktranslator.FormatAntigravity,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Execute() error = %v", err)
|
||||||
|
}
|
||||||
|
if len(resp.Payload) == 0 {
|
||||||
|
t.Fatal("Execute() returned empty payload")
|
||||||
|
}
|
||||||
|
|
||||||
|
mu.Lock()
|
||||||
|
defer mu.Unlock()
|
||||||
|
if len(requestBodies) != 2 {
|
||||||
|
t.Fatalf("request count = %d, want 2", len(requestBodies))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAntigravityExecute_SkipsCreditsRetryWhenAlreadyExhausted(t *testing.T) {
|
||||||
|
resetAntigravityCreditsRetryState()
|
||||||
|
t.Cleanup(resetAntigravityCreditsRetryState)
|
||||||
|
|
||||||
|
var requestCount int
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
requestCount++
|
||||||
|
w.WriteHeader(http.StatusTooManyRequests)
|
||||||
|
_, _ = w.Write([]byte(`{"error":{"status":"RESOURCE_EXHAUSTED","message":"QUOTA_EXHAUSTED"}}`))
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
exec := NewAntigravityExecutor(&config.Config{
|
||||||
|
QuotaExceeded: config.QuotaExceeded{AntigravityCredits: true},
|
||||||
|
})
|
||||||
|
auth := &cliproxyauth.Auth{
|
||||||
|
ID: "auth-credits-exhausted",
|
||||||
|
Attributes: map[string]string{
|
||||||
|
"base_url": server.URL,
|
||||||
|
},
|
||||||
|
Metadata: map[string]any{
|
||||||
|
"access_token": "token",
|
||||||
|
"project_id": "project-1",
|
||||||
|
"expired": time.Now().Add(1 * time.Hour).Format(time.RFC3339),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
markAntigravityCreditsExhausted(auth, time.Now())
|
||||||
|
|
||||||
|
_, err := exec.Execute(context.Background(), auth, cliproxyexecutor.Request{
|
||||||
|
Model: "gemini-2.5-flash",
|
||||||
|
Payload: []byte(`{"request":{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}}`),
|
||||||
|
}, cliproxyexecutor.Options{
|
||||||
|
SourceFormat: sdktranslator.FormatAntigravity,
|
||||||
|
})
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("Execute() error = nil, want 429")
|
||||||
|
}
|
||||||
|
sErr, ok := err.(statusErr)
|
||||||
|
if !ok {
|
||||||
|
t.Fatalf("Execute() error type = %T, want statusErr", err)
|
||||||
|
}
|
||||||
|
if got := sErr.StatusCode(); got != http.StatusTooManyRequests {
|
||||||
|
t.Fatalf("Execute() status code = %d, want %d", got, http.StatusTooManyRequests)
|
||||||
|
}
|
||||||
|
if requestCount != 1 {
|
||||||
|
t.Fatalf("request count = %d, want 1", requestCount)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAntigravityExecute_PrefersCreditsAfterSuccessfulFallback(t *testing.T) {
|
||||||
|
resetAntigravityCreditsRetryState()
|
||||||
|
t.Cleanup(resetAntigravityCreditsRetryState)
|
||||||
|
|
||||||
|
var (
|
||||||
|
mu sync.Mutex
|
||||||
|
requestBodies []string
|
||||||
|
)
|
||||||
|
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
body, _ := io.ReadAll(r.Body)
|
||||||
|
_ = r.Body.Close()
|
||||||
|
|
||||||
|
mu.Lock()
|
||||||
|
requestBodies = append(requestBodies, string(body))
|
||||||
|
reqNum := len(requestBodies)
|
||||||
|
mu.Unlock()
|
||||||
|
|
||||||
|
switch reqNum {
|
||||||
|
case 1:
|
||||||
|
w.WriteHeader(http.StatusTooManyRequests)
|
||||||
|
_, _ = w.Write([]byte(`{"error":{"status":"RESOURCE_EXHAUSTED","details":[{"@type":"type.googleapis.com/google.rpc.ErrorInfo","reason":"QUOTA_EXHAUSTED"},{"@type":"type.googleapis.com/google.rpc.RetryInfo","retryDelay":"10s"}]}}`))
|
||||||
|
case 2, 3:
|
||||||
|
if !strings.Contains(string(body), `"enabledCreditTypes":["GOOGLE_ONE_AI"]`) {
|
||||||
|
t.Fatalf("request %d body missing enabledCreditTypes: %s", reqNum, string(body))
|
||||||
|
}
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
_, _ = w.Write([]byte(`{"response":{"candidates":[{"content":{"role":"model","parts":[{"text":"OK"}]}}],"usageMetadata":{"promptTokenCount":1,"candidatesTokenCount":1,"totalTokenCount":2}}}`))
|
||||||
|
default:
|
||||||
|
t.Fatalf("unexpected request count %d", reqNum)
|
||||||
|
}
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
exec := NewAntigravityExecutor(&config.Config{
|
||||||
|
QuotaExceeded: config.QuotaExceeded{AntigravityCredits: true},
|
||||||
|
})
|
||||||
|
auth := &cliproxyauth.Auth{
|
||||||
|
ID: "auth-prefer-credits",
|
||||||
|
Attributes: map[string]string{
|
||||||
|
"base_url": server.URL,
|
||||||
|
},
|
||||||
|
Metadata: map[string]any{
|
||||||
|
"access_token": "token",
|
||||||
|
"project_id": "project-1",
|
||||||
|
"expired": time.Now().Add(1 * time.Hour).Format(time.RFC3339),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
request := cliproxyexecutor.Request{
|
||||||
|
Model: "gemini-2.5-flash",
|
||||||
|
Payload: []byte(`{"request":{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}}`),
|
||||||
|
}
|
||||||
|
opts := cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatAntigravity}
|
||||||
|
|
||||||
|
if _, err := exec.Execute(context.Background(), auth, request, opts); err != nil {
|
||||||
|
t.Fatalf("first Execute() error = %v", err)
|
||||||
|
}
|
||||||
|
if _, err := exec.Execute(context.Background(), auth, request, opts); err != nil {
|
||||||
|
t.Fatalf("second Execute() error = %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
mu.Lock()
|
||||||
|
defer mu.Unlock()
|
||||||
|
if len(requestBodies) != 3 {
|
||||||
|
t.Fatalf("request count = %d, want 3", len(requestBodies))
|
||||||
|
}
|
||||||
|
if strings.Contains(requestBodies[0], `"enabledCreditTypes":["GOOGLE_ONE_AI"]`) {
|
||||||
|
t.Fatalf("first request unexpectedly used credits: %s", requestBodies[0])
|
||||||
|
}
|
||||||
|
if !strings.Contains(requestBodies[1], `"enabledCreditTypes":["GOOGLE_ONE_AI"]`) {
|
||||||
|
t.Fatalf("fallback request missing credits: %s", requestBodies[1])
|
||||||
|
}
|
||||||
|
if !strings.Contains(requestBodies[2], `"enabledCreditTypes":["GOOGLE_ONE_AI"]`) {
|
||||||
|
t.Fatalf("preferred request missing credits: %s", requestBodies[2])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAntigravityExecute_PreservesBaseURLFallbackAfterCreditsRetryFailure(t *testing.T) {
|
||||||
|
resetAntigravityCreditsRetryState()
|
||||||
|
t.Cleanup(resetAntigravityCreditsRetryState)
|
||||||
|
|
||||||
|
var (
|
||||||
|
mu sync.Mutex
|
||||||
|
firstCount int
|
||||||
|
secondCount int
|
||||||
|
)
|
||||||
|
|
||||||
|
firstServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
body, _ := io.ReadAll(r.Body)
|
||||||
|
_ = r.Body.Close()
|
||||||
|
|
||||||
|
mu.Lock()
|
||||||
|
firstCount++
|
||||||
|
reqNum := firstCount
|
||||||
|
mu.Unlock()
|
||||||
|
|
||||||
|
switch reqNum {
|
||||||
|
case 1:
|
||||||
|
w.WriteHeader(http.StatusTooManyRequests)
|
||||||
|
_, _ = w.Write([]byte(`{"error":{"status":"RESOURCE_EXHAUSTED","details":[{"@type":"type.googleapis.com/google.rpc.ErrorInfo","reason":"QUOTA_EXHAUSTED"}]}}`))
|
||||||
|
case 2:
|
||||||
|
if !strings.Contains(string(body), `"enabledCreditTypes":["GOOGLE_ONE_AI"]`) {
|
||||||
|
t.Fatalf("credits retry missing enabledCreditTypes: %s", string(body))
|
||||||
|
}
|
||||||
|
w.WriteHeader(http.StatusForbidden)
|
||||||
|
_, _ = w.Write([]byte(`{"error":{"message":"permission denied"}}`))
|
||||||
|
default:
|
||||||
|
t.Fatalf("unexpected first server request count %d", reqNum)
|
||||||
|
}
|
||||||
|
}))
|
||||||
|
defer firstServer.Close()
|
||||||
|
|
||||||
|
secondServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
mu.Lock()
|
||||||
|
secondCount++
|
||||||
|
mu.Unlock()
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
_, _ = w.Write([]byte(`{"response":{"candidates":[{"content":{"role":"model","parts":[{"text":"ok"}]}}],"usageMetadata":{"promptTokenCount":1,"candidatesTokenCount":1,"totalTokenCount":2}}}`))
|
||||||
|
}))
|
||||||
|
defer secondServer.Close()
|
||||||
|
|
||||||
|
exec := NewAntigravityExecutor(&config.Config{
|
||||||
|
QuotaExceeded: config.QuotaExceeded{AntigravityCredits: true},
|
||||||
|
})
|
||||||
|
auth := &cliproxyauth.Auth{
|
||||||
|
ID: "auth-baseurl-fallback",
|
||||||
|
Attributes: map[string]string{
|
||||||
|
"base_url": firstServer.URL,
|
||||||
|
},
|
||||||
|
Metadata: map[string]any{
|
||||||
|
"access_token": "token",
|
||||||
|
"project_id": "project-1",
|
||||||
|
"expired": time.Now().Add(1 * time.Hour).Format(time.RFC3339),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
originalOrder := antigravityBaseURLFallbackOrder
|
||||||
|
defer func() { antigravityBaseURLFallbackOrder = originalOrder }()
|
||||||
|
antigravityBaseURLFallbackOrder = func(auth *cliproxyauth.Auth) []string {
|
||||||
|
return []string{firstServer.URL, secondServer.URL}
|
||||||
|
}
|
||||||
|
|
||||||
|
resp, err := exec.Execute(context.Background(), auth, cliproxyexecutor.Request{
|
||||||
|
Model: "gemini-2.5-flash",
|
||||||
|
Payload: []byte(`{"request":{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}}`),
|
||||||
|
}, cliproxyexecutor.Options{
|
||||||
|
SourceFormat: sdktranslator.FormatAntigravity,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Execute() error = %v", err)
|
||||||
|
}
|
||||||
|
if len(resp.Payload) == 0 {
|
||||||
|
t.Fatal("Execute() returned empty payload")
|
||||||
|
}
|
||||||
|
if firstCount != 2 {
|
||||||
|
t.Fatalf("first server request count = %d, want 2", firstCount)
|
||||||
|
}
|
||||||
|
if secondCount != 1 {
|
||||||
|
t.Fatalf("second server request count = %d, want 1", secondCount)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAntigravityExecute_DoesNotDirectInjectCreditsWhenFlagDisabled(t *testing.T) {
|
||||||
|
resetAntigravityCreditsRetryState()
|
||||||
|
t.Cleanup(resetAntigravityCreditsRetryState)
|
||||||
|
|
||||||
|
var requestBodies []string
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
body, _ := io.ReadAll(r.Body)
|
||||||
|
_ = r.Body.Close()
|
||||||
|
requestBodies = append(requestBodies, string(body))
|
||||||
|
w.WriteHeader(http.StatusTooManyRequests)
|
||||||
|
_, _ = w.Write([]byte(`{"error":{"status":"RESOURCE_EXHAUSTED","message":"QUOTA_EXHAUSTED"}}`))
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
exec := NewAntigravityExecutor(&config.Config{
|
||||||
|
QuotaExceeded: config.QuotaExceeded{AntigravityCredits: false},
|
||||||
|
})
|
||||||
|
auth := &cliproxyauth.Auth{
|
||||||
|
ID: "auth-flag-disabled",
|
||||||
|
Attributes: map[string]string{
|
||||||
|
"base_url": server.URL,
|
||||||
|
},
|
||||||
|
Metadata: map[string]any{
|
||||||
|
"access_token": "token",
|
||||||
|
"project_id": "project-1",
|
||||||
|
"expired": time.Now().Add(1 * time.Hour).Format(time.RFC3339),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
markAntigravityPreferCredits(auth, "gemini-2.5-flash", time.Now(), nil)
|
||||||
|
|
||||||
|
_, err := exec.Execute(context.Background(), auth, cliproxyexecutor.Request{
|
||||||
|
Model: "gemini-2.5-flash",
|
||||||
|
Payload: []byte(`{"request":{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}}`),
|
||||||
|
}, cliproxyexecutor.Options{
|
||||||
|
SourceFormat: sdktranslator.FormatAntigravity,
|
||||||
|
})
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("Execute() error = nil, want 429")
|
||||||
|
}
|
||||||
|
if len(requestBodies) != 1 {
|
||||||
|
t.Fatalf("request count = %d, want 1", len(requestBodies))
|
||||||
|
}
|
||||||
|
if strings.Contains(requestBodies[0], `"enabledCreditTypes":["GOOGLE_ONE_AI"]`) {
|
||||||
|
t.Fatalf("request unexpectedly used enabledCreditTypes with flag disabled: %s", requestBodies[0])
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -80,6 +80,9 @@ func BuildConfigChangeDetails(oldCfg, newCfg *config.Config) []string {
|
|||||||
if oldCfg.QuotaExceeded.SwitchPreviewModel != newCfg.QuotaExceeded.SwitchPreviewModel {
|
if oldCfg.QuotaExceeded.SwitchPreviewModel != newCfg.QuotaExceeded.SwitchPreviewModel {
|
||||||
changes = append(changes, fmt.Sprintf("quota-exceeded.switch-preview-model: %t -> %t", oldCfg.QuotaExceeded.SwitchPreviewModel, newCfg.QuotaExceeded.SwitchPreviewModel))
|
changes = append(changes, fmt.Sprintf("quota-exceeded.switch-preview-model: %t -> %t", oldCfg.QuotaExceeded.SwitchPreviewModel, newCfg.QuotaExceeded.SwitchPreviewModel))
|
||||||
}
|
}
|
||||||
|
if oldCfg.QuotaExceeded.AntigravityCredits != newCfg.QuotaExceeded.AntigravityCredits {
|
||||||
|
changes = append(changes, fmt.Sprintf("quota-exceeded.antigravity-credits: %t -> %t", oldCfg.QuotaExceeded.AntigravityCredits, newCfg.QuotaExceeded.AntigravityCredits))
|
||||||
|
}
|
||||||
|
|
||||||
if oldCfg.Routing.Strategy != newCfg.Routing.Strategy {
|
if oldCfg.Routing.Strategy != newCfg.Routing.Strategy {
|
||||||
changes = append(changes, fmt.Sprintf("routing.strategy: %s -> %s", oldCfg.Routing.Strategy, newCfg.Routing.Strategy))
|
changes = append(changes, fmt.Sprintf("routing.strategy: %s -> %s", oldCfg.Routing.Strategy, newCfg.Routing.Strategy))
|
||||||
|
|||||||
@@ -229,7 +229,7 @@ func TestBuildConfigChangeDetails_FlagsAndKeys(t *testing.T) {
|
|||||||
MaxRetryCredentials: 1,
|
MaxRetryCredentials: 1,
|
||||||
MaxRetryInterval: 1,
|
MaxRetryInterval: 1,
|
||||||
WebsocketAuth: false,
|
WebsocketAuth: false,
|
||||||
QuotaExceeded: config.QuotaExceeded{SwitchProject: false, SwitchPreviewModel: false},
|
QuotaExceeded: config.QuotaExceeded{SwitchProject: false, SwitchPreviewModel: false, AntigravityCredits: false},
|
||||||
ClaudeKey: []config.ClaudeKey{{APIKey: "c1"}},
|
ClaudeKey: []config.ClaudeKey{{APIKey: "c1"}},
|
||||||
CodexKey: []config.CodexKey{{APIKey: "x1"}},
|
CodexKey: []config.CodexKey{{APIKey: "x1"}},
|
||||||
AmpCode: config.AmpCode{UpstreamAPIKey: "keep", RestrictManagementToLocalhost: false},
|
AmpCode: config.AmpCode{UpstreamAPIKey: "keep", RestrictManagementToLocalhost: false},
|
||||||
@@ -253,7 +253,7 @@ func TestBuildConfigChangeDetails_FlagsAndKeys(t *testing.T) {
|
|||||||
MaxRetryCredentials: 3,
|
MaxRetryCredentials: 3,
|
||||||
MaxRetryInterval: 3,
|
MaxRetryInterval: 3,
|
||||||
WebsocketAuth: true,
|
WebsocketAuth: true,
|
||||||
QuotaExceeded: config.QuotaExceeded{SwitchProject: true, SwitchPreviewModel: true},
|
QuotaExceeded: config.QuotaExceeded{SwitchProject: true, SwitchPreviewModel: true, AntigravityCredits: true},
|
||||||
ClaudeKey: []config.ClaudeKey{
|
ClaudeKey: []config.ClaudeKey{
|
||||||
{APIKey: "c1", BaseURL: "http://new", ProxyURL: "http://p", Headers: map[string]string{"H": "1"}, ExcludedModels: []string{"a"}},
|
{APIKey: "c1", BaseURL: "http://new", ProxyURL: "http://p", Headers: map[string]string{"H": "1"}, ExcludedModels: []string{"a"}},
|
||||||
{APIKey: "c2"},
|
{APIKey: "c2"},
|
||||||
@@ -297,6 +297,7 @@ func TestBuildConfigChangeDetails_FlagsAndKeys(t *testing.T) {
|
|||||||
expectContains(t, details, "nonstream-keepalive-interval: 0 -> 5")
|
expectContains(t, details, "nonstream-keepalive-interval: 0 -> 5")
|
||||||
expectContains(t, details, "quota-exceeded.switch-project: false -> true")
|
expectContains(t, details, "quota-exceeded.switch-project: false -> true")
|
||||||
expectContains(t, details, "quota-exceeded.switch-preview-model: false -> true")
|
expectContains(t, details, "quota-exceeded.switch-preview-model: false -> true")
|
||||||
|
expectContains(t, details, "quota-exceeded.antigravity-credits: false -> true")
|
||||||
expectContains(t, details, "api-keys count: 1 -> 2")
|
expectContains(t, details, "api-keys count: 1 -> 2")
|
||||||
expectContains(t, details, "claude-api-key count: 1 -> 2")
|
expectContains(t, details, "claude-api-key count: 1 -> 2")
|
||||||
expectContains(t, details, "codex-api-key count: 1 -> 2")
|
expectContains(t, details, "codex-api-key count: 1 -> 2")
|
||||||
@@ -320,7 +321,7 @@ func TestBuildConfigChangeDetails_AllBranches(t *testing.T) {
|
|||||||
MaxRetryCredentials: 1,
|
MaxRetryCredentials: 1,
|
||||||
MaxRetryInterval: 1,
|
MaxRetryInterval: 1,
|
||||||
WebsocketAuth: false,
|
WebsocketAuth: false,
|
||||||
QuotaExceeded: config.QuotaExceeded{SwitchProject: false, SwitchPreviewModel: false},
|
QuotaExceeded: config.QuotaExceeded{SwitchProject: false, SwitchPreviewModel: false, AntigravityCredits: false},
|
||||||
GeminiKey: []config.GeminiKey{
|
GeminiKey: []config.GeminiKey{
|
||||||
{APIKey: "g-old", BaseURL: "http://g-old", ProxyURL: "http://gp-old", Headers: map[string]string{"A": "1"}},
|
{APIKey: "g-old", BaseURL: "http://g-old", ProxyURL: "http://gp-old", Headers: map[string]string{"A": "1"}},
|
||||||
},
|
},
|
||||||
@@ -374,7 +375,7 @@ func TestBuildConfigChangeDetails_AllBranches(t *testing.T) {
|
|||||||
MaxRetryCredentials: 3,
|
MaxRetryCredentials: 3,
|
||||||
MaxRetryInterval: 3,
|
MaxRetryInterval: 3,
|
||||||
WebsocketAuth: true,
|
WebsocketAuth: true,
|
||||||
QuotaExceeded: config.QuotaExceeded{SwitchProject: true, SwitchPreviewModel: true},
|
QuotaExceeded: config.QuotaExceeded{SwitchProject: true, SwitchPreviewModel: true, AntigravityCredits: true},
|
||||||
GeminiKey: []config.GeminiKey{
|
GeminiKey: []config.GeminiKey{
|
||||||
{APIKey: "g-new", BaseURL: "http://g-new", ProxyURL: "http://gp-new", Headers: map[string]string{"A": "2"}, ExcludedModels: []string{"x", "y"}},
|
{APIKey: "g-new", BaseURL: "http://g-new", ProxyURL: "http://gp-new", Headers: map[string]string{"A": "2"}, ExcludedModels: []string{"x", "y"}},
|
||||||
},
|
},
|
||||||
@@ -437,6 +438,7 @@ func TestBuildConfigChangeDetails_AllBranches(t *testing.T) {
|
|||||||
expectContains(t, changes, "ws-auth: false -> true")
|
expectContains(t, changes, "ws-auth: false -> true")
|
||||||
expectContains(t, changes, "quota-exceeded.switch-project: false -> true")
|
expectContains(t, changes, "quota-exceeded.switch-project: false -> true")
|
||||||
expectContains(t, changes, "quota-exceeded.switch-preview-model: false -> true")
|
expectContains(t, changes, "quota-exceeded.switch-preview-model: false -> true")
|
||||||
|
expectContains(t, changes, "quota-exceeded.antigravity-credits: false -> true")
|
||||||
expectContains(t, changes, "api-keys: values updated (count unchanged, redacted)")
|
expectContains(t, changes, "api-keys: values updated (count unchanged, redacted)")
|
||||||
expectContains(t, changes, "gemini[0].base-url: http://g-old -> http://g-new")
|
expectContains(t, changes, "gemini[0].base-url: http://g-old -> http://g-new")
|
||||||
expectContains(t, changes, "gemini[0].proxy-url: http://gp-old -> http://gp-new")
|
expectContains(t, changes, "gemini[0].proxy-url: http://gp-old -> http://gp-new")
|
||||||
|
|||||||
Reference in New Issue
Block a user