feat(webapp): server-side agent message quota - #4552
Conversation
Enforce the Free plan's agent-message allowance on the server, not just as a client hint. A per-(organizationId, period) counter lives in its own table, not joined to chats, so deleting a chat can no longer free quota inside the period. The create path and the .in append path each count one user message and refuse over the cap with a typed 403 the client renders as the upgrade block; wakes (action turns) never count. Fails open: an absent limit (self-hosted, or before the cloud side ships) or a counter read that throws means no cap. TRI-12863.
|
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
WalkthroughAdds monthly, organization-level agent message quotas with UTC billing periods. Stores usage in a new database table with atomic increment queries. Applies quota checks during chat creation and message forwarding, and records counted messages. Updates dashboard agent chat to process quota refusals, refresh usage, and display an upgrade block with the applicable limit. Adds PostgreSQL-backed tests for threshold, period, persistence, resolver, and error-handling behavior. 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
📝 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 |
…1166' into feat/agent-message-quota-tri-12863
@trigger.dev/build
trigger.dev
@trigger.dev/core
@trigger.dev/python
@trigger.dev/react-hooks
@trigger.dev/redis-worker
@trigger.dev/rsc
@trigger.dev/schema-to-json
@trigger.dev/sdk
commit: |
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (2)
apps/webapp/test/dashboardAgentQuota.test.ts (1)
162-174: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider covering a failing
readLimitas well.The fail-open test covers only a counter read that throws. The limit lookup is the other half of
Promise.alland can also reject. Add a case wherereadLimitrejects and assertundefined.♻️ Suggested extra case
it("fails open when the counter read throws", async () => {Add after that test:
it("fails open when the limit lookup throws", async () => { const result = await resolveAgentMessageQuota({} as unknown as DashboardAgentDb, { organizationId: ORG, readLimit: async () => { throw new Error("limit cache down"); }, }); expect(result).toBeUndefined(); });apps/webapp/app/services/dashboardAgentQuota.server.ts (1)
13-15: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueKeep the cast only while
agentMessagesis outsideLimits.This is the only non-node_modules reference to
AGENT_MESSAGE_LIMIT_KEY, andagentMessagesis not declared in the repo yet. Add a follow-up to remove the cast once@trigger.dev/platformdeclares it, or keep the currentLimitslookup via theagentMessageskey directly.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: fc824e93-d916-42f2-a87d-9e9793e73c9c
📒 Files selected for processing (12)
.server-changes/agent-message-quota.mdapps/webapp/app/components/dashboard-agent/DashboardAgentChat.tsxapps/webapp/app/components/dashboard-agent/useAgentMessageQuota.tsapps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.dashboard-agent.in.$.tsapps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.dashboard-agent.tsapps/webapp/app/services/dashboardAgentQuota.server.tsapps/webapp/test/dashboardAgentQuota.test.tsinternal-packages/dashboard-agent-db/drizzle/0004_stale_corsair.sqlinternal-packages/dashboard-agent-db/drizzle/meta/0004_snapshot.jsoninternal-packages/dashboard-agent-db/drizzle/meta/_journal.jsoninternal-packages/dashboard-agent-db/src/queries.tsinternal-packages/dashboard-agent-db/src/schema.ts
|
|
||
| if (countsAgainstQuota) { | ||
| await recordAgentMessageSent(dashboardAgentDb, { | ||
| organizationId: project.organizationId, | ||
| }); | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Count the message after the upstream send succeeds.
recordAgentMessageSent runs at Line 175, before the upstream fetch at Line 189. If that fetch throws, the route returns 502 and the message was never delivered, but the quota was already consumed. The same happens when the upstream returns a non-2xx status.
The sibling route records only after startDashboardAgentHeadStart resolves. Match that ordering here.
🐛 Proposed fix
body = JSON.stringify(parsed);
-
- if (countsAgainstQuota) {
- await recordAgentMessageSent(dashboardAgentDb, {
- organizationId: project.organizationId,
- });
- }
}
}Then record after the upstream call succeeds:
try {
const upstream = await fetch(upstreamUrl, { method: "POST", headers, body });
const text = await upstream.text();
+ // Only a delivered message consumes quota; a failed forward must not burn one.
+ if (countsAgainstQuota && upstream.ok) {
+ await recordAgentMessageSent(dashboardAgentDb, {
+ organizationId: project.organizationId,
+ });
+ }
return new Response(text, {countsAgainstQuota must move out of the if (parsed) block for this to compile; declare it as let countsAgainstQuota = false; above Line 124.
| // The per-period counter, org-wide: a deleted chat can't lower it within the period. | ||
| if (searchParams.get("quota") === "1") { | ||
| const used = await countUserMessages(dashboardAgentDb, { | ||
| const used = await getAgentMessageUsage(dashboardAgentDb, { | ||
| organizationId: project.organizationId, | ||
| userId, | ||
| excludeChatId: searchParams.get("chatId") ?? undefined, | ||
| period: currentAgentMessagePeriod(), | ||
| }); | ||
| return json({ used }); | ||
| } |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Server and client evaluate the quota against different limits. The server resolves the limit through resolveAgentMessageQuota, which defaults to UNLIMITED_AGENT_MESSAGES when the plan limit is absent. The quota poll drops that limit and returns only used, so the client falls back to FREE_PLAN_MESSAGE_LIMIT = 20. On a billing-wired install with no agentMessages plan limit configured, the client blocks the composer at 20 messages while the server still accepts sends. The documented fail-open behavior does not reach the UI.
apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.dashboard-agent.ts#L158-L165: callresolveAgentMessageQuotain the?quota=1branch and return{ used, limit }; return an empty payload when it resolves toundefined.apps/webapp/app/components/dashboard-agent/useAgentMessageQuota.ts#L26-L46: readlimitfrom the poll response, hold it in state, and pass it toresolveMessageQuotaso the client uses the server's limit instead of the hardcoded fallback.
📍 Affects 2 files
apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.dashboard-agent.ts#L158-L165(this comment)apps/webapp/app/components/dashboard-agent/useAgentMessageQuota.ts#L26-L46
| isFreePlan, | ||
| used: usedElsewhere === undefined ? undefined : usedElsewhere + countUserMessages(messages), | ||
| }); | ||
| return resolveMessageQuota({ isFreePlan, used }); |
There was a problem hiding this comment.
🟡 The message allowance shown and enforced in the browser ignores the plan's actual allowance
The browser decides whether the allowance is used up with a hardcoded number of 20 (resolveMessageQuota({ isFreePlan, used }) at apps/webapp/app/components/dashboard-agent/useAgentMessageQuota.ts:45) rather than the allowance the server actually enforces, so the two can disagree.
Impact: Users can be shown the wrong allowance and blocked from sending even when their plan permits more messages.
The `?quota=1` endpoint returns only `used`, never the limit
Server-side enforcement resolves the limit from the billing cache (AGENT_MESSAGE_LIMIT_KEY in apps/webapp/app/services/dashboardAgentQuota.server.ts:15,44-54), falling back to the unlimited sentinel. The loader branch for ?quota=1 (apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.dashboard-agent.ts:155-162) returns only used, so the client falls back to the local FREE_PLAN_MESSAGE_LIMIT = 20.
Today that means free-plan cloud orgs are blocked at 20 purely client-side while the server considers them unlimited; once the cloud agentMessages limit ships with any value other than 20, the composer will be hidden at the wrong point and the "N of 20 free messages left" copy will be wrong. Returning limit alongside used from the loader and using it in resolveMessageQuota keeps both sides on one number.
Prompt for agents
The server enforces the agent message cap using the per-plan `agentMessages` limit resolved in apps/webapp/app/services/dashboardAgentQuota.server.ts, but the `?quota=1` loader branch in the dashboard-agent resource route only returns `used`. The client (useAgentMessageQuota.ts) therefore falls back to the hardcoded FREE_PLAN_MESSAGE_LIMIT of 20 for both the 'x of N free messages left' copy and for hiding the composer. Have the loader also return the resolved limit (it already has it via resolveAgentMessageQuota) and thread it into resolveMessageQuota so the UI and the server enforcement can never disagree.
Was this helpful? React with 👍 or 👎 to provide feedback.
| const reached = parseQuotaReachedResponse(res.status, data); | ||
| if (reached) { | ||
| setCapReached(reached); | ||
| setActive(null); | ||
| return; | ||
| } |
There was a problem hiding this comment.
🔍 Draft view only learns about the cap after a failed create
capReached is set purely from a refused create response; the draft never polls ?quota=1. A free-plan user with no open chat therefore always sees the composer, types a message, and only then gets the upgrade block — the typed text is dropped (createChat never returns it). Worth confirming this is the intended first-run experience versus pre-resolving the quota when the panel opens with no active chat.
Was this helpful? React with 👍 or 👎 to provide feedback.
A failed upstream send (5xx/502) or a non-2xx response burned a quota message that never reached the agent. Record only after upstream.ok.
Draft submit and chat retry now bail when the message cap is reached, so a suggested prompt or retry over the cap no longer fires a silent 403. The capped draft keeps any open watch card. The quota re-reads when a turn settles instead of on optimistic append, so the count and cap no longer lag by one message.
What & why
The Free-plan agent-message allowance was only a client-side hint (
useIsFreePlan()was hardcoded toundefined, so nothing was enforced) and the running count was a liveCOUNT(*)overchat_messages— which meant deleting a chat silently freed quota. This PR makes the allowance a real server-side limit with a durable counter, so the cap holds regardless of the client and a deleted chat can't reclaim messages within the period.The cloud side that fills the actual per-plan number is a separate PR (TRI-12863 P0). Until it deploys the
agentMessageslimit key is absent, which resolves to the repo's unlimited sentinel — the correct fail-open default, and why this ships independently.What's inside
trigger_dashboard_agent.agent_message_usage, keyed(organization_id, period)whereperiodis a UTC calendar month"YYYY-MM". FK-free (pgSchemaconvention), plus a drizzle migration. Deliberately not joined to chats — that closes the delete-a-chat hole. Distinct from TRI-13068'sAiUsageEvent; neither is derived from the other.dashboardAgentQuota.server.ts: a purecheckAgentMessageQuota({ used, limit })(so the MCP send path can reuse the rule later) plus an org-scoped resolver that reads the period counter and the cached plan limit, and arecordAgentMessageSentincrement..inappend path each increment one user message. The append path counts only after the existingtrigger === "action"403, so wakes never count.403 { error: "message_quota_reached", limit }, which the client renders asAgentUpgradeBlockrather than a generic failure. Never a silent drop.?quota=1now reads the period counter;useIsFreePlan()is a real read gated on billing presence (no subscription → self-hosted → no cap, no upgrade UI), not on the plan value.Key decisions
100_000_000(neverInfinity— that serializes tonullin the Redis limit cache), and a counter read that throws returns "no cap". Self-hosted needs zero extra branching — it falls out of the fallback, with a test to prove it.(org, period)table closes the delete hole. A standalone counter, not a count over chat rows, so deleting a chat can't free quota inside the period.Testing
apps/webapp/test/dashboardAgentQuota.test.ts(testcontainers, no mocks):checkAgentMessageQuotaunder/at/over/unlimited (control-breaks the>=);agentTurnCountsAgainstQuotacounts a message but not a wake (control-breaks the wake exclusion);TRI-12863