feat: add periodic quota policies for API keys#5791
Open
zhukangfeng wants to merge 2 commits into
Open
Conversation
Add optional API key periodic quota policies with preset/custom windows, anchor time, strict or graceful boundary handling, reject-only or temporary-disable actions, auto recovery, and manual reset support. Integrate policy accounting into token pre-consume/post-consume billing paths while preserving existing user/token quota behavior, returning explicit token_quota_policy_exhausted errors for periodic quota exhaustion, and caching policy snapshots per relay request to avoid extra reads for normal tokens. Expose policy configuration and temporary-disabled status in the default UI with localized labels/tooltips/log rendering across supported locales, add common usage log descriptors, and synchronize the backend OpenAPI token schema and reset route. Add model/service/controller/E2E tests, PR screenshots, and verify with go test ./..., web/default bun run typecheck, web/default bun run build, and git diff --check.
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: ⛔ Files ignored due to path filters (2)
📒 Files selected for processing (47)
✅ Files skipped from review due to trivial changes (7)
🚧 Files skipped from review as they are similar to previous changes (40)
WalkthroughAdds token quota policy storage, enforcement, reset, API, frontend, localization, and tests across the stack. ChangesToken periodic quota policy
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 OpenGrep (1.23.0)service/billing_session.go┌──────────────┐ �[32m✔�[39m �[1mOpengrep OSS�[0m [00.20][ERROR]: unable to find a config; path service/quota.go┌──────────────┐ �[32m✔�[39m �[1mOpengrep OSS�[0m [00.36][ERROR]: unable to find a config; path service/pre_consume_quota.go┌──────────────┐ �[32m✔�[39m �[1mOpengrep OSS�[0m [01.16][ERROR]: unable to find a config; path
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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 9
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
controller/token.go (1)
264-273: 🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy liftMake token and quota-policy writes atomic.
AddTokeninsertscleanTokenbefore saving the policy, andUpdateTokenupdates the base token before saving the policy. IfSaveTokenQuotaPolicyForTokenfails, the API returns an error after partially committing token changes. Wrap each base-token write and policy save in one database transaction or model-layer helper that rolls back on policy failure.Also applies to: 354-369
🤖 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 `@controller/token.go` around lines 264 - 273, `AddToken` and `UpdateToken` are committing the base token write before `SaveTokenQuotaPolicyForToken`, so a policy failure can leave partial token state; make the token write and quota-policy save atomic by using one transaction or a model-layer helper that rolls back on failure. Update the logic in `AddToken` and `UpdateToken` (and the shared token/quota-policy write path they call) so `cleanToken.Insert`/the token update and `model.SaveTokenQuotaPolicyForToken` succeed or fail together, with errors propagated only after the transaction is aborted.service/billing_session.go (1)
50-62: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winBlock refunds after a delta-zero settle reaches policy finalization.
When
delta == 0,fundingSettledstays false. IfsettleTokenQuotaPolicyfails here,Refund()can still refundtokenConsumed/funding vianeedsRefundLocked, even thoughSettleis finalizing a completed request. Mark the funding side as committed before returning this settlement error.Suggested guard
tokenErr = settleTokenQuotaPolicy(s.relayInfo, actualQuota) if tokenErr != nil { + s.fundingSettled = true return tokenErr }🤖 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 `@service/billing_session.go` around lines 50 - 62, `Refund()` can still issue a refund after a zero-delta settlement fails because `fundingSettled` remains false when `s.funding.Settle(delta)` is skipped. Update the settlement flow in `BillingSession` so that once `settleTokenQuotaPolicy(s.relayInfo, actualQuota)` fails after a completed request, the funding side is marked committed before returning the error, preventing `needsRefundLocked` from treating it as refundable; use the existing `fundingSettled` flag and the `Settle`/`settleTokenQuotaPolicy` path to place this guard.
🧹 Nitpick comments (1)
web/default/src/features/keys/components/api-keys-cells.tsx (1)
45-51: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMatch the local prop-typing convention here.
Use an explicit props interface/type alias and access
props.apiKey/props.classNamedirectly instead of inline destructuring so this new component stays consistent with theweb/defaultcomponent pattern.As per coding guidelines, "Do not destructure objects unless necessary, especially component props; prefer direct property access such as props.xxx" and "Use function components and Hooks, and give component props an explicit interface or type alias."
🤖 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 `@web/default/src/features/keys/components/api-keys-cells.tsx` around lines 45 - 51, The ApiKeyStatusBadge component is using inline prop destructuring instead of the local web/default props convention. Introduce an explicit props interface or type alias for ApiKeyStatusBadge, accept a single props argument, and update the component to read props.apiKey and props.className directly so it matches the existing component pattern and coding guidelines.Source: Coding guidelines
🤖 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 `@middleware/auth.go`:
- Around line 373-378: The quota exhaustion branch in auth.go uses
MsgTokenQuotaPolicyExhaustedPending for all non-auto-resume cases, which reuses
the wrong locale copy for manual-reset policies. Update the message selection
logic around the quotaPolicyErr.AutoResume check to choose a
manual-reset-specific i18n key when AutoResume is false, and keep
MsgTokenQuotaPolicyExhausted only for auto-resume exhaustion with a future
NextResetAt. Make the change in the quota handling flow that sets messageKey and
resetTime so the returned text matches the manual-resume state.
In `@model/token_quota_policy.go`:
- Around line 223-240: The quota exhaustion check in TokenQuotaPolicy update
flow is using a strict greater-than comparison, which leaves exact quota usage
unmarked as exhausted. Update the logic in the token quota update path and the
final return check in the TokenQuotaPolicy methods so they consistently treat
used_quota >= quota as exhausted, matching the rest of the model and ensuring
disable_token is triggered immediately when quota is exactly reached.
- Around line 438-444: The TokenQuotaPolicy exhaustion update is overwriting
exhausted_token_status on repeated calls, which loses the originally enabled
status. Update the logic in the exhaustion path that uses token.Status so it
only sets exhausted_token_status the first time a token is marked exhausted, and
preserves the existing value on subsequent calls; use the TokenQuotaPolicy
update block as the place to guard this before calling Updates.
In `@service/quota.go`:
- Around line 439-448: The quota finalization in the non-playground branch can
fail after wallet/subscription funding has already been updated, leaving
accounting inconsistent. In the `settleTokenQuotaPolicy` and
`postConsumeTokenQuotaPolicy` path inside `quota.go`, either move policy
settlement before the funding mutation or add a compensating rollback that
restores the funding delta if either call returns an error. Keep the fix
localized around the settlement flow so the state changes in `relayInfo`/quota
remain atomic.
In `@service/token_quota_policy.go`:
- Around line 124-141: The strict pre-check path in token quota handling returns
TokenQuotaPolicyExhausted without honoring the disable-token action. Update the
logic in the token quota policy flow around tokenQuotaPolicySnapshotExhausted
and MarkTokenQuotaPolicyExhausted so that when ExhaustedAction is
TokenQuotaExhaustDisableToken, the token is marked disabled before returning the
exhaustion error, matching the non-strict branch behavior and the existing
model.ConsumeTokenQuotaPolicy handling.
In `@web/default/src/features/keys/components/api-keys-columns.tsx`:
- Around line 109-114: The Status column filtering in api-keys-columns.tsx is
using the raw status field, which no longer matches what ApiKeyStatusBadge
displays for quota-derived temporary states. Update the Status column’s filterFn
to filter by the same derived display state used by ApiKeyStatusBadge for each
row, so temporarily quota-disabled keys are excluded from the “Enabled” filter
and can be matched by their visible label.
In `@web/default/src/features/keys/lib/api-key-form.ts`:
- Line 126: The default for quota_policy_anchor_time is being captured at module
load, so the API key form can reuse a stale timestamp on later opens. Move the
Date creation out of API_KEY_FORM_DEFAULT_VALUES and compute it lazily inside
getApiKeyFormDefaultValues so each form open gets a fresh value; keep
transformFormDataToPayload unchanged since it will continue forwarding any
non-nil field.
In `@web/default/src/features/keys/lib/quota-policy-status.ts`:
- Around line 27-35: The quota status helper is treating every exhausted
periodic quota as a temporary disable state, which breaks reject-only exhaustion
handling. Update the logic in the quota status helper to distinguish real
temporary-disabled policies from reject-only exhaustion by checking the policy’s
exhaustion action or persisted temporary-disabled flag before returning an
active/recovering disable state. Use the existing helper/logic around the
quota-policy status computation to gate the “Temporarily disabled” result only
when the policy was actually paused.
In `@web/default/src/features/usage-logs/lib/format.ts`:
- Around line 423-428: The sentinel handling in formatSystemLogTimestamp should
treat numeric 0 and negative values the same way as formatTimestampToDate so
they render as “-” instead of raw numbers. Update the guard in
formatSystemLogTimestamp to return “-” for numeric timestamps that are less than
or equal to 0, while keeping the String(value) fallback only for non-numeric
input; use formatSystemLogTimestamp and formatTimestampToDate as the key
reference points.
---
Outside diff comments:
In `@controller/token.go`:
- Around line 264-273: `AddToken` and `UpdateToken` are committing the base
token write before `SaveTokenQuotaPolicyForToken`, so a policy failure can leave
partial token state; make the token write and quota-policy save atomic by using
one transaction or a model-layer helper that rolls back on failure. Update the
logic in `AddToken` and `UpdateToken` (and the shared token/quota-policy write
path they call) so `cleanToken.Insert`/the token update and
`model.SaveTokenQuotaPolicyForToken` succeed or fail together, with errors
propagated only after the transaction is aborted.
In `@service/billing_session.go`:
- Around line 50-62: `Refund()` can still issue a refund after a zero-delta
settlement fails because `fundingSettled` remains false when
`s.funding.Settle(delta)` is skipped. Update the settlement flow in
`BillingSession` so that once `settleTokenQuotaPolicy(s.relayInfo, actualQuota)`
fails after a completed request, the funding side is marked committed before
returning the error, preventing `needsRefundLocked` from treating it as
refundable; use the existing `fundingSettled` flag and the
`Settle`/`settleTokenQuotaPolicy` path to place this guard.
---
Nitpick comments:
In `@web/default/src/features/keys/components/api-keys-cells.tsx`:
- Around line 45-51: The ApiKeyStatusBadge component is using inline prop
destructuring instead of the local web/default props convention. Introduce an
explicit props interface or type alias for ApiKeyStatusBadge, accept a single
props argument, and update the component to read props.apiKey and
props.className directly so it matches the existing component pattern and coding
guidelines.
🪄 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: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: c2487986-ae7d-4364-ac85-62329a702501
⛔ Files ignored due to path filters (2)
docs/pr-screenshots/token-periodic-quota/01-api-keys-page.pngis excluded by!**/*.pngdocs/pr-screenshots/token-periodic-quota/02-api-key-periodic-quota-drawer.pngis excluded by!**/*.png
📒 Files selected for processing (46)
constant/context_key.gocontroller/token.gocontroller/token_periodic_quota_e2e_test.gocontroller/token_test.godocs/openapi/api.jsoni18n/keys.goi18n/locales/en.yamli18n/locales/zh-CN.yamli18n/locales/zh-TW.yamlmiddleware/auth.gomodel/log.gomodel/main.gomodel/system_task.gomodel/task_cas_test.gomodel/token.gomodel/token_quota_policy.gomodel/token_quota_policy_test.gorelay/common/relay_info.gorouter/api-router.goservice/billing_session.goservice/pre_consume_quota.goservice/quota.goservice/task_billing_test.goservice/token_quota_policy.goservice/token_quota_policy_test.gotypes/error.goweb/default/src/features/keys/api.tsweb/default/src/features/keys/components/api-keys-cells.tsxweb/default/src/features/keys/components/api-keys-columns.tsxweb/default/src/features/keys/components/api-keys-mutate-drawer.tsxweb/default/src/features/keys/components/api-keys-table.tsxweb/default/src/features/keys/components/data-table-row-actions.tsxweb/default/src/features/keys/lib/api-key-form.tsweb/default/src/features/keys/lib/quota-policy-status.tsweb/default/src/features/keys/types.tsweb/default/src/features/usage-logs/components/columns/common-logs-columns.tsxweb/default/src/features/usage-logs/components/dialogs/details-dialog.tsxweb/default/src/features/usage-logs/lib/format.tsweb/default/src/i18n/locales/en.jsonweb/default/src/i18n/locales/fr.jsonweb/default/src/i18n/locales/ja.jsonweb/default/src/i18n/locales/ru.jsonweb/default/src/i18n/locales/vi.jsonweb/default/src/i18n/locales/zh.jsonweb/default/src/i18n/static-keys.tsweb/default/src/lib/format.ts
Make token and quota policy saves atomic, preserve original token status during repeated exhaustion, and treat exact quota settlement as exhausted. Rollback funding when post-consume quota policy settlement fails, keep finalized billing sessions from double-refunding, and clarify manual-reset exhaustion errors. Tighten key status UI filtering, lazy form anchor defaults, reject-only disabled state handling, and log timestamp sentinel formatting.
zhukangfeng
force-pushed
the
codex/token-periodic-quota
branch
from
f28142f to
abcf6d5
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
service/quota.go (1)
440-461: 🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy liftThe final token-quota write can still leave accounting half-applied.
Lines 441-454 now compensate policy failures, but lines 455-461 still return after funding and policy usage have already been committed. If
DecreaseTokenQuota/IncreaseTokenQuotafails there, wallet/subscription state and periodic-policy state stay updated while the token’s own quota does not. This path needs the same atomicity guarantee as the earlier rollback case—either a single transactional helper or compensating rollback for both funding and policy on token-quota write failure.🤖 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 `@service/quota.go` around lines 440 - 461, The final token-quota update in the quota flow can still fail after funding and policy changes have already been applied. In the quota handling path in service/quota.go, update the logic around settleTokenQuotaPolicy/postConsumeTokenQuotaPolicy and the final model.DecreaseTokenQuota/model.IncreaseTokenQuota write so the whole sequence is atomic. If the token-quota write fails, perform compensating rollback for both the funding update and the policy state, or move the entire operation into a single transactional helper to keep relayInfo accounting consistent.
🤖 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 `@controller/token.go`:
- Around line 363-370: The update flow in cleanToken.Update() followed by
model.AttachTokenQuotaPolicy(cleanToken) can repopulate cache without the
existing quota policy, causing a race where cached token data may briefly lose
periodic-quota information. Fix the token update path in controller/token.go by
preserving the current quota policy before Update() runs, or by using a
post-write refresh helper that reloads the token from DB and re-caches the fully
populated record. Keep the existing error handling around cleanToken.Update()
and AttachTokenQuotaPolicy(), but ensure the cached token always includes the
prior policy when QuotaPolicy is not being changed.
In `@web/default/src/features/keys/constants.ts`:
- Around line 67-69: The new status option in constants.ts is using a raw
English label, which bypasses translation. Update the status entry in the keys
constants to store an i18n key instead of 'Temporarily disabled', following the
same label/value convention used by the other status options in this feature.
Keep the value unchanged and make sure the label symbol used here matches the
existing i18n key pattern for feature status/option labels so the UI can
translate it properly.
---
Outside diff comments:
In `@service/quota.go`:
- Around line 440-461: The final token-quota update in the quota flow can still
fail after funding and policy changes have already been applied. In the quota
handling path in service/quota.go, update the logic around
settleTokenQuotaPolicy/postConsumeTokenQuotaPolicy and the final
model.DecreaseTokenQuota/model.IncreaseTokenQuota write so the whole sequence is
atomic. If the token-quota write fails, perform compensating rollback for both
the funding update and the policy state, or move the entire operation into a
single transactional helper to keep relayInfo accounting consistent.
🪄 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: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 9bbbebfe-05bc-448e-a225-cba968d79f6d
📒 Files selected for processing (17)
controller/token.goi18n/keys.goi18n/locales/en.yamli18n/locales/zh-CN.yamli18n/locales/zh-TW.yamlmiddleware/auth.gomodel/token_quota_policy.gomodel/token_quota_policy_test.goservice/billing_session.goservice/quota.goservice/token_quota_policy.goweb/default/src/features/keys/components/api-keys-cells.tsxweb/default/src/features/keys/components/api-keys-columns.tsxweb/default/src/features/keys/constants.tsweb/default/src/features/keys/lib/api-key-form.tsweb/default/src/features/keys/lib/quota-policy-status.tsweb/default/src/features/usage-logs/lib/format.ts
✅ Files skipped from review due to trivial changes (2)
- i18n/locales/en.yaml
- i18n/locales/zh-CN.yaml
🚧 Files skipped from review as they are similar to previous changes (11)
- web/default/src/features/keys/components/api-keys-columns.tsx
- web/default/src/features/keys/lib/quota-policy-status.ts
- web/default/src/features/usage-logs/lib/format.ts
- middleware/auth.go
- i18n/keys.go
- web/default/src/features/keys/components/api-keys-cells.tsx
- service/token_quota_policy.go
- service/billing_session.go
- model/token_quota_policy_test.go
- web/default/src/features/keys/lib/api-key-form.ts
- model/token_quota_policy.go
Comment on lines
+67
to
+69
| { | ||
| label: 'Temporarily disabled', | ||
| value: 'quota_policy_disabled', |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Store the new status label as an i18n key, not a raw English literal.
'Temporarily disabled' is user-facing text in a feature constants.ts file, so adding it as a literal here will bypass translation in non-English locales. Use the same i18n-key convention this feature uses for status/option labels instead. As per coding guidelines, frontend UI text must support i18n and feature constants.ts status/option labels must store i18n keys consistently.
🤖 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 `@web/default/src/features/keys/constants.ts` around lines 67 - 69, The new
status option in constants.ts is using a raw English label, which bypasses
translation. Update the status entry in the keys constants to store an i18n key
instead of 'Temporarily disabled', following the same label/value convention
used by the other status options in this feature. Keep the value unchanged and
make sure the label symbol used here matches the existing i18n key pattern for
feature status/option labels so the UI can translate it properly.
Source: Coding guidelines
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Important
📝 变更描述 / Description
This PR adds optional periodic quota policies for API keys. A token can now have an additional budget window on top of the existing user quota and token total quota, so administrators can limit usage by 5-hour, daily, weekly, monthly, or custom-minute cycles.
Backend changes:
TokenQuotaPolicypersistence and migration with period mode, anchor time, quota, used quota, next reset time, boundary mode, exhausted action, and auto-resume metadata.token_quota_policy_exhaustederrors with approximate recovery time wording and no used/quota amounts in client-facing messages.Frontend changes:
Testing and review coverage:
docs/pr-screenshots/token-periodic-quota/.🚀 变更类型 / Type of change
🔗 关联任务 / Related Issue
✅ 提交前检查项 / Checklist
Bug fix，我已提交或关联对应 Issue，且不会将设计取舍、预期不一致或理解偏差直接归类为 bug。📸 运行证明 / Proof of Work
Local verification:
go test ./...Manual/UI proof files included in this PR:
docs/pr-screenshots/token-periodic-quota/01-api-keys-page.pngdocs/pr-screenshots/token-periodic-quota/02-api-key-periodic-quota-drawer.pngAdditional notes:
git log; this contribution was AI-assisted and reviewed locally before PR creation.Summary by CodeRabbit