fix(platform): fix budget scope override logic and enforcement#1392
Conversation
There was a problem hiding this comment.
Your free trial has ended. If you'd like to continue receiving code reviews, you can add a payment method here.
📝 WalkthroughWalkthroughThis PR implements a hierarchical budget governance system that enables scope-based budget rules (user, team, role, organization-wide) with independent limit resolution per scope tier, warning thresholds, and period-based enforcement. It replaces agent-config-derived team targeting with user/team/role-aware budget context resolution across all enforcement points. The UI editor now supports organization-wide budgets, warning threshold configuration, and internationalized labels. Core enforcement logic groups applicable rules by period, resolves effective limits with precedence (user > team > role > default), and checks personal, team aggregate, and organization aggregate usage independently. Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
🚥 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 docstrings
🧪 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 review |
✅ Actions performedReview triggered.
|
Resolve per-user role and team memberships at all checkBudget() call sites via a shared resolveBudgetContext() helper. Implement priority override model (user > team > role > default) with per-field resolution so that more specific scopes override broader ones. Add org-wide aggregate budget cap, daily/weekly period enforcement, warning threshold support, and team aggregate usage checks. Replace .collect() with async iteration per coding standards. Add comprehensive unit tests for rule matching, effective limit resolution, and period key generation. Internationalize all budget editor UI strings across en and de locales.
579b606 to
a74eb36
Compare
There was a problem hiding this comment.
Actionable comments posted: 6
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (7)
services/platform/convex/agent_tools/location/mutations.ts (1)
137-176:⚠️ Potential issue | 🟠 MajorBudget denial leaves the location flow stuck in
generating.This branch sets up the stream and marks the thread as generating before calling
checkBudget(), then throws on denial without resetting the thread back toidle. The UI will keep showing a running generation that never starts.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@services/platform/convex/agent_tools/location/mutations.ts` around lines 137 - 176, The code creates a stream via persistentStreaming.createStream and sets threadMetadata.generationStatus to 'generating' before calling checkBudget, so when checkBudget denies it throws and leaves the thread stuck; fix by moving budget checks to run before creating the stream and before patching threadMetadata (or, if stream must be created early, ensure the catch/denial path resets the metadata and deletes the stream): call resolveBudgetContext and checkBudget prior to persistentStreaming.createStream and the ctx.db.patch that sets generationStatus, or on budget denial call ctx.db.patch(threadMeta._id, { generationStatus: 'idle', streamId: null, updatedAt: Date.now() }) and clean up persistentStreaming (delete/close) to avoid leaking stream IDs.services/platform/convex/agent_tools/approval_shared.ts (1)
69-96:⚠️ Potential issue | 🟠 MajorReset
generationStatuswhen the budget gate rejects.
threadMetadatais already patched togeneratinghere, but a deniedcheckBudget()throws without restoringidle. That leaves completion-triggered agent replies stuck in a loading state.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@services/platform/convex/agent_tools/approval_shared.ts` around lines 69 - 96, The code sets threadMeta.generationStatus to 'generating' then calls checkBudget, but if budgetResult.allowed is false it throws without restoring the thread state; update the control flow so that when persistentStreaming.createStream and ctx.db.patch (threadMeta) have been applied and checkBudget returns allowed === false you perform a ctx.db.patch on threadMeta._id to reset generationStatus to 'idle' (and clear streamId/updatedAt as appropriate) before throwing the Error; reference persistentStreaming.createStream, threadMeta, ctx.db.patch and checkBudget to locate where to add the compensating patch (or wrap the budget check in try/finally that resets state on rejection).services/platform/convex/governance/internal_mutations.ts (2)
42-50:⚠️ Potential issue | 🔴 CriticalA single
teamIdcan’t represent the new multi-team budget model.
checkBudget()now matches all of a user’s team memberships, butusageLedgerstill stores only oneteamIdper user/period row. If a user belongs to two budgeted teams, only one team aggregate gets incremented, so the other team cap can be bypassed. This needs fan-out per effective team (or a multi-team representation plus matching aggregation), not a single optional field.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@services/platform/convex/governance/internal_mutations.ts` around lines 42 - 50, The usageLedger insert currently writes a single optional teamId (ctx.db.insert into 'usageLedger' using teamId / existing.teamId), which breaks the multi-team budget model used by checkBudget(); change the write to fan-out per effective team instead of a single teamId: determine the full set of effectiveTeamIds for the user/period (the same list checkBudget() matches), and for each teamId insert a separate usageLedger row (or alternatively switch to a multi-team column and update all aggregation queries), and update the backfill branch that conditionally used existing.teamId to handle multiple teamIds as well so each budgeted team gets its own increment/entry.
60-90:⚠️ Potential issue | 🔴 CriticalThe race-recovery merge still drops usage and
teamId.Inside the dedupe loop, every patch recomputes from the original
keepsnapshot, so three concurrent inserts end up askeep + lastDupinstead of the full sum. This branch also deletes a duplicate without carryingdup.teamIdontokeep, which can erase the only team attribution for that period.🩹 Suggested fix
- const keep = sorted[0]; + let keep = sorted[0]; if (!keep) return null; for (let i = 1; i < sorted.length; i++) { const dup = sorted[i]; if (!dup) continue; await ctx.db.patch(keep._id, { inputTokens: keep.inputTokens + dup.inputTokens, outputTokens: keep.outputTokens + dup.outputTokens, totalTokens: keep.totalTokens + dup.totalTokens, costEstimate: keep.costEstimate + dup.costEstimate, requestCount: keep.requestCount + dup.requestCount, + ...(keep.teamId ? {} : dup.teamId ? { teamId: dup.teamId } : {}), }); + keep = { + ...keep, + inputTokens: keep.inputTokens + dup.inputTokens, + outputTokens: keep.outputTokens + dup.outputTokens, + totalTokens: keep.totalTokens + dup.totalTokens, + costEstimate: keep.costEstimate + dup.costEstimate, + requestCount: keep.requestCount + dup.requestCount, + teamId: keep.teamId ?? dup.teamId, + }; await ctx.db.delete(dup._id); }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@services/platform/convex/governance/internal_mutations.ts` around lines 60 - 90, The dedupe loop currently recomputes patches from the original keep snapshot causing lost increments and also drops dup.teamId; fix by accumulating merged values (e.g., let accInput = keep.inputTokens, accOutput = keep.outputTokens, accTotal = keep.totalTokens, accCost = keep.costEstimate, accRequests = keep.requestCount, accTeamId = keep.teamId) as you iterate sorted[1..], adding each dup's numeric fields into these accumulators and setting accTeamId = accTeamId || dup.teamId (or otherwise merge team attribution) when dup.teamId exists, then after the loop perform a single ctx.db.patch(keep._id, {inputTokens: accInput, outputTokens: accOutput, totalTokens: accTotal, costEstimate: accCost, requestCount: accRequests, teamId: accTeamId}) and delete duplicates with ctx.db.delete(dup._id); ensure you reference the existing symbols (usageLedger query by_org_user_period, keep, dup, ctx.db.patch, ctx.db.delete) and handle null/undefined numeric fields safely when summing.services/platform/convex/agent_tools/human_input/mutations.ts (1)
170-210:⚠️ Potential issue | 🟠 MajorBudget denial leaves this thread stuck in
generating.This flow creates the stream and patches
threadMetadata.generationStatusbeforecheckBudget(), then throws on denial without restoringidle. The frontend spinner will stick until some unrelated write clears it. Either move the budget gate before Lines 170-182 or mirror the reset logic used inservices/platform/convex/lib/agent_chat/start_agent_chat.ts.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@services/platform/convex/agent_tools/human_input/mutations.ts` around lines 170 - 210, The code sets up a persistent stream and patches threadMetadata.generationStatus -> 'generating' before calling resolveBudgetContext/checkBudget, so when budget is denied the thread remains stuck in 'generating'; fix by either moving the budget gate (resolveBudgetContext/checkBudget) to run before persistentStreaming.createStream and before the threadMetadata patch, or if you prefer minimal change mirror the reset logic from services/platform/convex/lib/agent_chat/start_agent_chat.ts: catch the budget denial and immediately patch the same threadMetadata record (found via the by_threadId query) to restore generationStatus to 'idle' and clear streamId/updatedAt as appropriate, then throw the error; reference persistentStreaming.createStream, threadMetadata (withIndex('by_threadId')), resolveBudgetContext, and checkBudget to locate where to apply the change.services/platform/app/features/settings/governance/components/budget-editor.tsx (2)
34-40:⚠️ Potential issue | 🟡 MinorLocalize the option labels too.
The dialog now uses
t()for surrounding copy, but these dropdown options still render hardcoded English (Organization,Weekly,Daily), so the form stays mixed-language inde. Move the option arrays behindt()instead of keeping literal labels here.As per coding guidelines, "Do NOT hardcode text, use the translation hooks/functions instead for user-facing UI".
Also applies to: 53-56
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@services/platform/app/features/settings/governance/components/budget-editor.tsx` around lines 34 - 40, SCOPE_OPTIONS and the other option arrays (the arrays around lines 53-56) currently use hardcoded English labels; move these option definitions into the component scope (so translation hook is available) and replace label strings with calls to the translation function (e.g. t('settings.scope.default'), t('settings.scope.user'), etc.) so every option label is localized; update any usages of SCOPE_OPTIONS, frequency options, or similar arrays to reference the newly translated arrays and add corresponding translation keys to the i18n resource files.
109-116:⚠️ Potential issue | 🟠 MajorReset
scopeIdwhenever the scope kind changes.Changing an existing rule from
team/user/roleto another targeted scope keeps the previousscopeId, so it's possible to save a role rule with a team id (or vice versa) that will never match during enforcement.💡 Suggested fix
const updateDraft = useCallback((patch: Partial<BudgetRule>) => { setDraft((prev) => { const updated = { ...prev, ...patch }; - if (patch.scope === 'default' || patch.scope === 'org') { + if (patch.scope && patch.scope !== prev.scope) { delete updated.scopeId; } return updated; }); }, []);🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@services/platform/app/features/settings/governance/components/budget-editor.tsx` around lines 109 - 116, The updateDraft handler currently only deletes scopeId when patch.scope is 'default' or 'org', but it must also clear scopeId whenever the scope kind actually changes (e.g., team->role) to avoid carrying an invalid id; inside updateDraft's setDraft((prev) => { ... }) compare patch.scope to prev.scope and if patch.scope is present and different from prev.scope then delete updated.scopeId (in addition to the existing deletion for 'default'/'org'); this touches the updateDraft function and the BudgetRule fields scope and scopeId.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@services/platform/convex/governance/__tests__/helpers.test.ts`:
- Around line 32-52: Replace the weak regex assertions in the weekly tests for
buildPeriodKeyFromTimestamp with concrete expected period keys: change the
mid-year assertion to expect '2024-W24' for the timestamp 2024-06-15, keep the
start-of-year test expecting '2024-W01', and replace the year-end test with a
concrete assertion that verifies the exact ISO-week key for 2024-12-31 (use
buildPeriodKeyFromTimestamp('weekly', ts) to get the computed value and assert
it equals the known correct ISO week key), and add an extra test comparing
2024-12-31 and 2025-01-01 to assert the week boundary behavior so off-by-one
errors fail.
In `@services/platform/convex/governance/budget_enforcement.ts`:
- Around line 477-527: The team and org aggregate checks never run warning
collection, so team/org warning thresholds (e.g., orgWarningThresholdPercent)
are dead; update the team loop and the org check to call collectWarnings with
the same rule+usage pair used for violation checks (use teamRule and teamUsage
inside the for loop, and orgRule and orgUsage in the org block) and merge those
warnings into allWarnings (or include them in the returned violation object) so
configured teamWarningThresholdPercent/orgWarningThresholdPercent on limits are
honored; reference collectWarnings, effectiveTeamIds, getTeamPeriodUsage,
getOrgPeriodUsage, teamRule, orgRule, and checkRuleAgainstUsage to locate where
to add these calls and propagate the resulting warnings.
- Around line 480-497: The loop currently reuses the merged "limits" object for
every team (effectiveTeamIds) which causes team aggregates to be validated
against user/default fields; instead, for each teamId call the same resolution
logic that produced "limits" but scoped to that team (i.e., derive
per-teamLimits from the matching team-scoped rules before building teamRule),
then construct teamRule from those per-teamLimits (only including
maxTokens/maxCostCents/maxRequests if they actually originated from team/org
scope) and pass that to checkRuleAgainstUsage; use getTeamPeriodUsage as before
but replace the merged "limits" usage with the per-team resolved limits to
ensure team-scoped enforcement is correct.
In `@services/platform/convex/lib/agent_chat/__tests__/start_agent_chat.test.ts`:
- Around line 16-18: Add assertions to the existing test(s) to verify that
startAgentChat forwards the resolved budget context into checkBudget: mock or
spy on checkBudget (the function imported/used by startAgentChat), run the test
that uses resolveBudgetContext (the case where betterAuth adapter returns the
mocked findMany), and assert that checkBudget was called with an argument
containing the resolved userTeamIds and userRole returned by
resolveBudgetContext; apply the same assertion for the other test case(s) around
lines 89-96 so regressions to agent-derived team scope are caught.
In `@services/platform/lib/shared/schemas/governance.ts`:
- Around line 14-22: The schema currently allows saving a rule with only
warningThresholdPercent set and no ceiling; update budgetRuleSchema to validate
that if warningThresholdPercent is provided then at least one of maxTokens,
maxCostCents, or maxRequests is also provided (or alternatively enforce that a
rule must include at least one ceiling field overall). Implement this with a zod
refinement or superRefine on budgetRuleSchema that checks
warningThresholdPercent presence and throws a validation error if none of the
ceiling fields (maxTokens, maxCostCents, maxRequests) are set, referencing
budgetRuleSchema, warningThresholdPercent, maxTokens, maxCostCents, and
maxRequests to locate the change.
In `@services/platform/messages/en.json`:
- Line 3683: The help text for the "costLimitHelp" string currently hardcodes
model pricing ("GPT-4o ~ $10/1M tokens"); update the "costLimitHelp" value to
remove specific price info and instead use a generic phrasing (e.g., "Hard
spending cap in USD. Model pricing varies — see vendor docs for current rates.")
so it doesn't drift; locate the "costLimitHelp" JSON key and replace the message
accordingly.
---
Outside diff comments:
In
`@services/platform/app/features/settings/governance/components/budget-editor.tsx`:
- Around line 34-40: SCOPE_OPTIONS and the other option arrays (the arrays
around lines 53-56) currently use hardcoded English labels; move these option
definitions into the component scope (so translation hook is available) and
replace label strings with calls to the translation function (e.g.
t('settings.scope.default'), t('settings.scope.user'), etc.) so every option
label is localized; update any usages of SCOPE_OPTIONS, frequency options, or
similar arrays to reference the newly translated arrays and add corresponding
translation keys to the i18n resource files.
- Around line 109-116: The updateDraft handler currently only deletes scopeId
when patch.scope is 'default' or 'org', but it must also clear scopeId whenever
the scope kind actually changes (e.g., team->role) to avoid carrying an invalid
id; inside updateDraft's setDraft((prev) => { ... }) compare patch.scope to
prev.scope and if patch.scope is present and different from prev.scope then
delete updated.scopeId (in addition to the existing deletion for
'default'/'org'); this touches the updateDraft function and the BudgetRule
fields scope and scopeId.
In `@services/platform/convex/agent_tools/approval_shared.ts`:
- Around line 69-96: The code sets threadMeta.generationStatus to 'generating'
then calls checkBudget, but if budgetResult.allowed is false it throws without
restoring the thread state; update the control flow so that when
persistentStreaming.createStream and ctx.db.patch (threadMeta) have been applied
and checkBudget returns allowed === false you perform a ctx.db.patch on
threadMeta._id to reset generationStatus to 'idle' (and clear streamId/updatedAt
as appropriate) before throwing the Error; reference
persistentStreaming.createStream, threadMeta, ctx.db.patch and checkBudget to
locate where to add the compensating patch (or wrap the budget check in
try/finally that resets state on rejection).
In `@services/platform/convex/agent_tools/human_input/mutations.ts`:
- Around line 170-210: The code sets up a persistent stream and patches
threadMetadata.generationStatus -> 'generating' before calling
resolveBudgetContext/checkBudget, so when budget is denied the thread remains
stuck in 'generating'; fix by either moving the budget gate
(resolveBudgetContext/checkBudget) to run before
persistentStreaming.createStream and before the threadMetadata patch, or if you
prefer minimal change mirror the reset logic from
services/platform/convex/lib/agent_chat/start_agent_chat.ts: catch the budget
denial and immediately patch the same threadMetadata record (found via the
by_threadId query) to restore generationStatus to 'idle' and clear
streamId/updatedAt as appropriate, then throw the error; reference
persistentStreaming.createStream, threadMetadata (withIndex('by_threadId')),
resolveBudgetContext, and checkBudget to locate where to apply the change.
In `@services/platform/convex/agent_tools/location/mutations.ts`:
- Around line 137-176: The code creates a stream via
persistentStreaming.createStream and sets threadMetadata.generationStatus to
'generating' before calling checkBudget, so when checkBudget denies it throws
and leaves the thread stuck; fix by moving budget checks to run before creating
the stream and before patching threadMetadata (or, if stream must be created
early, ensure the catch/denial path resets the metadata and deletes the stream):
call resolveBudgetContext and checkBudget prior to
persistentStreaming.createStream and the ctx.db.patch that sets
generationStatus, or on budget denial call ctx.db.patch(threadMeta._id, {
generationStatus: 'idle', streamId: null, updatedAt: Date.now() }) and clean up
persistentStreaming (delete/close) to avoid leaking stream IDs.
In `@services/platform/convex/governance/internal_mutations.ts`:
- Around line 42-50: The usageLedger insert currently writes a single optional
teamId (ctx.db.insert into 'usageLedger' using teamId / existing.teamId), which
breaks the multi-team budget model used by checkBudget(); change the write to
fan-out per effective team instead of a single teamId: determine the full set of
effectiveTeamIds for the user/period (the same list checkBudget() matches), and
for each teamId insert a separate usageLedger row (or alternatively switch to a
multi-team column and update all aggregation queries), and update the backfill
branch that conditionally used existing.teamId to handle multiple teamIds as
well so each budgeted team gets its own increment/entry.
- Around line 60-90: The dedupe loop currently recomputes patches from the
original keep snapshot causing lost increments and also drops dup.teamId; fix by
accumulating merged values (e.g., let accInput = keep.inputTokens, accOutput =
keep.outputTokens, accTotal = keep.totalTokens, accCost = keep.costEstimate,
accRequests = keep.requestCount, accTeamId = keep.teamId) as you iterate
sorted[1..], adding each dup's numeric fields into these accumulators and
setting accTeamId = accTeamId || dup.teamId (or otherwise merge team
attribution) when dup.teamId exists, then after the loop perform a single
ctx.db.patch(keep._id, {inputTokens: accInput, outputTokens: accOutput,
totalTokens: accTotal, costEstimate: accCost, requestCount: accRequests, teamId:
accTeamId}) and delete duplicates with ctx.db.delete(dup._id); ensure you
reference the existing symbols (usageLedger query by_org_user_period, keep, dup,
ctx.db.patch, ctx.db.delete) and handle null/undefined numeric fields safely
when summing.
🪄 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: ASSERTIVE
Plan: Pro
Run ID: 0bf5cb04-daa8-409d-bd9e-d2db0430c4b6
⛔ Files ignored due to path filters (1)
services/platform/convex/_generated/api.d.tsis excluded by!**/_generated/**
📒 Files selected for processing (14)
services/platform/app/features/settings/governance/components/budget-editor.tsxservices/platform/convex/agent_tools/approval_shared.tsservices/platform/convex/agent_tools/human_input/mutations.tsservices/platform/convex/agent_tools/location/mutations.tsservices/platform/convex/governance/__tests__/budget_enforcement.test.tsservices/platform/convex/governance/__tests__/helpers.test.tsservices/platform/convex/governance/budget_enforcement.tsservices/platform/convex/governance/internal_mutations.tsservices/platform/convex/governance/resolve_budget_context.tsservices/platform/convex/lib/agent_chat/__tests__/start_agent_chat.test.tsservices/platform/convex/lib/agent_chat/start_agent_chat.tsservices/platform/lib/shared/schemas/governance.tsservices/platform/messages/de.jsonservices/platform/messages/en.json
| describe('weekly', () => { | ||
| it('returns YYYY-Www format', () => { | ||
| // 2024-06-15 is a Saturday in ISO week 24 | ||
| const ts = Date.UTC(2024, 5, 15, 12, 0, 0); | ||
| const result = buildPeriodKeyFromTimestamp('weekly', ts); | ||
| expect(result).toMatch(/^\d{4}-W\d{2}$/); | ||
| }); | ||
|
|
||
| it('returns week 01 for start of year', () => { | ||
| // 2024-01-01 is a Monday | ||
| const ts = Date.UTC(2024, 0, 1, 0, 0, 0); | ||
| const result = buildPeriodKeyFromTimestamp('weekly', ts); | ||
| expect(result).toBe('2024-W01'); | ||
| }); | ||
|
|
||
| it('returns correct week for end of year', () => { | ||
| // 2024-12-31 — day 366 of a leap year | ||
| const ts = Date.UTC(2024, 11, 31, 0, 0, 0); | ||
| const result = buildPeriodKeyFromTimestamp('weekly', ts); | ||
| expect(result).toMatch(/^\d{4}-W\d{2}$/); | ||
| }); |
There was a problem hiding this comment.
Weekly assertions are too weak for a budget boundary helper.
These cases only validate the YYYY-Www shape, so an off-by-one week calculation would still pass. Please assert concrete keys for representative mid-year and year-end dates; weekly bucketing is part of the enforcement path.
As per coding guidelines, "Tests should cover happy paths, edge cases, and error conditions."
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@services/platform/convex/governance/__tests__/helpers.test.ts` around lines
32 - 52, Replace the weak regex assertions in the weekly tests for
buildPeriodKeyFromTimestamp with concrete expected period keys: change the
mid-year assertion to expect '2024-W24' for the timestamp 2024-06-15, keep the
start-of-year test expecting '2024-W01', and replace the year-end test with a
concrete assertion that verifies the exact ISO-week key for 2024-12-31 (use
buildPeriodKeyFromTimestamp('weekly', ts) to get the computed value and assert
it equals the known correct ISO week key), and add an extra test comparing
2024-12-31 and 2025-01-01 to assert the week boundary behavior so off-by-one
errors fail.
| // Collect warnings for approaching limits | ||
| allWarnings.push(...collectWarnings(limits, userUsage, period)); | ||
|
|
||
| let usage: UsageTotals; | ||
| if (rule.scope === 'team' && rule.scopeId) { | ||
| usage = await getTeamPeriodUsage( | ||
| // Check team aggregate usage when limits came from team-scoped rules | ||
| for (const teamId of limits.effectiveTeamIds) { | ||
| const teamUsage = await getTeamPeriodUsage( | ||
| ctx, | ||
| organizationId, | ||
| rule.scopeId, | ||
| teamId, | ||
| periodKey, | ||
| ); | ||
| } else { | ||
| usage = await getUserPeriodUsage(ctx, organizationId, userId, periodKey); | ||
| const teamRule: BudgetRule = { | ||
| scope: 'team', | ||
| scopeId: teamId, | ||
| period: period, | ||
| maxTokens: limits.maxTokens, | ||
| maxCostCents: limits.maxCostCents, | ||
| maxRequests: limits.maxRequests, | ||
| }; | ||
| const teamViolation = checkRuleAgainstUsage(teamRule, teamUsage); | ||
| if (teamViolation) { | ||
| return { | ||
| ...teamViolation, | ||
| warnings: allWarnings.length > 0 ? allWarnings : undefined, | ||
| }; | ||
| } | ||
| } | ||
|
|
||
| const violation = checkRuleAgainstUsage(rule, usage); | ||
| if (violation) { | ||
| return violation; | ||
| // Check org-scoped limits against org-wide aggregate usage | ||
| if ( | ||
| limits.orgMaxTokens != null || | ||
| limits.orgMaxCostCents != null || | ||
| limits.orgMaxRequests != null | ||
| ) { | ||
| const orgUsage = await getOrgPeriodUsage(ctx, organizationId, periodKey); | ||
| const orgRule: BudgetRule = { | ||
| scope: 'org', | ||
| period: period, | ||
| maxTokens: limits.orgMaxTokens, | ||
| maxCostCents: limits.orgMaxCostCents, | ||
| maxRequests: limits.orgMaxRequests, | ||
| }; | ||
| const orgViolation = checkRuleAgainstUsage(orgRule, orgUsage); | ||
| if (orgViolation) { | ||
| return { | ||
| ...orgViolation, | ||
| reason: `Organization-wide ${orgViolation.reason}`, | ||
| warnings: allWarnings.length > 0 ? allWarnings : undefined, | ||
| }; | ||
| } | ||
| } |
There was a problem hiding this comment.
Aggregate warning thresholds are currently dead code.
checkBudget() only calls collectWarnings() for personal usage. Team aggregate checks never emit warnings, and orgWarningThresholdPercent is resolved but never applied to org usage, so configuring warning thresholds on shared team/org caps has no effect until the request is fully blocked.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@services/platform/convex/governance/budget_enforcement.ts` around lines 477 -
527, The team and org aggregate checks never run warning collection, so team/org
warning thresholds (e.g., orgWarningThresholdPercent) are dead; update the team
loop and the org check to call collectWarnings with the same rule+usage pair
used for violation checks (use teamRule and teamUsage inside the for loop, and
orgRule and orgUsage in the org block) and merge those warnings into allWarnings
(or include them in the returned violation object) so configured
teamWarningThresholdPercent/orgWarningThresholdPercent on limits are honored;
reference collectWarnings, effectiveTeamIds, getTeamPeriodUsage,
getOrgPeriodUsage, teamRule, orgRule, and checkRuleAgainstUsage to locate where
to add these calls and propagate the resulting warnings.
| // Check team aggregate usage when limits came from team-scoped rules | ||
| for (const teamId of limits.effectiveTeamIds) { | ||
| const teamUsage = await getTeamPeriodUsage( | ||
| ctx, | ||
| organizationId, | ||
| rule.scopeId, | ||
| teamId, | ||
| periodKey, | ||
| ); | ||
| } else { | ||
| usage = await getUserPeriodUsage(ctx, organizationId, userId, periodKey); | ||
| const teamRule: BudgetRule = { | ||
| scope: 'team', | ||
| scopeId: teamId, | ||
| period: period, | ||
| maxTokens: limits.maxTokens, | ||
| maxCostCents: limits.maxCostCents, | ||
| maxRequests: limits.maxRequests, | ||
| }; | ||
| const teamViolation = checkRuleAgainstUsage(teamRule, teamUsage); | ||
| if (teamViolation) { |
There was a problem hiding this comment.
Don't validate team aggregates against the merged personal limits.
limits is resolved across user > team > role > default, but you reuse that merged object to build every teamRule. That means a team aggregate can be blocked by a user/default field that was never defined for the team, and users in multiple teams get every team checked against the same merged cap instead of each team's own budget. This breaks team-scoped enforcement.
Resolve limits per teamId from that team's matching rules before calling checkRuleAgainstUsage, and only evaluate fields that actually came from team/org scope there.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@services/platform/convex/governance/budget_enforcement.ts` around lines 480 -
497, The loop currently reuses the merged "limits" object for every team
(effectiveTeamIds) which causes team aggregates to be validated against
user/default fields; instead, for each teamId call the same resolution logic
that produced "limits" but scoped to that team (i.e., derive per-teamLimits from
the matching team-scoped rules before building teamRule), then construct
teamRule from those per-teamLimits (only including
maxTokens/maxCostCents/maxRequests if they actually originated from team/org
scope) and pass that to checkRuleAgainstUsage; use getTeamPeriodUsage as before
but replace the merged "limits" usage with the per-team resolved limits to
ensure team-scoped enforcement is correct.
| betterAuth: { | ||
| adapter: { findMany: 'mock-betterAuth-findMany' }, | ||
| }, |
There was a problem hiding this comment.
The new budget-context plumbing still lacks a regression assertion.
These mocks make resolveBudgetContext() runnable, but none of the tests verify that startAgentChat() forwards the resolved userTeamIds and userRole into checkBudget(). A regression back to agent-derived team scope would still pass this file.
As per coding guidelines, "ALWAYS write tests for new features and bug fixes" and "Tests should cover happy paths, edge cases, and error conditions."
Also applies to: 89-96
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@services/platform/convex/lib/agent_chat/__tests__/start_agent_chat.test.ts`
around lines 16 - 18, Add assertions to the existing test(s) to verify that
startAgentChat forwards the resolved budget context into checkBudget: mock or
spy on checkBudget (the function imported/used by startAgentChat), run the test
that uses resolveBudgetContext (the case where betterAuth adapter returns the
mocked findMany), and assert that checkBudget was called with an argument
containing the resolved userTeamIds and userRole returned by
resolveBudgetContext; apply the same assertion for the other test case(s) around
lines 89-96 so regressions to agent-derived team scope are caught.
| export const budgetRuleSchema = z.object({ | ||
| scope: z.enum(['user', 'team', 'role', 'default']), | ||
| scope: z.enum(['user', 'team', 'role', 'org', 'default']), | ||
| scopeId: z.string().optional(), | ||
| period: z.enum(['daily', 'weekly', 'monthly']), | ||
| maxTokens: z.number().nonnegative().optional(), | ||
| maxCostCents: z.number().nonnegative().optional(), | ||
| maxRequests: z.number().nonnegative().optional(), | ||
| warningThresholdPercent: z.number().min(0).max(100).optional(), | ||
| }); |
There was a problem hiding this comment.
Reject warning-only budget rules.
warningThresholdPercent can now be persisted without any of maxTokens, maxCostCents, or maxRequests. That creates a rule with no ceiling to evaluate against, so admins can save a warning that will never fire.
🧩 Suggested schema refinement
-export const budgetRuleSchema = z.object({
- scope: z.enum(['user', 'team', 'role', 'org', 'default']),
- scopeId: z.string().optional(),
- period: z.enum(['daily', 'weekly', 'monthly']),
- maxTokens: z.number().nonnegative().optional(),
- maxCostCents: z.number().nonnegative().optional(),
- maxRequests: z.number().nonnegative().optional(),
- warningThresholdPercent: z.number().min(0).max(100).optional(),
-});
+export const budgetRuleSchema = z
+ .object({
+ scope: z.enum(['user', 'team', 'role', 'org', 'default']),
+ scopeId: z.string().optional(),
+ period: z.enum(['daily', 'weekly', 'monthly']),
+ maxTokens: z.number().nonnegative().optional(),
+ maxCostCents: z.number().nonnegative().optional(),
+ maxRequests: z.number().nonnegative().optional(),
+ warningThresholdPercent: z.number().min(0).max(100).optional(),
+ })
+ .superRefine((rule, ctx) => {
+ if (
+ rule.warningThresholdPercent != null &&
+ rule.maxTokens == null &&
+ rule.maxCostCents == null &&
+ rule.maxRequests == null
+ ) {
+ ctx.addIssue({
+ code: z.ZodIssueCode.custom,
+ path: ['warningThresholdPercent'],
+ message:
+ 'warningThresholdPercent requires maxTokens, maxCostCents, or maxRequests',
+ });
+ }
+ });📝 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.
| export const budgetRuleSchema = z.object({ | |
| scope: z.enum(['user', 'team', 'role', 'default']), | |
| scope: z.enum(['user', 'team', 'role', 'org', 'default']), | |
| scopeId: z.string().optional(), | |
| period: z.enum(['daily', 'weekly', 'monthly']), | |
| maxTokens: z.number().nonnegative().optional(), | |
| maxCostCents: z.number().nonnegative().optional(), | |
| maxRequests: z.number().nonnegative().optional(), | |
| warningThresholdPercent: z.number().min(0).max(100).optional(), | |
| }); | |
| export const budgetRuleSchema = z | |
| .object({ | |
| scope: z.enum(['user', 'team', 'role', 'org', 'default']), | |
| scopeId: z.string().optional(), | |
| period: z.enum(['daily', 'weekly', 'monthly']), | |
| maxTokens: z.number().nonnegative().optional(), | |
| maxCostCents: z.number().nonnegative().optional(), | |
| maxRequests: z.number().nonnegative().optional(), | |
| warningThresholdPercent: z.number().min(0).max(100).optional(), | |
| }) | |
| .superRefine((rule, ctx) => { | |
| if ( | |
| rule.warningThresholdPercent != null && | |
| rule.maxTokens == null && | |
| rule.maxCostCents == null && | |
| rule.maxRequests == null | |
| ) { | |
| ctx.addIssue({ | |
| code: z.ZodIssueCode.custom, | |
| path: ['warningThresholdPercent'], | |
| message: | |
| 'warningThresholdPercent requires maxTokens, maxCostCents, or maxRequests', | |
| }); | |
| } | |
| }); |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@services/platform/lib/shared/schemas/governance.ts` around lines 14 - 22, The
schema currently allows saving a rule with only warningThresholdPercent set and
no ceiling; update budgetRuleSchema to validate that if warningThresholdPercent
is provided then at least one of maxTokens, maxCostCents, or maxRequests is also
provided (or alternatively enforce that a rule must include at least one ceiling
field overall). Implement this with a zod refinement or superRefine on
budgetRuleSchema that checks warningThresholdPercent presence and throws a
validation error if none of the ceiling fields (maxTokens, maxCostCents,
maxRequests) are set, referencing budgetRuleSchema, warningThresholdPercent,
maxTokens, maxCostCents, and maxRequests to locate the change.
| "costLimitUsd": "Cost limit (USD)", | ||
| "requestsPerMinute": "Requests per minute", | ||
| "costLimitUsd": "Max cost (USD)", | ||
| "costLimitHelp": "Hard spending cap in USD. GPT-4o ~ $10/1M tokens.", |
There was a problem hiding this comment.
Avoid hardcoded model pricing in help text.
"GPT-4o ~ $10/1M tokens" will drift as vendor pricing changes and can mislead users configuring caps.
Suggested copy tweak
- "costLimitHelp": "Hard spending cap in USD. GPT-4o ~ $10/1M tokens.",
+ "costLimitHelp": "Hard spending cap in USD for the selected period.",🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@services/platform/messages/en.json` at line 3683, The help text for the
"costLimitHelp" string currently hardcodes model pricing ("GPT-4o ~ $10/1M
tokens"); update the "costLimitHelp" value to remove specific price info and
instead use a generic phrasing (e.g., "Hard spending cap in USD. Model pricing
varies — see vendor docs for current rates.") so it doesn't drift; locate the
"costLimitHelp" JSON key and replace the message accordingly.
Add a user-level budget status query (getMyBudgetStatus) that checks the current user's budget against governance rules and surfaces both hard limit violations and approaching-threshold warnings. Display a non-dismissible red banner when budget is exceeded and a dismissible yellow banner when usage approaches the warning threshold. The banner only appears on the new chat page (no threadId) to avoid noise in existing conversations. - Add getMyBudgetStatus query in governance/queries.ts - Add useMyBudgetStatus hook in governance hooks - Wire budget status banner into chat-interface.tsx - Include usage/limit details in cost limit exceeded reason - Add i18n strings for en and de locales
Summary
Closes #1366
resolveBudgetContext()helper that resolves user role and team memberships at allcheckBudget()call sites, replacing agent-team scope with user-team scope (R1, R2, R4 from task)resolveEffectiveLimits(), so more specific scopes override broader ones (R3)orgscope to budget rules for organization-wide aggregate caps, daily/weekly period enforcement, warning threshold support, and team aggregate usage checks (R4, R5, R8).collect()withfor awaitasync iteration in governance module per coding standards (R12)en.jsonandde.jsonlocales (R11)Test plan
bun run --filter @tale/platform lint)bun run --filter @tale/platform typecheck)Summary by CodeRabbit
Release Notes
New Features
Improvements
Tests