Skip to content

fix: sync codex plan type on reset#142

Merged
james-6-23 merged 1 commit into
james-6-23:mainfrom
DearGlory:sync-plan-on-reset
May 24, 2026
Merged

fix: sync codex plan type on reset#142
james-6-23 merged 1 commit into
james-6-23:mainfrom
DearGlory:sync-plan-on-reset

Conversation

@DearGlory
Copy link
Copy Markdown
Contributor

@DearGlory DearGlory commented May 18, 2026

Summary

  • sync x-codex-plan-type from upstream responses and persist any non-empty plan value
  • refresh plan metadata after reset-status and batch-reset-status
  • add regression tests for header sync and reset-triggered plan refresh

Validation

  • go test ./auth ./proxy ./admin -count=1

Summary by CodeRabbit

Release Notes

  • New Features

    • Account reset operations now synchronize Codex plan metadata from upstream headers during the reset process.
    • Plan type information is automatically updated and persisted when synchronizing account state.
  • Tests

    • Added test coverage for account reset with plan metadata synchronization.
    • Added test coverage for Codex usage state synchronization with plan type updates.

Review Change Stack

@coderabbitai
Copy link
Copy Markdown

coderabbitai Bot commented May 18, 2026

📝 Walkthrough

Walkthrough

This 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.

Changes

Codex Plan Type Sync on Account Reset

Layer / File(s) Summary
Sync Callback Hook Infrastructure
admin/handler.go
Handler struct adds syncAccountPlanOnReset callback field, initialized in NewHandler to the default implementation for pluggable testing and behavior customization.
Core Plan Sync Implementation
admin/handler.go
syncAccountPlanAfterReset is the guarded entry point (nil check, timeout, error logging); syncSingleAccountPlanOnReset executes a test request via proxy, extracts Codex plan-type headers, and updates runtime account state.
Plan Type Persistence in Auth Store
auth/store.go
Store.UpdateAccountPlanType normalizes plan strings, updates in-memory account state, recomputes scheduler scheduling, triggers fast refresh, and persists to database (with error logging).
Integration: Reset Endpoints and Proxy Handler
admin/handler.go, proxy/handler.go
ResetAccountStatus and BatchResetStatus call the sync hook after cache clear; SyncCodexUsageState guards UpdateAccountPlanType with nil-store check.
Test Coverage for Plan Sync Flow
admin/handler_test.go, proxy/handler_test.go
Admin tests verify hook invocation and cache clearing during single/batch resets; proxy test confirms in-memory and persisted plan-type updates from response headers.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Poem

🐰 A reset rings, the plans align—
A hook is born, code lines entwine.
From headers tall to stores below,
Codex types sync and swiftly flow.
Tests confirm the journey true,
Now sync hooks make accounts anew! 🎭

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title 'fix: sync codex plan type on reset' directly reflects the main change: adding plan type synchronization logic to the account reset flow.
Docstring Coverage ✅ Passed No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

Copy link
Copy Markdown

@coderabbitai coderabbitai Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

📥 Commits

Reviewing files that changed from the base of the PR and between 18de86a and 6b51f00.

📒 Files selected for processing (5)
  • admin/handler.go
  • admin/handler_test.go
  • auth/store.go
  • proxy/handler.go
  • proxy/handler_test.go

Comment thread admin/handler.go
Comment on lines 2480 to +2482
h.store.ClearCooldown(acc)
acc.ClearUsageCache()
h.syncAccountPlanAfterReset(c.Request.Context(), acc)
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

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.

Comment thread admin/handler.go
Comment on lines +2508 to +2510
if h == nil || h.store == nil || acc == nil || acc.IsOpenAIResponsesAPI() || acc.GetAccessToken() == "" {
return nil
}
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

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.

Comment thread auth/store.go
Comment on lines +3583 to +3591
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)
}
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

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.

Suggested change
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.

@james-6-23 james-6-23 merged commit cfc1d0d into james-6-23:main May 24, 2026
6 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants