diff --git a/web/src/app/admin/token-rate-limits/CreateRateLimitModal.test.tsx b/web/src/app/admin/token-rate-limits/CreateRateLimitModal.test.tsx new file mode 100644 index 0000000000..8310215ec3 --- /dev/null +++ b/web/src/app/admin/token-rate-limits/CreateRateLimitModal.test.tsx @@ -0,0 +1,19 @@ +import { render, screen, setupUser } from "@tests/setup/test-utils"; +import CreateRateLimitModal from "@/app/admin/token-rate-limits/CreateRateLimitModal"; + +test("requires at least one enforcement budget", async () => { + const user = setupUser(); + const onSubmit = jest.fn().mockResolvedValue(undefined); + + render( + + ); + + await user.type(screen.getByLabelText("Time Window (UTC Days)"), "1"); + await user.click(screen.getByRole("button", { name: "Create" })); + + expect( + await screen.findByText("Set a token budget and/or a cost budget") + ).toBeInTheDocument(); + expect(onSubmit).not.toHaveBeenCalled(); +}); diff --git a/web/src/app/admin/token-rate-limits/CreateRateLimitModal.tsx b/web/src/app/admin/token-rate-limits/CreateRateLimitModal.tsx index 06d50b6cdc..9e48483818 100644 --- a/web/src/app/admin/token-rate-limits/CreateRateLimitModal.tsx +++ b/web/src/app/admin/token-rate-limits/CreateRateLimitModal.tsx @@ -19,9 +19,10 @@ interface CreateRateLimitModalProps { onSubmit: ( target_scope: Scope, period_hours: number, - token_budget: number, + token_budget: number | null, + cost_budget_cents: number | null, group_id: number - ) => void; + ) => Promise; forSpecificScope?: Scope; forSpecificUserGroup?: number; } @@ -72,6 +73,7 @@ export default function CreateRateLimitModal({ enabled: true, period_days: "", token_budget: "", + cost_budget_dollars: "", target_scope: forSpecificScope || Scope.GLOBAL, user_group_id: forSpecificUserGroup, }} @@ -81,8 +83,26 @@ export default function CreateRateLimitModal({ .integer("Time Window must be a whole number of days") .min(1, "Time Window must be at least 1 day"), token_budget: Yup.number() - .required("Token Budget is a required field") - .min(1, "Token Budget must be at least 1"), + // Empty (no token budget) is allowed — a cost-only limit. Without + // this, "" coerces to NaN and trips .min(1) even when only a cost + // budget is set. Mirrors the cost_budget_dollars transform below. + .transform((value, original) => + original === "" ? undefined : value + ) + .min(1, "Token Budget must be at least 1") + .test( + "budget-required", + "Set a token budget and/or a cost budget", + (value, context) => + value != null || context.parent.cost_budget_dollars != null + ), + cost_budget_dollars: Yup.number() + // Empty (no cost budget) is allowed; a 0 would make the gate fire + // on the first request (cost_since >= 0 is always true). + .transform((value, original) => + original === "" ? undefined : value + ) + .moreThan(0, "Cost Budget must be greater than 0"), target_scope: Yup.string().required( "Target Scope is a required field" ), @@ -98,15 +118,22 @@ export default function CreateRateLimitModal({ } ), })} - onSubmit={async (values, formikHelpers) => { - formikHelpers.setSubmitting(true); - onSubmit( + onSubmit={async (values) => { + // Empty token field → null (cost-only); the gate skips a null budget. + // Sending 0 would mean "0-token limit" and block every request. + const tokenBudget = + values.token_budget === "" ? null : Number(values.token_budget); + const costBudgetCents = + values.cost_budget_dollars === "" + ? null + : Math.round(Number(values.cost_budget_dollars) * 100); + await onSubmit( values.target_scope, Number(values.period_days) * HOURS_PER_DAY, - Number(values.token_budget), + tokenBudget, + costBudgetCents, Number(values.user_group_id) ); - return formikHelpers.setSubmitting(false); }} > {({ isSubmitting, values, setFieldValue }) => ( @@ -147,7 +174,13 @@ export default function CreateRateLimitModal({ /> + diff --git a/web/src/app/admin/token-rate-limits/TokenRateLimitTables.test.tsx b/web/src/app/admin/token-rate-limits/TokenRateLimitTables.test.tsx new file mode 100644 index 0000000000..a760df0921 --- /dev/null +++ b/web/src/app/admin/token-rate-limits/TokenRateLimitTables.test.tsx @@ -0,0 +1,92 @@ +import { render, screen, setupUser, waitFor } from "@tests/setup/test-utils"; +import { formatPeriod } from "@/app/admin/token-rate-limits/TokenRateLimitTables"; +import { TokenRateLimitTable } from "@/app/admin/token-rate-limits/TokenRateLimitTables"; +import type { TokenRateLimitDisplay } from "@/app/admin/token-rate-limits/types"; + +const mockDelete = jest.fn(); +const mockMutate = jest.fn(); +const mockToastError = jest.fn(); +const mockUpdate = jest.fn(); + +jest.mock("@/app/admin/token-rate-limits/lib", () => ({ + deleteTokenRateLimit: (...args: unknown[]) => mockDelete(...args), + updateTokenRateLimit: (...args: unknown[]) => mockUpdate(...args), +})); + +jest.mock("@opal/layouts/toast/store", () => ({ + toast: { + error: (...args: unknown[]) => mockToastError(...args), + }, +})); + +jest.mock("swr", () => ({ + __esModule: true, + ...jest.requireActual("swr"), + mutate: (...args: unknown[]) => mockMutate(...args), +})); + +function tokenRateLimit( + periodHours: number, + tokenBudget: number | null, + costBudgetCents: number | null +): TokenRateLimitDisplay { + return { + token_id: 1, + enabled: true, + token_budget: tokenBudget, + period_hours: periodHours, + cost_budget_cents: costBudgetCents, + }; +} + +test.each([ + ["token-only", tokenRateLimit(24, 1, null), "1 UTC day"], + ["cost-only", tokenRateLimit(24, null, 100), "1 UTC day"], + ["dual", tokenRateLimit(48, 1, 100), "2 UTC days"], +])("%s limits display UTC-day windows", (_name, limit, expected) => { + expect(formatPeriod(limit)).toBe(expected); +}); + +describe("token rate limit mutation failures", () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + test("shows an update error instead of rejecting the click handler", async () => { + const user = setupUser(); + mockUpdate.mockRejectedValue(new Error("Update failed")); + render( + + ); + + await user.click(screen.getByRole("checkbox")); + + await waitFor(() => + expect(mockToastError).toHaveBeenCalledWith("Update failed") + ); + expect(mockMutate).not.toHaveBeenCalled(); + }); + + test("shows a delete error instead of rejecting the click handler", async () => { + const user = setupUser(); + mockDelete.mockRejectedValue(new Error("Delete failed")); + render( + + ); + + await user.click(screen.getByRole("button")); + + await waitFor(() => + expect(mockToastError).toHaveBeenCalledWith("Delete failed") + ); + expect(mockMutate).not.toHaveBeenCalled(); + }); +}); diff --git a/web/src/app/admin/token-rate-limits/TokenRateLimitTables.tsx b/web/src/app/admin/token-rate-limits/TokenRateLimitTables.tsx index 53698a5882..160e0f1c91 100644 --- a/web/src/app/admin/token-rate-limits/TokenRateLimitTables.tsx +++ b/web/src/app/admin/token-rate-limits/TokenRateLimitTables.tsx @@ -10,7 +10,7 @@ import { import Title from "@/components/ui/title"; import { DeleteButton } from "@/components/DeleteButton"; import { deleteTokenRateLimit, updateTokenRateLimit } from "./lib"; -import { PageLoader } from "@opal/layouts"; +import { PageLoader, toast } from "@opal/layouts"; import { TokenRateLimitDisplay } from "./types"; import { errorHandlingFetcher } from "@/lib/fetcher"; import useSWR, { mutate } from "swr"; @@ -21,16 +21,12 @@ import { Spacer } from "@opal/components"; const HOURS_PER_DAY = 24; const UTC_DAY_LABEL = "UTC day"; -const HOUR_LABEL = "hour"; -const COST_ONLY_LABEL = "Cost only"; +const UPDATE_ERROR_MESSAGE = "Failed to update token rate limit"; +const DELETE_ERROR_MESSAGE = "Failed to delete token rate limit"; -function formatPeriod(tokenRateLimit: TokenRateLimitDisplay): string { - const isTokenBudget = tokenRateLimit.token_budget !== null; - const value = isTokenBudget - ? tokenRateLimit.period_hours / HOURS_PER_DAY - : tokenRateLimit.period_hours; - const unit = isTokenBudget ? UTC_DAY_LABEL : HOUR_LABEL; - return `${value} ${unit}${value === 1 ? "" : "s"}`; +export function formatPeriod(tokenRateLimit: TokenRateLimitDisplay): string { + const days = tokenRateLimit.period_hours / HOURS_PER_DAY; + return `${days} ${UTC_DAY_LABEL}${days === 1 ? "" : "s"}`; } type TokenRateLimitTableArgs = { @@ -55,7 +51,7 @@ export const TokenRateLimitTable = ({ tokenRateLimits[0] !== undefined && tokenRateLimits[0].group_name !== undefined; - const handleEnabledChange = (id: number) => { + const handleEnabledChange = async (id: number) => { const tokenRateLimit = tokenRateLimits.find( (tokenRateLimit) => tokenRateLimit.token_id === id ); @@ -64,19 +60,31 @@ export const TokenRateLimitTable = ({ return; } - updateTokenRateLimit(id, { - token_budget: tokenRateLimit.token_budget, - period_hours: tokenRateLimit.period_hours, - enabled: !tokenRateLimit.enabled, - }).then(() => { - mutate(fetchUrl); - }); + try { + await updateTokenRateLimit(id, { + token_budget: tokenRateLimit.token_budget, + period_hours: tokenRateLimit.period_hours, + cost_budget_cents: tokenRateLimit.cost_budget_cents, + enabled: !tokenRateLimit.enabled, + }); + await mutate(fetchUrl); + } catch (error) { + toast.error( + error instanceof Error ? error.message : UPDATE_ERROR_MESSAGE + ); + } }; - const handleDelete = (id: number) => - deleteTokenRateLimit(id).then(() => { - mutate(fetchUrl); - }); + const handleDelete = async (id: number) => { + try { + await deleteTokenRateLimit(id); + await mutate(fetchUrl); + } catch (error) { + toast.error( + error instanceof Error ? error.message : DELETE_ERROR_MESSAGE + ); + } + }; if (tokenRateLimits.length === 0) { return ( @@ -116,7 +124,8 @@ export const TokenRateLimitTable = ({ Enabled {shouldRenderGroupName() && Group Name} Time Window - Token Budget (Thousands) + Token Budget + Cost Budget (USD) {isAdmin && Delete} @@ -162,9 +171,15 @@ export const TokenRateLimitTable = ({ )} {formatPeriod(tokenRateLimit)} - {tokenRateLimit.token_budget === null - ? COST_ONLY_LABEL - : `${tokenRateLimit.token_budget} thousand tokens`} + {tokenRateLimit.token_budget != null + ? (tokenRateLimit.token_budget * 1000).toLocaleString() + + " tokens" + : "—"} + + + {tokenRateLimit.cost_budget_cents != null + ? "$" + (tokenRateLimit.cost_budget_cents / 100).toFixed(2) + : "—"} {isAdmin && ( diff --git a/web/src/app/admin/token-rate-limits/lib.test.tsx b/web/src/app/admin/token-rate-limits/lib.test.tsx new file mode 100644 index 0000000000..bb76f99f15 --- /dev/null +++ b/web/src/app/admin/token-rate-limits/lib.test.tsx @@ -0,0 +1,61 @@ +import { + deleteTokenRateLimit, + insertGlobalTokenRateLimit, + insertGroupTokenRateLimit, + insertUserTokenRateLimit, + updateTokenRateLimit, +} from "@/app/admin/token-rate-limits/lib"; +import type { TokenRateLimitArgs } from "@/app/admin/token-rate-limits/types"; + +const tokenRateLimit: TokenRateLimitArgs = { + enabled: true, + token_budget: 1, + period_hours: 24, + cost_budget_cents: null, +}; +const ERROR_DETAIL = "Database unavailable"; + +const mutationCases: [string, () => Promise][] = [ + ["global create", () => insertGlobalTokenRateLimit(tokenRateLimit)], + ["user create", () => insertUserTokenRateLimit(tokenRateLimit)], + ["group create", () => insertGroupTokenRateLimit(tokenRateLimit, 1)], + ["update", () => updateTokenRateLimit(1, tokenRateLimit)], + ["delete", () => deleteTokenRateLimit(1)], +]; + +describe("token rate limit mutations", () => { + const originalFetch = global.fetch; + const fetchMock = jest.fn(); + + beforeAll(() => { + global.fetch = fetchMock; + }); + + afterEach(() => { + fetchMock.mockReset(); + }); + + afterAll(() => { + global.fetch = originalFetch; + }); + + test.each(mutationCases)( + "%s throws the backend detail", + async (_, mutate) => { + fetchMock.mockResolvedValue({ + ok: false, + json: jest.fn().mockResolvedValue({ detail: ERROR_DETAIL }), + } as unknown as Response); + + await expect(mutate()).rejects.toThrow(ERROR_DETAIL); + } + ); + + test("resolves after a successful response", async () => { + fetchMock.mockResolvedValue({ ok: true } as Response); + + await expect( + insertGlobalTokenRateLimit(tokenRateLimit) + ).resolves.toBeUndefined(); + }); +}); diff --git a/web/src/app/admin/token-rate-limits/lib.ts b/web/src/app/admin/token-rate-limits/lib.ts index f63db82c82..60e293c6e2 100644 --- a/web/src/app/admin/token-rate-limits/lib.ts +++ b/web/src/app/admin/token-rate-limits/lib.ts @@ -1,64 +1,89 @@ +import { parseErrorDetail } from "@/lib/fetcher"; import { TokenRateLimitArgs } from "./types"; const API_PREFIX = "/api/admin/token-rate-limits"; +const CREATE_ERROR_MESSAGE = "Failed to create token rate limit"; +const UPDATE_ERROR_MESSAGE = "Failed to update token rate limit"; +const DELETE_ERROR_MESSAGE = "Failed to delete token rate limit"; // Global Token Limits export const insertGlobalTokenRateLimit = async ( tokenRateLimit: TokenRateLimitArgs -) => { - return await fetch(`${API_PREFIX}/global`, { +): Promise => { + const response = await fetch(`${API_PREFIX}/global`, { method: "POST", headers: { "Content-Type": "application/json", }, body: JSON.stringify(tokenRateLimit), }); + if (!response.ok) { + throw new Error(await parseErrorDetail(response, CREATE_ERROR_MESSAGE)); + } }; // User Token Limits export const insertUserTokenRateLimit = async ( tokenRateLimit: TokenRateLimitArgs -) => { - return await fetch(`${API_PREFIX}/users`, { +): Promise => { + const response = await fetch(`${API_PREFIX}/users`, { method: "POST", headers: { "Content-Type": "application/json", }, body: JSON.stringify(tokenRateLimit), }); + if (!response.ok) { + throw new Error(await parseErrorDetail(response, CREATE_ERROR_MESSAGE)); + } }; // User Group Token Limits (EE Only) export const insertGroupTokenRateLimit = async ( tokenRateLimit: TokenRateLimitArgs, group_id: number -) => { - return await fetch(`${API_PREFIX}/user-group/${group_id}`, { +): Promise => { + const response = await fetch(`${API_PREFIX}/user-group/${group_id}`, { method: "POST", headers: { "Content-Type": "application/json", }, body: JSON.stringify(tokenRateLimit), }); + if (!response.ok) { + throw new Error(await parseErrorDetail(response, CREATE_ERROR_MESSAGE)); + } }; // Common Endpoints -export const deleteTokenRateLimit = async (token_rate_limit_id: number) => { - return await fetch(`${API_PREFIX}/rate-limit/${token_rate_limit_id}`, { - method: "DELETE", - }); +export const deleteTokenRateLimit = async ( + token_rate_limit_id: number +): Promise => { + const response = await fetch( + `${API_PREFIX}/rate-limit/${token_rate_limit_id}`, + { method: "DELETE" } + ); + if (!response.ok) { + throw new Error(await parseErrorDetail(response, DELETE_ERROR_MESSAGE)); + } }; export const updateTokenRateLimit = async ( token_rate_limit_id: number, tokenRateLimit: TokenRateLimitArgs -) => { - return await fetch(`${API_PREFIX}/rate-limit/${token_rate_limit_id}`, { - method: "PUT", - headers: { - "Content-Type": "application/json", - }, - body: JSON.stringify(tokenRateLimit), - }); +): Promise => { + const response = await fetch( + `${API_PREFIX}/rate-limit/${token_rate_limit_id}`, + { + method: "PUT", + headers: { + "Content-Type": "application/json", + }, + body: JSON.stringify(tokenRateLimit), + } + ); + if (!response.ok) { + throw new Error(await parseErrorDetail(response, UPDATE_ERROR_MESSAGE)); + } }; diff --git a/web/src/app/admin/token-rate-limits/page.tsx b/web/src/app/admin/token-rate-limits/page.tsx index 4fd1120659..eaa0f2f7c5 100644 --- a/web/src/app/admin/token-rate-limits/page.tsx +++ b/web/src/app/admin/token-rate-limits/page.tsx @@ -36,28 +36,38 @@ const USER_GROUP_DESCRIPTION = all users in the group will be temporarily blocked from spending tokens, regardless \ of their individual limits. If a user is in multiple groups, the most lenient limit \ will apply."; +const CREATE_ERROR_MESSAGE = "Failed to create token rate limit"; const handleCreateTokenRateLimit = async ( target_scope: Scope, period_hours: number, - token_budget: number, + token_budget: number | null, + cost_budget_cents: number | null, group_id: number = -1 -) => { +): Promise => { const tokenRateLimitArgs = { enabled: true, token_budget: token_budget, period_hours: period_hours, + cost_budget_cents: cost_budget_cents, }; if (target_scope === Scope.GLOBAL) { - return await insertGlobalTokenRateLimit(tokenRateLimitArgs); - } else if (target_scope === Scope.USER) { - return await insertUserTokenRateLimit(tokenRateLimitArgs); - } else if (target_scope === Scope.USER_GROUP) { - return await insertGroupTokenRateLimit(tokenRateLimitArgs, group_id); - } else { - throw new Error(`Invalid target_scope: ${target_scope}`); + await insertGlobalTokenRateLimit(tokenRateLimitArgs); + return; } + + if (target_scope === Scope.USER) { + await insertUserTokenRateLimit(tokenRateLimitArgs); + return; + } + + if (target_scope === Scope.USER_GROUP) { + await insertGroupTokenRateLimit(tokenRateLimitArgs, group_id); + return; + } + + throw new Error(`Invalid target_scope: ${target_scope}`); }; function Main() { @@ -79,26 +89,29 @@ function Main() { } }; - const handleSubmit = ( + const handleSubmit = async ( target_scope: Scope, period_hours: number, - token_budget: number, + token_budget: number | null, + cost_budget_cents: number | null, group_id: number = -1 - ) => { - handleCreateTokenRateLimit( - target_scope, - period_hours, - token_budget, - group_id - ) - .then(() => { - setModalIsOpen(false); - toast.success("Token rate limit created!"); - updateTable(target_scope); - }) - .catch((error) => { - toast.error(error.message); - }); + ): Promise => { + try { + await handleCreateTokenRateLimit( + target_scope, + period_hours, + token_budget, + cost_budget_cents, + group_id + ); + setModalIsOpen(false); + toast.success("Token rate limit created!"); + updateTable(target_scope); + } catch (error) { + toast.error( + error instanceof Error ? error.message : CREATE_ERROR_MESSAGE + ); + } }; return ( diff --git a/web/src/app/admin/token-rate-limits/types.ts b/web/src/app/admin/token-rate-limits/types.ts index 4fd68e1949..991d03fcb5 100644 --- a/web/src/app/admin/token-rate-limits/types.ts +++ b/web/src/app/admin/token-rate-limits/types.ts @@ -8,6 +8,8 @@ export interface TokenRateLimitArgs { enabled: boolean; token_budget: number | null; period_hours: number; + // Cents at the API boundary; the UI collects dollars and converts. + cost_budget_cents?: number | null; } export interface TokenRateLimit { @@ -15,6 +17,7 @@ export interface TokenRateLimit { enabled: boolean; token_budget: number | null; period_hours: number; + cost_budget_cents: number | null; } export interface TokenRateLimitDisplay extends TokenRateLimit { diff --git a/web/src/app/app/interfaces.ts b/web/src/app/app/interfaces.ts index 112d0cd6ca..fed327e91e 100644 --- a/web/src/app/app/interfaces.ts +++ b/web/src/app/app/interfaces.ts @@ -305,6 +305,18 @@ export interface StreamingError { details?: Record; } +// error_code emitted by the backend usage rate-limiter (429). Branch on this to +// show the dedicated usage-limit banner instead of the generic chat error. +export const RATE_LIMITED_ERROR_CODE = "RATE_LIMITED"; + +// Shape of StreamingError.details for a RATE_LIMITED error — mirrors the 429 +// JSON body so the banner can compute a human-friendly reset time. +export interface RateLimitDetails { + scope?: string; + reset_at?: string; // ISO timestamp + retry_after_seconds?: number; +} + export interface InputPrompt { id: number; prompt: string; diff --git a/web/src/app/app/message/Resubmit.test.tsx b/web/src/app/app/message/Resubmit.test.tsx new file mode 100644 index 0000000000..66e5fd36a5 --- /dev/null +++ b/web/src/app/app/message/Resubmit.test.tsx @@ -0,0 +1,32 @@ +import { act, render, screen } from "@tests/setup/test-utils"; +import { ErrorBanner } from "@/app/app/message/Resubmit"; +import { RATE_LIMITED_ERROR_CODE } from "@/app/app/interfaces"; + +describe("rate-limit error banner", () => { + beforeEach(() => { + jest.useFakeTimers(); + jest.setSystemTime(new Date("2026-07-28T12:00:00Z")); + }); + + afterEach(() => { + jest.useRealTimers(); + }); + + test("updates when the reset time is reached", () => { + render( + + ); + + expect(screen.getByText(/Resets in 1 minute/)).toBeInTheDocument(); + + act(() => { + jest.advanceTimersByTime(30_000); + }); + + expect(screen.getByText("You can try again now.")).toBeInTheDocument(); + }); +}); diff --git a/web/src/app/app/message/Resubmit.tsx b/web/src/app/app/message/Resubmit.tsx index e5eb643ca1..367e17f5df 100644 --- a/web/src/app/app/message/Resubmit.tsx +++ b/web/src/app/app/message/Resubmit.tsx @@ -1,9 +1,101 @@ -import { useState } from "react"; +import { useEffect, useMemo, useState } from "react"; import { Alert, AlertDescription, AlertTitle } from "@/components/ui/alert"; import { SvgChevronDown, SvgChevronRight } from "@opal/icons"; import { Button } from "@opal/components"; import { CopyButton } from "@opal/components"; import { getErrorIcon, getErrorTitle } from "./errorHelpers"; +import { + RateLimitDetails, + RATE_LIMITED_ERROR_CODE, +} from "@/app/app/interfaces"; + +const COUNTDOWN_TICK_MS = 1_000; + +function formatRateLimitReset(resetMs: number, nowMs: number): string { + const remainingMs = resetMs - nowMs; + if (remainingMs <= 0) return "You can try again now."; + + const plural = (n: number, unit: string) => + `${n} ${unit}${n === 1 ? "" : "s"}`; + const minutes = Math.ceil(remainingMs / 60_000); + const hours = Math.ceil(remainingMs / 3_600_000); + const days = Math.ceil(remainingMs / 86_400_000); + const relative = + minutes < 60 + ? plural(minutes, "minute") + : hours < 48 + ? plural(hours, "hour") + : plural(days, "day"); + const resetDate = new Date(resetMs); + // For multi-day resets a date is clearer than just a clock time. + const at = + days >= 2 + ? resetDate.toLocaleDateString(undefined, { + month: "short", + day: "numeric", + }) + : resetDate.toLocaleTimeString(undefined, { + hour: "numeric", + minute: "2-digit", + }); + return `Resets in ${relative} (${at}).`; +} + +function resolveResetMs( + resetAt?: string, + retryAfterSeconds?: number +): number | null { + if (resetAt) { + const parsed = Date.parse(resetAt); + if (!Number.isNaN(parsed)) return parsed; + } + if (typeof retryAfterSeconds === "number") { + return Date.now() + retryAfterSeconds * 1_000; + } + return null; +} + +interface RateLimitBannerProps { + error: string; + errorCode: string; + details: RateLimitDetails; +} + +function RateLimitBanner({ error, errorCode, details }: RateLimitBannerProps) { + const [nowMs, setNowMs] = useState(Date.now()); + const resetMs = useMemo( + () => resolveResetMs(details.reset_at, details.retry_after_seconds), + [details.reset_at, details.retry_after_seconds] + ); + + useEffect(() => { + if (resetMs === null) return; + + setNowMs(Date.now()); + const interval = window.setInterval( + () => setNowMs(Date.now()), + COUNTDOWN_TICK_MS + ); + return () => window.clearInterval(interval); + }, [resetMs]); + + const resetLine = + resetMs === null ? null : formatRateLimitReset(resetMs, nowMs); + return ( +
+ + {getErrorIcon(errorCode)} + {getErrorTitle(errorCode)} + + {error || "You've reached your usage limit."} + {resetLine && ( + {resetLine} + )} + + +
+ ); +} interface ResubmitProps { resubmit: () => void; @@ -37,6 +129,16 @@ export const ErrorBanner = ({ }) => { const [isStackTraceExpanded, setIsStackTraceExpanded] = useState(false); + if (errorCode === RATE_LIMITED_ERROR_CODE) { + return ( + + ); + } + return (
diff --git a/web/src/app/app/message/errorHelpers.tsx b/web/src/app/app/message/errorHelpers.tsx index c305c630e7..6b775c592a 100644 --- a/web/src/app/app/message/errorHelpers.tsx +++ b/web/src/app/app/message/errorHelpers.tsx @@ -6,6 +6,7 @@ import { AlertCircle, Clock, Lock, Wifi, Server } from "lucide-react"; export const getErrorIcon = (errorCode?: string) => { switch (errorCode) { case "RATE_LIMIT": + case "RATE_LIMITED": return ; case "AUTH_ERROR": case "PERMISSION_DENIED": @@ -28,6 +29,8 @@ export const getErrorTitle = (errorCode?: string) => { switch (errorCode) { case "RATE_LIMIT": return "Rate Limit Exceeded"; + case "RATE_LIMITED": + return "Usage limit reached"; case "AUTH_ERROR": return "Authentication Error"; case "PERMISSION_DENIED": diff --git a/web/src/app/app/services/lib.tsx b/web/src/app/app/services/lib.tsx index febfc176a7..43ac30a180 100644 --- a/web/src/app/app/services/lib.tsx +++ b/web/src/app/app/services/lib.tsx @@ -14,6 +14,8 @@ import { Message, MessageResponseIDInfo, MultiModelMessageResponseIDInfo, + RateLimitDetails, + RATE_LIMITED_ERROR_CODE, ResearchType, RetrievalType, StreamingError, @@ -215,6 +217,31 @@ export async function* sendMessage({ if (!response.ok) { const data = await response.json().catch(() => ({})); + + // Surface the usage rate-limit (429) as a structured StreamingError packet + // so the chat UI can render the dedicated usage-limit banner. Throwing a + // bare Error here would flatten error_code/reset_at into an opaque string. + if ( + response.status === 429 && + data.error_code === RATE_LIMITED_ERROR_CODE + ) { + const rateLimitDetails: RateLimitDetails = { + scope: data.scope, + reset_at: data.reset_at, + retry_after_seconds: data.retry_after_seconds, + }; + const streamingError: StreamingError = { + error: data.detail ?? "You've reached your usage limit.", + stack_trace: "", + error_code: RATE_LIMITED_ERROR_CODE, + // Regenerate would just re-trip the limit, so this is not retryable. + is_retryable: false, + details: rateLimitDetails, + }; + yield streamingError; + return; + } + throw new Error(data.detail ?? `HTTP error! status: ${response.status}`); } diff --git a/web/src/refresh-components/inputs/InputNumber.tsx b/web/src/refresh-components/inputs/InputNumber.tsx index b897693661..b0ddfda400 100644 --- a/web/src/refresh-components/inputs/InputNumber.tsx +++ b/web/src/refresh-components/inputs/InputNumber.tsx @@ -48,6 +48,7 @@ export interface InputNumberProps { min?: number; max?: number; step?: number; + decimalPlaces?: number; defaultValue?: number; showReset?: boolean; variant?: Variants; @@ -62,6 +63,7 @@ export default function InputNumber({ min, max, step = 1, + decimalPlaces = 0, defaultValue, showReset = false, variant = "primary", @@ -74,6 +76,8 @@ export default function InputNumber({ value === null ? "" : String(value) ); const isDisabled = disabled || variant === "disabled"; + const inputPattern = + decimalPlaces === 0 ? "[0-9]*" : `[0-9]*[.]?[0-9]{0,${decimalPlaces}}`; // Sync input value when external value changes (e.g., from stepper buttons or reset) React.useEffect(() => { @@ -86,17 +90,19 @@ export default function InputNumber({ value !== null && (min === undefined || effectiveValue > min); const canReset = showReset && defaultValue !== undefined && value !== defaultValue; + const normalizePrecision = (newValue: number) => + Number(newValue.toFixed(decimalPlaces)); const handleIncrement = () => { if (canIncrement) { - const newValue = effectiveValue + step; + const newValue = normalizePrecision(effectiveValue + step); onChange(max !== undefined ? Math.min(newValue, max) : newValue); } }; const handleDecrement = () => { if (canDecrement) { - const newValue = effectiveValue - step; + const newValue = normalizePrecision(effectiveValue - step); onChange(min !== undefined ? Math.max(newValue, min) : newValue); } }; @@ -119,8 +125,7 @@ export default function InputNumber({ const handleInputChange = (e: React.ChangeEvent) => { const rawValue = e.target.value; - // Only allow digits (and empty string) - if (rawValue !== "" && !/^\d+$/.test(rawValue)) { + if (rawValue !== "" && !new RegExp(`^${inputPattern}$`).test(rawValue)) { return; } @@ -131,7 +136,10 @@ export default function InputNumber({ return; } - const val = parseInt(rawValue, 10); + const val = Number(rawValue); + if (!Number.isFinite(val)) { + return; + } let newValue = val; if (min !== undefined) newValue = Math.max(newValue, min); if (max !== undefined) newValue = Math.min(newValue, max); @@ -150,8 +158,8 @@ export default function InputNumber({ ([]); const [selectedAgentIds, setSelectedAgentIds] = useState([]); const [tokenLimits, setTokenLimits] = useState([ - { tokenBudget: null, periodDays: null }, + { + tokenId: null, + enabled: true, + tokenBudget: null, + periodDays: null, + costBudgetDollars: null, + }, ]); const { rows: allRows, isLoading, error } = useGroupMemberCandidates(); diff --git a/web/src/views/admin/GroupsPage/EditGroupPage.tsx b/web/src/views/admin/GroupsPage/EditGroupPage.tsx index 573a3f8c3c..f7be6fca95 100644 --- a/web/src/views/admin/GroupsPage/EditGroupPage.tsx +++ b/web/src/views/admin/GroupsPage/EditGroupPage.tsx @@ -101,7 +101,13 @@ function EditGroupPage({ groupId }: EditGroupPageProps) { const [selectedDocSetIds, setSelectedDocSetIds] = useState([]); const [selectedAgentIds, setSelectedAgentIds] = useState([]); const [tokenLimits, setTokenLimits] = useState([ - { tokenBudget: null, periodDays: null }, + { + tokenId: null, + enabled: true, + tokenBudget: null, + periodDays: null, + costBudgetDollars: null, + }, ]); const [showDeleteModal, setShowDeleteModal] = useState(false); const [isDeleting, setIsDeleting] = useState(false); @@ -141,8 +147,12 @@ function EditGroupPage({ groupId }: EditGroupPageProps) { if (tokenRateLimits && tokenRateLimits.length > 0) { setTokenLimits( tokenRateLimits.map((trl) => ({ + tokenId: trl.token_id, + enabled: trl.enabled, tokenBudget: trl.token_budget, periodDays: trl.period_hours / HOURS_PER_DAY, + costBudgetDollars: + trl.cost_budget_cents != null ? trl.cost_budget_cents / 100 : null, })) ); } diff --git a/web/src/views/admin/GroupsPage/TokenLimitSection.test.tsx b/web/src/views/admin/GroupsPage/TokenLimitSection.test.tsx new file mode 100644 index 0000000000..d2a06efbce --- /dev/null +++ b/web/src/views/admin/GroupsPage/TokenLimitSection.test.tsx @@ -0,0 +1,50 @@ +import { useState } from "react"; +import { render, screen, setupUser } from "@tests/setup/test-utils"; +import TokenLimitSection from "@/views/admin/GroupsPage/TokenLimitSection"; +import type { TokenLimit } from "@/views/admin/GroupsPage/TokenLimitSection"; + +function TokenLimitSectionHarness() { + const [limits, setLimits] = useState([ + { + tokenId: null, + enabled: true, + tokenBudget: null, + periodDays: null, + costBudgetDollars: null, + }, + ]); + + return ; +} + +test("accepts a cost limit with cents", async () => { + const user = setupUser(); + render(); + + const costLimit = screen.getByPlaceholderText("Cost limit"); + await user.type(costLimit, "1.23"); + await user.click(screen.getByRole("button", { name: "Add Limit" })); + + expect(costLimit).toHaveValue("1.23"); + expect(screen.getAllByPlaceholderText("Cost limit")).toHaveLength(2); +}); + +test("enforces a positive cost limit", async () => { + const user = setupUser(); + render(); + + const costLimit = screen.getByPlaceholderText("Cost limit"); + await user.type(costLimit, "0"); + + expect(costLimit).toHaveValue("0.01"); +}); + +test("enforces a positive token limit", async () => { + const user = setupUser(); + render(); + + const tokenLimit = screen.getByPlaceholderText("Token limit (thousands)"); + await user.type(tokenLimit, "0"); + + expect(tokenLimit).toHaveValue("1"); +}); diff --git a/web/src/views/admin/GroupsPage/TokenLimitSection.tsx b/web/src/views/admin/GroupsPage/TokenLimitSection.tsx index c49ee4ca08..ab3ac819ec 100644 --- a/web/src/views/admin/GroupsPage/TokenLimitSection.tsx +++ b/web/src/views/admin/GroupsPage/TokenLimitSection.tsx @@ -18,10 +18,15 @@ import SimpleCollapsible from "@/refresh-components/SimpleCollapsible"; // --------------------------------------------------------------------------- export interface TokenLimit { + tokenId: number | null; + enabled: boolean; tokenBudget: number | null; periodDays: number | null; + costBudgetDollars: number | null; } +type TokenLimitValueField = "tokenBudget" | "periodDays" | "costBudgetDollars"; + interface TokenLimitSectionProps { limits: TokenLimit[]; onLimitsChange: (limits: TokenLimit[]) => void; @@ -54,12 +59,24 @@ function TokenLimitSection({ function addLimit() { const emptyIndex = limits.findIndex( - (l) => l.tokenBudget === null && l.periodDays === null + (l) => + l.tokenBudget === null && + l.periodDays === null && + l.costBudgetDollars === null ); if (emptyIndex !== -1) return; const key = nextKeyRef.current++; keysRef.current = [...keysRef.current, key]; - onLimitsChange([...limits, { tokenBudget: null, periodDays: null }]); + onLimitsChange([ + ...limits, + { + tokenId: null, + enabled: true, + tokenBudget: null, + periodDays: null, + costBudgetDollars: null, + }, + ]); } function removeLimit(index: number) { @@ -69,7 +86,7 @@ function TokenLimitSection({ function updateLimit( index: number, - field: keyof TokenLimit, + field: TokenLimitValueField, value: number | null ) { onLimitsChange( @@ -103,7 +120,15 @@ function TokenLimitSection({ Token Limit - (thousand tokens) + (thousands) + +
+
+ + Cost Limit + + + (USD)
@@ -126,8 +151,18 @@ function TokenLimitSection({ updateLimit(i, "tokenBudget", v)} - min={0} - placeholder="Token limit in thousands" + min={1} + placeholder="Token limit (thousands)" + /> +
+
+ updateLimit(i, "costBudgetDollars", v)} + min={0.01} + step={0.01} + decimalPlaces={2} + placeholder="Cost limit" />
diff --git a/web/src/views/admin/GroupsPage/groupTokenLimitsSvc.test.tsx b/web/src/views/admin/GroupsPage/groupTokenLimitsSvc.test.tsx new file mode 100644 index 0000000000..e9302d94b9 --- /dev/null +++ b/web/src/views/admin/GroupsPage/groupTokenLimitsSvc.test.tsx @@ -0,0 +1,103 @@ +import { saveTokenLimits } from "@/views/admin/GroupsPage/svc"; + +describe("group token-limit persistence", () => { + const originalFetch = global.fetch; + const fetchMock = jest.fn(); + + beforeAll(() => { + global.fetch = fetchMock; + }); + + afterEach(() => { + fetchMock.mockReset(); + }); + + afterAll(() => { + global.fetch = originalFetch; + }); + + test.each([ + [ + "token", + { tokenBudget: 0, periodDays: 1, costBudgetDollars: null }, + { token_budget: 0, cost_budget_cents: null }, + ], + [ + "cost", + { tokenBudget: null, periodDays: 1, costBudgetDollars: 0 }, + { token_budget: null, cost_budget_cents: 0 }, + ], + ])( + "preserves an invalid zero %s budget for API validation", + async (_, limit, expected) => { + fetchMock.mockResolvedValue({ ok: true } as Response); + + await saveTokenLimits(7, [limit], []); + + const [url, options] = fetchMock.mock.calls[0] as [string, RequestInit]; + expect(url).toBe("/api/admin/token-rate-limits/user-group/7"); + expect(JSON.parse(String(options.body))).toEqual({ + enabled: true, + ...expected, + period_hours: 24, + }); + } + ); + + test("preserves limit identity when a middle row is removed", async () => { + fetchMock.mockResolvedValue({ ok: true } as Response); + const existing = [ + { + token_id: 10, + enabled: true, + token_budget: 100, + period_hours: 24, + cost_budget_cents: null, + }, + { + token_id: 20, + enabled: false, + token_budget: 200, + period_hours: 24, + cost_budget_cents: null, + }, + { + token_id: 30, + enabled: true, + token_budget: 300, + period_hours: 24, + cost_budget_cents: null, + }, + ]; + + await saveTokenLimits( + 7, + [ + { + tokenId: 10, + enabled: true, + tokenBudget: 100, + periodDays: 1, + costBudgetDollars: null, + }, + { + tokenId: 30, + enabled: true, + tokenBudget: 300, + periodDays: 1, + costBudgetDollars: null, + }, + ], + existing + ); + + expect(fetchMock.mock.calls.map(([url]) => url)).toEqual([ + "/api/admin/token-rate-limits/rate-limit/10", + "/api/admin/token-rate-limits/rate-limit/30", + "/api/admin/token-rate-limits/rate-limit/20", + ]); + expect( + fetchMock.mock.calls.map(([, options]) => (options as RequestInit).method) + ).toEqual(["PUT", "PUT", "DELETE"]); + }); +}); diff --git a/web/src/views/admin/GroupsPage/interfaces.ts b/web/src/views/admin/GroupsPage/interfaces.ts index b113a7306b..48140e3b5d 100644 --- a/web/src/views/admin/GroupsPage/interfaces.ts +++ b/web/src/views/admin/GroupsPage/interfaces.ts @@ -17,6 +17,8 @@ export interface MemberRow extends UserRow { export interface TokenRateLimitDisplay { token_id: number; enabled: boolean; - token_budget: number; + // null for a cost-only limit (matches the backend's nullable column). + token_budget: number | null; period_hours: number; + cost_budget_cents: number | null; } diff --git a/web/src/views/admin/GroupsPage/svc.ts b/web/src/views/admin/GroupsPage/svc.ts index 792eb8bd81..2097fc9258 100644 --- a/web/src/views/admin/GroupsPage/svc.ts +++ b/web/src/views/admin/GroupsPage/svc.ts @@ -202,15 +202,27 @@ async function updateDocSetGroupSharing( const HOURS_PER_DAY = 24; interface TokenLimitPayload { + tokenId?: number | null; + enabled?: boolean; tokenBudget: number | null; periodDays: number | null; + costBudgetDollars: number | null; } interface ExistingTokenLimit { token_id: number; enabled: boolean; - token_budget: number; + token_budget: number | null; period_hours: number; + cost_budget_cents: number | null; +} + +interface ValidTokenLimit { + tokenId: number | null; + enabled: boolean; + tokenBudget: number | null; + periodDays: number; + costBudgetCents: number | null; } async function saveTokenLimits( @@ -218,48 +230,56 @@ async function saveTokenLimits( limits: TokenLimitPayload[], existing: ExistingTokenLimit[] ): Promise { - // Filter to only valid (non-null) limits - const validLimits = limits.filter( - (l): l is { tokenBudget: number; periodDays: number } => - l.tokenBudget != null && l.periodDays != null - ); + const validLimits: ValidTokenLimit[] = limits + .map((l) => { + const costBudgetCents = + l.costBudgetDollars != null + ? Math.round(l.costBudgetDollars * 100) + : null; + return { + tokenId: l.tokenId ?? null, + enabled: l.enabled ?? true, + tokenBudget: l.tokenBudget, + periodDays: l.periodDays, + costBudgetCents, + }; + }) + .filter( + (l): l is ValidTokenLimit => + l.periodDays != null && + (l.tokenBudget != null || l.costBudgetCents != null) + ); - // Update existing limits (match by index position) - const toUpdate = Math.min(validLimits.length, existing.length); - for (let i = 0; i < toUpdate; i++) { - const limit = validLimits[i]!; - const existingLimit = existing[i]!; + for (const limit of validLimits.filter((item) => item.tokenId !== null)) { const updateRes = await fetch( - `/api/admin/token-rate-limits/rate-limit/${existingLimit.token_id}`, + `/api/admin/token-rate-limits/rate-limit/${limit.tokenId}`, { method: "PUT", headers: { "Content-Type": "application/json" }, body: JSON.stringify({ - enabled: existingLimit.enabled, + enabled: limit.enabled, token_budget: limit.tokenBudget, period_hours: limit.periodDays * HOURS_PER_DAY, + cost_budget_cents: limit.costBudgetCents, }), } ); if (!updateRes.ok) { - throw new Error( - `Failed to update token rate limit ${existingLimit.token_id}` - ); + throw new Error(`Failed to update token rate limit ${limit.tokenId}`); } } - // Create new limits beyond existing count - for (let i = toUpdate; i < validLimits.length; i++) { - const limit = validLimits[i]!; + for (const limit of validLimits.filter((item) => item.tokenId === null)) { const createRes = await fetch( `/api/admin/token-rate-limits/user-group/${groupId}`, { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify({ - enabled: true, + enabled: limit.enabled, token_budget: limit.tokenBudget, period_hours: limit.periodDays * HOURS_PER_DAY, + cost_budget_cents: limit.costBudgetCents, }), } ); @@ -268,9 +288,13 @@ async function saveTokenLimits( } } - // Delete excess existing limits - for (let i = toUpdate; i < existing.length; i++) { - const existingLimit = existing[i]!; + const retainedIds = new Set( + validLimits.flatMap((limit) => + limit.tokenId === null ? [] : [limit.tokenId] + ) + ); + for (const existingLimit of existing) { + if (retainedIds.has(existingLimit.token_id)) continue; const deleteRes = await fetch( `/api/admin/token-rate-limits/rate-limit/${existingLimit.token_id}`, { method: "DELETE" } diff --git a/web/tests/e2e/admin/token_rate_limit.spec.ts b/web/tests/e2e/admin/token_rate_limit.spec.ts new file mode 100644 index 0000000000..f5d2ce62e5 --- /dev/null +++ b/web/tests/e2e/admin/token_rate_limit.spec.ts @@ -0,0 +1,49 @@ +import { test, expect } from "@playwright/test"; +import { ChatPage } from "@tests/e2e/chat/ChatPage"; + +/** + * Rate-limit enforcement is a server-side gate, so this is a REAL e2e: it does + * NOT mock the chat endpoint (mocking would bypass the gate). It creates a tiny + * per-user token budget, sends real chats until usage accrues past it, and + * asserts the structured 429 renders as the "usage limit reached" banner. + * + * Requires a working LLM provider in the e2e environment (chats must actually + * produce token usage). The same admin user creates the limit and sends the + * chats, so the USER scope applies to them. + */ +test.use({ storageState: "admin_auth.json" }); + +const RATE_LIMIT_API = "/api/admin/token-rate-limits"; + +test.describe("usage budget blocks chat", () => { + let limitId: number | undefined; + + test.afterEach(async ({ page }) => { + if (limitId != null) { + await page.request.delete(`${RATE_LIMIT_API}/rate-limit/${limitId}`); + limitId = undefined; + } + }); + + test("a per-user token budget surfaces the usage-limit banner", async ({ + page, + }) => { + // 1) Tiny per-user token budget (1 => 1,000 tokens enforced). + const res = await page.request.post(`${RATE_LIMIT_API}/users`, { + data: { + enabled: true, + token_budget: 1, + period_hours: 168, + cost_budget_cents: null, + }, + }); + expect(res.ok()).toBeTruthy(); + const body = await res.json(); + limitId = body.token_id ?? body.id; + + const chat = new ChatPage(page); + await chat.goto(); + await chat.sendUntilUsageLimit(8); + await chat.expectAccountUsageLimit(); + }); +}); diff --git a/web/tests/e2e/chat/ChatPage.ts b/web/tests/e2e/chat/ChatPage.ts index 0f41b8f49b..061c9046b9 100644 --- a/web/tests/e2e/chat/ChatPage.ts +++ b/web/tests/e2e/chat/ChatPage.ts @@ -20,6 +20,7 @@ export class ChatPage { // Message collections readonly humanMessages: Locator; readonly aiMessages: Locator; + readonly usageLimitBanner: Locator; constructor(page: Page) { this.page = page; @@ -28,6 +29,7 @@ export class ChatPage { this.scrollContainer = page.getByTestId("chat-scroll-container"); this.humanMessages = page.locator("#onyx-human-message"); this.aiMessages = page.getByTestId("onyx-ai-message"); + this.usageLimitBanner = page.getByText(/you've reached the usage budget/i); } humanMessage(index = 0): Locator { @@ -86,4 +88,28 @@ export class ChatPage { async expectNoHumanMessages(): Promise { await expect(this.humanMessages).toHaveCount(0); } + + async sendUntilUsageLimit(maxTurns: number): Promise { + for ( + let turn = 0; + turn < maxTurns && !(await this.usageLimitBanner.isVisible()); + turn++ + ) { + await this.inputBar.fill(`write a few sentences about topic ${turn}`); + await this.inputBar.send(); + await Promise.race([ + this.usageLimitBanner + .waitFor({ state: "visible", timeout: 45_000 }) + .catch(() => {}), + this.aiMessage(turn) + .waitFor({ state: "visible", timeout: 45_000 }) + .catch(() => {}), + ]); + } + } + + async expectAccountUsageLimit(): Promise { + await expect(this.usageLimitBanner).toBeVisible(); + await expect(this.page.getByText(/your account/i)).toBeVisible(); + } }