fix: sync codex plan type on reset#142
Conversation
📝 WalkthroughWalkthroughThis PR integrates Codex plan-type synchronization into the account reset flow. When accounts are reset, a new callback hook now triggers a test request to retrieve upstream plan metadata, persists the plan-type to the store and database, and updates runtime state. The proxy handler guards against nil-store crashes. Admin and proxy handler tests validate the feature. ChangesCodex Plan Type Sync on Account Reset
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@admin/handler.go`:
- Around line 2508-2510: The early return skips plan sync for RT-backed accounts
with empty AccessToken; change the logic so that when acc.GetAccessToken() == ""
you first attempt to refresh the account (using the account/store refresh API
available in your codebase—e.g. call h.store.RefreshAccount(ctx, acc) or
acc.Refresh(ctx) if such methods exist), handle errors, and only return if the
refresh fails or the token remains empty; keep the existing guards for h,
h.store, and acc.IsOpenAIResponsesAPI() but run the probe after a successful
refresh so RT-backed accounts get their plan_type updated.
In `@auth/store.go`:
- Around line 3583-3591: The code currently returns early when changed==false
and thus skips persisting plan to DB; change the logic so we only skip when s.db
== nil (keep the nil guard) but still call s.db.UpdateCredentials(ctx, acc.DBID,
map[string]interface{}{"plan_type": plan}) even if changed==false so a repeated
upstream header can repair a blank/stale plan_type; locate the block using
symbols s.db, UpdateCredentials, acc.DBID, plan and move/remove the changed
check so the DB update runs whenever s.db != nil, preserving the context timeout
and error logging.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 027277f1-bde0-4e00-ae3a-23f3140a1d83
📒 Files selected for processing (5)
admin/handler.goadmin/handler_test.goauth/store.goproxy/handler.goproxy/handler_test.go
| h.store.ClearCooldown(acc) | ||
| acc.ClearUsageCache() | ||
| h.syncAccountPlanAfterReset(c.Request.Context(), acc) |
There was a problem hiding this comment.
Batch reset now serially waits on one upstream probe per account.
syncAccountPlanAfterReset can block for up to 15 seconds, and this loop invokes it inline for every ID. That makes batch-reset-status latency scale with account count, so a modest batch can exceed proxy/request timeouts even though the actual reset work already succeeded. This should be decoupled from the response path or run through a bounded async worker pool.
| if h == nil || h.store == nil || acc == nil || acc.IsOpenAIResponsesAPI() || acc.GetAccessToken() == "" { | ||
| return nil | ||
| } |
There was a problem hiding this comment.
Reset-time plan sync is skipped for accounts that most need it.
The early return on empty AccessToken makes this feature a no-op for RT-backed accounts whose AT has expired or been cleared. Those are common reset candidates, so plan_type stays stale until some unrelated refresh/request happens later. Refresh the account first when it still has refresh credentials, then run the probe.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@admin/handler.go` around lines 2508 - 2510, The early return skips plan sync
for RT-backed accounts with empty AccessToken; change the logic so that when
acc.GetAccessToken() == "" you first attempt to refresh the account (using the
account/store refresh API available in your codebase—e.g. call
h.store.RefreshAccount(ctx, acc) or acc.Refresh(ctx) if such methods exist),
handle errors, and only return if the refresh fails or the token remains empty;
keep the existing guards for h, h.store, and acc.IsOpenAIResponsesAPI() but run
the probe after a successful refresh so RT-backed accounts get their plan_type
updated.
| if s.db == nil || !changed { | ||
| return changed | ||
| } | ||
|
|
||
| ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) | ||
| defer cancel() | ||
| if err := s.db.UpdateCredentials(ctx, acc.DBID, map[string]interface{}{"plan_type": plan}); err != nil { | ||
| log.Printf("[璐﹀彿 %d] 鎸佷箙鍖?plan_type 澶辫触: %v", acc.DBID, err) | ||
| } |
There was a problem hiding this comment.
Persist the observed plan even when memory already matches it.
This method currently skips UpdateCredentials when changed == false, so a repeated upstream header cannot repair a blank/stale plan_type in the database after a previous write failure or any other memory-only update. That means reset-triggered sync can still be lost on restart.
Suggested fix
- if s.db == nil || !changed {
+ if s.db == nil {
return changed
}
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
defer cancel()
if err := s.db.UpdateCredentials(ctx, acc.DBID, map[string]interface{}{"plan_type": plan}); err != nil {
log.Printf("[璐﹀彿 %d] 鎸佷箙鍖?plan_type 澶辫触: %v", acc.DBID, err)
}
return changed📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if s.db == nil || !changed { | |
| return changed | |
| } | |
| ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) | |
| defer cancel() | |
| if err := s.db.UpdateCredentials(ctx, acc.DBID, map[string]interface{}{"plan_type": plan}); err != nil { | |
| log.Printf("[璐﹀彿 %d] 鎸佷箙鍖?plan_type 澶辫触: %v", acc.DBID, err) | |
| } | |
| if s.db == nil { | |
| return changed | |
| } | |
| ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) | |
| defer cancel() | |
| if err := s.db.UpdateCredentials(ctx, acc.DBID, map[string]interface{}{"plan_type": plan}); err != nil { | |
| log.Printf("[璐﹀彿 %d] 鎸佷箙鍖?plan_type 澶辫触: %v", acc.DBID, err) | |
| } | |
| return changed |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@auth/store.go` around lines 3583 - 3591, The code currently returns early
when changed==false and thus skips persisting plan to DB; change the logic so we
only skip when s.db == nil (keep the nil guard) but still call
s.db.UpdateCredentials(ctx, acc.DBID, map[string]interface{}{"plan_type": plan})
even if changed==false so a repeated upstream header can repair a blank/stale
plan_type; locate the block using symbols s.db, UpdateCredentials, acc.DBID,
plan and move/remove the changed check so the DB update runs whenever s.db !=
nil, preserving the context timeout and error logging.
Summary
x-codex-plan-typefrom upstream responses and persist any non-empty plan valuereset-statusandbatch-reset-statusValidation
go test ./auth ./proxy ./admin -count=1Summary by CodeRabbit
Release Notes
New Features
Tests