Skip to content

fix(platform): fix budget scope override logic and enforcement#1392

Merged
larryro merged 2 commits into
mainfrom
fix/issue-1366-budget-limits
Apr 11, 2026
Merged

fix(platform): fix budget scope override logic and enforcement#1392
larryro merged 2 commits into
mainfrom
fix/issue-1366-budget-limits

Conversation

@yannickmonney
Copy link
Copy Markdown
Contributor

@yannickmonney yannickmonney commented Apr 11, 2026

Summary

Closes #1366

  • Add resolveBudgetContext() helper that resolves user role and team memberships at all checkBudget() call sites, replacing agent-team scope with user-team scope (R1, R2, R4 from task)
  • Implement priority override model (user > team > role > default) with per-field resolution via resolveEffectiveLimits(), so more specific scopes override broader ones (R3)
  • Add org scope to budget rules for organization-wide aggregate caps, daily/weekly period enforcement, warning threshold support, and team aggregate usage checks (R4, R5, R8)
  • Replace .collect() with for await async iteration in governance module per coding standards (R12)
  • Add unit tests for rule matching, effective limit resolution, period key generation, and boundary conditions (R9)
  • Internationalize all budget editor UI strings across en.json and de.json locales (R11)

Test plan

  • Lint passes (bun run --filter @tale/platform lint)
  • Typecheck passes (bun run --filter @tale/platform typecheck)
  • All 158 server test files pass (1996 tests)
  • All 403 client test files pass (3848 tests)
  • Verify budget editor UI renders correctly with new org scope and warning threshold fields
  • Verify daily/weekly period enforcement blocks requests when limits are exceeded
  • Verify user-scoped rule overrides default for the same limit field
  • Verify org-wide cap blocks all users when aggregate usage exceeds the limit

Summary by CodeRabbit

Release Notes

  • New Features

    • Organization-level budget rule scopes for setting cross-user limits
    • Weekly and daily billing periods now available for selection (previously disabled)
    • Warning thresholds to alert when usage approaches configured limits
  • Improvements

    • Budget enforcement now uses resolved user team memberships and roles instead of agent configuration
    • Enhanced budget rule configuration interface with improved translations and accessibility features
  • Tests

    • Added comprehensive test coverage for budget enforcement and period handling logic

Copy link
Copy Markdown

@greptile-apps greptile-apps Bot left a comment

Choose a reason for hiding this comment

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

Your free trial has ended. If you'd like to continue receiving code reviews, you can add a payment method here.

@coderabbitai
Copy link
Copy Markdown

coderabbitai Bot commented Apr 11, 2026

📝 Walkthrough

Walkthrough

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

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 52.63% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title directly addresses the main change: fixing budget scope override logic and enforcement across user/team/role/org scopes with a priority model.
Linked Issues check ✅ Passed The PR fully implements all issue #1366 requirements: per-user budgets, per-team/group budgets, per-scope limits, a priority override model (user>team>role>default), and consistent enforcement across all AI requests.
Out of Scope Changes check ✅ Passed All changes align with issue #1366 objectives. The PR adds org scope, warning thresholds, resolves budget context at enforcement sites, implements effective limit resolution, and includes supporting tests and i18n.

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

✨ Finishing Touches
📝 Generate docstrings
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch fix/issue-1366-budget-limits

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.

@yannickmonney
Copy link
Copy Markdown
Contributor Author

@coderabbitai review

@coderabbitai
Copy link
Copy Markdown

coderabbitai Bot commented Apr 11, 2026

✅ Actions performed

Review triggered.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

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.
@larryro larryro force-pushed the fix/issue-1366-budget-limits branch from 579b606 to a74eb36 Compare April 11, 2026 09:17
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: 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 | 🟠 Major

Budget 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 to idle. 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 | 🟠 Major

Reset generationStatus when the budget gate rejects.

threadMetadata is already patched to generating here, but a denied checkBudget() throws without restoring idle. 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 | 🔴 Critical

A single teamId can’t represent the new multi-team budget model.

checkBudget() now matches all of a user’s team memberships, but usageLedger still stores only one teamId per 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 | 🔴 Critical

The race-recovery merge still drops usage and teamId.

Inside the dedupe loop, every patch recomputes from the original keep snapshot, so three concurrent inserts end up as keep + lastDup instead of the full sum. This branch also deletes a duplicate without carrying dup.teamId onto keep, 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 | 🟠 Major

Budget denial leaves this thread stuck in generating.

This flow creates the stream and patches threadMetadata.generationStatus before checkBudget(), then throws on denial without restoring idle. 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 in services/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 | 🟡 Minor

Localize 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 in de. Move the option arrays behind t() 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 | 🟠 Major

Reset scopeId whenever the scope kind changes.

Changing an existing rule from team/user/role to another targeted scope keeps the previous scopeId, 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

📥 Commits

Reviewing files that changed from the base of the PR and between 638418f and a74eb36.

⛔ Files ignored due to path filters (1)
  • services/platform/convex/_generated/api.d.ts is excluded by !**/_generated/**
📒 Files selected for processing (14)
  • services/platform/app/features/settings/governance/components/budget-editor.tsx
  • services/platform/convex/agent_tools/approval_shared.ts
  • services/platform/convex/agent_tools/human_input/mutations.ts
  • services/platform/convex/agent_tools/location/mutations.ts
  • services/platform/convex/governance/__tests__/budget_enforcement.test.ts
  • services/platform/convex/governance/__tests__/helpers.test.ts
  • services/platform/convex/governance/budget_enforcement.ts
  • services/platform/convex/governance/internal_mutations.ts
  • services/platform/convex/governance/resolve_budget_context.ts
  • services/platform/convex/lib/agent_chat/__tests__/start_agent_chat.test.ts
  • services/platform/convex/lib/agent_chat/start_agent_chat.ts
  • services/platform/lib/shared/schemas/governance.ts
  • services/platform/messages/de.json
  • services/platform/messages/en.json

Comment on lines +32 to +52
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}$/);
});
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 | 🟡 Minor

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.

Comment on lines +477 to 527
// 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,
};
}
}
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

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.

Comment on lines +480 to +497
// 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) {
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 | 🔴 Critical

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.

Comment on lines +16 to +18
betterAuth: {
adapter: { findMany: 'mock-betterAuth-findMany' },
},
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 | 🟡 Minor

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.

Comment on lines 14 to 22
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(),
});
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

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.

Suggested change
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.",
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 | 🟡 Minor

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
@larryro larryro merged commit 7e23494 into main Apr 11, 2026
24 checks passed
@larryro larryro deleted the fix/issue-1366-budget-limits branch April 11, 2026 09:48
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.

Unable to configure Budgets and Limits per user, team, role, org-wide

2 participants