Add projects
This commit is contained in:
parent
2d3a9ad623
commit
8b607dd700
SSH key fingerprint:
SHA256:GBzT2DXvAuGV9XIV5W3WrzVpjU54FThmxHXdbz95J24
1802 changed files with 503346 additions and 2 deletions
|
|
@ -0,0 +1,775 @@
|
|||
package executor
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"math/rand"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
homekv "github.com/router-for-me/CLIProxyAPI/v7/internal/home"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps"
|
||||
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"github.com/tidwall/gjson"
|
||||
"github.com/tidwall/sjson"
|
||||
"golang.org/x/sync/singleflight"
|
||||
)
|
||||
|
||||
type antigravity429Category string
|
||||
|
||||
type antigravityCreditsFailureState struct {
|
||||
PermanentlyDisabled bool
|
||||
ExplicitBalanceExhausted bool
|
||||
}
|
||||
|
||||
type antigravity429DecisionKind string
|
||||
|
||||
const (
|
||||
antigravity429Unknown antigravity429Category = "unknown"
|
||||
antigravity429RateLimited antigravity429Category = "rate_limited"
|
||||
antigravity429QuotaExhausted antigravity429Category = "quota_exhausted"
|
||||
antigravity429SoftRateLimit antigravity429Category = "soft_rate_limit"
|
||||
antigravity429DecisionSoftRetry antigravity429DecisionKind = "soft_retry"
|
||||
antigravity429DecisionInstantRetrySameAuth antigravity429DecisionKind = "instant_retry_same_auth"
|
||||
antigravity429DecisionShortCooldownSwitchAuth antigravity429DecisionKind = "short_cooldown_switch_auth"
|
||||
antigravity429DecisionFullQuotaExhausted antigravity429DecisionKind = "full_quota_exhausted"
|
||||
)
|
||||
|
||||
type antigravity429Decision struct {
|
||||
kind antigravity429DecisionKind
|
||||
retryAfter *time.Duration
|
||||
reason string
|
||||
}
|
||||
|
||||
var (
|
||||
randSource = rand.New(rand.NewSource(time.Now().UnixNano()))
|
||||
randSourceMutex sync.Mutex
|
||||
antigravityCreditsFailureByAuth sync.Map
|
||||
antigravityShortCooldownByAuth sync.Map
|
||||
antigravityCreditsBalanceByAuth sync.Map // auth.ID → antigravityCreditsBalance
|
||||
antigravityCreditsHintRefreshByID sync.Map // auth.ID → *antigravityCreditsHintRefreshState
|
||||
antigravityRefreshGroup singleflight.Group
|
||||
antigravityQuotaExhaustedKeywords = []string{
|
||||
"quota_exhausted",
|
||||
"quota exhausted",
|
||||
}
|
||||
)
|
||||
|
||||
type antigravityKVClient interface {
|
||||
KVGet(ctx context.Context, key string) ([]byte, bool, error)
|
||||
KVSet(ctx context.Context, key string, value []byte, opts homekv.KVSetOptions) (bool, error)
|
||||
KVSetNX(ctx context.Context, key string, value []byte, ttl time.Duration) (bool, error)
|
||||
KVDel(ctx context.Context, keys ...string) (int64, error)
|
||||
}
|
||||
|
||||
var currentAntigravityKVClient = func() (antigravityKVClient, bool, error) {
|
||||
return homekv.CurrentKVClient()
|
||||
}
|
||||
|
||||
type antigravityCreditsBalance struct {
|
||||
CreditAmount float64
|
||||
MinCreditAmount float64
|
||||
PaidTierID string
|
||||
Known bool
|
||||
}
|
||||
|
||||
type antigravityCreditsHintRefreshState struct {
|
||||
mu sync.Mutex
|
||||
lastAttempt time.Time
|
||||
}
|
||||
|
||||
type antigravityTokenRefreshData struct {
|
||||
AccessToken string `json:"access_token"`
|
||||
RefreshToken string `json:"refresh_token"`
|
||||
ExpiresIn int64 `json:"expires_in"`
|
||||
TokenType string `json:"token_type"`
|
||||
}
|
||||
|
||||
func antigravityAuthHasCredits(auth *cliproxyauth.Auth) bool {
|
||||
ok, err := antigravityAuthHasCreditsRequired(context.Background(), auth)
|
||||
if err != nil {
|
||||
log.Errorf("antigravity executor: home kv credits check error: %v", err)
|
||||
return false
|
||||
}
|
||||
return ok
|
||||
}
|
||||
|
||||
func antigravityAuthHasCreditsRequired(ctx context.Context, auth *cliproxyauth.Auth) (bool, error) {
|
||||
if auth == nil || strings.TrimSpace(auth.ID) == "" {
|
||||
return false, nil
|
||||
}
|
||||
authID := strings.TrimSpace(auth.ID)
|
||||
if hint, ok, errHint := cliproxyauth.GetAntigravityCreditsHintRequired(ctx, authID); errHint != nil {
|
||||
return false, errHint
|
||||
} else if ok && hint.Known {
|
||||
return hint.Available, nil
|
||||
}
|
||||
|
||||
client, homeMode, errClient := currentAntigravityKVClient()
|
||||
if homeMode {
|
||||
if errClient != nil {
|
||||
return false, errClient
|
||||
}
|
||||
raw, found, errBalance := client.KVGet(ctx, antigravityCreditsBalanceKey(authID))
|
||||
if errBalance != nil {
|
||||
return false, errBalance
|
||||
}
|
||||
if !found {
|
||||
return true, nil
|
||||
}
|
||||
var homeBalance antigravityCreditsBalance
|
||||
if errUnmarshal := json.Unmarshal(raw, &homeBalance); errUnmarshal != nil {
|
||||
return false, errUnmarshal
|
||||
}
|
||||
return antigravityCreditsBalanceAvailable(authID, homeBalance), nil
|
||||
}
|
||||
|
||||
val, ok := antigravityCreditsBalanceByAuth.Load(authID)
|
||||
if !ok {
|
||||
return true, nil // optimistic: assume credits available when balance unknown
|
||||
}
|
||||
bal, valid := val.(antigravityCreditsBalance)
|
||||
if !valid {
|
||||
antigravityCreditsBalanceByAuth.Delete(authID)
|
||||
return false, nil
|
||||
}
|
||||
return antigravityCreditsBalanceAvailable(authID, bal), nil
|
||||
}
|
||||
|
||||
func antigravityCreditsBalanceAvailable(authID string, bal antigravityCreditsBalance) bool {
|
||||
if !bal.Known {
|
||||
return false
|
||||
}
|
||||
available := bal.CreditAmount >= bal.MinCreditAmount
|
||||
cliproxyauth.SetAntigravityCreditsHint(strings.TrimSpace(authID), 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
|
||||
}
|
||||
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 {
|
||||
switch decideAntigravity429(body).kind {
|
||||
case antigravity429DecisionInstantRetrySameAuth, antigravity429DecisionShortCooldownSwitchAuth:
|
||||
return antigravity429RateLimited
|
||||
case antigravity429DecisionFullQuotaExhausted:
|
||||
return antigravity429QuotaExhausted
|
||||
case antigravity429DecisionSoftRetry:
|
||||
return antigravity429SoftRateLimit
|
||||
default:
|
||||
return antigravity429Unknown
|
||||
}
|
||||
}
|
||||
|
||||
func decideAntigravity429(body []byte) antigravity429Decision {
|
||||
decision := antigravity429Decision{kind: antigravity429DecisionSoftRetry}
|
||||
if len(body) == 0 {
|
||||
return decision
|
||||
}
|
||||
|
||||
if retryAfter, parseErr := helps.ParseRetryDelay(body); parseErr == nil && retryAfter != nil {
|
||||
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) {
|
||||
decision.kind = antigravity429DecisionFullQuotaExhausted
|
||||
decision.reason = "quota_exhausted"
|
||||
return decision
|
||||
}
|
||||
}
|
||||
|
||||
decision.kind = antigravity429DecisionSoftRetry
|
||||
return decision
|
||||
}
|
||||
|
||||
func antigravityCreditsRetryEnabled(cfg *config.Config) bool {
|
||||
return cfg != nil && cfg.QuotaExceeded.AntigravityCredits
|
||||
}
|
||||
|
||||
func clearAntigravityCreditsFailureState(auth *cliproxyauth.Auth) {
|
||||
if auth == nil || strings.TrimSpace(auth.ID) == "" {
|
||||
return
|
||||
}
|
||||
antigravityCreditsFailureByAuth.Delete(strings.TrimSpace(auth.ID))
|
||||
}
|
||||
func markAntigravityCreditsPermanentlyDisabled(auth *cliproxyauth.Auth) {
|
||||
if auth == nil || strings.TrimSpace(auth.ID) == "" {
|
||||
return
|
||||
}
|
||||
authID := strings.TrimSpace(auth.ID)
|
||||
state := antigravityCreditsFailureState{
|
||||
PermanentlyDisabled: true,
|
||||
ExplicitBalanceExhausted: true,
|
||||
}
|
||||
antigravityCreditsFailureByAuth.Store(authID, state)
|
||||
bal := antigravityCreditsBalance{
|
||||
CreditAmount: 0,
|
||||
MinCreditAmount: 1,
|
||||
Known: true,
|
||||
}
|
||||
storeAntigravityCreditsBalanceBestEffort(authID, bal)
|
||||
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 {
|
||||
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
|
||||
}
|
||||
reason := strings.TrimSpace(detail.Get("reason").String())
|
||||
if strings.EqualFold(reason, "INSUFFICIENT_G1_CREDITS_BALANCE") {
|
||||
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 := helps.ParseRetryDelay(body); parseErr == nil && retryAfter != nil {
|
||||
err.retryAfter = retryAfter
|
||||
}
|
||||
}
|
||||
return err
|
||||
}
|
||||
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
|
||||
}
|
||||
|
||||
if client, homeMode, errClient := currentAntigravityKVClient(); homeMode {
|
||||
if errClient != nil {
|
||||
log.Errorf("antigravity executor: home kv best-effort refresh lock failed prefix=cpa:antigravity:*: %v", errClient)
|
||||
return
|
||||
}
|
||||
written, errSetNX := client.KVSetNX(context.Background(), antigravityCreditsRefreshLockKey(authID), []byte("1"), antigravityCreditsHintRefreshInterval)
|
||||
if errSetNX != nil {
|
||||
log.Errorf("antigravity executor: home kv best-effort refresh lock failed prefix=cpa:antigravity:*: %v", errSetNX)
|
||||
return
|
||||
}
|
||||
if !written {
|
||||
return
|
||||
}
|
||||
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(auth *cliproxyauth.Auth, token string) {
|
||||
defer cancel()
|
||||
e.updateAntigravityCreditsBalance(refreshCtx, auth, token)
|
||||
}(authCopy, 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) 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
|
||||
}
|
||||
|
||||
userAgent := resolveUserAgent(auth)
|
||||
loadReqBody, errMarshal := json.Marshal(map[string]any{
|
||||
"metadata": map[string]string{
|
||||
"ideType": "ANTIGRAVITY",
|
||||
},
|
||||
})
|
||||
if errMarshal != nil {
|
||||
log.Debugf("antigravity executor: marshal loadCodeAssist request error: %v", errMarshal)
|
||||
return
|
||||
}
|
||||
baseURL := antigravityLoadCodeAssistBaseURL(auth)
|
||||
endpointURL := strings.TrimSuffix(baseURL, "/") + "/v1internal:loadCodeAssist"
|
||||
httpReq, errReq := http.NewRequestWithContext(ctx, http.MethodPost, endpointURL, bytes.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("Accept", "*/*")
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
httpReq.Header.Set("User-Agent", userAgent)
|
||||
|
||||
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,
|
||||
}
|
||||
storeAntigravityCreditsBalanceBestEffort(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 antigravityShouldRetryNoCapacity(statusCode int, body []byte) bool {
|
||||
if statusCode != http.StatusServiceUnavailable {
|
||||
return false
|
||||
}
|
||||
if len(body) == 0 {
|
||||
return false
|
||||
}
|
||||
msg := strings.ToLower(string(body))
|
||||
return strings.Contains(msg, "no capacity available")
|
||||
}
|
||||
|
||||
func antigravityShouldRetryTransientResourceExhausted429(statusCode int, body []byte) bool {
|
||||
if statusCode != http.StatusTooManyRequests {
|
||||
return false
|
||||
}
|
||||
if len(body) == 0 {
|
||||
return false
|
||||
}
|
||||
if classifyAntigravity429(body) != antigravity429Unknown {
|
||||
return false
|
||||
}
|
||||
status := strings.TrimSpace(gjson.GetBytes(body, "error.status").String())
|
||||
if !strings.EqualFold(status, "RESOURCE_EXHAUSTED") {
|
||||
return false
|
||||
}
|
||||
msg := strings.ToLower(string(body))
|
||||
return strings.Contains(msg, "resource has been exhausted")
|
||||
}
|
||||
|
||||
func antigravityShouldRetrySoftRateLimit(statusCode int, body []byte) bool {
|
||||
if statusCode != http.StatusTooManyRequests {
|
||||
return false
|
||||
}
|
||||
return decideAntigravity429(body).kind == antigravity429DecisionSoftRetry
|
||||
}
|
||||
|
||||
func antigravityShouldBypassShortCooldown(ctx context.Context, cfg *config.Config) bool {
|
||||
return cliproxyauth.AntigravityCreditsRequested(ctx) && antigravityCreditsRetryEnabled(cfg)
|
||||
}
|
||||
|
||||
func antigravitySoftRateLimitDelay(attempt int) time.Duration {
|
||||
if attempt < 0 {
|
||||
attempt = 0
|
||||
}
|
||||
base := time.Duration(attempt+1) * 500 * time.Millisecond
|
||||
if base > 3*time.Second {
|
||||
base = 3 * time.Second
|
||||
}
|
||||
return base
|
||||
}
|
||||
|
||||
func antigravityShortCooldownKey(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 + "|sc"
|
||||
}
|
||||
|
||||
func antigravityCreditsBalanceKey(authID string) string {
|
||||
return "cpa:antigravity:credits-balance:" + strings.TrimSpace(authID)
|
||||
}
|
||||
|
||||
func antigravityCreditsRefreshLockKey(authID string) string {
|
||||
return "cpa:antigravity:credits-refresh-lock:" + strings.TrimSpace(authID)
|
||||
}
|
||||
|
||||
func antigravityShortCooldownKVKey(auth *cliproxyauth.Auth, modelName string) string {
|
||||
if auth == nil {
|
||||
return ""
|
||||
}
|
||||
authID := strings.TrimSpace(auth.ID)
|
||||
modelName = strings.TrimSpace(modelName)
|
||||
if authID == "" || modelName == "" {
|
||||
return ""
|
||||
}
|
||||
return "cpa:antigravity:short-cooldown:" + authID + ":" + homekv.HashKeyPart(modelName)
|
||||
}
|
||||
|
||||
func antigravityIsInShortCooldown(auth *cliproxyauth.Auth, modelName string, now time.Time) (bool, time.Duration) {
|
||||
inCooldown, remaining, errCooldown := antigravityIsInShortCooldownRequired(context.Background(), auth, modelName, now)
|
||||
if errCooldown != nil {
|
||||
log.Errorf("antigravity executor: home kv cooldown read error: %v", errCooldown)
|
||||
return false, 0
|
||||
}
|
||||
return inCooldown, remaining
|
||||
}
|
||||
|
||||
func antigravityIsInShortCooldownRequired(ctx context.Context, auth *cliproxyauth.Auth, modelName string, now time.Time) (bool, time.Duration, error) {
|
||||
kvKey := antigravityShortCooldownKVKey(auth, modelName)
|
||||
client, homeMode, errClient := currentAntigravityKVClient()
|
||||
if homeMode {
|
||||
if errClient != nil {
|
||||
return false, 0, errClient
|
||||
}
|
||||
if kvKey == "" {
|
||||
return false, 0, nil
|
||||
}
|
||||
raw, found, errGet := client.KVGet(ctx, kvKey)
|
||||
if errGet != nil || !found {
|
||||
return false, 0, errGet
|
||||
}
|
||||
untilNano, errParse := strconv.ParseInt(strings.TrimSpace(string(raw)), 10, 64)
|
||||
if errParse != nil {
|
||||
return false, 0, errParse
|
||||
}
|
||||
remaining := time.Unix(0, untilNano).Sub(now)
|
||||
if remaining <= 0 {
|
||||
if _, errDel := client.KVDel(ctx, kvKey); errDel != nil {
|
||||
return false, 0, errDel
|
||||
}
|
||||
return false, 0, nil
|
||||
}
|
||||
return true, remaining, nil
|
||||
}
|
||||
|
||||
key := antigravityShortCooldownKey(auth, modelName)
|
||||
if key == "" {
|
||||
return false, 0, nil
|
||||
}
|
||||
value, ok := antigravityShortCooldownByAuth.Load(key)
|
||||
if !ok {
|
||||
return false, 0, nil
|
||||
}
|
||||
until, ok := value.(time.Time)
|
||||
if !ok || until.IsZero() {
|
||||
antigravityShortCooldownByAuth.Delete(key)
|
||||
return false, 0, nil
|
||||
}
|
||||
remaining := until.Sub(now)
|
||||
if remaining <= 0 {
|
||||
antigravityShortCooldownByAuth.Delete(key)
|
||||
return false, 0, nil
|
||||
}
|
||||
return true, remaining, nil
|
||||
}
|
||||
|
||||
func markAntigravityShortCooldown(auth *cliproxyauth.Auth, modelName string, now time.Time, duration time.Duration) {
|
||||
if errMark := markAntigravityShortCooldownRequired(context.Background(), auth, modelName, now, duration); errMark != nil {
|
||||
log.Errorf("antigravity executor: home kv cooldown write error: %v", errMark)
|
||||
}
|
||||
}
|
||||
|
||||
func markAntigravityShortCooldownRequired(ctx context.Context, auth *cliproxyauth.Auth, modelName string, now time.Time, duration time.Duration) error {
|
||||
kvKey := antigravityShortCooldownKVKey(auth, modelName)
|
||||
client, homeMode, errClient := currentAntigravityKVClient()
|
||||
if homeMode {
|
||||
if errClient != nil {
|
||||
return errClient
|
||||
}
|
||||
if kvKey == "" || duration <= 0 {
|
||||
return nil
|
||||
}
|
||||
until := now.Add(duration)
|
||||
written, errSet := client.KVSet(ctx, kvKey, []byte(strconv.FormatInt(until.UnixNano(), 10)), homekv.KVSetOptions{EX: duration + 5*time.Second})
|
||||
if errSet != nil {
|
||||
return errSet
|
||||
}
|
||||
if !written {
|
||||
return fmt.Errorf("home kv store unavailable")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
key := antigravityShortCooldownKey(auth, modelName)
|
||||
if key == "" {
|
||||
return nil
|
||||
}
|
||||
antigravityShortCooldownByAuth.Store(key, now.Add(duration))
|
||||
return nil
|
||||
}
|
||||
|
||||
func storeAntigravityCreditsBalanceBestEffort(authID string, bal antigravityCreditsBalance) {
|
||||
authID = strings.TrimSpace(authID)
|
||||
if authID == "" {
|
||||
return
|
||||
}
|
||||
if client, homeMode, errClient := currentAntigravityKVClient(); homeMode {
|
||||
if errClient != nil {
|
||||
log.Errorf("antigravity executor: home kv best-effort credits balance set failed prefix=cpa:antigravity:*: %v", errClient)
|
||||
return
|
||||
}
|
||||
raw, errMarshal := json.Marshal(bal)
|
||||
if errMarshal != nil {
|
||||
log.Errorf("antigravity executor: home kv best-effort credits balance set failed prefix=cpa:antigravity:*: %v", errMarshal)
|
||||
return
|
||||
}
|
||||
if _, errSet := client.KVSet(context.Background(), antigravityCreditsBalanceKey(authID), raw, homekv.KVSetOptions{EX: 30 * time.Minute}); errSet != nil {
|
||||
log.Errorf("antigravity executor: home kv best-effort credits balance set failed prefix=cpa:antigravity:*: %v", errSet)
|
||||
}
|
||||
return
|
||||
}
|
||||
antigravityCreditsBalanceByAuth.Store(authID, bal)
|
||||
}
|
||||
|
||||
func homeKVUnavailableStatusErr(cause error) statusErr {
|
||||
if cause == nil {
|
||||
return statusErr{code: http.StatusServiceUnavailable, msg: "home kv store unavailable"}
|
||||
}
|
||||
return statusErr{code: http.StatusServiceUnavailable, msg: fmt.Sprintf("home kv store unavailable: %v", cause)}
|
||||
}
|
||||
|
||||
func antigravityNoCapacityRetryDelay(attempt int) time.Duration {
|
||||
if attempt < 0 {
|
||||
attempt = 0
|
||||
}
|
||||
delay := time.Duration(attempt+1) * 250 * time.Millisecond
|
||||
if delay > 2*time.Second {
|
||||
delay = 2 * time.Second
|
||||
}
|
||||
return delay
|
||||
}
|
||||
|
||||
func antigravityTransient429RetryDelay(attempt int) time.Duration {
|
||||
if attempt < 0 {
|
||||
attempt = 0
|
||||
}
|
||||
delay := time.Duration(attempt+1) * 100 * time.Millisecond
|
||||
if delay > 500*time.Millisecond {
|
||||
delay = 500 * time.Millisecond
|
||||
}
|
||||
return delay
|
||||
}
|
||||
|
||||
func antigravityInstantRetryDelay(wait time.Duration) time.Duration {
|
||||
if wait <= 0 {
|
||||
return 0
|
||||
}
|
||||
return wait + 800*time.Millisecond
|
||||
}
|
||||
|
||||
func antigravityWait(ctx context.Context, wait time.Duration) error {
|
||||
if wait <= 0 {
|
||||
return nil
|
||||
}
|
||||
timer := time.NewTimer(wait)
|
||||
defer timer.Stop()
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case <-timer.C:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue