Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
152 changes: 152 additions & 0 deletions backend/internal/service/antigravity_cooldown_clamp_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,152 @@
//go:build unit

package service

import (
"testing"
"time"

"github.com/stretchr/testify/require"
)

func TestClampAntigravityRateLimitDuration(t *testing.T) {
tests := []struct {
name string
input time.Duration
want time.Duration
wantClamp bool
}{
{
name: "short delay passes through",
input: 500 * time.Millisecond,
want: 500 * time.Millisecond,
wantClamp: false,
},
{
name: "30s passes through",
input: 30 * time.Second,
want: 30 * time.Second,
wantClamp: false,
},
{
name: "exactly the cap is not clamped",
input: antigravityMaxRateLimitCooldown,
want: antigravityMaxRateLimitCooldown,
wantClamp: false,
},
{
name: "5h clamped down to cap",
input: 5 * time.Hour,
want: antigravityMaxRateLimitCooldown,
wantClamp: true,
},
{
name: "24h clamped down to cap",
input: 24 * time.Hour,
want: antigravityMaxRateLimitCooldown,
wantClamp: true,
},
{
name: "7d clamped down to cap",
input: 7 * 24 * time.Hour,
want: antigravityMaxRateLimitCooldown,
wantClamp: true,
},
}

for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got, did := clampAntigravityRateLimitDuration(tt.input)
require.Equal(t, tt.want, got)
require.Equal(t, tt.wantClamp, did)
})
}
}

func TestClampAntigravityRateLimitResetAt(t *testing.T) {
now := time.Date(2026, 5, 9, 12, 0, 0, 0, time.UTC)

tests := []struct {
name string
resetAt time.Time
want time.Time
wantClamp bool
}{
{
name: "near future passes through",
resetAt: now.Add(30 * time.Minute),
want: now.Add(30 * time.Minute),
wantClamp: false,
},
{
name: "exactly the cap is not clamped",
resetAt: now.Add(antigravityMaxRateLimitCooldown),
want: now.Add(antigravityMaxRateLimitCooldown),
wantClamp: false,
},
{
name: "5h clamped to now+cap",
resetAt: now.Add(5 * time.Hour),
want: now.Add(antigravityMaxRateLimitCooldown),
wantClamp: true,
},
{
name: "7d clamped to now+cap",
resetAt: now.Add(7 * 24 * time.Hour),
want: now.Add(antigravityMaxRateLimitCooldown),
wantClamp: true,
},
{
name: "past time passes through (caller decides)",
resetAt: now.Add(-1 * time.Hour),
want: now.Add(-1 * time.Hour),
wantClamp: false,
},
}

for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got, did := clampAntigravityRateLimitResetAt(tt.resetAt, now)
require.Equal(t, tt.want, got)
require.Equal(t, tt.wantClamp, did)
})
}
}

func TestParseAntigravitySmartRetryInfo_ClampsExcessiveRetryDelay(t *testing.T) {
// 模拟上游配额耗尽时返回 24h retryDelay 的情况:解析后应被截断到本地上限。
body := `{
"error": {
"code": 429,
"status": "RESOURCE_EXHAUSTED",
"details": [
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "claude-sonnet-4-6"}, "reason": "RATE_LIMIT_EXCEEDED"},
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "86400s"}
]
}
}`

info := parseAntigravitySmartRetryInfo([]byte(body))
require.NotNil(t, info)
require.Equal(t, "claude-sonnet-4-6", info.ModelName)
require.Equal(t, antigravityMaxRateLimitCooldown, info.RetryDelay,
"24h retryDelay should be clamped to local cap")
}

