mirror of
https://github.com/onyx-dot-app/onyx.git
synced 2026-10-22 03:14:06 -05:00
feat(usage): add cost limits and enforcement UI (#13600)
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Evan LohnandCursor authored and GitHub committed
1 parent
0ebb611c4c
commit
8b6288931f
23 files changed
+870
-120
No files matched your search
@@ -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(
|
||||
<CreateRateLimitModal isOpen setIsOpen={jest.fn()} onSubmit={onSubmit} />
|
||||
);
|
||||
|
||||
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();
|
||||
});
|
||||
@@ -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<void>;
|
||||
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({
|
||||
/>
|
||||
<TextFormField
|
||||
name="token_budget"
|
||||
label="Token Budget (Thousands)"
|
||||
label="Token Budget (Thousands, optional)"
|
||||
type="number"
|
||||
placeholder=""
|
||||
/>
|
||||
<TextFormField
|
||||
name="cost_budget_dollars"
|
||||
label="Cost Budget (USD per period, optional)"
|
||||
type="number"
|
||||
placeholder=""
|
||||
/>
|
||||
|
||||
@@ -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(
|
||||
<TokenRateLimitTable
|
||||
tokenRateLimits={[tokenRateLimit(24, 1, null)]}
|
||||
fetchUrl="/api/test-limits"
|
||||
isAdmin
|
||||
/>
|
||||
);
|
||||
|
||||
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(
|
||||
<TokenRateLimitTable
|
||||
tokenRateLimits={[tokenRateLimit(24, 1, null)]}
|
||||
fetchUrl="/api/test-limits"
|
||||
isAdmin
|
||||
/>
|
||||
);
|
||||
|
||||
await user.click(screen.getByRole("button"));
|
||||
|
||||
await waitFor(() =>
|
||||
expect(mockToastError).toHaveBeenCalledWith("Delete failed")
|
||||
);
|
||||
expect(mockMutate).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
@@ -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 = ({
|
||||
<TableHead>Enabled</TableHead>
|
||||
{shouldRenderGroupName() && <TableHead>Group Name</TableHead>}
|
||||
<TableHead>Time Window</TableHead>
|
||||
<TableHead>Token Budget (Thousands)</TableHead>
|
||||
<TableHead>Token Budget</TableHead>
|
||||
<TableHead>Cost Budget (USD)</TableHead>
|
||||
{isAdmin && <TableHead>Delete</TableHead>}
|
||||
</TableRow>
|
||||
</TableHeader>
|
||||
@@ -162,9 +171,15 @@ export const TokenRateLimitTable = ({
|
||||
)}
|
||||
<TableCell>{formatPeriod(tokenRateLimit)}</TableCell>
|
||||
<TableCell>
|
||||
{tokenRateLimit.token_budget === null
|
||||
? COST_ONLY_LABEL
|
||||
: `${tokenRateLimit.token_budget} thousand tokens`}
|
||||
{tokenRateLimit.token_budget != null
|
||||
? (tokenRateLimit.token_budget * 1000).toLocaleString() +
|
||||
" tokens"
|
||||
: "—"}
|
||||
</TableCell>
|
||||
<TableCell>
|
||||
{tokenRateLimit.cost_budget_cents != null
|
||||
? "$" + (tokenRateLimit.cost_budget_cents / 100).toFixed(2)
|
||||
: "—"}
|
||||
</TableCell>
|
||||
{isAdmin && (
|
||||
<TableCell>
|
||||
|
||||
@@ -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<void>][] = [
|
||||
["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();
|
||||
});
|
||||
});
|
||||
@@ -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<void> => {
|
||||
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<void> => {
|
||||
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<void> => {
|
||||
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<void> => {
|
||||
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<void> => {
|
||||
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));
|
||||
}
|
||||
};
|
||||
@@ -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<void> => {
|
||||
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<void> => {
|
||||
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 (
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -305,6 +305,18 @@ export interface StreamingError {
|
||||
details?: Record<string, any>;
|
||||
}
|
||||
|
||||
// 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;
|
||||
|
||||
@@ -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(
|
||||
<ErrorBanner
|
||||
error="You've reached the usage budget for your account."
|
||||
errorCode={RATE_LIMITED_ERROR_CODE}
|
||||
details={{ reset_at: "2026-07-28T12:00:30Z" }}
|
||||
/>
|
||||
);
|
||||
|
||||
expect(screen.getByText(/Resets in 1 minute/)).toBeInTheDocument();
|
||||
|
||||
act(() => {
|
||||
jest.advanceTimersByTime(30_000);
|
||||
});
|
||||
|
||||
expect(screen.getByText("You can try again now.")).toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
@@ -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 (
|
||||
<div className="text-red-700 mt-4 text-sm my-auto">
|
||||
<Alert variant="broken">
|
||||
{getErrorIcon(errorCode)}
|
||||
<AlertTitle>{getErrorTitle(errorCode)}</AlertTitle>
|
||||
<AlertDescription className="flex flex-col gap-y-1">
|
||||
<span>{error || "You've reached your usage limit."}</span>
|
||||
{resetLine && (
|
||||
<span className="text-xs text-muted-foreground">{resetLine}</span>
|
||||
)}
|
||||
</AlertDescription>
|
||||
</Alert>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
interface ResubmitProps {
|
||||
resubmit: () => void;
|
||||
@@ -37,6 +129,16 @@ export const ErrorBanner = ({
|
||||
}) => {
|
||||
const [isStackTraceExpanded, setIsStackTraceExpanded] = useState(false);
|
||||
|
||||
if (errorCode === RATE_LIMITED_ERROR_CODE) {
|
||||
return (
|
||||
<RateLimitBanner
|
||||
error={error}
|
||||
errorCode={errorCode}
|
||||
details={(details as RateLimitDetails) ?? {}}
|
||||
/>
|
||||
);
|
||||
}
|
||||
|
||||
return (
|
||||
<div className="text-red-700 mt-4 text-sm my-auto">
|
||||
<Alert variant="broken">
|
||||
|
||||
@@ -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 <Clock className="h-4 w-4" />;
|
||||
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":
|
||||
|
||||
@@ -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}`);
|
||||
}
|
||||
|
||||
|
||||
@@ -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<HTMLInputElement>) => {
|
||||
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({
|
||||
<input
|
||||
ref={inputRef}
|
||||
type="text"
|
||||
inputMode="numeric"
|
||||
pattern="[0-9]*"
|
||||
inputMode={decimalPlaces === 0 ? "numeric" : "decimal"}
|
||||
pattern={inputPattern}
|
||||
disabled={isDisabled}
|
||||
value={inputValue}
|
||||
placeholder={placeholder}
|
||||
|
||||
@@ -32,7 +32,13 @@ function CreateGroupPage() {
|
||||
const [selectedDocSetIds, setSelectedDocSetIds] = useState<number[]>([]);
|
||||
const [selectedAgentIds, setSelectedAgentIds] = useState<number[]>([]);
|
||||
const [tokenLimits, setTokenLimits] = useState<TokenLimit[]>([
|
||||
{ tokenBudget: null, periodDays: null },
|
||||
{
|
||||
tokenId: null,
|
||||
enabled: true,
|
||||
tokenBudget: null,
|
||||
periodDays: null,
|
||||
costBudgetDollars: null,
|
||||
},
|
||||
]);
|
||||
|
||||
const { rows: allRows, isLoading, error } = useGroupMemberCandidates();
|
||||
|
||||
@@ -101,7 +101,13 @@ function EditGroupPage({ groupId }: EditGroupPageProps) {
|
||||
const [selectedDocSetIds, setSelectedDocSetIds] = useState<number[]>([]);
|
||||
const [selectedAgentIds, setSelectedAgentIds] = useState<number[]>([]);
|
||||
const [tokenLimits, setTokenLimits] = useState<TokenLimit[]>([
|
||||
{ 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,
|
||||
}))
|
||||
);
|
||||
}
|
||||
|
||||
@@ -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<TokenLimit[]>([
|
||||
{
|
||||
tokenId: null,
|
||||
enabled: true,
|
||||
tokenBudget: null,
|
||||
periodDays: null,
|
||||
costBudgetDollars: null,
|
||||
},
|
||||
]);
|
||||
|
||||
return <TokenLimitSection limits={limits} onLimitsChange={setLimits} />;
|
||||
}
|
||||
|
||||
test("accepts a cost limit with cents", async () => {
|
||||
const user = setupUser();
|
||||
render(<TokenLimitSectionHarness />);
|
||||
|
||||
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(<TokenLimitSectionHarness />);
|
||||
|
||||
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(<TokenLimitSectionHarness />);
|
||||
|
||||
const tokenLimit = screen.getByPlaceholderText("Token limit (thousands)");
|
||||
await user.type(tokenLimit, "0");
|
||||
|
||||
expect(tokenLimit).toHaveValue("1");
|
||||
});
|
||||
@@ -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
|
||||
</Text>
|
||||
<Text mainUiMuted text03 className="ml-0.5">
|
||||
(thousand tokens)
|
||||
(thousands)
|
||||
</Text>
|
||||
</div>
|
||||
<div className="flex-1 flex items-center min-w-[160px]">
|
||||
<Text mainUiAction text04>
|
||||
Cost Limit
|
||||
</Text>
|
||||
<Text mainUiMuted text03 className="ml-0.5">
|
||||
(USD)
|
||||
</Text>
|
||||
</div>
|
||||
<div className="flex-1 flex items-center min-w-[160px]">
|
||||
@@ -126,8 +151,18 @@ function TokenLimitSection({
|
||||
<InputNumber
|
||||
value={limit.tokenBudget}
|
||||
onChange={(v) => updateLimit(i, "tokenBudget", v)}
|
||||
min={0}
|
||||
placeholder="Token limit in thousands"
|
||||
min={1}
|
||||
placeholder="Token limit (thousands)"
|
||||
/>
|
||||
</div>
|
||||
<div className="flex-1">
|
||||
<InputNumber
|
||||
value={limit.costBudgetDollars}
|
||||
onChange={(v) => updateLimit(i, "costBudgetDollars", v)}
|
||||
min={0.01}
|
||||
step={0.01}
|
||||
decimalPlaces={2}
|
||||
placeholder="Cost limit"
|
||||
/>
|
||||
</div>
|
||||
<div className="flex-1">
|
||||
|
||||
@@ -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"]);
|
||||
});
|
||||
});
|
||||
@@ -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;
|
||||
}
|
||||
@@ -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<void> {
|
||||
// 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" }
|
||||
|
||||
@@ -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();
|
||||
});
|
||||
});
|
||||
@@ -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<void> {
|
||||
await expect(this.humanMessages).toHaveCount(0);
|
||||
}
|
||||
|
||||
async sendUntilUsageLimit(maxTurns: number): Promise<void> {
|
||||
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<void> {
|
||||
await expect(this.usageLimitBanner).toBeVisible();
|
||||
await expect(this.page.getByText(/your account/i)).toBeVisible();
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user