feat: add daily token limit for channels#5043
Open
Ylsssq926 wants to merge 4 commits into
Open
Conversation
Allow administrators to configure a per-channel daily token cap
(stored in Channel.Setting JSON). When a channel reaches its
configured daily_token_limit, it is filtered out of routing
candidates until usage resets at the server local midnight via
Redis TTL.
- DTO: add DailyTokenLimit to dto.ChannelSettings (no schema migration).
- Counter: Redis-backed channel_daily_tokens:{id}:{yyyymmdd} with
atomic INCRBY+EXPIREAT, falls back to in-process counter when
Redis is disabled (single-instance only).
- Routing: GetRandomSatisfiedChannel and ability.GetChannel skip
channels at their daily cap. Affinity hits also fall through to
the standard selection path. Token-specified channels return 429
when their cap is reached.
- Counting: text/audio/realtime success paths increment the daily
counter using prompt+completion totals (raw tokens, not quota).
- UI: web/default channel edit drawer exposes a 'Daily Token Limit'
input under Extra Settings.
Closes QuantumNous#2489
- Remove unused getChannelQuery / getPriority in model/ability.go after the GetChannel rewrite no longer needed them. - Fold service.IsChannelDailyTokenAvailable into a thin wrapper that delegates to model.IsChannelDailyTokenAvailable, so the selection-time logic only has one source of truth. - Drop service.GetDailyTokenUsage / GetDailyTokenLimit which had no callers in the first PR; they can come back when an admin endpoint exposing usage is added.
- Localise the daily-limit error path in distributor.go via a new i18n key MsgDistributorChannelDailyTokenLimitExceeded so the routing layer stays consistent with the rest of the project. - Make the affinity miss caused by daily-limit observable by emitting a structured log line and clearing the cached affinity for the request, so the affinity state machine no longer silently swallows a meaningful event. - Fall back to the in-process counter when Redis writes fail in IncreaseChannelDailyTokenUsage, and emit an explicit warning so the daily limit cannot be silently bypassed by transient Redis outages. - Align WSS / audio counters with the text path by using PromptTokens + CompletionTokens (or InputTokens + OutputTokens for WSS), so daily limit semantics no longer drift between channel types.
This commit is a drive-by cleanup found while making `go vet ./...`
output clean for the daily-token-limit work; reviewers can drop
this commit if they prefer to handle these in a separate PR.
- common/custom-event.go: drop the value-receiver Mutex that vet
flags as a lock-by-value when CustomEvent is passed around.
- relay/channel/{baidu,cloudflare,cohere,dify,jina,mistral,mokaai,
palm,tencent,xunfei,zhipu}/adaptor.go: remove unreachable returns
after `panic("implement me")` stubs that vet's unreachable-code
check rejects.
No behavioural change.
Contributor
WalkthroughThis PR implements per-channel daily token usage limits with Redis-backed tracking and in-process fallback. The middleware enforces limits during channel selection, quota services record usage, and channel selection filters by availability. UI forms support configuration, and relay adapters are hardened to panic on unimplemented Claude conversions. ChangesPer-Channel Daily Token Limit Feature
CustomEvent Mutex Removal
Claude Conversion Panic Enforcement
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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 |
![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: 1
🤖 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 `@service/text_quota.go`:
- Around line 371-374: The code currently builds dailyTokenUsage from
summary.PromptTokens + summary.CompletionTokens (which may be normalized) before
calling IncreaseChannelDailyTokenUsage; change it to use the raw token total
(e.g. summary.TotalTokens or the original raw usage totals variable) so the
daily counter reflects unnormalized usage—update the construction of
dailyTokenUsage and the call to IncreaseChannelDailyTokenUsage to pass the raw
total instead of the normalized prompt/completion sum.
🪄 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: 3dc27095-d044-4bee-9ef4-acec6e403de3
📒 Files selected for processing (34)
common/channel_daily_limit.gocommon/custom-event.gocommon/redis.goconstant/cache_key.godto/channel_settings.goi18n/keys.goi18n/locales/en.yamli18n/locales/zh-CN.yamli18n/locales/zh-TW.yamlmiddleware/distributor.gomodel/ability.gomodel/channel_cache.gomodel/channel_daily_limit.gorelay/channel/baidu/adaptor.gorelay/channel/cloudflare/adaptor.gorelay/channel/cohere/adaptor.gorelay/channel/dify/adaptor.gorelay/channel/jina/adaptor.gorelay/channel/mistral/adaptor.gorelay/channel/mokaai/adaptor.gorelay/channel/palm/adaptor.gorelay/channel/tencent/adaptor.gorelay/channel/xunfei/adaptor.gorelay/channel/zhipu/adaptor.goservice/channel_affinity.goservice/channel_daily_limit.goservice/quota.goservice/text_quota.gotypes/error.goweb/default/src/features/channels/components/drawers/channel-mutate-drawer.tsxweb/default/src/features/channels/lib/channel-form.tsweb/default/src/features/channels/types.tsweb/default/src/i18n/locales/en.jsonweb/default/src/i18n/locales/zh.json
💤 Files with no reviewable changes (12)
- relay/channel/jina/adaptor.go
- relay/channel/tencent/adaptor.go
- relay/channel/baidu/adaptor.go
- relay/channel/xunfei/adaptor.go
- relay/channel/zhipu/adaptor.go
- relay/channel/cohere/adaptor.go
- relay/channel/mokaai/adaptor.go
- common/custom-event.go
- relay/channel/dify/adaptor.go
- relay/channel/palm/adaptor.go
- relay/channel/cloudflare/adaptor.go
- relay/channel/mistral/adaptor.go
Comment on lines
+371
to
+374
| dailyTokenUsage := int64(summary.PromptTokens + summary.CompletionTokens) | ||
| if err := IncreaseChannelDailyTokenUsage(relayInfo.ChannelId, dailyTokenUsage); err != nil { | ||
| logger.LogError(ctx, "channel daily token usage NOT recorded, daily limit may not take effect: "+err.Error()) | ||
| } |
Contributor
There was a problem hiding this comment.
Use raw token totals when incrementing daily usage.
This path uses summary.PromptTokens, which can be normalized (reduced) for specific billing semantics, causing daily usage undercount and delayed limit enforcement. Use raw totals (summary.TotalTokens or raw usage totals) for the daily counter.
💡 Suggested fix
- dailyTokenUsage := int64(summary.PromptTokens + summary.CompletionTokens)
+ dailyTokenUsage := int64(summary.TotalTokens)
if err := IncreaseChannelDailyTokenUsage(relayInfo.ChannelId, dailyTokenUsage); err != nil {
logger.LogError(ctx, "channel daily token usage NOT recorded, daily limit may not take effect: "+err.Error())
}📝 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
| dailyTokenUsage := int64(summary.PromptTokens + summary.CompletionTokens) | |
| if err := IncreaseChannelDailyTokenUsage(relayInfo.ChannelId, dailyTokenUsage); err != nil { | |
| logger.LogError(ctx, "channel daily token usage NOT recorded, daily limit may not take effect: "+err.Error()) | |
| } | |
| dailyTokenUsage := int64(summary.TotalTokens) | |
| if err := IncreaseChannelDailyTokenUsage(relayInfo.ChannelId, dailyTokenUsage); err != nil { | |
| logger.LogError(ctx, "channel daily token usage NOT recorded, daily limit may not take effect: "+err.Error()) | |
| } |
🤖 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/text_quota.go` around lines 371 - 374, The code currently builds
dailyTokenUsage from summary.PromptTokens + summary.CompletionTokens (which may
be normalized) before calling IncreaseChannelDailyTokenUsage; change it to use
the raw token total (e.g. summary.TotalTokens or the original raw usage totals
variable) so the daily counter reflects unnormalized usage—update the
construction of dailyTokenUsage and the call to IncreaseChannelDailyTokenUsage
to pass the raw total instead of the normalized prompt/completion sum.
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.
What
Adds a per-channel daily token limit: administrators can configure a daily cap (in raw tokens) on each channel via the existing
Channel.SettingJSON. When a channel reaches its configured cap, it is filtered out of routing candidates until usage automatically resets at the server-local midnight.Closes #2489
Why
OpenAI and several other upstream services hand out a daily free token allowance. Without an in-product cap, operators have no way to "spend the free tier first" - once daily traffic exceeds the free quota the same channel just keeps getting billed. The issue asked for a way to express "channel A: at most 1M tokens per day" without having to manually disable/enable the channel each night.
Design
Storage
daily_token_limitondto.ChannelSettings(which is whatChannel.Settingalready serialises to). 0 means unlimited.omitempty).Counter
channel_daily_tokens:{channel_id}:{yyyymmdd}, TTL set to expire at the next server-local midnight via an atomic LuaINCRBY + EXPIREATscript (added ascommon.RedisIncrByWithExpireAt).SysLogwarning that this is not multi-instance consistent.Routing
model.GetRandomSatisfiedChannelandmodel.GetChannel(DB path, used when memory cache is disabled) skip channels that have hit their daily cap.middleware/distributor.go, an affinity hit on a capped channel falls through to the standard selection path (and the affinity cache for that request is cleared, so the affinity state machine no longer silently swallows the event).429 Too Many Requests(i18n keyMsgDistributorChannelDailyTokenLimitExceeded, en/zh-CN/zh-TW translations included).Channel.Statusorabilities.enabledfor daily limits - those are persistent admin states; daily limits are a soft, time-boxed filter.Counting
service/text_quota.go): increments bysummary.PromptTokens + summary.CompletionTokenson success.service/quota.go): increments byusage.InputTokens + usage.OutputTokens.service/quota.go): increments byusage.PromptTokens + usage.CompletionTokens.UI
In
web/default's channel edit drawer, "Channel Extra Settings" gets a new "Daily Token Limit" number input (label, description, hint about midnight reset, all i18n'd inen.json/zh.json). The form schema, defaults, parser, andbuildSettingJSONare all updated together.Out of scope (deliberate)
service/task_billing.goand I'd rather land that as a follow-up than balloon this PR.Verification
go build ./...andgo vet ./...both pass cleanly on the branch.daily_token_limit = 1000. Issue a request that consumes ~800 tokens - should succeed.429if the request was pinned to this channel via token binding.Commits
feat: add daily token limit for channels- main changefix(channel-daily-limit): cleanup unused helpers after reviewfix(channel-daily-limit): address review feedback- i18n, Redis fallback, affinity log, token alignmentchore(vet): clean up pre-existing go vet warnings- drive-by; feel free to drop this commit if you'd prefer a separate PRThe fourth commit is genuinely a drive-by -
go vet ./...already had a few unreachablereturnstatements afterpanicstubs and aMutexcopied via value receiver. I cleaned them up so vet stays quiet on this branch, but I'm happy to drop that commit and re-submit it as its own PR if you'd prefer.Risks
SysLogwarning makes this visible at startup; documentation could call it out more loudly if you'd like.Happy to iterate on any of the above.
Summary by CodeRabbit
Release Notes
New Features
Improvements