feat: add adaptive Codex quota cooldown#4400
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces adaptive Codex cooldowns and non-consuming quota probes that periodically query the ChatGPT WHAM usage endpoint to check rate limits. It also includes resource cleanup improvements in stream handlers. Key feedback highlights a critical spin lock risk in the quota probe loop upon context cancellation, a potential nil pointer dereference when handling stream retry results, and potential deadlocks in test channel communications.
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.
| func (m *Manager) runCodexQuotaProbeLoop(ctx context.Context, wake <-chan struct{}) { | ||
| for { | ||
| now := time.Now() | ||
| due, next := m.collectCodexQuotaProbeBuckets(now) | ||
| if len(due) > 0 { | ||
| m.probeDueCodexQuotaBuckets(ctx, due) | ||
| continue | ||
| } |
There was a problem hiding this comment.
If ctx is canceled while there are due buckets, probeDueCodexQuotaBuckets will return immediately without probing or updating the buckets' NextProbeAt. The loop will then continue and immediately find the same buckets due, calling probeDueCodexQuotaBuckets again in an infinite busy loop, spinning the CPU at 100%. Adding a non-blocking check for ctx.Done() at the start of the loop prevents this spin lock.
func (m *Manager) runCodexQuotaProbeLoop(ctx context.Context, wake <-chan struct{}) {
for {
select {
case <-ctx.Done():
return
default:
}
now := time.Now()
due, next := m.collectCodexQuotaProbeBuckets(now)
if len(due) > 0 {
m.probeDueCodexQuotaBuckets(ctx, due)
continue
}| retryResult, retryErr := h.AuthManager.ExecuteStream(ctx, providers, req, opts) | ||
| if retryErr != nil { | ||
| pendingChunks[0].Err = enrichAuthSelectionError(retryErr, providers, normalizedModel) | ||
| break | ||
| } | ||
| rawStreamHeaders = cloneHeader(retryResult.Headers) |
There was a problem hiding this comment.
If retryErr is nil but retryResult is also nil, dereferencing retryResult.Headers will cause a nil pointer dereference panic. Adding a nil check for retryResult ensures robust error handling.
| retryResult, retryErr := h.AuthManager.ExecuteStream(ctx, providers, req, opts) | |
| if retryErr != nil { | |
| pendingChunks[0].Err = enrichAuthSelectionError(retryErr, providers, normalizedModel) | |
| break | |
| } | |
| rawStreamHeaders = cloneHeader(retryResult.Headers) | |
| retryResult, retryErr := h.AuthManager.ExecuteStream(ctx, providers, req, opts) | |
| if retryErr != nil { | |
| pendingChunks[0].Err = enrichAuthSelectionError(retryErr, providers, normalizedModel) | |
| break | |
| } | |
| if retryResult == nil { | |
| pendingChunks[0].Err = &interfaces.ErrorMessage{StatusCode: http.StatusInternalServerError, Error: fmt.Errorf("empty stream result")} | |
| break | |
| } | |
| rawStreamHeaders = cloneHeader(retryResult.Headers) |
| func (s *stubStore) PersistConfig(context.Context) error { | ||
| atomic.AddInt32(&s.cfgPersisted, 1) | ||
| if s.cfgDone != nil { | ||
| s.cfgDone <- struct{}{} | ||
| } | ||
| return nil | ||
| } |
There was a problem hiding this comment.
If PersistConfig is called multiple times concurrently or sequentially in tests, sending to s.cfgDone can block indefinitely if the channel buffer is full and no receiver is active. Using a non-blocking select send prevents potential deadlocks in tests.
func (s *stubStore) PersistConfig(context.Context) error {
atomic.AddInt32(&s.cfgPersisted, 1)
if s.cfgDone != nil {
select {
case s.cfgDone <- struct{}{}:
default:
}
}
return nil
}| func (s *stubStore) PersistAuthFiles(_ context.Context, message string, paths ...string) error { | ||
| atomic.AddInt32(&s.authPersisted, 1) | ||
| s.mu.Lock() | ||
| s.lastAuthMessage = message | ||
| s.lastAuthPaths = paths | ||
| s.lastAuthPaths = append([]string(nil), paths...) | ||
| s.mu.Unlock() | ||
| if s.authDone != nil { | ||
| s.authDone <- struct{}{} | ||
| } | ||
| return nil | ||
| } |
There was a problem hiding this comment.
Similarly to PersistConfig, sending to s.authDone can block indefinitely if called multiple times. Using a non-blocking select send ensures the test stub remains robust and deadlock-free.
func (s *stubStore) PersistAuthFiles(_ context.Context, message string, paths ...string) error {
atomic.AddInt32(&s.authPersisted, 1)
s.mu.Lock()
s.lastAuthMessage = message
s.lastAuthPaths = append([]string(nil), paths...)
s.mu.Unlock()
if s.authDone != nil {
select {
case s.authDone <- struct{}{}:
default:
}
}
return nil
}
|
Closing this duplicate. The first create request succeeded as draft PR #4399 and was automatically retargeted to dev by the repository workflow. |
1 participant
Summary
Safety
Verification
go test -count=1 ./..., and server buildgo vet ./...go test -race -count=1 -p 2 ./...d0e328e020d1cbe6b4b508919293fc4f75ed2e9158d45b12d52e8d239af1ba28f70e47ea2e11158d2033ee4893a1aa4494dd3e5022a3503e93c5a7d909938cea