func TestParseAntigravitySmartRetryInfo_KeepsShortRetryDelay(t *testing.T) {
// 短 retryDelay 应原样保留,不受 clamp 影响。
body := `{
"error": {
"code": 429,
"status": "RESOURCE_EXHAUSTED",
"details": [
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "claude-sonnet-4-6"}, "reason": "RATE_LIMIT_EXCEEDED"},
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.4s"}
]
}
}`

info := parseAntigravitySmartRetryInfo([]byte(body))
require.NotNil(t, info)
require.Equal(t, 400*time.Millisecond, info.RetryDelay)
}
44 changes: 41 additions & 3 deletions backend/internal/service/antigravity_gateway_service.go
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,12 @@ const (

// MODEL_CAPACITY_EXHAUSTED 全局去重:重试全部失败后的 cooldown 时间
antigravityModelCapacityCooldown = 10 * time.Second

// antigravityMaxRateLimitCooldown 是上游 retryDelay / reset 时间戳的本地上限。
// 部分场景下上游会返回小时甚至天级的等待时间(地理限制、配额耗尽等),
// 直接采用会让账号长时间不可调度;统一截断到此值,与 OpenAI 路径
// (maxRateLimit429CooldownSeconds=7200) 保持同一量级。
antigravityMaxRateLimitCooldown = 2 * time.Hour
)

// antigravityPassthroughErrorMessages 透传给客户端的错误消息白名单(小写)
Expand Down Expand Up @@ -2579,6 +2585,25 @@ func antigravityFallbackCooldownSeconds() (time.Duration, bool) {
return time.Duration(seconds) * time.Second, true
}

// clampAntigravityRateLimitDuration 将上游解析得到的 retryDelay 截断到本地上限。
// 大于上限时返回上限并返回 true(调用方可记录日志)。
func clampAntigravityRateLimitDuration(d time.Duration) (time.Duration, bool) {
if d > antigravityMaxRateLimitCooldown {
return antigravityMaxRateLimitCooldown, true
}
return d, false
}

// clampAntigravityRateLimitResetAt 将上游解析得到的绝对重置时间点截断到 now+上限。
// 截断时返回 true(调用方可记录日志)。
func clampAntigravityRateLimitResetAt(resetAt, now time.Time) (time.Time, bool) {
cap := now.Add(antigravityMaxRateLimitCooldown)
if resetAt.After(cap) {
return cap, true
}
return resetAt, false
}

// antigravitySmartRetryInfo 智能重试所需的信息
type antigravitySmartRetryInfo struct {
RetryDelay time.Duration // 重试延迟时间
Expand Down Expand Up @@ -2697,6 +2722,12 @@ func parseAntigravitySmartRetryInfo(body []byte) *antigravitySmartRetryInfo {
retryDelay = antigravityDefaultRateLimitDuration
}

// 防止上游异常返回(如 24h/7d)导致账号长时间不可调度。
if clamped, did := clampAntigravityRateLimitDuration(retryDelay); did {
logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity] retry_delay_clamped original=%v clamped=%v model=%s", retryDelay, clamped, modelName)
retryDelay = clamped
}

return &antigravitySmartRetryInfo{
RetryDelay: retryDelay,
ModelName: modelName,
Expand Down Expand Up @@ -2957,12 +2988,19 @@ func (s *AntigravityGatewayService) getDefaultRateLimitDuration() time.Duration
return defaultDur
}

// resolveResetTime 根据解析的重置时间或默认时长计算重置时间点
// resolveResetTime 根据解析的重置时间或默认时长计算重置时间点。
// 当上游给出的 reset 时间戳超过本地上限时(antigravityMaxRateLimitCooldown),
// 截断到上限以避免账号被长时间锁定。
func (s *AntigravityGatewayService) resolveResetTime(resetAt *int64, defaultDur time.Duration) time.Time {
now := time.Now()
if resetAt != nil {
return time.Unix(*resetAt, 0)
t, did := clampAntigravityRateLimitResetAt(time.Unix(*resetAt, 0), now)
if did {
logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity] reset_at_clamped upstream=%v clamped=%v", time.Unix(*resetAt, 0).Format(time.RFC3339), t.Format(time.RFC3339))
}
return t
}
return time.Now().Add(defaultDur)
return now.Add(defaultDur)
}

type antigravityStreamResult struct {
Expand Down
Loading