mirror of
https://mirror.skon.top/github.com/router-for-me/CLIProxyAPI
synced 2026-04-30 16:20:23 +08:00
Merge pull request #2971 from sususu98/feat/antigravity-credits-fallback
feat(antigravity): conductor-level credits fallback for Claude models
This commit is contained in:
@@ -98,7 +98,7 @@ disable-cooling: false
|
||||
quota-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
|
||||
antigravity-credits: true # Whether to retry Antigravity quota_exhausted 429s once with enabledCreditTypes=["GOOGLE_ONE_AI"]
|
||||
antigravity-credits: true # Whether to use credits as last-resort fallback when all free-tier auths are exhausted for Claude models
|
||||
|
||||
# Routing strategy for selecting credentials when multiple match.
|
||||
routing:
|
||||
|
||||
@@ -206,8 +206,9 @@ type QuotaExceeded struct {
|
||||
// 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"`
|
||||
|
||||
// AntigravityCredits indicates whether to retry Antigravity quota_exhausted 429s once
|
||||
// on the same credential with enabledCreditTypes=["GOOGLE_ONE_AI"].
|
||||
// AntigravityCredits enables credits-based last-resort fallback for Claude models.
|
||||
// When all free-tier auths are exhausted (429/503), the conductor retries with
|
||||
// an auth that has available Google One AI credits.
|
||||
AntigravityCredits bool `yaml:"antigravity-credits" json:"antigravity-credits"`
|
||||
}
|
||||
|
||||
|
||||
@@ -52,8 +52,8 @@ const (
|
||||
defaultAntigravityAgent = "antigravity/1.21.9 darwin/arm64" // fallback only; overridden at runtime by misc.AntigravityUserAgent()
|
||||
antigravityAuthType = "antigravity"
|
||||
refreshSkew = 3000 * time.Second
|
||||
antigravityCreditsRetryTTL = 5 * time.Hour
|
||||
antigravityCreditsAutoDisableDuration = 5 * time.Hour
|
||||
antigravityCreditsHintRefreshInterval = 10 * time.Minute
|
||||
antigravityCreditsHintRefreshTimeout = 5 * time.Second
|
||||
antigravityShortQuotaCooldownThreshold = 5 * time.Minute
|
||||
antigravityInstantRetryThreshold = 3 * time.Second
|
||||
// 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**"
|
||||
@@ -62,8 +62,6 @@ const (
|
||||
type antigravity429Category string
|
||||
|
||||
type antigravityCreditsFailureState struct {
|
||||
Count int
|
||||
DisabledUntil time.Time
|
||||
PermanentlyDisabled bool
|
||||
ExplicitBalanceExhausted bool
|
||||
}
|
||||
@@ -91,28 +89,85 @@ var (
|
||||
randSource = rand.New(rand.NewSource(time.Now().UnixNano()))
|
||||
randSourceMutex sync.Mutex
|
||||
antigravityCreditsFailureByAuth sync.Map
|
||||
antigravityPreferCreditsByModel sync.Map
|
||||
antigravityShortCooldownByAuth sync.Map
|
||||
antigravityCreditsBalanceByAuth sync.Map // auth.ID → antigravityCreditsBalance
|
||||
antigravityCreditsHintRefreshByID sync.Map // auth.ID → *antigravityCreditsHintRefreshState
|
||||
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",
|
||||
}
|
||||
)
|
||||
|
||||
type antigravityCreditsBalance struct {
|
||||
CreditAmount float64
|
||||
MinCreditAmount float64
|
||||
PaidTierID string
|
||||
Known bool
|
||||
}
|
||||
|
||||
type antigravityCreditsHintRefreshState struct {
|
||||
mu sync.Mutex
|
||||
lastAttempt time.Time
|
||||
}
|
||||
|
||||
func antigravityAuthHasCredits(auth *cliproxyauth.Auth) bool {
|
||||
if auth == nil || strings.TrimSpace(auth.ID) == "" {
|
||||
return false
|
||||
}
|
||||
if hint, ok := cliproxyauth.GetAntigravityCreditsHint(auth.ID); ok && hint.Known {
|
||||
return hint.Available
|
||||
}
|
||||
val, ok := antigravityCreditsBalanceByAuth.Load(strings.TrimSpace(auth.ID))
|
||||
if !ok {
|
||||
return true // optimistic: assume credits available when balance unknown
|
||||
}
|
||||
bal, valid := val.(antigravityCreditsBalance)
|
||||
if !valid {
|
||||
antigravityCreditsBalanceByAuth.Delete(strings.TrimSpace(auth.ID))
|
||||
return false
|
||||
}
|
||||
if !bal.Known {
|
||||
return false
|
||||
}
|
||||
available := bal.CreditAmount >= bal.MinCreditAmount
|
||||
cliproxyauth.SetAntigravityCreditsHint(strings.TrimSpace(auth.ID), cliproxyauth.AntigravityCreditsHint{
|
||||
Known: true,
|
||||
Available: available,
|
||||
CreditAmount: bal.CreditAmount,
|
||||
MinCreditAmount: bal.MinCreditAmount,
|
||||
PaidTierID: bal.PaidTierID,
|
||||
UpdatedAt: time.Now(),
|
||||
})
|
||||
return available
|
||||
}
|
||||
|
||||
// parseMetaFloat extracts a float64 from auth.Metadata (handles string and numeric types).
|
||||
func parseMetaFloat(metadata map[string]any, key string) (float64, bool) {
|
||||
v, ok := metadata[key]
|
||||
if !ok {
|
||||
return 0, false
|
||||
}
|
||||
switch typed := v.(type) {
|
||||
case float64:
|
||||
return typed, true
|
||||
case int:
|
||||
return float64(typed), true
|
||||
case int64:
|
||||
return float64(typed), true
|
||||
case uint64:
|
||||
return float64(typed), true
|
||||
case json.Number:
|
||||
if f, err := typed.Float64(); err == nil {
|
||||
return f, true
|
||||
}
|
||||
case string:
|
||||
if f, err := strconv.ParseFloat(strings.TrimSpace(typed), 64); err == nil {
|
||||
return f, true
|
||||
}
|
||||
}
|
||||
return 0, false
|
||||
}
|
||||
|
||||
// AntigravityExecutor proxies requests to the antigravity upstream.
|
||||
type AntigravityExecutor struct {
|
||||
cfg *config.Config
|
||||
@@ -189,7 +244,7 @@ func validateAntigravityRequestSignatures(from sdktranslator.Format, rawJSON []b
|
||||
if from.String() != "claude" {
|
||||
return rawJSON, nil
|
||||
}
|
||||
// Always strip thinking blocks with empty signatures (proxy-generated).
|
||||
// Always strip thinking blocks with invalid signatures (empty or non-Claude-format).
|
||||
rawJSON = antigravityclaude.StripEmptySignatureThinkingBlocks(rawJSON)
|
||||
if cache.SignatureCacheEnabled() {
|
||||
return rawJSON, nil
|
||||
@@ -298,6 +353,41 @@ func decideAntigravity429(body []byte) antigravity429Decision {
|
||||
decision.retryAfter = retryAfter
|
||||
}
|
||||
|
||||
status := strings.TrimSpace(gjson.GetBytes(body, "error.status").String())
|
||||
if !strings.EqualFold(status, "RESOURCE_EXHAUSTED") {
|
||||
return decision
|
||||
}
|
||||
|
||||
details := gjson.GetBytes(body, "error.details")
|
||||
if details.Exists() && details.IsArray() {
|
||||
for _, detail := range details.Array() {
|
||||
if detail.Get("@type").String() != "type.googleapis.com/google.rpc.ErrorInfo" {
|
||||
continue
|
||||
}
|
||||
reason := strings.TrimSpace(detail.Get("reason").String())
|
||||
decision.reason = reason
|
||||
switch {
|
||||
case strings.EqualFold(reason, "QUOTA_EXHAUSTED"):
|
||||
decision.kind = antigravity429DecisionFullQuotaExhausted
|
||||
return decision
|
||||
case strings.EqualFold(reason, "RATE_LIMIT_EXCEEDED"):
|
||||
if decision.retryAfter == nil {
|
||||
decision.kind = antigravity429DecisionSoftRetry
|
||||
return decision
|
||||
}
|
||||
switch {
|
||||
case *decision.retryAfter < antigravityInstantRetryThreshold:
|
||||
decision.kind = antigravity429DecisionInstantRetrySameAuth
|
||||
case *decision.retryAfter < antigravityShortQuotaCooldownThreshold:
|
||||
decision.kind = antigravity429DecisionShortCooldownSwitchAuth
|
||||
default:
|
||||
decision.kind = antigravity429DecisionFullQuotaExhausted
|
||||
}
|
||||
return decision
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
lowerBody := strings.ToLower(string(body))
|
||||
for _, keyword := range antigravityQuotaExhaustedKeywords {
|
||||
if strings.Contains(lowerBody, keyword) {
|
||||
@@ -307,123 +397,14 @@ func decideAntigravity429(body []byte) antigravity429Decision {
|
||||
}
|
||||
}
|
||||
|
||||
status := strings.TrimSpace(gjson.GetBytes(body, "error.status").String())
|
||||
if !strings.EqualFold(status, "RESOURCE_EXHAUSTED") {
|
||||
return decision
|
||||
}
|
||||
|
||||
details := gjson.GetBytes(body, "error.details")
|
||||
if !details.Exists() || !details.IsArray() {
|
||||
decision.kind = antigravity429DecisionSoftRetry
|
||||
return decision
|
||||
}
|
||||
|
||||
for _, detail := range details.Array() {
|
||||
if detail.Get("@type").String() != "type.googleapis.com/google.rpc.ErrorInfo" {
|
||||
continue
|
||||
}
|
||||
reason := strings.TrimSpace(detail.Get("reason").String())
|
||||
decision.reason = reason
|
||||
switch {
|
||||
case strings.EqualFold(reason, "QUOTA_EXHAUSTED"):
|
||||
decision.kind = antigravity429DecisionFullQuotaExhausted
|
||||
return decision
|
||||
case strings.EqualFold(reason, "RATE_LIMIT_EXCEEDED"):
|
||||
if decision.retryAfter == nil {
|
||||
decision.kind = antigravity429DecisionSoftRetry
|
||||
return decision
|
||||
}
|
||||
switch {
|
||||
case *decision.retryAfter < antigravityInstantRetryThreshold:
|
||||
decision.kind = antigravity429DecisionInstantRetrySameAuth
|
||||
case *decision.retryAfter < antigravityShortQuotaCooldownThreshold:
|
||||
decision.kind = antigravity429DecisionShortCooldownSwitchAuth
|
||||
default:
|
||||
decision.kind = antigravity429DecisionFullQuotaExhausted
|
||||
}
|
||||
return decision
|
||||
}
|
||||
}
|
||||
|
||||
decision.kind = antigravity429DecisionSoftRetry
|
||||
return decision
|
||||
}
|
||||
|
||||
func antigravityHasQuotaResetDelayOrModelInfo(body []byte) bool {
|
||||
if len(body) == 0 {
|
||||
return false
|
||||
}
|
||||
details := gjson.GetBytes(body, "error.details")
|
||||
if !details.Exists() || !details.IsArray() {
|
||||
return false
|
||||
}
|
||||
for _, detail := range details.Array() {
|
||||
if detail.Get("@type").String() != "type.googleapis.com/google.rpc.ErrorInfo" {
|
||||
continue
|
||||
}
|
||||
if strings.TrimSpace(detail.Get("metadata.quotaResetDelay").String()) != "" {
|
||||
return true
|
||||
}
|
||||
if strings.TrimSpace(detail.Get("metadata.model").String()) != "" {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func antigravityCreditsRetryEnabled(cfg *config.Config) bool {
|
||||
return cfg != nil && cfg.QuotaExceeded.AntigravityCredits
|
||||
}
|
||||
|
||||
func antigravityCreditsFailureStateForAuth(auth *cliproxyauth.Auth) (string, antigravityCreditsFailureState, bool) {
|
||||
if auth == nil || strings.TrimSpace(auth.ID) == "" {
|
||||
return "", antigravityCreditsFailureState{}, false
|
||||
}
|
||||
authID := strings.TrimSpace(auth.ID)
|
||||
value, ok := antigravityCreditsFailureByAuth.Load(authID)
|
||||
if !ok {
|
||||
return authID, antigravityCreditsFailureState{}, true
|
||||
}
|
||||
state, ok := value.(antigravityCreditsFailureState)
|
||||
if !ok {
|
||||
antigravityCreditsFailureByAuth.Delete(authID)
|
||||
return authID, antigravityCreditsFailureState{}, true
|
||||
}
|
||||
return authID, state, true
|
||||
}
|
||||
|
||||
func antigravityCreditsDisabled(auth *cliproxyauth.Auth, now time.Time) bool {
|
||||
authID, state, ok := antigravityCreditsFailureStateForAuth(auth)
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
if state.PermanentlyDisabled {
|
||||
return true
|
||||
}
|
||||
if state.DisabledUntil.IsZero() {
|
||||
return false
|
||||
}
|
||||
if state.DisabledUntil.After(now) {
|
||||
return true
|
||||
}
|
||||
antigravityCreditsFailureByAuth.Delete(authID)
|
||||
return false
|
||||
}
|
||||
|
||||
func recordAntigravityCreditsFailure(auth *cliproxyauth.Auth, now time.Time) {
|
||||
authID, state, ok := antigravityCreditsFailureStateForAuth(auth)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if state.PermanentlyDisabled {
|
||||
antigravityCreditsFailureByAuth.Store(authID, state)
|
||||
return
|
||||
}
|
||||
state.Count++
|
||||
state.DisabledUntil = now.Add(antigravityCreditsAutoDisableDuration)
|
||||
antigravityCreditsFailureByAuth.Store(authID, state)
|
||||
}
|
||||
|
||||
func clearAntigravityCreditsFailureState(auth *cliproxyauth.Auth) {
|
||||
if auth == nil || strings.TrimSpace(auth.ID) == "" {
|
||||
return
|
||||
@@ -440,6 +421,25 @@ func markAntigravityCreditsPermanentlyDisabled(auth *cliproxyauth.Auth) {
|
||||
ExplicitBalanceExhausted: true,
|
||||
}
|
||||
antigravityCreditsFailureByAuth.Store(authID, state)
|
||||
antigravityCreditsBalanceByAuth.Store(authID, antigravityCreditsBalance{
|
||||
CreditAmount: 0,
|
||||
MinCreditAmount: 1,
|
||||
Known: true,
|
||||
})
|
||||
cliproxyauth.SetAntigravityCreditsHint(authID, cliproxyauth.AntigravityCreditsHint{
|
||||
Known: true,
|
||||
Available: false,
|
||||
CreditAmount: 0,
|
||||
MinCreditAmount: 1,
|
||||
UpdatedAt: time.Now(),
|
||||
})
|
||||
}
|
||||
|
||||
func clearAntigravityCreditsPermanentlyDisabled(auth *cliproxyauth.Auth) {
|
||||
if auth == nil || strings.TrimSpace(auth.ID) == "" {
|
||||
return
|
||||
}
|
||||
antigravityCreditsFailureByAuth.Delete(strings.TrimSpace(auth.ID))
|
||||
}
|
||||
|
||||
func antigravityHasExplicitCreditsBalanceExhaustedReason(body []byte) bool {
|
||||
@@ -462,81 +462,6 @@ func antigravityHasExplicitCreditsBalanceExhaustedReason(body []byte) bool {
|
||||
return false
|
||||
}
|
||||
|
||||
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) {
|
||||
if keyword == "resource has been exhausted" &&
|
||||
statusCode == http.StatusTooManyRequests &&
|
||||
decideAntigravity429(body).kind == antigravity429DecisionSoftRetry &&
|
||||
!antigravityHasQuotaResetDelayOrModelInfo(body) {
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func newAntigravityStatusErr(statusCode int, body []byte) statusErr {
|
||||
err := statusErr{code: statusCode, msg: string(body)}
|
||||
if statusCode == http.StatusTooManyRequests {
|
||||
@@ -547,129 +472,6 @@ func newAntigravityStatusErr(statusCode int, body []byte) statusErr {
|
||||
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 decideAntigravity429(originalBody).kind != antigravity429DecisionFullQuotaExhausted {
|
||||
return nil, false
|
||||
}
|
||||
now := time.Now()
|
||||
if shouldForcePermanentDisableCredits(originalBody) {
|
||||
clearAntigravityPreferCredits(auth, modelName)
|
||||
markAntigravityCreditsPermanentlyDisabled(auth)
|
||||
return nil, false
|
||||
}
|
||||
|
||||
if antigravityHasExplicitCreditsBalanceExhaustedReason(originalBody) {
|
||||
clearAntigravityPreferCredits(auth, modelName)
|
||||
markAntigravityCreditsPermanentlyDisabled(auth)
|
||||
return nil, false
|
||||
}
|
||||
|
||||
if antigravityCreditsDisabled(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 {
|
||||
helps.RecordAPIResponseError(ctx, e.cfg, errReq)
|
||||
clearAntigravityPreferCredits(auth, modelName)
|
||||
recordAntigravityCreditsFailure(auth, now)
|
||||
return nil, true
|
||||
}
|
||||
httpResp, errDo := httpClient.Do(httpReq)
|
||||
if errDo != nil {
|
||||
helps.RecordAPIResponseError(ctx, e.cfg, errDo)
|
||||
clearAntigravityPreferCredits(auth, modelName)
|
||||
recordAntigravityCreditsFailure(auth, now)
|
||||
return nil, true
|
||||
}
|
||||
if httpResp.StatusCode >= http.StatusOK && httpResp.StatusCode < http.StatusMultipleChoices {
|
||||
retryAfter, _ := parseRetryDelay(originalBody)
|
||||
markAntigravityPreferCredits(auth, modelName, now, retryAfter)
|
||||
clearAntigravityCreditsFailureState(auth)
|
||||
return httpResp, true
|
||||
}
|
||||
|
||||
helps.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 {
|
||||
helps.RecordAPIResponseError(ctx, e.cfg, errRead)
|
||||
clearAntigravityPreferCredits(auth, modelName)
|
||||
recordAntigravityCreditsFailure(auth, now)
|
||||
return nil, true
|
||||
}
|
||||
helps.AppendAPIResponseChunk(ctx, e.cfg, bodyBytes)
|
||||
if shouldForcePermanentDisableCredits(bodyBytes) {
|
||||
clearAntigravityPreferCredits(auth, modelName)
|
||||
markAntigravityCreditsPermanentlyDisabled(auth)
|
||||
return nil, true
|
||||
}
|
||||
|
||||
if antigravityHasExplicitCreditsBalanceExhaustedReason(bodyBytes) {
|
||||
clearAntigravityPreferCredits(auth, modelName)
|
||||
markAntigravityCreditsPermanentlyDisabled(auth)
|
||||
return nil, true
|
||||
}
|
||||
|
||||
clearAntigravityPreferCredits(auth, modelName)
|
||||
recordAntigravityCreditsFailure(auth, now)
|
||||
return nil, true
|
||||
}
|
||||
|
||||
func (e *AntigravityExecutor) handleDirectCreditsFailure(ctx context.Context, auth *cliproxyauth.Auth, modelName string, reqErr error) {
|
||||
if reqErr != nil {
|
||||
if shouldForcePermanentDisableCredits(reqErrBody(reqErr)) {
|
||||
clearAntigravityPreferCredits(auth, modelName)
|
||||
markAntigravityCreditsPermanentlyDisabled(auth)
|
||||
return
|
||||
}
|
||||
|
||||
if antigravityHasExplicitCreditsBalanceExhaustedReason(reqErrBody(reqErr)) {
|
||||
clearAntigravityPreferCredits(auth, modelName)
|
||||
markAntigravityCreditsPermanentlyDisabled(auth)
|
||||
return
|
||||
}
|
||||
|
||||
helps.RecordAPIResponseError(ctx, e.cfg, reqErr)
|
||||
}
|
||||
clearAntigravityPreferCredits(auth, modelName)
|
||||
recordAntigravityCreditsFailure(auth, time.Now())
|
||||
}
|
||||
func reqErrBody(reqErr error) []byte {
|
||||
if reqErr == nil {
|
||||
return nil
|
||||
}
|
||||
msg := reqErr.Error()
|
||||
if strings.TrimSpace(msg) == "" {
|
||||
return nil
|
||||
}
|
||||
return []byte(msg)
|
||||
}
|
||||
|
||||
func shouldForcePermanentDisableCredits(body []byte) bool {
|
||||
return antigravityHasExplicitCreditsBalanceExhaustedReason(body)
|
||||
}
|
||||
|
||||
// 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) {
|
||||
if opts.Alt == "responses/compact" {
|
||||
@@ -721,6 +523,8 @@ func (e *AntigravityExecutor) Execute(ctx context.Context, auth *cliproxyauth.Au
|
||||
requestedModel := helps.PayloadRequestedModel(opts, req.Model)
|
||||
translated = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, "antigravity", "request", translated, originalTranslated, requestedModel)
|
||||
|
||||
useCredits := cliproxyauth.AntigravityCreditsRequested(ctx) && antigravityCreditsRetryEnabled(e.cfg)
|
||||
|
||||
baseURLs := antigravityBaseURLFallbackOrder(auth)
|
||||
httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0)
|
||||
attempts := antigravityRetryAttempts(auth, e.cfg)
|
||||
@@ -733,11 +537,10 @@ attemptLoop:
|
||||
|
||||
for idx, baseURL := range baseURLs {
|
||||
requestPayload := translated
|
||||
usedCreditsDirect := false
|
||||
if antigravityCreditsRetryEnabled(e.cfg) && antigravityShouldPreferCredits(auth, baseModel, time.Now()) {
|
||||
if creditsPayload := injectEnabledCreditTypes(translated); len(creditsPayload) > 0 {
|
||||
requestPayload = creditsPayload
|
||||
usedCreditsDirect = true
|
||||
if useCredits {
|
||||
if cp := injectEnabledCreditTypes(translated); len(cp) > 0 {
|
||||
requestPayload = cp
|
||||
helps.MarkCreditsUsed(ctx)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -785,7 +588,6 @@ attemptLoop:
|
||||
wait := antigravityInstantRetryDelay(*decision.retryAfter)
|
||||
log.Debugf("antigravity executor: instant retry for model %s, waiting %s", baseModel, wait)
|
||||
if errWait := antigravityWait(ctx, wait); errWait != nil {
|
||||
|
||||
return resp, errWait
|
||||
}
|
||||
}
|
||||
@@ -794,34 +596,13 @@ attemptLoop:
|
||||
case antigravity429DecisionShortCooldownSwitchAuth:
|
||||
if decision.retryAfter != nil && *decision.retryAfter > 0 {
|
||||
markAntigravityShortCooldown(auth, baseModel, time.Now(), *decision.retryAfter)
|
||||
log.Debugf("antigravity executor: short quota cooldown (%s) for model %s, recorded cooldown and skipping credits fallback", *decision.retryAfter, baseModel)
|
||||
log.Debugf("antigravity executor: short quota cooldown (%s) for model %s, recorded cooldown", *decision.retryAfter, baseModel)
|
||||
}
|
||||
case antigravity429DecisionFullQuotaExhausted:
|
||||
if usedCreditsDirect {
|
||||
clearAntigravityPreferCredits(auth, baseModel)
|
||||
recordAntigravityCreditsFailure(auth, time.Now())
|
||||
} else {
|
||||
creditsResp, _ := e.attemptCreditsFallback(ctx, auth, httpClient, token, baseModel, translated, false, opts.Alt, baseURL, bodyBytes)
|
||||
if creditsResp != nil {
|
||||
helps.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 {
|
||||
helps.RecordAPIResponseError(ctx, e.cfg, errCreditsRead)
|
||||
err = errCreditsRead
|
||||
return resp, err
|
||||
}
|
||||
helps.AppendAPIResponseChunk(ctx, e.cfg, creditsBody)
|
||||
reporter.Publish(ctx, helps.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 useCredits && antigravityHasExplicitCreditsBalanceExhaustedReason(bodyBytes) {
|
||||
markAntigravityCreditsPermanentlyDisabled(auth)
|
||||
}
|
||||
// No credits logic - just fall through to error return below
|
||||
}
|
||||
}
|
||||
|
||||
@@ -870,6 +651,10 @@ attemptLoop:
|
||||
return resp, err
|
||||
}
|
||||
|
||||
// Success
|
||||
if useCredits {
|
||||
clearAntigravityCreditsFailureState(auth)
|
||||
}
|
||||
reporter.Publish(ctx, helps.ParseAntigravityUsage(bodyBytes))
|
||||
var param any
|
||||
converted := sdktranslator.TranslateNonStream(ctx, to, from, req.Model, opts.OriginalRequest, translated, bodyBytes, ¶m)
|
||||
@@ -935,6 +720,8 @@ func (e *AntigravityExecutor) executeClaudeNonStream(ctx context.Context, auth *
|
||||
requestedModel := helps.PayloadRequestedModel(opts, req.Model)
|
||||
translated = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, "antigravity", "request", translated, originalTranslated, requestedModel)
|
||||
|
||||
useCredits := cliproxyauth.AntigravityCreditsRequested(ctx) && antigravityCreditsRetryEnabled(e.cfg)
|
||||
|
||||
baseURLs := antigravityBaseURLFallbackOrder(auth)
|
||||
httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0)
|
||||
|
||||
@@ -948,11 +735,10 @@ attemptLoop:
|
||||
|
||||
for idx, baseURL := range baseURLs {
|
||||
requestPayload := translated
|
||||
usedCreditsDirect := false
|
||||
if antigravityCreditsRetryEnabled(e.cfg) && antigravityShouldPreferCredits(auth, baseModel, time.Now()) {
|
||||
if creditsPayload := injectEnabledCreditTypes(translated); len(creditsPayload) > 0 {
|
||||
requestPayload = creditsPayload
|
||||
usedCreditsDirect = true
|
||||
if useCredits {
|
||||
if cp := injectEnabledCreditTypes(translated); len(cp) > 0 {
|
||||
requestPayload = cp
|
||||
helps.MarkCreditsUsed(ctx)
|
||||
}
|
||||
}
|
||||
httpReq, errReq := e.buildRequest(ctx, auth, token, baseModel, requestPayload, true, opts.Alt, baseURL)
|
||||
@@ -1014,7 +800,6 @@ attemptLoop:
|
||||
wait := antigravityInstantRetryDelay(*decision.retryAfter)
|
||||
log.Debugf("antigravity executor: instant retry for model %s, waiting %s", baseModel, wait)
|
||||
if errWait := antigravityWait(ctx, wait); errWait != nil {
|
||||
|
||||
return resp, errWait
|
||||
}
|
||||
}
|
||||
@@ -1023,25 +808,16 @@ attemptLoop:
|
||||
case antigravity429DecisionShortCooldownSwitchAuth:
|
||||
if decision.retryAfter != nil && *decision.retryAfter > 0 {
|
||||
markAntigravityShortCooldown(auth, baseModel, time.Now(), *decision.retryAfter)
|
||||
log.Debugf("antigravity executor: short quota cooldown (%s) for model %s, recorded cooldown and skipping credits fallback", *decision.retryAfter, baseModel)
|
||||
log.Debugf("antigravity executor: short quota cooldown (%s) for model %s, recorded cooldown", *decision.retryAfter, baseModel)
|
||||
}
|
||||
case antigravity429DecisionFullQuotaExhausted:
|
||||
if usedCreditsDirect {
|
||||
clearAntigravityPreferCredits(auth, baseModel)
|
||||
recordAntigravityCreditsFailure(auth, time.Now())
|
||||
} else {
|
||||
creditsResp, _ := e.attemptCreditsFallback(ctx, auth, httpClient, token, baseModel, translated, true, opts.Alt, baseURL, bodyBytes)
|
||||
if creditsResp != nil {
|
||||
httpResp = creditsResp
|
||||
helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone())
|
||||
}
|
||||
if useCredits && antigravityHasExplicitCreditsBalanceExhaustedReason(bodyBytes) {
|
||||
markAntigravityCreditsPermanentlyDisabled(auth)
|
||||
}
|
||||
// No credits logic - just fall through to error return below
|
||||
}
|
||||
}
|
||||
|
||||
if httpResp.StatusCode >= http.StatusOK && httpResp.StatusCode < http.StatusMultipleChoices {
|
||||
goto streamSuccessClaudeNonStream
|
||||
}
|
||||
lastStatus = httpResp.StatusCode
|
||||
lastBody = append([]byte(nil), bodyBytes...)
|
||||
lastErr = nil
|
||||
@@ -1085,7 +861,10 @@ attemptLoop:
|
||||
return resp, err
|
||||
}
|
||||
|
||||
streamSuccessClaudeNonStream:
|
||||
// Stream success
|
||||
if useCredits {
|
||||
clearAntigravityCreditsFailureState(auth)
|
||||
}
|
||||
out := make(chan cliproxyexecutor.StreamChunk)
|
||||
go func(resp *http.Response) {
|
||||
defer close(out)
|
||||
@@ -1389,6 +1168,7 @@ func (e *AntigravityExecutor) ExecuteStream(ctx context.Context, auth *cliproxya
|
||||
if updatedAuth != nil {
|
||||
auth = updatedAuth
|
||||
}
|
||||
|
||||
originalTranslated := sdktranslator.TranslateRequest(from, to, baseModel, originalPayload, true)
|
||||
translated := sdktranslator.TranslateRequest(from, to, baseModel, req.Payload, true)
|
||||
|
||||
@@ -1400,6 +1180,8 @@ func (e *AntigravityExecutor) ExecuteStream(ctx context.Context, auth *cliproxya
|
||||
requestedModel := helps.PayloadRequestedModel(opts, req.Model)
|
||||
translated = helps.ApplyPayloadConfigWithRoot(e.cfg, baseModel, "antigravity", "request", translated, originalTranslated, requestedModel)
|
||||
|
||||
useCredits := cliproxyauth.AntigravityCreditsRequested(ctx) && antigravityCreditsRetryEnabled(e.cfg)
|
||||
|
||||
baseURLs := antigravityBaseURLFallbackOrder(auth)
|
||||
httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0)
|
||||
|
||||
@@ -1413,11 +1195,10 @@ attemptLoop:
|
||||
|
||||
for idx, baseURL := range baseURLs {
|
||||
requestPayload := translated
|
||||
usedCreditsDirect := false
|
||||
if antigravityCreditsRetryEnabled(e.cfg) && antigravityShouldPreferCredits(auth, baseModel, time.Now()) {
|
||||
if creditsPayload := injectEnabledCreditTypes(translated); len(creditsPayload) > 0 {
|
||||
requestPayload = creditsPayload
|
||||
usedCreditsDirect = true
|
||||
if useCredits {
|
||||
if cp := injectEnabledCreditTypes(translated); len(cp) > 0 {
|
||||
requestPayload = cp
|
||||
helps.MarkCreditsUsed(ctx)
|
||||
}
|
||||
}
|
||||
httpReq, errReq := e.buildRequest(ctx, auth, token, baseModel, requestPayload, true, opts.Alt, baseURL)
|
||||
@@ -1478,7 +1259,6 @@ attemptLoop:
|
||||
wait := antigravityInstantRetryDelay(*decision.retryAfter)
|
||||
log.Debugf("antigravity executor: instant retry for model %s, waiting %s", baseModel, wait)
|
||||
if errWait := antigravityWait(ctx, wait); errWait != nil {
|
||||
|
||||
return nil, errWait
|
||||
}
|
||||
}
|
||||
@@ -1487,25 +1267,16 @@ attemptLoop:
|
||||
case antigravity429DecisionShortCooldownSwitchAuth:
|
||||
if decision.retryAfter != nil && *decision.retryAfter > 0 {
|
||||
markAntigravityShortCooldown(auth, baseModel, time.Now(), *decision.retryAfter)
|
||||
log.Debugf("antigravity executor: short quota cooldown (%s) for model %s, recorded cooldown and skipping credits fallback", *decision.retryAfter, baseModel)
|
||||
log.Debugf("antigravity executor: short quota cooldown (%s) for model %s recorded", *decision.retryAfter, baseModel)
|
||||
}
|
||||
case antigravity429DecisionFullQuotaExhausted:
|
||||
if usedCreditsDirect {
|
||||
clearAntigravityPreferCredits(auth, baseModel)
|
||||
recordAntigravityCreditsFailure(auth, time.Now())
|
||||
} else {
|
||||
creditsResp, _ := e.attemptCreditsFallback(ctx, auth, httpClient, token, baseModel, translated, true, opts.Alt, baseURL, bodyBytes)
|
||||
if creditsResp != nil {
|
||||
httpResp = creditsResp
|
||||
helps.RecordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone())
|
||||
}
|
||||
if useCredits && antigravityHasExplicitCreditsBalanceExhaustedReason(bodyBytes) {
|
||||
markAntigravityCreditsPermanentlyDisabled(auth)
|
||||
}
|
||||
// No credits logic - just fall through to error return below
|
||||
}
|
||||
}
|
||||
|
||||
if httpResp.StatusCode >= http.StatusOK && httpResp.StatusCode < http.StatusMultipleChoices {
|
||||
goto streamSuccessExecuteStream
|
||||
}
|
||||
lastStatus = httpResp.StatusCode
|
||||
lastBody = append([]byte(nil), bodyBytes...)
|
||||
lastErr = nil
|
||||
@@ -1549,7 +1320,10 @@ attemptLoop:
|
||||
return nil, err
|
||||
}
|
||||
|
||||
streamSuccessExecuteStream:
|
||||
// Stream success
|
||||
if useCredits {
|
||||
clearAntigravityCreditsFailureState(auth)
|
||||
}
|
||||
out := make(chan cliproxyexecutor.StreamChunk)
|
||||
go func(resp *http.Response) {
|
||||
defer close(out)
|
||||
@@ -1792,6 +1566,7 @@ func (e *AntigravityExecutor) ensureAccessToken(ctx context.Context, auth *clipr
|
||||
accessToken := metaStringValue(auth.Metadata, "access_token")
|
||||
expiry := tokenExpiry(auth.Metadata)
|
||||
if accessToken != "" && expiry.After(time.Now().Add(refreshSkew)) {
|
||||
e.maybeRefreshAntigravityCreditsHint(ctx, auth, accessToken)
|
||||
return accessToken, nil, nil
|
||||
}
|
||||
refreshCtx := context.Background()
|
||||
@@ -1807,6 +1582,63 @@ func (e *AntigravityExecutor) ensureAccessToken(ctx context.Context, auth *clipr
|
||||
return metaStringValue(updated.Metadata, "access_token"), updated, nil
|
||||
}
|
||||
|
||||
func (e *AntigravityExecutor) maybeRefreshAntigravityCreditsHint(ctx context.Context, auth *cliproxyauth.Auth, accessToken string) {
|
||||
if e == nil || auth == nil || !antigravityCreditsRetryEnabled(e.cfg) {
|
||||
return
|
||||
}
|
||||
if ctx != nil && ctx.Err() != nil {
|
||||
return
|
||||
}
|
||||
authID := strings.TrimSpace(auth.ID)
|
||||
if authID == "" {
|
||||
return
|
||||
}
|
||||
if hint, ok := cliproxyauth.GetAntigravityCreditsHint(authID); ok && hint.Known {
|
||||
return
|
||||
}
|
||||
if strings.TrimSpace(accessToken) == "" {
|
||||
accessToken = metaStringValue(auth.Metadata, "access_token")
|
||||
}
|
||||
if strings.TrimSpace(accessToken) == "" {
|
||||
return
|
||||
}
|
||||
|
||||
state := &antigravityCreditsHintRefreshState{}
|
||||
if existing, loaded := antigravityCreditsHintRefreshByID.LoadOrStore(authID, state); loaded {
|
||||
if cast, ok := existing.(*antigravityCreditsHintRefreshState); ok && cast != nil {
|
||||
state = cast
|
||||
} else {
|
||||
antigravityCreditsHintRefreshByID.Delete(authID)
|
||||
antigravityCreditsHintRefreshByID.Store(authID, state)
|
||||
}
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
if !state.mu.TryLock() {
|
||||
return
|
||||
}
|
||||
if !state.lastAttempt.IsZero() && now.Sub(state.lastAttempt) < antigravityCreditsHintRefreshInterval {
|
||||
state.mu.Unlock()
|
||||
return
|
||||
}
|
||||
state.lastAttempt = now
|
||||
|
||||
refreshCtx := context.Background()
|
||||
if ctx != nil {
|
||||
if rt, ok := ctx.Value("cliproxy.roundtripper").(http.RoundTripper); ok && rt != nil {
|
||||
refreshCtx = context.WithValue(refreshCtx, "cliproxy.roundtripper", rt)
|
||||
}
|
||||
}
|
||||
refreshCtx, cancel := context.WithTimeout(refreshCtx, antigravityCreditsHintRefreshTimeout)
|
||||
authCopy := auth.Clone()
|
||||
|
||||
go func(state *antigravityCreditsHintRefreshState, auth *cliproxyauth.Auth, token string) {
|
||||
defer cancel()
|
||||
defer state.mu.Unlock()
|
||||
e.updateAntigravityCreditsBalance(refreshCtx, auth, token)
|
||||
}(state, authCopy, accessToken)
|
||||
}
|
||||
|
||||
func (e *AntigravityExecutor) refreshToken(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) {
|
||||
if auth == nil {
|
||||
return nil, statusErr{code: http.StatusUnauthorized, msg: "missing auth"}
|
||||
@@ -1882,6 +1714,7 @@ func (e *AntigravityExecutor) refreshToken(ctx context.Context, auth *cliproxyau
|
||||
if errProject := e.ensureAntigravityProjectID(ctx, auth, tokenResp.AccessToken); errProject != nil {
|
||||
log.Warnf("antigravity executor: ensure project id failed: %v", errProject)
|
||||
}
|
||||
e.updateAntigravityCreditsBalance(ctx, auth, tokenResp.AccessToken)
|
||||
return auth, nil
|
||||
}
|
||||
|
||||
@@ -1918,6 +1751,94 @@ func (e *AntigravityExecutor) ensureAntigravityProjectID(ctx context.Context, au
|
||||
return nil
|
||||
}
|
||||
|
||||
func (e *AntigravityExecutor) updateAntigravityCreditsBalance(ctx context.Context, auth *cliproxyauth.Auth, accessToken string) {
|
||||
if auth == nil || strings.TrimSpace(auth.ID) == "" {
|
||||
return
|
||||
}
|
||||
token := strings.TrimSpace(accessToken)
|
||||
if token == "" {
|
||||
token = metaStringValue(auth.Metadata, "access_token")
|
||||
}
|
||||
if token == "" {
|
||||
return
|
||||
}
|
||||
|
||||
loadReqBody := `{"metadata":{"ideType":"ANTIGRAVITY","platform":"PLATFORM_UNSPECIFIED","pluginType":"GEMINI"}}`
|
||||
endpointURL := "https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist"
|
||||
httpReq, errReq := http.NewRequestWithContext(ctx, http.MethodPost, endpointURL, strings.NewReader(loadReqBody))
|
||||
if errReq != nil {
|
||||
log.Debugf("antigravity executor: create loadCodeAssist request error: %v", errReq)
|
||||
return
|
||||
}
|
||||
httpReq.Header.Set("Authorization", "Bearer "+token)
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
httpReq.Header.Set("User-Agent", "google-api-nodejs-client/9.15.1")
|
||||
|
||||
httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0)
|
||||
httpResp, errDo := httpClient.Do(httpReq)
|
||||
if errDo != nil {
|
||||
log.Debugf("antigravity executor: loadCodeAssist request error: %v", errDo)
|
||||
return
|
||||
}
|
||||
defer func() {
|
||||
if errClose := httpResp.Body.Close(); errClose != nil {
|
||||
log.Errorf("antigravity executor: close loadCodeAssist response body error: %v", errClose)
|
||||
}
|
||||
}()
|
||||
|
||||
bodyBytes, errRead := io.ReadAll(httpResp.Body)
|
||||
if errRead != nil || httpResp.StatusCode < http.StatusOK || httpResp.StatusCode >= http.StatusMultipleChoices {
|
||||
log.Debugf("antigravity executor: loadCodeAssist returned status %d, err=%v", httpResp.StatusCode, errRead)
|
||||
return
|
||||
}
|
||||
|
||||
authID := strings.TrimSpace(auth.ID)
|
||||
paidTierID := strings.TrimSpace(gjson.GetBytes(bodyBytes, "paidTier.id").String())
|
||||
|
||||
credits := gjson.GetBytes(bodyBytes, "paidTier.availableCredits")
|
||||
if !credits.IsArray() {
|
||||
cliproxyauth.SetAntigravityCreditsHint(authID, cliproxyauth.AntigravityCreditsHint{
|
||||
Known: true,
|
||||
Available: false,
|
||||
PaidTierID: paidTierID,
|
||||
UpdatedAt: time.Now(),
|
||||
})
|
||||
return
|
||||
}
|
||||
for _, credit := range credits.Array() {
|
||||
if !strings.EqualFold(credit.Get("creditType").String(), "GOOGLE_ONE_AI") {
|
||||
continue
|
||||
}
|
||||
creditAmount, errCA := strconv.ParseFloat(strings.TrimSpace(credit.Get("creditAmount").String()), 64)
|
||||
if errCA != nil {
|
||||
continue
|
||||
}
|
||||
minAmount, errMA := strconv.ParseFloat(strings.TrimSpace(credit.Get("minimumCreditAmountForUsage").String()), 64)
|
||||
if errMA != nil {
|
||||
continue
|
||||
}
|
||||
bal := antigravityCreditsBalance{
|
||||
CreditAmount: creditAmount,
|
||||
MinCreditAmount: minAmount,
|
||||
PaidTierID: paidTierID,
|
||||
Known: true,
|
||||
}
|
||||
antigravityCreditsBalanceByAuth.Store(authID, bal)
|
||||
cliproxyauth.SetAntigravityCreditsHint(authID, cliproxyauth.AntigravityCreditsHint{
|
||||
Known: true,
|
||||
Available: creditAmount >= minAmount,
|
||||
CreditAmount: creditAmount,
|
||||
MinCreditAmount: minAmount,
|
||||
PaidTierID: paidTierID,
|
||||
UpdatedAt: time.Now(),
|
||||
})
|
||||
if creditAmount >= minAmount {
|
||||
clearAntigravityCreditsPermanentlyDisabled(auth)
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
func (e *AntigravityExecutor) buildRequest(ctx context.Context, auth *cliproxyauth.Auth, token, modelName string, payload []byte, stream bool, alt, baseURL string) (*http.Request, error) {
|
||||
if token == "" {
|
||||
return nil, statusErr{code: http.StatusUnauthorized, msg: "missing access token"}
|
||||
|
||||
@@ -18,8 +18,9 @@ import (
|
||||
|
||||
func resetAntigravityCreditsRetryState() {
|
||||
antigravityCreditsFailureByAuth = sync.Map{}
|
||||
antigravityPreferCreditsByModel = sync.Map{}
|
||||
antigravityShortCooldownByAuth = sync.Map{}
|
||||
antigravityCreditsBalanceByAuth = sync.Map{}
|
||||
antigravityCreditsHintRefreshByID = sync.Map{}
|
||||
}
|
||||
|
||||
func TestClassifyAntigravity429(t *testing.T) {
|
||||
@@ -30,6 +31,43 @@ func TestClassifyAntigravity429(t *testing.T) {
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("standard antigravity rate limit with ui message stays rate limited", func(t *testing.T) {
|
||||
body := []byte(`{
|
||||
"error": {
|
||||
"code": 429,
|
||||
"message": "You have exhausted your capacity on this model. Your quota will reset after 0s.",
|
||||
"status": "RESOURCE_EXHAUSTED",
|
||||
"details": [
|
||||
{
|
||||
"@type": "type.googleapis.com/google.rpc.ErrorInfo",
|
||||
"reason": "RATE_LIMIT_EXCEEDED",
|
||||
"domain": "cloudcode-pa.googleapis.com",
|
||||
"metadata": {
|
||||
"model": "claude-opus-4-6-thinking",
|
||||
"quotaResetDelay": "479.417207ms",
|
||||
"quotaResetTimeStamp": "2026-04-20T09:19:49Z",
|
||||
"uiMessage": "true"
|
||||
}
|
||||
},
|
||||
{
|
||||
"@type": "type.googleapis.com/google.rpc.RetryInfo",
|
||||
"retryDelay": "0.479417207s"
|
||||
}
|
||||
]
|
||||
}
|
||||
}`)
|
||||
if got := classifyAntigravity429(body); got != antigravity429RateLimited {
|
||||
t.Fatalf("classifyAntigravity429() = %q, want %q", got, antigravity429RateLimited)
|
||||
}
|
||||
decision := decideAntigravity429(body)
|
||||
if decision.kind != antigravity429DecisionInstantRetrySameAuth {
|
||||
t.Fatalf("decideAntigravity429().kind = %q, want %q", decision.kind, antigravity429DecisionInstantRetrySameAuth)
|
||||
}
|
||||
if decision.retryAfter == nil {
|
||||
t.Fatal("decideAntigravity429().retryAfter = nil")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("structured rate limit", func(t *testing.T) {
|
||||
body := []byte(`{
|
||||
"error": {
|
||||
@@ -67,8 +105,31 @@ func TestClassifyAntigravity429(t *testing.T) {
|
||||
})
|
||||
}
|
||||
|
||||
func TestAntigravityShouldRetryNoCapacity_Standard503(t *testing.T) {
|
||||
body := []byte(`{
|
||||
"error": {
|
||||
"code": 503,
|
||||
"message": "No capacity available for model gemini-3.1-flash-image on the server",
|
||||
"status": "UNAVAILABLE",
|
||||
"details": [
|
||||
{
|
||||
"@type": "type.googleapis.com/google.rpc.ErrorInfo",
|
||||
"reason": "MODEL_CAPACITY_EXHAUSTED",
|
||||
"domain": "cloudcode-pa.googleapis.com",
|
||||
"metadata": {
|
||||
"model": "gemini-3.1-flash-image"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
}`)
|
||||
if !antigravityShouldRetryNoCapacity(http.StatusServiceUnavailable, body) {
|
||||
t.Fatal("antigravityShouldRetryNoCapacity() = false, want true")
|
||||
}
|
||||
}
|
||||
|
||||
func TestInjectEnabledCreditTypes(t *testing.T) {
|
||||
body := []byte(`{"model":"gemini-2.5-flash","request":{}}`)
|
||||
body := []byte(`{"model":"claude-sonnet-4-6","request":{}}`)
|
||||
got := injectEnabledCreditTypes(body)
|
||||
if got == nil {
|
||||
t.Fatal("injectEnabledCreditTypes() returned nil")
|
||||
@@ -82,34 +143,18 @@ func TestInjectEnabledCreditTypes(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestShouldMarkAntigravityCreditsExhausted(t *testing.T) {
|
||||
t.Run("credit errors are marked", func(t *testing.T) {
|
||||
for _, body := range [][]byte{
|
||||
[]byte(`{"error":{"message":"Insufficient GOOGLE_ONE_AI credits"}}`),
|
||||
[]byte(`{"error":{"message":"minimumCreditAmountForUsage requirement not met"}}`),
|
||||
} {
|
||||
if !shouldMarkAntigravityCreditsExhausted(http.StatusForbidden, body, nil) {
|
||||
t.Fatalf("shouldMarkAntigravityCreditsExhausted(%s) = false, want true", string(body))
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("transient 429 resource exhausted is not marked", func(t *testing.T) {
|
||||
body := []byte(`{"error":{"code":429,"message":"Resource has been exhausted (e.g. check quota).","status":"RESOURCE_EXHAUSTED"}}`)
|
||||
if shouldMarkAntigravityCreditsExhausted(http.StatusTooManyRequests, body, nil) {
|
||||
t.Fatalf("shouldMarkAntigravityCreditsExhausted(%s) = true, want false", string(body))
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("resource exhausted with quota metadata is still marked", func(t *testing.T) {
|
||||
body := []byte(`{"error":{"code":429,"message":"Resource has been exhausted","status":"RESOURCE_EXHAUSTED","details":[{"@type":"type.googleapis.com/google.rpc.ErrorInfo","metadata":{"quotaResetDelay":"1h","model":"claude-sonnet-4-6"}}]}}`)
|
||||
if !shouldMarkAntigravityCreditsExhausted(http.StatusTooManyRequests, 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 TestParseRetryDelay_HumanReadableDuration(t *testing.T) {
|
||||
body := []byte(`{"error":{"message":"You have exhausted your capacity on this model. Your quota will reset after 1h43m56s."}}`)
|
||||
retryAfter, err := parseRetryDelay(body)
|
||||
if err != nil {
|
||||
t.Fatalf("parseRetryDelay() error = %v", err)
|
||||
}
|
||||
if retryAfter == nil {
|
||||
t.Fatal("parseRetryDelay() returned nil")
|
||||
}
|
||||
want := time.Hour + 43*time.Minute + 56*time.Second
|
||||
if *retryAfter != want {
|
||||
t.Fatalf("parseRetryDelay() = %v, want %v", *retryAfter, want)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -147,7 +192,7 @@ func TestAntigravityExecute_RetriesTransient429ResourceExhausted(t *testing.T) {
|
||||
}
|
||||
|
||||
resp, err := exec.Execute(context.Background(), auth, cliproxyexecutor.Request{
|
||||
Model: "gemini-2.5-flash",
|
||||
Model: "claude-sonnet-4-6",
|
||||
Payload: []byte(`{"request":{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}}`),
|
||||
}, cliproxyexecutor.Options{
|
||||
SourceFormat: sdktranslator.FormatAntigravity,
|
||||
@@ -163,32 +208,18 @@ func TestAntigravityExecute_RetriesTransient429ResourceExhausted(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestAntigravityExecute_RetriesQuotaExhaustedWithCredits(t *testing.T) {
|
||||
func TestAntigravityExecute_CreditsInjectedWhenConductorRequests(t *testing.T) {
|
||||
resetAntigravityCreditsRetryState()
|
||||
t.Cleanup(resetAntigravityCreditsRetryState)
|
||||
|
||||
var (
|
||||
mu sync.Mutex
|
||||
requestBodies []string
|
||||
)
|
||||
|
||||
var 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))
|
||||
t.Fatalf("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}}}`))
|
||||
@@ -199,7 +230,7 @@ func TestAntigravityExecute_RetriesQuotaExhaustedWithCredits(t *testing.T) {
|
||||
QuotaExceeded: config.QuotaExceeded{AntigravityCredits: true},
|
||||
})
|
||||
auth := &cliproxyauth.Auth{
|
||||
ID: "auth-credits-ok",
|
||||
ID: "auth-credits-conductor",
|
||||
Attributes: map[string]string{
|
||||
"base_url": server.URL,
|
||||
},
|
||||
@@ -210,8 +241,11 @@ func TestAntigravityExecute_RetriesQuotaExhaustedWithCredits(t *testing.T) {
|
||||
},
|
||||
}
|
||||
|
||||
resp, err := exec.Execute(context.Background(), auth, cliproxyexecutor.Request{
|
||||
Model: "gemini-2.5-flash",
|
||||
// Simulate conductor setting credits requested flag in context
|
||||
ctx := cliproxyauth.WithAntigravityCredits(context.Background())
|
||||
|
||||
resp, err := exec.Execute(ctx, auth, cliproxyexecutor.Request{
|
||||
Model: "claude-sonnet-4-6",
|
||||
Payload: []byte(`{"request":{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}}`),
|
||||
}, cliproxyexecutor.Options{
|
||||
SourceFormat: sdktranslator.FormatAntigravity,
|
||||
@@ -222,227 +256,12 @@ func TestAntigravityExecute_RetriesQuotaExhaustedWithCredits(t *testing.T) {
|
||||
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))
|
||||
if len(requestBodies) != 1 {
|
||||
t.Fatalf("request count = %d, want 1", 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),
|
||||
},
|
||||
}
|
||||
recordAntigravityCreditsFailure(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) {
|
||||
func TestAntigravityExecute_NoCreditsWithoutConductorFlag(t *testing.T) {
|
||||
resetAntigravityCreditsRetryState()
|
||||
t.Cleanup(resetAntigravityCreditsRetryState)
|
||||
|
||||
@@ -457,10 +276,10 @@ func TestAntigravityExecute_DoesNotDirectInjectCreditsWhenFlagDisabled(t *testin
|
||||
defer server.Close()
|
||||
|
||||
exec := NewAntigravityExecutor(&config.Config{
|
||||
QuotaExceeded: config.QuotaExceeded{AntigravityCredits: false},
|
||||
QuotaExceeded: config.QuotaExceeded{AntigravityCredits: true},
|
||||
})
|
||||
auth := &cliproxyauth.Auth{
|
||||
ID: "auth-flag-disabled",
|
||||
ID: "auth-no-conductor-flag",
|
||||
Attributes: map[string]string{
|
||||
"base_url": server.URL,
|
||||
},
|
||||
@@ -470,10 +289,10 @@ func TestAntigravityExecute_DoesNotDirectInjectCreditsWhenFlagDisabled(t *testin
|
||||
"expired": time.Now().Add(1 * time.Hour).Format(time.RFC3339),
|
||||
},
|
||||
}
|
||||
markAntigravityPreferCredits(auth, "gemini-2.5-flash", time.Now(), nil)
|
||||
|
||||
// No conductor credits flag set in context
|
||||
_, err := exec.Execute(context.Background(), auth, cliproxyexecutor.Request{
|
||||
Model: "gemini-2.5-flash",
|
||||
Model: "claude-sonnet-4-6",
|
||||
Payload: []byte(`{"request":{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}}`),
|
||||
}, cliproxyexecutor.Options{
|
||||
SourceFormat: sdktranslator.FormatAntigravity,
|
||||
@@ -484,7 +303,156 @@ func TestAntigravityExecute_DoesNotDirectInjectCreditsWhenFlagDisabled(t *testin
|
||||
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])
|
||||
// Should NOT contain credits since conductor didn't request them
|
||||
if strings.Contains(requestBodies[0], `"enabledCreditTypes"`) {
|
||||
t.Fatalf("request should not contain enabledCreditTypes without conductor flag: %s", requestBodies[0])
|
||||
}
|
||||
}
|
||||
|
||||
func TestAntigravityAuthHasCredits(t *testing.T) {
|
||||
t.Run("sufficient balance", func(t *testing.T) {
|
||||
resetAntigravityCreditsRetryState()
|
||||
auth := &cliproxyauth.Auth{ID: "test-sufficient"}
|
||||
antigravityCreditsBalanceByAuth.Store("test-sufficient", antigravityCreditsBalance{
|
||||
CreditAmount: 25000,
|
||||
MinCreditAmount: 50,
|
||||
Known: true,
|
||||
})
|
||||
if !antigravityAuthHasCredits(auth) {
|
||||
t.Fatal("antigravityAuthHasCredits() = false, want true")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("insufficient balance", func(t *testing.T) {
|
||||
resetAntigravityCreditsRetryState()
|
||||
auth := &cliproxyauth.Auth{ID: "test-insufficient"}
|
||||
antigravityCreditsBalanceByAuth.Store("test-insufficient", antigravityCreditsBalance{
|
||||
CreditAmount: 30,
|
||||
MinCreditAmount: 50,
|
||||
Known: true,
|
||||
})
|
||||
if antigravityAuthHasCredits(auth) {
|
||||
t.Fatal("antigravityAuthHasCredits() = true, want false")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("no balance stored returns true (optimistic)", func(t *testing.T) {
|
||||
resetAntigravityCreditsRetryState()
|
||||
auth := &cliproxyauth.Auth{ID: "test-no-balance"}
|
||||
if !antigravityAuthHasCredits(auth) {
|
||||
t.Fatal("antigravityAuthHasCredits() = false with no balance stored, want true (optimistic default)")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("nil auth returns false", func(t *testing.T) {
|
||||
if antigravityAuthHasCredits(nil) {
|
||||
t.Fatal("antigravityAuthHasCredits(nil) = true, want false")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("empty ID returns false", func(t *testing.T) {
|
||||
auth := &cliproxyauth.Auth{}
|
||||
if antigravityAuthHasCredits(auth) {
|
||||
t.Fatal("antigravityAuthHasCredits(empty ID) = true, want false")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("unknown balance returns false", func(t *testing.T) {
|
||||
resetAntigravityCreditsRetryState()
|
||||
auth := &cliproxyauth.Auth{ID: "test-unknown"}
|
||||
antigravityCreditsBalanceByAuth.Store("test-unknown", antigravityCreditsBalance{
|
||||
Known: false,
|
||||
})
|
||||
if antigravityAuthHasCredits(auth) {
|
||||
t.Fatal("antigravityAuthHasCredits() = true for unknown balance, want false")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
type roundTripperFunc func(*http.Request) (*http.Response, error)
|
||||
|
||||
func (f roundTripperFunc) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||
return f(req)
|
||||
}
|
||||
|
||||
func TestEnsureAccessToken_WarmTokenLoadsCreditsHint(t *testing.T) {
|
||||
resetAntigravityCreditsRetryState()
|
||||
t.Cleanup(resetAntigravityCreditsRetryState)
|
||||
|
||||
exec := NewAntigravityExecutor(&config.Config{
|
||||
QuotaExceeded: config.QuotaExceeded{AntigravityCredits: true},
|
||||
})
|
||||
auth := &cliproxyauth.Auth{
|
||||
ID: "auth-warm-token-credits",
|
||||
Metadata: map[string]any{
|
||||
"access_token": "token",
|
||||
"expired": time.Now().Add(1 * time.Hour).Format(time.RFC3339),
|
||||
},
|
||||
}
|
||||
ctx := context.WithValue(context.Background(), "cliproxy.roundtripper", roundTripperFunc(func(req *http.Request) (*http.Response, error) {
|
||||
if req.URL.String() != "https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist" {
|
||||
t.Fatalf("unexpected request url %s", req.URL.String())
|
||||
}
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: make(http.Header),
|
||||
Body: io.NopCloser(strings.NewReader(`{"paidTier":{"id":"tier-1","availableCredits":[{"creditType":"GOOGLE_ONE_AI","creditAmount":"25000","minimumCreditAmountForUsage":"50"}]}}`)),
|
||||
}, nil
|
||||
}))
|
||||
|
||||
token, updatedAuth, err := exec.ensureAccessToken(ctx, auth)
|
||||
if err != nil {
|
||||
t.Fatalf("ensureAccessToken() error = %v", err)
|
||||
}
|
||||
if token != "token" {
|
||||
t.Fatalf("ensureAccessToken() token = %q, want %q", token, "token")
|
||||
}
|
||||
if updatedAuth != nil {
|
||||
t.Fatalf("ensureAccessToken() updatedAuth = %v, want nil", updatedAuth)
|
||||
}
|
||||
deadline := time.Now().Add(2 * time.Second)
|
||||
for time.Now().Before(deadline) && !cliproxyauth.HasKnownAntigravityCreditsHint(auth.ID) {
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
if !cliproxyauth.HasKnownAntigravityCreditsHint(auth.ID) {
|
||||
t.Fatal("expected credits hint to be populated for warm token auth")
|
||||
}
|
||||
hint, ok := cliproxyauth.GetAntigravityCreditsHint(auth.ID)
|
||||
if !ok {
|
||||
t.Fatal("expected credits hint lookup to succeed")
|
||||
}
|
||||
if !hint.Available {
|
||||
t.Fatalf("hint.Available = %v, want true", hint.Available)
|
||||
}
|
||||
if hint.CreditAmount != 25000 || hint.MinCreditAmount != 50 {
|
||||
t.Fatalf("hint amounts = (%v, %v), want (25000, 50)", hint.CreditAmount, hint.MinCreditAmount)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseMetaFloat(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
value any
|
||||
wantVal float64
|
||||
wantOK bool
|
||||
}{
|
||||
{"string", "25000", 25000, true},
|
||||
{"float64", float64(100), 100, true},
|
||||
{"int", int(50), 50, true},
|
||||
{"int64", int64(75), 75, true},
|
||||
{"empty string", "", 0, false},
|
||||
{"invalid string", "abc", 0, false},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
meta := map[string]any{"key": tt.value}
|
||||
got, ok := parseMetaFloat(meta, "key")
|
||||
if ok != tt.wantOK {
|
||||
t.Fatalf("parseMetaFloat() ok = %v, want %v", ok, tt.wantOK)
|
||||
}
|
||||
if ok && got != tt.wantVal {
|
||||
t.Fatalf("parseMetaFloat() = %f, want %f", got, tt.wantVal)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -898,7 +898,14 @@ func parseRetryDelay(errorBody []byte) (*time.Duration, error) {
|
||||
if matches := re.FindStringSubmatch(message); len(matches) > 1 {
|
||||
seconds, err := strconv.Atoi(matches[1])
|
||||
if err == nil {
|
||||
return new(time.Duration(seconds) * time.Second), nil
|
||||
duration := time.Duration(seconds) * time.Second
|
||||
return &duration, nil
|
||||
}
|
||||
}
|
||||
reHuman := regexp.MustCompile(`after\s+((?:\d+h)?(?:\d+m)?(?:\d+s)?)\.?`)
|
||||
if matches := reHuman.FindStringSubmatch(strings.ToLower(message)); len(matches) > 1 {
|
||||
if duration, err := time.ParseDuration(matches[1]); err == nil && duration > 0 {
|
||||
return &duration, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -24,6 +24,7 @@ const (
|
||||
apiRequestKey = "API_REQUEST"
|
||||
apiResponseKey = "API_RESPONSE"
|
||||
apiWebsocketTimelineKey = "API_WEBSOCKET_TIMELINE"
|
||||
creditsUsedKey = "__antigravity_credits_used__"
|
||||
)
|
||||
|
||||
// UpstreamRequestLog captures the outbound upstream request details for logging.
|
||||
@@ -568,3 +569,24 @@ func LogWithRequestID(ctx context.Context) *log.Entry {
|
||||
}
|
||||
return log.WithField("request_id", requestID)
|
||||
}
|
||||
|
||||
// MarkCreditsUsed flags the request as having used AI credits for billing.
|
||||
func MarkCreditsUsed(ctx context.Context) {
|
||||
ginCtx := ginContextFrom(ctx)
|
||||
if ginCtx != nil {
|
||||
ginCtx.Set(creditsUsedKey, true)
|
||||
}
|
||||
}
|
||||
|
||||
// CreditsUsed returns true if the request used AI credits.
|
||||
func CreditsUsed(ctx context.Context) bool {
|
||||
ginCtx := ginContextFrom(ctx)
|
||||
if ginCtx != nil {
|
||||
if val, exists := ginCtx.Get(creditsUsedKey); exists {
|
||||
if b, ok := val.(bool); ok {
|
||||
return b
|
||||
}
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
90
sdk/cliproxy/auth/antigravity_credits.go
Normal file
90
sdk/cliproxy/auth/antigravity_credits.go
Normal file
@@ -0,0 +1,90 @@
|
||||
package auth
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
type antigravityUseCreditsContextKey struct{}
|
||||
|
||||
// WithAntigravityCredits returns a child context that signals the executor to
|
||||
// inject enabledCreditTypes into the request payload.
|
||||
func WithAntigravityCredits(ctx context.Context) context.Context {
|
||||
return context.WithValue(ctx, antigravityUseCreditsContextKey{}, true)
|
||||
}
|
||||
|
||||
// AntigravityCreditsRequested reports whether the context carries the credits flag.
|
||||
func AntigravityCreditsRequested(ctx context.Context) bool {
|
||||
if ctx == nil {
|
||||
return false
|
||||
}
|
||||
v, _ := ctx.Value(antigravityUseCreditsContextKey{}).(bool)
|
||||
return v
|
||||
}
|
||||
|
||||
// AntigravityCreditsHint stores the latest known AI credits state for one auth.
|
||||
type AntigravityCreditsHint struct {
|
||||
Known bool
|
||||
Available bool
|
||||
CreditAmount float64
|
||||
MinCreditAmount float64
|
||||
PaidTierID string
|
||||
UpdatedAt time.Time
|
||||
}
|
||||
|
||||
var antigravityCreditsHintByAuth sync.Map
|
||||
|
||||
// SetAntigravityCreditsHint updates the latest known AI credits state for an auth.
|
||||
func SetAntigravityCreditsHint(authID string, hint AntigravityCreditsHint) {
|
||||
authID = strings.TrimSpace(authID)
|
||||
if authID == "" {
|
||||
return
|
||||
}
|
||||
if hint.UpdatedAt.IsZero() {
|
||||
hint.UpdatedAt = time.Now()
|
||||
}
|
||||
antigravityCreditsHintByAuth.Store(authID, hint)
|
||||
}
|
||||
|
||||
// GetAntigravityCreditsHint returns the latest known AI credits state for an auth.
|
||||
func GetAntigravityCreditsHint(authID string) (AntigravityCreditsHint, bool) {
|
||||
authID = strings.TrimSpace(authID)
|
||||
if authID == "" {
|
||||
return AntigravityCreditsHint{}, false
|
||||
}
|
||||
value, ok := antigravityCreditsHintByAuth.Load(authID)
|
||||
if !ok {
|
||||
return AntigravityCreditsHint{}, false
|
||||
}
|
||||
hint, ok := value.(AntigravityCreditsHint)
|
||||
if !ok {
|
||||
antigravityCreditsHintByAuth.Delete(authID)
|
||||
return AntigravityCreditsHint{}, false
|
||||
}
|
||||
return hint, true
|
||||
}
|
||||
|
||||
// HasKnownAntigravityCreditsHint reports whether credits state has been discovered for an auth.
|
||||
func HasKnownAntigravityCreditsHint(authID string) bool {
|
||||
hint, ok := GetAntigravityCreditsHint(authID)
|
||||
return ok && hint.Known
|
||||
}
|
||||
|
||||
func antigravityCreditsAvailableForModel(auth *Auth, model string) bool {
|
||||
if auth == nil {
|
||||
return false
|
||||
}
|
||||
if !strings.EqualFold(strings.TrimSpace(auth.Provider), "antigravity") {
|
||||
return false
|
||||
}
|
||||
if !strings.Contains(strings.ToLower(strings.TrimSpace(model)), "claude") {
|
||||
return false
|
||||
}
|
||||
hint, ok := GetAntigravityCreditsHint(auth.ID)
|
||||
if !ok || !hint.Known {
|
||||
return false
|
||||
}
|
||||
return hint.Available
|
||||
}
|
||||
62
sdk/cliproxy/auth/antigravity_credits_test.go
Normal file
62
sdk/cliproxy/auth/antigravity_credits_test.go
Normal file
@@ -0,0 +1,62 @@
|
||||
package auth
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
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)
|
||||
}
|
||||
}
|
||||
@@ -1202,12 +1202,16 @@ func (m *Manager) Execute(ctx context.Context, providers []string, req cliproxye
|
||||
}
|
||||
}
|
||||
if lastErr != nil {
|
||||
if shouldAttemptAntigravityCreditsFallback(m, lastErr, normalized) {
|
||||
if resp, ok := m.tryAntigravityCreditsExecute(ctx, req, opts); ok {
|
||||
return resp, nil
|
||||
}
|
||||
}
|
||||
return cliproxyexecutor.Response{}, lastErr
|
||||
}
|
||||
return cliproxyexecutor.Response{}, &Error{Code: "auth_not_found", Message: "no auth available"}
|
||||
}
|
||||
|
||||
// ExecuteCount performs a non-streaming execution using the configured selector and executor.
|
||||
// It supports multiple providers for the same model and round-robins the starting provider per model.
|
||||
func (m *Manager) ExecuteCount(ctx context.Context, providers []string, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
||||
normalized := m.normalizeProviders(providers)
|
||||
@@ -1264,6 +1268,11 @@ func (m *Manager) ExecuteStream(ctx context.Context, providers []string, req cli
|
||||
}
|
||||
}
|
||||
if lastErr != nil {
|
||||
if shouldAttemptAntigravityCreditsFallback(m, lastErr, normalized) {
|
||||
if result, ok := m.tryAntigravityCreditsExecuteStream(ctx, req, opts); ok {
|
||||
return result, nil
|
||||
}
|
||||
}
|
||||
return nil, lastErr
|
||||
}
|
||||
return nil, &Error{Code: "auth_not_found", Message: "no auth available"}
|
||||
@@ -2319,7 +2328,8 @@ func retryAfterFromError(err error) *time.Duration {
|
||||
if retryAfter == nil {
|
||||
return nil
|
||||
}
|
||||
return new(*retryAfter)
|
||||
value := *retryAfter
|
||||
return &value
|
||||
}
|
||||
|
||||
func statusCodeFromResult(err *Error) int {
|
||||
@@ -2409,11 +2419,18 @@ func isRequestInvalidError(err error) bool {
|
||||
status := statusCodeFromError(err)
|
||||
switch status {
|
||||
case http.StatusBadRequest:
|
||||
return strings.Contains(err.Error(), "invalid_request_error")
|
||||
msg := err.Error()
|
||||
return strings.Contains(msg, "invalid_request_error") ||
|
||||
strings.Contains(msg, "INVALID_ARGUMENT") ||
|
||||
strings.Contains(msg, "FAILED_PRECONDITION")
|
||||
case http.StatusNotFound:
|
||||
return isRequestScopedNotFoundMessage(err.Error())
|
||||
case http.StatusUnprocessableEntity:
|
||||
return true
|
||||
case http.StatusInternalServerError:
|
||||
msg := err.Error()
|
||||
return strings.Contains(msg, "\"status\":\"UNKNOWN\"") ||
|
||||
strings.Contains(msg, "\"status\": \"UNKNOWN\"")
|
||||
default:
|
||||
return false
|
||||
}
|
||||
@@ -2886,6 +2903,175 @@ func (m *Manager) pickNextMixed(ctx context.Context, providers []string, model s
|
||||
return authCopy, executor, providerKey, nil
|
||||
}
|
||||
|
||||
func (m *Manager) findAllAntigravityCreditsCandidateAuths(routeModel string, opts cliproxyexecutor.Options) []creditsCandidateEntry {
|
||||
if m == nil {
|
||||
return nil
|
||||
}
|
||||
pinnedAuthID := pinnedAuthIDFromMetadata(opts.Metadata)
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
var known []creditsCandidateEntry
|
||||
var unknown []creditsCandidateEntry
|
||||
for _, auth := range m.auths {
|
||||
if auth == nil || auth.Disabled || auth.Status == StatusDisabled {
|
||||
continue
|
||||
}
|
||||
if pinnedAuthID != "" && auth.ID != pinnedAuthID {
|
||||
continue
|
||||
}
|
||||
if !strings.EqualFold(strings.TrimSpace(auth.Provider), "antigravity") {
|
||||
continue
|
||||
}
|
||||
if !strings.Contains(strings.ToLower(strings.TrimSpace(routeModel)), "claude") {
|
||||
continue
|
||||
}
|
||||
providerKey := strings.TrimSpace(strings.ToLower(auth.Provider))
|
||||
executor, ok := m.executors[providerKey]
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
|
||||
hint, okHint := GetAntigravityCreditsHint(auth.ID)
|
||||
if okHint && hint.Known {
|
||||
if !hint.Available {
|
||||
continue
|
||||
}
|
||||
known = append(known, creditsCandidateEntry{
|
||||
auth: auth.Clone(),
|
||||
executor: executor,
|
||||
provider: providerKey,
|
||||
})
|
||||
continue
|
||||
}
|
||||
unknown = append(unknown, creditsCandidateEntry{
|
||||
auth: auth.Clone(),
|
||||
executor: executor,
|
||||
provider: providerKey,
|
||||
})
|
||||
}
|
||||
sort.Slice(known, func(i, j int) bool {
|
||||
return known[i].auth.ID < known[j].auth.ID
|
||||
})
|
||||
sort.Slice(unknown, func(i, j int) bool {
|
||||
return unknown[i].auth.ID < unknown[j].auth.ID
|
||||
})
|
||||
return append(known, unknown...)
|
||||
}
|
||||
|
||||
type creditsCandidateEntry struct {
|
||||
auth *Auth
|
||||
executor ProviderExecutor
|
||||
provider string
|
||||
}
|
||||
|
||||
func shouldAttemptAntigravityCreditsFallback(m *Manager, lastErr error, providers []string) bool {
|
||||
if m == nil || lastErr == nil {
|
||||
return false
|
||||
}
|
||||
if len(providers) > 0 {
|
||||
hasAntigravity := false
|
||||
for _, p := range providers {
|
||||
if strings.EqualFold(strings.TrimSpace(p), "antigravity") {
|
||||
hasAntigravity = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !hasAntigravity {
|
||||
return false
|
||||
}
|
||||
}
|
||||
cfg, _ := m.runtimeConfig.Load().(*internalconfig.Config)
|
||||
if cfg == nil || !cfg.QuotaExceeded.AntigravityCredits {
|
||||
return false
|
||||
}
|
||||
status := statusCodeFromError(lastErr)
|
||||
switch status {
|
||||
case http.StatusTooManyRequests, http.StatusServiceUnavailable:
|
||||
return true
|
||||
case 0:
|
||||
var authErr *Error
|
||||
if errors.As(lastErr, &authErr) && authErr != nil {
|
||||
return authErr.Code == "auth_not_found" || authErr.Code == "auth_unavailable" || authErr.Code == "model_cooldown"
|
||||
}
|
||||
var cooldownErr *modelCooldownError
|
||||
if errors.As(lastErr, &cooldownErr) {
|
||||
return true
|
||||
}
|
||||
return false
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func (m *Manager) tryAntigravityCreditsExecute(ctx context.Context, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, bool) {
|
||||
routeModel := req.Model
|
||||
candidates := m.findAllAntigravityCreditsCandidateAuths(routeModel, opts)
|
||||
for _, c := range candidates {
|
||||
if ctx.Err() != nil {
|
||||
return cliproxyexecutor.Response{}, false
|
||||
}
|
||||
creditsCtx := WithAntigravityCredits(ctx)
|
||||
if rt := m.roundTripperFor(c.auth); rt != nil {
|
||||
creditsCtx = context.WithValue(creditsCtx, roundTripperContextKey{}, rt)
|
||||
creditsCtx = context.WithValue(creditsCtx, "cliproxy.roundtripper", rt)
|
||||
}
|
||||
creditsOpts := ensureRequestedModelMetadata(opts, routeModel)
|
||||
publishSelectedAuthMetadata(creditsOpts.Metadata, c.auth.ID)
|
||||
models := m.executionModelCandidates(c.auth, routeModel)
|
||||
if len(models) == 0 {
|
||||
continue
|
||||
}
|
||||
for _, upstreamModel := range models {
|
||||
resultModel := m.stateModelForExecution(c.auth, routeModel, upstreamModel, len(models) > 1)
|
||||
execReq := req
|
||||
execReq.Model = upstreamModel
|
||||
resp, errExec := c.executor.Execute(creditsCtx, c.auth, execReq, creditsOpts)
|
||||
result := Result{AuthID: c.auth.ID, Provider: c.provider, Model: resultModel, Success: errExec == nil}
|
||||
if errExec != nil {
|
||||
result.Error = &Error{Message: errExec.Error()}
|
||||
if se, ok := errors.AsType[cliproxyexecutor.StatusError](errExec); ok && se != nil {
|
||||
result.Error.HTTPStatus = se.StatusCode()
|
||||
}
|
||||
if ra := retryAfterFromError(errExec); ra != nil {
|
||||
result.RetryAfter = ra
|
||||
}
|
||||
m.MarkResult(creditsCtx, result)
|
||||
continue
|
||||
}
|
||||
m.MarkResult(creditsCtx, result)
|
||||
return resp, true
|
||||
}
|
||||
}
|
||||
return cliproxyexecutor.Response{}, false
|
||||
}
|
||||
|
||||
func (m *Manager) tryAntigravityCreditsExecuteStream(ctx context.Context, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, bool) {
|
||||
routeModel := req.Model
|
||||
candidates := m.findAllAntigravityCreditsCandidateAuths(routeModel, opts)
|
||||
for _, c := range candidates {
|
||||
if ctx.Err() != nil {
|
||||
return nil, false
|
||||
}
|
||||
creditsCtx := WithAntigravityCredits(ctx)
|
||||
if rt := m.roundTripperFor(c.auth); rt != nil {
|
||||
creditsCtx = context.WithValue(creditsCtx, roundTripperContextKey{}, rt)
|
||||
creditsCtx = context.WithValue(creditsCtx, "cliproxy.roundtripper", rt)
|
||||
}
|
||||
creditsOpts := ensureRequestedModelMetadata(opts, routeModel)
|
||||
publishSelectedAuthMetadata(creditsOpts.Metadata, c.auth.ID)
|
||||
models := m.executionModelCandidates(c.auth, routeModel)
|
||||
if len(models) == 0 {
|
||||
continue
|
||||
}
|
||||
result, errStream := m.executeStreamWithModelPool(creditsCtx, c.executor, c.auth, c.provider, req, creditsOpts, routeModel, models, len(models) > 1)
|
||||
if errStream != nil {
|
||||
continue
|
||||
}
|
||||
return result, true
|
||||
}
|
||||
return nil, false
|
||||
}
|
||||
|
||||
func (m *Manager) persist(ctx context.Context, auth *Auth) error {
|
||||
if m.store == nil || auth == nil {
|
||||
return nil
|
||||
@@ -3200,14 +3386,15 @@ func (m *Manager) refreshAuth(ctx context.Context, id string) {
|
||||
m.mu.RLock()
|
||||
auth := m.auths[id]
|
||||
var exec ProviderExecutor
|
||||
var cloned *Auth
|
||||
if auth != nil {
|
||||
exec = m.executors[auth.Provider]
|
||||
cloned = auth.Clone()
|
||||
}
|
||||
m.mu.RUnlock()
|
||||
if auth == nil || exec == nil {
|
||||
return
|
||||
}
|
||||
cloned := auth.Clone()
|
||||
updated, err := exec.Refresh(ctx, cloned)
|
||||
if err != nil && errors.Is(err, context.Canceled) {
|
||||
log.Debugf("refresh canceled for %s, %s", auth.Provider, auth.ID)
|
||||
|
||||
61
sdk/cliproxy/auth/conductor_credits_candidates_test.go
Normal file
61
sdk/cliproxy/auth/conductor_credits_candidates_test.go
Normal file
@@ -0,0 +1,61 @@
|
||||
package auth
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor"
|
||||
)
|
||||
|
||||
func TestFindAllAntigravityCreditsCandidateAuths_PrefersKnownCreditsThenUnknown(t *testing.T) {
|
||||
m := &Manager{
|
||||
auths: map[string]*Auth{
|
||||
"zz-credits": {ID: "zz-credits", Provider: "antigravity"},
|
||||
"aa-unknown": {ID: "aa-unknown", Provider: "antigravity"},
|
||||
"mm-no": {ID: "mm-no", Provider: "antigravity"},
|
||||
},
|
||||
executors: map[string]ProviderExecutor{
|
||||
"antigravity": schedulerTestExecutor{},
|
||||
},
|
||||
}
|
||||
|
||||
SetAntigravityCreditsHint("zz-credits", AntigravityCreditsHint{
|
||||
Known: true,
|
||||
Available: true,
|
||||
UpdatedAt: time.Now(),
|
||||
})
|
||||
SetAntigravityCreditsHint("mm-no", AntigravityCreditsHint{
|
||||
Known: true,
|
||||
Available: false,
|
||||
UpdatedAt: time.Now(),
|
||||
})
|
||||
|
||||
opts := cliproxyexecutor.Options{}
|
||||
|
||||
candidates := m.findAllAntigravityCreditsCandidateAuths("claude-sonnet-4-6", opts)
|
||||
if len(candidates) != 2 {
|
||||
t.Fatalf("candidates len = %d, want 2", len(candidates))
|
||||
}
|
||||
if candidates[0].auth.ID != "zz-credits" {
|
||||
t.Fatalf("candidates[0].auth.ID = %q, want %q", candidates[0].auth.ID, "zz-credits")
|
||||
}
|
||||
if candidates[1].auth.ID != "aa-unknown" {
|
||||
t.Fatalf("candidates[1].auth.ID = %q, want %q", candidates[1].auth.ID, "aa-unknown")
|
||||
}
|
||||
|
||||
nonClaude := m.findAllAntigravityCreditsCandidateAuths("gemini-3-flash", opts)
|
||||
if len(nonClaude) != 0 {
|
||||
t.Fatalf("nonClaude len = %d, want 0", len(nonClaude))
|
||||
}
|
||||
|
||||
pinnedOpts := cliproxyexecutor.Options{
|
||||
Metadata: map[string]any{cliproxyexecutor.PinnedAuthMetadataKey: "aa-unknown"},
|
||||
}
|
||||
pinned := m.findAllAntigravityCreditsCandidateAuths("claude-sonnet-4-6", pinnedOpts)
|
||||
if len(pinned) != 1 {
|
||||
t.Fatalf("pinned len = %d, want 1", len(pinned))
|
||||
}
|
||||
if pinned[0].auth.ID != "aa-unknown" {
|
||||
t.Fatalf("pinned[0].auth.ID = %q, want %q", pinned[0].auth.ID, "aa-unknown")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user