feat: add adaptive Codex quota cooldown#4399
Conversation
|
This pull request targeted The base branch has been automatically changed to |
There was a problem hiding this comment.
Code Review
This pull request introduces adaptive Codex cooldowns and a non-consuming quota-status probe loop that queries the WHAM usage endpoint to check for limit resets. It also refactors several tests to eliminate race conditions and arbitrary sleeps, fixes context cancellation ownership in stream callbacks, and updates FileBodySource.WriteTo to report written bytes. The review feedback highlights a potential infinite loop and log spam in the quota probe loop if ResetQuota fails, suggesting a backoff schedule for failed resets.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| if outcome.recovered && outcome.err == nil { | ||
| ids := m.codexQuotaAuthIDsForBucket(bucketKey) | ||
| for _, authID := range ids { | ||
| if _, _, errReset := m.ResetQuota(ctx, authID); errReset != nil { | ||
| log.Warnf("failed to reset recovered Codex quota state for %s: %v", authID, errReset) | ||
| } | ||
| } | ||
| return | ||
| } |
There was a problem hiding this comment.
If ResetQuota fails, the quota state remains exceeded and NextProbeAt is not updated. In the next iteration of the probe loop, this auth will be collected as due again immediately, leading to an infinite loop of rapid probes and log spam. To prevent this, we should back off the probe schedule for any auths that failed to reset.
if outcome.recovered && outcome.err == nil {
ids := m.codexQuotaAuthIDsForBucket(bucketKey)
var failedIDs []string
for _, authID := range ids {
if _, _, errReset := m.ResetQuota(ctx, authID); errReset != nil {
log.Warnf("failed to reset recovered Codex quota state for %s: %v", authID, errReset)
failedIDs = append(failedIDs, authID)
}
}
if len(failedIDs) > 0 {
m.mu.Lock()
for _, authID := range failedIDs {
if auth := m.auths[authID]; auth != nil {
auth.Quota.NextProbeAt = now.Add(quotaProbeNearInterval)
for _, state := range auth.ModelStates {
if state != nil {
state.Quota.NextProbeAt = now.Add(quotaProbeNearInterval)
}
}
}
}
m.mu.Unlock()
}
return
}
Summary
Safety
Verification
go test -count=1 ./..., and server buildgo vet ./...go test -race -count=1 -p 2 ./...d0e328e020d1cbe6b4b508919293fc4f75ed2e9158d45b12d52e8d239af1ba28f70e47ea2e11158d2033ee4893a1aa4494dd3e5022a3503e93c5a7d909938cea