diff --git a/.server-changes/agent-message-quota.md b/.server-changes/agent-message-quota.md new file mode 100644 index 0000000000..296b082de7 --- /dev/null +++ b/.server-changes/agent-message-quota.md @@ -0,0 +1,6 @@ +--- +area: webapp +type: feature +--- + +The dashboard agent now comes with a monthly message allowance. A message that fails to send doesn't count against it. diff --git a/apps/webapp/app/components/dashboard-agent/AgentUpgradeGate.tsx b/apps/webapp/app/components/dashboard-agent/AgentUpgradeGate.tsx index 4b795edde9..498a043157 100644 --- a/apps/webapp/app/components/dashboard-agent/AgentUpgradeGate.tsx +++ b/apps/webapp/app/components/dashboard-agent/AgentUpgradeGate.tsx @@ -1,9 +1,10 @@ import { Link } from "@remix-run/react"; +import { AgentMonoLogo } from "~/components/primitives/AgentDotMatrix"; import { LinkButton } from "~/components/primitives/Buttons"; import { useOrganization } from "~/hooks/useOrganizations"; -import { cn } from "~/utils/cn"; import { v3BillingPath } from "~/utils/pathBuilder"; -import { AgentIcon, AGENT_ICON_ACCENT_CLASS, ASK_AGENT_LABEL } from "./agent-identity"; +import { ASK_AGENT_LABEL } from "./agent-identity"; +import { messageQuotaReachedCopy } from "./message-quota"; // Matches the composer's outer geometry so the replacement lands in the same place. const SLOT = "flex shrink-0 flex-col bg-background-bright px-3 pb-3 pt-1"; @@ -22,14 +23,12 @@ export function AgentUpgradeBlock({ {context}
- + Upgrade to unlock {ASK_AGENT_LABEL}
-

- You've used all {limit} messages included on the Free plan. Your chats stay here to read. -

+

{messageQuotaReachedCopy(limit)}

Upgrade diff --git a/apps/webapp/app/components/dashboard-agent/DashboardAgentChat.tsx b/apps/webapp/app/components/dashboard-agent/DashboardAgentChat.tsx index 19e65a570c..aaecd89e40 100644 --- a/apps/webapp/app/components/dashboard-agent/DashboardAgentChat.tsx +++ b/apps/webapp/app/components/dashboard-agent/DashboardAgentChat.tsx @@ -17,6 +17,7 @@ import { DashboardAgentContextBanner } from "./DashboardAgentContextBanner"; import { DashboardAgentHero } from "./DashboardAgentHero"; import { DashboardAgentMessages, type TurnActivity } from "./DashboardAgentMessages"; import { MESSAGE_TOO_LARGE_ERROR } from "./message-limits"; +import { FREE_PLAN_MESSAGE_LIMIT, parseQuotaReachedResponse } from "./message-quota"; import { createTranscriptOrder, orderTranscript } from "./message-order"; import { navigateDestination } from "./navigate-target"; import { pendingNavigateIntents, pendingWatchIntents } from "./pending-intents"; @@ -102,6 +103,9 @@ export function DashboardAgentChat({ onActivityChange?: (chatId: string, activity: TurnActivity | null) => void; }) { const [input, setInput] = useState(""); + // Set when the server refuses a send over the cap, so the block shows at once rather than + // waiting for the next quota poll. + const [quotaReached, setQuotaReached] = useState<{ limit: number } | null>(null); const navigate = useNavigate(); const location = useLocation(); const toast = useToast(); @@ -128,6 +132,18 @@ export function DashboardAgentChat({ .catch(() => null)) as { error?: string } | null; throw new Error(data?.error ?? MESSAGE_TOO_LARGE_ERROR); } + // Over the message cap: show the upgrade block instead of a generic turn error. + if (res.status === 403) { + const data = (await res + .clone() + .json() + .catch(() => null)) as { error?: string; limit?: number } | null; + const reached = parseQuotaReachedResponse(res.status, data); + if (reached) { + setQuotaReached(reached); + throw new Error("You've reached your message limit."); + } + } return res; }, clientData, @@ -185,9 +201,12 @@ export function DashboardAgentChat({ const orderRef = useRef(createTranscriptOrder(initialMessages)); const messages = orderTranscript(rawMessages, orderRef.current); - // Counted here, not in the panel, so it includes the turn just sent. - const quota = useAgentMessageQuota({ actionPath, chatId, messages }); - const atMessageCap = quota.kind === "reached"; + // Read here, not in the panel, so it re-reads as each turn settles. + const quota = useAgentMessageQuota({ actionPath, chatId, status }); + // Either the poll saw the cap, or a send was just refused over it. + const atMessageCap = quota.kind === "reached" || quotaReached !== null; + const messageCapLimit = + quotaReached?.limit ?? (quota.kind === "reached" ? quota.limit : FREE_PLAN_MESSAGE_LIMIT); const isStreaming = status === "streaming"; // From status, not the last part: the indicator must stay up through silent tool calls. @@ -252,6 +271,8 @@ export function DashboardAgentChat({ }, [sendRequest, submit, canSend]); const retry = useCallback(() => { + // Over the cap, a retry only earns another 403 — same guard as `submit`. + if (atMessageCap) return; // A watch's consent record is a user message nobody typed, so retry never treats it as one. const action = retryAction( messages.filter((m) => !(m.role === "user" && isWatchRequestMessageId(m.id))) @@ -264,7 +285,7 @@ export function DashboardAgentChat({ return; } void sendMessage({ text: action.text, messageId: action.messageId }); - }, [messages, sendMessage, regenerate, clearError]); + }, [messages, sendMessage, regenerate, clearError, atMessageCap]); const resolveUri = useTriggerUriResolver(actionPath); @@ -414,9 +435,9 @@ export function DashboardAgentChat({ /> )} {watchCard ?
{watchCard}
: null} - {quota.kind === "reached" ? ( + {atMessageCap ? ( void; projectSlug: string; @@ -24,6 +26,7 @@ export function DashboardAgentDraft({ pageContext?: AgentPageContext; promotedPrompt?: SuggestedPrompt; watchCard?: React.ReactNode; + capReached?: { limit: number } | null; }) { const [input, setInput] = useState(""); @@ -43,12 +46,14 @@ export function DashboardAgentDraft({ const submit = useCallback( (text: string) => { + // Suggested prompts reach here via the hero, bypassing the composer's cap guard. + if (capReached) return; const trimmed = text.trim(); if (!trimmed) return; setInput(""); onSubmit(trimmed); }, - [onSubmit] + [onSubmit, capReached] ); return ( @@ -57,25 +62,41 @@ export function DashboardAgentDraft({ pageContext={pageContext} promoted={promotedPrompt} composer={ -
- {watchCard} - submit(input)} - onStop={() => {}} - isStreaming={false} - placeholderSuggestion={watchCard ? undefined : placeholderSuggestion} - context={ - - } - /> -
+ capReached ? ( +
+ {watchCard} + + } + /> +
+ ) : ( +
+ {watchCard} + submit(input)} + onStop={() => {}} + isStreaming={false} + placeholderSuggestion={watchCard ? undefined : placeholderSuggestion} + context={ + + } + /> +
+ ) } /> ); diff --git a/apps/webapp/app/components/dashboard-agent/DashboardAgentPanel.tsx b/apps/webapp/app/components/dashboard-agent/DashboardAgentPanel.tsx index d4d7c76ccd..06f596256b 100644 --- a/apps/webapp/app/components/dashboard-agent/DashboardAgentPanel.tsx +++ b/apps/webapp/app/components/dashboard-agent/DashboardAgentPanel.tsx @@ -24,6 +24,7 @@ import { writeLastChat, } from "./last-chat-storage"; import { DashboardAgentDraft } from "./DashboardAgentDraft"; +import { parseQuotaReachedResponse } from "./message-quota"; import { WatchCard } from "./WatchCard"; import { watchDraftFor } from "./watch-card"; import { NO_WATCH_CARD, watchCardReducer } from "./watch-card-state"; @@ -109,6 +110,8 @@ export function DashboardAgentPanel({ // Until the list has arrived, the page load's server count is the better answer. const [chatsLoaded, setChatsLoaded] = useState(false); const [active, setActive] = useState(null); + // A refused `create` over the cap: the draft shows the upgrade block instead of a raw toast. + const [capReached, setCapReached] = useState<{ limit: number } | null>(null); // Starts true so an `openWith` request waits for the restore instead of racing it. const [loading, setLoading] = useState( () => readLastChat(storageKey)?.path === location.pathname @@ -245,14 +248,22 @@ export function DashboardAgentPanel({ publicAccessToken?: string; headStarted?: boolean; error?: string; + limit?: number; }; if (seq !== openChatRequestSeq.current) return; if (!res.ok || !data.chatId || !data.publicAccessToken) { + const reached = parseQuotaReachedResponse(res.status, data); + if (reached) { + setCapReached(reached); + setActive(null); + return; + } console.error(`Dashboard agent: failed to create chat (${res.status})`, data.error); toast.error(data.error ?? "We couldn't start that chat. Try again in a moment."); setActive(null); return; } + setCapReached(null); setActive({ chatId: data.chatId, organizationId: organization.id, @@ -294,6 +305,7 @@ export function DashboardAgentPanel({ panelOrg.current = organization.id; claimChatSlot(); setActive(null); + setCapReached(null); setLoading(false); setChats([]); setChatsLoaded(false); @@ -621,6 +633,7 @@ export function DashboardAgentPanel({ pageContext={pageContext} promotedPrompt={promotedPrompt} watchCard={watchCardElement} + capReached={capReached} /> )} diff --git a/apps/webapp/app/components/dashboard-agent/message-quota.test.ts b/apps/webapp/app/components/dashboard-agent/message-quota.test.ts index 5a252db7ac..3b47395e99 100644 --- a/apps/webapp/app/components/dashboard-agent/message-quota.test.ts +++ b/apps/webapp/app/components/dashboard-agent/message-quota.test.ts @@ -1,5 +1,12 @@ import { describe, expect, it } from "vitest"; -import { countUserMessages, FREE_PLAN_MESSAGE_LIMIT, resolveMessageQuota } from "./message-quota"; +import { + countUserMessages, + FREE_PLAN_MESSAGE_LIMIT, + MESSAGE_QUOTA_REACHED_ERROR, + messageQuotaReachedCopy, + parseQuotaReachedResponse, + resolveMessageQuota, +} from "./message-quota"; describe("resolveMessageQuota", () => { it("caps a Free plan at the limit", () => { @@ -42,6 +49,37 @@ describe("resolveMessageQuota", () => { }); }); +describe("parseQuotaReachedResponse", () => { + it("maps a create/in 403 cap body to the limit", () => { + // Both the create path and the `in` transport refuse with this exact body. + expect( + parseQuotaReachedResponse(403, { error: MESSAGE_QUOTA_REACHED_ERROR, limit: 20 }) + ).toEqual({ limit: 20 }); + }); + + it("falls back to the free limit when the body omits it", () => { + expect(parseQuotaReachedResponse(403, { error: MESSAGE_QUOTA_REACHED_ERROR })).toEqual({ + limit: FREE_PLAN_MESSAGE_LIMIT, + }); + }); + + it("ignores other errors and non-403 statuses so they surface normally", () => { + expect(parseQuotaReachedResponse(403, { error: "something_else" })).toBeNull(); + expect(parseQuotaReachedResponse(500, { error: MESSAGE_QUOTA_REACHED_ERROR })).toBeNull(); + expect(parseQuotaReachedResponse(403, null)).toBeNull(); + }); +}); + +describe("messageQuotaReachedCopy", () => { + it("is a friendly sentence naming the limit, never the raw code", () => { + const copy = messageQuotaReachedCopy(20); + expect(copy).toContain("all 20 messages"); + expect(copy).toContain("Free plan"); + // Control break: if the mapping leaked the server code, this fails. + expect(copy).not.toContain(MESSAGE_QUOTA_REACHED_ERROR); + }); +}); + describe("countUserMessages", () => { it("counts only what the user sent", () => { expect( diff --git a/apps/webapp/app/components/dashboard-agent/message-quota.ts b/apps/webapp/app/components/dashboard-agent/message-quota.ts index f65481c870..c603a836a2 100644 --- a/apps/webapp/app/components/dashboard-agent/message-quota.ts +++ b/apps/webapp/app/components/dashboard-agent/message-quota.ts @@ -27,6 +27,28 @@ export function resolveMessageQuota({ : { kind: "within", used, limit, remaining }; } +// The server code both the create and `in` paths refuse with. The client owns the copy, +// so this code must never reach the UI as text. +export const MESSAGE_QUOTA_REACHED_ERROR = "message_quota_reached"; + +// Maps a 403 refusal body to the cap signal, or null for any other error. Both paths use +// this so a `message_quota_reached` code routes to the upgrade block, never a raw toast. +export function parseQuotaReachedResponse( + status: number, + data: { error?: string; limit?: number } | null | undefined +): { limit: number } | null { + if (status === 403 && data?.error === MESSAGE_QUOTA_REACHED_ERROR) { + return { limit: data.limit ?? FREE_PLAN_MESSAGE_LIMIT }; + } + return null; +} + +// The upgrade block's sentence. Pure so the copy is asserted directly, and so the raw +// server code can never be what the user reads. +export function messageQuotaReachedCopy(limit: number): string { + return `You've used all ${limit} messages included on the Free plan. Your chats stay here to read.`; +} + // A watch's consent record is a user message the person never typed, so it is // excluded here exactly as the stored count excludes it. export function countUserMessages(messages: { role: string; id?: string }[]): number { diff --git a/apps/webapp/app/components/dashboard-agent/useAgentMessageQuota.ts b/apps/webapp/app/components/dashboard-agent/useAgentMessageQuota.ts index da6cc3d83b..958d8dfcc1 100644 --- a/apps/webapp/app/components/dashboard-agent/useAgentMessageQuota.ts +++ b/apps/webapp/app/components/dashboard-agent/useAgentMessageQuota.ts @@ -1,46 +1,55 @@ -import type { UIMessage } from "@ai-sdk/react"; -import { useEffect, useState } from "react"; -import { countUserMessages, resolveMessageQuota, type MessageQuota } from "./message-quota"; +import { useEffect, useRef, useState } from "react"; +import { useCurrentPlan } from "~/routes/_app.orgs.$organizationSlug/route"; +import { resolveMessageQuota, type MessageQuota } from "./message-quota"; -// Always undefined until billing supplies plan detection, which means no cap. +// Gated on billing PRESENCE, not the plan value: no subscription means billing isn't wired +// up (self-hosted), so there is no cap and no upgrade UI. A wired-up, non-paying plan is free. function useIsFreePlan(): boolean | undefined { - return undefined; + const subscription = useCurrentPlan()?.v3Subscription; + if (!subscription) return undefined; + return subscription.isPaying === false; } -// Counted in two halves: the server aggregates other chats, this chat's own count -// comes from the live transcript so the message just sent counts immediately. +// `used` is the server's per-period count for the org. Re-read once a turn settles — the +// server increment happens mid-turn in the `.in` proxy, so reading on optimistic append +// would lag the count by one message and show the cap a message late. export function useAgentMessageQuota({ actionPath, chatId, - messages, + status, }: { actionPath: string; chatId: string; - messages: UIMessage[]; + status: string; }): MessageQuota { const isFreePlan = useIsFreePlan(); - const [usedElsewhere, setUsedElsewhere] = useState(undefined); + const [used, setUsed] = useState(undefined); + + // Bumped each time the status leaves streaming/submitted, which drives the re-read. + const [settleTick, setSettleTick] = useState(0); + const prevStatus = useRef(status); + useEffect(() => { + const wasInFlight = prevStatus.current === "streaming" || prevStatus.current === "submitted"; + const nowSettled = status === "ready" || status === "error"; + prevStatus.current = status; + if (wasInFlight && nowSettled) setSettleTick((tick) => tick + 1); + }, [status]); useEffect(() => { if (isFreePlan !== true) return; const controller = new AbortController(); void (async () => { try { - const res = await fetch(`${actionPath}?quota=1&chatId=${encodeURIComponent(chatId)}`, { - signal: controller.signal, - }); + const res = await fetch(`${actionPath}?quota=1`, { signal: controller.signal }); if (!res.ok) return; const data = (await res.json()) as { used?: number }; - if (typeof data.used === "number") setUsedElsewhere(data.used); + if (typeof data.used === "number") setUsed(data.used); } catch { // Leave the count unknown, which means no cap. See `resolveMessageQuota`. } })(); return () => controller.abort(); - }, [isFreePlan, actionPath, chatId]); + }, [isFreePlan, actionPath, chatId, settleTick]); - return resolveMessageQuota({ - isFreePlan, - used: usedElsewhere === undefined ? undefined : usedElsewhere + countUserMessages(messages), - }); + return resolveMessageQuota({ isFreePlan, used }); } diff --git a/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.dashboard-agent.in.$.ts b/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.dashboard-agent.in.$.ts index 5939e836db..655da1d9aa 100644 --- a/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.dashboard-agent.in.$.ts +++ b/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.dashboard-agent.in.$.ts @@ -7,6 +7,7 @@ import { MESSAGE_TOO_LARGE_CODE, MESSAGE_TOO_LARGE_ERROR, } from "~/components/dashboard-agent/message-limits"; +import { MESSAGE_QUOTA_REACHED_ERROR } from "~/components/dashboard-agent/message-quota"; import { findProjectBySlug } from "~/models/project.server"; import { findEnvironmentBySlug } from "~/models/runtimeEnvironment.server"; import { @@ -15,6 +16,12 @@ import { resolveDashboardAgentRepoSnapshot, } from "~/services/dashboardAgent.server"; import { dashboardAgentEnvironmentAddress } from "~/services/dashboardAgentEnvironmentAddress.server"; +import { dashboardAgentDb } from "~/services/dashboardAgentDb.server"; +import { + agentTurnCountsAgainstQuota, + recordAgentMessageSent, + resolveAgentMessageQuota, +} from "~/services/dashboardAgentQuota.server"; import { logger } from "~/services/logger.server"; import { requireUser } from "~/services/session.server"; import { readBoundedBodyText } from "~/utils/boundedRequestBody.server"; @@ -115,6 +122,9 @@ export async function action({ request, params }: ActionFunctionArgs) { parsed = undefined; } + // Hoisted so it is visible after the fetch: quota is charged only once the send succeeds. + let countsAgainstQuota = false; + if (parsed) { // Actions are placed by the server only, and this proxy is the one path a browser // can reach `.in` through. @@ -127,6 +137,17 @@ export async function action({ request, params }: ActionFunctionArgs) { return tooLarge(); } + // Only a real user message consumes quota; action turns were refused above. + countsAgainstQuota = agentTurnCountsAgainstQuota(parsed); + if (countsAgainstQuota) { + const quota = await resolveAgentMessageQuota(dashboardAgentDb, { + organizationId: project.organizationId, + }); + if (quota?.reached) { + return json({ error: MESSAGE_QUOTA_REACHED_ERROR, limit: quota.limit }, { status: 403 }); + } + } + let userActorToken: string; try { userActorToken = await mintDashboardAgentUserActorToken(user.id, { @@ -165,6 +186,13 @@ export async function action({ request, params }: ActionFunctionArgs) { try { const upstream = await fetch(upstreamUrl, { method: "POST", headers, body }); const text = await upstream.text(); + // Charge quota only for a delivered message: a non-2xx upstream (or a throw below) + // must not burn a send that never reached the agent. + if (countsAgainstQuota && upstream.ok) { + await recordAgentMessageSent(dashboardAgentDb, { + organizationId: project.organizationId, + }); + } return new Response(text, { status: upstream.status, headers: { "content-type": upstream.headers.get("content-type") ?? "application/json" }, diff --git a/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.dashboard-agent.ts b/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.dashboard-agent.ts index ae6a93142f..8e51c04937 100644 --- a/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.dashboard-agent.ts +++ b/apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.dashboard-agent.ts @@ -3,7 +3,7 @@ import { chatExists, countUnreadWatchWakes, countChatsWithUnreadWork, - countUserMessages, + getAgentMessageUsage, createChat, getChatMessages, getSession, @@ -28,6 +28,7 @@ import { MESSAGE_TOO_LARGE_CODE, MESSAGE_TOO_LARGE_ERROR, } from "~/components/dashboard-agent/message-limits"; +import { MESSAGE_QUOTA_REACHED_ERROR } from "~/components/dashboard-agent/message-quota"; import { MAX_URIS_PER_RESOLVE_REQUEST } from "~/components/dashboard-agent/resolve-uris"; import { $replica } from "~/db.server"; import { env } from "~/env.server"; @@ -52,6 +53,11 @@ import { import { dashboardAgentEnvironmentAddress } from "~/services/dashboardAgentEnvironmentAddress.server"; import { startDashboardAgentHeadStart } from "~/services/dashboardAgentHeadStart.server"; import { dashboardAgentDb } from "~/services/dashboardAgentDb.server"; +import { + currentAgentMessagePeriod, + recordAgentMessageSent, + resolveAgentMessageQuota, +} from "~/services/dashboardAgentQuota.server"; import { logger } from "~/services/logger.server"; import { resolveTriggerUri } from "~/services/resolveTriggerUri.server"; import { requireUser } from "~/services/session.server"; @@ -150,13 +156,11 @@ export const loader = async ({ request, params }: LoaderFunctionArgs) => { const project = await findProjectBySlug(organizationSlug, projectParam, userId); if (!project) return json({ error: "Project not found" }, { status: 404 }); - // The open chat is excluded and counted from the live transcript instead, so an - // unpersisted turn still counts against the cap. + // 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 }); } @@ -290,6 +294,13 @@ export const action = async ({ request, params }: ActionFunctionArgs) => { return messageTooLarge(); } + const quota = await resolveAgentMessageQuota(dashboardAgentDb, { + organizationId: project.organizationId, + }); + if (quota?.reached) { + return json({ error: MESSAGE_QUOTA_REACHED_ERROR, limit: quota.limit }, { status: 403 }); + } + let clientData: Record | undefined; try { clientData = parsed.data.clientData @@ -387,6 +398,14 @@ export const action = async ({ request, params }: ActionFunctionArgs) => { throw error; } + // Only the head start dispatches the first message here; a cold start sends it through + // the `in` proxy, which counts it there. Counting both would double-count. + if (headStarted) { + await recordAgentMessageSent(dashboardAgentDb, { + organizationId: project.organizationId, + }); + } + let publicAccessToken: string; try { publicAccessToken = await mintDashboardAgentToken(chatId); diff --git a/apps/webapp/app/services/dashboardAgentQuota.server.ts b/apps/webapp/app/services/dashboardAgentQuota.server.ts new file mode 100644 index 0000000000..4bfca8808b --- /dev/null +++ b/apps/webapp/app/services/dashboardAgentQuota.server.ts @@ -0,0 +1,101 @@ +import type { Limits } from "@trigger.dev/platform"; +import { + getAgentMessageUsage, + incrementAgentMessageUsage, + type DashboardAgentDb, +} from "@internal/dashboard-agent-db"; +import { getCachedLimit } from "./platform.v3.server"; +import { logger } from "./logger.server"; + +// The repo's unlimited sentinel. Never Infinity: it serializes to null in the limit cache. +export const UNLIMITED_AGENT_MESSAGES = 100_000_000; + +// Filled by cloud billing (TRI-12863 P0). Absent until then, and always on self-hosted, +// so the fallback applies and the cap is effectively off. +const AGENT_MESSAGE_LIMIT_KEY = "agentMessages" as keyof Limits; + +/** The billing period the counter is scoped to: a UTC calendar month, "YYYY-MM". */ +export function currentAgentMessagePeriod(now: Date = new Date()): string { + return `${now.getUTCFullYear()}-${String(now.getUTCMonth() + 1).padStart(2, "0")}`; +} + +/** Pure so the send routes and, later, the MCP path share one rule. */ +export function checkAgentMessageQuota({ used, limit }: { used: number; limit: number }): { + reached: boolean; +} { + return { reached: used >= limit }; +} + +export type AgentMessageQuota = { reached: boolean; used: number; limit: number }; + +/** + * The period counter and the cached plan limit for one org. Fails open: an absent limit + * (self-hosted, or before the cloud side ships) resolves to the unlimited sentinel, and a + * counter read that throws returns `undefined` — either way there is no cap. + */ +export async function resolveAgentMessageQuota( + db: DashboardAgentDb, + params: { + organizationId: string; + now?: Date; + readLimit?: (organizationId: string) => Promise; + } +): Promise { + const readLimit = + params.readLimit ?? + (async (organizationId: string) => { + const cached = await getCachedLimit( + organizationId, + AGENT_MESSAGE_LIMIT_KEY, + UNLIMITED_AGENT_MESSAGES + ); + // A cache error leaves `val` empty; fall open to unlimited. + return cached.val ?? UNLIMITED_AGENT_MESSAGES; + }); + try { + const [limit, used] = await Promise.all([ + readLimit(params.organizationId), + getAgentMessageUsage(db, { + organizationId: params.organizationId, + period: currentAgentMessagePeriod(params.now), + }), + ]); + return { ...checkAgentMessageQuota({ used, limit }), used, limit }; + } catch (error) { + logger.error("Failed to resolve dashboard agent message quota", { + organizationId: params.organizationId, + error, + }); + return undefined; + } +} + +/** Record one sent user message. Swallows errors: the cap is a nudge, never a send blocker. */ +export async function recordAgentMessageSent( + db: DashboardAgentDb, + params: { organizationId: string; now?: Date } +): Promise { + try { + await incrementAgentMessageUsage(db, { + organizationId: params.organizationId, + period: currentAgentMessagePeriod(params.now), + }); + } catch (error) { + logger.error("Failed to record a dashboard agent message against the quota", { + organizationId: params.organizationId, + error, + }); + } +} + +/** + * Whether an agent turn consumes quota. Only a genuine new user message counts: the transport + * tags it `trigger: "submit-message"`. A retry/regenerate re-runs the agent from its own history + * without a new message (`trigger: "regenerate-message"`), and a wake is `"action"` — neither is + * something the user typed, so neither counts. + */ +export function agentTurnCountsAgainstQuota( + turn: { kind?: string; payload?: { trigger?: string } } | undefined +): boolean { + return turn?.kind === "message" && turn.payload?.trigger === "submit-message"; +} diff --git a/apps/webapp/test/dashboardAgentQuota.test.ts b/apps/webapp/test/dashboardAgentQuota.test.ts new file mode 100644 index 0000000000..21d5e46ffe --- /dev/null +++ b/apps/webapp/test/dashboardAgentQuota.test.ts @@ -0,0 +1,184 @@ +import { + createChat, + createDashboardAgentDb, + getAgentMessageUsage, + incrementAgentMessageUsage, + softDeleteChat, + type DashboardAgentDb, + type DashboardAgentDbClient, +} from "@internal/dashboard-agent-db"; +import { postgresTest } from "@internal/testcontainers"; +import type { PrismaClient } from "@trigger.dev/database"; +import { readdirSync, readFileSync } from "node:fs"; +import path from "node:path"; +import { afterEach, describe, expect, it } from "vitest"; +import { + agentTurnCountsAgainstQuota, + checkAgentMessageQuota, + currentAgentMessagePeriod, + resolveAgentMessageQuota, + UNLIMITED_AGENT_MESSAGES, +} from "~/services/dashboardAgentQuota.server"; + +/** + * Server-side agent message quota (TRI-12863): a per-(org, period) counter that a deleted chat + * can't lower, a pure at/over/under rule, and a resolver that fails open when the limit is + * absent (self-hosted) or the counter read throws. + */ + +const MIGRATIONS = path.resolve(__dirname, "../../../internal-packages/dashboard-agent-db/drizzle"); + +async function applyAgentSchema(prisma: PrismaClient) { + for (const name of readdirSync(MIGRATIONS) + .filter((file) => file.endsWith(".sql")) + .sort()) { + const sql = readFileSync(path.join(MIGRATIONS, name), "utf8"); + for (const statement of sql.split("--> statement-breakpoint")) { + const trimmed = statement.trim(); + if (trimmed.length > 0) await prisma.$executeRawUnsafe(trimmed); + } + } +} + +const ORG = "org_quota"; +const USER = "user_quota"; + +let agentDbClient: DashboardAgentDbClient | undefined; + +async function boot(prisma: PrismaClient, connectionUri: string): Promise { + await applyAgentSchema(prisma); + agentDbClient = createDashboardAgentDb(connectionUri, { max: 4 }); + return agentDbClient.db; +} + +afterEach(async () => { + await agentDbClient?.close(); + agentDbClient = undefined; +}); + +describe("checkAgentMessageQuota", () => { + it("is not reached under the limit", () => { + expect(checkAgentMessageQuota({ used: 5, limit: 20 })).toEqual({ reached: false }); + }); + + it("is reached at the limit", () => { + // Control break: `>=`. Flip to `>` and this fails. + expect(checkAgentMessageQuota({ used: 20, limit: 20 })).toEqual({ reached: true }); + }); + + it("is reached over the limit", () => { + expect(checkAgentMessageQuota({ used: 21, limit: 20 })).toEqual({ reached: true }); + }); + + it("is never reached against the unlimited sentinel", () => { + expect(checkAgentMessageQuota({ used: 10_000, limit: UNLIMITED_AGENT_MESSAGES })).toEqual({ + reached: false, + }); + }); +}); + +describe("agentTurnCountsAgainstQuota", () => { + it("counts a genuine new user message (submit-message)", () => { + expect( + agentTurnCountsAgainstQuota({ kind: "message", payload: { trigger: "submit-message" } }) + ).toBe(true); + }); + + it("does not count a retry/regenerate", () => { + // Control break: a regenerate re-runs from history with no new message, so it must not + // burn quota. Widen the rule back to `!== "action"` and this fails. + expect( + agentTurnCountsAgainstQuota({ kind: "message", payload: { trigger: "regenerate-message" } }) + ).toBe(false); + }); + + it("does not count a wake (action turn)", () => { + expect(agentTurnCountsAgainstQuota({ kind: "message", payload: { trigger: "action" } })).toBe( + false + ); + }); + + it("does not count a non-message turn or a missing body", () => { + expect(agentTurnCountsAgainstQuota({ kind: "action" })).toBe(false); + expect(agentTurnCountsAgainstQuota(undefined)).toBe(false); + }); +}); + +describe("currentAgentMessagePeriod", () => { + it("is a zero-padded UTC calendar month", () => { + expect(currentAgentMessagePeriod(new Date(Date.UTC(2026, 7, 9)))).toBe("2026-08"); + expect(currentAgentMessagePeriod(new Date(Date.UTC(2026, 0, 1)))).toBe("2026-01"); + }); +}); + +describe("the per-(org, period) counter", () => { + postgresTest( + "accumulates and a deleted chat cannot free quota within the period", + async ({ prisma, postgresContainer }) => { + const db = await boot(prisma, postgresContainer.getConnectionUri()); + const period = "2026-08"; + + // The create path and then an append: two messages, same period. + expect(await incrementAgentMessageUsage(db, { organizationId: ORG, period })).toBe(1); + expect(await incrementAgentMessageUsage(db, { organizationId: ORG, period })).toBe(2); + expect(await getAgentMessageUsage(db, { organizationId: ORG, period })).toBe(2); + + // Deleting a chat must not move the counter: it is not joined to chats. + await createChat(db, { id: "chat_del", organizationId: ORG, userId: USER }); + await softDeleteChat(db, { chatId: "chat_del", userId: USER, organizationId: ORG }); + expect(await getAgentMessageUsage(db, { organizationId: ORG, period })).toBe(2); + + // The next period and other orgs start fresh. + expect(await getAgentMessageUsage(db, { organizationId: ORG, period: "2026-09" })).toBe(0); + expect(await getAgentMessageUsage(db, { organizationId: "org_other", period })).toBe(0); + } + ); +}); + +describe("resolveAgentMessageQuota", () => { + postgresTest( + "reports reached over the limit, and never reached when unlimited", + async ({ prisma, postgresContainer }) => { + const db = await boot(prisma, postgresContainer.getConnectionUri()); + const now = new Date(); + const period = currentAgentMessagePeriod(now); + for (let i = 0; i < 3; i++) { + await incrementAgentMessageUsage(db, { organizationId: ORG, period }); + } + + expect( + await resolveAgentMessageQuota(db, { organizationId: ORG, now, readLimit: async () => 3 }) + ).toEqual({ + reached: true, + used: 3, + limit: 3, + }); + expect( + await resolveAgentMessageQuota(db, { organizationId: ORG, now, readLimit: async () => 20 }) + ).toEqual({ reached: false, used: 3, limit: 20 }); + + // Self-hosted: the limit is absent, so the fallback (unlimited sentinel) applies and there + // is no cap — no extra branching, it falls out of the fallback. + const selfHosted = await resolveAgentMessageQuota(db, { + organizationId: ORG, + now, + readLimit: async () => UNLIMITED_AGENT_MESSAGES, + }); + expect(selfHosted?.reached).toBe(false); + } + ); + + it("fails open when the counter read throws", async () => { + const throwingDb = { + select: () => { + throw new Error("db down"); + }, + } as unknown as DashboardAgentDb; + + const result = await resolveAgentMessageQuota(throwingDb, { + organizationId: ORG, + readLimit: async () => 5, + }); + expect(result).toBeUndefined(); + }); +}); diff --git a/internal-packages/dashboard-agent-db/drizzle/0004_stale_corsair.sql b/internal-packages/dashboard-agent-db/drizzle/0004_stale_corsair.sql new file mode 100644 index 0000000000..d581b4fec3 --- /dev/null +++ b/internal-packages/dashboard-agent-db/drizzle/0004_stale_corsair.sql @@ -0,0 +1,8 @@ +CREATE TABLE "trigger_dashboard_agent"."agent_message_usage" ( + "organization_id" text NOT NULL, + "period" text NOT NULL, + "count" integer DEFAULT 0 NOT NULL, + "created_at" timestamp with time zone DEFAULT now() NOT NULL, + "updated_at" timestamp with time zone DEFAULT now() NOT NULL, + CONSTRAINT "agent_message_usage_organization_id_period_pk" PRIMARY KEY("organization_id","period") +); diff --git a/internal-packages/dashboard-agent-db/drizzle/meta/0004_snapshot.json b/internal-packages/dashboard-agent-db/drizzle/meta/0004_snapshot.json new file mode 100644 index 0000000000..5ee40ae24f --- /dev/null +++ b/internal-packages/dashboard-agent-db/drizzle/meta/0004_snapshot.json @@ -0,0 +1,1344 @@ +{ + "id": "f7cbfef4-7fc8-4deb-8da2-59248b242a60", + "prevId": "efb6f8b8-af9f-4ba7-9e38-bafd1f430b28", + "version": "7", + "dialect": "postgresql", + "tables": { + "trigger_dashboard_agent.agent_message_usage": { + "name": "agent_message_usage", + "schema": "trigger_dashboard_agent", + "columns": { + "organization_id": { + "name": "organization_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "period": { + "name": "period", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "count": { + "name": "count", + "type": "integer", + "primaryKey": false, + "notNull": true, + "default": 0 + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "updated_at": { + "name": "updated_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": { + "agent_message_usage_organization_id_period_pk": { + "name": "agent_message_usage_organization_id_period_pk", + "columns": ["organization_id", "period"] + } + }, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "trigger_dashboard_agent.chat_messages": { + "name": "chat_messages", + "schema": "trigger_dashboard_agent", + "columns": { + "chat_id": { + "name": "chat_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "message_id": { + "name": "message_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "position": { + "name": "position", + "type": "integer", + "primaryKey": false, + "notNull": true + }, + "role": { + "name": "role", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "message": { + "name": "message", + "type": "jsonb", + "primaryKey": false, + "notNull": true + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "chat_messages_chat_user_role_idx": { + "name": "chat_messages_chat_user_role_idx", + "columns": [ + { + "expression": "chat_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "message_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"trigger_dashboard_agent\".\"chat_messages\".\"role\" = 'user'", + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": {}, + "compositePrimaryKeys": { + "chat_messages_chat_id_message_id_pk": { + "name": "chat_messages_chat_id_message_id_pk", + "columns": ["chat_id", "message_id"] + } + }, + "uniqueConstraints": { + "chat_messages_chat_position_key": { + "name": "chat_messages_chat_position_key", + "nullsNotDistinct": false, + "columns": ["chat_id", "position"] + } + }, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "trigger_dashboard_agent.chat_sessions": { + "name": "chat_sessions", + "schema": "trigger_dashboard_agent", + "columns": { + "chat_id": { + "name": "chat_id", + "type": "text", + "primaryKey": true, + "notNull": true + }, + "public_access_token": { + "name": "public_access_token", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "last_event_id": { + "name": "last_event_id", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "run_id": { + "name": "run_id", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "updated_at": { + "name": "updated_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "trigger_dashboard_agent.chat_turn_evals": { + "name": "chat_turn_evals", + "schema": "trigger_dashboard_agent", + "columns": { + "chat_id": { + "name": "chat_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "turn": { + "name": "turn", + "type": "integer", + "primaryKey": false, + "notNull": true + }, + "organization_id": { + "name": "organization_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "user_id": { + "name": "user_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "agent_run_id": { + "name": "agent_run_id", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "eval_run_id": { + "name": "eval_run_id", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "project_ref": { + "name": "project_ref", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "environment": { + "name": "environment", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "current_page": { + "name": "current_page", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "model": { + "name": "model", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "prompt_slug": { + "name": "prompt_slug", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "prompt_version": { + "name": "prompt_version", + "type": "integer", + "primaryKey": false, + "notNull": false + }, + "tools_used": { + "name": "tools_used", + "type": "jsonb", + "primaryKey": false, + "notNull": true, + "default": "'[]'::jsonb" + }, + "tool_error": { + "name": "tool_error", + "type": "boolean", + "primaryKey": false, + "notNull": true, + "default": false + }, + "judge_model": { + "name": "judge_model", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "score_grounded": { + "name": "score_grounded", + "type": "smallint", + "primaryKey": false, + "notNull": false + }, + "score_answered": { + "name": "score_answered", + "type": "smallint", + "primaryKey": false, + "notNull": false + }, + "score_concise": { + "name": "score_concise", + "type": "smallint", + "primaryKey": false, + "notNull": false + }, + "passed": { + "name": "passed", + "type": "boolean", + "primaryKey": false, + "notNull": false + }, + "intent_category": { + "name": "intent_category", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "outcome": { + "name": "outcome", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "sentiment": { + "name": "sentiment", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "capability_gap": { + "name": "capability_gap", + "type": "boolean", + "primaryKey": false, + "notNull": true, + "default": false + }, + "docs_gap": { + "name": "docs_gap", + "type": "boolean", + "primaryKey": false, + "notNull": true, + "default": false + }, + "support_opportunity": { + "name": "support_opportunity", + "type": "boolean", + "primaryKey": false, + "notNull": true, + "default": false + }, + "feature_request": { + "name": "feature_request", + "type": "boolean", + "primaryKey": false, + "notNull": true, + "default": false + }, + "topics": { + "name": "topics", + "type": "jsonb", + "primaryKey": false, + "notNull": true, + "default": "'[]'::jsonb" + }, + "signals": { + "name": "signals", + "type": "jsonb", + "primaryKey": false, + "notNull": true, + "default": "'[]'::jsonb" + }, + "summary": { + "name": "summary", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "user_text": { + "name": "user_text", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "judge": { + "name": "judge", + "type": "jsonb", + "primaryKey": false, + "notNull": false + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "chat_turn_evals_org_created_idx": { + "name": "chat_turn_evals_org_created_idx", + "columns": [ + { + "expression": "organization_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "created_at", + "isExpression": false, + "asc": false, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + }, + "chat_turn_evals_created_idx": { + "name": "chat_turn_evals_created_idx", + "columns": [ + { + "expression": "created_at", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + }, + "chat_turn_evals_org_opps_idx": { + "name": "chat_turn_evals_org_opps_idx", + "columns": [ + { + "expression": "organization_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "created_at", + "isExpression": false, + "asc": false, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"trigger_dashboard_agent\".\"chat_turn_evals\".\"capability_gap\" or \"trigger_dashboard_agent\".\"chat_turn_evals\".\"docs_gap\" or \"trigger_dashboard_agent\".\"chat_turn_evals\".\"support_opportunity\" or \"trigger_dashboard_agent\".\"chat_turn_evals\".\"feature_request\"", + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": {}, + "compositePrimaryKeys": { + "chat_turn_evals_chat_id_turn_pk": { + "name": "chat_turn_evals_chat_id_turn_pk", + "columns": ["chat_id", "turn"] + } + }, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "trigger_dashboard_agent.chats": { + "name": "chats", + "schema": "trigger_dashboard_agent", + "columns": { + "id": { + "name": "id", + "type": "text", + "primaryKey": true, + "notNull": true + }, + "organization_id": { + "name": "organization_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "user_id": { + "name": "user_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "title": { + "name": "title", + "type": "text", + "primaryKey": false, + "notNull": true, + "default": "'New chat'" + }, + "metadata": { + "name": "metadata", + "type": "jsonb", + "primaryKey": false, + "notNull": true, + "default": "'{}'::jsonb" + }, + "messages": { + "name": "messages", + "type": "jsonb", + "primaryKey": false, + "notNull": true, + "default": "'[]'::jsonb" + }, + "pinned_at": { + "name": "pinned_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "last_read_at": { + "name": "last_read_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "deleted_at": { + "name": "deleted_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "last_message_at": { + "name": "last_message_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "next_message_position": { + "name": "next_message_position", + "type": "integer", + "primaryKey": false, + "notNull": true, + "default": 1 + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "updated_at": { + "name": "updated_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "chats_org_user_last_msg_idx": { + "name": "chats_org_user_last_msg_idx", + "columns": [ + { + "expression": "organization_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "user_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "last_message_at", + "isExpression": false, + "asc": false, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"trigger_dashboard_agent\".\"chats\".\"deleted_at\" is null", + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "trigger_dashboard_agent.investigations": { + "name": "investigations", + "schema": "trigger_dashboard_agent", + "columns": { + "id": { + "name": "id", + "type": "text", + "primaryKey": true, + "notNull": true + }, + "chat_id": { + "name": "chat_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "project_ref": { + "name": "project_ref", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "environment_ref": { + "name": "environment_ref", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "revision": { + "name": "revision", + "type": "integer", + "primaryKey": false, + "notNull": true, + "default": 0 + }, + "state": { + "name": "state", + "type": "jsonb", + "primaryKey": false, + "notNull": true + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "updated_at": { + "name": "updated_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "investigations_chat_idx": { + "name": "investigations_chat_idx", + "columns": [ + { + "expression": "chat_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + }, + "investigations_open_updated_idx": { + "name": "investigations_open_updated_idx", + "columns": [ + { + "expression": "updated_at", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"trigger_dashboard_agent\".\"investigations\".\"state\"->>'outcome' = 'in_progress'", + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "trigger_dashboard_agent.watch_batches": { + "name": "watch_batches", + "schema": "trigger_dashboard_agent", + "columns": { + "environment_id": { + "name": "environment_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "cadence_minutes": { + "name": "cadence_minutes", + "type": "integer", + "primaryKey": false, + "notNull": true + }, + "epoch": { + "name": "epoch", + "type": "integer", + "primaryKey": false, + "notNull": true, + "default": 0 + }, + "generation": { + "name": "generation", + "type": "integer", + "primaryKey": false, + "notNull": true, + "default": 0 + }, + "status": { + "name": "status", + "type": "text", + "primaryKey": false, + "notNull": true, + "default": "'running'" + }, + "armed_at": { + "name": "armed_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "last_tick_at": { + "name": "last_tick_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + } + }, + "indexes": {}, + "foreignKeys": {}, + "compositePrimaryKeys": { + "watch_batches_environment_id_cadence_minutes_pk": { + "name": "watch_batches_environment_id_cadence_minutes_pk", + "columns": ["environment_id", "cadence_minutes"] + } + }, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "trigger_dashboard_agent.watch_submissions": { + "name": "watch_submissions", + "schema": "trigger_dashboard_agent", + "columns": { + "chat_id": { + "name": "chat_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "client_request_id": { + "name": "client_request_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "organization_id": { + "name": "organization_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "user_id": { + "name": "user_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "project_id": { + "name": "project_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "environment_id": { + "name": "environment_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "draft_hash": { + "name": "draft_hash", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "draft": { + "name": "draft", + "type": "jsonb", + "primaryKey": false, + "notNull": true + }, + "state": { + "name": "state", + "type": "text", + "primaryKey": false, + "notNull": true, + "default": "'pending'" + }, + "watch_id": { + "name": "watch_id", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "unavailable": { + "name": "unavailable", + "type": "boolean", + "primaryKey": false, + "notNull": true, + "default": false + }, + "external_notification_status": { + "name": "external_notification_status", + "type": "text", + "primaryKey": false, + "notNull": true, + "default": "'not_requested'" + }, + "external_notification_reason": { + "name": "external_notification_reason", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "immediate_result": { + "name": "immediate_result", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "refusal_code": { + "name": "refusal_code", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "refusal_error": { + "name": "refusal_error", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "refusal_existing_id": { + "name": "refusal_existing_id", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "updated_at": { + "name": "updated_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "watch_submissions_created_idx": { + "name": "watch_submissions_created_idx", + "columns": [ + { + "expression": "created_at", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": {}, + "compositePrimaryKeys": { + "watch_submissions_chat_id_client_request_id_pk": { + "name": "watch_submissions_chat_id_client_request_id_pk", + "columns": ["chat_id", "client_request_id"] + } + }, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "trigger_dashboard_agent.watches": { + "name": "watches", + "schema": "trigger_dashboard_agent", + "columns": { + "id": { + "name": "id", + "type": "text", + "primaryKey": true, + "notNull": true + }, + "chat_id": { + "name": "chat_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "identity": { + "name": "identity", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "spec": { + "name": "spec", + "type": "jsonb", + "primaryKey": false, + "notNull": true + }, + "status": { + "name": "status", + "type": "text", + "primaryKey": false, + "notNull": true, + "default": "'active'" + }, + "delivery_status": { + "name": "delivery_status", + "type": "text", + "primaryKey": false, + "notNull": true, + "default": "'not_required'" + }, + "cancel_reason": { + "name": "cancel_reason", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "resolution": { + "name": "resolution", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "observed_outcome": { + "name": "observed_outcome", + "type": "jsonb", + "primaryKey": false, + "notNull": false + }, + "investigate_on_attention": { + "name": "investigate_on_attention", + "type": "boolean", + "primaryKey": false, + "notNull": true, + "default": false + }, + "organization_id": { + "name": "organization_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "project_id": { + "name": "project_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "project_ref": { + "name": "project_ref", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "environment_id": { + "name": "environment_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "user_id": { + "name": "user_id", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "expires_at": { + "name": "expires_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true + }, + "last_checked_at": { + "name": "last_checked_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "last_attempted_at": { + "name": "last_attempted_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "fired_at": { + "name": "fired_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "delivery_claimed_at": { + "name": "delivery_claimed_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "delivery_claim_id": { + "name": "delivery_claim_id", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "delivered_at": { + "name": "delivered_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "cancelled_at": { + "name": "cancelled_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "last_result": { + "name": "last_result", + "type": "jsonb", + "primaryKey": false, + "notNull": false + }, + "tick_count": { + "name": "tick_count", + "type": "integer", + "primaryKey": false, + "notNull": true, + "default": 0 + }, + "alert_dispatch_key": { + "name": "alert_dispatch_key", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "retention_at": { + "name": "retention_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false, + "generated": { + "as": "greatest(delivered_at, cancelled_at, fired_at, last_checked_at, created_at)", + "type": "stored" + } + }, + "cadence_minutes": { + "name": "cadence_minutes", + "type": "integer", + "primaryKey": false, + "notNull": false, + "generated": { + "as": "((spec ->> 'checkEveryMinutes')::int)", + "type": "stored" + } + } + }, + "indexes": { + "watches_chat_idx": { + "name": "watches_chat_idx", + "columns": [ + { + "expression": "chat_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + }, + "watches_chat_active_identity_key": { + "name": "watches_chat_active_identity_key", + "columns": [ + { + "expression": "chat_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "project_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "environment_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "identity", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "where": "\"trigger_dashboard_agent\".\"watches\".\"status\" = 'active'", + "concurrently": false, + "method": "btree", + "with": {} + }, + "watches_status_expires_idx": { + "name": "watches_status_expires_idx", + "columns": [ + { + "expression": "status", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "expires_at", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + }, + "watches_pending_delivery_idx": { + "name": "watches_pending_delivery_idx", + "columns": [ + { + "expression": "fired_at", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "last_checked_at", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"trigger_dashboard_agent\".\"watches\".\"delivery_status\" in ('pending', 'delivering')", + "concurrently": false, + "method": "btree", + "with": {} + }, + "watches_org_user_wake_idx": { + "name": "watches_org_user_wake_idx", + "columns": [ + { + "expression": "organization_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "user_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "coalesce(\"fired_at\", \"last_checked_at\") desc", + "asc": true, + "isExpression": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"trigger_dashboard_agent\".\"watches\".\"delivery_status\" = 'delivered' and \"trigger_dashboard_agent\".\"watches\".\"status\" in ('fired', 'expired')", + "concurrently": false, + "method": "btree", + "with": {} + }, + "watches_org_user_active_idx": { + "name": "watches_org_user_active_idx", + "columns": [ + { + "expression": "organization_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "user_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"trigger_dashboard_agent\".\"watches\".\"status\" = 'active'", + "concurrently": false, + "method": "btree", + "with": {} + }, + "watches_active_env_cadence_idx": { + "name": "watches_active_env_cadence_idx", + "columns": [ + { + "expression": "environment_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "cadence_minutes", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "coalesce(\"last_attempted_at\", \"last_checked_at\", \"created_at\")", + "asc": true, + "isExpression": true, + "nulls": "last" + }, + { + "expression": "expires_at", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"trigger_dashboard_agent\".\"watches\".\"status\" = 'active'", + "concurrently": false, + "method": "btree", + "with": {} + }, + "watches_env_cadence_delivery_idx": { + "name": "watches_env_cadence_delivery_idx", + "columns": [ + { + "expression": "environment_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "cadence_minutes", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "delivery_status", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "coalesce(\"fired_at\", \"last_checked_at\")", + "asc": true, + "isExpression": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"trigger_dashboard_agent\".\"watches\".\"status\" in ('fired', 'expired') and \"trigger_dashboard_agent\".\"watches\".\"delivery_status\" in ('pending', 'delivering')", + "concurrently": false, + "method": "btree", + "with": {} + }, + "watches_retention_idx": { + "name": "watches_retention_idx", + "columns": [ + { + "expression": "retention_at", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "where": "\"trigger_dashboard_agent\".\"watches\".\"status\" in ('fired', 'expired', 'cancelled') and \"trigger_dashboard_agent\".\"watches\".\"delivery_status\" in ('not_required', 'delivered')", + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": {}, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + } + }, + "enums": {}, + "schemas": { + "trigger_dashboard_agent": "trigger_dashboard_agent" + }, + "sequences": {}, + "roles": {}, + "policies": {}, + "views": {}, + "_meta": { + "columns": {}, + "schemas": {}, + "tables": {} + } +} diff --git a/internal-packages/dashboard-agent-db/drizzle/meta/_journal.json b/internal-packages/dashboard-agent-db/drizzle/meta/_journal.json index 9efe7bc1a1..1f33e4ddf8 100644 --- a/internal-packages/dashboard-agent-db/drizzle/meta/_journal.json +++ b/internal-packages/dashboard-agent-db/drizzle/meta/_journal.json @@ -29,6 +29,13 @@ "when": 1786264383741, "tag": "0003_backfill_chat_last_read_at", "breakpoints": true + }, + { + "idx": 4, + "version": "7", + "when": 1786359241538, + "tag": "0004_stale_corsair", + "breakpoints": true } ] } diff --git a/internal-packages/dashboard-agent-db/src/queries.ts b/internal-packages/dashboard-agent-db/src/queries.ts index 5565b18271..4cf9329a79 100644 --- a/internal-packages/dashboard-agent-db/src/queries.ts +++ b/internal-packages/dashboard-agent-db/src/queries.ts @@ -8,6 +8,7 @@ import type { DashboardAgentDb } from "./client.js"; import { generateInvestigationId } from "./ids.js"; import { lockChatForWatches, type DashboardAgentDbOrTx } from "./internal.js"; import { + agentMessageUsage, chatMessages, chats, chatSessions, @@ -122,6 +123,45 @@ export async function countUserMessages( return rows[0]?.count ?? 0; } +/** + * The message count for one org in one billing period. Reads the standalone counter, + * never the chat rows, so a deleted chat can't lower it within the period. `period` is + * a UTC calendar month, "YYYY-MM"; the caller chooses it. + */ +export async function getAgentMessageUsage( + db: DashboardAgentDb, + params: { organizationId: string; period: string } +): Promise { + const rows = await db + .select({ count: agentMessageUsage.count }) + .from(agentMessageUsage) + .where( + and( + eq(agentMessageUsage.organizationId, params.organizationId), + eq(agentMessageUsage.period, params.period) + ) + ) + .limit(1); + return rows[0]?.count ?? 0; +} + +/** Bump the counter by one, creating the period row on first use. Returns the new count. */ +export async function incrementAgentMessageUsage( + db: DashboardAgentDb, + params: { organizationId: string; period: string; by?: number } +): Promise { + const by = params.by ?? 1; + const rows = await db + .insert(agentMessageUsage) + .values({ organizationId: params.organizationId, period: params.period, count: by }) + .onConflictDoUpdate({ + target: [agentMessageUsage.organizationId, agentMessageUsage.period], + set: { count: sql`${agentMessageUsage.count} + ${by}`, updatedAt: sql`now()` }, + }) + .returning({ count: agentMessageUsage.count }); + return rows[0]?.count ?? by; +} + /** * Chats whose transcript moved on after their owner last looked. A watch wake is one way * that happens; an answer that landed while the panel was closed is another, and the panel diff --git a/internal-packages/dashboard-agent-db/src/schema.ts b/internal-packages/dashboard-agent-db/src/schema.ts index 0943f838d1..73dafc64b0 100644 --- a/internal-packages/dashboard-agent-db/src/schema.ts +++ b/internal-packages/dashboard-agent-db/src/schema.ts @@ -183,6 +183,23 @@ export const investigations = dashboardAgentSchema.table( ] ); +/** + * Per-(org, period) message counter. Deliberately not joined to chats: deleting a chat + * must not free quota inside the period. `period` is a UTC calendar month, "YYYY-MM". + * Org id is a main-DB id with no FK. + */ +export const agentMessageUsage = dashboardAgentSchema.table( + "agent_message_usage", + { + organizationId: text("organization_id").notNull(), + period: text("period").notNull(), + count: integer("count").notNull().default(0), + createdAt: timestamp("created_at", { withTimezone: true }).notNull().defaultNow(), + updatedAt: timestamp("updated_at", { withTimezone: true }).notNull().defaultNow(), + }, + (t) => [primaryKey({ columns: [t.organizationId, t.period] })] +); + export type Chat = typeof chats.$inferSelect; export type NewChat = typeof chats.$inferInsert; export type ChatMessage = typeof chatMessages.$inferSelect; @@ -193,3 +210,5 @@ export type ChatTurnEval = typeof chatTurnEvals.$inferSelect; export type NewChatTurnEval = typeof chatTurnEvals.$inferInsert; export type Investigation = typeof investigations.$inferSelect; export type NewInvestigation = typeof investigations.$inferInsert; +export type AgentMessageUsage = typeof agentMessageUsage.$inferSelect; +export type NewAgentMessageUsage = typeof agentMessageUsage.$inferInsert;